Spaces:
Running on Zero
Running on Zero
| """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 | |
| 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 | |