Update app.py

app.py

@@ -22,8 +22,6 @@ 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"
@@ -38,7 +36,6 @@ DESC = """
 - 参考：`_colormnet_tmp/ref/<视频名>/ref.png`  
 - 输出：`_colormnet_tmp/output/<视频名>/*.png`  
 - 合成视频：`_colormnet_tmp/<视频名>.mp4`
-
 **English**  
 Upload a **B&W video** and a **reference image**, then click “Start Coloring”.  
 This app runs **ZeroGPU scheduling in `app.py`** and calls `test.py` **in-process**.  
@@ -52,10 +49,8 @@ Temp workspace layout:
 PAPER = """
 ### 论文 / Paper
 **ECCV 2024 — ColorMNet: A Memory-based Deep Spatial-Temporal Feature Propagation Network for Video Colorization**  
-
 如果你喜欢这个项目，欢迎到 GitHub 点个 ⭐ Star：  
 **GitHub**: https://github.com/yyang181/colormnet
-
 **BibTeX 引用 / BibTeX Citation**
 ```bibtex
 @inproceedings{yang2024colormnet,
@@ -81,12 +76,9 @@ BADGES_HTML = """
 # ----------------- REFERENCE FRAME GUIDE (NO CROPPING) -----------------
 REF_GUIDE_MD = r"""
 ## 参考帧制作指南 / Reference Frame Guide
-
 **目的 / Goal**  
 为模型提供一张与你的视频关键帧在**姿态、光照、构图**尽量接近的**彩色参考图**，用来指导整段视频的着色风格与主体颜色。
-
 ---
-
 ### 中文步骤
 1. **挑帧**：从视频里挑一帧（或相近角度的照片），尽量与要着色的镜头在**姿态 / 光照 / 场景**一致。  
 2. **上色方式**：若你只有黑白参考图、但需要彩色参考，可用 **通义千问·图像编辑（Qwen-Image）**：  
@@ -97,15 +89,12 @@ REF_GUIDE_MD = r"""
    - 可按需多次编辑（如补充「衣服为复古蓝、肤色自然、不要锐化」）  
 3. **保存格式**：PNG/JPG 均可；推荐分辨率 ≥ **480px**（短边）。  
 4. **文件放置**：本应用会自动放置为 `ref/<视频名>/ref.png`。  
-
 **注意事项（Do/Don’t）**  
 - ✅ 主体清晰、颜色干净，不要过曝或强滤镜。  
 - ✅ 关键区域（衣服、皮肤、头发、天空等）颜色与目标风格一致。  
 - ❌ 不要更改几何结构（如人脸形状/姿态），**只修改颜色**。  
 - ❌ 避免文字、贴纸、重度风格化滤镜。
-
 ---
-
 ### English Steps
 1. **Pick a frame** (or a similar photo) that matches the target shot in **pose / lighting / composition**.  
 2. **Colorizing if your reference is B&W** — use **Qwen-Image (Image Editing)**:  
@@ -115,7 +104,6 @@ REF_GUIDE_MD = r"""
    - Iterate if needed (e.g., “vintage blue jacket, natural skin tone; avoid sharpening”).  
 3. **Format**: PNG/JPG; recommended short side ≥ **480px**.  
 4. **File placement**: The app will place it as `ref/<video_stem>/ref.png`.
-
 **Do / Don’t**
 - ✅ Clean subject and palette; avoid overexposure/harsh filters.  
 - ✅ Ensure key regions (clothes/skin/hair/sky) match the intended colors.  
@@ -301,25 +289,6 @@ 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,
@@ -328,7 +297,9 @@ 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:
@@ -422,10 +393,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 根
@@ -434,19 +405,17 @@ 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
-
-    log = inference_with_gpu(args_list)
+    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
 
     # 在合成 mp4 之前：清空 CUDA（防止显存占用）
     try: