Update app.py

app.py

@@ -2,113 +2,55 @@ import numpy as np
 import torch
 from transformers import AutoTokenizer, AutoModelForCausalLM
 import gradio as gr
-from PIL import Image
-import io
 
 # Model setup
-device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')  # Use GPU if available
-dtype = torch.float32  # Adjust based on model requirements
-model_name_or_path = 'GoodBaiBai88/M3D-LaMed-Phi-3-4B'
+device = torch.device("cpu")  # Use 'cuda' if GPU is available
+dtype = torch.float32
+model_name_or_path = "GoodBaiBai88/M3D-LaMed-Phi-3-4B"
 proj_out_num = 256
 
 # Load model and tokenizer
-try:
-    model = AutoModelForCausalLM.from_pretrained(
-        model_name_or_path,
-        torch_dtype=dtype,
-        device_map=device,
-        trust_remote_code=True
-    )
-    print("Model loaded successfully!")
-except Exception as e:
-    print(f"Error loading model: {e}")
-    raise
+model = AutoModelForCausalLM.from_pretrained(
+    model_name_or_path,
+    torch_dtype=torch.float32,
+    device_map='cpu',
+    trust_remote_code=True
+)
 
-try:
-    tokenizer = AutoTokenizer.from_pretrained(
-        model_name_or_path,
-        model_max_length=512,
-        padding_side="right",
-        use_fast=False,
-        trust_remote_code=True
-    )
-    print("Tokenizer loaded successfully!")
-except Exception as e:
-    print(f"Error loading tokenizer: {e}")
-    raise
+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
+current_image = None
 
-# Function to convert .npy to JPEG
-def npy_to_jpeg(image_np):
-    # Handle multi-dimensional .npy files
-    if image_np.ndim == 4:  # If batch dimension is present (e.g., (1, 256, 256, 3))
-        image_np = image_np.squeeze(0)  # Remove batch dimension
-    elif image_np.ndim == 2:  # Grayscale image (e.g., (256, 256))
-        image_np = np.expand_dims(image_np, axis=-1)  # Add channel dimension
-    
-    # Normalize and convert to uint8
-    image_np = (image_np - image_np.min()) / (image_np.max() - image_np.min()) * 255
-    image_np = image_np.astype(np.uint8)
-    
-    # Convert to PIL Image
-    if image_np.shape[-1] == 1:  # Grayscale
-        image_np = image_np.squeeze()
-        image = Image.fromarray(image_np, mode='L')
-    else:  # RGB
-        image = Image.fromarray(image_np, mode='RGB')
-    
-    # Save to bytes
-    buf = io.BytesIO()
-    image.save(buf, format='JPEG')
-    buf.seek(0)
-    return buf
-
-# Function to process image and generate response
 def process_image(question):
     global current_image
     if current_image is None:
-        return "Please upload an image first.", 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)
     
-    try:
-        # Load the stored .npy image
-        image_np = np.load(current_image)
-        if image_np.shape[-1] != proj_out_num:  # Ensure image matches expected dimensions
-            return f"Invalid image dimensions. Expected {proj_out_num} patches, got {image_np.shape[-1]}.", None
-        
-        # Convert .npy to JPEG
-        jpeg_image = npy_to_jpeg(image_np)
-        
-        # Prepare image tokens and input text
-        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, dtype=torch.long)  # Ensure input_id is LongTensor
-        
-        # Prepare image tensor
-        image_pt = torch.from_numpy(image_np).unsqueeze(0).to(dtype=dtype, device=device)
-        
-        # Generate response
-        generation = model.generate(
-            inputs=image_pt,
-            input_ids=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], jpeg_image
-    except Exception as e:
-        return f"Error processing image: {str(e)}", None
+    # Generate response
+    generation = model.generate(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, jpeg_image = process_image(question)
-    if jpeg_image:
-        chat_history.append((None, (jpeg_image,)))  # Display image in chat
+    response = process_image(question)
     chat_history.append((question, response))
     return chat_history
 
@@ -118,37 +60,20 @@ def upload_image(image):
     current_image = image.name
     return "Image uploaded successfully!"
 
-# Function to clear chat history
-def clear_chat():
-    global chat_history, current_image
-    chat_history = []
-    current_image = None
-    return [], "Chat history cleared."
-
 # Gradio UI
 with gr.Blocks(theme=gr.themes.Soft()) as chat_ui:
     gr.Markdown("# 🏥 Medical Image Analysis Chatbot")
-    
-    # File upload section
-    with gr.Row():
-        upload_button = gr.UploadButton(label="📁 Upload .npy Image", file_types=[".npy"])
-        upload_status = gr.Textbox(label="Status", interactive=False)
-    
-    # Chat interface
-    with gr.Row():
-        chat_list = gr.Chatbot(value=[], label="Chat History", elem_id="chat-history")
-    
-    # Question input and buttons
     with gr.Row():
-        question_input = gr.Textbox(label="Ask a question", placeholder="Ask something about the image...", lines=2)
-        submit_button = gr.Button("Send")
-        clear_button = gr.Button("Clear Chat")
+        with gr.Column(scale=1, min_width=200):
+            chat_list = gr.Chatbot(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")
     
-    # Event handlers
-    upload_button.upload(upload_image, upload_button, upload_status)
+    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)
-    clear_button.click(clear_chat, outputs=[chat_list, upload_status])
 
-# Launch the app
-chat_ui.launch()
+chat_ui.launch()