Update app.py

app.py

@@ -1,86 +1,83 @@
 import gradio as gr
+import torch
+import numpy as np
 from transformers import (
     pipeline,
-    AutoProcessor,
     AutoModelForCausalLM,
+    AutoProcessor,
     AutoTokenizer,
     GenerationConfig,
-    TextStreamer,
-    set_seed
 )
-import torch
-import numpy as np
-import requests
-import io
 from datasets import load_dataset
+from PIL import Image
 
-# Set device
+# Set device and dtype
 device = "cuda" if torch.cuda.is_available() else "cpu"
 torch_dtype = torch.float16 if torch.cuda.is_available() else torch.float32
-set_seed(42)
+print(f"Device set to use {device}")
 
-# Image Captioning (BLIP)
-caption_model = pipeline("image-to-text", model="Salesforce/blip-image-captioning-base", device=0 if torch.cuda.is_available() else -1)
+# Load image captioning model
+caption_model = pipeline("image-to-text", model="Salesforce/blip-image-captioning-base", device=device)
 
-# Text-to-Speech (SpeechT5)
-synthesiser = pipeline("text-to-speech", model="microsoft/speecht5_tts", device=0 if torch.cuda.is_available() else -1)
+# Load text-to-speech model
+synthesiser = pipeline("text-to-speech", model="microsoft/speecht5_tts")
 
-# OCR using Florence-2-base
-ocr_model = AutoModelForCausalLM.from_pretrained("microsoft/Florence-2-base", torch_dtype=torch_dtype, trust_remote_code=True).to(device)
+# Load 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)
 
-# Load SmallDoge for context generation
-doge_tokenizer = AutoTokenizer.from_pretrained("SmallDoge/Doge-320M-Instruct")
+# 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_config = GenerationConfig(
+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,
 )
 
-# Load speaker embedding (600-dim)
-SPEAKER_EMBEDDING_URL = "https://huggingface.co/datasets/Matthijs/cmu-arctic-xvectors/resolve/main/spkemb/fn0012.npy"
-response = requests.get(SPEAKER_EMBEDDING_URL)
-buffer = io.BytesIO(response.content)
-speaker_embedding = torch.tensor(np.load(buffer, allow_pickle=True)).unsqueeze(0)  # Shape: [1, 600]
+# 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.")
 
-# Main function
+
 def process_image(image):
     try:
-        # Generate caption
+        # Generate image caption
         caption = caption_model(image)[0]['generated_text']
 
-        # OCR extraction
+        # Run OCR using Florence-2-base
         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"],
-            max_new_tokens=512,
+            max_new_tokens=1024,
             do_sample=False,
-            num_beams=3
+            num_beams=3,
         )
         extracted_text = ocr_processor.batch_decode(generated_ids, skip_special_tokens=True)[0]
 
-        # Generate context using SmallDoge
-        prompt = f"Determine the context of this image based on the caption and extracted text.\nCaption: {caption}\nExtracted text: {extracted_text}\nContext:"
+        # Generate context using Doge model
+        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}]
-        inputs = doge_tokenizer.apply_chat_template(
+        doge_inputs = doge_tokenizer.apply_chat_template(
             conversation=conversation, tokenize=True, return_tensors="pt"
         ).to(device)
 
-        outputs = doge_model.generate(inputs, generation_config=doge_config)
-        context = doge_tokenizer.decode(outputs[0], skip_special_tokens=True)
-
-        # Text-to-Speech
-        speech = synthesiser(
-            context,
-            forward_params={"speaker_embeddings": speaker_embedding}
+        doge_outputs = doge_model.generate(
+            input_ids=doge_inputs,
+            generation_config=generation_config,
         )
+        context = doge_tokenizer.decode(doge_outputs[0], skip_special_tokens=True)
+
+        # Generate speech from context
+        speech = synthesiser(context, forward_params={"speaker_embeddings": speaker_embedding})
         audio = np.array(speech["audio"])
         rate = speech["sampling_rate"]
 
@@ -89,6 +86,7 @@ def process_image(image):
     except Exception as e:
         return None, f"Error: {str(e)}", "", ""
 
+
 # Gradio UI
 iface = gr.Interface(
     fn=process_image,
@@ -99,8 +97,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 with Florence-2-base, contextualize with Doge-320M-Instruct, and hear it with SpeechT5."
+    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."
 )
 
 iface.launch(share=True)