From bbb1f9781b918f407e0da85dcaaaeb9bb091a2ad Mon Sep 17 00:00:00 2001 From: Enderfga Date: Tue, 19 May 2026 16:00:40 +0800 Subject: [PATCH] [demo] AnyFlow 14B T2V inference demo (NFE=4 + NFE=50) Single-rank demo script that loads nvidia/AnyFlow-Wan2.1-T2V-14B-Diffusers into FastVideo's WanTransformer3DModel via param_names_mapping, then samples 81 frames at 480x832 with the new FlowMapEulerDiscreteScheduler at both NFE=4 (~30s) and NFE=50 (~5min) and decodes each via the Wan VAE (tiling enabled). Uses guidance_scale=1.0 since the on-policy distilled checkpoint has fuse_guidance_scale=3.0 baked into the weights, matching the AnyFlow paper's official demo.py default. Memory tactics for single H200 (141 GB HBM): - Encode prompts with UMT5 first, free the text encoder. - Build 14B transformer (~28 GB bf16), load AnyFlow shards. - Sample at both NFEs, free transformer, then VAE decode with tiling. Peak GPU usage ~57 GB on a single H200. --- scripts/demo_anyflow_14b.py | 279 ++++++++++++++++++++++++++++++++++++ 1 file changed, 279 insertions(+) create mode 100644 scripts/demo_anyflow_14b.py diff --git a/scripts/demo_anyflow_14b.py b/scripts/demo_anyflow_14b.py new file mode 100644 index 000000000..129dbb372 --- /dev/null +++ b/scripts/demo_anyflow_14b.py @@ -0,0 +1,279 @@ +#!/usr/bin/env python3 +# SPDX-License-Identifier: Apache-2.0 +"""FastVideo-side AnyFlow 14B T2V demo at NFE=4 and NFE=50. + +Loads ``nvidia/AnyFlow-Wan2.1-T2V-14B-Diffusers`` into FastVideo's +``WanTransformer3DModel`` (with the ``param_names_mapping`` regex +handling the ``delta_embedder`` rename), runs the +``FlowMapEulerDiscreteScheduler`` for the requested NFE schedule, and +saves the decoded video as MP4. Matches the prompt / shift / guidance +recipe used by the parallel FastGen demo so the videos are directly +comparable. + +Memory tactics (single H200, 141 GB HBM): +- Encode prompts with UMT5, free the encoder. +- Build the FastVideo Wan-14B transformer, load AnyFlow safetensor + shards via param_names_mapping translation. +- Sample at both NFEs without re-loading the transformer. +- Free the transformer, then load the Wan VAE with tiling for decode. + +Run via:: + + srun --jobid=304 bash -c ' + source ~/.venvs/anyflow/bin/activate + export CPATH=/home/guian/.local/share/uv/python/cpython-3.10.20-linux-x86_64-gnu/include/python3.10 + cd ~/projects/anyflow/FastVideo + PYTHONPATH=$PWD python scripts/demo_anyflow_14b.py + ' +""" + +from __future__ import annotations + +import gc +import os +import re +import sys +import time +from pathlib import Path + +import torch +from safetensors.torch import load_file + + +ANYFLOW_LOCAL = Path("/home/guian/projects/anyflow/anyflow-14b") +OUT_DIR = Path("/home/guian/projects/anyflow/demo_videos") +OUT_DIR.mkdir(exist_ok=True) + +SEED = 0 +DEVICE = torch.device("cuda") +DTYPE = torch.bfloat16 +NUM_FRAMES = 81 +HEIGHT, WIDTH = 480, 832 +PROMPT = ( + "CG game concept digital art, a majestic elephant with a vibrant tusk and sleek fur " + "running swiftly towards a herd of its kind. The elephant has a calm yet determined " + "expression, with its ears flapping slightly as it moves at high speed. The herd consists " + "of several other elephants of various ages and sizes, all moving in unison. The landscape " + "is vast savanna with rolling hills, tall grasses, and scattered acacia trees. The sun " + "sets behind the horizon, casting a warm golden glow over the scene. Low-angle view, focus " + "on the elephant as it accelerates towards the herd." +) +NEG_PROMPT = "blurry, low quality, distorted" +# The published nvidia/AnyFlow-* checkpoints are on-policy distilled with +# fuse_guidance_scale=3.0 baked into the weights — inference uses 1.0 +# (single conditional forward, no CFG; matches AnyFlow's official demo.py). +GUIDANCE = 1.0 + + +def banner(msg: str) -> None: + print("\n" + "=" * 80) + print(msg) + print("=" * 80) + + +def free_gpu() -> None: + gc.collect() + torch.cuda.empty_cache() + torch.cuda.synchronize() + used = torch.cuda.memory_allocated() / 1e9 + print(f" [mem] allocated {used:.1f} GB after free") + + +def init_single_rank() -> None: + os.environ.setdefault("MASTER_ADDR", "127.0.0.1") + os.environ.setdefault("MASTER_PORT", "29571") + os.environ.setdefault("WORLD_SIZE", "1") + os.environ.setdefault("RANK", "0") + os.environ.setdefault("LOCAL_RANK", "0") + from fastvideo.distributed import ( + init_distributed_environment, + initialize_model_parallel, + ) + init_distributed_environment(world_size=1, rank=0, local_rank=0, backend="nccl") + initialize_model_parallel( + tensor_model_parallel_size=1, + sequence_model_parallel_size=1, + data_parallel_size=1, + ) + + +def translate_keys(raw: dict, *, mapping: dict[str, str]) -> dict: + out: dict = {} + for k, v in raw.items(): + new_k = k + for pat, repl in mapping.items(): + new_k = re.sub(pat, repl, new_k) + out[new_k] = v + return out + + +def encode_prompts(): + banner("(1) Encode prompts via UMT5") + from transformers import AutoTokenizer, UMT5EncoderModel + + tok = AutoTokenizer.from_pretrained(str(ANYFLOW_LOCAL), subfolder="tokenizer", use_fast=False) + enc = UMT5EncoderModel.from_pretrained( + str(ANYFLOW_LOCAL), subfolder="text_encoder", torch_dtype=DTYPE, + ).to(DEVICE).eval() + + @torch.no_grad() + def encode_one(prompts): + out = tok( + prompts, padding="max_length", max_length=512, truncation=True, + return_attention_mask=True, return_tensors="pt") + ids = out.input_ids.to(DEVICE) + mask = out.attention_mask.to(DEVICE) + seq_lens = mask.gt(0).sum(dim=1).long() + embeds = enc(ids, mask).last_hidden_state + padded = [] + for i, l in enumerate(seq_lens): + e = embeds[i, :l] + pad = torch.zeros(512 - e.size(0), e.size(1), device=DEVICE, dtype=embeds.dtype) + padded.append(torch.cat([e, pad], dim=0)) + return torch.stack(padded, dim=0) + + text_e = encode_one([PROMPT]) + neg_e = encode_one([NEG_PROMPT]) + print(f" prompts encoded: text={tuple(text_e.shape)} neg={tuple(neg_e.shape)}") + del enc, tok + free_gpu() + return text_e, neg_e + + +def load_transformer(): + banner("(2) Build FastVideo Wan-14B + load AnyFlow weights") + from fastvideo.configs.models.dits import WanVideoConfig + from fastvideo.models.dits.wanvideo import WanTransformer3DModel + + cfg = WanVideoConfig() + arch = cfg.arch_config + # Wan2.1-T2V-14B arch (per AnyFlow checkpoint config.json). + arch.num_attention_heads = 40 + arch.attention_head_dim = 128 + arch.num_layers = 40 + arch.ffn_dim = 13824 + arch.r_embedder = True + arch.r_embedder_fusion = "gated" + arch.r_embedder_gate_value = 0.25 + arch.r_embedder_deltatime_type = "r" + arch.__post_init__() + + t0 = time.time() + model = WanTransformer3DModel(config=cfg, hf_config={}).to(DEVICE, dtype=DTYPE).eval() + print(f" transformer built in {time.time() - t0:.1f}s; " + f"params: {sum(p.numel() for p in model.parameters())/1e9:.2f}B") + + ckpt_dir = ANYFLOW_LOCAL / "transformer" + sd: dict = {} + for shard in sorted(ckpt_dir.glob("diffusion_pytorch_model-*.safetensors")): + sd.update(load_file(str(shard), device="cpu")) + print(f" AnyFlow state dict: {len(sd)} tensors loaded") + sd = translate_keys(sd, mapping=arch.param_names_mapping) + info = model.load_state_dict(sd, strict=False) + print(f" load: missing={len(info.missing_keys)} unexpected={len(info.unexpected_keys)}") + del sd + free_gpu() + return model + + +@torch.no_grad() +def sample(model, text_e, neg_e, nfe: int) -> torch.Tensor: + banner(f"(3) Sample 14B NFE={nfe}") + from fastvideo.forward_context import set_forward_context + from fastvideo.models.schedulers.scheduling_flow_map_euler_discrete import ( + FlowMapEulerDiscreteScheduler, ) + + scheduler = FlowMapEulerDiscreteScheduler(num_train_timesteps=1000, shift=5.0) + scheduler.set_timesteps(num_inference_steps=nfe, device=DEVICE) + timesteps = scheduler.timesteps.to(DEVICE, dtype=DTYPE) + + B, C = 1, 16 + F = (NUM_FRAMES - 1) // 4 + 1 # temporal VAE compression = 4 (81 → 21) + H_l, W_l = HEIGHT // 8, WIDTH // 8 + g = torch.Generator(device=DEVICE).manual_seed(SEED) + x = torch.randn(B, C, F, H_l, W_l, device=DEVICE, dtype=DTYPE, generator=g) + + t0 = time.time() + for i, (t_cur, t_next) in enumerate(zip(timesteps[:-1], timesteps[1:])): + t_in = t_cur.expand(B).to(DTYPE) + r_in = t_next.expand(B).to(DTYPE) + with set_forward_context(current_timestep=t_in, attn_metadata=None): + flow_cond = model( + hidden_states=x, encoder_hidden_states=text_e, + timestep=t_in, r_timestep=r_in) + if GUIDANCE != 1.0: + flow_uncond = model( + hidden_states=x, encoder_hidden_states=neg_e, + timestep=t_in, r_timestep=r_in) + flow = flow_uncond + GUIDANCE * (flow_cond - flow_uncond) + else: + flow = flow_cond + x = scheduler.step( + flow, sample=x, + timestep=t_cur.repeat(B), r_timestep=t_next.repeat(B)) + print(f" NFE={nfe} sample time: {time.time() - t0:.1f}s " + f"({(time.time() - t0) / nfe:.1f}s/step)") + xf = x.float() + print(f" latents mean={xf.mean().item():+.3f} std={xf.std().item():.3f} " + f"range=[{xf.min().item():+.2f}, {xf.max().item():+.2f}] " + f"finite={torch.isfinite(xf).all().item()}") + return x.detach() + + +@torch.no_grad() +def decode_to_mp4(latents: torch.Tensor, out_path: Path) -> tuple[Path, tuple]: + from diffusers.models.autoencoders.autoencoder_kl_wan import AutoencoderKLWan + import imageio.v3 as iio + + vae = AutoencoderKLWan.from_pretrained( + str(ANYFLOW_LOCAL), subfolder="vae", torch_dtype=DTYPE, + ).to(DEVICE).eval() + try: + vae.enable_tiling() + print(f" VAE tiling enabled") + except Exception: + pass + + mean = torch.tensor(vae.config.latents_mean, device=DEVICE, dtype=DTYPE).view(1, -1, 1, 1, 1) + std = torch.tensor(vae.config.latents_std, device=DEVICE, dtype=DTYPE).view(1, -1, 1, 1, 1) + latents_unscaled = latents * std + mean + t0 = time.time() + frames = vae.decode(latents_unscaled, return_dict=False)[0] + print(f" VAE decode time: {time.time() - t0:.1f}s") + frames = (frames.clamp(-1, 1) + 1) / 2 + frames = frames[0].permute(1, 2, 3, 0).float().cpu().numpy() + frames = (frames * 255).astype("uint8") + iio.imwrite(str(out_path), frames, fps=16, codec="libx264", quality=8) + del vae + free_gpu() + return out_path, frames.shape + + +def main() -> None: + torch.manual_seed(SEED) + torch.cuda.manual_seed_all(SEED) + init_single_rank() + + text_e, neg_e = encode_prompts() + model = load_transformer() + + latents_list = [] + for nfe in [4, 50]: + lat = sample(model, text_e, neg_e, nfe=nfe) + latents_list.append((nfe, lat)) + + del model, text_e, neg_e + free_gpu() + + for nfe, lat in latents_list: + banner(f"(4) Decode 14B NFE={nfe}") + out_path = OUT_DIR / f"fastvideo_anyflow_14b_nfe{nfe}_seed{SEED}.mp4" + path, shape = decode_to_mp4(lat, out_path) + print(f" decoded {shape}") + print(f" saved: {path} ({path.stat().st_size / 1e6:.2f} MB)") + + banner("DONE") + + +if __name__ == "__main__": + main()