Update app.py

app.py

@@ -1,21 +1,33 @@
 import gradio as gr
-from transformers import pipeline
+import torch
+from transformers import AutoProcessor, AutoModelForVision2Seq
+from PIL import Image
 
-# This line downloads the model from Hugging Face
-pipe = pipeline("image-text-to-text", model="ibm-granite/granite-docling-258M")
+# Load model and processor
+model_id = "ibm-granite/granite-docling-258M"
+# We must use trust_remote_code=True for this specific model architecture
+processor = AutoProcessor.from_pretrained(model_id, trust_remote_code=True)
+model = AutoModelForVision2Seq.from_pretrained(model_id, trust_remote_code=True, torch_dtype=torch.float32)
 
 def process_image(image):
-    # This runs the model on the uploaded image
-    result = pipe(image)
-    # This returns the text extracted from the document
-    return str(result)
+    # Convert image to RGB
+    image = image.convert("RGB")
+    
+    # Prepare the inputs
+    messages = [{"role": "user", "content": [{"type": "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
 
-# This creates the user interface
+# Create interface
 demo = gr.Interface(
     fn=process_image,
     inputs=gr.Image(type="pil"),
     outputs="text"
 )
 
-# This launches the app
 demo.launch()