Update app.py

app.py

@@ -1,5 +1,12 @@
 import gradio as gr
-from transformers import pipeline, AutoProcessor, AutoModelForCausalLM, AutoTokenizer, set_seed
+from transformers import (
+    pipeline,
+    AutoProcessor,
+    AutoModelForCausalLM,
+    AutoTokenizer,
+    GenerationConfig,
+    set_seed
+)
 from datasets import load_dataset
 import torch
 import numpy as np
@@ -14,25 +21,43 @@ caption_model = pipeline("image-to-text", model="Salesforce/blip-image-captionin
 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_device = "cuda" 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')
+# Load Doge-320M-Instruct model for context generation
+doge_tokenizer = AutoTokenizer.from_pretrained("SmallDoge/Doge-320M-Instruct")
+doge_model = AutoModelForCausalLM.from_pretrained(
+    "SmallDoge/Doge-320M-Instruct",
+    torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32,
+    trust_remote_code=True
+).to("cuda" if torch.cuda.is_available() else "cpu")
+
+doge_generation_config = GenerationConfig(
+    max_new_tokens=100,
+    use_cache=True,
+    do_sample=True,
+    temperature=0.8,
+    top_p=0.9,
+    repetition_penalty=1.0
+)
 
-# Load speaker embedding
+# Load speaker embedding for SpeechT5
 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
+        # Step 1: Generate caption
         caption = caption_model(image)[0]['generated_text']
 
-        # Extract text (OCR) using Florence-2
+        # Step 2: OCR to extract text
         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,22 +68,31 @@ 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']
+        # Step 3: Generate context using Doge model
+        prompt = f"Determine the context of this image based on the caption and extracted text.\nCaption: {caption}\nExtracted text: {extracted_text}\nContext:"
+        conversation = [{"role": "user", "content": prompt}]
+        doge_inputs = doge_tokenizer.apply_chat_template(
+            conversation=conversation,
+            tokenize=True,
+            return_tensors="pt"
+        ).to(doge_model.device)
+
+        doge_outputs = doge_model.generate(
+            doge_inputs,
+            generation_config=doge_generation_config
+        )
+
+        context = doge_tokenizer.decode(doge_outputs[0], skip_special_tokens=True).strip()
 
-        # Convert context to speech
+        # Step 4: 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 context
         return (rate, audio), caption, extracted_text, context
 
     except Exception as e:
@@ -75,8 +109,8 @@ iface = gr.Interface(
         gr.Textbox(label="Extracted Text (OCR)"),
         gr.Textbox(label="Generated 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="SeeSay Contextualizer with Doge-320M",
+    description="Upload an image to generate a caption, extract text (OCR), generate context using Doge, and turn it into speech using SpeechT5."
 )
 
-iface.launch()
+iface.launch()