app.py

app.py

@@ -0,0 +1,78 @@
+import gradio as gr
+from PIL import Image
+import torch
+
+# Lingshu-7B imports
+from transformers import Qwen2_5_VLForConditionalGeneration, AutoProcessor
+
+# MedGemma imports
+from transformers import pipeline
+
+def load_lingshu_model():
+    model = Qwen2_5_VLForConditionalGeneration.from_pretrained(
+        "lingshu-medical-mllm/Lingshu-7B",
+        torch_dtype=torch.bfloat16,
+        attn_implementation="flash_attention_2",
+        device_map="auto"
+    )
+    processor = AutoProcessor.from_pretrained("lingshu-medical-mllm/Lingshu-7B")
+    return model, processor
+
+def load_medgemma_model():
+    pipe = pipeline(
+        "image-text-to-text",
+        model="google/medgemma-27b-it",
+        torch_dtype=torch.bfloat16,
+        device="cuda"
+    )
+    return pipe
+
+lingshu_model, lingshu_processor = None, None
+medgemma_pipe = None
+
+def setup_models(selected_model):
+    global lingshu_model, lingshu_processor, medgemma_pipe
+    if selected_model == "Lingshu-7B" and lingshu_model is None:
+        lingshu_model, lingshu_processor = load_lingshu_model()
+    if selected_model == "MedGemma-27B-IT" and medgemma_pipe is None:
+        medgemma_pipe = load_medgemma_model()
+
+def med_ai_inference(img, prompt, model_type):
+    setup_models(model_type)
+    if model_type == "Lingshu-7B":
+        messages = [
+            {
+                "role": "user",
+                "content": [
+                    {"type": "image", "image": img},
+                    {"type": "text", "text": prompt}
+                ]
+            }
+        ]
+        text = lingshu_processor.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
+        inputs = lingshu_processor(text=[text], images=[img], padding=True, return_tensors="pt").to(lingshu_model.device)
+        with torch.no_grad():
+            generated_ids = lingshu_model.generate(**inputs, max_new_tokens=128)
+            trim_ids = generated_ids[:, inputs.input_ids.shape[1]:]
+            out_text = lingshu_processor.batch_decode(trim_ids, skip_special_tokens=True)
+        return out_text[0]
+    if model_type == "MedGemma-27B-IT":
+        # MedGemma expects messages
+        messages = [
+            {"role": "system", "content": [{"type": "text", "text": "You are a medical expert."}]},
+            {"role": "user", "content": [{"type": "text", "text": prompt}, {"type": "image", "image": img}]}
+        ]
+        res = medgemma_pipe(text=messages, max_new_tokens=200)
+        return res[0]["generated_text"][-1]["content"]
+
+with gr.Blocks() as demo:
+    gr.Markdown("# Medical AI Companion")
+    gr.Markdown("Upload a medical image, type your medical question or prompt, and select a model for automated report/answer.")
+    model_radio = gr.Radio(label="Model", choices=["Lingshu-7B", "MedGemma-27B-IT"], value="Lingshu-7B")
+    img_input = gr.Image(type="pil", label="Medical Image")
+    text_input = gr.Textbox(lines=2, label="Prompt", value="Describe this image.")
+    outbox = gr.Textbox(lines=10, label="AI Report / Answer", interactive=False)
+    run_btn = gr.Button("Analyze")
+    run_btn.click(med_ai_inference, [img_input, text_input, model_radio], outbox)
+
+demo.launch()