Update app.py

app.py

@@ -1,11 +1,7 @@
-
-
-# Get API token from environment variable
-#api_token = os.getenv("HF_TOKEN").strip()
-
 import gradio as gr
 from transformers import AutoModel, AutoTokenizer
 import torch
+from PIL import Image
 
 # Load the model and tokenizer
 model_name = "ContactDoctor/Bio-Medical-MultiModal-Llama-3-8B-V1"
@@ -13,20 +9,41 @@ model = AutoModel.from_pretrained(model_name, trust_remote_code=True, device_map
 tokenizer = AutoTokenizer.from_pretrained(model_name, trust_remote_code=True)
 
 def process_query(image, question):
-    inputs = {"question": question}
-    if image:
-        inputs["image"] = image
-
-    # Process the inputs and generate a response
-    response = model.chat(image=inputs.get("image"), msgs=[{"role": "user", "content": question}], tokenizer=tokenizer)
-    return response
-
+    try:
+        # Construct the messages for the model
+        msgs = [{"role": "user", "content": question}]
+        
+        # Handle cases with and without an image
+        if image is not None:
+            # Convert the image to the required format
+            image_input = [Image.fromarray(image).convert("RGB")]
+            response = model.chat(
+                image=image_input,
+                msgs=msgs,
+                tokenizer=tokenizer,
+            )
+        else:
+            # For text-only queries, omit the `image` parameter
+            response = model.chat(
+                image=None,
+                msgs=msgs,
+                tokenizer=tokenizer,
+            )
+        
+        return response
+    except Exception as e:
+        return f"Error: {str(e)}"
+
+# Gradio interface
 iface = gr.Interface(
     fn=process_query,
-    inputs=[gr.Image(label="Upload Medical Image"), gr.Textbox(label="Question")],
+    inputs=[
+        gr.Image(type="numpy", label="Upload Medical Image (Optional)"),
+        gr.Textbox(label="Enter Your Medical Question"),
+    ],
     outputs="text",
     title="Medical Multimodal Assistant",
-    description="Upload a medical image and ask your question."
+    description="Upload a medical image and/or ask a question for AI-powered assistance.",
 )
 
 iface.launch()