app.py

"""SwiftVR: Real-Time One-Step Generative Video Restoration.

Gradio Space demo that loads the SwiftVR model from H-oliday/SwiftVR on the
Hugging Face Hub and restores / upscales user-uploaded videos in real time.
"""

import os
os.environ.setdefault("PYTORCH_CUDA_ALLOC_CONF", "expandable_segments:True")

import spaces
import sys
import tempfile
import traceback
import math

# Make the bundled swiftvr package importable
sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))

import torch
import torch.nn.functional as F
import gradio as gr
import numpy as np
import imageio
from PIL import Image
from huggingface_hub import snapshot_download

import decord
decord.bridge.set_bridge("torch")

from swiftvr import SwiftVRPipeline
from swiftvr.io import (
    get_video_info, iter_video_clips_fixed_scheme,
    preprocess_clip_uint8, crop_spatial_padding_ntchw,
    ntchw_to_uint8_frames, open_stream_video_writer,
)
from swiftvr.streaming.chunk import ChunkType

CHECKPOINT_REPO = "H-oliday/SwiftVR"
CHECKPOINT_DIR = "/tmp/swiftvr_checkpoints"

# --- Load model at module scope (ZeroGPU rule) --- #

os.makedirs(CHECKPOINT_DIR, exist_ok=True)

print("Downloading SwiftVR checkpoints...")
snapshot_download(
    repo_id=CHECKPOINT_REPO,
    local_dir=CHECKPOINT_DIR,
    repo_type="model",
)
print("Checkpoints downloaded. Loading model...")

pipe = SwiftVRPipeline.from_pretrained(CHECKPOINT_DIR)
pipe.to("cuda", dtype="bfloat16", attention_backend="sdpa")
print("SwiftVR model loaded and moved to CUDA with SDPA attention.")


def _aligned_pad(size, multiple=32):
    return (multiple - size % multiple) % multiple


@spaces.GPU(duration=90)
def restore_video(
    input_video: str,
    upscale: int = 2,
    clip_len: int = 24,
    quality: int = 60,
):
    """Restore a low-quality video using SwiftVR.

    Args:
        input_video: Path to the input low-quality video file.
        upscale: Upscale factor (2, 4, or 8).
        clip_len: Processing chunk size (must be multiple of 4).
        quality: Output quality 0-100 (maps to x265 CRF).

    Returns:
        Path to the restored output video, or an error message.
    """
    if input_video is None:
        return None, "Please upload or select a video first."

    tmp_dir = tempfile.mkdtemp()
    output_path = os.path.join(tmp_dir, "restored.mp4")

    try:
        # Single-threaded pipeline to avoid ZeroGPU fork conflicts
        device = pipe.device
        dtype = pipe.dtype

        # Get video info
        raw_total, lq_h, lq_w, src_fps = get_video_info(input_video, fallback_fps=30)
        total_frames = 4 * ((raw_total - 1) // 4) + 1

        # Compute target size
        out_h = lq_h * upscale
        out_w = lq_w * upscale
        pad_h = _aligned_pad(out_h)
        pad_w = _aligned_pad(out_w)

        # Reset streaming state
        pipe.tae_stream.reset()
        pipe.dit_stream.reset()
        pipe.dit_stream.overlap = 0

        n_lat = clip_len // 4
        prev_dit_out_cpu = None

        # Open video writer
        writer = open_stream_video_writer(
            output_path, fps=src_fps, video_format="",
            preset="", quality=quality)

        # Process chunks (single-threaded, no pipeline threads)
        clips = iter_video_clips_fixed_scheme(
            input_video, clip_len=clip_len,
            total_frames=total_frames,
            crop_h=lq_h, crop_w=lq_w)

        total_written = 0
        for spec, cpu_rgb in clips:
            # Move to GPU
            gpu_rgb = cpu_rgb.to(device=device)

            # Preprocess
            clip_rgb = preprocess_clip_uint8(
                gpu_rgb, out_h=out_h, out_w=out_w,
                mode=pipe.upscale_mode, pad_h=pad_h, pad_w=pad_w,
                dtype=dtype)

            # Encode
            z = pipe.tae_stream.encode_chunk_fixed(clip_rgb, spec)

            # Denoise
            if spec.ctype == ChunkType.LAST:
                z_ntchw = pipe.dit_stream.denoise_last_chunk(
                    z, spec, pipe.prompt_emb, prev_dit_out_cpu,
                    n_lat, device, dtype)
            else:
                z_bcfhw = z.permute(0, 2, 1, 3, 4).contiguous()
                z_den = pipe.dit_stream.denoise(z_bcfhw, pipe.prompt_emb)
                z_ntchw = z_den.permute(0, 2, 1, 3, 4).contiguous()
                prev_dit_out_cpu = z_bcfhw[:, :, -n_lat:].detach().cpu().clone()

            # Decode
            rgb_out = pipe.tae_stream.decode_chunk_fixed(z_ntchw, spec)
            if rgb_out is not None and rgb_out.shape[1] > 0:
                rgb_out = crop_spatial_padding_ntchw(rgb_out, pad_h, pad_w).detach()
                frames = ntchw_to_uint8_frames(rgb_out)
                if frames is not None:
                    for frame in frames:
                        writer.append_data(frame)
                    total_written += frames.shape[0]

            # Clean up GPU memory
            del gpu_rgb, clip_rgb, z, z_ntchw, rgb_out
            torch.cuda.empty_cache()

        writer.close()
        msg = f"Restored {total_written} frames from {raw_total} input frames at {src_fps:.1f} fps"
        return output_path, msg

    except Exception as e:
        err_msg = f"Error: {e}\n{traceback.format_exc()}"
        print(err_msg)
        return None, err_msg


CSS = """
#col-container { max-width: 1100px; margin: 0 auto; }
.dark .gradio-container { color: var(--body-text-color); }
"""

with gr.Blocks() as demo:
    with gr.Column(elem_id="col-container"):
        gr.Markdown(
            """
            # SwiftVR: Real-Time One-Step Generative Video Restoration
            Upload a low-quality video to restore and upscale it in real time using the SwiftVR model.
            [Paper](https://arxiv.org/abs/2606.09516) | [GitHub](https://github.com/H-oliday/SwiftVR) | [Model](https://huggingface.co/H-oliday/SwiftVR)
            """
        )

        with gr.Row():
            input_video = gr.Video(label="Input Video (low quality)", sources=["upload"])
            output_video = gr.Video(label="Restored Video", interactive=False)

        status = gr.Textbox(label="Status", interactive=False)

        with gr.Accordion("Advanced Settings", open=False):
            upscale = gr.Slider(
                minimum=1, maximum=4, step=1, value=2,
                label="Upscale Factor",
                info="How many times to upscale the input video resolution. Use 2 for faster processing."
            )
            clip_len = gr.Slider(
                minimum=4, maximum=48, step=4, value=24,
                label="Clip Length",
                info="Processing chunk size (must be multiple of 4). Larger = more context, slower."
            )
            quality = gr.Slider(
                minimum=0, maximum=100, step=5, value=60,
                label="Output Quality",
                info="0-100, maps to x265 CRF. Higher = better quality, larger file."
            )

        run_btn = gr.Button("Restore Video", variant="primary")

    gr.Examples(
        examples=[
            ["examples/cat.mp4", 2, 24, 60],
        ],
        inputs=[input_video, upscale, clip_len, quality],
        outputs=[output_video, status],
        fn=restore_video,
        cache_examples=True,
        cache_mode="lazy",
    )

    run_btn.click(
        fn=restore_video,
        inputs=[input_video, upscale, clip_len, quality],
        outputs=[output_video, status],
    )

if __name__ == "__main__":
    demo.launch(mcp_server=True, show_error=True)