# SeedVR2-3B video restoration / upscaling + image upscaling.
#
# 2026-10-26: an IMAGE endpoint (upscale_image, fn_index 1) was merged in
# from the Upsampler v3 image recipe (~/workspace/tmp/seedvr2-src/upsampler).
# Both endpoints share this tree, the config, the weights and the memoized
# runner (the image path decodes t=1, where the runner's disabled causal
# slicing is inactive anyway). The owner asked to raise the image working
# area to 2560*1440 (3.7MP) to test detail retention; it was deployed and
# probed twice that day and the 3.7MP decode died ~10 s in both times
# (ZeroGPU RuntimeError, both allocator configs) -- same ceiling the v3
# author measured. IMAGE_WORK_AREA therefore defaults to the working
# 1920*1080 and stays env-overridable for future re-probes.
# expandable_segments stays OFF process-wide (see the note below); the
# image function applies the v3 recipe's runtime-only set instead.
#
# Built on the Upsampler seedvr2-3b packaging of ByteDance's SeedVR2 code:
# same trimmed tree, same pinned requirements, same patched config (pure
# PyTorch rms/layer norms instead of the apex fused paths that cannot
# install on ZeroGPU's Python 3.12, SDPA fallback instead of a hard
# flash-attn dependency, torch 2.11 instead of upstream's 2.4.0 which
# ZeroGPU does not support). What this app adds back is the VIDEO path
# (the Upsampler port is image-only), restructured for ZeroGPU:
#
# - ONE @spaces.GPU entry with a dynamic duration (upstream booked three
#   separate 100 s GPU allocations per request).
# - The runner is built lazily inside the first GPU call and memoized.
# - Frames are processed in temporal chunks of 41 ((t - 1) % 4 == 0 for
#   the VAE's 4x temporal stride), last frame repeated to pad the tail
#   chunk, padded frames trimmed from the output. Chunking bounds the
#   VAE decode peak, which is what trips the ZeroGPU NVML assert at
#   large working areas.
# - Modes: "restore" re-renders at the input's own area (deblur /
#   denoise / detail rebuild, no resolution change -- this is what the
#   upstream demo does by default); "1.5x" / "2x" scale the working
#   area by 2.25x / 4x, capped at 2560x1440 area.
# - cond_noise_scale 0.0 (upstream's own inference script value; the
#   demo's 0.1 was measured to scatter bright blemishes).
# - Optional "sliding" window mode (restore only): windows of 5 frames
#   advance by 4; every window after the first opens with the previous
#   window's finalized restored last frame as a causal anchor at
#   position 0 (the causal VAE encodes frame 0 independently, so it
#   anchors what follows), and that frame's re-render is averaged
#   0.5/0.5 with its finalized version. Same 5-frame cap (latent t=2);
#   ~25% more windows than chunked.
# - Original audio is stream-copied onto the result.

import os

# NOTE: no expandable_segments here. The image port uses it, but on this
# runtime the video VAE decoder (list-based causal inflation ops) trips
# the CUDACachingAllocator NVML assert with it enabled -- measured
# 2026-10-18, two runs, CHUNK 41 and CHUNK 9 alike, always inside the
# decoder's up blocks. Trying the native allocator for the video path.

import gc
import subprocess
import tempfile
from pathlib import Path

import cv2
import gradio as gr
import imageio
import imageio_ffmpeg
import numpy as np
import spaces
import torch
from einops import rearrange
from huggingface_hub import hf_hub_download
from omegaconf import OmegaConf
from PIL import Image
from torchvision.transforms import Compose, Lambda, Normalize
import torchvision.transforms as T

from data.image.transforms.divisible_crop import DivisibleCrop
from data.image.transforms.na_resize import NaResize
from data.video.transforms.rearrange import Rearrange
from common.config import load_config
from common.distributed import init_torch
from common.seed import set_seed
from projects.video_diffusion_sr.infer import VideoDiffusionInfer

try:
    from projects.video_diffusion_sr.color_fix import wavelet_reconstruction

    USE_COLOR_FIX = True
except ImportError:
    USE_COLOR_FIX = False
    print("color fix unavailable")

WEIGHTS_REPO = "ByteDance-Seed/SeedVR2-3B"
CKPT_DIR = Path("./ckpts")
CKPT_DIR.mkdir(exist_ok=True)


def _fetch(filename: str, target: Path) -> str:
    if target.exists():
        return str(target)
    path = hf_hub_download(repo_id=WEIGHTS_REPO, filename=filename)
    target.symlink_to(path)
    return str(target)


DIT_CKPT = _fetch("seedvr2_ema_3b.pth", CKPT_DIR / "seedvr2_ema_3b.pth")
VAE_CKPT = _fetch("ema_vae.pth", CKPT_DIR / "ema_vae.pth")
POS_EMB = _fetch("pos_emb.pt", Path("./pos_emb.pt"))
NEG_EMB = _fetch("neg_emb.pt", Path("./neg_emb.pt"))

INPUT_NOISE = 0.0
# Image endpoint working area (see header note). 2560*1440 = upstream's own
# image value; the public Upsampler Space caps at 1920*1080 because at
# 3.7MP their VAE decode tripped the ZeroGPU NVML assert. Re-probed on
# this Space 2026-10-26 at the owner's request: two runs at 2560*1440
# (with and without runtime expandable_segments) both died ~10 s in with
# the same opaque ZeroGPU RuntimeError as video '2x' -- the 3.7MP decode
# does not survive on this runtime in either allocator config, matching
# the v3 author's measurement. Default therefore stays at the measured
# ceiling 1920*1080; raise SEEDVR2_WORK_AREA to re-probe.
IMAGE_WORK_AREA = int(os.environ.get("SEEDVR2_WORK_AREA", 1920 * 1080))
IMAGE_DEFAULT_SEED = int(os.environ.get("SEEDVR2_SEED", 0))
OUTPUT_WEBP_QUALITY = 95
WEBP_MAX_SIDE = 16383  # libwebp hard limit; larger results fall back to PNG
# Measured threshold on this ZeroGPU runtime (2026-10-18, card119):
# a chunk of 5 frames (latent t=2) decodes fine; 9 frames (latent t=3)
# and up die in the VAE decoder with the CUDACachingAllocator NVML
# assert, regardless of chunk slicing, causal slicing on/off, or
# allocator settings. So chunks are 5 frames: each is decoded as an
# independent causal tensor (no cross-chunk cache), which is also why
# causal slicing is disabled in _ensure_runner. Per-chunk cost measured
# ~4.2 s at 720x1280 restore on a cold worker.
CHUNK = 5
SLIDE_STEP = CHUNK - 1  # sliding windows overlap by one frame
MAX_FRAMES = 481
MAX_TARGET_AREA = 2560 * 1440
MODE_AREA_SCALE = {"restore": 1.0, "1.5x": 2.25, "2x": 4.0}

os.environ.setdefault("MASTER_ADDR", "127.0.0.1")
os.environ.setdefault("MASTER_PORT", "12355")
os.environ.setdefault("RANK", "0")
os.environ.setdefault("WORLD_SIZE", "1")
os.environ.setdefault("LOCAL_RANK", "0")

_runner = None


def _ensure_runner():
    global _runner
    if _runner is not None:
        return _runner
    if not torch.distributed.is_initialized():
        init_torch(cudnn_benchmark=False)
    runner = VideoDiffusionInfer(load_config(os.path.join("./configs_3b", "main.yaml")))
    OmegaConf.set_readonly(runner.config, False)
    runner.configure_dit_model(device="cuda", checkpoint=DIT_CKPT)
    runner.configure_vae_model()
    # Disable the VAE's causal slicing for decode: every multi-frame run
    # so far died inside the sliced decode path (memory-state cache
    # machinery in the decoder up blocks), while the image port -- which
    # always decodes with slicing inactive (t=1) -- runs fine on this
    # runtime. With slicing off, each chunk decodes as one causal tensor,
    # the same code path the image port exercises. Conv-level memory
    # limits (0.15 GiB) still apply.
    try:
        runner.vae.set_causal_slicing(split_size=None, memory_device="same")
        print("[vae] causal slicing disabled", flush=True)
    except Exception as exc:
        print(f"[vae] could not disable slicing: {exc}", flush=True)
    if hasattr(runner.vae, "set_memory_limit"):
        runner.vae.set_memory_limit(**runner.config.vae.memory_limit)
    _runner = runner
    return _runner


def _probe(video_path):
    cap = cv2.VideoCapture(video_path)
    # CAP_PROP_FRAME_COUNT trusts container duration metadata, which is
    # wrong for stream-copied excerpts (it reported ~120 for a 48-frame
    # clip and inflated the ZeroGPU duration request). Count by grabbing.
    frames = 0
    while cap.grab():
        frames += 1
    if frames == 0:
        frames = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
    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 24.0
    cap.release()
    return frames, w, h, fps


def _target_area(w, h, mode):
    return int(min(w * h * MODE_AREA_SCALE[mode], MAX_TARGET_AREA))


def _duration(video, mode="restore", steps=1, seed=666, window="chunked", progress=None) -> int:
    # ZeroGPU's per-request maximum tracks the account's remaining GPU
    # quota; after a heavy session it drops below what an honest estimate
    # for even a short clip requests, and every job is rejected up front.
    # The dynamic estimate also never reliably saw the real frame count
    # at duration-eval time. Request the floor instead: overruns only
    # produce a warning (observed on the 384-frame run), not a kill.
    return 60


def _transform(target_area):
    return Compose(
        [
            NaResize(
                resolution=target_area**0.5, mode="area", downsample_only=False
            ),
            Lambda(lambda x: torch.clamp(x, 0.0, 1.0)),
            DivisibleCrop((16, 16)),
            Normalize(0.5, 0.5),
            Rearrange("t c h w -> c t h w"),
        ]
    )


def _read_frames(video_path, max_frames):
    cap = cv2.VideoCapture(video_path)
    fps = cap.get(cv2.CAP_PROP_FPS) or 24.0
    frames = []
    while len(frames) < max_frames:
        ok, frame = cap.read()
        if not ok:
            break
        frames.append(cv2.cvtColor(frame, cv2.COLOR_BGR2RGB))
    cap.release()
    if not frames:
        raise gr.Error("Could not read any frames from the uploaded video.")
    return np.stack(frames), fps  # (t, h, w, c) uint8


def _pad_chunk(chunk):
    """Pad a (t, c, h, w) tensor by repeating the last frame until
    (t - 1) % 4 == 0. Returns (padded, original_t)."""
    t = chunk.size(0)
    if t <= 4:
        need = 5 - t
    elif (t - 1) % 4 == 0:
        return chunk, t
    else:
        need = 4 - ((t - 1) % 4)
    pad = chunk[-1:].repeat(need, 1, 1, 1)
    return torch.cat([chunk, pad], dim=0), t


def _process_window(runner, raw_frames, transform, text_embeds):
    """Run one window of uint8 frames (t, h, w, c) through
    encode -> DiT -> decode; return restored uint8 frames
    (t, h_out, w_out, c) with the same frame count as the input."""
    tensor = torch.from_numpy(np.ascontiguousarray(raw_frames)).float().div_(255.0)
    tensor = tensor.permute(0, 3, 1, 2)
    tensor, ori_t = _pad_chunk(tensor)
    cond = transform(tensor.to("cuda"))
    reference = rearrange(cond, "c t h w -> t c h w")

    latents = runner.vae_encode([cond])
    noise = [torch.randn_like(latent) for latent in latents]
    aug_noise = [torch.randn_like(latent) for latent in latents]

    def _add_noise(x, aug):
        t = torch.tensor([1000.0], device="cuda") * INPUT_NOISE
        shape = torch.tensor(x.shape[1:], device="cuda")[None]
        return runner.schedule.forward(x, aug, runner.timestep_transform(t, shape))

    conditions = [
        runner.get_condition(n, task="sr", latent_blur=_add_noise(latent, a))
        for n, a, latent in zip(noise, aug_noise, latents)
    ]
    with torch.autocast("cuda", torch.bfloat16, enabled=True):
        videos = runner.inference(
            noises=noise, conditions=conditions, dit_offload=False, **text_embeds
        )
    sample = videos[0]
    sample = (
        rearrange(sample[:, None], "c t h w -> t c h w")
        if sample.ndim == 3
        else rearrange(sample, "c t h w -> t c h w")
    )
    if USE_COLOR_FIX:
        sample = wavelet_reconstruction(
            sample.to("cpu"), reference[: sample.size(0)].to("cpu")
        )
    else:
        sample = sample.to("cpu")
    sample = sample[:ori_t]
    sample = rearrange(sample, "t c h w -> t h w c")
    sample = sample.clip(-1, 1).mul_(0.5).add_(0.5).mul_(255).round().to(torch.uint8).numpy()
    del tensor, cond, reference, latents, noise, aug_noise, conditions, videos
    return sample


def _run_sliding(runner, frames_np, transform, text_embeds):
    """Sliding-window pass (restore mode only). Windows of CHUNK frames
    advance by SLIDE_STEP. Every window after the first opens with the
    previous window's last restored frame -- restored pixels, resized
    back to the input's frame size so the window runs through the exact
    same transform/inference path as chunked mode -- at position 0,
    where the causal VAE encodes it independently as the anchor for the
    frames that follow. That frame's re-render is averaged 0.5/0.5 with
    its already-finalized version; each window's new last frame is held
    back as the next window's anchor."""
    total = frames_np.shape[0]
    h_in, w_in = frames_np.shape[1], frames_np.shape[2]
    results = []
    pending = None  # last restored frame (uint8, output size), unfinalized
    next_src = 0
    win_no = 0
    while next_src < total:
        if pending is None:
            take = frames_np[next_src : next_src + CHUNK]
            valid_new = take.shape[0]
            win = list(take)
        else:
            ctx = cv2.resize(pending, (w_in, h_in), interpolation=cv2.INTER_LINEAR)
            take = frames_np[next_src : next_src + SLIDE_STEP]
            valid_new = take.shape[0]
            win = [ctx] + list(take)
        while len(win) < CHUNK:
            win.append(win[-1])
        out = _process_window(runner, np.stack(win), transform, text_embeds)
        last = next_src + valid_new >= total
        if pending is None:
            if last:
                results.extend(out[:valid_new])
            else:
                results.extend(out[: valid_new - 1])
                pending = out[valid_new - 1]
        else:
            averaged = np.clip(
                np.rint(
                    0.5 * out[0].astype(np.float32)
                    + 0.5 * pending.astype(np.float32)
                ),
                0,
                255,
            ).astype(np.uint8)
            results.append(averaged)
            if last:
                results.extend(out[1 : 1 + valid_new])
                pending = None
            else:
                results.extend(out[1:valid_new])
                pending = out[valid_new]
        next_src += valid_new
        win_no += 1
        print(f"[job] sliding window {win_no} done ({len(results)} frames finalized)", flush=True)
        gc.collect()
        torch.cuda.empty_cache()
    assert len(results) == total, f"sliding produced {len(results)} frames for {total} input frames"
    return np.stack(results)


@spaces.GPU(duration=_duration)
@torch.no_grad()
def restore_video(video, mode="restore", steps=1, seed=666, window="chunked", progress=gr.Progress(track_tqdm=True)):
    if video is None:
        raise gr.Error("Upload a video first.")
    in_path = video["video"]["path"] if isinstance(video, dict) else video

    print(f"[alloc] PYTORCH_CUDA_ALLOC_CONF={os.environ.get('PYTORCH_CUDA_ALLOC_CONF')!r}", flush=True)

    frames_np, fps = _read_frames(in_path, MAX_FRAMES)
    total = frames_np.shape[0]
    h, w = frames_np.shape[1], frames_np.shape[2]
    target_area = _target_area(w, h, mode)
    window_mode = window
    if window_mode == "sliding" and mode != "restore":
        print(f"[job] sliding applies to restore mode only; falling back to chunked for mode={mode}", flush=True)
        window_mode = "chunked"
    print(f"[job] {total} frames {w}x{h}@{fps:.3f} mode={mode} window={window_mode} target_area={target_area}", flush=True)

    runner = _ensure_runner()
    runner.config.diffusion.cfg.scale = 1.0
    runner.config.diffusion.cfg.rescale = 0.0
    runner.config.diffusion.timesteps.sampling.steps = int(steps)
    runner.configure_diffusion()
    set_seed(int(seed) % (2**32), same_across_ranks=True)

    text_embeds = {
        "texts_pos": [torch.load(POS_EMB).to("cuda")],
        "texts_neg": [torch.load(NEG_EMB).to("cuda")],
    }
    transform = _transform(target_area)

    if window_mode == "sliding":
        out_frames = _run_sliding(runner, frames_np, transform, text_embeds)
    else:
        out_chunks = []
        for start in range(0, total, CHUNK):
            raw = frames_np[start : start + CHUNK]
            sample = _process_window(runner, raw, transform, text_embeds)
            out_chunks.append(sample)
            print(f"[job] chunk {start}-{start + sample.shape[0]} done", flush=True)
            gc.collect()
            torch.cuda.empty_cache()
        out_frames = np.concatenate(out_chunks, axis=0)
    workdir = tempfile.mkdtemp(prefix="seedvr2-out-")
    silent_path = os.path.join(workdir, "video.mp4")
    writer = imageio.get_writer(
        silent_path,
        fps=fps,
        codec="libx264",
        quality=8,
        macro_block_size=None,
        ffmpeg_params=["-pix_fmt", "yuv420p", "-movflags", "+faststart"],
    )
    for frame in out_frames:
        writer.append_data(frame)
    writer.close()

    # Stream-copy the original audio track if there is one (no -shortest:
    # it truncates video tails against audio duration).
    final_path = os.path.join(workdir, "seedvr2-out.mp4")
    ffmpeg = imageio_ffmpeg.get_ffmpeg_exe()
    r = subprocess.run(
        [ffmpeg, "-y", "-i", silent_path, "-i", in_path,
         "-map", "0:v:0", "-map", "1:a?", "-c", "copy", final_path],
        capture_output=True, text=True,
    )
    if r.returncode != 0 or not os.path.exists(final_path):
        print(f"[job] audio mux failed, returning silent: {r.stderr[-300:]}", flush=True)
        final_path = silent_path
    print(f"[job] done -> {final_path} ({out_frames.shape})", flush=True)
    return final_path


# ---------------------------------------------------------------------------
# Image endpoint (merged from the Upsampler v3 image recipe, 2026-10-26).
# Differences from the v3 original, all deliberate: it shares this app's
# _ensure_runner (which disables VAE causal slicing -- inactive at t=1
# anyway) and it does NOT set expandable_segments, neither via env at
# import nor via torch.cuda.memory._set_allocator_settings at runtime:
# on this runtime that setting trips the NVML assert in the video VAE
# decoder, and the v3 author measured it made no difference there either.
# ---------------------------------------------------------------------------

def _image_transform():
    return Compose(
        [
            NaResize(resolution=IMAGE_WORK_AREA**0.5, mode="area", downsample_only=False),
            Lambda(lambda x: torch.clamp(x, 0.0, 1.0)),
            DivisibleCrop((16, 16)),
            Normalize(0.5, 0.5),
            Rearrange("t c h w -> c t h w"),
        ]
    )


def _image_duration(image, steps=1, seed=0, progress=None) -> int:
    # Every request restores at the same working resolution, so cost tracks
    # the step count. v3 formula: 12 s/step measured at 2MP, scaled by the
    # working area; at 3.7MP and 1 step this requests 35 s.
    area_scale = IMAGE_WORK_AREA / (1920 * 1080)
    return int(min(90, 14 + 12 * int(steps) * area_scale))


def _save_result(image):
    # Gradio would save a PIL/array result as WebP at quality 80, throwing
    # away the fine detail the model just produced. Write it ourselves at
    # quality 95 and return the path; Gradio serves a path as is.
    if not isinstance(image, Image.Image):
        image = Image.fromarray(image)
    out_dir = tempfile.mkdtemp(prefix="seedvr2-img-")
    if max(image.size) > WEBP_MAX_SIDE:
        path = os.path.join(out_dir, "upscaled.png")
        image.save(path, "PNG")
    else:
        path = os.path.join(out_dir, "upscaled.webp")
        image.save(path, "WEBP", quality=OUTPUT_WEBP_QUALITY, method=4)
    return path


@spaces.GPU(duration=_image_duration)
@torch.no_grad()
def upscale_image(image, steps=1, seed=0, progress=gr.Progress(track_tqdm=True)):
    if image is None:
        raise gr.Error("Upload an image first.")

    print(f"[alloc] PYTORCH_CUDA_ALLOC_CONF={os.environ.get('PYTORCH_CUDA_ALLOC_CONF')!r}", flush=True)
    # The v3 recipe's runtime allocator set, kept in its original form
    # (image function only, never at import -- the env form breaks the
    # video decoder). The 2026-10-26 3.7MP probe showed it does not change
    # the decode ceiling either way; at the 2MP default it matches the
    # public Space's known-good configuration.
    try:
        torch.cuda.memory._set_allocator_settings("expandable_segments:True")
        print("[alloc] expandable_segments set at runtime", flush=True)
    except Exception as exc:
        print(f"[alloc] runtime allocator setting unavailable: {exc}", flush=True)
    print(f"[img] work area {IMAGE_WORK_AREA} px", flush=True)

    runner = _ensure_runner()

    runner.config.diffusion.cfg.scale = 1.0
    runner.config.diffusion.cfg.rescale = 0.0
    runner.config.diffusion.timesteps.sampling.steps = int(steps)
    runner.configure_diffusion()
    # Explicit API seed wins; 0 falls back to the pinned default or random.
    import random

    chosen = int(seed) or IMAGE_DEFAULT_SEED or random.randrange(1, 2**31)
    print(f"[img] seed {chosen}", flush=True)
    set_seed(chosen, same_across_ranks=True)

    img = Image.open(image).convert("RGB") if isinstance(image, str) else image.convert("RGB")
    tensor = T.ToTensor()(img).unsqueeze(0)  # (t=1, c, h, w)

    cond = _image_transform()(tensor.to("cuda"))
    original = cond
    latents = runner.vae_encode([cond])

    text_embeds = {
        "texts_pos": [torch.load(POS_EMB).to("cuda")],
        "texts_neg": [torch.load(NEG_EMB).to("cuda")],
    }

    noise = [torch.randn_like(latent) for latent in latents]
    aug_noise = [torch.randn_like(latent) for latent in latents]

    def _add_noise(x, aug):
        t = torch.tensor([1000.0], device="cuda") * INPUT_NOISE
        shape = torch.tensor(x.shape[1:], device="cuda")[None]
        return runner.schedule.forward(x, aug, runner.timestep_transform(t, shape))

    conditions = [
        runner.get_condition(n, task="sr", latent_blur=_add_noise(latent, a))
        for n, a, latent in zip(noise, aug_noise, latents)
    ]

    with torch.autocast("cuda", torch.bfloat16, enabled=True):
        videos = runner.inference(
            noises=noise, conditions=conditions, dit_offload=False, **text_embeds
        )

    sample = videos[0]
    sample = (
        rearrange(sample[:, None], "c t h w -> t c h w")
        if sample.ndim == 3
        else rearrange(sample, "c t h w -> t c h w")
    )
    reference = (
        rearrange(original[:, None], "c t h w -> t c h w")
        if original.ndim == 3
        else rearrange(original, "c t h w -> t c h w")
    )
    if USE_COLOR_FIX:
        sample = wavelet_reconstruction(sample.to("cpu"), reference[: sample.size(0)].to("cpu"))
    else:
        sample = sample.to("cpu")

    sample = rearrange(sample, "t c h w -> t h w c")
    sample = sample.clip(-1, 1).mul_(0.5).add_(0.5).mul_(255).round().to(torch.uint8).numpy()

    del latents, conditions, videos
    gc.collect()
    torch.cuda.empty_cache()

    return _save_result(sample[0])


with gr.Blocks(title="SeedVR2 Video + Image") as demo:
    gr.Markdown(
        "## SeedVR2-3B Restoration\n"
        "**Video**: one-step diffusion restoration. **Restore** re-renders at the "
        "input's own resolution (deblur / detail rebuild); **1.5x / 2x** also "
        "enlarge. Up to 481 frames per job; longer clips are trimmed. "
        "**Window**: *chunked* processes independent 5-frame windows; "
        "*sliding* (restore only) carries each window's last restored frame "
        "into the next window as context and averages the overlap, for "
        "better temporal continuity at ~25% more compute.\n\n"
        "**Image**: single-image restoration at a 1920x1080 working area."
    )
    with gr.Tab("Video"):
        with gr.Row():
            with gr.Column():
                video_in = gr.Video(label="Video", sources=["upload"])
                mode = gr.Dropdown(
                    choices=["restore", "1.5x", "2x"], value="restore", label="Mode"
                )
                window = gr.Radio(
                    choices=["chunked", "sliding"],
                    value="chunked",
                    label="Window",
                    info="chunked: independent 5-frame windows. sliding (restore only): each window carries the previous restored frame as context; the overlap frame is averaged.",
                )
                steps = gr.Slider(1, 4, value=1, step=1, label="Steps")
                seed = gr.Number(value=666, precision=0, label="Seed")
                run = gr.Button("Restore Video", variant="primary")
            with gr.Column():
                video_out = gr.Video(label="Result")

        # Registered first so /restore stays fn_index 0.
        run.click(
            restore_video,
            inputs=[video_in, mode, steps, seed, window],
            outputs=[video_out],
            api_name="restore",
        )
    with gr.Tab("Image"):
        with gr.Row():
            with gr.Column():
                image_in = gr.Image(label="Image", type="filepath")
                img_steps = gr.Slider(1, 4, value=1, step=1, label="Steps")
                img_seed = gr.Number(value=0, precision=0, visible=False)
                img_run = gr.Button("Upscale Image", variant="primary")
            with gr.Column():
                image_out = gr.Image(label="Result", type="filepath")

        # No leading slash: gradio prefixes it.
        img_run.click(
            upscale_image,
            inputs=[image_in, img_steps, img_seed],
            outputs=[image_out],
            api_name="upscale_image",
        )

demo.launch(ssr_mode=False, show_error=True)