Update app.py

app.py

@@ -273,19 +273,22 @@ def run_inference(image, task, dataset_name, h_type, mode_type):
 #    成功 → 打印结果范围，并把输出图存到 Demo1_result.png
 #    失败 → 打印完整 traceback，方便定位错误
 # ══════════════════════════════════════════════════════════════
-_DEMO_IMG_PATH = "Demo1.png"
-print(f"\n{'='*60}")
-print(f"🧪 Startup inference test — {_DEMO_IMG_PATH} (device=cpu)")
-print(f"{'='*60}")
-try:
-    if not os.path.exists(_DEMO_IMG_PATH):
-        print(f"⚠️  {_DEMO_IMG_PATH} not found, skipping test.")
-    else:
+@spaces.GPU(duration=120)
+def _startup_gpu_test():
+    """必须包在 @spaces.GPU 里，ZeroGPU 才会分配真实 GPU 给这段代码"""
+    _DEMO_IMG_PATH = "Demo1.png"
+    print(f"\n{'='*60}")
+    print(f"🧪 Startup inference test — {_DEMO_IMG_PATH} (device=cuda)")
+    print(f"{'='*60}")
+    try:
+        if not os.path.exists(_DEMO_IMG_PATH):
+            print(f"⚠️  {_DEMO_IMG_PATH} not found, skipping test.")
+            return
+
         _test_img = Image.open(_DEMO_IMG_PATH)
         print(f"   Image size : {_test_img.size}, mode: {_test_img.mode}")
 
-        # 把模型组件移到 CPU（此时本来就在 CPU，仅做显式确认）
-        _model.to("cuda")
+        _model.to("cuda")   # 在 @spaces.GPU 上下文内移动，ZeroGPU 才认得
 
         _out_img, _info = _run_inference_core(
             _model, "cuda",
@@ -300,11 +303,18 @@ try:
         print(f"   Info       : {_info}")
         print(f"   Saved to   : Demo1_result.png")
 
-except Exception:
-    print("❌ Test FAILED — full traceback below:")
-    traceback.print_exc()
+    except Exception:
+        print("❌ Test FAILED — full traceback below:")
+        traceback.print_exc()
+
+    finally:
+        _model.to("cpu")           # 用完归还 GPU 内存
+        torch.cuda.empty_cache()
+        print(f"{'='*60}\n")
+
 
-print(f"{'='*60}\n")
+# 调用启动测试（此时已在 @spaces.GPU 上下文内，GPU 正常可用）
+_startup_gpu_test()
 
 
 # ══════════════════════════════════════════════════════════════