Update app.py

app.py

@@ -1,10 +1,7 @@
 import os
 import torch
-from transformers import AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig
+from transformers import AutoModel, AutoTokenizer, BitsAndBytesConfig
 import gradio as gr
-from PIL import Image
-import base64
-import io
 
 # Get API token from environment variable
 api_token = os.getenv("HF_TOKEN").strip()
@@ -17,84 +14,64 @@ bnb_config = BitsAndBytesConfig(
     bnb_4bit_compute_dtype=torch.float16
 )
 
-# Load model for causal language modeling
-model = AutoModelForCausalLM.from_pretrained(
+# Load the model and tokenizer
+model = AutoModel.from_pretrained(
     "ContactDoctor/Bio-Medical-MultiModal-Llama-3-8B-V1",
     quantization_config=bnb_config,
     device_map="auto",
     torch_dtype=torch.float16,
     trust_remote_code=True,
-    token=api_token,
-    revision="main"
+    token=api_token
 )
 
 tokenizer = AutoTokenizer.from_pretrained(
     "ContactDoctor/Bio-Medical-MultiModal-Llama-3-8B-V1",
     trust_remote_code=True,
-    token=api_token,
-    revision="main"
+    token=api_token
 )
 
-def analyze_input(image_data, question):
+def analyze_input(image, question):
     try:
-        if not question.strip():
-            return {
-                "status": "error",
-                "message": "Question is required."
+        # Prepare inputs
+        if image:
+            prompt = f"Given the medical image and question: {question}\nPlease provide a detailed analysis."
+            # Convert image to RGB
+            image = image.convert('RGB')
+            # Custom model_inputs for multimodal generation
+            model_inputs = {
+                "input_ids": tokenizer(prompt, return_tensors="pt").input_ids.to(model.device),
+                "images": [image]
             }
-        
-        # Handle the input image (if any)
-        if image_data is not None:
-            return {
-                "status": "error",
-                "message": "Image support is not implemented yet."
+        else:
+            prompt = f"Medical question: {question}\nAnswer:"
+            model_inputs = {
+                "input_ids": tokenizer(prompt, return_tensors="pt").input_ids.to(model.device),
+                "images": None
             }
-        
-        # Prepare prompt for text-only input
-        prompt = f"Medical question: {question}\nAnswer: "
 
-        # Tokenize input
-        inputs = tokenizer(prompt, return_tensors="pt")
-        input_ids = inputs.input_ids.to(model.device)
+        # Generate response using model's custom method
+        outputs = model.generate(model_inputs=model_inputs, max_new_tokens=256)
 
-        # Generate response
-        outputs = model.generate(
-            input_ids=input_ids,
-            max_length=256,  # Limit the length of the generated text
-            eos_token_id=tokenizer.eos_token_id,  # Ensure generation stops correctly
-            pad_token_id=tokenizer.pad_token_id,
-            temperature=0.7,  # Control randomness
-            top_p=0.9,  # Nucleus sampling
-            top_k=50  # Top-k sampling
-        )
-
-        # Decode response
+        # Decode and clean response
         response = tokenizer.decode(outputs[0], skip_special_tokens=True)
+        return {"status": "success", "response": response}
 
-        return {
-            "status": "success",
-            "response": response
-        }
     except Exception as e:
-        return {
-            "status": "error",
-            "message": str(e)
-        }
+        return {"status": "error", "message": str(e)}
 
 # Create Gradio interface
 demo = gr.Interface(
     fn=analyze_input,
     inputs=[
-        gr.Image(type="numpy", label="Medical Image (Optional)"),
-        gr.Textbox(label="Question", placeholder="Enter your medical query...")
+        gr.Image(type="pil", label="Upload Medical Image (Optional)"),
+        gr.Textbox(label="Medical Question")
     ],
     outputs=gr.JSON(label="Analysis"),
-    title="Medical Query Analysis",
-    description="Ask medical questions. For now, please focus on text-based queries.",
-    flagging_mode="never"
+    title="ContactDoctor Medical Assistant",
+    description="Upload a medical image and/or enter a question to receive detailed AI-powered responses."
 )
 
-# Launch the interface
+# Launch the Gradio app
 demo.launch(
     share=True,
     server_name="0.0.0.0",