y

app.py

@@ -3,68 +3,77 @@ import torch
 from transformers import AutoTokenizer, AutoModelForCausalLM
 import gradio as gr
 
-# Set device & model details
-device = torch.device('cpu')
-dtype = torch.float32
+# 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  # Number of projection outputs required
+proj_out_num = 256
 
-# Load model & tokenizer
+# Load model and tokenizer
 model = AutoModelForCausalLM.from_pretrained(
-    model_name_or_path, torch_dtype=dtype, device_map='cpu', trust_remote_code=True
+    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
+    model_name_or_path,
+    model_max_length=512,
+    padding_side="right",
+    use_fast=False,
+    trust_remote_code=True
 )
 
-# Image placeholder (to maintain session context)
-uploaded_image = None
+# Chat history storage
+chat_history = []
+current_image = None  # To store the uploaded image
 
-def process_image(question, history):
-    global uploaded_image
-    if uploaded_image is None:
-        return "⚠️ Please upload an image first!"
-
-    # Load the .npy image
-    image_np = np.load(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)
-
-    # Convert image to tensor
+    
+    # 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):
-    """ Stores the uploaded image path to be used in chat """
-    global uploaded_image
-    uploaded_image = image.name
-    return f"✅ Image uploaded successfully: {image.name}"
+    global current_image
+    current_image = image.name
+    return "Image uploaded successfully!"
 
-# Chat Interface with File Upload
-with gr.Blocks(theme="soft") as chat_ui:
+# 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=2):
-            gr.Markdown("# 🏥 Medical Image Chatbot")
-            uploaded_file = gr.File(label="Upload .npy Image", type="filepath")
-            upload_button = gr.Button("Upload")
-            status = gr.Markdown("")
-            chat = gr.Chatbot(height=400)
-
-        with gr.Column(scale=3):
-            input_box = gr.Textbox(placeholder="Ask something about the image...")
-            send_button = gr.Button("Send ✉️")
-
-    # Handle image upload
-    upload_button.click(upload_image, inputs=[uploaded_file], outputs=[status])
-
-    # Handle chat interaction
-    send_button.click(process_image, inputs=[input_box, chat], outputs=[chat])
+        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()