Update app.py

app.py

@@ -1,32 +1,37 @@
 import gradio as gr
-from transformers import (
-    pipeline, AutoModelForCausalLM, AutoProcessor, AutoTokenizer,
-    GenerationConfig, TextStreamer
-)
-from datasets import load_dataset
 import torch
 import numpy as np
 from PIL import Image
+from transformers import (
+    pipeline,
+    AutoModelForCausalLM,
+    AutoProcessor,
+    AutoTokenizer,
+    GenerationConfig,
+    TextStreamer,
+)
+from datasets import load_dataset
 
-# 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
+# Use CPU if no GPU is available
+device = "cuda" if torch.cuda.is_available() else "cpu"
+dtype = torch.float16 if torch.cuda.is_available() else torch.float32
+print(f"Device set to use {device}")
 
-# Caption model (BLIP)
-caption_model = pipeline("image-to-text", model="Salesforce/blip-image-captioning-base")
+# Load image captioning model (BLIP)
+caption_model = pipeline("image-to-text", model="Salesforce/blip-image-captioning-base", device=device)
 
-# Florence-2-base model for OCR
+# Load OCR model (Florence-2-base)
 ocr_model = AutoModelForCausalLM.from_pretrained(
-    "microsoft/Florence-2-base",
-    torch_dtype=torch_dtype,
-    trust_remote_code=True
+    "microsoft/Florence-2-base", trust_remote_code=True, torch_dtype=dtype
 ).to(device)
 ocr_processor = AutoProcessor.from_pretrained("microsoft/Florence-2-base", trust_remote_code=True)
 
-# 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(
+# Load SmallDoge 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,
@@ -34,50 +39,42 @@ generation_config = GenerationConfig(
     top_p=0.9,
     repetition_penalty=1.0
 )
-streamer = TextStreamer(tokenizer=context_tokenizer, skip_prompt=True)
 
-# SpeechT5 for TTS
-tts = pipeline("text-to-speech", model="microsoft/speecht5_tts")
+# Load SpeechT5 for TTS
+synthesiser = pipeline("text-to-speech", model="microsoft/speecht5_tts", device=device)
 
-# 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.")
+# Use known compatible 600-dim speaker embedding
+embeddings_dataset = load_dataset("Matthijs/cmu-arctic-xvectors", split="validation")
+speaker_embedding = torch.tensor(embeddings_dataset[7306]["xvector"]).unsqueeze(0)  # Shape: [1, 600]
 
 def process_image(image):
     try:
-        # Generate caption
+        # Caption generation
         caption = caption_model(image)[0]['generated_text']
 
-        # 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"],
-            pixel_values=inputs["pixel_values"],
+        # OCR extraction
+        ocr_inputs = ocr_processor(text="<OCR>", images=image, return_tensors="pt").to(device)
+        ocr_outputs = ocr_model.generate(
+            input_ids=ocr_inputs["input_ids"],
+            pixel_values=ocr_inputs["pixel_values"],
             max_new_tokens=1024,
-            do_sample=False,
             num_beams=3,
+            do_sample=False,
         )
-        extracted_text = ocr_processor.batch_decode(generated_ids, skip_special_tokens=True)[0]
+        extracted_text = ocr_processor.batch_decode(ocr_outputs, skip_special_tokens=True)[0]
 
-        # 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:"
+        # Context generation using Doge
+        prompt = f"Determine the context of this image based on the caption and extracted text.\nCaption: {caption}\nExtracted text: {extracted_text}\nContext:"
         conversation = [{"role": "user", "content": prompt}]
-        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 = context_tokenizer.decode(output[0], skip_special_tokens=True)
+        doge_inputs = doge_tokenizer.apply_chat_template(conversation, tokenize=True, return_tensors="pt").to(device)
+        doge_output = doge_model.generate(doge_inputs, generation_config=doge_config)
+        context = doge_tokenizer.decode(doge_output[0], skip_special_tokens=True)
 
         # Convert context to speech
-        speech = tts(context, forward_params={"speaker_embeddings": speaker_embedding})
+        speech = synthesiser(
+            context,
+            forward_params={"speaker_embeddings": speaker_embedding}
+        )
         audio = np.array(speech["audio"])
         rate = speech["sampling_rate"]
 
@@ -86,7 +83,7 @@ def process_image(image):
     except Exception as e:
         return None, f"Error: {str(e)}", "", ""
 
-# Gradio Interface
+# Gradio UI
 iface = gr.Interface(
     fn=process_image,
     inputs=gr.Image(type='pil', label="Upload an Image"),
@@ -96,8 +93,8 @@ iface = gr.Interface(
         gr.Textbox(label="Extracted Text (OCR)"),
         gr.Textbox(label="Generated Context")
     ],
-    title="SeeSay Contextualizer",
-    description="Upload an image to generate a caption, extract text, generate context with Doge, and convert to speech."
+    title="SeeSay Contextualizer with Doge + SpeechT5",
+    description="Upload an image to generate a caption, extract OCR text, determine context with Doge-320M, and hear it with SpeechT5."
 )
 
 iface.launch(share=True)