Update app.py

app.py

@@ -2,6 +2,8 @@ import os
 import torch
 from transformers import AutoModel, AutoTokenizer, BitsAndBytesConfig
 import gradio as gr
+from PIL import Image
+from torchvision.transforms import ToTensor
 
 # Get API token from environment variable
 api_token = os.getenv("HF_TOKEN").strip()
@@ -30,27 +32,34 @@ tokenizer = AutoTokenizer.from_pretrained(
     token=api_token
 )
 
+# Preprocess image
+def preprocess_image(image):
+    transform = ToTensor()
+    return transform(image).unsqueeze(0).to(model.device)
+
 def analyze_input(image, question):
     try:
         # Prepare inputs
         if image:
-            prompt = f"Given the medical image and question: {question}\nPlease provide a detailed analysis."
-            # Convert image to RGB
+            # Process image
             image = image.convert('RGB')
-            # Custom model_inputs for multimodal generation
+            pixel_values = preprocess_image(image)
+            prompt = f"Given the medical image and question: {question}\nPlease provide a detailed analysis."
+            
+            # Model inputs for multimodal processing
             model_inputs = {
                 "input_ids": tokenizer(prompt, return_tensors="pt").input_ids.to(model.device),
-                "images": [image]
+                "pixel_values": pixel_values
             }
         else:
+            # Text-only processing
             prompt = f"Medical question: {question}\nAnswer:"
             model_inputs = {
-                "input_ids": tokenizer(prompt, return_tensors="pt").input_ids.to(model.device),
-                "images": None
+                "input_ids": tokenizer(prompt, return_tensors="pt").input_ids.to(model.device)
             }
 
         # Generate response using model's custom method
-        outputs = model.generate(model_inputs=model_inputs, max_new_tokens=256)
+        outputs = model.generate(**model_inputs, max_new_tokens=256)
 
         # Decode and clean response
         response = tokenizer.decode(outputs[0], skip_special_tokens=True)