Update app.py

app.py

@@ -1,10 +1,12 @@
-import gradio as gr
+import os
 import torch
 from transformers import AutoModel, AutoTokenizer, BitsAndBytesConfig
 from PIL import Image
-import os
+import gradio as gr
+import base64
+import io
 
-# Get API token from environment variables
+# Get API token from environment variable
 api_token = os.getenv("HF_TOKEN").strip()
 
 # Quantization configuration
@@ -15,7 +17,7 @@ bnb_config = BitsAndBytesConfig(
     bnb_4bit_compute_dtype=torch.float16
 )
 
-# Load the model and tokenizer
+# Load model
 model = AutoModel.from_pretrained(
     "ContactDoctor/Bio-Medical-MultiModal-Llama-3-8B-V1", 
     quantization_config=bnb_config, 
@@ -31,42 +33,54 @@ tokenizer = AutoTokenizer.from_pretrained(
     token=api_token
 )
 
-# Function to handle inputs
-def process_query(image, question):
+def analyze_input(image_data=None, question=""):
     try:
-        if image:
-            # Process image and text
-            image = image.convert('RGB')
+        # 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')
+        # Handle direct image input
+        elif image_data is not None:
+            image = Image.fromarray(image_data).convert('RGB')
+        else:
+            image = None
+
+        # 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]
             )
-            outputs = model.generate(**inputs, max_new_tokens=256)
         else:
-            # Process text-only
             inputs = tokenizer(question, return_tensors="pt")
-            outputs = model.generate(**inputs, max_new_tokens=256)
 
-        # Decode response
+        outputs = model.generate(**inputs, max_new_tokens=256)
         response = tokenizer.decode(outputs[0], skip_special_tokens=True)
-        return response
-
+        
+        return {
+            "status": "success",
+            "response": response
+        }
     except Exception as e:
-        return f"Error: {str(e)}"
+        return {
+            "status": "error",
+            "message": str(e)
+        }
 
-# Define Gradio interface
-interface = gr.Interface(
-    fn=process_query,
+# Create Gradio interface
+demo = gr.Interface(
+    fn=analyze_input,
     inputs=[
-        gr.Image(type="pil", label="Upload an Image (Optional)"),
-        gr.Textbox(label="Enter a Question")
+        gr.Image(type="numpy", label="Medical Image (Optional)", optional=True),
+        gr.Textbox(label="Question", placeholder="Enter your medical query...")
     ],
-    outputs="text",
-    title="ContactDoctor Multimodal Medical Assistant",
-    description="Provide an image and/or question to get AI-powered medical advice.",
-    enable_api=True  # Enable API for external calls
+    outputs=gr.JSON(label="Analysis"),
+    title="Bio-Medical MultiModal Analysis",
+    description="Ask questions with or without an image",
+    allow_flagging="never",
 )
 
-# Launch the app
-if __name__ == "__main__":
-    interface.launch()
+# Launch with API access enabled
+demo.launch(share=True, server_name="0.0.0.0", server_port=7860, enable_queue=True)