update

.gitignore

@@ -9,6 +9,8 @@ wandb/
 pretrain/
 Pytorch-Correlation-extension/
 result
+src/
+DINOv2FeatureV6_LocalAtten_s2_154000.pth
 
 # Byte-compiled / optimized / DLL files
 __pycache__/

app.py

@@ -1,147 +1,535 @@
+# 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
+
+import os
+import sys
+import shutil
+import subprocess
+import uuid
+import urllib.request
+import warnings
+from os import path
+
+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 gradio as gr
+import spaces  # ZeroGPU decorator
 import numpy as np
-import random
-from diffusers import DiffusionPipeline
+from PIL import Image
+import cv2
+import traceback
+
 import torch
-import spaces
-
-device = "cuda" if torch.cuda.is_available() else "cpu"
-
-if torch.cuda.is_available():
-    torch.cuda.max_memory_allocated(device=device)
-    pipe = DiffusionPipeline.from_pretrained("stabilityai/sdxl-turbo", torch_dtype=torch.float16, variant="fp16", use_safetensors=True)
-    pipe.enable_xformers_memory_efficient_attention()
-    pipe = pipe.to(device)
-else: 
-    pipe = DiffusionPipeline.from_pretrained("stabilityai/sdxl-turbo", use_safetensors=True)
-    pipe = pipe.to(device)
-
-MAX_SEED = np.iinfo(np.int32).max
-MAX_IMAGE_SIZE = 1024
-
-def infer(prompt, negative_prompt, seed, randomize_seed, width, height, guidance_scale, num_inference_steps):
-
-    if randomize_seed:
-        seed = random.randint(0, MAX_SEED)
-        
-    generator = torch.Generator().manual_seed(seed)
-    
-    image = pipe(
-        prompt = prompt, 
-        negative_prompt = negative_prompt,
-        guidance_scale = guidance_scale, 
-        num_inference_steps = num_inference_steps, 
-        width = width, 
-        height = height,
-        generator = generator
-    ).images[0] 
-    
-    return image
-
-examples = [
-    "Astronaut in a jungle, cold color palette, muted colors, detailed, 8k",
-    "An astronaut riding a green horse",
-    "A delicious ceviche cheesecake slice",
-]
-
-css="""
-#col-container {
-    margin: 0 auto;
-    max-width: 520px;
-}
+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 -----------------
+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`
 """
 
-if torch.cuda.is_available():
-    power_device = "GPU"
-else:
-    power_device = "CPU"
-
-with gr.Blocks(css=css) as demo:
-    
-    with gr.Column(elem_id="col-container"):
-        gr.Markdown(f"""
-        # Text-to-Image Gradio Template
-        Currently running on {power_device}.
-        """)
-        
-        with gr.Row():
-            
-            prompt = gr.Text(
-                label="Prompt",
-                show_label=False,
-                max_lines=1,
-                placeholder="Enter your prompt",
-                container=False,
-            )
-            
-            run_button = gr.Button("Run", scale=0)
-        
-        result = gr.Image(label="Result", show_label=False)
-
-        with gr.Accordion("Advanced Settings", open=False):
-            
-            negative_prompt = gr.Text(
-                label="Negative prompt",
-                max_lines=1,
-                placeholder="Enter a negative prompt",
-                visible=False,
-            )
-            
-            seed = gr.Slider(
-                label="Seed",
-                minimum=0,
-                maximum=MAX_SEED,
-                step=1,
-                value=0,
-            )
-            
-            randomize_seed = gr.Checkbox(label="Randomize seed", value=True)
-            
-            with gr.Row():
-                
-                width = gr.Slider(
-                    label="Width",
-                    minimum=256,
-                    maximum=MAX_IMAGE_SIZE,
-                    step=32,
-                    value=512,
-                )
-                
-                height = gr.Slider(
-                    label="Height",
-                    minimum=256,
-                    maximum=MAX_IMAGE_SIZE,
-                    step=32,
-                    value=512,
-                )
-            
-            with gr.Row():
-                
-                guidance_scale = gr.Slider(
-                    label="Guidance scale",
-                    minimum=0.0,
-                    maximum=10.0,
-                    step=0.1,
-                    value=0.0,
-                )
-                
-                num_inference_steps = gr.Slider(
-                    label="Number of inference steps",
-                    minimum=1,
-                    maximum=12,
-                    step=1,
-                    value=2,
-                )
-        
-        gr.Examples(
-            examples = examples,
-            inputs = [prompt]
+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)
+    os.makedirs(d, exist_ok=True)
+
+# ----------------- MISC UTILS -----------------
+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)
+
+# ---------- extract frames: dataset-root/<video_stem>/00000.png ----------
+def video_to_dataset_root(video_path: str, dataset_root: str):
+    """
+    将单个视频抽帧到 dataset_root/<video_stem>/00000.png...
+    返回: (subdir_path, video_stem, width, height, fps, frame_count)
+    """
+    ensure_clean_dir(dataset_root)
+    basename = path.basename(video_path)
+    stem, _ = path.splitext(basename)
+    subdir = path.join(dataset_root, stem)
+    ensure_clean_dir(subdir)
+
+    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
+
+    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)
+
+        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, stem, 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
+
+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."
+
+    first = cv2.imread(path.join(frames_dir, 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))
+        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:
+    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 CWD
+    base_run_dir = path.join(os.getcwd(), 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": False,
+        "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,
+    }
+    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_list = meta_dataset.get_datasets()
+
+    target_reader = None
+    for vr in meta_list:
+        if getattr(vr, "vid_name", None) == vid_stem:
+            target_reader = vr
+            break
+    if target_reader is None:
+        if len(meta_list) == 1:
+            target_reader = meta_list[0]
+        else:
+            raise RuntimeError(f"未在数据集中找到目标视频子目录：{vid_stem}；可用={ [getattr(v, 'vid_name', '?') for v in meta_list] }")
+
+    # 输出路径规则（与 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
+
+    # 6) 推理（逐帧；内部逻辑与 main.py 对齐；保留调试打印）
+    vid_reader = target_reader
+    # 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])
+        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)
+    shutil.rmtree(base_run_dir, ignore_errors=True)
+
+    return final_mp4
+
+# ----------------- GRADIO HANDLERS -----------------
+@spaces.GPU(duration=1200)
+def gradio_infer(
+    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
+):
+    if not torch.cuda.is_available():
+        return None, "ZeroGPU 未分配到 GPU，请重试（或检查 Space 硬件是否为 ZeroGPU）。"
+
+    if bw_video is None:
+        return None, "请上传黑白视频。"
+    if ref_image is None:
+        return None, "请上传参考图像。"
+
+    # Video path
+    if isinstance(bw_video, dict) and "name" in bw_video:
+        bw_video_path = bw_video["name"]
+    elif isinstance(bw_video, str):
+        bw_video_path = bw_video
+    else:
+        return None, "无法读取视频输入。"
+
+    # Ref path
+    if isinstance(ref_image, Image.Image):
+        tmp_ref_path = path.join(os.getcwd(), f"ref_{uuid.uuid4().hex}.png")
+        ref_image.save(tmp_ref_path)
+        ref_path = tmp_ref_path
+    elif isinstance(ref_image, str):
+        ref_path = ref_image
+    else:
+        return None, "无法读取参考图像输入。"
+
+    default_config = {
+        "FirstFrameIsNotExemplar": True,
+        "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,
+    }
+
+    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"],
+        "split": str(split) if split else default_config["split"],
+        "save_all": bool(save_all) if save_all is not None else default_config["save_all"],
+        "benchmark": bool(benchmark) if benchmark is not None else default_config["benchmark"],
+        "disable_long_term": bool(disable_long_term) if disable_long_term is not None else default_config["disable_long_term"],
+        "max_mid_term_frames": int(max_mid) if max_mid is not None else default_config["max_mid_term_frames"],
+        "min_mid_term_frames": int(min_mid) if min_mid is not None else default_config["min_mid_term_frames"],
+        "max_long_term_elements": int(max_long) if max_long is not None else default_config["max_long_term_elements"],
+        "num_prototypes": int(num_proto) if num_proto is not None else default_config["num_prototypes"],
+        "top_k": int(top_k) if top_k is not None else default_config["top_k"],
+        "mem_every": int(mem_every) if mem_every is not None else default_config["mem_every"],
+        "deep_update_every": int(deep_update) if deep_update is not None else default_config["deep_update_every"],
+        "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"],
+    }
+
+    try:
+        out_mp4 = run_pipeline_cuda(
+            bw_video_path, ref_path, user_config, debug_shapes=bool(debug_shapes)
         )
+        return out_mp4, "完成 ✅"
+    except subprocess.CalledProcessError as e:
+        return None, f"运行时错误：\n{e}"
+    except Exception as e:
+        return None, f"{e}"
+
+# ----------------- UI -----------------
+with gr.Blocks() as demo:
+    gr.Markdown(f"# {TITLE}")
+    gr.Markdown(DESC)
+
+    debug_shapes = gr.Checkbox(label="调试日志（打印形状与完整Traceback）", value=False)
+
+    with gr.Row():
+        inp_video = gr.Video(label="黑白视频（mp4/webm/avi）", interactive=True)
+        inp_ref = gr.Image(label="参考图像（RGB）", type="pil")
+
+    with gr.Accordion("高级参数设置（与 main.py 对齐）", open=False):
+        with gr.Row():
+            first_not_exemplar = gr.Checkbox(label="FirstFrameIsNotExemplar", 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)
+        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)
+        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_button.click(
-        fn = infer,
-        inputs = [prompt, negative_prompt, seed, randomize_seed, width, height, guidance_scale, num_inference_steps],
-        outputs = [result]
+    run_btn = gr.Button("开始着色（ZeroGPU 推理）")
+    with gr.Row():
+        out_video = gr.Video(label="输出视频（着色结果）")
+        status = gr.Textbox(label="状态 / 调试输出", interactive=False, lines=12)
+
+    run_btn.click(
+        fn=gradio_infer,
+        inputs=[
+            debug_shapes,
+            inp_video, inp_ref,
+            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
+        ],
+        outputs=[out_video, status]
     )
 
-demo.queue().launch()
+if __name__ == "__main__":
+    try:
+        ensure_checkpoint()
+    except Exception as e:
+        print(f"[WARN] 预下载权重失败（首次推理会再试）: {e}")
+
+    demo.queue(max_size=32).launch(server_name="0.0.0.0", server_port=7860)

inference/data/video_reader.py

@@ -82,6 +82,10 @@ class VideoReader_221128_TransColorization(Dataset):
         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)
             mask_ab = mask[1:3,:,:]
             data['mask'] = mask_ab

requirements.txt

@@ -82,6 +82,7 @@ tb-nightly
 tensorboard
 tensorboard-data-server
 -e git+https://github.com/cheind/py-thin-plate-spline.git@f6995795397118b7d0ac01aecd3f39ffbfad9dee#egg=thinplate
+# -e git+https://github.com/ClementPinard/Pytorch-Correlation-extension.git#egg=spatial_correlation_sampler
 tifffile
 tomli
 tqdm
@@ -91,4 +92,11 @@ urllib3
 wandb
 Werkzeug
 yapf
-zipp
+zipp
+gradio
+torch
+opencv-python
+numpy
+pillow
+scikit-image
+spaces   # <<< 关键：提供 @spaces.GPU 装饰器