Update app.py

app.py

@@ -1,42 +1,38 @@
 import gradio as gr
-from transformers import pipeline, AutoProcessor, AutoModelForCausalLM
+from transformers import pipeline, AutoProcessor, AutoModelForCausalLM, AutoTokenizer, set_seed
+from datasets import load_dataset
 import torch
 import numpy as np
-from datasets import load_dataset
-from PIL import Image
 
-# 1) IMAGE CAPTIONING MODEL
+# Set seed for reproducibility
+set_seed(42)
+
+# Load BLIP model for image captioning
 caption_model = pipeline("image-to-text", model="Salesforce/blip-image-captioning-base")
 
-# 2) OCR MODEL (Florence-2)
+# Load SpeechT5 model for text-to-speech
+synthesiser = pipeline("text-to-speech", model="microsoft/speecht5_tts")
+
+# Load Florence-2 model for OCR
 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)
 
-# 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 GPT-2 (124M) model for text generation
+gpt2_generator = pipeline('text-generation', model='gpt2')
 
 # 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:
-        # 1) Generate caption from the image
+        # Generate caption from the image
         caption = caption_model(image)[0]['generated_text']
 
-        # 2) Extract text from the image using Florence-2
+        # Extract text (OCR) 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"],
@@ -47,32 +43,28 @@ def process_image(image):
         )
         extracted_text = ocr_processor.batch_decode(generated_ids, skip_special_tokens=True)[0]
 
-        # 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)
+        # 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']
 
-        # 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,
+        # Convert context to speech
+        speech = synthesiser(
+            context,
             forward_params={"speaker_embeddings": speaker_embedding}
         )
 
-        # Prepare audio data for Gradio
-        audio = np.array(speech_data["audio"])
-        rate = speech_data["sampling_rate"]
+        # Prepare audio data
+        audio = np.array(speech["audio"])
+        rate = speech["sampling_rate"]
 
-        # Return audio, caption, extracted text, and final context
-        return (rate, audio), caption, extracted_text, final_context
+        # Return audio, caption, extracted text, and context
+        return (rate, audio), caption, extracted_text, context
 
     except Exception as e:
         return None, f"Error: {str(e)}", "", ""
 
+
 # Gradio Interface
 iface = gr.Interface(
     fn=process_image,
@@ -81,15 +73,10 @@ iface = gr.Interface(
         gr.Audio(label="Generated Audio"),
         gr.Textbox(label="Generated Caption"),
         gr.Textbox(label="Extracted Text (OCR)"),
-        gr.Textbox(label="QA-derived Context")
+        gr.Textbox(label="Generated Context")
     ],
-    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."
-    ),
+    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."
 )
 
 iface.launch()