Update app.py

app.py

@@ -1,58 +1,61 @@
 import gradio as gr
-import torch
-import numpy as np
 from transformers import (
-    pipeline,
-    AutoModelForCausalLM,
-    AutoProcessor,
-    AutoTokenizer,
-    GenerationConfig,
+    pipeline, AutoModelForCausalLM, AutoProcessor, AutoTokenizer,
+    GenerationConfig, TextStreamer
 )
 from datasets import load_dataset
+import torch
+import numpy as np
 from PIL import Image
 
-# Set device and dtype
-device = "cuda" if torch.cuda.is_available() else "cpu"
+# Device and dtype setup
+device = "cuda:0" if torch.cuda.is_available() else "cpu"
 torch_dtype = torch.float16 if torch.cuda.is_available() else torch.float32
-print(f"Device set to use {device}")
-
-# Load image captioning model
-caption_model = pipeline("image-to-text", model="Salesforce/blip-image-captioning-base", device=device)
 
-# Load text-to-speech model
-synthesiser = pipeline("text-to-speech", model="microsoft/speecht5_tts")
+# Caption model (BLIP)
+caption_model = pipeline("image-to-text", model="Salesforce/blip-image-captioning-base")
 
-# Load Florence-2-base for OCR
+# Florence-2-base model for OCR
 ocr_model = AutoModelForCausalLM.from_pretrained(
-    "microsoft/Florence-2-base", torch_dtype=torch_dtype, trust_remote_code=True
+    "microsoft/Florence-2-base",
+    torch_dtype=torch_dtype,
+    trust_remote_code=True
 ).to(device)
 ocr_processor = AutoProcessor.from_pretrained("microsoft/Florence-2-base", trust_remote_code=True)
 
-# Load Doge model for context generation
-doge_tokenizer = AutoTokenizer.from_pretrained("SmallDoge/Doge-320M-Instruct", trust_remote_code=True)
-doge_model = AutoModelForCausalLM.from_pretrained("SmallDoge/Doge-320M-Instruct", trust_remote_code=True).to(device)
+# Doge model for context generation
+context_tokenizer = AutoTokenizer.from_pretrained("SmallDoge/Doge-320M-Instruct")
+context_model = AutoModelForCausalLM.from_pretrained("SmallDoge/Doge-320M-Instruct", trust_remote_code=True).to(device)
 generation_config = GenerationConfig(
     max_new_tokens=100,
     use_cache=True,
     do_sample=True,
     temperature=0.8,
     top_p=0.9,
-    repetition_penalty=1.0,
+    repetition_penalty=1.0
 )
+streamer = TextStreamer(tokenizer=context_tokenizer, skip_prompt=True)
 
-# Load speaker embedding
-embedding_dataset = load_dataset("Matthijs/cmu-arctic-xvectors", split="validation")
-speaker_embedding = torch.tensor(embedding_dataset[7306]["xvector"]).unsqueeze(0)
-if speaker_embedding.shape[1] < 600:
-    raise ValueError("No suitable speaker embedding of at least 600 dimensions found.")
+# SpeechT5 for TTS
+tts = pipeline("text-to-speech", model="microsoft/speecht5_tts")
 
+# Load valid 600-dim speaker embedding
+speaker_embedding = None
+embedding_dataset = load_dataset("Matthijs/cmu-arctic-xvectors", split="validation")
+for item in embedding_dataset:
+    vec = torch.tensor(item["xvector"])
+    if vec.shape[0] == 600:
+        speaker_embedding = vec.unsqueeze(0)
+        break
+if speaker_embedding is None:
+    raise ValueError("No suitable speaker embedding of 600 dimensions found.")
 
 def process_image(image):
     try:
-        # Generate image caption
+        # Generate caption
         caption = caption_model(image)[0]['generated_text']
 
-        # Run OCR using Florence-2-base
+        # Extract text using Florence-2
         inputs = ocr_processor(text="<OCR>", images=image, return_tensors="pt").to(device, torch_dtype)
         generated_ids = ocr_model.generate(
             input_ids=inputs["input_ids"],
@@ -63,21 +66,18 @@ def process_image(image):
         )
         extracted_text = ocr_processor.batch_decode(generated_ids, skip_special_tokens=True)[0]
 
-        # Generate context using Doge model
+        # Generate context using Doge
         prompt = f"Determine the context of this image based on the caption and extracted text. Caption: {caption}. Extracted text: {extracted_text}. Context:"
         conversation = [{"role": "user", "content": prompt}]
-        doge_inputs = doge_tokenizer.apply_chat_template(
-            conversation=conversation, tokenize=True, return_tensors="pt"
-        ).to(device)
-
-        doge_outputs = doge_model.generate(
-            input_ids=doge_inputs,
+        context_inputs = context_tokenizer.apply_chat_template(conversation=conversation, tokenize=True, return_tensors="pt").to(device)
+        output = context_model.generate(
+            context_inputs,
             generation_config=generation_config,
         )
-        context = doge_tokenizer.decode(doge_outputs[0], skip_special_tokens=True)
+        context = context_tokenizer.decode(output[0], skip_special_tokens=True)
 
-        # Generate speech from context
-        speech = synthesiser(context, forward_params={"speaker_embeddings": speaker_embedding})
+        # Convert context to speech
+        speech = tts(context, forward_params={"speaker_embeddings": speaker_embedding})
         audio = np.array(speech["audio"])
         rate = speech["sampling_rate"]
 
@@ -86,8 +86,7 @@ def process_image(image):
     except Exception as e:
         return None, f"Error: {str(e)}", "", ""
 
-
-# Gradio UI
+# Gradio Interface
 iface = gr.Interface(
     fn=process_image,
     inputs=gr.Image(type='pil', label="Upload an Image"),
@@ -97,8 +96,8 @@ iface = gr.Interface(
         gr.Textbox(label="Extracted Text (OCR)"),
         gr.Textbox(label="Generated Context")
     ],
-    title="SeeSay Contextualizer with Doge and Florence-2",
-    description="Upload an image to generate a caption, extract text, determine the context using Doge, and convert context to speech."
+    title="SeeSay Contextualizer",
+    description="Upload an image to generate a caption, extract text, generate context with Doge, and convert to speech."
 )
 
 iface.launch(share=True)