Update app.py

app.py

@@ -1,5 +1,5 @@
 import gradio as gr
-from transformers import pipeline, AutoProcessor, AutoModelForCausalLM
+from transformers import pipeline, AutoProcessor, AutoModelForCausalLM, AutoTokenizer
 from datasets import load_dataset
 import torch
 import numpy as np
@@ -16,6 +16,11 @@ 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 speaker embedding
 embeddings_dataset = load_dataset("Matthijs/cmu-arctic-xvectors", split="validation")
 speaker_embedding = torch.tensor(embeddings_dataset[7306]["xvector"]).unsqueeze(0)
@@ -43,15 +48,21 @@ 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)
+
         # Prepare audio data
         audio = np.array(speech["audio"])
         rate = speech["sampling_rate"]
 
-        # Return audio, caption, and extracted text
-        return (rate, audio), caption, extracted_text
+        # Return audio, caption, extracted text, and generated context
+        return (rate, audio), caption, extracted_text, context
 
     except Exception as e:
-        return None, f"Error: {str(e)}", ""
+        return None, f"Error: {str(e)}", "", ""
 
 
 # Gradio Interface
@@ -61,10 +72,12 @@ iface = gr.Interface(
     outputs=[
         gr.Audio(label="Generated Audio"),
         gr.Textbox(label="Generated Caption"),
-        gr.Textbox(label="Extracted Text (OCR)")
+        gr.Textbox(label="Extracted Text (OCR)"),
+        gr.Textbox(label="Generated Context")
     ],
-    title="SeeSay with SpeechT5 and Florence-2 OCR",
-    description="Upload an image to generate a caption, hear it described with SpeechT5's speech synthesis, and extract text using Florence-2 OCR."
+    title="SeeSay Contextualizer",
+    description="Upload an image to generate a caption, extract text, create audio, and determine the context using TxGemma."
 )
 
 iface.launch()
+