optimize zero-gpu usage

app.py

@@ -22,6 +22,8 @@ from PIL import Image
 import cv2
 import torch  # used for cuda sync & empty_cache
 
+import test
+
 # ----------------- BASIC INFO -----------------
 CHECKPOINT_URL = "https://github.com/yyang181/colormnet/releases/download/v0.1/DINOv2FeatureV6_LocalAtten_s2_154000.pth"
 CHECKPOINT_LOCAL = "DINOv2FeatureV6_LocalAtten_s2_154000.pth"
@@ -299,6 +301,25 @@ def ensure_correlation_extension_installed():
 
 # ----------------- GRADIO HANDLER -----------------
 @spaces.GPU(duration=120)   # 确保 CUDA 初始化在此函数体内
+def inference_with_gpu(args_list):
+    # <<< ZeroGPU 分配后：按需安装 Pytorch-Correlation-extension >>>
+    ensure_correlation_extension_installed()
+    # --------------------------------------------------------------
+
+    buf = io.StringIO()
+    try:
+        with redirect_stdout(buf), redirect_stderr(buf):
+            entry = getattr(test, "run_cli", None)
+            if entry is None or not callable(entry):
+                raise RuntimeError("test.py 未提供可调用的 run_cli(args_list) 接口。")
+            entry(args_list)
+        log = f"Args: {' '.join(args_list)}\n\n{buf.getvalue()}"
+    except Exception as e:
+        log = f"Args: {' '.join(args_list)}\n\n{buf.getvalue()}\n\nERROR: {e}"
+
+    return log
+
+
 def gradio_infer(
     debug_shapes,
     bw_video, ref_image,
@@ -307,9 +328,7 @@ def gradio_infer(
     num_proto, top_k, mem_every, deep_update,
     save_scores, flip, size, reverse
 ):
-    # <<< ZeroGPU 分配后：按需安装 Pytorch-Correlation-extension >>>
-    ensure_correlation_extension_installed()
-    # --------------------------------------------------------------
+
 
     # 1) 基本校验与临时目录
     if bw_video is None:
@@ -403,10 +422,10 @@ def gradio_infer(
     ensure_checkpoint()
 
     # 8) 同进程调用 test.py
-    try:
-        import test  # 确保 test.py 同目录且提供 run_cli(args_list)
-    except Exception as e:
-        return None, f"导入 test.py 失败 / Failed to import test.py：\n{e}"
+    # try:
+    #     import test  # 确保 test.py 同目录且提供 run_cli(args_list)
+    # except Exception as e:
+    #     return None, f"导入 test.py 失败 / Failed to import test.py：\n{e}"
 
     args_list = build_args_list_for_test(
         d16_batch_path=input_root,   # 指向 input_video 根
@@ -415,17 +434,19 @@ def gradio_infer(
         cfg=user_config
     )
 
-    buf = io.StringIO()
-    try:
-        with redirect_stdout(buf), redirect_stderr(buf):
-            entry = getattr(test, "run_cli", None)
-            if entry is None or not callable(entry):
-                raise RuntimeError("test.py 未提供可调用的 run_cli(args_list) 接口。")
-            entry(args_list)
-        log = f"Args: {' '.join(args_list)}\n\n{buf.getvalue()}"
-    except Exception as e:
-        log = f"Args: {' '.join(args_list)}\n\n{buf.getvalue()}\n\nERROR: {e}"
-        return None, log
+    # buf = io.StringIO()
+    # try:
+    #     with redirect_stdout(buf), redirect_stderr(buf):
+    #         entry = getattr(test, "run_cli", None)
+    #         if entry is None or not callable(entry):
+    #             raise RuntimeError("test.py 未提供可调用的 run_cli(args_list) 接口。")
+    #         entry(args_list)
+    #     log = f"Args: {' '.join(args_list)}\n\n{buf.getvalue()}"
+    # except Exception as e:
+    #     log = f"Args: {' '.join(args_list)}\n\n{buf.getvalue()}\n\nERROR: {e}"
+    #     return None, log
+
+    log = inference_with_gpu(args_list)
 
     # 在合成 mp4 之前：清空 CUDA（防止显存占用）
     try: