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)