Update app.py

app.py

@@ -1,80 +1,60 @@
 import gradio as gr
-import torch
-import numpy as np
-from PIL import Image
 from transformers import (
-    pipeline,
-    AutoModelForCausalLM,
-    AutoProcessor,
-    AutoTokenizer,
-    GenerationConfig,
-    TextStreamer,
+    pipeline, 
+    AutoProcessor, 
+    AutoModelForCausalLM, 
+    AutoTokenizer, 
+    set_seed
 )
 from datasets import load_dataset
+import torch
+import numpy as np
 
-# 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}")
+# Set seed
+set_seed(42)
 
-# Load image captioning model (BLIP)
-caption_model = pipeline("image-to-text", model="Salesforce/blip-image-captioning-base", device=device)
+# Captioning model
+caption_model = pipeline("image-to-text", model="Salesforce/blip-image-captioning-base")
 
-# Load OCR model (Florence-2-base)
-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)
+# GPT-2 model for context generation
+gpt2_generator = pipeline("text-generation", model="gpt2")
 
-# 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,
-    temperature=0.8,
-    top_p=0.9,
-    repetition_penalty=1.0
-)
+# SpeechT5 for text-to-speech
+synthesiser = pipeline("text-to-speech", model="microsoft/speecht5_tts")
 
-# Load SpeechT5 for TTS
-synthesiser = pipeline("text-to-speech", model="microsoft/speecht5_tts", device=device)
+# Load Florence-2-base for OCR
+ocr_device = "cuda" if torch.cuda.is_available() else "cpu"
+ocr_dtype = torch.float16 if torch.cuda.is_available() else torch.float32
+ocr_model = AutoModelForCausalLM.from_pretrained("microsoft/Florence-2-base", torch_dtype=ocr_dtype, trust_remote_code=True).to(ocr_device)
+ocr_processor = AutoProcessor.from_pretrained("microsoft/Florence-2-base", trust_remote_code=True)
 
-# Use known compatible 600-dim speaker embedding
+# Load 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]
+speaker_embedding = torch.tensor(embeddings_dataset[7306]["xvector"]).unsqueeze(0)
 
 def process_image(image):
     try:
-        # Caption generation
+        # Generate caption
         caption = caption_model(image)[0]['generated_text']
 
-        # 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"],
+        # Extract OCR text
+        inputs = ocr_processor(text="<OCR>", images=image, return_tensors="pt").to(ocr_device, ocr_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(ocr_outputs, skip_special_tokens=True)[0]
+        extracted_text = ocr_processor.batch_decode(generated_ids, skip_special_tokens=True)[0]
 
-        # 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}]
-        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)
+        # Generate context with GPT-2
+        prompt = f"Determine the context of this image based on the caption and extracted text. Caption: {caption}. Extracted text: {extracted_text}. Context:"
+        context_output = gpt2_generator(prompt, max_length=100, num_return_sequences=1)
+        context = context_output[0]['generated_text']
 
-        # Convert context to speech
-        speech = synthesiser(
-            context,
-            forward_params={"speaker_embeddings": speaker_embedding}
-        )
+        # Text-to-speech
+        speech = synthesiser(context, forward_params={"speaker_embeddings": speaker_embedding})
         audio = np.array(speech["audio"])
         rate = speech["sampling_rate"]
 
@@ -93,8 +73,8 @@ iface = gr.Interface(
         gr.Textbox(label="Extracted Text (OCR)"),
         gr.Textbox(label="Generated Context")
     ],
-    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."
+    title="SeeSay Contextualizer",
+    description="Upload an image to generate a caption, extract text, create audio from context, and determine the context using GPT-2 and Florence-2-base."
 )
 
-iface.launch(share=True)
+iface.launch()