"""Wave Forcing — faithful single-GPU demo using the paper's own naive_rt pipeline. Why not plain diffusers: the WaveForcing-1.3B checkpoints are *block-causal* DMD distillations of Wan2.1-T2V-1.3B. Run through the standard bidirectional WanPipeline they produce broken video — the weights expect rolling-window causal attention with a KV cache. This Space therefore clones the authors' repo (Apache-2.0) and runs their single-GPU reference pipeline, naive_rt/rolling_forcing (the same pipeline scripts/naive_ref.py uses), with the released 4-step / 5-step checkpoints symlinked in as the DiT. Paper: https://arxiv.org/abs/2611.00859 Repo: https://github.com/SJTU-DENG-Lab/WaveForcing Wts: https://huggingface.co/SJTU-DENG-Lab/WaveForcing-1.3B The headline 134.8 fps streaming numbers come from the authors' multi-GPU WaveRT runtime (8xH200); this Space is the faithful single-GPU quality reference for the same released weights. """ import os import random import shutil import subprocess import sys 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 import gradio as gr import imageio from huggingface_hub import hf_hub_download, snapshot_download from omegaconf import OmegaConf # --------------------------------------------------------------------------- # Constants # --------------------------------------------------------------------------- REPO_URL = "https://github.com/SJTU-DENG-Lab/WaveForcing" REPO_COMMIT = "fde159c" # pinned tree this Space was built against WORK_DIR = "/tmp/waveforcing" # writable scratch: repo clone + staged Wan weights REPO_DIR = os.path.join(WORK_DIR, "WaveForcing") WAN_REPO = "Wan-AI/Wan2.1-T2V-1.3B" # flat layout: T5 encoder, tokenizer, VAE CKPT_REPO = "SJTU-DENG-Lab/WaveForcing-1.3B" # distilled DiT safetensors # 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) 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 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](https://arxiv.org/abs/2611.00859)** (Hugging Face Daily Papers, Nov 3 2026): block-causal rolling-window video diffusion on a Wan2.1-1.3B backbone, distilled to **4–5 denoising steps per block**. This demo runs the authors' own single-GPU reference pipeline (`naive_rt`, [WaveForcing 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 (Apache-2.0). The paper's 134.8 fps streaming numbers come from their 8×H200 WaveRT runtime — this Space is the faithful single-GPU quality reference for the same weights. """ # --------------------------------------------------------------------------- # Setup: clone pinned repo, stage Wan base, load both checkpoint variants # --------------------------------------------------------------------------- def _ensure_repo() -> str: """Shallow-fetch the pinned WaveForcing commit (Apache-2.0) into WORK_DIR.""" if os.path.isdir(os.path.join(REPO_DIR, "naive_rt")): return REPO_DIR if shutil.which("git") is None: raise RuntimeError("git is not available in this Space image") os.makedirs(REPO_DIR, exist_ok=True) def run(*args): subprocess.run(args, check=True) # NOTE: GitHub does not serve bare short SHAs as fetch refs, so fetch the # branch head instead (REPO_COMMIT below documents the tree this was built against). run("git", "clone", "-q", "--depth", "1", REPO_URL, REPO_DIR) return REPO_DIR def _stage_wan_dir() -> str: """Materialize wan_models/Wan2.1-T2V-1.3B/ in the 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. Only the non-DiT components come from the Wan base repo; the DiT safetensors is symlinked to the released WaveForcing EMA checkpoint instead. """ 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.""" import naive_rt.rolling_forcing.configs as _cfgs cfg_dir = os.path.dirname(_cfgs.__file__) 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.""" from naive_rt.rolling_forcing.pipeline.rolling_forcing_inference import ( CausalInferencePipeline, ) _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 pipe = pipe.to(dtype=torch.bfloat16) # in-place for nn.Module pipes[label] = pipe first = pipes["4-step"] first.text_encoder.to("cuda") first.vae.to("cuda") for pipe in pipes.values(): pipe.generator.to("cuda") return pipes os.makedirs(WORK_DIR, exist_ok=True) os.chdir(WORK_DIR) # wan_wrapper hard-codes *relative* wan_models/... paths _ensure_repo() sys.path.insert(0, REPO_DIR) 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. Window count is nf/3 + steps - 1; wall ~= 3 + 2.2 * windows on an H200 slice; 15% headroom added.""" 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, variant="4-step", num_frames=DEFAULT_NUM_FRAMES, seed=-1, progress=gr.Progress()): variant = "4-step" if str(variant).startswith("4") else "5-step" num_frames = int(num_frames) 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) if not prompt or not prompt.strip(): return None, "Type a prompt first." pipe = PIPES[variant] # torch GradMode is THREAD-LOCAL and ZeroGPU runs @spaces.GPU tasks in a # worker thread — enforce no_grad inside the task itself or every rolling # window's activation graph accumulates and OOMs. with torch.no_grad(): # Seed the global CUDA RNG too: the rolling-forcing loop draws # intermediate noise with torch.randn_like (global RNG), so a fixed # seed only reproduces the INITIAL noise without this. torch.cuda.manual_seed(seed) 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, {n_pix / FPS:.1f}s) · {elapsed:.1f}s · **{n_pix / elapsed:.1f} fps** single-GPU" ) 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 of a golden retriever sprinting after a red sports car speeding down the road.", "4-step", 24], ["A high-speed action shot of a cheetah sprinting at full speed across the savanna.", "4-step", 24], ["A handheld camera follows a dog running joyfully through a park, bouncing and tilting with its movements.", "4-step", 33], ["A dynamic hip-hop dance scene in a vibrant urban style, featuring a girl in a bright yellow T-shirt and white pants.", "5-step", 24], ] with gr.Blocks(title="Wave Forcing — streaming text-to-video") as demo: 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( ["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( MIN_NUM_FRAMES, 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, inputs=[prompt, variant_radio, num_frames_slider], label="Example prompts") run_btn.click(generate, [prompt, variant_radio, num_frames_slider, seed_input], [output, info]) gr.Markdown( """ --- *Weights: [SJTU-DENG-Lab/WaveForcing-1.3B](https://huggingface.co/SJTU-DENG-Lab/WaveForcing-1.3B) (Apache-2.0) · Pipeline: the authors' `naive_rt` single-GPU reference from the [WaveForcing repo](https://github.com/SJTU-DENG-Lab/WaveForcing) · Base components: [Wan-AI/Wan2.1-T2V-1.3B](https://huggingface.co/Wan-AI/Wan2.1-T2V-1.3B) · Paper: [arXiv:2611.00859](https://arxiv.org/abs/2611.00859).* """ ) if __name__ == "__main__": demo.launch()