app.py

"""FlashVSR+ — one-step video super-resolution demo (turbo preset).

Loads the FlashVSR+ turbo checkpoint (hybrid sliding attention + Triton
block-sparse stream attention + TinyDecoder) and super-resolves an uploaded
video to a target resolution that the user picks. The requested height/width
is the OUTPUT resolution: the input is center-cropped and resized to it, and
the output video is written at that size with the input's FPS (and audio,
when present).
"""

import spaces  # noqa: F401  # must precede any torch / CUDA import
import os
import shutil
import tempfile
import time

import gradio as gr
import torch
from huggingface_hub import snapshot_download

WEIGHTS = snapshot_download(
    "ruchangyao/FlashVSR_plus",
    allow_patterns=["turbo/*", "TCDecoder.ckpt"],
)
CONTEXT_PATH = os.path.join(os.path.dirname(os.path.abspath(__file__)), "prompt_tensor", "posi_prompt.pth")

from inference.presets import stream_config_for_mode  # noqa: E402
from inference.loading import ModelPaths  # noqa: E402
from inference.pipeline import build_inference_pipeline  # noqa: E402
from inference.request import InferenceRequest  # noqa: E402

# Output resolutions (both sides divisible by 128 — a hard requirement).
RESOLUTIONS = [
    ("1024 × 1920 (README base example)", 1024, 1920),
    ("512 × 896", 512, 896),
    ("384 × 640", 384, 640),
    ("2048 × 3840 (4K, README turbo example)", 2048, 3840),
    ("2176 × 3840 (4K, README incremental example)", 2176, 3840),
]

DEFAULT_HEIGHT, DEFAULT_WIDTH = 512, 896

# Incremental input/decoding keeps RAM bounded at high resolutions: only one
# decode segment of frames is resident at a time (README's 4K path).
DECODE_EVERY = 28


def _build_pipeline():
    config = stream_config_for_mode(
        "turbo",
        device="cuda",
        seed=2026,
        decode_every=DECODE_EVERY,
        temporal_compression=4,
        attn_kernel="auto",
        color_correction=False,
    )
    paths = ModelPaths(
        dit=os.path.join(WEIGHTS, "turbo", "dit.safetensors"),
        lq_projection=os.path.join(WEIGHTS, "turbo", "lq_proj_in.safetensors"),
        context=CONTEXT_PATH,
        tiny_decoder=os.path.join(WEIGHTS, "TCDecoder.ckpt"),
    )
    return build_inference_pipeline(config, paths)


print("Loading FlashVSR+ turbo weights (~7 GB) ...", flush=True)
t_start = time.perf_counter()
pipeline = _build_pipeline()
print(f"FlashVSR+ pipeline ready in {time.perf_counter() - t_start:.1f}s", flush=True)


# Measured warm-worker GPU cost on ZeroGPU (turbo preset, compression 4):
# (megapixels, fixed seconds, per-frame seconds). Runtime is dominated by a
# fixed per-resolution cost up to ~1 MPix, then a per-frame term that grows
# with resolution (the TinyDecoder decode segments get expensive at 4K).
#   384x640 / 65f  -> 5.6 s    512x896 / 33,97f -> 9.7, 10.0 s
#   1024x1920 / 33f -> 15.4 s  2048x3840 / 33,97f -> 53.3, 112.3 s
_PIXEL_ANCHORS = [
    (384 * 640 / 1e6, 6.3, 0.002),
    (512 * 896 / 1e6, 10.2, 0.008),
    (1024 * 1920 / 1e6, 14.0, 0.073),
    (2048 * 3840 / 1e6, 22.9, 0.922),
    (2176 * 3840 / 1e6, 24.5, 0.98),
]


def _gpu_cost_seconds(height: int, width: int, frames: int, compression: int) -> float:
    """Interpolate the measured GPU cost of one upscale call, in seconds."""
    pixels = height * width / 1e6
    fixed, per_frame = _PIXEL_ANCHORS[-1][1], _PIXEL_ANCHORS[-1][2]
    for i in range(len(_PIXEL_ANCHORS) - 1):
        p0, f0, pf0 = _PIXEL_ANCHORS[i]
        p1, f1, pf1 = _PIXEL_ANCHORS[i + 1]
        if pixels <= p1:
            span = (pixels - p0) / (p1 - p0)
            fixed = f0 + (f1 - f0) * span
            per_frame = pf0 + (pf1 - pf0) * span
            break
    # Compression 1 measured ~1.34x the cost of compression 4 at 512x896;
    # (4/c)^0.25 = 1.41 at c=1, a close, slightly conservative proxy.
    estimate = (fixed + per_frame * frames) * (4.0 / compression) ** 0.25
    # 15% slack + 2 s covers run-to-run noise and a cold worker streaming
    # weights into VRAM (measured cold-vs-warm delta was < 1 s).
    return min(max(estimate * 1.15 + 2.0, 12.0), 300.0)


def _estimate_duration(video, resolution, temporal_compression, color_correction, max_frames, *args, **kwargs):
    """Duration callable for @spaces.GPU: scales with measured cost.

    ZeroGPU calls this with the exact request arguments before scheduling, so
    the quota charged to a visitor tracks what their request really costs —
    calibrated from smoke tests on this Space's hardware.
    """
    del video, color_correction, args, kwargs
    _, height, width = RESOLUTIONS[resolution]
    frames = max(int(max_frames), 17) if max_frames else 128
    return round(_gpu_cost_seconds(height, width, frames, int(temporal_compression)))


@spaces.GPU(duration=_estimate_duration)
def upscale_video(
    video,
    resolution: int,
    temporal_compression: int,
    color_correction: bool,
    max_frames: int,
    progress: gr.Progress = gr.Progress(),
):
    """Super-resolve a video with FlashVSR+ (one-step, turbo preset).

    Args:
        video: Input video path (any resolution/FPS; audio is preserved).
        resolution: Index into RESOLUTIONS; the picked height/width is the
            OUTPUT resolution (input is center-cropped and resized to it).
        temporal_compression: Input frames per latent — 4 (fast, default),
            2 or 1 for large motion (slower).
        color_correction: Match output low-frequency color to the LQ input.
        max_frames: Cap on input frames processed (minimum 9 for one latent).

    Returns:
        The super-resolved MP4 path.
    """
    if video is None:
        raise gr.Error("Please upload or pick an input video first.")
    _, height, width = RESOLUTIONS[resolution]

    def report(message: str) -> None:
        progress(0.5, desc=message)

    frame_cap = max(int(max_frames), 17)
    # Worker threads fork COW from this process; config is rebuilt per call to
    # keep the fields deterministic (seed/compression are request-scoped).
    pipeline.config.seed = 2026
    pipeline.config.temporal_compression = int(temporal_compression)
    pipeline.config.decoder.decode_every = DECODE_EVERY
    if color_correction:
        from inference.postprocessing import WaveletColorCorrector

        if not isinstance(pipeline.color_corrector, WaveletColorCorrector):
            pipeline.color_corrector = WaveletColorCorrector(levels=5, chunk_size=16)
    else:
        pipeline.color_corrector = None

    out_dir = tempfile.mkdtemp(prefix="flashvsr_out_")
    output_path = os.path.join(out_dir, f"flashvsr_{width}x{height}.mp4")
    try:
        result = pipeline.infer(
            InferenceRequest(
                input_path=video,
                output_path=output_path,
                height=height,
                width=width,
                max_frames=frame_cap,
            ),
            progress=report,
            video_preset="veryfast",
        )
    except Exception as error:
        shutil.rmtree(out_dir, ignore_errors=True)
        raise gr.Error(f"Super-resolution failed: {error}") from error
    progress(1.0, desc="Done")
    return result.output_path


WARNING_MD = (
    "**Heads-up:** higher resolutions and lower temporal compression cost "
    "real GPU time. Start with the default (512 × 896, compression 4) or use "
    "the examples; 4K targets can take minutes on ZeroGPU."
)

with gr.Blocks(title="FlashVSR+ Video Super-Resolution") as demo:
    gr.Markdown(
        "# ⚡ FlashVSR+ — Efficient Video Super-Resolution\n"
        "One-step diffusion VSR with hybrid attention and a TinyDecoder. "
        "Pick a target **output** resolution; your video is center-cropped, "
        "resized, and super-resolved to that size at the input's FPS (audio kept)."
    )
    with gr.Row():
        with gr.Column():
            video_in = gr.Video(label="Input video (low resolution)", sources=["upload"])
            resolution = gr.Dropdown(
                choices=[r[0] for r in RESOLUTIONS],
                type="index",
                value=RESOLUTIONS[1][0],
                label="Output resolution",
                info="Target output size (both sides must be multiples of 128); input is center-cropped and resized to it",
            )
            with gr.Accordion("Advanced options", open=False):
                temporal_compression = gr.Radio(
                    choices=[("4 (fast, default)", 4), ("2 (large motion)", 2), ("1 (largest motion, slowest)", 1)],
                    value=4,
                    type="value",
                    label="Temporal compression",
                    info="Input frames per latent frame — lower preserves motion detail at the cost of speed",
                )
                color_correction = gr.Checkbox(
                    value=False,
                    label="Match low-frequency color to input",
                )
                max_frames = gr.Slider(
                    minimum=17,
                    maximum=129,
                    value=97,
                    step=4,
                    label="Max input frames",
                    info="Caps processing time; long videos are truncated to this many frames (17 ≈ 1 latent chunk, minimum)",
                )
            submit = gr.Button("Super-resolve", variant="primary")
            gr.Markdown(WARNING_MD)
        with gr.Column():
            video_out = gr.Video(label="Super-resolved video", autoplay=False)

    submit.click(
        fn=upscale_video,
        inputs=[video_in, resolution, temporal_compression, color_correction, max_frames],
        outputs=video_out,
        api_name="upscale",
    )

    gr.Examples(
        examples=[
            ["examples/case_1.mp4", RESOLUTIONS[1][0], 4, False, 33],
            ["examples/case_0.mp4", RESOLUTIONS[1][0], 4, False, 97],
            ["examples/case_2.mp4", RESOLUTIONS[2][0], 4, False, 65],
        ],
        inputs=[video_in, resolution, temporal_compression, color_correction, max_frames],
        outputs=video_out,
        fn=upscale_video,
        cache_examples=True,
        cache_mode="lazy",
        run_on_click=True,
        label="Example inputs (from the FlashVSR+ repo)",
    )

demo.launch(theme=gr.themes.Citrus(), mcp_server=True)