Update app.py

app.py

@@ -1,38 +1,42 @@
 import gradio as gr
-from transformers import pipeline, AutoProcessor, AutoModelForCausalLM, AutoTokenizer, set_seed
-from datasets import load_dataset
+from transformers import pipeline, AutoProcessor, AutoModelForCausalLM
 import torch
 import numpy as np
+from datasets import load_dataset
+from PIL import Image
 
-# Set seed for reproducibility
-set_seed(42)
-
-# Load BLIP model for image captioning
+# 1) IMAGE CAPTIONING MODEL
 caption_model = pipeline("image-to-text", model="Salesforce/blip-image-captioning-base")
 
-# Load SpeechT5 model for text-to-speech
-synthesiser = pipeline("text-to-speech", model="microsoft/speecht5_tts")
-
-# Load Florence-2 model for OCR
+# 2) OCR MODEL (Florence-2)
 ocr_device = "cuda:0" 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-large", torch_dtype=ocr_dtype, trust_remote_code=True).to(ocr_device)
+ocr_model = AutoModelForCausalLM.from_pretrained(
+    "microsoft/Florence-2-large",
+    torch_dtype=ocr_dtype,
+    trust_remote_code=True
+).to(ocr_device)
 ocr_processor = AutoProcessor.from_pretrained("microsoft/Florence-2-large", trust_remote_code=True)
 
-# Load GPT-2 (124M) model for text generation
-gpt2_generator = pipeline('text-generation', model='gpt2')
+# 3) QUESTION-ANSWERING MODEL
+qa_model = pipeline(
+    "question-answering",
+    model="timpal0l/mdeberta-v3-base-squad2"
+)
+
+# 4) TEXT-TO-SPEECH MODEL
+tts_pipeline = pipeline("text-to-speech", model="microsoft/speecht5_tts")
 
 # Load speaker embedding
 embeddings_dataset = load_dataset("Matthijs/cmu-arctic-xvectors", split="validation")
 speaker_embedding = torch.tensor(embeddings_dataset[7306]["xvector"]).unsqueeze(0)
 
-
 def process_image(image):
     try:
-        # Generate caption from the image
+        # 1) Generate caption from the image
         caption = caption_model(image)[0]['generated_text']
 
-        # Extract text (OCR) using Florence-2
+        # 2) Extract text from the image using Florence-2
         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"],
@@ -43,28 +47,32 @@ def process_image(image):
         )
         extracted_text = ocr_processor.batch_decode(generated_ids, skip_special_tokens=True)[0]
 
-        # Generate context using GPT-2 (124M)
-        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']
+        # 3) Use QA model to derive context from caption + extracted text
+        # We treat the "context" string as the knowledge base and ask a question about it
+        question = "What is the context of this image?"
+        combined_context = f"Caption: {caption}\nExtracted Text: {extracted_text}"
+        qa_result = qa_model(question=question, context=combined_context)
 
-        # Convert context to speech
-        speech = synthesiser(
-            context,
+        # The QA model returns an extracted "answer" from the combined context
+        # If the model can't find a direct span, it may return an empty string or a short phrase
+        final_context = qa_result["answer"]
+
+        # 4) Convert the final context to speech
+        speech_data = tts_pipeline(
+            final_context,
             forward_params={"speaker_embeddings": speaker_embedding}
         )
 
-        # Prepare audio data
-        audio = np.array(speech["audio"])
-        rate = speech["sampling_rate"]
+        # Prepare audio data for Gradio
+        audio = np.array(speech_data["audio"])
+        rate = speech_data["sampling_rate"]
 
-        # Return audio, caption, extracted text, and context
-        return (rate, audio), caption, extracted_text, context
+        # Return audio, caption, extracted text, and final context
+        return (rate, audio), caption, extracted_text, final_context
 
     except Exception as e:
         return None, f"Error: {str(e)}", "", ""
 
-
 # Gradio Interface
 iface = gr.Interface(
     fn=process_image,
@@ -73,10 +81,15 @@ iface = gr.Interface(
         gr.Audio(label="Generated Audio"),
         gr.Textbox(label="Generated Caption"),
         gr.Textbox(label="Extracted Text (OCR)"),
-        gr.Textbox(label="Generated Context")
+        gr.Textbox(label="QA-derived Context")
     ],
-    title="SeeSay Contextualizer with GPT-2 (124M)",
-    description="Upload an image to generate a caption, extract text, create audio from context, and determine the context using GPT-2."
+    title="Contextual Image QA with SpeechT5",
+    description=(
+        "1) Generate a caption via BLIP. "
+        "2) Extract text using Florence-2. "
+        "3) Use QA with mDeBERTa to find a 'context' from caption + text. "
+        "4) Convert it to audio via SpeechT5."
+    ),
 )
 
 iface.launch()