Update app.py

app.py

@@ -1,35 +1,38 @@
 import gradio as gr
 from transformers import (
     pipeline,
-    AutoTokenizer,
+    AutoProcessor,
     AutoModelForCausalLM,
+    AutoTokenizer,
     GenerationConfig,
-    set_seed
+    TextStreamer
 )
+from datasets import load_dataset
 import torch
 import numpy as np
-import pytesseract
-from PIL import Image
-from datasets import load_dataset
-
-set_seed(42)
 
-# Device
-device = "cuda" if torch.cuda.is_available() else "cpu"
+# Set device and dtype
+device = "cuda:0" if torch.cuda.is_available() else "cpu"
+torch_dtype = torch.float16 if torch.cuda.is_available() else torch.float32
 
-# Image Captioning (BLIP)
+# Image Captioning
 caption_model = pipeline("image-to-text", model="Salesforce/blip-image-captioning-base")
 
-# Text-to-Speech without speaker embeddings
+# Text-to-Speech
 synthesiser = pipeline("text-to-speech", model="microsoft/speecht5_tts")
 
-# Doge-320M-Instruct for Context Generation
-doge_tokenizer = AutoTokenizer.from_pretrained("SmallDoge/Doge-320M-Instruct")
-doge_model = AutoModelForCausalLM.from_pretrained(
-    "SmallDoge/Doge-320M-Instruct", trust_remote_code=True
+# Florence-2-base for OCR
+ocr_model = AutoModelForCausalLM.from_pretrained(
+    "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)
 
-doge_generation_config = GenerationConfig(
+# Doge model for context generation
+doge_tokenizer = AutoTokenizer.from_pretrained("SmallDoge/Doge-320M-Instruct")
+doge_model = AutoModelForCausalLM.from_pretrained("SmallDoge/Doge-320M-Instruct", trust_remote_code=True).to(device)
+doge_config = GenerationConfig(
     max_new_tokens=100,
     use_cache=True,
     do_sample=True,
@@ -38,22 +41,35 @@ doge_generation_config = GenerationConfig(
     repetition_penalty=1.0
 )
 
-def extract_text_with_tesseract(image):
-    return pytesseract.image_to_string(image)
+# Speaker embedding (600-dim)
+embeddings_dataset = load_dataset("Matthijs/cmu-arctic-xvectors", split="validation")
+embedding = None
+for entry in embeddings_dataset:
+    vector = torch.tensor(entry["xvector"]).unsqueeze(0)
+    if vector.shape[1] >= 600:
+        embedding = vector[:, :600]
+        break
+if embedding is None:
+    raise ValueError("No suitable speaker embedding of at least 600 dimensions found.")
 
 def process_image(image):
     try:
-        # 1. Caption
+        # Caption
         caption = caption_model(image)[0]['generated_text']
 
-        # 2. OCR
-        extracted_text = extract_text_with_tesseract(image)
-
-        # 3. Context with Doge (truncate input)
-        prompt = (
-            f"Determine the context of this image.\n"
-            f"Caption: {caption[:200]}\nExtracted text: {extracted_text[:200]}\nContext:"
+        # OCR
+        ocr_inputs = ocr_processor(text="<OCR>", images=image, return_tensors="pt").to(device, torch_dtype)
+        generated_ids = ocr_model.generate(
+            input_ids=ocr_inputs["input_ids"],
+            pixel_values=ocr_inputs["pixel_values"],
+            max_new_tokens=1024,
+            num_beams=3,
+            do_sample=False
         )
+        extracted_text = ocr_processor.batch_decode(generated_ids, skip_special_tokens=True)[0]
+
+        # Doge context generation
+        prompt = f"Determine the context of this image. Caption: {caption} Extracted text: {extracted_text}"
         conversation = [{"role": "user", "content": prompt}]
         doge_inputs = doge_tokenizer.apply_chat_template(
             conversation=conversation,
@@ -61,14 +77,17 @@ def process_image(image):
             return_tensors="pt"
         ).to(device)
 
-        doge_output = doge_model.generate(
-            input_ids=doge_inputs,
-            generation_config=doge_generation_config
+        outputs = doge_model.generate(
+            doge_inputs,
+            generation_config=doge_config
         )
-        context = doge_tokenizer.decode(doge_output[0], skip_special_tokens=True).strip()
+        context = doge_tokenizer.decode(outputs[0], skip_special_tokens=True)
 
-        # 4. Text-to-Speech (no embeddings)
-        speech = synthesiser(context)
+        # TTS
+        speech = synthesiser(
+            context,
+            forward_params={"speaker_embeddings": embedding}
+        )
         audio = np.array(speech["audio"])
         rate = speech["sampling_rate"]
 
@@ -77,6 +96,7 @@ def process_image(image):
     except Exception as e:
         return None, f"Error: {str(e)}", "", ""
 
+# Gradio Interface
 iface = gr.Interface(
     fn=process_image,
     inputs=gr.Image(type='pil', label="Upload an Image"),
@@ -86,8 +106,8 @@ iface = gr.Interface(
         gr.Textbox(label="Extracted Text (OCR)"),
         gr.Textbox(label="Generated Context")
     ],
-    title="SeeSay Contextualizer (Optimized)",
-    description="Upload an image to generate a caption, extract text (OCR), generate context, and hear it spoken."
+    title="SeeSay Contextualizer with Doge & BLIP",
+    description="Upload an image to generate a caption, extract text, determine context, and convert it to audio."
 )
 
-iface.launch(share=True)
+iface.launch()