"""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)