fix package error

app.py

@@ -12,6 +12,9 @@ import urllib.request
 from os import path
 import io
 from contextlib import redirect_stdout, redirect_stderr
+import subprocess
+import tempfile
+import importlib
 
 import gradio as gr
 import spaces
@@ -239,6 +242,61 @@ def build_args_list_for_test(d16_batch_path: str,
             args.extend([flag, str(v)])
     return args
 
+# ===== 新增：ZeroGPU 后按需安装 Pytorch-Correlation-extension =====
+_CORR_OK_STAMP = path.join(os.getcwd(), ".corr_ext_installed")
+
+def ensure_correlation_extension_installed():
+    """
+    在 ZeroGPU 分配后调用（位于 @spaces.GPU 函数体内）：
+    - 若已能 import 或存在本地 stamp，则直接返回
+    - 否则执行：
+        git clone https://github.com/ClementPinard/Pytorch-Correlation-extension.git
+        cd Pytorch-Correlation-extension && python setup.py install && cd ..
+    """
+    # 1) 尝试直接导入
+    try:
+        import spatial_correlation_sampler  # noqa: F401
+        return
+    except Exception:
+        pass
+    # 2) 之前成功过（打过 stamp）
+    if path.exists(_CORR_OK_STAMP):
+        return
+
+    repo_url = "https://github.com/ClementPinard/Pytorch-Correlation-extension.git"
+    workdir = tempfile.mkdtemp(prefix="corr_ext_")
+    repo_dir = path.join(workdir, "Pytorch-Correlation-extension")
+    try:
+        print("[INFO] Installing Pytorch-Correlation-extension ...")
+        # clone
+        subprocess.run(
+            ["git", "clone", "--depth", "1", repo_url],
+            cwd=workdir, check=True
+        )
+        # build & install
+        subprocess.run(
+            [sys.executable, "setup.py", "install"],
+            cwd=repo_dir, check=True
+        )
+        # 验证
+        importlib.invalidate_caches()
+        import spatial_correlation_sampler  # noqa: F401
+        # 打 stamp，避免下次重复
+        with open(_CORR_OK_STAMP, "w") as f:
+            f.write("ok")
+        print("[INFO] Pytorch-Correlation-extension installed successfully.")
+    except subprocess.CalledProcessError as e:
+        print(f"[WARN] Failed to build/install correlation extension: {e}")
+        print("You can still proceed if your pipeline doesn't use it.")
+    except Exception as e:
+        print(f"[WARN] Correlation extension install check failed: {e}")
+    finally:
+        # 清理临时目录
+        try:
+            shutil.rmtree(workdir, ignore_errors=True)
+        except Exception:
+            pass
+
 # ----------------- GRADIO HANDLER -----------------
 @spaces.GPU(duration=600)   # 确保 CUDA 初始化在此函数体内
 def gradio_infer(
@@ -249,6 +307,10 @@ 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:
         return None, "请上传黑白视频 / Please upload a B&W video."