Update app.py

app.py

@@ -1,79 +1,81 @@
 import gradio as gr
+import torch
+import numpy as np
+import requests
+import io
+from PIL import Image
 from transformers import (
     pipeline,
     AutoProcessor,
     AutoModelForCausalLM,
     AutoTokenizer,
-    GenerationConfig
+    GenerationConfig,
 )
-import torch
-import numpy as np
-from PIL import Image
-import requests
-import io
+from datasets import load_dataset
 
-# Device setup
-device = "cuda:0" if torch.cuda.is_available() else "cpu"
-torch_dtype = torch.float16 if torch.cuda.is_available() else torch.float32
+# Set device and dtype
+device = "cuda" if torch.cuda.is_available() else "cpu"
+dtype = torch.float16 if torch.cuda.is_available() else torch.float32
 
-# Load image captioning model (BLIP)
+# Load BLIP for image captioning
 caption_model = pipeline("image-to-text", model="Salesforce/blip-image-captioning-base")
 
-# Load TTS model (SpeechT5)
-synthesiser = pipeline("text-to-speech", model="microsoft/speecht5_tts")
-
 # 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_model = AutoModelForCausalLM.from_pretrained(
+    "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)
 
-# Load Doge-320M-Instruct for contextual reasoning
+# Load SmallDoge 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)
 
-generation_config = GenerationConfig(
+doge_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 SpeechT5 for TTS
+synthesiser = pipeline("text-to-speech", model="microsoft/speecht5_tts")
+
+# Load speaker embedding from .npy using BytesIO
+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)).unsqueeze(0)  # Shape: [1, 600]
+
+
 def process_image(image):
     try:
         # Captioning
-        caption = caption_model(image)[0]['generated_text']
+        caption = caption_model(image)[0]["generated_text"]
 
         # OCR
-        inputs = ocr_processor(text="<OCR>", images=image, return_tensors="pt").to(device, torch_dtype)
+        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=1024,
             num_beams=3,
-            do_sample=False
+            do_sample=False,
         )
         extracted_text = ocr_processor.batch_decode(generated_ids, skip_special_tokens=True)[0]
 
-        # Context generation
-        prompt = f"Determine the context of this image based on the caption and extracted text. Caption: {caption}. Extracted text: {extracted_text}. Context:"
+        # Context generation with 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}]
-        doge_inputs = doge_tokenizer.apply_chat_template(
-            conversation=conversation,
-            tokenize=True,
-            return_tensors="pt"
-        ).to(device)
-
-        output_ids = doge_model.generate(
-            doge_inputs,
-            generation_config=generation_config
-        )
-
+        inputs = doge_tokenizer.apply_chat_template(conversation, tokenize=True, return_tensors="pt").to(device)
+        output_ids = doge_model.generate(inputs, generation_config=doge_config)
         context = doge_tokenizer.decode(output_ids[0], skip_special_tokens=True)
 
-        # Text-to-speech (no speaker embedding required)
-        speech = synthesiser(context)
-
+        # TTS
+        speech = synthesiser(
+            context,
+            forward_params={"speaker_embeddings": speaker_embedding}
+        )
         audio = np.array(speech["audio"])
         rate = speech["sampling_rate"]
 
@@ -82,18 +84,19 @@ 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"),
+    inputs=gr.Image(type="pil", label="Upload an Image"),
     outputs=[
         gr.Audio(label="Generated Audio"),
         gr.Textbox(label="Generated Caption"),
         gr.Textbox(label="Extracted Text (OCR)"),
-        gr.Textbox(label="Generated Context")
+        gr.Textbox(label="Generated Context"),
     ],
-    title="SeeSay",
-    description="Upload an image to generate a caption, extract text, convert context to speech, and understand the image using Doge-320M."
+    title="SeeSay Contextualizer",
+    description="Upload an image to generate a caption, extract text (OCR), generate context using Doge, and convert to audio with SpeechT5.",
 )
 
-iface.launch(share=True)
+iface.launch()