Update app.py

app.py

@@ -25,21 +25,20 @@ model = AutoModel.from_pretrained(
     torch_dtype=torch.float16,
     trust_remote_code=True,
     token=api_token,
-    revision="main"  # Pin to specific revision
+    revision="main"
 )
 
 tokenizer = AutoTokenizer.from_pretrained(
     "ContactDoctor/Bio-Medical-MultiModal-Llama-3-8B-V1",
     trust_remote_code=True,
     token=api_token,
-    revision="main"  # Pin to specific revision
+    revision="main"
 )
 
 def analyze_input(image_data, question):
     try:
         # Handle base64 image if provided
         if isinstance(image_data, str) and image_data.startswith('data:image'):
-            # Extract base64 data after the comma
             base64_data = image_data.split(',')[1]
             image_bytes = base64.b64decode(base64_data)
             image = Image.open(io.BytesIO(image_bytes)).convert('RGB')
@@ -51,14 +50,30 @@ def analyze_input(image_data, question):
 
         # Process with or without image
         if image is not None:
-            inputs = model.prepare_inputs_for_generation(
-                input_ids=tokenizer(question, return_tensors="pt").input_ids,
-                images=[image]
-            )
+            # Prepare inputs for multimodal generation
+            model_inputs = {
+                "input_ids": tokenizer(question, return_tensors="pt").input_ids.to(model.device),
+                "images": [image]
+            }
         else:
-            inputs = tokenizer(question, return_tensors="pt")
+            # Prepare inputs for text-only generation
+            model_inputs = {
+                "input_ids": tokenizer(question, return_tensors="pt").input_ids.to(model.device)
+            }
 
-        outputs = model.generate(**inputs, max_new_tokens=256)
+        # Generate response with proper inputs
+        generation_config = {
+            "max_new_tokens": 256,
+            "do_sample": True,
+            "temperature": 0.7,
+            "top_p": 0.9,
+        }
+        
+        outputs = model.generate(
+            model_inputs=model_inputs,
+            **generation_config
+        )
+        
         response = tokenizer.decode(outputs[0], skip_special_tokens=True)
         
         return {
@@ -81,10 +96,10 @@ demo = gr.Interface(
     outputs=gr.JSON(label="Analysis"),
     title="Bio-Medical MultiModal Analysis",
     description="Ask questions with or without an image",
-    flagging_mode="never"  # Updated from allow_flagging
+    flagging_mode="never"
 )
 
-# Launch with simplified parameters
+# Launch the interface
 demo.launch(
     share=True,
     server_name="0.0.0.0",