app.py

app.py

@@ -3,34 +3,26 @@ from PIL import Image
 import torch
 import os
 
-# Your Hugging Face token for gated model access
-import os
-HF_TOKEN = os.getenv("HF_TOKEN")  # Secure load from Space secret!
-
-# Lingshu-7B imports
-from transformers import Qwen2_5_VLForConditionalGeneration, AutoProcessor
+# Load Hugging Face token securely from Space Secrets
+HF_TOKEN = os.getenv("HF_TOKEN")
 
-# MedGemma imports
-from transformers import pipeline
+from transformers import Qwen2_5_VLForConditionalGeneration, AutoProcessor, pipeline
 
-# Caching models and processors to avoid repeat loading
-lingshu_model, lingshu_processor = None, None
+lingshu_model = None
+lingshu_processor = None
 medgemma_pipe = None
 
-# Load Lingshu-7B
 def load_lingshu():
     global lingshu_model, lingshu_processor
     if lingshu_model is None or lingshu_processor is None:
         lingshu_model = Qwen2_5_VLForConditionalGeneration.from_pretrained(
             "lingshu-medical-mllm/Lingshu-7B",
             torch_dtype=torch.bfloat16,
-            attn_implementation="flash_attention_2",
             device_map="auto"
         )
         lingshu_processor = AutoProcessor.from_pretrained("lingshu-medical-mllm/Lingshu-7B")
     return lingshu_model, lingshu_processor
 
-# Load MedGemma-27B-IT with token for gated access
 def load_medgemma():
     global medgemma_pipe
     if medgemma_pipe is None:
@@ -44,9 +36,8 @@ def load_medgemma():
     return medgemma_pipe
 
 def inference(image, question, selected_model):
-    # Check image and question validity
     if image is None or question is None or question.strip() == "":
-        return "Please upload a medical image and enter your question/prompt."
+        return "Please upload a medical image and enter your question or prompt."
     if selected_model == "Lingshu-7B":
         model, processor = load_lingshu()
         messages = [