Update app.py

app.py

@@ -1,79 +1,67 @@
 import gradio as gr
-from llava.model.builder import load_pretrained_model
-from llava.mm_utils import process_images, tokenizer_image_token
-from llava.constants import IMAGE_TOKEN_INDEX
-import torch
+from llava_med import LlavaMedProcessor, LlavaMedForCausalLM
 from PIL import Image
+import torch
 
-model_path = "microsoft/llava-med-v1.5-mistral-7b"
-tokenizer, model, image_processor, _ = load_pretrained_model(
-    model_path=model_path,
-    model_base=None,
-    model_name="llava-med-v1.5-mistral-7b",
-    load_4bit=False,  # Disable 4-bit quantization for CPU
-    device_map="cpu"   # Force CPU usage
+# Load model and processor
+model = LlavaMedForCausalLM.from_pretrained(
+    "microsoft/llava-med-v1.5-mistral-7b",
+    torch_dtype=torch.float32,  # Use float32 for CPU stability
+    low_cpu_mem_usage=True,
+    device_map="cpu"
+)
+processor = LlavaMedProcessor.from_pretrained(
+    "microsoft/llava-med-v1.5-mistral-7b"
 )
-model.to('cpu')
 
 def analyze_medical_image(image, question):
-    if isinstance(image, str):
-        image = Image.open(image)
-    else:
-        image = Image.fromarray(image)
+    # Prepare inputs
+    prompt = f"Question: {question} Answer:"
     
-    image_tensor = process_images([image], image_processor, model.config)[0]
-    prompt = f"USER: <image>\n{question}\nASSISTANT:"
+    # Process inputs
+    inputs = processor(
+        text=prompt,
+        images=image,
+        return_tensors="pt",
+        padding=True
+    ).to("cpu")
     
-    input_ids = tokenizer_image_token(
-        prompt, 
-        tokenizer,
-        IMAGE_TOKEN_INDEX,
-        return_tensors='pt'
-    ).unsqueeze(0)
-    
-    with torch.inference_mode():
-        output_ids = model.generate(
-            input_ids,
-            images=image_tensor.unsqueeze(0),
-            max_new_tokens=512,
+    # Generate response
+    with torch.no_grad():
+        outputs = model.generate(
+            **inputs,
+            max_new_tokens=256,
             do_sample=True,
             temperature=0.7,
-            use_cache=True
+            top_p=0.9
         )
-
-    response = tokenizer.decode(
-        output_ids[0][input_ids.shape[1]:], 
+    
+    # Decode response
+    response = processor.batch_decode(
+        outputs, 
         skip_special_tokens=True
-    ).strip()
+    )[0].split("Answer:")[-1].strip()
     
     return response
 
+# Gradio interface
 with gr.Blocks() as demo:
-    gr.Markdown("# LLaVA-Med Medical Image Analysis")
-    gr.Markdown("Ask questions about medical images using Microsoft's LLaVA-Med 1.5-Mistral-7B")
+    gr.Markdown("# LLaVA-Med Medical Analysis (CPU)")
+    gr.Markdown("Official Microsoft LLaVA-Med 1.5-Mistral-7B implementation")
     
     with gr.Row():
         with gr.Column():
-            image_input = gr.Image(label="Upload Medical Image", type="pil")
-            question_input = gr.Textbox(label="Question", placeholder="Ask about the medical image...")
+            image_input = gr.Image(label="Medical Image", type="pil")
+            question_input = gr.Textbox(label="Clinical Question", placeholder="Enter your medical question...")
             submit_btn = gr.Button("Analyze")
         
         with gr.Column():
-            output_text = gr.Textbox(label="Analysis Result", interactive=False)
-    
-    examples = gr.Examples(
-        examples=[
-            ["examples/chest_xray.jpg", "What abnormalities are present in this chest X-ray?"],
-            ["examples/retina_scan.jpg", "Are there any signs of diabetic retinopathy?"]
-        ],
-        inputs=[image_input, question_input],
-        label="Example Queries"
-    )
-    
+            output_text = gr.Textbox(label="Clinical Analysis", interactive=False)
+
     submit_btn.click(
         fn=analyze_medical_image,
         inputs=[image_input, question_input],
         outputs=output_text
     )
 
-demo.queue(max_size=10).launch()
+demo.queue(max_size=5).launch()