Update app.py

app.py

@@ -1,77 +1,82 @@
 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,
+    TextStreamer,
+    set_seed
 )
+import torch
+import numpy as np
+import requests
+import io
 from datasets import load_dataset
 
-# Set device and dtype
+# Set device
 device = "cuda" if torch.cuda.is_available() else "cpu"
-dtype = torch.float16 if torch.cuda.is_available() else torch.float32
+torch_dtype = torch.float16 if torch.cuda.is_available() else torch.float32
+set_seed(42)
+
+# 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 BLIP for image captioning
-caption_model = pipeline("image-to-text", model="Salesforce/blip-image-captioning-base")
+# Text-to-Speech (SpeechT5)
+synthesiser = pipeline("text-to-speech", model="microsoft/speecht5_tts", device=0 if torch.cuda.is_available() else -1)
 
-# Load Florence-2-base for OCR
-ocr_model = AutoModelForCausalLM.from_pretrained(
-    "microsoft/Florence-2-base", trust_remote_code=True, torch_dtype=dtype
-).to(device)
+# OCR using Florence-2-base
+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")
 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,
     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
+# 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)).unsqueeze(0)  # Shape: [1, 600]
-
+speaker_embedding = torch.tensor(np.load(buffer, allow_pickle=True)).unsqueeze(0)  # Shape: [1, 600]
+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:
-        # Captioning
-        caption = caption_model(image)[0]["generated_text"]
+        # Generate caption
+        caption = caption_model(image)[0]['generated_text']
 
-        # OCR
-        inputs = ocr_processor(text="<OCR>", images=image, return_tensors="pt").to(device, dtype)
+        # OCR extraction
+        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=1024,
-            num_beams=3,
+            max_new_tokens=512,
             do_sample=False,
+            num_beams=3
         )
         extracted_text = ocr_processor.batch_decode(generated_ids, skip_special_tokens=True)[0]
 
-        # Context generation with Doge
+        # 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:"
         conversation = [{"role": "user", "content": prompt}]
-        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)
+        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)
 
-        # TTS
+        # Text-to-Speech
         speech = synthesiser(
             context,
             forward_params={"speaker_embeddings": speaker_embedding}
@@ -84,19 +89,18 @@ 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"),
+    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 Contextualizer",
-    description="Upload an image to generate a caption, extract text (OCR), generate context using Doge, and convert to audio with SpeechT5.",
+    description="Upload an image to generate a caption, extract text with Florence-2-base, contextualize with Doge-320M-Instruct, and hear it with SpeechT5."
 )
 
-iface.launch()
+iface.launch(share=True)