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