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"
model = AutoModel.from_pretrained(model_name, trust_remote_code=True, device_map="auto", torch_dtype=torch.float16)
tokenizer = AutoTokenizer.from_pretrained(model_name, trust_remote_code=True)

def process_query(image, question):
    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(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/or ask a question for AI-powered assistance.",
)

iface.launch()