Update app.py

app.py

@@ -12,14 +12,20 @@ import spaces
 import torch
 from loguru import logger
 from PIL import Image
-from transformers import AutoProcessor, AutoModelForImageTextToText, TextIteratorStreamer
+from transformers import AutoProcessor, AutoModelForImageTextToText, TextIteratorStreamer, Qwen2_5_VLForConditionalGeneration
+from qwen_vl_utils import process_vision_info
 
-model_id = os.getenv("MODEL_ID", "google/medgemma-4b-it")
-processor = AutoProcessor.from_pretrained(model_id)
-model = AutoModelForImageTextToText.from_pretrained(
-    model_id, device_map="auto", torch_dtype=torch.bfloat16
+
+model = Qwen2_5_VLForConditionalGeneration.from_pretrained(
+    "lingshu-medical-mllm/Lingshu-32B",
+    torch_dtype=torch.bfloat16,
+    attn_implementation="flash_attention_2",
+    device_map="auto",
 )
 
+processor = AutoProcessor.from_pretrained("lingshu-medical-mllm/Lingshu-32B")
+
+
 MAX_NUM_IMAGES = int(os.getenv("MAX_NUM_IMAGES", "5"))
 
 
@@ -175,13 +181,19 @@ def run(message: dict, history: list[dict], system_prompt: str = "", max_new_tok
     messages.extend(process_history(history))
     messages.append({"role": "user", "content": process_new_user_message(message)})
 
-    inputs = processor.apply_chat_template(
-        messages,
-        add_generation_prompt=True,
-        tokenize=True,
-        return_dict=True,
+    # Preparation for inference
+    text = processor.apply_chat_template(
+        messages, tokenize=False, add_generation_prompt=True
+    )
+    image_inputs, video_inputs = process_vision_info(messages)
+    inputs = processor(
+        text=[text],
+        images=image_inputs,
+        videos=video_inputs,
+        padding=True,
         return_tensors="pt",
-    ).to(device=model.device, dtype=torch.bfloat16)
+    )
+    inputs = inputs.to(model.device)
 
     streamer = TextIteratorStreamer(processor, timeout=30.0, skip_prompt=True, skip_special_tokens=True)
     generate_kwargs = dict(