397 lines
16 KiB
Python
397 lines
16 KiB
Python
#!/usr/bin/env python3
|
|
# SPDX-License-Identifier: Apache-2.0
|
|
"""FastVideo↔AnyFlow numerical parity verification.
|
|
|
|
Two checks, both on a single H200:
|
|
|
|
(A) Forward parity — load nvidia/AnyFlow-Wan2.1-T2V-1.3B-Diffusers into
|
|
FastVideo's WanTransformer3DModel (with r_embedder enabled +
|
|
param_names_mapping handling the delta_embedder rename), forward
|
|
it on identical inputs against AnyFlow's reference loader, and
|
|
compare. Expectation: rel-mean diff < 10% in bf16 (bf16 kernel noise).
|
|
|
|
(B) Any-step end-to-end sampling — run the new
|
|
FlowMapEulerDiscreteScheduler through 4 Euler-flow steps on the
|
|
same loaded weights, confirm the final latent is finite and
|
|
well-scaled.
|
|
|
|
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-1.3B-Diffusers``
|
|
(default ``./anyflow-1.3b``)
|
|
ANYFLOW_REF — path to the NVlabs/AnyFlow reference repo, used to
|
|
import its loader (default ``./anyflow-ref``)
|
|
|
|
Run via::
|
|
|
|
PYTHONPATH=$PWD python scripts/verify_anyflow_fastvideo_parity.py
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
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-1.3b")).expanduser()
|
|
ANYFLOW_REF = Path(os.environ.get("ANYFLOW_REF", "./anyflow-ref")).expanduser()
|
|
sys.path.insert(0, str(ANYFLOW_REF))
|
|
|
|
SEED = 1234
|
|
DEVICE = torch.device("cuda")
|
|
DTYPE = torch.bfloat16
|
|
|
|
|
|
def banner(msg: str) -> None:
|
|
print("\n" + "=" * 80)
|
|
print(msg)
|
|
print("=" * 80)
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Distributed bootstrap (single-rank). Required by WanTransformer3DModel
|
|
# which calls get_sp_world_size() and uses ReplicatedLinear.
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
def init_single_rank() -> None:
|
|
os.environ.setdefault("MASTER_ADDR", "127.0.0.1")
|
|
os.environ.setdefault("MASTER_PORT", "29551")
|
|
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,
|
|
)
|
|
print(" single-rank distributed environment + TP/SP/DP groups initialized")
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Translate AnyFlow HF safetensor keys onto FastVideo's WanTransformer3DModel
|
|
# internal layout, applying the regex from WanVideoArchConfig.param_names_mapping.
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
def translate_keys(
|
|
raw: dict[str, torch.Tensor],
|
|
*,
|
|
mapping: dict[str, str],
|
|
) -> dict[str, torch.Tensor]:
|
|
out: dict[str, torch.Tensor] = {}
|
|
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
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Build FastVideo WanTransformer3DModel with AnyFlow weights loaded.
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
def build_fastvideo_transformer():
|
|
banner("(1) Build FastVideo WanTransformer3DModel + load AnyFlow weights")
|
|
from fastvideo.configs.models.dits import WanVideoConfig
|
|
from fastvideo.models.dits.wanvideo import WanTransformer3DModel
|
|
|
|
cfg = WanVideoConfig()
|
|
arch = cfg.arch_config
|
|
# AnyFlow Wan2.1-T2V-1.3B arch dims.
|
|
arch.num_attention_heads = 12
|
|
arch.attention_head_dim = 128
|
|
arch.num_layers = 30
|
|
arch.ffn_dim = 8960
|
|
# AnyFlow dual-timestep.
|
|
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__()
|
|
|
|
# WanTransformer3DModel takes (config, hf_config). hf_config is
|
|
# diffusers-style — we provide a minimal dict; only fields the model
|
|
# actually reads matter.
|
|
hf_config: dict = {}
|
|
t0 = time.time()
|
|
model = WanTransformer3DModel(config=cfg, hf_config=hf_config)
|
|
model = model.to(DEVICE, dtype=DTYPE).eval()
|
|
print(f" Wan transformer built in {time.time() - t0:.1f}s; "
|
|
f"params: {sum(p.numel() for p in model.parameters())/1e9:.2f}B")
|
|
|
|
# Load AnyFlow checkpoint.
|
|
af_path = ANYFLOW_LOCAL / "transformer" / "diffusion_pytorch_model.safetensors"
|
|
af_raw = load_file(str(af_path), device="cpu")
|
|
print(f" AnyFlow checkpoint: {len(af_raw)} tensors")
|
|
|
|
translated = translate_keys(af_raw, mapping=arch.param_names_mapping)
|
|
info = model.load_state_dict(translated, strict=False)
|
|
miss, unex = info.missing_keys, info.unexpected_keys
|
|
print(f" missing_keys : {len(miss)} (first 5: {miss[:5]})")
|
|
print(f" unexpected_keys : {len(unex)} (first 5: {unex[:5]})")
|
|
return model, len(miss), len(unex)
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Build AnyFlow reference net.
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
def build_anyflow_reference():
|
|
banner("(2) Build AnyFlow reference loader")
|
|
from far.models import build_model
|
|
|
|
af_net = build_model("FAR_Wan_Transformer3DModel").from_pretrained(
|
|
str(ANYFLOW_LOCAL),
|
|
subfolder="transformer",
|
|
chunk_partition=None,
|
|
full_chunk_limit=0,
|
|
compressed_patch_size=[1, 4, 4],
|
|
).to(DEVICE, dtype=DTYPE).eval()
|
|
print(f" AnyFlow {type(af_net).__name__} ready")
|
|
return af_net
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Forward parity test.
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
def forward_compare(fv_model, af_net):
|
|
banner("(3) Forward output comparison")
|
|
B, C, F, H, W = 1, 16, 21, 60, 104
|
|
SEQ, DIM = 32, 4096
|
|
g = torch.Generator(device=DEVICE).manual_seed(SEED)
|
|
x = torch.randn(B, C, F, H, W, device=DEVICE, dtype=DTYPE, generator=g)
|
|
enc = torch.randn(B, SEQ, DIM, device=DEVICE, dtype=DTYPE, generator=g)
|
|
|
|
# AnyFlow expects per-frame (t, r) [B, F]; FastVideo T2V expects a
|
|
# single scalar per sample [B] (timestep.dim()==2 is reserved for
|
|
# Wan2.2 ti2v's per-token timestep schedule). For shared-t T2V the two
|
|
# are semantically equivalent.
|
|
t_per_frame = torch.full((B, F), 500.0, device=DEVICE, dtype=DTYPE)
|
|
r_per_frame = torch.full((B, F), 200.0, device=DEVICE, dtype=DTYPE)
|
|
t_per_sample = torch.full((B, ), 500.0, device=DEVICE, dtype=DTYPE)
|
|
r_per_sample = torch.full((B, ), 200.0, device=DEVICE, dtype=DTYPE)
|
|
|
|
with torch.no_grad():
|
|
# AnyFlow native: takes [B, F, C, H, W] + [B, F] (t, r).
|
|
x_af = x.permute(0, 2, 1, 3, 4).contiguous()
|
|
af_out = af_net(
|
|
x_af,
|
|
timestep=t_per_frame,
|
|
r_timestep=r_per_frame,
|
|
encoder_hidden_states=enc,
|
|
return_dict=False,
|
|
is_causal=False,
|
|
)[0]
|
|
if af_out.shape[1] != C and af_out.shape[2] == C:
|
|
af_out = af_out.permute(0, 2, 1, 3, 4).contiguous()
|
|
|
|
# FastVideo: [B, C, F, H, W] + [B] (t, r). Must wrap in
|
|
# set_forward_context so the attention layer can pick up the
|
|
# current timestep / attn_metadata.
|
|
from fastvideo.forward_context import set_forward_context
|
|
with set_forward_context(
|
|
current_timestep=t_per_sample,
|
|
attn_metadata=None,
|
|
):
|
|
fv_out = fv_model(
|
|
hidden_states=x,
|
|
encoder_hidden_states=enc,
|
|
timestep=t_per_sample,
|
|
r_timestep=r_per_sample,
|
|
)
|
|
|
|
print(f" AnyFlow out: shape={tuple(af_out.shape)} dtype={af_out.dtype}")
|
|
print(f" FastVideo : shape={tuple(fv_out.shape)} dtype={fv_out.dtype}")
|
|
if af_out.shape != fv_out.shape:
|
|
print(" ❌ shape mismatch")
|
|
return False
|
|
diff = (af_out.float() - fv_out.float()).abs()
|
|
ref = af_out.float().abs().mean().item() + 1e-12
|
|
print(f" max abs diff : {diff.max().item():.3e}")
|
|
print(f" mean abs diff: {diff.mean().item():.3e}")
|
|
print(f" rel mean diff: {diff.mean().item() / ref:.3e}")
|
|
ok = diff.mean().item() / ref < 0.10
|
|
print(f" >>> {'PASS' if ok else 'FAIL'} (target rel diff < 10%, bf16 noise)")
|
|
return ok
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Any-step sampling smoke via FlowMapEulerDiscreteScheduler.
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
def sample_anystep(fv_model):
|
|
banner("(4) Any-step 4-step Euler-flow sampling")
|
|
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=4, device=DEVICE)
|
|
timesteps = scheduler.timesteps.to(dtype=DTYPE)
|
|
|
|
B, C, F, H, W = 1, 16, 21, 60, 104
|
|
SEQ, DIM = 32, 4096
|
|
g = torch.Generator(device=DEVICE).manual_seed(SEED)
|
|
x = torch.randn(B, C, F, H, W, device=DEVICE, dtype=DTYPE, generator=g)
|
|
enc = torch.randn(B, SEQ, DIM, device=DEVICE, dtype=DTYPE, generator=g)
|
|
|
|
from fastvideo.forward_context import set_forward_context
|
|
t0 = time.time()
|
|
with torch.no_grad():
|
|
for t_cur, t_next in 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):
|
|
v = fv_model(
|
|
hidden_states=x,
|
|
encoder_hidden_states=enc,
|
|
timestep=t_in,
|
|
r_timestep=r_in,
|
|
)
|
|
x = scheduler.step(
|
|
v,
|
|
sample=x,
|
|
timestep=t_cur.repeat(B),
|
|
r_timestep=t_next.repeat(B),
|
|
)
|
|
elapsed = time.time() - t0
|
|
xf = x.float()
|
|
print(f" elapsed: {elapsed:.1f}s")
|
|
print(f" final latent: 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()}")
|
|
ok = torch.isfinite(xf).all().item() and 0.01 < xf.std().item() < 30
|
|
print(f" >>> {'PASS' if ok else 'FAIL'}")
|
|
return ok
|
|
|
|
|
|
NUM_TRAIN_TIMESTEPS = 1000
|
|
EPSILON = 5.0 # AnyFlow paper default
|
|
|
|
|
|
@torch.no_grad()
|
|
def training_step_compare(fv_model, af_net) -> bool:
|
|
"""Inline replica of AnyFlow's train_bidirection central-difference loss
|
|
on both code paths with identical synthetic (real, noise, t, r) inputs.
|
|
|
|
Compares scalar loss + intermediate flow_pred / target tensors.
|
|
"""
|
|
banner("(5) Training-step loss comparison (central-difference target)")
|
|
from fastvideo.forward_context import set_forward_context
|
|
|
|
B, C, F, H, W = 1, 16, 21, 60, 104
|
|
SEQ, DIM = 32, 4096
|
|
g = torch.Generator(device=DEVICE).manual_seed(SEED)
|
|
real = torch.randn(B, C, F, H, W, device=DEVICE, dtype=DTYPE, generator=g)
|
|
enc = torch.randn(B, SEQ, DIM, device=DEVICE, dtype=DTYPE, generator=g)
|
|
noise = torch.randn_like(real)
|
|
t_abs_pf = torch.full((B, F), 500.0, device=DEVICE, dtype=DTYPE)
|
|
r_abs_pf = torch.full((B, F), 200.0, device=DEVICE, dtype=DTYPE)
|
|
t_abs_ps = torch.full((B, ), 500.0, device=DEVICE, dtype=DTYPE)
|
|
r_abs_ps = torch.full((B, ), 200.0, device=DEVICE, dtype=DTYPE)
|
|
|
|
# AnyFlow inline replica: uses [B, F, C, H, W] layout.
|
|
real_btchw = real.permute(0, 2, 1, 3, 4).contiguous()
|
|
noise_btchw = noise.permute(0, 2, 1, 3, 4).contiguous()
|
|
t_norm_pf = (t_abs_pf / NUM_TRAIN_TIMESTEPS).view(B, F, 1, 1, 1).to(DTYPE)
|
|
noisy_btchw = t_norm_pf * noise_btchw + (1 - t_norm_pf) * real_btchw
|
|
|
|
def u_func_af(x_in, t_in, r_in):
|
|
return af_net(x_in,
|
|
timestep=t_in,
|
|
r_timestep=r_in,
|
|
encoder_hidden_states=enc,
|
|
return_dict=False,
|
|
is_causal=False)[0]
|
|
|
|
v_pred = noise_btchw - real_btchw
|
|
eps = EPSILON
|
|
F_plus = u_func_af(noisy_btchw + v_pred * (eps / NUM_TRAIN_TIMESTEPS), t_abs_pf + eps, r_abs_pf)
|
|
F_minus = u_func_af(noisy_btchw - v_pred * (eps / NUM_TRAIN_TIMESTEPS), t_abs_pf - eps, r_abs_pf)
|
|
dF_dt_af = (F_plus - F_minus) / (2 * eps)
|
|
target_af = ((noise_btchw - real_btchw) - (t_abs_pf - r_abs_pf).view(B, F, 1, 1, 1) * dF_dt_af)
|
|
flow_af = u_func_af(noisy_btchw, t_abs_pf, r_abs_pf)
|
|
loss_af = (flow_af.float() - target_af.float()).pow(2).reshape(B, -1).mean(-1)
|
|
|
|
# FastVideo inline replica: uses [B, C, F, H, W] layout + [B] t/r.
|
|
real_bcfhw = real
|
|
noise_bcfhw = noise
|
|
t_norm_ps = (t_abs_ps / NUM_TRAIN_TIMESTEPS).view(B, 1, 1, 1, 1).to(DTYPE)
|
|
noisy_bcfhw = t_norm_ps * noise_bcfhw + (1 - t_norm_ps) * real_bcfhw
|
|
|
|
def u_func_fv(x_in, t_in, r_in):
|
|
with set_forward_context(current_timestep=t_in, attn_metadata=None):
|
|
return fv_model(
|
|
hidden_states=x_in,
|
|
encoder_hidden_states=enc,
|
|
timestep=t_in,
|
|
r_timestep=r_in,
|
|
)
|
|
|
|
v_pred_fv = noise_bcfhw - real_bcfhw
|
|
F_plus_fv = u_func_fv(noisy_bcfhw + v_pred_fv * (eps / NUM_TRAIN_TIMESTEPS), t_abs_ps + eps, r_abs_ps)
|
|
F_minus_fv = u_func_fv(noisy_bcfhw - v_pred_fv * (eps / NUM_TRAIN_TIMESTEPS), t_abs_ps - eps, r_abs_ps)
|
|
dF_dt_fv = (F_plus_fv - F_minus_fv) / (2 * eps)
|
|
target_fv = ((noise_bcfhw - real_bcfhw) - (t_abs_ps - r_abs_ps).view(B, 1, 1, 1, 1) * dF_dt_fv)
|
|
flow_fv = u_func_fv(noisy_bcfhw, t_abs_ps, r_abs_ps)
|
|
loss_fv = (flow_fv.float() - target_fv.float()).pow(2).reshape(B, -1).mean(-1)
|
|
|
|
# Compare. Align AnyFlow's [B, F, C, H, W] → [B, C, F, H, W].
|
|
flow_af_aligned = flow_af.permute(0, 2, 1, 3, 4)
|
|
target_af_aligned = target_af.permute(0, 2, 1, 3, 4)
|
|
flow_diff = (flow_fv.float() - flow_af_aligned.float()).abs()
|
|
target_diff = (target_fv.float() - target_af_aligned.float()).abs()
|
|
|
|
af_loss_v = loss_af.mean().item()
|
|
fv_loss_v = loss_fv.mean().item()
|
|
abs_diff = abs(af_loss_v - fv_loss_v)
|
|
rel_diff = abs_diff / abs(af_loss_v + 1e-12)
|
|
print(f" AnyFlow loss : {af_loss_v:.6f}")
|
|
print(f" FastVideo loss: {fv_loss_v:.6f}")
|
|
print(f" abs diff : {abs_diff:.3e}")
|
|
print(f" rel diff : {rel_diff:.3e}")
|
|
print(f" flow_pred : max abs {flow_diff.max().item():.3e} "
|
|
f"mean {flow_diff.mean().item():.3e}")
|
|
print(f" target : max abs {target_diff.max().item():.3e} "
|
|
f"mean {target_diff.mean().item():.3e}")
|
|
ok = rel_diff < 0.20
|
|
print(f" >>> {'PASS' if ok else 'FAIL'} (target rel loss diff < 20%)")
|
|
return ok
|
|
|
|
|
|
def main() -> None:
|
|
torch.manual_seed(SEED)
|
|
torch.cuda.manual_seed_all(SEED)
|
|
|
|
init_single_rank()
|
|
fv_model, n_miss, n_unex = build_fastvideo_transformer()
|
|
af_net = build_anyflow_reference()
|
|
forward_ok = forward_compare(fv_model, af_net)
|
|
sample_ok = sample_anystep(fv_model)
|
|
train_ok = training_step_compare(fv_model, af_net)
|
|
banner(f"SUMMARY: missing_keys={n_miss} unexpected_keys={n_unex} "
|
|
f"forward_parity={forward_ok} sample_smoke={sample_ok} "
|
|
f"training_parity={train_ok}")
|
|
sys.exit(0 if (forward_ok and sample_ok and train_ok) else 1)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|