Update app.py

app.py

@@ -23,7 +23,7 @@ def load_model():
         )
 
     # Load the processor and model using the correct identifier
-    model_id = "google/paligemma2-10b-pt-448"
+    model_id = "google/paligemma2-28b-pt-896"
     processor = PaliGemmaProcessor.from_pretrained(model_id, use_auth_token=token)
     device = "cuda" if torch.cuda.is_available() else "cpu"
     model = PaliGemmaForConditionalGeneration.from_pretrained(
@@ -34,7 +34,7 @@ def load_model():
 
 
 @spaces.GPU(duration=120)  # Increased timeout to 120 seconds
-def process_image_and_text(image_pil, text_input, num_beams, temperature, seed):
+def process_image_and_text(image_pil, num_beams, temperature, seed):
     """Extract text from image using PaliGemma2."""
     try:
         processor, model = load_model()
@@ -43,18 +43,21 @@ def process_image_and_text(image_pil, text_input, num_beams, temperature, seed):
         # Load the image using load_image
         image = load_image(image_pil)
 
-        prompt = ""
+        # Add <image> token to the beginning of the text prompt
+        text_input = " "
 
         # Use the provided text input
-        model_inputs = processor(text=prompt, images=image, return_tensors="pt").to(
+        model_inputs = processor(text=text_input, images=image, return_tensors="pt").to(
             device, dtype=torch.bfloat16
         )
         input_len = model_inputs["input_ids"].shape[-1]
-        
-        torch.manual_seed(seed) # Set random seed for reproducibility
+
+        # Set random seed for reproducibility, only if a seed is provided
+        if seed is not None:
+            torch.manual_seed(int(seed))
 
         with torch.inference_mode():
-            generation = model.generate(**model_inputs, max_new_tokens=500, do_sample=True, num_beams=num_beams, temperature=temperature)
+            generation = model.generate(**model_inputs, max_new_tokens=200, do_sample=True, num_beams=num_beams, temperature=temperature)
             generation = generation[0][input_len:]
             decoded = processor.decode(generation, skip_special_tokens=True)
 
@@ -71,10 +74,10 @@ if __name__ == "__main__":
             gr.Image(type="pil", label="Upload an image"),
             gr.Slider(minimum=1, maximum=10, step=1, value=1, label="Number of Beams"),
             gr.Slider(minimum=0.1, maximum=2.0, step=0.1, value=1.0, label="Temperature"),
-            gr.Number(label="Random Seed", value=0, precision=0),
+            gr.Number(label="Random Seed", value=42, precision=0,  allow_none=True),
         ],
         outputs=gr.Textbox(label="Generated Text"),
-        title="PaliGemma2 Image and Text to Text",
+        title="PaliGemma2 Image to Text",
         description="Upload an image and enter a text prompt. The model will generate text based on both.",
     )
     iface.launch()