alibaba--mnn
wehub-resource-sync
1041 行
42 KiB
Python
1041 行
42 KiB
Python
import argparse
|
|
import importlib.util
|
|
import json
|
|
import sys
|
|
import time
|
|
from dataclasses import dataclass
|
|
from datetime import datetime
|
|
from pathlib import Path
|
|
|
|
import numpy as np
|
|
|
|
|
|
MODULE_ORDER = ("vae", "t5", "dit", "e2e")
|
|
MODULE_ALIASES = {
|
|
"text_encoder": "t5",
|
|
"transformer": "dit",
|
|
"vae_decoder": "vae",
|
|
}
|
|
DEFAULT_PROMPT = "A calm ocean wave rolling onto a black sand beach at sunrise."
|
|
DEFAULT_NEGATIVE_PROMPT = ""
|
|
DEFAULT_E2E_WIDTH = 256
|
|
DEFAULT_E2E_HEIGHT = 256
|
|
DEFAULT_E2E_FRAMES = 9
|
|
DEFAULT_E2E_STEPS = 2
|
|
DEFAULT_CFG_SCALE = 5.0
|
|
DEFAULT_TEXT_LEN = 512
|
|
MANIFEST_VERSION = 1
|
|
|
|
|
|
def utc_now():
|
|
return datetime.utcnow().replace(microsecond=0).isoformat() + "Z"
|
|
|
|
|
|
def normalize_module_name(name):
|
|
key = name.lower()
|
|
return MODULE_ALIASES.get(key, key)
|
|
|
|
|
|
def normalize_modules(values):
|
|
modules = []
|
|
for value in values or MODULE_ORDER:
|
|
name = normalize_module_name(value)
|
|
if name == "all":
|
|
return list(MODULE_ORDER)
|
|
if name not in MODULE_ORDER:
|
|
raise ValueError(
|
|
"Unsupported module '{}'. Expected one of: all, {}.".format(
|
|
value, ", ".join(MODULE_ORDER)
|
|
)
|
|
)
|
|
if name not in modules:
|
|
modules.append(name)
|
|
return modules or list(MODULE_ORDER)
|
|
|
|
|
|
def add_common_case_args(parser, include_model_path):
|
|
if include_model_path:
|
|
parser.add_argument("--model_path", required=True, help="Local Wan checkpoint or official Wan source root.")
|
|
parser.add_argument("--output_dir", required=True, help="Directory for golden artifacts and reports.")
|
|
parser.add_argument("--prompt", default=DEFAULT_PROMPT, help="Positive prompt for tokenizer/T5 and e2e.")
|
|
parser.add_argument("--negative_prompt", default=DEFAULT_NEGATIVE_PROMPT, help="Negative prompt for CFG.")
|
|
parser.add_argument("--seed", type=int, default=42, help="Seed for deterministic latent and tensor generation.")
|
|
parser.add_argument("--width", type=int, default=DEFAULT_E2E_WIDTH, help="Video width.")
|
|
parser.add_argument("--height", type=int, default=DEFAULT_E2E_HEIGHT, help="Video height.")
|
|
parser.add_argument("--frames", type=int, default=DEFAULT_E2E_FRAMES, help="Video frame count.")
|
|
parser.add_argument("--steps", type=int, default=DEFAULT_E2E_STEPS, help="Low-spec e2e denoise steps.")
|
|
parser.add_argument("--cfg_scale", type=float, default=DEFAULT_CFG_SCALE, help="CFG guidance scale.")
|
|
parser.add_argument("--text_len", type=int, default=DEFAULT_TEXT_LEN, help="Tokenizer/T5 max token length.")
|
|
parser.add_argument(
|
|
"--dtype",
|
|
choices=("fp32", "float32", "fp16", "float16"),
|
|
default="fp32",
|
|
help="Reference dtype for torch capture or backend execution metadata.",
|
|
)
|
|
parser.add_argument(
|
|
"--device",
|
|
default=None,
|
|
help="Torch device for reference capture/e2e. Defaults to cuda for fp16 when available, else cpu.",
|
|
)
|
|
parser.add_argument(
|
|
"--module",
|
|
action="append",
|
|
default=[],
|
|
help="Subset to run: vae, t5, dit, e2e, or all. Repeatable.",
|
|
)
|
|
|
|
|
|
def build_parser():
|
|
parser = argparse.ArgumentParser(
|
|
description="Wan numerical alignment helper for reference, ONNX, and MNN module/e2e comparisons."
|
|
)
|
|
subparsers = parser.add_subparsers(dest="command", required=True)
|
|
|
|
capture = subparsers.add_parser(
|
|
"capture",
|
|
help="Run the local Wan reference path and save golden inputs/outputs plus manifest.json.",
|
|
)
|
|
add_common_case_args(capture, include_model_path=True)
|
|
|
|
compare_onnx = subparsers.add_parser(
|
|
"compare-onnx",
|
|
help="Run exported ONNX models with onnxruntime and compare against an existing manifest.",
|
|
)
|
|
compare_onnx.add_argument("--artifact_dir", required=True, help="Directory containing manifest.json and golden npy files.")
|
|
compare_onnx.add_argument("--onnx_root", required=True, help="Root produced by wan_onnx_export.py.")
|
|
compare_onnx.add_argument("--module", action="append", default=[], help="Subset: vae, t5, dit, e2e, or all.")
|
|
compare_onnx.add_argument("--atol", type=float, default=None, help="Override absolute tolerance.")
|
|
compare_onnx.add_argument("--rtol", type=float, default=None, help="Override relative tolerance.")
|
|
|
|
compare_mnn = subparsers.add_parser(
|
|
"compare-mnn",
|
|
help="Run converted MNN modules with PyMNN and compare against an existing manifest.",
|
|
)
|
|
compare_mnn.add_argument("--artifact_dir", required=True, help="Directory containing manifest.json and golden npy files.")
|
|
compare_mnn.add_argument("--mnn_root", required=True, help="Root produced by wan_convert_mnn.py.")
|
|
compare_mnn.add_argument("--module", action="append", default=[], help="Subset: vae, t5, dit, e2e, or all.")
|
|
compare_mnn.add_argument("--atol", type=float, default=None, help="Override absolute tolerance.")
|
|
compare_mnn.add_argument("--rtol", type=float, default=None, help="Override relative tolerance.")
|
|
compare_mnn.add_argument("--thread_num", type=int, default=1, help="PyMNN runtime thread count.")
|
|
compare_mnn.add_argument(
|
|
"--mnn_backend",
|
|
choices=("CPU", "CUDA"),
|
|
default="CPU",
|
|
help="MNN runtime backend (default CPU).",
|
|
)
|
|
|
|
e2e = subparsers.add_parser(
|
|
"e2e",
|
|
help="Run the fixed small e2e case on reference, ONNX, or MNN, dump intermediates, and optionally compare.",
|
|
)
|
|
e2e.add_argument("--artifact_dir", required=True, help="Directory containing manifest.json and output reports.")
|
|
e2e.add_argument(
|
|
"--backend",
|
|
required=True,
|
|
choices=("reference", "onnx", "mnn"),
|
|
help="Execution backend for the low-spec e2e rerun.",
|
|
)
|
|
e2e.add_argument("--model_path", help="Local Wan checkpoint root for backend=reference.")
|
|
e2e.add_argument("--onnx_root", help="ONNX root for backend=onnx.")
|
|
e2e.add_argument("--mnn_root", help="MNN root for backend=mnn.")
|
|
e2e.add_argument("--compare", action="store_true", help="Compare the backend e2e dump against the reference manifest.")
|
|
e2e.add_argument("--atol", type=float, default=None, help="Override absolute tolerance.")
|
|
e2e.add_argument("--rtol", type=float, default=None, help="Override relative tolerance.")
|
|
e2e.add_argument("--thread_num", type=int, default=1, help="PyMNN runtime thread count for backend=mnn.")
|
|
e2e.add_argument(
|
|
"--mnn_backend",
|
|
choices=("CPU", "CUDA"),
|
|
default="CPU",
|
|
help="MNN runtime backend for backend=mnn (default CPU).",
|
|
)
|
|
e2e.add_argument(
|
|
"--dtype",
|
|
choices=("fp32", "float32", "fp16", "float16"),
|
|
default="fp32",
|
|
help="Reference dtype for backend=reference.",
|
|
)
|
|
e2e.add_argument("--device", default=None, help="Torch device for backend=reference.")
|
|
return parser
|
|
|
|
|
|
def require_file(path, description):
|
|
path = Path(path)
|
|
if not path.exists():
|
|
raise FileNotFoundError("{} not found: {}".format(description, path))
|
|
return path
|
|
|
|
|
|
def load_json(path):
|
|
with Path(path).open("r", encoding="utf-8") as fp:
|
|
return json.load(fp)
|
|
|
|
|
|
def save_json(path, data):
|
|
path = Path(path)
|
|
path.parent.mkdir(parents=True, exist_ok=True)
|
|
with path.open("w", encoding="utf-8") as fp:
|
|
json.dump(data, fp, indent=2, sort_keys=False)
|
|
fp.write("\n")
|
|
|
|
|
|
def write_npy(root, relative_path, array):
|
|
root = Path(root)
|
|
path = root / relative_path
|
|
path.parent.mkdir(parents=True, exist_ok=True)
|
|
np.save(path.as_posix(), np.asarray(array))
|
|
array = np.asarray(array)
|
|
return {
|
|
"path": relative_path,
|
|
"shape": list(array.shape),
|
|
"dtype": str(array.dtype),
|
|
}
|
|
|
|
|
|
def load_npy(root, entry):
|
|
return np.load((Path(root) / entry["path"]).as_posix(), allow_pickle=False)
|
|
|
|
|
|
def cosine_similarity(lhs, rhs):
|
|
lhs = np.asarray(lhs, dtype=np.float64).reshape(-1)
|
|
rhs = np.asarray(rhs, dtype=np.float64).reshape(-1)
|
|
lhs_norm = np.linalg.norm(lhs)
|
|
rhs_norm = np.linalg.norm(rhs)
|
|
if lhs_norm == 0.0 and rhs_norm == 0.0:
|
|
return 1.0
|
|
if lhs_norm == 0.0 or rhs_norm == 0.0:
|
|
return 0.0
|
|
return float(np.dot(lhs, rhs) / (lhs_norm * rhs_norm))
|
|
|
|
|
|
def compare_arrays(reference, candidate, atol, rtol):
|
|
reference = np.asarray(reference)
|
|
candidate = np.asarray(candidate)
|
|
if reference.shape != candidate.shape:
|
|
raise ValueError("Shape mismatch: reference={} candidate={}".format(reference.shape, candidate.shape))
|
|
diff = np.abs(reference.astype(np.float64) - candidate.astype(np.float64))
|
|
return {
|
|
"shape": list(reference.shape),
|
|
"dtype_reference": str(reference.dtype),
|
|
"dtype_candidate": str(candidate.dtype),
|
|
"max_abs": float(diff.max()) if diff.size else 0.0,
|
|
"mean_abs": float(diff.mean()) if diff.size else 0.0,
|
|
"cosine": cosine_similarity(reference, candidate),
|
|
"allclose": bool(np.allclose(reference, candidate, atol=atol, rtol=rtol)),
|
|
"atol": float(atol),
|
|
"rtol": float(rtol),
|
|
}
|
|
|
|
|
|
def module_tolerance(module_name, atol_override=None, rtol_override=None):
|
|
base = {
|
|
"vae": (1e-3, 1e-3),
|
|
"t5": (1e-3, 1e-3),
|
|
"dit": (1e-2, 1e-3),
|
|
"e2e": (1e-2, 1e-3),
|
|
}
|
|
atol, rtol = base[module_name]
|
|
if atol_override is not None:
|
|
atol = atol_override
|
|
if rtol_override is not None:
|
|
rtol = rtol_override
|
|
return atol, rtol
|
|
|
|
|
|
def latent_frame_count(frames):
|
|
return max(1, (frames - 1) // 4 + 1)
|
|
|
|
|
|
def resolve_torch_dtype(dtype_name, torch):
|
|
if dtype_name in ("fp16", "float16"):
|
|
return torch.float16
|
|
return torch.float32
|
|
|
|
|
|
def pick_torch_device(dtype_name, explicit_device, torch):
|
|
if explicit_device:
|
|
return explicit_device
|
|
if dtype_name in ("fp16", "float16") and torch.cuda.is_available():
|
|
return "cuda"
|
|
return "cpu"
|
|
|
|
|
|
def load_wan_export_module():
|
|
module_path = Path(__file__).resolve().parent / "wan_onnx_export.py"
|
|
spec = importlib.util.spec_from_file_location("wan_onnx_export_local", module_path.as_posix())
|
|
if spec is None or spec.loader is None:
|
|
raise RuntimeError("Failed to load helper module: {}".format(module_path))
|
|
module = importlib.util.module_from_spec(spec)
|
|
spec.loader.exec_module(module)
|
|
return module
|
|
|
|
|
|
def build_prompt_batch(tokenizer, prompt, negative_prompt, text_len, torch, device):
|
|
if tokenizer is None:
|
|
raise RuntimeError("Wan loader did not expose a tokenizer; capture requires a local tokenizer.")
|
|
prompts = [negative_prompt, prompt]
|
|
encoded = tokenizer(
|
|
prompts,
|
|
padding="max_length",
|
|
truncation=True,
|
|
max_length=text_len,
|
|
return_tensors="pt",
|
|
)
|
|
input_ids = encoded["input_ids"].to(device=device)
|
|
attention_mask = encoded["attention_mask"].to(device=device)
|
|
if input_ids.shape[0] != 2:
|
|
raise RuntimeError("Expected CFG prompt batch=2, got {}".format(tuple(input_ids.shape)))
|
|
return input_ids, attention_mask
|
|
|
|
|
|
def torch_to_numpy(tensor):
|
|
return tensor.detach().cpu().numpy()
|
|
|
|
|
|
def deterministic_latent(shape, seed, dtype=np.float32):
|
|
rng = np.random.default_rng(seed)
|
|
return rng.standard_normal(shape, dtype=np.float32).astype(dtype)
|
|
|
|
|
|
def deterministic_timesteps(steps, shift=3.0):
|
|
"""Build the FlowMatch timestep schedule used by Wan2.1 inference.
|
|
|
|
This must mirror WanDiffusion::runVideo in C++ exactly so the alignment
|
|
harness reflects real on-device behaviour: linear t in [1.0, 0.001]
|
|
re-mapped through the shifted scheduler then scaled to [0, 1000].
|
|
"""
|
|
if steps <= 0:
|
|
raise ValueError("--steps must be > 0")
|
|
if steps == 1:
|
|
return np.array([1000.0], dtype=np.float32)
|
|
values = np.empty(steps, dtype=np.float32)
|
|
for i in range(steps):
|
|
t_linear = 1.0 + i * (0.001 - 1.0) / (steps - 1)
|
|
t_shifted = (shift * t_linear) / (1.0 + (shift - 1.0) * t_linear)
|
|
values[i] = t_shifted * 1000.0
|
|
return values
|
|
|
|
|
|
@dataclass
|
|
class ReferenceRunner:
|
|
torch: object
|
|
device: str
|
|
dtype: object
|
|
latent_channels: int
|
|
tokenizer: object
|
|
source: str
|
|
text_encoder: object
|
|
transformer: object
|
|
vae_decoder: object
|
|
|
|
def run_t5(self, input_ids):
|
|
tensor = self.torch.from_numpy(np.asarray(input_ids)).to(device=self.device, dtype=self.torch.int32)
|
|
# Free GPU memory: T5 (~25GB fp32) competes with DiT/VAE on a single 32GB card.
|
|
# Move to device just-in-time and offload back to CPU after use.
|
|
is_cuda = isinstance(self.device, str) and self.device.startswith("cuda")
|
|
inner = getattr(self.text_encoder, "inner", None)
|
|
module = getattr(self.text_encoder, "module", None)
|
|
if is_cuda:
|
|
if inner is not None and hasattr(inner, "to"):
|
|
inner.to(device=self.device, dtype=self.dtype)
|
|
if module is not None and module is not inner and hasattr(module, "to"):
|
|
module.to(device=self.device, dtype=self.dtype)
|
|
with self.torch.no_grad():
|
|
result = torch_to_numpy(self.text_encoder(tensor))
|
|
if is_cuda:
|
|
if inner is not None and hasattr(inner, "to"):
|
|
inner.to("cpu")
|
|
if module is not None and module is not inner and hasattr(module, "to"):
|
|
module.to("cpu")
|
|
self.torch.cuda.empty_cache()
|
|
return result
|
|
|
|
def run_dit(self, hidden_states, timestep, encoder_hidden_states, encoder_attention_mask):
|
|
hidden_states = self.torch.from_numpy(np.asarray(hidden_states)).to(device=self.device, dtype=self.dtype)
|
|
timestep = self.torch.from_numpy(np.asarray(timestep)).to(device=self.device, dtype=self.dtype)
|
|
encoder_hidden_states = self.torch.from_numpy(np.asarray(encoder_hidden_states)).to(
|
|
device=self.device, dtype=self.dtype
|
|
)
|
|
encoder_attention_mask = self.torch.from_numpy(np.asarray(encoder_attention_mask)).to(
|
|
device=self.device, dtype=self.torch.int32
|
|
)
|
|
with self.torch.no_grad():
|
|
return torch_to_numpy(
|
|
self.transformer(hidden_states, timestep, encoder_hidden_states, encoder_attention_mask)
|
|
)
|
|
|
|
def run_vae(self, latent_sample):
|
|
latent_sample = self.torch.from_numpy(np.asarray(latent_sample)).to(device=self.device, dtype=self.dtype)
|
|
with self.torch.no_grad():
|
|
return torch_to_numpy(self.vae_decoder(latent_sample))
|
|
|
|
|
|
@dataclass
|
|
class OnnxRunner:
|
|
ort: object
|
|
text_session: object = None
|
|
dit_session: object = None
|
|
vae_session: object = None
|
|
|
|
@classmethod
|
|
def from_root(cls, onnx_root, modules=None):
|
|
try:
|
|
import onnxruntime as ort
|
|
except Exception as exc:
|
|
raise RuntimeError("compare-onnx requires onnxruntime: {}".format(exc))
|
|
onnx_root = Path(onnx_root)
|
|
modules = set(modules or ("t5", "dit", "vae"))
|
|
providers = ["CPUExecutionProvider"]
|
|
return cls(
|
|
ort=ort,
|
|
text_session=ort.InferenceSession(
|
|
require_file(onnx_root / "text_encoder" / "model.onnx", "text_encoder ONNX").as_posix(),
|
|
providers=providers,
|
|
)
|
|
if "t5" in modules
|
|
else None,
|
|
dit_session=ort.InferenceSession(
|
|
require_file(onnx_root / "transformer" / "model.onnx", "transformer ONNX").as_posix(),
|
|
providers=providers,
|
|
)
|
|
if "dit" in modules
|
|
else None,
|
|
vae_session=ort.InferenceSession(
|
|
require_file(onnx_root / "vae_decoder" / "model.onnx", "vae_decoder ONNX").as_posix(),
|
|
providers=providers,
|
|
)
|
|
if "vae" in modules
|
|
else None,
|
|
)
|
|
|
|
@staticmethod
|
|
def _run_single(session, feeds):
|
|
outputs = session.run(None, feeds)
|
|
if not outputs:
|
|
raise RuntimeError("ONNX session produced no outputs for {}".format(session))
|
|
return np.asarray(outputs[0])
|
|
|
|
def run_t5(self, input_ids):
|
|
return self._run_single(self.text_session, {"input_ids": np.asarray(input_ids, dtype=np.int32)})
|
|
|
|
def run_dit(self, hidden_states, timestep, encoder_hidden_states, encoder_attention_mask):
|
|
feeds = {
|
|
"hidden_states": np.asarray(hidden_states),
|
|
"timestep": np.asarray(timestep),
|
|
"encoder_hidden_states": np.asarray(encoder_hidden_states),
|
|
"encoder_attention_mask": np.asarray(encoder_attention_mask, dtype=np.int32),
|
|
}
|
|
return self._run_single(self.dit_session, feeds)
|
|
|
|
def run_vae(self, latent_sample):
|
|
return self._run_single(self.vae_session, {"latent_sample": np.asarray(latent_sample)})
|
|
|
|
|
|
@dataclass
|
|
class MnnRunner:
|
|
MNN: object
|
|
F: object
|
|
nn: object
|
|
text_module: object = None
|
|
dit_module: object = None
|
|
vae_module: object = None
|
|
|
|
@classmethod
|
|
def from_root(cls, mnn_root, thread_num=1, modules=None, backend="CPU"):
|
|
try:
|
|
import MNN
|
|
import MNN.expr as F
|
|
import MNN.nn as nn
|
|
except Exception as exc:
|
|
raise RuntimeError("compare-mnn requires PyMNN in PYTHONPATH: {}".format(exc))
|
|
mnn_root = Path(mnn_root)
|
|
modules = set(modules or ("t5", "dit", "vae"))
|
|
runtime_manager = None
|
|
if backend != "CPU":
|
|
runtime_manager = nn.create_runtime_manager([
|
|
{
|
|
"backend": backend,
|
|
"precision": "high",
|
|
"numThread": thread_num,
|
|
}
|
|
])
|
|
kwargs = {
|
|
"dynamic": False,
|
|
"shape_mutable": False,
|
|
"thread_num": thread_num,
|
|
}
|
|
if runtime_manager is not None:
|
|
kwargs["runtime_manager"] = runtime_manager
|
|
return cls(
|
|
MNN=MNN,
|
|
F=F,
|
|
nn=nn,
|
|
text_module=nn.load_module_from_file(
|
|
require_file(mnn_root / "text_encoder.mnn", "text_encoder MNN").as_posix(),
|
|
["input_ids"],
|
|
["last_hidden_state"],
|
|
**kwargs
|
|
)
|
|
if "t5" in modules
|
|
else None,
|
|
dit_module=nn.load_module_from_file(
|
|
require_file(mnn_root / "transformer.mnn", "transformer MNN").as_posix(),
|
|
["hidden_states", "timestep", "encoder_hidden_states", "encoder_attention_mask"],
|
|
["noise_pred"],
|
|
**kwargs
|
|
)
|
|
if "dit" in modules
|
|
else None,
|
|
vae_module=nn.load_module_from_file(
|
|
require_file(mnn_root / "vae_decoder.mnn", "vae_decoder MNN").as_posix(),
|
|
["latent_sample"],
|
|
["sample"],
|
|
**kwargs
|
|
)
|
|
if "vae" in modules
|
|
else None,
|
|
)
|
|
|
|
def _dtype_to_expr(self, array):
|
|
dtype = np.asarray(array).dtype
|
|
if dtype == np.float64:
|
|
return self.F.double
|
|
if dtype == np.float32 or dtype == np.float16:
|
|
return self.F.float
|
|
if dtype == np.int64:
|
|
return self.F.int64
|
|
if dtype == np.uint8:
|
|
return self.F.uint8
|
|
if dtype == np.int32 or dtype == np.int16 or dtype == np.int8:
|
|
return self.F.int
|
|
raise TypeError("Unsupported numpy dtype for PyMNN input: {}".format(dtype))
|
|
|
|
def _make_input_var(self, array, expected_var):
|
|
array = np.ascontiguousarray(array)
|
|
var = self.F.const(array, list(array.shape), self.F.NCHW, self._dtype_to_expr(array))
|
|
if getattr(expected_var, "data_format", None) == self.F.NC4HW4:
|
|
var = self.F.convert(var, self.F.NC4HW4)
|
|
return var
|
|
|
|
def _run_single(self, module, arrays):
|
|
info = module.get_info()
|
|
input_vars = info["inputs"]
|
|
if len(input_vars) != len(arrays):
|
|
raise RuntimeError(
|
|
"Input arity mismatch for MNN module. expected {} got {}".format(len(input_vars), len(arrays))
|
|
)
|
|
vars_in = [self._make_input_var(arrays[i], input_vars[i]) for i in range(len(arrays))]
|
|
outputs = module.forward(vars_in if len(vars_in) != 1 else vars_in[0])
|
|
if not isinstance(outputs, (list, tuple)):
|
|
outputs = [outputs]
|
|
output = outputs[0]
|
|
if getattr(output, "data_format", None) == self.F.NC4HW4:
|
|
output = self.F.convert(output, self.F.NCHW)
|
|
data = np.array(output.read(), copy=True)
|
|
return data.reshape(output.shape)
|
|
|
|
def run_t5(self, input_ids):
|
|
return self._run_single(self.text_module, [np.asarray(input_ids, dtype=np.int32)])
|
|
|
|
def run_dit(self, hidden_states, timestep, encoder_hidden_states, encoder_attention_mask):
|
|
return self._run_single(
|
|
self.dit_module,
|
|
[
|
|
np.asarray(hidden_states),
|
|
np.asarray(timestep),
|
|
np.asarray(encoder_hidden_states),
|
|
np.asarray(encoder_attention_mask, dtype=np.int32),
|
|
],
|
|
)
|
|
|
|
def run_vae(self, latent_sample):
|
|
return self._run_single(self.vae_module, [np.asarray(latent_sample)])
|
|
|
|
|
|
def create_reference_runner(args):
|
|
try:
|
|
import torch
|
|
except Exception as exc:
|
|
raise RuntimeError("capture/reference e2e requires torch: {}".format(exc))
|
|
wan_export = load_wan_export_module()
|
|
torch_dtype = resolve_torch_dtype(args.dtype, torch)
|
|
device = pick_torch_device(args.dtype, getattr(args, "device", None), torch)
|
|
if torch_dtype == torch.float16 and device == "cpu":
|
|
raise ValueError("fp16 reference execution requires --device cuda or an available CUDA device")
|
|
components = wan_export.load_wan_components(Path(args.model_path).resolve().as_posix(), torch_dtype, device)
|
|
for module in (components.text_encoder, components.transformer, components.vae_decoder):
|
|
wan_export.move_to(module, device=device, dtype=torch_dtype)
|
|
latent_channels = 16
|
|
if hasattr(wan_export, "infer_latent_channels"):
|
|
latent_channels = int(wan_export.infer_latent_channels(components))
|
|
elif getattr(components.vae_decoder, "config", None) is not None and hasattr(components.vae_decoder.config, "z_dim"):
|
|
latent_channels = int(components.vae_decoder.config.z_dim)
|
|
return ReferenceRunner(
|
|
torch=torch,
|
|
device=device,
|
|
dtype=torch_dtype,
|
|
latent_channels=latent_channels,
|
|
tokenizer=components.tokenizer,
|
|
source=components.source,
|
|
text_encoder=wan_export.TextEncoderWrapper(torch, components.text_encoder),
|
|
transformer=wan_export.TransformerWrapper(torch, components.transformer),
|
|
vae_decoder=wan_export.VaeDecoderWrapper(components.vae_decoder),
|
|
)
|
|
|
|
|
|
def build_case_metadata(args, latent_channels):
|
|
latent_h = max(1, args.height // 8)
|
|
latent_w = max(1, args.width // 8)
|
|
latent_t = latent_frame_count(args.frames)
|
|
return {
|
|
"prompt": args.prompt,
|
|
"negative_prompt": args.negative_prompt,
|
|
"seed": int(args.seed),
|
|
"width": int(args.width),
|
|
"height": int(args.height),
|
|
"frames": int(args.frames),
|
|
"steps": int(args.steps),
|
|
"cfg_scale": float(args.cfg_scale),
|
|
"text_len": int(args.text_len),
|
|
"dtype": args.dtype,
|
|
"latent_channels": int(latent_channels),
|
|
"latent_shape": [1, int(latent_channels), int(latent_t), int(latent_h), int(latent_w)],
|
|
"cfg_batch_order": ["negative", "positive"],
|
|
}
|
|
|
|
|
|
def run_e2e_case(runner, metadata, input_ids, attention_mask):
|
|
text = runner.run_t5(input_ids)
|
|
latent_shape = metadata["latent_shape"]
|
|
sample = deterministic_latent(latent_shape, metadata["seed"])
|
|
timesteps = deterministic_timesteps(metadata["steps"])
|
|
steps = []
|
|
current = sample.astype(np.float32, copy=True)
|
|
for index, timestep_value in enumerate(timesteps):
|
|
sample_before = current.copy()
|
|
hidden_states = np.repeat(sample_before, 2, axis=0)
|
|
timestep = np.full((2,), timestep_value, dtype=np.float32)
|
|
noise_pred = runner.run_dit(hidden_states, timestep, text, attention_mask)
|
|
if noise_pred.shape[0] != 2:
|
|
raise RuntimeError("Expected DiT noise prediction batch=2 for CFG, got {}".format(noise_pred.shape))
|
|
noise_pred_uncond = noise_pred[0:1].astype(np.float32, copy=True)
|
|
noise_pred_cond = noise_pred[1:2].astype(np.float32, copy=True)
|
|
guided_noise = noise_pred_uncond + np.float32(metadata["cfg_scale"]) * (noise_pred_cond - noise_pred_uncond)
|
|
# Match WanDiffusion::stepFlowMatch in C++:
|
|
# sample + modelOutput * (nextT - t) / 1000.0
|
|
next_timestep_value = float(timesteps[index + 1]) if index + 1 < len(timesteps) else 0.0
|
|
delta = np.float32((next_timestep_value - float(timestep_value)) / 1000.0)
|
|
sample_after = sample_before + delta * guided_noise
|
|
steps.append(
|
|
{
|
|
"index": index,
|
|
"timestep": float(timestep_value),
|
|
"sample_before": sample_before,
|
|
"noise_pred": noise_pred.astype(np.float32, copy=True),
|
|
"noise_pred_uncond": noise_pred_uncond,
|
|
"noise_pred_cond": noise_pred_cond,
|
|
"guided_noise": guided_noise.astype(np.float32, copy=True),
|
|
"sample_after": sample_after.astype(np.float32, copy=True),
|
|
}
|
|
)
|
|
current = sample_after.astype(np.float32, copy=True)
|
|
final_video = runner.run_vae(current)
|
|
return {
|
|
"text_embeddings": text.astype(np.float32, copy=True),
|
|
"initial_latent": sample.astype(np.float32, copy=True),
|
|
"steps": steps,
|
|
"final_latent": current.astype(np.float32, copy=True),
|
|
"final_video": np.asarray(final_video),
|
|
}
|
|
|
|
|
|
def capture_command(args):
|
|
modules = normalize_modules(args.module)
|
|
output_dir = Path(args.output_dir).resolve()
|
|
runner = create_reference_runner(args)
|
|
metadata = build_case_metadata(args, latent_channels=runner.latent_channels)
|
|
input_ids_torch, attention_mask_torch = build_prompt_batch(
|
|
runner.tokenizer, args.prompt, args.negative_prompt, args.text_len, runner.torch, runner.device
|
|
)
|
|
input_ids = torch_to_numpy(input_ids_torch).astype(np.int32, copy=False)
|
|
attention_mask = torch_to_numpy(attention_mask_torch).astype(np.int32, copy=False)
|
|
|
|
output_dir.mkdir(parents=True, exist_ok=True)
|
|
manifest = {
|
|
"version": MANIFEST_VERSION,
|
|
"created_at": utc_now(),
|
|
"reference": {
|
|
"source": runner.source,
|
|
"model_path": Path(args.model_path).resolve().as_posix(),
|
|
"device": runner.device,
|
|
"dtype": args.dtype,
|
|
},
|
|
"case": metadata,
|
|
"modules": {},
|
|
"reports": {},
|
|
}
|
|
|
|
manifest["modules"]["t5"] = {
|
|
"inputs": {
|
|
"input_ids": write_npy(output_dir, "t5/input_ids.npy", input_ids),
|
|
"attention_mask": write_npy(output_dir, "t5/attention_mask.npy", attention_mask),
|
|
}
|
|
}
|
|
|
|
if "t5" in modules or "dit" in modules or "e2e" in modules:
|
|
last_hidden_state = runner.run_t5(input_ids).astype(np.float32, copy=True)
|
|
manifest["modules"]["t5"]["outputs"] = {
|
|
"last_hidden_state": write_npy(output_dir, "t5/last_hidden_state.npy", last_hidden_state)
|
|
}
|
|
|
|
if "dit" in modules:
|
|
dit_hidden_states = deterministic_latent(metadata["latent_shape"], args.seed + 101)
|
|
dit_hidden_states = np.repeat(dit_hidden_states, 2, axis=0)
|
|
dit_timestep = np.full((2,), np.float32(0.75), dtype=np.float32)
|
|
dit_output = runner.run_dit(
|
|
dit_hidden_states,
|
|
dit_timestep,
|
|
last_hidden_state,
|
|
attention_mask,
|
|
).astype(np.float32, copy=True)
|
|
manifest["modules"]["dit"] = {
|
|
"inputs": {
|
|
"hidden_states": write_npy(output_dir, "dit/hidden_states.npy", dit_hidden_states),
|
|
"timestep": write_npy(output_dir, "dit/timestep.npy", dit_timestep),
|
|
"encoder_hidden_states": write_npy(output_dir, "dit/encoder_hidden_states.npy", last_hidden_state),
|
|
"encoder_attention_mask": write_npy(output_dir, "dit/encoder_attention_mask.npy", attention_mask),
|
|
},
|
|
"outputs": {
|
|
"noise_pred": write_npy(output_dir, "dit/noise_pred.npy", dit_output),
|
|
},
|
|
}
|
|
|
|
if "vae" in modules:
|
|
vae_latent = deterministic_latent(metadata["latent_shape"], args.seed + 202)
|
|
vae_output = runner.run_vae(vae_latent)
|
|
manifest["modules"]["vae"] = {
|
|
"inputs": {
|
|
"latent_sample": write_npy(output_dir, "vae/latent_sample.npy", vae_latent),
|
|
},
|
|
"outputs": {
|
|
"sample": write_npy(output_dir, "vae/sample.npy", vae_output),
|
|
},
|
|
}
|
|
|
|
if "e2e" in modules:
|
|
e2e = run_e2e_case(runner, metadata, input_ids, attention_mask)
|
|
e2e_manifest = {
|
|
"inputs": {
|
|
"input_ids": write_npy(output_dir, "e2e/input_ids.npy", input_ids),
|
|
"attention_mask": write_npy(output_dir, "e2e/attention_mask.npy", attention_mask),
|
|
"initial_latent": write_npy(output_dir, "e2e/initial_latent.npy", e2e["initial_latent"]),
|
|
"text_embeddings": write_npy(output_dir, "e2e/text_embeddings.npy", e2e["text_embeddings"]),
|
|
},
|
|
"steps": [],
|
|
"outputs": {
|
|
"final_latent": write_npy(output_dir, "e2e/final_latent.npy", e2e["final_latent"]),
|
|
"final_video": write_npy(output_dir, "e2e/final_video.npy", e2e["final_video"]),
|
|
},
|
|
}
|
|
for step in e2e["steps"]:
|
|
step_dir = "e2e/steps/step_{:02d}".format(step["index"])
|
|
e2e_manifest["steps"].append(
|
|
{
|
|
"index": step["index"],
|
|
"timestep": step["timestep"],
|
|
"sample_before": write_npy(output_dir, step_dir + "/sample_before.npy", step["sample_before"]),
|
|
"noise_pred": write_npy(output_dir, step_dir + "/noise_pred.npy", step["noise_pred"]),
|
|
"noise_pred_uncond": write_npy(
|
|
output_dir, step_dir + "/noise_pred_uncond.npy", step["noise_pred_uncond"]
|
|
),
|
|
"noise_pred_cond": write_npy(
|
|
output_dir, step_dir + "/noise_pred_cond.npy", step["noise_pred_cond"]
|
|
),
|
|
"guided_noise": write_npy(output_dir, step_dir + "/guided_noise.npy", step["guided_noise"]),
|
|
"sample_after": write_npy(output_dir, step_dir + "/sample_after.npy", step["sample_after"]),
|
|
}
|
|
)
|
|
manifest["modules"]["e2e"] = e2e_manifest
|
|
|
|
save_json(output_dir / "manifest.json", manifest)
|
|
save_json(
|
|
output_dir / "reports" / "capture_report.json",
|
|
{
|
|
"created_at": utc_now(),
|
|
"output_dir": output_dir.as_posix(),
|
|
"modules": modules,
|
|
"reference": manifest["reference"],
|
|
"case": manifest["case"],
|
|
},
|
|
)
|
|
print("Captured Wan alignment artifacts to {}".format(output_dir))
|
|
|
|
|
|
def load_manifest(artifact_dir):
|
|
artifact_dir = Path(artifact_dir).resolve()
|
|
manifest = load_json(require_file(artifact_dir / "manifest.json", "manifest.json"))
|
|
return artifact_dir, manifest
|
|
|
|
|
|
def manifest_modules(manifest, include_e2e):
|
|
modules = [name for name in MODULE_ORDER if name in manifest.get("modules", {})]
|
|
if not include_e2e:
|
|
modules = [name for name in modules if name != "e2e"]
|
|
return modules
|
|
|
|
|
|
def compare_module_output(artifact_dir, manifest, module_name, candidate_outputs, atol_override=None, rtol_override=None):
|
|
atol, rtol = module_tolerance(module_name, atol_override, rtol_override)
|
|
module_manifest = manifest["modules"].get(module_name)
|
|
if module_manifest is None:
|
|
raise RuntimeError("Manifest does not contain module '{}'".format(module_name))
|
|
outputs = module_manifest["outputs"]
|
|
report = {
|
|
"module": module_name,
|
|
"tensors": {},
|
|
"pass": True,
|
|
}
|
|
for output_name, entry in outputs.items():
|
|
reference = load_npy(artifact_dir, entry)
|
|
candidate = candidate_outputs[output_name]
|
|
metrics = compare_arrays(reference, candidate, atol=atol, rtol=rtol)
|
|
report["tensors"][output_name] = metrics
|
|
report["pass"] = report["pass"] and metrics["allclose"]
|
|
return report
|
|
|
|
|
|
def run_backend_module(artifact_dir, manifest, runner, module_name):
|
|
module = manifest["modules"][module_name]
|
|
if module_name == "t5":
|
|
input_ids = load_npy(artifact_dir, module["inputs"]["input_ids"])
|
|
return {"last_hidden_state": runner.run_t5(input_ids)}
|
|
if module_name == "dit":
|
|
return {
|
|
"noise_pred": runner.run_dit(
|
|
load_npy(artifact_dir, module["inputs"]["hidden_states"]),
|
|
load_npy(artifact_dir, module["inputs"]["timestep"]),
|
|
load_npy(artifact_dir, module["inputs"]["encoder_hidden_states"]),
|
|
load_npy(artifact_dir, module["inputs"]["encoder_attention_mask"]),
|
|
)
|
|
}
|
|
if module_name == "vae":
|
|
return {"sample": runner.run_vae(load_npy(artifact_dir, module["inputs"]["latent_sample"]))}
|
|
raise ValueError("Unsupported module backend run: {}".format(module_name))
|
|
|
|
|
|
def build_runner_from_compare_args(command, args, modules):
|
|
if command == "compare-onnx":
|
|
return OnnxRunner.from_root(args.onnx_root, modules=modules)
|
|
if command == "compare-mnn":
|
|
return MnnRunner.from_root(
|
|
args.mnn_root,
|
|
thread_num=args.thread_num,
|
|
modules=modules,
|
|
backend=getattr(args, "mnn_backend", "CPU"),
|
|
)
|
|
raise ValueError("Unsupported compare command: {}".format(command))
|
|
|
|
|
|
def compare_command(command, args):
|
|
artifact_dir, manifest = load_manifest(args.artifact_dir)
|
|
requested_modules = normalize_modules(args.module) if args.module else manifest_modules(manifest, include_e2e=True)
|
|
modules = [m for m in requested_modules if m != "e2e"]
|
|
runner = build_runner_from_compare_args(command, args, requested_modules)
|
|
report = {
|
|
"created_at": utc_now(),
|
|
"command": command,
|
|
"artifact_dir": artifact_dir.as_posix(),
|
|
"modules": {},
|
|
"pass": True,
|
|
}
|
|
for module_name in modules:
|
|
start = time.time()
|
|
outputs = run_backend_module(artifact_dir, manifest, runner, module_name)
|
|
module_report = compare_module_output(
|
|
artifact_dir, manifest, module_name, outputs, atol_override=args.atol, rtol_override=args.rtol
|
|
)
|
|
module_report["elapsed_ms"] = round((time.time() - start) * 1000.0, 3)
|
|
report["modules"][module_name] = module_report
|
|
report["pass"] = report["pass"] and module_report["pass"]
|
|
|
|
if "e2e" in requested_modules:
|
|
e2e_report = e2e_compare(
|
|
artifact_dir,
|
|
manifest,
|
|
runner,
|
|
backend_name="onnx" if command == "compare-onnx" else "mnn",
|
|
atol_override=args.atol,
|
|
rtol_override=args.rtol,
|
|
write_outputs=True,
|
|
)
|
|
report["modules"]["e2e"] = e2e_report
|
|
report["pass"] = report["pass"] and e2e_report["pass"]
|
|
|
|
report_path = artifact_dir / "reports" / "{}.json".format(command.replace("-", "_"))
|
|
save_json(report_path, report)
|
|
print("Wrote {}".format(report_path))
|
|
if not report["pass"]:
|
|
raise SystemExit(1)
|
|
|
|
|
|
def materialize_e2e_outputs(artifact_dir, backend_name, e2e_result):
|
|
output_root = Path(artifact_dir) / "reports" / "e2e" / backend_name
|
|
output_root.mkdir(parents=True, exist_ok=True)
|
|
dumped = {
|
|
"text_embeddings": write_npy(output_root, "text_embeddings.npy", e2e_result["text_embeddings"]),
|
|
"initial_latent": write_npy(output_root, "initial_latent.npy", e2e_result["initial_latent"]),
|
|
"final_latent": write_npy(output_root, "final_latent.npy", e2e_result["final_latent"]),
|
|
"final_video": write_npy(output_root, "final_video.npy", e2e_result["final_video"]),
|
|
"steps": [],
|
|
}
|
|
for step in e2e_result["steps"]:
|
|
step_dir = "steps/step_{:02d}".format(step["index"])
|
|
dumped["steps"].append(
|
|
{
|
|
"index": step["index"],
|
|
"timestep": step["timestep"],
|
|
"sample_before": write_npy(output_root, step_dir + "/sample_before.npy", step["sample_before"]),
|
|
"noise_pred": write_npy(output_root, step_dir + "/noise_pred.npy", step["noise_pred"]),
|
|
"noise_pred_uncond": write_npy(output_root, step_dir + "/noise_pred_uncond.npy", step["noise_pred_uncond"]),
|
|
"noise_pred_cond": write_npy(output_root, step_dir + "/noise_pred_cond.npy", step["noise_pred_cond"]),
|
|
"guided_noise": write_npy(output_root, step_dir + "/guided_noise.npy", step["guided_noise"]),
|
|
"sample_after": write_npy(output_root, step_dir + "/sample_after.npy", step["sample_after"]),
|
|
}
|
|
)
|
|
return dumped
|
|
|
|
|
|
def e2e_compare(artifact_dir, manifest, runner, backend_name, atol_override=None, rtol_override=None, write_outputs=False):
|
|
module = manifest["modules"].get("e2e")
|
|
if module is None:
|
|
raise RuntimeError("Manifest does not contain e2e artifacts. Re-run capture with --module e2e or --module all.")
|
|
metadata = manifest["case"]
|
|
input_ids = load_npy(artifact_dir, module["inputs"]["input_ids"])
|
|
attention_mask = load_npy(artifact_dir, module["inputs"]["attention_mask"])
|
|
result = run_e2e_case(runner, metadata, input_ids, attention_mask)
|
|
dumped = materialize_e2e_outputs(artifact_dir, backend_name, result) if write_outputs else None
|
|
atol, rtol = module_tolerance("e2e", atol_override, rtol_override)
|
|
report = {
|
|
"module": "e2e",
|
|
"backend": backend_name,
|
|
"pass": True,
|
|
"tensors": {
|
|
"text_embeddings": compare_arrays(
|
|
load_npy(artifact_dir, module["inputs"]["text_embeddings"]), result["text_embeddings"], atol, rtol
|
|
),
|
|
"initial_latent": compare_arrays(
|
|
load_npy(artifact_dir, module["inputs"]["initial_latent"]), result["initial_latent"], atol, rtol
|
|
),
|
|
"final_latent": compare_arrays(
|
|
load_npy(artifact_dir, module["outputs"]["final_latent"]), result["final_latent"], atol, rtol
|
|
),
|
|
"final_video": compare_arrays(
|
|
load_npy(artifact_dir, module["outputs"]["final_video"]), result["final_video"], atol, rtol
|
|
),
|
|
},
|
|
"steps": [],
|
|
}
|
|
for name in ("text_embeddings", "initial_latent", "final_latent", "final_video"):
|
|
report["pass"] = report["pass"] and report["tensors"][name]["allclose"]
|
|
if len(module["steps"]) != len(result["steps"]):
|
|
raise RuntimeError(
|
|
"E2E step count mismatch: manifest={} backend={}".format(len(module["steps"]), len(result["steps"]))
|
|
)
|
|
for step_ref, step_result in zip(module["steps"], result["steps"]):
|
|
step_report = {
|
|
"index": step_ref["index"],
|
|
"timestep": step_ref["timestep"],
|
|
"tensors": {},
|
|
"pass": True,
|
|
}
|
|
for tensor_name in (
|
|
"sample_before",
|
|
"noise_pred",
|
|
"noise_pred_uncond",
|
|
"noise_pred_cond",
|
|
"guided_noise",
|
|
"sample_after",
|
|
):
|
|
metrics = compare_arrays(
|
|
load_npy(artifact_dir, step_ref[tensor_name]),
|
|
step_result[tensor_name],
|
|
atol=atol,
|
|
rtol=rtol,
|
|
)
|
|
step_report["tensors"][tensor_name] = metrics
|
|
step_report["pass"] = step_report["pass"] and metrics["allclose"]
|
|
report["steps"].append(step_report)
|
|
report["pass"] = report["pass"] and step_report["pass"]
|
|
if dumped is not None:
|
|
report["dumped_outputs"] = dumped
|
|
return report
|
|
|
|
|
|
def e2e_command(args):
|
|
artifact_dir, manifest = load_manifest(args.artifact_dir)
|
|
if args.backend == "reference":
|
|
if not args.model_path:
|
|
raise ValueError("--model_path is required for e2e --backend reference")
|
|
runner = create_reference_runner(args)
|
|
elif args.backend == "onnx":
|
|
if not args.onnx_root:
|
|
raise ValueError("--onnx_root is required for e2e --backend onnx")
|
|
runner = OnnxRunner.from_root(args.onnx_root)
|
|
else:
|
|
if not args.mnn_root:
|
|
raise ValueError("--mnn_root is required for e2e --backend mnn")
|
|
runner = MnnRunner.from_root(
|
|
args.mnn_root,
|
|
thread_num=args.thread_num,
|
|
backend=getattr(args, "mnn_backend", "CPU"),
|
|
)
|
|
report = e2e_compare(
|
|
artifact_dir,
|
|
manifest,
|
|
runner,
|
|
backend_name=args.backend,
|
|
atol_override=args.atol,
|
|
rtol_override=args.rtol,
|
|
write_outputs=True,
|
|
)
|
|
out_path = artifact_dir / "reports" / "e2e_{}.json".format(args.backend)
|
|
save_json(
|
|
out_path,
|
|
{
|
|
"created_at": utc_now(),
|
|
"artifact_dir": artifact_dir.as_posix(),
|
|
"report": report,
|
|
},
|
|
)
|
|
print("Wrote {}".format(out_path))
|
|
if args.compare and not report["pass"]:
|
|
raise SystemExit(1)
|
|
|
|
|
|
def main():
|
|
parser = build_parser()
|
|
args = parser.parse_args()
|
|
try:
|
|
if args.command == "capture":
|
|
capture_command(args)
|
|
elif args.command == "compare-onnx":
|
|
compare_command("compare-onnx", args)
|
|
elif args.command == "compare-mnn":
|
|
compare_command("compare-mnn", args)
|
|
elif args.command == "e2e":
|
|
e2e_command(args)
|
|
else:
|
|
parser.error("Unknown command: {}".format(args.command))
|
|
except KeyboardInterrupt:
|
|
raise
|
|
except SystemExit:
|
|
raise
|
|
except Exception as exc:
|
|
import traceback
|
|
traceback.print_exc()
|
|
print("ERROR:", exc, file=sys.stderr)
|
|
raise SystemExit(2)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|