CorridorKeyModule/ck2.py

"""CorridorKey v2 adapter: assets, CPU-resident model, CPU pre/post, one GPU call.



The vendored `corridorkey` package is the unmodified v2 runtime. This module only

decides what runs where, so the time spent inside the ZeroGPU window is model work:



    CPU  probe + downscale the clip with ffmpeg

    GPU  run_inference (RVM prepass, Hann stage 1, stage 2, ffmpeg writers)

    CPU  browser-safe previews + zip



v2 needs CUDA. There is no CPU path; v1 in app.py keeps that role.

"""

from __future__ import annotations

import base64
import hashlib
import logging
import math
import os
import shutil
import subprocess
import sys
import tempfile
import time
import urllib.request
import zipfile
from pathlib import Path

import cv2

try:
    import spaces
    HAS_SPACES = True
except ImportError:
    HAS_SPACES = False

logger = logging.getLogger("ck2")

ROOT = Path(__file__).resolve().parent
# Weights and foundation models live here. CK2_ASSET_ROOT lets a local run reuse an existing download.
ASSETS = Path(os.environ.get("CK2_ASSET_ROOT") or ROOT)

CK2_REPO = "nikopueringer/CorridorKey2.0"
CK2_FILES = ("corridorkey_v2.safetensors", "foundation_stats.safetensors")
CRADIO_REPO = "nvidia/C-RADIOv4-SO400M"
CRADIO_REVISION = "c0457f5dc26ca145f954cd4fc5bb6114e5705ad8"
MOGE_REPO = "Ruicheng/moge-2-vitb-normal"
MOGE_REVISION = "ca5f0e07ff01d3e5a364c1d954ed12ee1814b368"
RVM_FILE = "pretrain_rvm_small16_256_204031069.npz"
RVM_URL = f"https://storage.googleapis.com/representations4d/checkpoints/{RVM_FILE}"
RVM_MD5_B64 = "ZJb6kj8+tasiJbVeplXu7Q=="

QUALITY_CHOICES = ["540p", "720p", "1080p", "Native"]
DEFAULT_QUALITY = "720p"
TEMPORAL_CHOICES = [("Auto", "auto"), ("Off (1 frame)", "1"), ("8 frames", "8"), ("24 frames", "24")]
DEFAULT_TEMPORAL = "auto"
STAGE2_BATCH = 4

GPU_OVERHEAD_S = 25.0
GPU_SAFETY = 1.35

_model = None
_load_error = ""
_ready = False


def seconds_per_frame(width: int, height: int) -> float:
    """Measured on an RTX 5090 (89 frames, 540p/720p/1080p): about 0.05 + 0.22 s per megapixel."""
    return 0.05 + 0.22 * (width * height / 1e6)


def estimate_gpu_seconds(frames: int, width: int, height: int) -> int:
    """GPU time to request. Uncapped on purpose: ZeroGPU refuses what the account's tier cannot run."""
    estimate = GPU_OVERHEAD_S + GPU_SAFETY * frames * seconds_per_frame(width, height)
    return int(max(30, math.ceil(estimate)))


def working_size(src_w: int, src_h: int, quality: str) -> tuple[int, int]:
    """Largest size not above the quality's short side, even dimensions, source aspect kept."""
    if quality == "Native":
        return src_w, src_h
    target = int(quality.rstrip("p"))
    short = min(src_w, src_h)
    if short <= target:
        return src_w - src_w % 2, src_h - src_h % 2
    scale = target / short
    return max(2, round(src_w * scale / 2) * 2), max(2, round(src_h * scale / 2) * 2)


def resolve_window(choice: str, width: int, height: int) -> tuple[int, int]:
    """Hann window (chunk, stride). Auto takes the widest window that stays close to the cost of 8

    frames: 80 up to 720p, 24 up to 1080p, 8 above. Measured on an RTX 5090: 80 frames at 1080p is

    about 2.5x slower than 8, and 24 frames at native 4K is over 10x slower."""
    if str(choice) == "auto":
        pixels = width * height
        chunk = 80 if pixels <= 1_000_000 else 24 if pixels <= 2_500_000 else 8
    elif str(choice) in ("1", "8", "24"):
        chunk = int(choice)
    else:
        raise RuntimeError(f"Unsupported temporal context {choice!r}; choose auto, 1, 8 or 24.")
    return chunk, max(1, chunk // 2)


def _md5_b64(path: Path) -> str:
    digest = hashlib.md5()
    with open(path, "rb") as f:
        for chunk in iter(lambda: f.read(1 << 20), b""):
            digest.update(chunk)
    return base64.b64encode(digest.digest()).decode()


def _retry(label: str, fn, attempts: int = 3):
    for attempt in range(1, attempts + 1):
        try:
            return fn()
        except Exception as exc:
            logger.warning("[ck2] %s failed (%d/%d): %s", label, attempt, attempts, exc)
            if attempt == attempts:
                raise
            time.sleep(3 * attempt)


def ensure_assets() -> None:
    """Fetch the v2 weights and the three frozen foundation models, skipping anything present."""
    from huggingface_hub import hf_hub_download, snapshot_download

    for name in CK2_FILES:
        _retry(name, lambda n=name: hf_hub_download(CK2_REPO, n, local_dir=ASSETS / "weights"))
    cradio = ASSETS / "vfms/cradio/C-RADIOv4-SO400M"
    _retry("C-RADIO", lambda: snapshot_download(
        CRADIO_REPO, revision=CRADIO_REVISION, local_dir=cradio, ignore_patterns=["*.pth.tar"]))
    _retry("MoGe", lambda: hf_hub_download(
        MOGE_REPO, "model.pt", revision=MOGE_REVISION, local_dir=ASSETS / "vfms/moge"))

    dst = ASSETS / "vfms/rvm" / RVM_FILE
    if dst.is_file() and _md5_b64(dst) == RVM_MD5_B64:
        return
    dst.parent.mkdir(parents=True, exist_ok=True)
    tmp = dst.with_suffix(dst.suffix + ".part")

    def fetch():
        with urllib.request.urlopen(RVM_URL, timeout=60) as response, open(tmp, "wb") as f:
            shutil.copyfileobj(response, f, length=1 << 20)
        if _md5_b64(tmp) != RVM_MD5_B64:
            tmp.unlink()
            raise RuntimeError("checksum mismatch downloading RVM weights")
        tmp.replace(dst)

    _retry("RVM", fetch)


def init() -> bool:
    """Download assets and build the model on CPU. Never raises: v1 must keep working."""
    global _model, _ready, _load_error
    try:
        t0 = time.time()
        ensure_assets()
        t1 = time.time()
        saved = {k: os.environ.get(k) for k in ("HF_HUB_OFFLINE", "TRANSFORMERS_OFFLINE")}
        try:
            from CorridorKeyModule.corridorkey_v2 import runtime as rt
            _model = rt.load_model(rt.RuntimePaths.from_package(ASSETS), device="cpu")
        finally:
            for key, value in saved.items():
                if value is None:
                    os.environ.pop(key, None)
                else:
                    os.environ[key] = value
        _ready = True
        logger.info("[ck2] ready: assets %.0fs, model load %.0fs", t1 - t0, time.time() - t1)
    except Exception as exc:
        _load_error = f"{type(exc).__name__}: {exc}"
        logger.exception("[ck2] v2 unavailable")
    return _ready


def is_ready() -> bool:
    return _ready


def load_error() -> str:
    return _load_error


def _run(cmd: list[str], what: str, timeout: int = 600) -> None:
    try:
        subprocess.run(cmd, capture_output=True, timeout=timeout, check=True)
    except subprocess.CalledProcessError as exc:
        raise RuntimeError(f"{what} failed: {exc.stderr.decode(errors='ignore')[-300:]}") from exc
    except (FileNotFoundError, subprocess.TimeoutExpired) as exc:
        raise RuntimeError(f"{what} failed: {exc}") from exc


def _probe(path: str) -> tuple[int, int, float, int]:
    cap = cv2.VideoCapture(path)
    try:
        w = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
        h = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
        fps = cap.get(cv2.CAP_PROP_FPS) or 30.0
        frames = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
    finally:
        cap.release()
    if frames <= 0 or w <= 0 or h <= 0:
        raise RuntimeError("Could not read video frames.")
    return w, h, fps, frames


def prepare_input(video_path: str, quality: str, workdir: str) -> dict:
    """CPU stage: downscale the clip once (Native keeps the source as is)."""
    src_w, src_h, fps, frames = _probe(video_path)
    w, h = working_size(src_w, src_h, quality)
    if (w, h) == (src_w, src_h):
        return {"work": video_path, "w": w, "h": h, "src_w": src_w, "src_h": src_h, "fps": fps, "frames": frames}
    work = os.path.join(workdir, "clip.mp4")
    cmd = ["ffmpeg", "-y", "-i", video_path, "-an", "-vf", f"scale={w}:{h}:flags=lanczos",
           "-c:v", "libx264", "-preset", "veryfast", "-crf", "8", "-pix_fmt", "yuv420p", work]
    _run(cmd, "Downscale")
    return {"work": work, "w": w, "h": h, "src_w": src_w, "src_h": src_h, "fps": fps, "frames": frames}


def _gpu_duration(work_video, out_dir, frames, width, height, *args, **kwargs):
    return estimate_gpu_seconds(int(frames), int(width), int(height))


def _gpu_decorator(fn):
    if HAS_SPACES:
        return spaces.GPU(duration=_gpu_duration)(fn)
    return fn


@_gpu_decorator
def _gpu_phase(work_video, out_dir, frames, width, height, chunk, stride, despill_strength):
    """All GPU work. The model is moved in and moved back out so every call starts clean."""
    import torch
    from CorridorKeyModule.corridorkey_v2 import runtime as rt

    if not torch.cuda.is_available():
        raise RuntimeError("CorridorKey v2 requires a CUDA GPU.")
    t0 = time.time()
    props = torch.cuda.get_device_properties(0)
    gpu = f"{props.name} {props.total_memory / 2**30:.0f}GB"
    tmp = Path(tempfile.mkdtemp(prefix="ck2_tmp_"))
    try:
        _model.to("cuda")
        cfg = rt.InferenceConfig(
            hann_chunk=int(chunk), hann_stride=int(stride), auto_vram=False,
            stage2_batch=STAGE2_BATCH, outputs={"alpha", "fg", "checker", "cutout"},
            despill_strength=float(despill_strength), temp_dir=tmp)
        writer = rt.WriterConfig(preset="veryfast", threads=4)
        rt.run_inference(_model, work_video, out_dir, cfg, writer,
                         initial_hint_path=None, hint_video_path=None)
    finally:
        _model.to("cpu")
        torch.cuda.empty_cache()
        shutil.rmtree(tmp, ignore_errors=True)
    return time.time() - t0, gpu


def _find(out_dir: str, suffix: str) -> str | None:
    for name in sorted(os.listdir(out_dir)):
        if name.endswith(".mp4") and suffix in name:
            return os.path.join(out_dir, name)
    return None


def _preview(src: str, dst: str, size: tuple[int, int] | None) -> str:
    vf = "format=yuv420p" if size is None else f"scale={size[0]}:{size[1]}:flags=lanczos,format=yuv420p"
    _run(["ffmpeg", "-y", "-i", src, "-vf", vf, "-c:v", "libx264", "-preset", "fast", "-crf", "18",
          "-movflags", "+faststart", dst], "Preview encode")
    return dst


def process(video_path, quality=DEFAULT_QUALITY, temporal=DEFAULT_TEMPORAL, despill_val=5,

            upscale=False, progress=None):
    """Run CorridorKey v2. Returns (composite, matte, zip, status) like the v1 pipeline."""
    if not _ready:
        raise RuntimeError(f"CorridorKey v2 is not available ({_load_error or 'still loading'}). "
                           "Switch to v1 or retry in a minute.")
    if video_path is None:
        raise RuntimeError("Please upload a video.")
    if quality not in QUALITY_CHOICES:
        raise RuntimeError(f"Unsupported quality {quality!r}; choose one of {QUALITY_CHOICES}.")
    def step(frac, text):
        if progress is not None:
            progress(frac, desc=text)

    t_total = time.time()
    workdir = tempfile.mkdtemp(prefix="ck2_job_")
    out_dir = os.path.join(workdir, "out")
    os.makedirs(out_dir)

    step(0.05, "Preparing clip (CPU)...")
    t0 = time.time()
    info = prepare_input(video_path, quality, workdir)
    t_prep = time.time() - t0

    chunk, stride = resolve_window(temporal, info["w"], info["h"])
    step(0.2, f"CorridorKey v2 on GPU ({info['w']}x{info['h']}, {info['frames']} frames)...")
    t_gpu, gpu = _gpu_phase(info["work"], out_dir, info["frames"], info["w"], info["h"],
                       chunk, stride, float(despill_val) / 10.0)

    step(0.85, "Encoding previews (CPU)...")
    t0 = time.time()
    alpha = _find(out_dir, "_alpha")
    fg = _find(out_dir, "_fg")
    comp = _find(out_dir, "_checker")
    if not (alpha and fg and comp):
        raise RuntimeError("CorridorKey v2 finished without all expected outputs.")
    size = None
    if upscale and (info["src_w"], info["src_h"]) != (info["w"], info["h"]):
        size = (info["src_w"] - info["src_w"] % 2, info["src_h"] - info["src_h"] % 2)
    comp_preview = _preview(comp, os.path.join(workdir, "comp_preview.mp4"), size)
    matte_preview = _preview(alpha, os.path.join(workdir, "matte_preview.mp4"), size)
    zip_path = os.path.join(workdir, "CorridorKey2_Output.zip")
    with zipfile.ZipFile(zip_path, "w", zipfile.ZIP_STORED) as zf:
        for path in (alpha, fg, comp):
            zf.write(path, f"Output/{os.path.basename(path)}")
        settings = [p for p in os.listdir(out_dir) if p.endswith("_settings.txt")]
        for name in settings:
            zf.write(os.path.join(out_dir, name), f"Output/{name}")
    t_post = time.time() - t0

    step(1.0, "Done!")
    out_res = f"{size[0]}x{size[1]}" if size else f"{info['w']}x{info['h']}"
    status = (f"CorridorKey v2 | {info['frames']} frames at {quality} ({out_res}) | "
              f"temporal {chunk} | prep {t_prep:.0f}s, GPU {t_gpu:.0f}s ({gpu}), post {t_post:.0f}s, "
              f"{time.time() - t_total:.0f}s total")
    return comp_preview, matte_preview, zip_path, status