Update app.py

app.py

@@ -1,29 +1,40 @@
 import gradio as gr
 import torch
-from transformers import AutoProcessor, AutoModel
+from transformers import AutoProcessor, Idefics3ForConditionalGeneration
 from PIL import Image
 
-# Load model and processor using generic AutoModel
-model_id = "ibm-granite/granite-docling-258M"
-# We must use trust_remote_code=True for this specific model
+# Load model and processor
+model_id = "HuggingFaceM4/Idefics3-8B-Llama3" # Ensure this matches your model
 processor = AutoProcessor.from_pretrained(model_id, trust_remote_code=True)
-model = AutoModel.from_pretrained(model_id, trust_remote_code=True, torch_dtype=torch.float32)
+model = Idefics3ForConditionalGeneration.from_pretrained(
+    model_id, 
+    trust_remote_code=True, 
+    torch_dtype=torch.float32
+)
 
 def process_image(image):
-    # Convert image to RGB
+    # Safety check for empty input
+    if image is None:
+        return "Please upload an image first."
+    
+    # Ensure image is PIL format
+    if not isinstance(image, Image.Image):
+        image = Image.fromarray(image)
+        
     image = image.convert("RGB")
     
-    # Prepare the inputs
-    messages = [{"role": "user", "content": [{"type": "image"}]}]
+    # Prepare inputs
+    messages = [{"role": "user", "content": [{"type": "image"}, {"type": "text", "text": "Describe this image."}]}]
     prompt = processor.apply_chat_template(messages, add_generation_prompt=True)
     inputs = processor(text=prompt, images=image, return_tensors="pt")
     
-    # Generate output
-    output = model.generate(**inputs, max_new_tokens=500)
-    result = processor.decode(output[0], skip_special_tokens=True)
-    return result
+    # Generate
+    generated_ids = model.generate(**inputs, max_new_tokens=500)
+    result = processor.batch_decode(generated_ids, skip_special_tokens=True)
+    
+    return result[0]
 
-# Create interface
+# UI Setup
 demo = gr.Interface(
     fn=process_image,
     inputs=gr.Image(type="pil"),