Blame for app.py

import os
import random
import tempfile
import time

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

import spaces  # ZeroGPU: must be imported before any CUDA-touching import

import torch
torch.set_grad_enabled(False)  # reference scripts/naive_ref.py runs fully no-grad

from huggingface_hub import hf_hub_download, snapshot_download
from omegaconf import OmegaConf

import gradio as gr
import imageio

# vendored single-GPU reference runtime shipped with the paper repo
import naive_rt.rolling_forcing.configs as _rf_configs_pkg
from naive_rt.rolling_forcing.pipeline.rolling_forcing_inference import (
    CausalInferencePipeline,
)

# ---------------------------------------------------------------------------
# Paths & constants
# ---------------------------------------------------------------------------
WAN_REPO = "Wan-AI/Wan2.1-T2V-1.3B"
CKPT_REPO = "SJTU-DENG-Lab/WaveForcing-1.3B"
_CFG_DIR = os.path.dirname(_rf_configs_pkg.__file__)

# latent geometry: (B, C, T, H, W) = (1, 16, T, 60, 104) -> 832x480 pixels
LATENT_C, LATENT_H, LATENT_W = 16, 60, 104
FRAMES_PER_LATENT = 4          # causal VAE: 4 pixel frames per latent frame
FPS = 16
FRAMES_PER_BLOCK = 3           # num_frame_per_block (must divide num_frames)
# UI label -> (repo subdir [no hyphen!], denoising-step list)
VARIANTS = {
    "4-step": ("4step", [1000, 750, 500, 250]),
    "5-step": ("5step", [1000, 800, 600, 400, 200]),
}
# rolling forcing needs at least one full window: n_blocks >= n_steps
# 5-step needs 5 blocks of 3 latent frames = 15; that minimum is valid for 4-step too
MIN_NUM_FRAMES = 15
DEFAULT_NUM_FRAMES = 24        # latent frames -> 96 pixel frames, ~6s @ 16fps
MAX_NUM_FRAMES = 48            # KV cache supports 24 blocks; keep 48 for headroom
MAX_SECONDS = 210

DESCRIPTION = """\
# Wave Forcing — Streaming Text-to-Video

**Wave Forcing: Towards Speed-of-Light Streaming Video Generation** — block-causal rolling-window diffusion on the
Wan2.1-T2V-1.3B backbone, distilled to 4 or 5 denoising steps per block. This demo runs the authors' single-GPU
reference pipeline (`naive_rt`) from the [paper repo](https://github.com/SJTU-DENG-Lab/WaveForcing) with the released
[SJTU-DENG-Lab/WaveForcing-1.3B](https://huggingface.co/SJTU-DENG-Lab/WaveForcing-1.3B) checkpoints. The paper's
speed-of-light numbers come from their multi-GPU WaveRT runtime; this Space is the faithful single-GPU quality
reference for the same weights.
"""


# ---------------------------------------------------------------------------
# Model loading (module scope, eager)
# ---------------------------------------------------------------------------
def _stage_wan_dir() -> str:
    """Materialize "wan_models/Wan2.1-T2V-1.3B/" in the local layout naive_rt expects.

    WanDiffusionWrapper hard-codes relative paths ("wan_models/Wan2.1-T2V-1.3B/...")
    for the T5 encoder, tokenizer, VAE and the CausalWanModel.from_pretrained dir.
    The WaveForcing checkpoints have the exact same 825 keys as the official base
    DiT, so the DiT safetensors inside the staged dir is a symlink to the released
    WaveForcing EMA checkpoint instead of the 5.68 GB base download.
    """
    return snapshot_download(
        WAN_REPO,
        allow_patterns=[
            "models_t5_umt5-xxl-enc-bf16.pth",
            "google/umt5-xxl/*",
            "Wan2.1_VAE.pth",
            "config.json",
        ],
        local_dir="wan_models/Wan2.1-T2V-1.3B",
    )


def _make_config(step_list):
    """Build the merged inference config exactly like scripts/naive_ref.py."""
    cfg = OmegaConf.merge(
        OmegaConf.load(os.path.join(_CFG_DIR, "default_config.yaml")),
        OmegaConf.load(os.path.join(_CFG_DIR, "rolling_forcing_dmd.yaml")),
    )
    cfg.denoising_step_list = step_list
    return cfg


def _link_checkpoint(repo_subdir: str) -> None:
    """Point the staged DiT file at the WaveForcing checkpoint for this variant."""
    dst = os.path.join(
        "wan_models", "Wan2.1-T2V-1.3B", "diffusion_pytorch_model.safetensors"
    )
    src = hf_hub_download(
        CKPT_REPO, f"{repo_subdir}/model.safetensors", repo_type="model"
    )
    os.makedirs(os.path.dirname(dst), exist_ok=True)
    try:
        os.remove(dst)
    except FileNotFoundError:
        pass
    os.symlink(src, dst)


def _load_all():
    """Build both step-list variants sharing one T5 encoder + VAE."""
    _stage_wan_dir()
    text_encoder = None
    vae = None
    pipes = {}
    for label, (repo_subdir, step_list) in VARIANTS.items():
        _link_checkpoint(repo_subdir)
        cfg = _make_config(step_list)
        pipe = CausalInferencePipeline(
            cfg, device="cuda", text_encoder=text_encoder, vae=vae
        )
        if text_encoder is None:
            # first variant constructs the shared components; later ones reuse
            text_encoder = pipe.text_encoder
            vae = pipe.vae
        # cast to bf16 like the reference (in-place for nn.Module)
        pipe = pipe.to(dtype=torch.bfloat16)
        pipes[label] = pipe
    # move shared components + both generators to GPU
    first = pipes["4-step"]
    first.text_encoder.to("cuda")
    first.vae.to("cuda")
    for pipe in pipes.values():
        pipe.generator.to("cuda")
    return pipes


print("[app] downloading weights & building pipelines ...", flush=True)
_t0 = time.perf_counter()
PIPES = _load_all()
print(f"[app] pipelines ready in {time.perf_counter() - _t0:.1f}s", flush=True)


# ---------------------------------------------------------------------------
# Inference
# ---------------------------------------------------------------------------
def _save_video(video: torch.Tensor, path: str) -> str:
    """Write a [B,F,C,H,W] float tensor in [0,1] to an mp4 (imageio, H.264)."""
    v = video.detach().float().cpu()
    if v.dim() == 5:
        v = v[0]
    v = v.permute(0, 2, 3, 1)  # [F,H,W,C]
    if v.max() <= 1.5:
        v = v * 255.0
    v = v.clamp(0, 255).to(torch.uint8).numpy()
    imageio.mimsave(path, list(v), fps=FPS, codec="libx264",
                    macro_block_size=None, quality=8)
    return path


def _duration_estimator(prompt, variant, num_frames, *args, **kwargs):
    """ZeroGPU duration estimate, fitted from measured walls on this Space.

    Measured: 4-step/24f = 25.2s cold (19.8s warm), 5-step/24f = 22.7s,
    4-step/48f = 40.6s. Window count is nf/3 + steps - 1, giving
    wall ~= 3 + 2.2 * windows; add 15% headroom for cold-start attach.
    """
    try:
        nf = int(num_frames)
    except Exception:
        nf = DEFAULT_NUM_FRAMES
    steps = 4 if str(variant).startswith("4") else 5
    windows = nf // FRAMES_PER_BLOCK + steps - 1
    est = 3 + 2.2 * windows
    return min(MAX_SECONDS, int(est * 1.15) + 2)


@spaces.GPU(duration=_duration_estimator)
def generate(
    prompt: str,
    variant: str = "4-step",
    num_frames: int = DEFAULT_NUM_FRAMES,
    seed: int = -1,
    progress=gr.Progress(),
):
    """Generate a video from a text prompt with Wave Forcing.

    Args:
        prompt: Text description of the video to generate.
        variant: "4-step" or "5-step" distilled checkpoint.
        num_frames: Number of latent frames; 4 pixel frames each, must be a multiple of 3.
        seed: RNG seed; -1 randomizes.
        progress: Gradio progress bar.
    """
    variant = "4-step" if str(variant).startswith("4") else "5-step"
    num_frames = int(num_frames)
    # clamp to a multiple of 3 in [3*steps, MAX]
    num_frames = max(FRAMES_PER_BLOCK * len(VARIANTS[variant][1]),
                     (num_frames // FRAMES_PER_BLOCK) * FRAMES_PER_BLOCK)
    num_frames = min(num_frames, MAX_NUM_FRAMES)
    if seed is None or int(seed) < 0:
        seed = random.randint(0, 2**31 - 1)
    seed = int(seed)

    pipe = PIPES[variant]

    # The reference runner (scripts/naive_ref.py) disables grad globally, but
    # torch GradMode is THREAD-LOCAL and ZeroGPU executes @spaces.GPU tasks in
    # a worker thread, where grad is re-enabled. With grad on, every rolling
    # window's activation graph accumulates (via `output[:, a:b] = pred`) and
    # OOMs by the second window. So enforce no_grad inside the task itself.
    print(f"[app] grad enabled in task thread: {torch.is_grad_enabled()}",
          flush=True)
    with torch.no_grad():
        # Seed the global CUDA RNG too: the rolling-forcing loop's intermediate
        # add_noise draws use torch.randn_like (global RNG), so without this a
        # fixed seed only reproduces the INITIAL noise. With both seeded, a
        # fixed seed is reproducible run-to-run.
        torch.cuda.manual_seed(seed)
        # draw noise exactly like scripts/naive_ref.py: (B,C,T,H,W) then permute
        shape_bcthw = [1, LATENT_C, num_frames, LATENT_H, LATENT_W]
        g = torch.Generator(device="cuda").manual_seed(seed)
        noise = torch.randn(shape_bcthw, generator=g, device="cuda",
                            dtype=torch.bfloat16).permute(0, 2, 1, 3, 4)

        t0 = time.perf_counter()
        video = pipe.inference_rolling_forcing(noise=noise, text_prompts=[prompt])
        elapsed = time.perf_counter() - t0

    out_path = tempfile.NamedTemporaryFile(suffix=".mp4", delete=False).name
    _save_video(video, out_path)
    n_pix = num_frames * FRAMES_PER_LATENT
    info = (f"seed {seed} · {variant} · {num_frames} latent frames "
            f"({n_pix} px frames) · {elapsed:.1f}s · {n_pix / elapsed:.1f} fps")
    return out_path, info


# ---------------------------------------------------------------------------
# UI
# ---------------------------------------------------------------------------
EXAMPLES = [
    ["A cinematic shot of a fluffy corgi running on a sunny beach, waves in the background.", "4-step", 24],
    ["A dynamic action shot in the style of a high-energy sports magazine spread, featuring a golden retriever sprinting with all its might after a red sports car speeding down the road.", "4-step", 24],
    ["A high-speed action shot of a cheetah in its natural habitat, sprinting at full speed while chasing its prey across the savanna.", "4-step", 24],
    ["A handheld camera captures a dog running through a park with a joyful exploration, the camera following the dog closely and bouncing and tilting with its movements.", "4-step", 33],
    ["A dynamic hip-hop dance scene in a vibrant urban style, featuring an Asian girl in a bright yellow T-shirt and white pants.", "5-step", 24],
]

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(DESCRIPTION)
        with gr.Row():
            prompt = gr.Textbox(
                label="Prompt",
                placeholder="Describe the video you want ...",
                lines=3,
                scale=4,
            )
            run_btn = gr.Button("Generate", variant="primary", scale=1)
        output = gr.Video(label="Generated video")
        info = gr.Markdown()
        with gr.Accordion("Advanced settings", open=False):
            variant_radio = gr.Radio(
                choices=["4-step", "5-step"],
                value="4-step",
                label="Distilled checkpoint",
                info="4-step [1000,750,500,250] or 5-step [1000,800,600,400,200] DMD steps per block",
            )
            num_frames_slider = gr.Slider(
                minimum=MIN_NUM_FRAMES, maximum=MAX_NUM_FRAMES,
                value=DEFAULT_NUM_FRAMES, step=3,
                label="Latent frames",
                info=f"multiples of {FRAMES_PER_BLOCK}; each latent frame = {FRAMES_PER_LATENT} pixel frames @ {FPS} fps",
            )
            seed_input = gr.Number(label="Seed", value=-1, precision=0,
                                  info="-1 = random")

        gr.Examples(
            examples=EXAMPLES,
            inputs=[prompt, variant_radio, num_frames_slider],
            fn=generate,
            outputs=[output, info],
            cache_examples=True,
            cache_mode="lazy",
        )

    run_btn.click(
        generate,
        inputs=[prompt, variant_radio, num_frames_slider, seed_input],
        outputs=[output, info],
        api_name="generate",
        queue=True,
    )

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