Update app.py

app.py

@@ -1,16 +1,12 @@
-import os
+import gradio as gr
 import torch
-from flask import Flask, request, jsonify
 from transformers import AutoModel, AutoTokenizer, BitsAndBytesConfig
 from PIL import Image
-import io
-import base64
+import os
 
-# Get API token from environment variable
+# Get API token from environment variables
 api_token = os.getenv("HF_TOKEN").strip()
 
-app = Flask(__name__)
-
 # Quantization configuration
 bnb_config = BitsAndBytesConfig(
     load_in_4bit=True, 
@@ -19,7 +15,7 @@ bnb_config = BitsAndBytesConfig(
     bnb_4bit_compute_dtype=torch.float16
 )
 
-# Load model without Flash Attention
+# Load the model and tokenizer
 model = AutoModel.from_pretrained(
     "ContactDoctor/Bio-Medical-MultiModal-Llama-3-8B-V1", 
     quantization_config=bnb_config, 
@@ -35,45 +31,42 @@ tokenizer = AutoTokenizer.from_pretrained(
     token=api_token
 )
 
-def decode_base64_image(base64_string):
-    # Decode base64 image
-    image_data = base64.b64decode(base64_string)
-    image = Image.open(io.BytesIO(image_data)).convert('RGB')
-    return image
-
-@app.route('/analyze', methods=['POST'])
-def analyze_input():
-    data = request.json
-    question = data.get('question', '')
-    base64_image = data.get('image', None)
-
+# Function to handle inputs
+def process_query(image, question):
     try:
-        # Process with image if provided
-        if base64_image:
-            image = decode_base64_image(base64_image)
+        if image:
+            # Process image and text
+            image = image.convert('RGB')
             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:
-            # Text-only processing
+            # Process text-only
             inputs = tokenizer(question, return_tensors="pt")
             outputs = model.generate(**inputs, max_new_tokens=256)
 
         # Decode response
         response = tokenizer.decode(outputs[0], skip_special_tokens=True)
-
-        return jsonify({
-            'status': 'success', 
-            'response': response
-        })
+        return response
 
     except Exception as e:
-        return jsonify({
-            'status': 'error', 
-            'message': str(e)
-        }), 500
+        return f"Error: {str(e)}"
+
+# Define Gradio interface
+interface = gr.Interface(
+    fn=process_query,
+    inputs=[
+        gr.Image(type="pil", label="Upload an Image (Optional)"),
+        gr.Textbox(label="Enter a Question")
+    ],
+    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
+)
 
-if __name__ == '__main__':
-    app.run(debug=True)
+# Launch the app
+if __name__ == "__main__":
+    interface.launch()