import numpy as np
import torch
from transformers import AutoTokenizer, AutoModelForCausalLM
import gradio as gr

# Model setup
device = torch.device('cpu')  # Use 'cuda' if GPU is available
dtype = torch.float32  # Data type for model processing
model_name_or_path = 'GoodBaiBai88/M3D-LaMed-Phi-3-4B'
proj_out_num = 256

# Load model and tokenizer
model = AutoModelForCausalLM.from_pretrained(
    model_name_or_path,
    torch_dtype=torch.float32,
    device_map='cpu',
    trust_remote_code=True
)

tokenizer = AutoTokenizer.from_pretrained(
    model_name_or_path,
    model_max_length=512,
    padding_side="right",
    use_fast=False,
    trust_remote_code=True
)

# Chat history storage
chat_history = []
current_image = None  # To store the uploaded image

def process_image(question):
    global current_image
    if current_image is None:
        return "Please upload an image first."
    
    image_np = np.load(current_image)  # Load the stored .npy image
    image_tokens = "<im_patch>" * proj_out_num
    input_txt = image_tokens + question
    input_id = tokenizer(input_txt, return_tensors="pt")['input_ids'].to(device=device)
    
    # Prepare image for model
    image_pt = torch.from_numpy(image_np).unsqueeze(0).to(dtype=dtype, device=device)
    
    # Generate response
    generation = model.generate(image_pt, input_id, max_new_tokens=256, do_sample=True, top_p=0.9, temperature=1.0)
    generated_texts = tokenizer.batch_decode(generation, skip_special_tokens=True)
    return generated_texts[0]

# Function to update chat
def chat_interface(question):
    global chat_history
    response = process_image(question)
    chat_history.append((question, response))
    return chat_history

# Function to handle image upload
def upload_image(image):
    global current_image
    current_image = image.name
    return "Image uploaded successfully!"

# Gradio UI
with gr.Blocks(theme=gr.themes.Soft()) as chat_ui:
    gr.Markdown("# 🏥 Medical Image Analysis Chatbot")
    with gr.Row():
        with gr.Column(scale=1, min_width=200):
            chat_list = gr.Chatbot(value=[], label="Chat History", elem_id="chat-history")
        with gr.Column(scale=4):
            uploaded_image = gr.File(label="Upload .npy Image", type="filepath")
            upload_status = gr.Textbox(label="Status", interactive=False)
            question_input = gr.Textbox(label="Ask a question", placeholder="Ask something about the image...")
            submit_button = gr.Button("Send")
    
    uploaded_image.upload(upload_image, uploaded_image, upload_status)
    submit_button.click(chat_interface, question_input, chat_list)
    question_input.submit(chat_interface, question_input, chat_list)

chat_ui.launch()