Update app.py

app.py

@@ -1,67 +1,59 @@
 import gradio as gr
-from llava_med import LlavaMedProcessor, LlavaMedForCausalLM
+from llava.model.builder import load_pretrained_model
+from llava.mm_utils import get_model_name_from_path
+from llava.eval.run_llava import eval_model
 from PIL import Image
 import torch
 
-# 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"
+# Load model configuration
+model_path = "microsoft/llava-med-v1.5-mistral-7b"
+model_name = get_model_name_from_path(model_path)
+tokenizer, model, image_processor, _ = load_pretrained_model(
+    model_path=model_path,
+    model_base=None,
+    model_name=model_name,
+    device_map="cpu",
+    load_4bit=False
 )
 
 def analyze_medical_image(image, question):
-    # Prepare inputs
-    prompt = f"Question: {question} Answer:"
-    
-    # Process inputs
-    inputs = processor(
-        text=prompt,
-        images=image,
-        return_tensors="pt",
-        padding=True
-    ).to("cpu")
+    # Convert Gradio input to PIL Image
+    if isinstance(image, str):
+        image = Image.open(image)
+    else:
+        image = Image.fromarray(image)
     
-    # Generate response
-    with torch.no_grad():
-        outputs = model.generate(
-            **inputs,
-            max_new_tokens=256,
-            do_sample=True,
-            temperature=0.7,
-            top_p=0.9
-        )
+    # Prepare prompt
+    prompt = f"<image>\nUSER: {question}\nASSISTANT:"
     
-    # Decode response
-    response = processor.batch_decode(
-        outputs, 
-        skip_special_tokens=True
-    )[0].split("Answer:")[-1].strip()
+    # Run inference
+    args = type('Args', (), {
+        "model_name": model_name,
+        "query": prompt,
+        "conv_mode": None,
+        "image_file": image,
+        "sep": ",",
+        "temperature": 0.2,
+        "top_p": None,
+        "num_beams": 1,
+        "max_new_tokens": 512
+    })()
     
-    return response
+    return eval_model(args, tokenizer, model, image_processor)
 
 # Gradio interface
 with gr.Blocks() as demo:
-    gr.Markdown("# LLaVA-Med Medical Analysis (CPU)")
-    gr.Markdown("Official Microsoft LLaVA-Med 1.5-Mistral-7B implementation")
-    
+    gr.Markdown("# LLaVA-Med Medical Analysis")
     with gr.Row():
-        with gr.Column():
-            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="Clinical Analysis", interactive=False)
-
-    submit_btn.click(
-        fn=analyze_medical_image,
-        inputs=[image_input, question_input],
-        outputs=output_text
-    )
-
-demo.queue(max_size=5).launch()
+        gr.Image(type="pil", label="Input Image", source="upload", elem_id="image")
+        gr.Textbox(label="Question", placeholder="Ask about the medical image...")
+        gr.Textbox(label="Analysis Result", interactive=False)
+    
+    examples = [
+        ["examples/xray.jpg", "Are there any signs of pneumonia in this chest X-ray?"],
+        ["examples/mri.jpg", "What abnormalities are visible in this brain MRI?"]
+    ]
+    
+    gr.Examples(examples=examples, inputs=[image, question])
+    
+demo.launch()