Update app.py

app.py

@@ -3,58 +3,78 @@
 #api_token = os.getenv("HF_TOKEN").strip()
 
 import torch
-from PIL import Image
+from flask import Flask, request, jsonify
 from transformers import AutoModel, AutoTokenizer, BitsAndBytesConfig
-import gradio as gr
+from PIL import Image
+import io
+import base64
 
+app = Flask(__name__)
 
-# Configuration for 4-bit quantization
+# Quantization configuration
 bnb_config = BitsAndBytesConfig(
-    load_in_4bit=True,
-    bnb_4bit_quant_type="nf4",
-    bnb_4bit_use_double_quant=True,
-    bnb_4bit_compute_dtype=torch.float16,
+    load_in_4bit=True, 
+    bnb_4bit_quant_type="nf4", 
+    bnb_4bit_use_double_quant=True, 
+    bnb_4bit_compute_dtype=torch.float16
 )
 
-# Load the model without flash-attn
+# Load model
 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,
-    attn_implementation=None,  # Disable flash-attn
+    "ContactDoctor/Bio-Medical-MultiModal-Llama-3-8B-V1", 
+    quantization_config=bnb_config, 
+    device_map="auto", 
+    torch_dtype=torch.float16, 
+    trust_remote_code=True, 
+    attn_implementation="flash_attention_2"
 )
 
 tokenizer = AutoTokenizer.from_pretrained(
-    "ContactDoctor/Bio-Medical-MultiModal-Llama-3-8B-V1",
+    "ContactDoctor/Bio-Medical-MultiModal-Llama-3-8B-V1", 
     trust_remote_code=True
 )
 
+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)
+
+    try:
+        # Process with image if provided
+        if base64_image:
+            image = decode_base64_image(base64_image)
+            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
+            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
+        })
+
+    except Exception as e:
+        return jsonify({
+            'status': 'error', 
+            'message': str(e)
+        }), 500
 
-# Define the function to handle the input
-def process_input(image, question):
-    image = Image.open(image).convert("RGB")
-    msgs = [{'role': 'user', 'content': [image, question]}]
-    res = model.chat(image=image, msgs=msgs, tokenizer=tokenizer, sampling=True, temperature=0.95, stream=True)
-
-    generated_text = ""
-    for new_text in res:
-        generated_text += new_text
-    return generated_text
-
-# Gradio interface
-iface = gr.Interface(
-    fn=process_input,
-    inputs=[
-        gr.Image(type="file", label="Upload Image"),
-        gr.Textbox(lines=2, label="Question")
-    ],
-    outputs=gr.Textbox(label="Generated Response"),
-    title="BioMedical MultiModal Llama",
-    description="Upload an image and ask a medical question."
-)
 
-if __name__ == "__main__":
-    iface.launch()
+if __name__ == '__main__':
+    app.run(debug=True)