281 lines
10 KiB
Python
281 lines
10 KiB
Python
#!/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.
|
|
|
|
Configure local checkout paths via env vars (defaults assume a sibling
|
|
layout next to this repo):
|
|
|
|
ANYFLOW_LOCAL — path to ``nvidia/AnyFlow-Wan2.1-T2V-14B-Diffusers``
|
|
(default ``./anyflow-14b``)
|
|
ANYFLOW_DEMO_OUT — output directory for the rendered MP4s
|
|
(default ``./demo_videos``)
|
|
|
|
Run via::
|
|
|
|
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(os.environ.get("ANYFLOW_LOCAL", "./anyflow-14b")).expanduser()
|
|
OUT_DIR = Path(os.environ.get("ANYFLOW_DEMO_OUT", "./demo_videos")).expanduser()
|
|
OUT_DIR.mkdir(parents=True, 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()
|