Update app.py

app.py

@@ -233,60 +233,23 @@ def build_args_list_for_test(d16_batch_path: str,
     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
 
+  def ensure_correlation_extension_installed():
+      """Install spatial_correlation_sampler at runtime (needs CUDA)."""
+      try:
+          import spatial_correlation_sampler
+          return
+      except ImportError:
+          pass
+      print("[INFO] Installing spatial-correlation-sampler via pip (needs CUDA)...")
+      subprocess.run(
+          [sys.executable, "-m", "pip", "install", "spatial-correlation-sampler"],
+          check=True
+      )
+      importlib.invalidate_caches()
+      import spatial_correlation_sampler  # verify
+      print("[INFO] spatial-correlation-sampler installed successfully.")
+      
 # ----------------- GRADIO HANDLER -----------------
 @spaces.GPU(duration=120)   # 确保 CUDA 初始化在此函数体内
 def gradio_infer(