Spaces:
Running on Zero
Running on Zero
| """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) | |
| 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() | |