Update app.py

app.py

@@ -16,10 +16,9 @@ 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 TxGemma model for text generation
-text_gen_device = "cuda:0" if torch.cuda.is_available() else "cpu"
-text_gen_tokenizer = AutoTokenizer.from_pretrained("google/txgemma-9b-predict")
-text_gen_model = AutoModelForCausalLM.from_pretrained("google/txgemma-9b-predict", device_map="auto")
+# Load Llama 3.2 model for text generation
+llama_tokenizer = AutoTokenizer.from_pretrained("meta-llama/Llama-3.2-3B-Instruct")
+llama_model = AutoModelForCausalLM.from_pretrained("meta-llama/Llama-3.2-3B-Instruct", device_map="auto", torch_dtype=torch.bfloat16)
 
 # Load speaker embedding
 embeddings_dataset = load_dataset("Matthijs/cmu-arctic-xvectors", split="validation")
@@ -31,12 +30,6 @@ def process_image(image):
         # Generate caption from the image
         caption = caption_model(image)[0]['generated_text']
 
-        # Convert caption to speech
-        speech = synthesiser(
-            caption,
-            forward_params={"speaker_embeddings": speaker_embedding}
-        )
-
         # 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(
@@ -48,17 +41,23 @@ def process_image(image):
         )
         extracted_text = ocr_processor.batch_decode(generated_ids, skip_special_tokens=True)[0]
 
-        # Generate context from caption and extracted text using TxGemma
-        prompt = f"Instructions: Determine the context of the image based on the caption and extracted text.\nCaption: {caption}\nExtracted Text: {extracted_text}\nContext:"
-        input_ids = text_gen_tokenizer(prompt, return_tensors="pt").to(text_gen_device)
-        outputs = text_gen_model.generate(**input_ids, max_new_tokens=50)
-        context = text_gen_tokenizer.decode(outputs[0], skip_special_tokens=True)
+        # Generate context using Llama 3.2
+        llama_prompt = f"Determine the context of this image. Caption: {caption}. Extracted text: {extracted_text}. Context:"
+        inputs_llama = llama_tokenizer(llama_prompt, return_tensors="pt").to(llama_model.device)
+        llama_output_ids = llama_model.generate(**inputs_llama, max_new_tokens=100)
+        context = llama_tokenizer.decode(llama_output_ids[0], skip_special_tokens=True)
+
+        # Convert context to speech
+        speech = synthesiser(
+            context,
+            forward_params={"speaker_embeddings": speaker_embedding}
+        )
 
         # Prepare audio data
         audio = np.array(speech["audio"])
         rate = speech["sampling_rate"]
 
-        # Return audio, caption, extracted text, and generated context
+        # Return audio, caption, extracted text, and context
         return (rate, audio), caption, extracted_text, context
 
     except Exception as e:
@@ -75,9 +74,10 @@ iface = gr.Interface(
         gr.Textbox(label="Extracted Text (OCR)"),
         gr.Textbox(label="Generated Context")
     ],
-    title="SeeSay Contextualizer",
-    description="Upload an image to generate a caption, extract text, create audio, and determine the context using TxGemma."
+    title="SeeSay Contextualizer with Llama 3.2",
+    description="Upload an image to generate a caption, extract text, create audio from context, and determine the context using Llama 3.2."
 )
 
 iface.launch()
 
+