Update app.py

app.py

@@ -1,39 +1,34 @@
 import gradio as gr
 from transformers import (
     pipeline,
-    AutoProcessor,
-    AutoModelForCausalLM,
     AutoTokenizer,
+    AutoModelForCausalLM,
     GenerationConfig,
     set_seed
 )
-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 and dtype
+# Device
 device = "cuda" if torch.cuda.is_available() else "cpu"
-dtype = torch.float16 if torch.cuda.is_available() else torch.float32
 
-# Load image captioning model
+# Image Captioning (BLIP)
 caption_model = pipeline("image-to-text", model="Salesforce/blip-image-captioning-base")
 
-# Load SpeechT5 text-to-speech model
+# Text-to-Speech without speaker embeddings
 synthesiser = pipeline("text-to-speech", model="microsoft/speecht5_tts")
 
-# Load OCR model (Florence-2)
-ocr_model = AutoModelForCausalLM.from_pretrained(
-    "microsoft/Florence-2-large", torch_dtype=dtype, trust_remote_code=True
-).to(device)
-ocr_processor = AutoProcessor.from_pretrained("microsoft/Florence-2-large", trust_remote_code=True)
-
-# Load Doge-320M-Instruct for context generation
+# 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
 ).to(device)
+
 doge_generation_config = GenerationConfig(
     max_new_tokens=100,
     use_cache=True,
@@ -43,44 +38,22 @@ doge_generation_config = GenerationConfig(
     repetition_penalty=1.0
 )
 
-# Load speaker embedding with fallback
-speaker_embedding = None
-embedding_data = load_dataset("Matthijs/cmu-arctic-xvectors", split="validation")
-
-for entry in embedding_data:
-    vec = entry["xvector"]
-    if len(vec) >= 600:
-        speaker_embedding = torch.tensor(vec[:600], dtype=torch.float32).unsqueeze(0)
-        break
-
-# Fallback: use a zero vector if none found
-if speaker_embedding is None:
-    print("⚠️ No suitable speaker embedding found. Using default 600-dim zero vector.")
-    speaker_embedding = torch.zeros(1, 600, dtype=torch.float32)
-
-# Ensure correct shape
-assert speaker_embedding.shape == (1, 600), f"Expected shape (1, 600), got {speaker_embedding.shape}"
-
+def extract_text_with_tesseract(image):
+    return pytesseract.image_to_string(image)
 
 def process_image(image):
     try:
-        # 1. Caption the image
+        # 1. Caption
         caption = caption_model(image)[0]['generated_text']
 
-        # 2. OCR with Florence-2
-        inputs = ocr_processor(text="<OCR>", images=image, return_tensors="pt").to(device, dtype)
-        generated_ids = ocr_model.generate(
-            input_ids=inputs["input_ids"],
-            pixel_values=inputs["pixel_values"],
-            max_new_tokens=4096,
-            num_beams=3,
-            do_sample=False
-        )
-        extracted_text = ocr_processor.batch_decode(generated_ids, skip_special_tokens=True)[0]
+        # 2. OCR
+        extracted_text = extract_text_with_tesseract(image)
 
-        # 3. Prompt Doge model for context generation
-        prompt = f"Determine the context of this image based on the caption and extracted text.\nCaption: {caption}\nExtracted text: {extracted_text}\nContext:"
-        prompt = prompt[:600]  # Prevent tensor mismatch error
+        # 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:"
+        )
         conversation = [{"role": "user", "content": prompt}]
         doge_inputs = doge_tokenizer.apply_chat_template(
             conversation=conversation,
@@ -94,12 +67,8 @@ def process_image(image):
         )
         context = doge_tokenizer.decode(doge_output[0], skip_special_tokens=True).strip()
 
-        # 4. Convert context to speech
-        speech = synthesiser(
-            context,
-            forward_params={"speaker_embeddings": speaker_embedding}
-        )
-
+        # 4. Text-to-Speech (no embeddings)
+        speech = synthesiser(context)
         audio = np.array(speech["audio"])
         rate = speech["sampling_rate"]
 
@@ -108,8 +77,6 @@ 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"),
@@ -119,8 +86,8 @@ iface = gr.Interface(
         gr.Textbox(label="Extracted Text (OCR)"),
         gr.Textbox(label="Generated Context")
     ],
-    title="SeeSay Contextualizer with Doge-320M",
-    description="Upload an image to caption it, extract text, generate context, and hear the result as speech."
+    title="SeeSay Contextualizer (Optimized)",
+    description="Upload an image to generate a caption, extract text (OCR), generate context, and hear it spoken."
 )
 
 iface.launch(share=True)