Update app.py

app.py

@@ -25,21 +25,21 @@ import spaces
 # def make_example(path):
 #     return [path, []]
 
+# --------- 全局状态 ---------
 MODEL = None
 DEVICE = torch.device("cpu")
 CUDA_READY = False
 
-# --------- 1) 先在 CPU 上加载权重（禁止这里触发 CUDA）---------
-def load_model_cpu(checkpoint_path):
-    # 请确保 inference.load_model 不会在里面 .to("cuda")
-    model, _ = load_model(checkpoint_path)   # 如果你里面会 .to(device)，请改掉
-    model = model.to("cpu")                  # 强制在 CPU
-    model.eval()
-    return model
+def load_model_cpu(checkpoint_path: str):
+    global MODEL, DEVICE
+    # 要求 inference.load_model 不要在内部 .to("cuda")
+    MODEL, _ = load_model(checkpoint_path)  # 或者直接返回 model
+    MODEL = MODEL.to("cpu")
+    MODEL.eval()
+    DEVICE = torch.device("cpu")
 
-MODEL = load_model_cpu("medsam_vit_b.pth")
+load_model_cpu("medsam_vit_b.pth")
 
-# --------- 2) 仅在 GPU 函数中迁移到 CUDA（不要在主进程里调用任何 CUDA）---------
 @spaces.GPU
 def prepare_cuda():
     global MODEL, DEVICE, CUDA_READY
@@ -47,8 +47,7 @@ def prepare_cuda():
         MODEL.to("cuda")
         DEVICE = torch.device("cuda")
         CUDA_READY = True
-        # 可选：做一次极小张量 warmup，仍在 GPU 函数里，安全
-        _ = torch.zeros(1, device=DEVICE)
+        _ = torch.zeros(1, device=DEVICE)  # 可选warmup
 
 def parse_first_bbox(bboxes):
     """
@@ -99,8 +98,8 @@ def segment(annot_value):
     box_np = np.array([[xmin, ymin, xmax, ymax]], dtype=float)
     box_1024 = box_np / np.array([W, H, W, H]) * 1024.0
 
-    embedding = get_embedding(model, img_np, device)
-    mask = run(model, embedding, box_1024, H, W)  # (H, W) 0/1
+    embedding = get_embedding(MODEL, img_np, DEVICE)
+    mask = run(MODEL, embedding, box_1024, H, W)  # (H, W) 0/1
 
     # 黑白 mask（白=前景）
     mask_rgb = np.stack([mask * 255] * 3, axis=-1).astype(np.uint8)