Update app.py

app.py

@@ -1,9 +1,12 @@
 import gradio as gr
-from transformers import pipeline, AutoProcessor, AutoModelForCausalLM, AutoTokenizer
+from transformers import pipeline, AutoProcessor, AutoModelForCausalLM, AutoTokenizer, set_seed
 from datasets import load_dataset
 import torch
 import numpy as np
 
+# 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")
 
@@ -16,10 +19,8 @@ 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_processor = AutoProcessor.from_pretrained("microsoft/Florence-2-large", trust_remote_code=True)
 
-# Load QwQ model for text generation
-qwq_model_name = "Qwen/QwQ-32B"
-qwq_model = AutoModelForCausalLM.from_pretrained(qwq_model_name, torch_dtype="auto", device_map="auto")
-qwq_tokenizer = AutoTokenizer.from_pretrained(qwq_model_name)
+# Load GPT-2 XL model for text generation
+gpt2_generator = pipeline('text-generation', model='gpt2-xl')
 
 # Load speaker embedding
 embeddings_dataset = load_dataset("Matthijs/cmu-arctic-xvectors", split="validation")
@@ -42,12 +43,10 @@ def process_image(image):
         )
         extracted_text = ocr_processor.batch_decode(generated_ids, skip_special_tokens=True)[0]
 
-        # Generate context using QwQ
-        messages = [{"role": "user", "content": f"Determine the context of this image based on the caption and extracted text. Caption: {caption}. Extracted text: {extracted_text}."}]
-        qwq_prompt = qwq_tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
-        inputs_qwq = qwq_tokenizer(qwq_prompt, return_tensors="pt").to(qwq_model.device)
-        qwq_output_ids = qwq_model.generate(**inputs_qwq, max_new_tokens=100)
-        context = qwq_tokenizer.batch_decode(qwq_output_ids[:, inputs_qwq.input_ids.shape[-1]:], skip_special_tokens=True)[0]
+        # Generate context using GPT-2 XL
+        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=150, num_return_sequences=1)
+        context = context_output[0]['generated_text']
 
         # Convert context to speech
         speech = synthesiser(
@@ -76,9 +75,8 @@ iface = gr.Interface(
         gr.Textbox(label="Extracted Text (OCR)"),
         gr.Textbox(label="Generated Context")
     ],
-    title="SeeSay Contextualizer with QwQ-32B",
-    description="Upload an image to generate a caption, extract text, create audio from context, and determine the context using QwQ-32B."
+    title="SeeSay Contextualizer with GPT-2 XL",
+    description="Upload an image to generate a caption, extract text, create audio from context, and determine the context using GPT-2 XL."
 )
 
 iface.launch()
-