backup

.gitignore

@@ -11,7 +11,7 @@ Pytorch-Correlation-extension/
 result
 src/
 DINOv2FeatureV6_LocalAtten_s2_154000.pth
-example/
+_colormnet_tmp/
 
 # Byte-compiled / optimized / DLL files
 __pycache__/

app.py

@@ -1,432 +1,177 @@
-# app.py (aligned to main.py logic; keeps debug hooks; Gradio-safe DataLoader)
-# Inputs: (1) Black-and-white video (mp4/webm/avi) (2) Reference image (RGB)
-# Output: Colored video (mp4)
-#
-# Model checkpoint is HARD-CODED as required:
-# https://github.com/yyang181/colormnet/releases/download/v0.1/DINOv2FeatureV6_LocalAtten_s2_154000.pth
+# app.py — Gradio front-end that calls test.py IN-PROCESS (ZeroGPU-safe)
+# Folder layout per run (under TEMP_ROOT):
+#   input_video/<video_stem>/00000.png ...
+#   ref/<video_stem>/ref.png
+#   output/<video_stem>/*.png
+# Final mp4: TEMP_ROOT/<video_stem>.mp4
 
 import os
 import sys
 import shutil
-import subprocess
-import uuid
 import urllib.request
-import warnings
 from os import path
-from progressbar import progressbar
-import gc
-
-# # 1) 完全禁止 PyTorch 调用 NVML（ZeroGPU/MIG 下经常拿不到 NVML 句柄）
-# os.environ.setdefault("PYTORCH_NO_NVML", "1")
-# # 2) 用 cudaMallocAsync 后端，降低碎片/避免旧分配器的 NVML 路径
-# os.environ.setdefault(
-#     "PYTORCH_CUDA_ALLOC_CONF",
-#     "backend:cudaMallocAsync,expandable_segments:True,garbage_collection_threshold:0.9,max_split_size_mb:64"
-# )
-# # （可选）定位更准：同步执行
-# os.environ.setdefault("CUDA_LAUNCH_BLOCKING", "1")
-# warnings.filterwarnings("ignore", message="The detected CUDA version .* minor version mismatch")
-# warnings.filterwarnings("ignore", message="There are no g\\+\\+ version bounds defined for CUDA version.*")
-# warnings.filterwarnings("ignore", category=UserWarning, module="torch.utils.cpp_extension")
-# os.environ.setdefault("TORCH_COMPILE_DISABLE", "1")
-# os.environ.setdefault("MAX_JOBS", "1")
+import io
+from contextlib import redirect_stdout, redirect_stderr
 
 import gradio as gr
-import spaces  # ZeroGPU decorator
-import numpy as np
+import spaces
 from PIL import Image
 import cv2
-import traceback
 
-import torch
-import torch.nn.functional as F
-from torch.utils.data import DataLoader
-
-# ---- Project imports ----
-from inference.data.test_datasets import DAVISTestDataset_221128_TransColorization_batch
-from inference.data.mask_mapper import MaskMapper
-from model.network import ColorMNet
-from inference.inference_core import InferenceCore
-from dataset.range_transform import inv_lll2rgb_trans
-from skimage import color
-
-# ----------------- CONFIG -----------------
+# ----------------- 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"
 
 TITLE = "ColorMNet — ZeroGPU (CUDA-only) Video Colorization with Reference Image"
 DESC = """
 上传**黑白视频**与**参考图像**，点击“开始着色”。  
-本 Space **仅在 ZeroGPU（CUDA）** 上运行；若未分配到 GPU，会报错提示。  
-模型权重已固定链接（如需修改，请编辑 `CHECKPOINT_URL`）。  
-**数据集结构：**
-- 抽帧 -> `./colormnet_run_<UUID>/input_video/<视频名不含扩展>/00000.png...`
-- 参考图 -> `./colormnet_run_<UUID>/input_ref/<视频名不含扩展>/ref.png`
+此版本在 **app.py 中调度 ZeroGPU**，并**在同一进程**调用 `test.py` 的入口函数。  
+临时工作目录结构：
+- 抽帧：`_colormnet_tmp/input_video/<视频名>/00000.png ...`
+- 参考：`_colormnet_tmp/ref/<视频名>/ref.png`
+- 输出：`_colormnet_tmp/output/<视频名>/*.png`
+- 合成视频：`_colormnet_tmp/<视频名>.mp4`
 """
 
 # ----------------- TEMP WORKDIR -----------------
 TEMP_ROOT = path.join(os.getcwd(), "_colormnet_tmp")
+INPUT_DIR = "input_video"
+REF_DIR = "ref"
+OUTPUT_DIR = "output"
 
 def reset_temp_root():
     """每次运行前清空并重建临时工作目录。"""
     if path.isdir(TEMP_ROOT):
         shutil.rmtree(TEMP_ROOT, ignore_errors=True)
     os.makedirs(TEMP_ROOT, exist_ok=True)
+    for sub in (INPUT_DIR, REF_DIR, OUTPUT_DIR):
+        os.makedirs(path.join(TEMP_ROOT, sub), exist_ok=True)
 
-torch.set_grad_enabled(False)
-
-# ----------------- DEBUG (kept) -----------------
-def _enable_runtime_debug():
-    os.environ["CUDA_LAUNCH_BLOCKING"] = "1"          # 同步执行，定位准确
-    os.environ["TORCH_SHOW_CPP_STACKTRACES"] = "1"     # 显示 C++ 栈
-    os.environ["PYTORCH_JIT"] = "0"                    # 关闭 JIT
-    try:
-        torch.autograd.set_detect_anomaly(True)        # 捕捉无效 op/grad
-    except Exception:
-        pass
-
-# ----------------- PATH/DIR UTILS -----------------
-def ensure_clean_dir(d: str):
-    if path.exists(d):
-        if path.isdir(d):
-            return
-        else:
-            os.remove(d)
+def ensure_dir(d: str):
     os.makedirs(d, exist_ok=True)
 
-# ----------------- MISC UTILS -----------------
+# ----------------- CHECKPOINT (可选) -----------------
 def ensure_checkpoint():
-    if not path.exists(CHECKPOINT_LOCAL):
-        print(f"[INFO] Downloading checkpoint from: {CHECKPOINT_URL}")
-        urllib.request.urlretrieve(CHECKPOINT_URL, CHECKPOINT_LOCAL)
-        print("[INFO] Checkpoint downloaded:", CHECKPOINT_LOCAL)
-
-def detach_to_cpu(x: torch.Tensor) -> torch.Tensor:
-    return x.detach().cpu()
-
-def tensor_to_np_float(image: torch.Tensor) -> np.ndarray:
-    image_np = image.numpy().astype("float32")
-    return image_np
-
-def lab2rgb_transform_PIL(mask: torch.Tensor) -> np.ndarray:
-    mask_d = detach_to_cpu(mask)
-    mask_d = inv_lll2rgb_trans(mask_d)
-    im = tensor_to_np_float(mask_d)
-    if len(im.shape) == 3:
-        im = im.transpose((1, 2, 0))
-    else:
-        im = im[:, :, None]
-    im = color.lab2rgb(im)
-    return im.clip(0, 1)
+    """若 test.py 会在当前目录加载权重，可提前预下载，避免首次拉取超时。"""
+    try:
+        if not path.exists(CHECKPOINT_LOCAL):
+            print(f"[INFO] Downloading checkpoint from: {CHECKPOINT_URL}")
+            urllib.request.urlretrieve(CHECKPOINT_URL, CHECKPOINT_LOCAL)
+            print("[INFO] Checkpoint downloaded:", CHECKPOINT_LOCAL)
+    except Exception as e:
+        print(f"[WARN] 预下载权重失败（首次推理会再试）: {e}")
 
-# ---------- extract frames: dataset-root/<video_stem>/00000.png ----------
-def video_to_dataset_root(video_path: str, dataset_root: str):
+# ----------------- VIDEO UTILS -----------------
+def video_to_frames_dir(video_path: str, frames_dir: str):
     """
-    将单个视频抽帧到 dataset_root/<video_stem>/00000.png...
-    返回: (subdir_path, video_stem, width, height, fps, frame_count)
+    抽帧到 frames_dir/00000.png ...
+    返回: (w, h, fps, n_frames)
     """
-    ensure_clean_dir(dataset_root)
-    basename = path.basename(video_path)
-    stem, _ = path.splitext(basename)
-    subdir = path.join(dataset_root, stem)
-    ensure_clean_dir(subdir)
-
+    ensure_dir(frames_dir)
     cap = cv2.VideoCapture(video_path)
     assert cap.isOpened(), f"Cannot open video: {video_path}"
-
-    fps = cap.get(cv2.CAP_PROP_FPS)
-    if not fps or fps <= 0:
-        fps = 25.0
-
+    fps = cap.get(cv2.CAP_PROP_FPS) or 25.0
     idx = 0
     w = h = None
-
     while True:
         ret, frame = cap.read()
         if not ret:
             break
         if frame is None:
             continue
-
         h, w = frame.shape[:2]
-        out_path = path.join(subdir, f"{idx:05d}.png")
-
-        parent = path.dirname(out_path)
-        if not path.isdir(parent):
-            if path.exists(parent):
-                os.remove(parent)
-            os.makedirs(parent, exist_ok=True)
-
+        out_path = path.join(frames_dir, f"{idx:05d}.png")
         ok = cv2.imwrite(out_path, frame)
         if not ok:
             raise RuntimeError(f"写入抽帧失败: {out_path}")
         idx += 1
-
     cap.release()
     if idx == 0:
         raise RuntimeError("Input video has no readable frames.")
-
-    return subdir, path.splitext(path.basename(video_path))[0], w, h, fps, idx
-
-# ---------- place ref image into ref_root/<video_stem>/ref.png ----------
-def ref_to_dataset_root(ref_image_path: str, ref_root: str, video_stem: str):
-    ensure_clean_dir(ref_root)
-    subdir = path.join(ref_root, video_stem)
-    ensure_clean_dir(subdir)
-
-    img = Image.open(ref_image_path).convert("RGB")
-    out_path = path.join(subdir, "ref.png")
-    img.save(out_path)
-    return subdir
+    return w, h, fps, idx
 
 def encode_frames_to_video(frames_dir: str, out_path: str, fps: float):
     frames = sorted([f for f in os.listdir(frames_dir) if f.lower().endswith(".png")])
-    assert len(frames) > 0, "No frames to encode."
-
+    if not frames:
+        raise RuntimeError(f"No frames found in {frames_dir}")
     first = cv2.imread(path.join(frames_dir, frames[0]))
+    if first is None:
+        raise RuntimeError(f"Failed to read first frame {frames[0]}")
     h, w = first.shape[:2]
-
     fourcc = cv2.VideoWriter_fourcc(*"mp4v")
     vw = cv2.VideoWriter(out_path, fourcc, fps, (w, h))
     for f in frames:
         img = cv2.imread(path.join(frames_dir, f))
+        if img is None:
+            continue
         vw.write(img)
     vw.release()
 
-# ----------------- MAIN PIPELINE (CUDA-only) -----------------
-def run_pipeline_cuda(bw_video_path: str, ref_image_path: str, user_config: dict, debug_shapes: bool) -> str:
-    print(bw_video_path, ref_image_path)
-    if not torch.cuda.is_available():
-        raise RuntimeError("未检测到 GPU。此 Space 仅支持 ZeroGPU (CUDA)。")
-
-    if debug_shapes:
-        _enable_runtime_debug()
-
-    ensure_checkpoint()
-
-    DEVICE = torch.device("cuda")
-
-    # Workspace in TEMP_ROOT
-    base_run_dir = path.join(TEMP_ROOT, f"colormnet_run_{uuid.uuid4().hex}")
-    input_video_root = path.join(base_run_dir, "input_video")
-    input_ref_root   = path.join(base_run_dir, "input_ref")
-    output_dir       = path.join(base_run_dir, "result")
-
-    for p in (base_run_dir, input_video_root, input_ref_root, output_dir):
-        ensure_clean_dir(p)
-
-    # 1) 抽帧（把抽帧输出到临时目录中）
-    vid_subdir, vid_stem, w, h, fps, n_frames = video_to_dataset_root(bw_video_path, input_video_root)
-    assert n_frames > 0, "Input video has no frames."
-
-    # 2) 参考图（存到临时目录）
-    _ = ref_to_dataset_root(ref_image_path, input_ref_root, vid_stem)
-
-    # 3) 配置（字段与 main.py 一致；值从 UI 合并）
-    default_config = {
-        "FirstFrameIsNotExemplar": True,
-        "d16_batch_path": "input",   # parity only
-        "ref_path": "ref",           # parity only
-        "output": "result",          # parity only
-        "generic_path": None,
-        "dataset": "D16_batch",
-        "split": "val",
-        "save_all": True,
-        "benchmark": False,
-        "disable_long_term": False,
-        "max_mid_term_frames": 10,
-        "min_mid_term_frames": 5,
-        "max_long_term_elements": 10000,
-        "num_prototypes": 128,
-        "top_k": 30,
-        "mem_every": 5,
-        "deep_update_every": -1,
-        "save_scores": False,
-        "flip": False,
-        "size": -1,
-        "reverse": False,
-    }
-    config = {**default_config, **(user_config or {})}
-    config["enable_long_term"] = not config["disable_long_term"]
-
-    # 4) 构建数据集（只选本视频 reader）
-    meta_dataset = DAVISTestDataset_221128_TransColorization_batch(
-        input_video_root, imset=input_ref_root, size=config["size"]
-    )
-    meta_loader = meta_dataset.get_datasets()
-
-    # 输出路径规则（与 main.py 一致）
-    is_youtube = str(config["dataset"]).startswith("Y")
-    is_davis   = str(config["dataset"]).startswith("D")
-    is_lv      = str(config["dataset"]).startswith("LV")
-
-    app_output_root = output_dir
-    if is_youtube or config["save_scores"]:
-        out_path = path.join(app_output_root, "Annotations")
-    else:
-        out_path = app_output_root
-
-    # 5) 模型（保持 app 的 URL 权重加载方式）
-    network = ColorMNet(config, CHECKPOINT_LOCAL).to(DEVICE).eval()
-    model_weights = torch.load(CHECKPOINT_LOCAL, map_location="cuda")
-    network.load_weights(model_weights, init_as_zero_if_needed=True)
-
-    total_process_time = 0.0
-    total_frames = 0
-
-    for vid_reader in progressbar(meta_loader, max_value=len(meta_dataset), redirect_stdout=True):
-        # 6) 推理（逐帧；内部逻辑与 main.py 对齐；保留调试打印）
-        # Gradio/Spaces 环境禁止子进程：num_workers=0（否则会触发 daemonic processes 错误）
-        loader = DataLoader(vid_reader, batch_size=1, shuffle=False, num_workers=0, pin_memory=True)
-        vid_name = vid_reader.vid_name
-        vid_length = len(loader)
-
-        # 长时记忆触发逻辑：按 main.py 原样（无除零保护）
-        config['enable_long_term_count_usage'] = (
-            config['enable_long_term'] and
-            (vid_length
-                / (config['max_mid_term_frames'] - config['min_mid_term_frames'])
-                * config['num_prototypes'])
-            >= config['max_long_term_elements']
-        )
-
-        mapper = MaskMapper()
-        processor = InferenceCore(network, config=config)
-        first_mask_loaded = False
-
-        for ti, data in enumerate(loader):
-            try:
-                with torch.cuda.amp.autocast(enabled=not config["benchmark"]):
-                    rgb = data['rgb'].cuda()[0]
-                    msk = data.get('mask')
-                    if not config['FirstFrameIsNotExemplar']:
-                        msk = msk[:, 1:3, :, :] if msk is not None else None
-
-                    info = data['info']
-                    frame = info['frame'][0]
-                    shape = info['shape']
-                    need_resize = info['need_resize'][0]
-
-                    if debug_shapes:
-                        print(f"[Loop] frame={ti} rgb={tuple(rgb.shape)} "
-                            f"msk={None if msk is None else tuple(msk.shape)}", flush=True)
-
-                    # timing 与 main.py 一致
-                    start = torch.cuda.Event(enable_timing=True)
-                    end = torch.cuda.Event(enable_timing=True)
-                    start.record()
-
-                    if not first_mask_loaded:
-                        if msk is not None:
-                            first_mask_loaded = True
-                        else:
-                            continue
-
-                    if config['flip']:
-                        rgb = torch.flip(rgb, dims=[-1])
-                        msk = torch.flip(msk, dims=[-1]) if msk is not None else None
-
-                    if msk is not None:
-                        msk = torch.Tensor(msk[0]).cuda()
-                        if need_resize:
-                            msk = vid_reader.resize_mask(msk.unsqueeze(0))[0]
-                        processor.set_all_labels(list(range(1, 3)))
-                        labels = range(1, 3)
-                    else:
-                        labels = None
-
-                    if config['FirstFrameIsNotExemplar']:
-                        prob = processor.step_AnyExemplar(
-                            rgb,
-                            msk[:1, :, :].repeat(3, 1, 1) if msk is not None else None,
-                            msk[1:3, :, :] if msk is not None else None,
-                            labels,
-                            end=(ti == vid_length - 1)
-                        )
-                    else:
-                        prob = processor.step(rgb, msk, labels, end=(ti == vid_length - 1))
-
-                    if need_resize:
-                        prob = F.interpolate(prob.unsqueeze(1), shape, mode='bilinear', align_corners=False)[:, 0]
-
-                    end.record()
-                    torch.cuda.synchronize()
-                    total_process_time += (start.elapsed_time(end) / 1000.0)
-                    total_frames += 1
-
-                    if config['flip']:
-                        prob = torch.flip(prob, dims=[-1])
-
-                    if debug_shapes:
-                        try:
-                            print(f"[Loop] prob={tuple(prob.shape)}", flush=True)
-                        except Exception:
-                            pass
-
-                    if config['save_scores']:
-                        prob = (prob.detach().cpu().numpy() * 255).astype(np.uint8)
-
-                    if config['save_all'] or info['save'][0]:
-                        this_out_path = path.join(out_path, vid_name)
-                        os.makedirs(this_out_path, exist_ok=True)
-
-                        out_mask_final = lab2rgb_transform_PIL(torch.cat([rgb[:1, :, :], prob], dim=0))
-                        out_mask_final = (out_mask_final * 255).astype(np.uint8)
-                        Image.fromarray(out_mask_final).save(os.path.join(this_out_path, frame[:-4] + '.png'))
-
-            except Exception as _e:
-                # 保留完整 traceback，方便定位
-                raise RuntimeError("FRAME_ERROR:\n" + traceback.format_exc())
-
-    if total_process_time > 0:
-        print(f'Total processing time: {total_process_time}')
-        print(f'Total processed frames: {total_frames}')
-        print(f'FPS: {total_frames / total_process_time}')
-        print(f'Max allocated memory (MB): {torch.cuda.max_memory_allocated() / (2**20)}')
-
-    # 7) 合成 mp4（按 main.py 的 out_path 规则找帧目录）
-    frames_dir = path.join(out_path, vid_stem if path.isdir(path.join(out_path, vid_stem)) else vid_name)
-    if not path.isdir(frames_dir):
-        subs = [d for d in os.listdir(out_path) if path.isdir(path.join(out_path, d))]
-        if len(subs) == 1:
-            frames_dir = path.join(out_path, subs[0])
+# ----------------- CLI MAPPING -----------------
+CONFIG_TO_CLI = {
+    "FirstFrameIsNotExemplar": "--FirstFrameIsNotExemplar",  # bool
+    "dataset": "--dataset",
+    "split": "--split",
+    "save_all": "--save_all",                                # bool
+    "benchmark": "--benchmark",                              # bool
+    "disable_long_term": "--disable_long_term",              # bool
+    "max_mid_term_frames": "--max_mid_term_frames",
+    "min_mid_term_frames": "--min_mid_term_frames",
+    "max_long_term_elements": "--max_long_term_elements",
+    "num_prototypes": "--num_prototypes",
+    "top_k": "--top_k",
+    "mem_every": "--mem_every",
+    "deep_update_every": "--deep_update_every",
+    "save_scores": "--save_scores",                          # bool
+    "flip": "--flip",                                        # bool
+    "size": "--size",
+    "reverse": "--reverse",                                  # bool
+}
+
+def build_args_list_for_test(d16_batch_path: str,
+                             out_path: str,
+                             ref_root: str,
+                             cfg: dict):
+    """
+    构造传给 test.run_cli(args_list) 的参数列表。
+    - 必传：--d16_batch_path <input_video_root>、--ref_path <ref_root>、--output <output_root>
+    """
+    args = [
+        "--d16_batch_path", d16_batch_path,
+        "--ref_path", ref_root,
+        "--output", out_path,
+    ]
+    for k, v in cfg.items():
+        if k not in CONFIG_TO_CLI:
+            continue
+        flag = CONFIG_TO_CLI[k]
+        if isinstance(v, bool):
+            if v:
+                args.append(flag)          # store_true
+        elif v is None:
+            continue
         else:
-            frames_dir = path.join(output_dir, vid_stem)
-
-    colored_mp4 = path.join(base_run_dir, "colored_output.mp4")
-    encode_frames_to_video(frames_dir, colored_mp4, fps=fps)
-
-    # 8) 输出视频到 CWD（只保留最终文件）
-    final_mp4 = path.join(os.getcwd(), "result.mp4")
-    shutil.move(colored_mp4, final_mp4)
+            args.extend([flag, str(v)])
+    return args
 
-    # 清理本次 run 的中间目录；(注：上传的原视频/参考帧位于 TEMP_ROOT，将在下次运行开头被 reset_temp_root 清掉)
-    shutil.rmtree(base_run_dir, ignore_errors=True)
-
-    return final_mp4
-
-# ----------------- GRADIO HANDLERS -----------------
-@spaces.GPU(duration=600)
+# ----------------- GRADIO HANDLER -----------------
+@spaces.GPU(duration=160)   # 确保 CUDA 初始化在此函数体内
 def gradio_infer(
-    debug_shapes,  # 调���开关（保留）
+    debug_shapes,
     bw_video, ref_image,
     first_not_exemplar, dataset, split, save_all, benchmark,
     disable_long_term, max_mid, min_mid, max_long,
     num_proto, top_k, mem_every, deep_update,
-    save_scores, flip, size, reverse  # 新增
+    save_scores, flip, size, reverse
 ):
-    if not torch.cuda.is_available():
-        return None, "ZeroGPU 未分配到 GPU，请重试（或检查 Space 硬件是否为 ZeroGPU）。"
-
+    # 1) 基本校验与临时目录
     if bw_video is None:
         return None, "请上传黑白视频。"
     if ref_image is None:
         return None, "请上传参考图像。"
-
-    # —— 每次运行先重置临时目录 —— #
     reset_temp_root()
 
-    # Video path -> 拷贝到临时目录
+    # 2) 解析视频源路径 & 目标 <video_stem>
     if isinstance(bw_video, dict) and "name" in bw_video:
         src_video_path = bw_video["name"]
     elif isinstance(bw_video, str):
@@ -434,28 +179,40 @@ def gradio_infer(
     else:
         return None, "无法读取视频输入。"
 
-    tmp_video_ext = path.splitext(src_video_path)[1] or ".mp4"
-    tmp_video_path = path.join(TEMP_ROOT, "input_video" + tmp_video_ext)
+    video_stem = path.splitext(path.basename(src_video_path))[0]
+
+    # 3) 生成临时路径
+    input_root = path.join(TEMP_ROOT, INPUT_DIR)     # _colormnet_tmp/input_video
+    ref_root   = path.join(TEMP_ROOT, REF_DIR)       # _colormnet_tmp/ref
+    output_root= path.join(TEMP_ROOT, OUTPUT_DIR)    # _colormnet_tmp/output
+    input_frames_dir = path.join(input_root, video_stem)
+    ref_dir = path.join(ref_root, video_stem)
+    out_frames_dir = path.join(output_root, video_stem)
+    for d in (input_root, ref_root, output_root, input_frames_dir, ref_dir, out_frames_dir):
+        ensure_dir(d)
+
+    # 4) 抽帧 -> input_video/<stem>/
     try:
-        shutil.copy2(src_video_path, tmp_video_path)
+        _w, _h, fps, _n = video_to_frames_dir(src_video_path, input_frames_dir)
     except Exception as e:
-        return None, f"复制视频到临时目录失败：{e}"
+        return None, f"抽帧失败：\n{e}"
 
-    # Ref path -> 保存/拷贝到临时目录
-    tmp_ref_path = path.join(TEMP_ROOT, "ref.png")
+    # 5) 参考帧 -> ref/<stem>/ref.png
+    ref_png_path = path.join(ref_dir, "ref.png")
     if isinstance(ref_image, Image.Image):
         try:
-            ref_image.save(tmp_ref_path)
+            ref_image.save(ref_png_path)
         except Exception as e:
-            return None, f"保存参考图像到临时目录失败：{e}"
+            return None, f"保存参考图像失败：\n{e}"
     elif isinstance(ref_image, str):
         try:
-            shutil.copy2(ref_image, tmp_ref_path)
+            shutil.copy2(ref_image, ref_png_path)
         except Exception as e:
-            return None, f"复制参考图像到临时目录失败：{e}"
+            return None, f"复制参考图像失败：\n{e}"
     else:
         return None, "无法读取参考图像输入。"
 
+    # 6) 收集 UI 配置
     default_config = {
         "FirstFrameIsNotExemplar": True,
         "dataset": "D16_batch",
@@ -473,8 +230,8 @@ def gradio_infer(
         "save_scores": False,
         "flip": False,
         "size": -1,
+        "reverse": False,
     }
-
     user_config = {
         "FirstFrameIsNotExemplar": bool(first_not_exemplar) if first_not_exemplar is not None else default_config["FirstFrameIsNotExemplar"],
         "dataset": str(dataset) if dataset else default_config["dataset"],
@@ -492,71 +249,92 @@ def gradio_infer(
         "save_scores": bool(save_scores) if save_scores is not None else default_config["save_scores"],
         "flip": bool(flip) if flip is not None else default_config["flip"],
         "size": int(size) if size is not None else default_config["size"],
-        "reverse": bool(reverse) if reverse is not None else False,
+        "reverse": bool(reverse) if reverse is not None else default_config["reverse"],
     }
 
+    # 7) 预下载权重（可选）
+    ensure_checkpoint()
+
+    # 8) 同进程调用 test.py
     try:
-        out_mp4 = run_pipeline_cuda(
-            tmp_video_path, tmp_ref_path, user_config, debug_shapes=bool(debug_shapes)
-        )
-        return out_mp4, "完成 ✅"
-    except subprocess.CalledProcessError as e:
-        # 出错也可以顺手清一下临时目录（可选）
-        try: shutil.rmtree(TEMP_ROOT, ignore_errors=True)
-        except: pass
-        return None, f"运行时错误：\n{e}"
+        import test  # 确保 test.py 同目录且有 run_cli 函数
     except Exception as e:
-        try: shutil.rmtree(TEMP_ROOT, ignore_errors=True)
-        except: pass
-        return None, f"{e}"
+        return None, f"导入 test.py 失败：\n{e}"
+
+    args_list = build_args_list_for_test(
+        d16_batch_path=input_root,   # 指向 input_video 根
+        out_path=output_root,        # 指向 output 根（test.py 写 output/<stem>/*.png）
+        ref_root=ref_root,           # 指向 ref 根（test.py 读 ref/<stem>/ref.png��
+        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
+
+    # 9) 合成 mp4：从 output/<stem>/ 帧合成 -> TEMP_ROOT/<stem>.mp4
+    out_frames = path.join(output_root, video_stem)
+    if not path.isdir(out_frames):
+        return None, f"未找到输出帧目录：{out_frames}\n\n{log}"
+    final_mp4 = path.join(TEMP_ROOT, f"{video_stem}.mp4")
+    try:
+        encode_frames_to_video(out_frames, final_mp4, fps=fps)
+    except Exception as e:
+        return None, f"合成视频失败：\n{e}\n\n{log}"
+
+    return final_mp4, f"完成 ✅\n\n{log}"
 
 # ----------------- UI -----------------
 with gr.Blocks() as demo:
     gr.Markdown(f"# {TITLE}")
     gr.Markdown(DESC)
 
-    debug_shapes = gr.Checkbox(label="调试日志（打印形状与完整Traceback）", value=False)
+    debug_shapes = gr.Checkbox(label="调试日志（仅用于显示更完整日志）", value=False)
 
     with gr.Row():
         inp_video = gr.Video(label="黑白视频（mp4/webm/avi）", interactive=True)
         inp_ref = gr.Image(label="参考图像（RGB）", type="pil")
-
         gr.Examples(
             label="示例输入",
-            examples=[
-                ["./example/4.mp4", "./example/4.png"],
-            ],
+            examples=[["./example/4.mp4", "./example/4.png"]],
             inputs=[inp_video, inp_ref],
-            # 不缓存，避免把推理结果当静态示例
             cache_examples=False,
         )
 
-    with gr.Accordion("高级参数设置（与 main.py 对齐）", open=False):
+    with gr.Accordion("高级参数设置（传给 test.py）", open=False):
         with gr.Row():
-            first_not_exemplar = gr.Checkbox(label="FirstFrameIsNotExemplar", value=True)
-            reverse = gr.Checkbox(label="reverse", value=False)
-            dataset = gr.Textbox(label="dataset", value="D16_batch")
-            split = gr.Textbox(label="split", value="val")
-            save_all = gr.Checkbox(label="save_all", value=True)
-            benchmark = gr.Checkbox(label="benchmark", value=False)
+            first_not_exemplar = gr.Checkbox(label="FirstFrameIsNotExemplar (--FirstFrameIsNotExemplar)", value=True)
+            reverse = gr.Checkbox(label="reverse (--reverse)", value=False)
+            dataset = gr.Textbox(label="dataset (--dataset)", value="D16_batch")
+            split = gr.Textbox(label="split (--split)", value="val")
+            save_all = gr.Checkbox(label="save_all (--save_all)", value=True)
+            benchmark = gr.Checkbox(label="benchmark (--benchmark)", value=False)
         with gr.Row():
-            disable_long_term = gr.Checkbox(label="disable_long_term", value=False)
-            max_mid = gr.Number(label="max_mid_term_frames", value=10, precision=0)
-            min_mid = gr.Number(label="min_mid_term_frames", value=5, precision=0)
-            max_long = gr.Number(label="max_long_term_elements", value=10000, precision=0)
-            num_proto = gr.Number(label="num_prototypes", value=128, precision=0)
+            disable_long_term = gr.Checkbox(label="disable_long_term (--disable_long_term)", value=False)
+            max_mid = gr.Number(label="max_mid_term_frames (--max_mid_term_frames)", value=10, precision=0)
+            min_mid = gr.Number(label="min_mid_term_frames (--min_mid_term_frames)", value=5, precision=0)
+            max_long = gr.Number(label="max_long_term_elements (--max_long_term_elements)", value=10000, precision=0)
+            num_proto = gr.Number(label="num_prototypes (--num_prototypes)", value=128, precision=0)
         with gr.Row():
-            top_k = gr.Number(label="top_k", value=30, precision=0)
-            mem_every = gr.Number(label="mem_every", value=5, precision=0)
-            deep_update = gr.Number(label="deep_update_every", value=-1, precision=0)
-            save_scores = gr.Checkbox(label="save_scores", value=False)
-            flip = gr.Checkbox(label="flip", value=False)
-            size = gr.Number(label="size", value=-1, precision=0)
-
-    run_btn = gr.Button("开始着色（ZeroGPU 推理）")
+            top_k = gr.Number(label="top_k (--top_k)", value=30, precision=0)
+            mem_every = gr.Number(label="mem_every (--mem_every)", value=5, precision=0)
+            deep_update = gr.Number(label="deep_update_every (--deep_update_every)", value=-1, precision=0)
+            save_scores = gr.Checkbox(label="save_scores (--save_scores)", value=False)
+            flip = gr.Checkbox(label="flip (--flip)", value=False)
+            size = gr.Number(label="size (--size)", value=-1, precision=0)
+
+    run_btn = gr.Button("开始着色（同进程调用 test.py）")
     with gr.Row():
         out_video = gr.Video(label="输出视频（着色结果）")
-        status = gr.Textbox(label="状态 / 调试输出", interactive=False, lines=12)
+        status = gr.Textbox(label="状态 / 日志输出（test.py stdout/stderr）", interactive=False, lines=16)
 
     run_btn.click(
         fn=gradio_infer,
@@ -566,7 +344,7 @@ with gr.Blocks() as demo:
             first_not_exemplar, dataset, split, save_all, benchmark,
             disable_long_term, max_mid, min_mid, max_long,
             num_proto, top_k, mem_every, deep_update,
-            save_scores, flip, size, reverse   # reverse 已接入
+            save_scores, flip, size, reverse
         ],
         outputs=[out_video, status]
     )
@@ -576,5 +354,4 @@ if __name__ == "__main__":
         ensure_checkpoint()
     except Exception as e:
         print(f"[WARN] 预下载权重失败（首次推理会再试）: {e}")
-
-    demo.queue().launch(server_name="0.0.0.0", server_port=7860)
+    demo.queue(max_size=32).launch(server_name="0.0.0.0", server_port=7860)

inference/data/test_datasets.py

@@ -1,36 +1,116 @@
 import os
 from os import path
-import json
 
-from inference.data.video_reader import VideoReader_221128_TransColorization
+from torch.utils.data.dataset import Dataset
+from torchvision import transforms
+from torchvision.transforms import InterpolationMode
+import torch.nn.functional as Ff
+from PIL import Image
+import numpy as np
 
-class DAVISTestDataset_221128_TransColorization_batch:
-    def __init__(self, data_root, imset='2017/val.txt', size=-1, args=None):
-        self.image_dir = data_root
-        self.mask_dir = imset
-        self.size_dir = data_root
+from dataset.range_transform import im_normalization, im_rgb2lab_normalization, ToTensor, RGB2Lab
+
+class VideoReader_221128_TransColorization(Dataset):
+    """
+    This class is used to read a video, one frame at a time
+    """
+    def __init__(self, vid_name, image_dir, mask_dir, size=-1, to_save=None, use_all_mask=False, size_dir=None, args=None):
+        """
+        image_dir - points to a directory of jpg images
+        mask_dir - points to a directory of png masks
+        size - resize min. side to size. Does nothing if <0.
+        to_save - optionally contains a list of file names without extensions 
+            where the segmentation mask is required
+        use_all_mask - when true, read all available mask in mask_dir.
+            Default false. Set to true for YouTubeVOS validation.
+        """
+        self.vid_name = vid_name
+        self.image_dir = image_dir
+        self.mask_dir = mask_dir
+        self.to_save = to_save
+        self.use_all_mask = use_all_mask
+        # print('use_all_mask', use_all_mask);assert 1==0
+        if size_dir is None:
+            self.size_dir = self.image_dir
+        else:
+            self.size_dir = size_dir
+
+        # flag_reverse = args.getattr('reverse', False) if args is not None else False
+        flag_reverse = False
+        self.frames = [img for img in sorted(os.listdir(self.image_dir), reverse=flag_reverse) if (img.endswith('.jpg') or img.endswith('.png')) and not img.startswith('.')]
+        self.palette = Image.open(path.join(mask_dir, sorted([msk for msk in os.listdir(mask_dir) if not msk.startswith('.')])[0])).getpalette()
+        self.first_gt_path = path.join(self.mask_dir, sorted([msk for msk in os.listdir(self.mask_dir) if not msk.startswith('.')])[0])
+        self.suffix = self.first_gt_path.split('.')[-1]
+
+        if size < 0:
+            self.im_transform = transforms.Compose([
+                RGB2Lab(),
+                ToTensor(),
+                im_rgb2lab_normalization,
+            ])
+        else:
+            self.im_transform = transforms.Compose([
+                transforms.ToTensor(),
+                im_normalization,
+                transforms.Resize(size, interpolation=InterpolationMode.BILINEAR),
+            ])
         self.size = size
 
-        self.vid_list =  [clip_name for clip_name in sorted(os.listdir(data_root)) if clip_name != '.DS_Store' and not clip_name.startswith('.')]
-        self.ref_img_list = [clip_name for clip_name in sorted(os.listdir(imset)) if clip_name != '.DS_Store' and not clip_name.startswith('.')]
 
-        self.args = args
+    def __getitem__(self, idx):
+        frame = self.frames[idx]
+        info = {}
+        data = {}
+        info['frame'] = frame
+        info['vid_name'] = self.vid_name
+        info['save'] = (self.to_save is None) or (frame[:-4] in self.to_save)
+
+        im_path = path.join(self.image_dir, frame)
+        img = Image.open(im_path).convert('RGB')
+
+        if self.image_dir == self.size_dir:
+            shape = np.array(img).shape[:2]
+        else:
+            size_path = path.join(self.size_dir, frame)
+            size_im = Image.open(size_path).convert('RGB')
+            shape = np.array(size_im).shape[:2]
+
+        gt_path = path.join(self.mask_dir, sorted(os.listdir(self.mask_dir))[idx]) if idx < len(os.listdir(self.mask_dir)) else None 
+
+        img = self.im_transform(img)
+        img_l = img[:1,:,:]
+        img_lll = img_l.repeat(3,1,1)
+
+        load_mask = self.use_all_mask or (gt_path == self.first_gt_path)
+        if load_mask and path.exists(gt_path):
+            mask = Image.open(gt_path).convert('RGB')
+            
+            # 用 PIL 先 resize 成和 img 尺寸一致
+            mask = mask.resize((img.shape[2], img.shape[1]), Image.BILINEAR)
+
+            mask = self.im_transform(mask)
+
+            # keep L channel of reference image in case First frame is not exemplar
+            # mask_ab = mask[1:3,:,:]
+            # data['mask'] = mask_ab
+            data['mask'] = mask
+
+        info['shape'] = shape
+        info['need_resize'] = not (self.size < 0)
+        data['rgb'] = img_lll
+        data['info'] = info
 
-        # print(lst, len(lst), self.vid_list, self.vid_list_DAVIS2016, path.join(data_root, 'ImageSets', imset));assert 1==0
+        return data
 
-    def get_datasets(self):
-        for video in self.vid_list:
-            if video not in self.ref_img_list:
-                continue
+    def resize_mask(self, mask):
+        # mask transform is applied AFTER mapper, so we need to post-process it in eval.py
+        h, w = mask.shape[-2:]
+        min_hw = min(h, w)
+        return Ff.interpolate(mask, (int(h/min_hw*self.size), int(w/min_hw*self.size)), 
+                    mode='nearest')
 
-            # print(self.image_dir, video, path.join(self.image_dir, video));assert 1==0
-            yield VideoReader_221128_TransColorization(video, 
-                path.join(self.image_dir, video), 
-                path.join(self.mask_dir, video),
-                size=self.size,
-                size_dir=path.join(self.size_dir, video),
-                args=self.args
-            )
+    def get_palette(self):
+        return self.palette
 
     def __len__(self):
-        return len(self.vid_list)
+        return len(self.frames)

test.py

@@ -1,8 +1,15 @@
+# test.py — In-process callable version (for ZeroGPU stateless)
+# Keep original logic; add build_parser(), run_cli(args_list), and run_inference(args)
+# Do NOT initialize CUDA at import-time.
+
 import os
 from os import path
 from argparse import ArgumentParser
 import shutil
 
+# 不在这里做 @spaces.GPU 装饰，避免与 app.py 的 @spaces.GPU 双重调度
+# import spaces
+
 import torch
 import torch.nn.functional as F
 from torch.utils.data import DataLoader
@@ -27,47 +34,7 @@ except ImportError:
     print('Failed to import hickle. Fine if not using multi-scale testing.')
 
 
-"""
-Arguments loading
-"""
-parser = ArgumentParser()
-parser.add_argument('--model', default='saves/DINOv2FeatureV6_LocalAtten_s2_154000.pth')
-
-# dataset setting
-parser.add_argument('--d16_batch_path', default='input')
-parser.add_argument('--deoldify_path', default='ref')
-parser.add_argument('--output', default='result')
-
-# For generic (G) evaluation, point to a folder that contains "JPEGImages" and "Annotations"
-parser.add_argument('--generic_path')
-parser.add_argument('--dataset', help='D16/D17/Y18/Y19/LV1/LV3/G', default='D16_batch')
-parser.add_argument('--split', help='val/test', default='val')
-parser.add_argument('--save_all', action='store_true', 
-            help='Save all frames. Useful only in YouTubeVOS/long-time video', )
-parser.add_argument('--benchmark', action='store_true', help='enable to disable amp for FPS benchmarking')
-        
-# Long-term memory options
-parser.add_argument('--disable_long_term', action='store_true')
-parser.add_argument('--max_mid_term_frames', help='T_max in paper, decrease to save memory', type=int, default=10)
-parser.add_argument('--min_mid_term_frames', help='T_min in paper, decrease to save memory', type=int, default=5)
-parser.add_argument('--max_long_term_elements', help='LT_max in paper, increase if objects disappear for a long time', 
-                                                type=int, default=10000)
-parser.add_argument('--num_prototypes', help='P in paper', type=int, default=128)
-
-parser.add_argument('--top_k', type=int, default=30)
-parser.add_argument('--mem_every', help='r in paper. Increase to improve running speed.', type=int, default=5)
-parser.add_argument('--deep_update_every', help='Leave -1 normally to synchronize with mem_every', type=int, default=-1)
-
-# Multi-scale options
-parser.add_argument('--save_scores', action='store_true')
-parser.add_argument('--flip', action='store_true')
-parser.add_argument('--size', default=-1, type=int, 
-            help='Resize the shorter side to this size. -1 to use original resolution. ')
-
-args = parser.parse_args()
-config = vars(args)
-config['enable_long_term'] = not config['disable_long_term']
-
+# ----------------- small utils -----------------
 def detach_to_cpu(x):
     return x.detach().cpu()
 
@@ -89,142 +56,236 @@ def lab2rgb_transform_PIL(mask):
 
     return im.clip(0, 1)
 
-if args.output is None:
-    args.output = f'.output/{args.dataset}_{args.split}'
-    print(f'Output path not provided. Defaulting to {args.output}')
-
-"""
-Data preparation
-"""
-is_youtube = args.dataset.startswith('Y')
-is_davis = args.dataset.startswith('D')
-is_lv = args.dataset.startswith('LV')
-
-if is_youtube or args.save_scores:
-    out_path = path.join(args.output, 'Annotations')
-else:
-    out_path = args.output
-
-if args.split == 'val':
-    # Set up Dataset, a small hack to use the image set in the 2017 folder because the 2016 one is of a different format
-    meta_dataset = DAVISTestDataset_221128_TransColorization_batch(args.d16_batch_path, imset=args.deoldify_path, size=args.size)
-else:
-    raise NotImplementedError
-palette = None
-
-torch.autograd.set_grad_enabled(False)
-
-# Set up loader
-meta_loader = meta_dataset.get_datasets()
-
-# Load our checkpoint
-network = ColorMNet(config, args.model).cuda().eval()
-if args.model is not None:
-    model_weights = torch.load(args.model)
-    network.load_weights(model_weights, init_as_zero_if_needed=True)
-else:
-    print('No model loaded.')
-
-total_process_time = 0
-total_frames = 0
-
-# Start eval
-for vid_reader in progressbar(meta_loader, max_value=len(meta_dataset), redirect_stdout=True):
-
-    loader = DataLoader(vid_reader, batch_size=1, shuffle=False, num_workers=2)
-    vid_name = vid_reader.vid_name
-    vid_length = len(loader)
-    # no need to count usage for LT if the video is not that long anyway
-    config['enable_long_term_count_usage'] = (
-        config['enable_long_term'] and
-        (vid_length
-            / (config['max_mid_term_frames']-config['min_mid_term_frames'])
-            * config['num_prototypes'])
-        >= config['max_long_term_elements']
+
+# ----------------- argparse -----------------
+def build_parser() -> ArgumentParser:
+    parser = ArgumentParser()
+    parser.add_argument('--model', default='saves/DINOv2FeatureV6_LocalAtten_s2_154000.pth')
+    parser.add_argument('--FirstFrameIsNotExemplar', help='Whether the provided reference frame is exactly the first input frame', action='store_true')
+
+    # dataset setting
+    parser.add_argument('--d16_batch_path', default='input', help='Point to folder A/ which contains <video_name>/00000.png etc.')
+    parser.add_argument('--ref_path', default='ref', help='Kept for parity; dataset will also read ref.png under each video folder when args provided')
+    parser.add_argument('--output', default='result', help='Directory to save results')
+
+    parser.add_argument('--reverse', default=False, action='store_true', help='whether to reverse the frame order')
+    parser.add_argument('--allow_resume', action='store_true', 
+                        help='skip existing videos that have been colorized')
+
+    # For generic (G) evaluation, point to a folder that contains "JPEGImages" and "Annotations"
+    parser.add_argument('--generic_path')
+    parser.add_argument('--dataset', help='D16/D17/Y18/Y19/LV1/LV3/G', default='D16_batch')
+    parser.add_argument('--split', help='val/test', default='val')
+    parser.add_argument('--save_all', action='store_true', 
+                        help='Save all frames. Useful only in YouTubeVOS/long-time video')
+    parser.add_argument('--benchmark', action='store_true', help='enable to disable amp for FPS benchmarking')
+            
+    # Long-term memory options
+    parser.add_argument('--disable_long_term', action='store_true')
+    parser.add_argument('--max_mid_term_frames', help='T_max in paper, decrease to save memory', type=int, default=10)
+    parser.add_argument('--min_mid_term_frames', help='T_min in paper, decrease to save memory', type=int, default=5)
+    parser.add_argument('--max_long_term_elements', help='LT_max in paper, increase if objects disappear for a long time', 
+                                                    type=int, default=10000)
+    parser.add_argument('--num_prototypes', help='P in paper', type=int, default=128)
+
+    parser.add_argument('--top_k', type=int, default=30)
+    parser.add_argument('--mem_every', help='r in paper. Increase to improve running speed.', type=int, default=5)
+    parser.add_argument('--deep_update_every', help='Leave -1 normally to synchronize with mem_every', type=int, default=-1)
+
+    # Multi-scale options
+    parser.add_argument('--save_scores', action='store_true')
+    parser.add_argument('--flip', action='store_true')
+    parser.add_argument('--size', default=-1, type=int, 
+                        help='Resize the shorter side to this size. -1 to use original resolution. ')
+    return parser
+
+
+# ----------------- core inference -----------------
+def run_inference(args):
+    """
+    真正的推理流程。必须在 ZeroGPU 的调度上下文里被调用（由 app.py 的 @spaces.GPU 包裹）。
+    不要在导入模块时做任何 CUDA 初始化。
+    """
+    config = vars(args)
+    config['enable_long_term'] = not config['disable_long_term']
+
+    if args.output is None:
+        args.output = f'.output/{args.dataset}_{args.split}'
+        print(f'Output path not provided. Defaulting to {args.output}')
+
+    # ----- Data preparation -----
+    is_youtube = args.dataset.startswith('Y')
+    is_davis = args.dataset.startswith('D')
+    is_lv = args.dataset.startswith('LV')
+
+    if is_youtube or args.save_scores:
+        out_path = path.join(args.output, 'Annotations')
+    else:
+        out_path = args.output
+
+    if args.split != 'val':
+        raise NotImplementedError('Only split=val is supported in this script.')
+
+    # 数据集：支持 A/<video>/00000.png ... 且读取 A/<video>/ref.png
+    meta_dataset = DAVISTestDataset_221128_TransColorization_batch(
+        args.d16_batch_path, imset=args.ref_path, size=args.size, args=args
     )
+    palette = None  # 兼容保留
+
+    torch.autograd.set_grad_enabled(False)
 
-    mapper = MaskMapper()
-    processor = InferenceCore(network, config=config)
-    first_mask_loaded = False
-
-    for ti, data in enumerate(loader):
-        with torch.cuda.amp.autocast(enabled=not args.benchmark):
-            rgb = data['rgb'].cuda()[0]
-            msk = data.get('mask')
-            info = data['info']
-            frame = info['frame'][0]
-            shape = info['shape']
-            need_resize = info['need_resize'][0]
-
-            """
-            For timing see https://discuss.pytorch.org/t/how-to-measure-time-in-pytorch/26964
-            Seems to be very similar in testing as my previous timing method 
-            with two cuda sync + time.time() in STCN though 
-            """
-            start = torch.cuda.Event(enable_timing=True)
-            end = torch.cuda.Event(enable_timing=True)
-            start.record()
-
-            if not first_mask_loaded:
+    # Set up loader list (video readers)
+    meta_loader = meta_dataset.get_datasets()
+
+    # Load checkpoint/model
+    device = torch.device('cuda') if torch.cuda.is_available() else torch.device('cpu')
+    network = ColorMNet(config, args.model).to(device).eval()
+    if args.model is not None:
+        # map_location 不指定，按默认走（ZeroGPU 下会在被调度的设备上加载）
+        model_weights = torch.load(args.model, map_location=device)
+        network.load_weights(model_weights, init_as_zero_if_needed=True)
+    else:
+        print('No model loaded.')
+
+    total_process_time = 0.0
+    total_frames = 0
+
+    # ----- Start eval over videos -----
+    for vid_reader in progressbar(meta_loader, max_value=len(meta_dataset), redirect_stdout=True):
+        # 注意：ZeroGPU/Spaces 环境不允许子进程多线程加载，保持 num_workers=0
+        loader = DataLoader(vid_reader, batch_size=1, shuffle=False, num_workers=0, pin_memory=False)
+        vid_name = vid_reader.vid_name
+        vid_length = len(loader)
+
+        # LT usage check per original logic
+        config['enable_long_term_count_usage'] = (
+            config['enable_long_term'] and
+            (vid_length
+                / (config['max_mid_term_frames'] - config['min_mid_term_frames'])
+                * config['num_prototypes'])
+            >= config['max_long_term_elements']
+        )
+
+        mapper = MaskMapper()
+        processor = InferenceCore(network, config=config)
+        first_mask_loaded = False
+
+        # skip existing videos 
+        if args.allow_resume:
+            this_out_path = path.join(out_path, vid_name)
+            if path.exists(this_out_path):
+                print(f'Skipping {this_out_path} because output already exists.')
+                continue
+
+        for ti, data in enumerate(loader):
+            with torch.cuda.amp.autocast(enabled=not args.benchmark):
+                rgb = data['rgb'].to(device)[0]
+
+                msk = data.get('mask')
+                if not config['FirstFrameIsNotExemplar']:
+                    msk = msk[:, 1:3, :, :] if msk is not None else None
+                    
+                info = data['info']
+                frame = info['frame'][0]
+                shape = info['shape']
+                need_resize = info['need_resize'][0]
+
+                start = torch.cuda.Event(enable_timing=True)
+                end = torch.cuda.Event(enable_timing=True)
+                start.record()
+
+                # 第一次必须有 mask
+                if not first_mask_loaded:
+                    if msk is not None:
+                        first_mask_loaded = True
+                    else:
+                        continue
+
+                if args.flip:
+                    rgb = torch.flip(rgb, dims=[-1])
+                    msk = torch.flip(msk, dims=[-1]) if msk is not None else None
+
+                # Map possibly non-continuous labels to continuous ones
                 if msk is not None:
-                    first_mask_loaded = True
+                    msk = torch.Tensor(msk[0]).to(device)
+                    if need_resize:
+                        msk = vid_reader.resize_mask(msk.unsqueeze(0))[0]
+                    processor.set_all_labels(list(range(1, 3)))
+                    labels = range(1, 3)
                 else:
-                    # no point to do anything without a mask
-                    continue
-
-            if args.flip:
-                rgb = torch.flip(rgb, dims=[-1])
-                msk = torch.flip(msk, dims=[-1]) if msk is not None else None
+                    labels = None
+        
+                # Run the model on this frame
+                if config['FirstFrameIsNotExemplar']:
+                    prob = processor.step_AnyExemplar(
+                        rgb,
+                        msk[:1, :, :].repeat(3, 1, 1) if msk is not None else None,
+                        msk[1:3, :, :] if msk is not None else None,
+                        labels,
+                        end=(ti == vid_length - 1)
+                    )
+                else:
+                    prob = processor.step(rgb, msk, labels, end=(ti == vid_length - 1))
 
-            # Map possibly non-continuous labels to continuous ones
-            if msk is not None:
-                msk = torch.Tensor(msk[0]).cuda()
+                # Upsample to original size if needed
                 if need_resize:
-                    msk = vid_reader.resize_mask(msk.unsqueeze(0))[0]
-                processor.set_all_labels(list(range(1,3)))
-                labels = range(1,3)
-            else:
-                labels = None
-    
-            # Run the model on this frame
-            prob = processor.step(rgb, msk, labels, end=(ti==vid_length-1))
-
-            # Upsample to original size if needed
-            if need_resize:
-                prob = F.interpolate(prob.unsqueeze(1), shape, mode='bilinear', align_corners=False)[:,0]
-
-            end.record()
-            torch.cuda.synchronize()
-            total_process_time += (start.elapsed_time(end)/1000)
-            total_frames += 1
-
-            if args.flip:
-                prob = torch.flip(prob, dims=[-1])
-
-            if args.save_scores:
-                prob = (prob.detach().cpu().numpy()*255).astype(np.uint8)
-
-            # Save the mask
-            if args.save_all or info['save'][0]:
-                this_out_path = path.join(out_path, vid_name)
-                os.makedirs(this_out_path, exist_ok=True)
-
-                out_mask_final = lab2rgb_transform_PIL(torch.cat([rgb[:1,:,:], prob], dim=0))
-                out_mask_final = out_mask_final * 255
-                out_mask_final = out_mask_final.astype(np.uint8)
-
-                out_img = Image.fromarray(out_mask_final)
-                out_img.save(os.path.join(this_out_path, frame[:-4]+'.png'))
-
-print(f'Total processing time: {total_process_time}')
-print(f'Total processed frames: {total_frames}')
-print(f'FPS: {total_frames / total_process_time}')
-print(f'Max allocated memory (MB): {torch.cuda.max_memory_allocated() / (2**20)}')
-
-if not args.save_scores:
-    if is_youtube:
-        print('Making zip for YouTubeVOS...')
-        shutil.make_archive(path.join(args.output, path.basename(args.output)), 'zip', args.output, 'Annotations')
-    elif is_davis and args.split == 'test':
-        print('Making zip for DAVIS test-dev...')
-        shutil.make_archive(args.output, 'zip', args.output)
+                    prob = F.interpolate(prob.unsqueeze(1), shape, mode='bilinear', align_corners=False)[:, 0]
+
+                end.record()
+                torch.cuda.synchronize()
+                total_process_time += (start.elapsed_time(end)/1000)
+                total_frames += 1
+
+                if args.flip:
+                    prob = torch.flip(prob, dims=[-1])
+
+                if args.save_scores:
+                    prob = (prob.detach().cpu().numpy() * 255).astype(np.uint8)
+
+                # Save the mask
+                if args.save_all or info['save'][0]:
+                    this_out_path = path.join(out_path, vid_name)
+                    os.makedirs(this_out_path, exist_ok=True)
+
+                    out_mask_final = lab2rgb_transform_PIL(torch.cat([rgb[:1, :, :], prob], dim=0))
+                    out_mask_final = (out_mask_final * 255).astype(np.uint8)
+
+                    out_img = Image.fromarray(out_mask_final)
+                    out_img.save(os.path.join(this_out_path, frame[:-4] + '.png'))
+
+    print(f'Total processing time: {total_process_time}')
+    print(f'Total processed frames: {total_frames}')
+    print(f'FPS: {total_frames / total_process_time}')
+    print(f'Max allocated memory (MB): {torch.cuda.max_memory_allocated() / (2**20)}')
+
+    # 与原版一致：只在 save_scores=False 且特定数据集/子集时打 zip
+    if not args.save_scores:
+        if is_youtube:
+            print('Making zip for YouTubeVOS...')
+            shutil.make_archive(path.join(args.output, path.basename(args.output)), 'zip', args.output, 'Annotations')
+        elif is_davis and args.split == 'test':
+            print('Making zip for DAVIS test-dev...')
+            shutil.make_archive(args.output, 'zip', args.output)
+
+
+# ----------------- public entrypoints -----------------
+def run_cli(args_list=None):
+    """
+    供 app.py 同进程调用：test.run_cli(args_list)
+    """
+    parser = build_parser()
+    args = parser.parse_args(args_list)
+    return run_inference(args)
+
+
+def main():
+    """
+    保留命令行可运行：python test.py --d16_batch_path A --output result ...
+    注意：若在 Hugging Face Spaces/ZeroGPU 无状态环境下直接 run main()，
+    需要由上层（如 app.py 的 @spaces.GPU）提供调度上下文。
+    """
+    run_cli()
+
+
+if __name__ == '__main__':
+    main()