# SPDX-License-Identifier: AGPL-3.0-only # Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 """Speed + memory lever benchmark for the VIDEO diffusion backend (B200). The video default path already stacks several optimisations (verified in video.py): TE auto-quant, DiT auto-quant when it fits resident, VAE auto (skipped for the fp32-VAE Wan families), regional torch.compile (speed_mode "default"), cuDNN fused attention, and First-Block-Cache auto-engaged at >= 20 steps. This benchmark drives the SAME real lever functions the loader calls -- ``quantize_text_encoders`` / ``quantize_vae`` / ``quantize_transformer`` / ``apply_speed_optims`` / ``apply_attention_backend`` / ``apply_step_cache`` + ``maybe_toggle_step_cache`` -- with the loader's own default arguments, so each measured configuration reflects a real load, not a synthetic one. For each configuration it loads the full pipeline fresh (quant/compile mutate irreversibly), warms up (to pay the one-time compile), then measures a short clip generation: total latency, median per-step ms, peak resident GB, steady weight GB, and per-frame LPIPS(AlexNet) averaged against the bit-exact reference config (everything off/dense/native). This isolates each lever's contribution and answers: how much does the shipped default win, and do ``max`` compile / flash4 attention leave speed on the table for video. Memory/timing idiom lifted from scripts/quant_speedmem_bench.py (reset_peak -> load -> memory_allocated / max_memory_allocated with synchronize + perf_counter). Example: CUDA_VISIBLE_DEVICES=0 python scripts/video_speedmem_bench.py --family wan2.2-ti2v-5b \\ --configs reference,compile,cudnn,fbcache,ditquant,shipped,speedmax,flash4 \\ --steps 30 --num-frames 25 --width 512 --height 320 --iters 3 """ from __future__ import annotations import argparse import gc import json import os import sys import time import types from pathlib import Path from typing import Any, Optional os.environ.setdefault("BITSANDBYTES_NOWELCOME", "1") _REPO_ROOT = Path(__file__).resolve().parent.parent _BACKEND_ROOT = _REPO_ROOT / "studio" / "backend" for _p in (str(_BACKEND_ROOT), str(_REPO_ROOT / "scripts")): if _p not in sys.path: sys.path.insert(0, _p) PROMPT = ( "A cinematic drone shot flying over a misty mountain valley at sunrise, " "golden light, volumetric fog, highly detailed, smooth camera motion" ) _FAMILIES: dict[str, dict[str, Any]] = { "wan2.2-ti2v-5b": { "repo": "Wan-AI/Wan2.2-TI2V-5B-Diffusers", "vae_force_fp32": True, "guidance": 5.0, }, "ltx-2": {"repo": "Lightricks/LTX-2", "vae_force_fp32": False, "guidance": 4.0}, "hunyuanvideo-1.5": { "repo": "hunyuanvideo-community/HunyuanVideo-1.5-Diffusers-480p_t2v", "vae_force_fp32": False, "guidance": 6.0, }, "hunyuanvideo-1.5-720p": { "repo": "hunyuanvideo-community/HunyuanVideo-1.5-Diffusers-720p_t2v", "vae_force_fp32": False, "guidance": 6.0, }, # Wan2.2-A14B is a dual-expert MoE (transformer + transformer_2); _apply_levers quantizes both. "wan2.2-t2v-a14b": { "repo": "Wan-AI/Wan2.2-T2V-A14B-Diffusers", "vae_force_fp32": True, "guidance": 5.0, }, } # ── cuda memory / timing helpers ─────────────────────────────────────────────── def _sync() -> None: import torch if torch.cuda.is_available(): torch.cuda.synchronize() def _reset_peak() -> None: import torch if torch.cuda.is_available(): torch.cuda.reset_peak_memory_stats() def _alloc_gb() -> float: import torch return torch.cuda.memory_allocated() / 1e9 if torch.cuda.is_available() else 0.0 def _peak_gb() -> float: import torch return torch.cuda.max_memory_allocated() / 1e9 if torch.cuda.is_available() else 0.0 def _empty() -> None: import torch gc.collect() if torch.cuda.is_available(): torch.cuda.empty_cache() def _median(xs: list[float]) -> float: return sorted(xs)[len(xs) // 2] if xs else 0.0 _LP: dict = {} def _lpips_alex(ref_arr, arr): """LPIPS(AlexNet) between two HxWx3 uint8 frames (net on CPU). None if lpips missing.""" try: import lpips import torch fn = _LP.get("fn") if fn is None: fn = lpips.LPIPS(net = "alex", verbose = False).eval() _LP["fn"] = fn def _t(a): import torch as _torch return _torch.from_numpy(a).float().permute(2, 0, 1).unsqueeze(0) / 127.5 - 1.0 with torch.no_grad(): return float(fn(_t(ref_arr), _t(arr)).item()) except Exception: return None def _frames_to_arrays(output) -> list: """Normalize a video pipeline output to a list of HxWx3 uint8 numpy frames.""" import numpy as np frames = getattr(output, "frames", None) if frames is None: return [] batch0 = frames[0] arrs = [] for fr in batch0: if hasattr(fr, "convert"): # PIL image arrs.append(np.array(fr.convert("RGB"))) else: a = np.asarray(fr) if a.dtype != np.uint8: a = np.clip(a * (255.0 if a.max() <= 1.0 else 1.0), 0, 255).astype(np.uint8) arrs.append(a) return arrs def _mean_luma(arrs: list) -> Optional[float]: """Mean Rec.601 luma over all frames (0-255). ~0 == black frames (the fp8 failure signal).""" import numpy as np if not arrs: return None vals = [] for a in arrs: a = np.asarray(a).astype(np.float32) if a.ndim == 3 and a.shape[-1] >= 3: luma = 0.299 * a[..., 0] + 0.587 * a[..., 1] + 0.114 * a[..., 2] else: luma = a vals.append(float(luma.mean())) return round(sum(vals) / len(vals), 3) if vals else None def _mean_lpips(ref_arrs: list, arrs: list) -> Optional[float]: """Mean per-frame LPIPS over the min common frame count.""" if not ref_arrs or not arrs: return None n = min(len(ref_arrs), len(arrs)) vals = [] for i in range(n): v = _lpips_alex(ref_arrs[i], arrs[i]) if v is not None: vals.append(v) return round(sum(vals) / len(vals), 4) if vals else None def _import_diffusers(): import torch # noqa: F401 import torchao # noqa: F401 import diffusers.utils.import_utils as iu iu._bitsandbytes_available = False import diffusers return diffusers def _target(): """The object the real casters/optimisers read: a stand-in for DiffusionDeviceTarget. supports_default_torch_compile must be True or compile_eligible() bails.""" import torch return types.SimpleNamespace( device = "cuda", dtype = torch.bfloat16, supports_default_torch_compile = True, ) # ── config matrix ────────────────────────────────────────────────────────────── # Each config names the lever settings applied to a fresh pipe. Built UP from the # bit-exact reference so each successive config isolates one lever's contribution; # `shipped` is the current default; `speedmax`/`flash4` probe untapped headroom. _CONFIGS: dict[str, dict[str, Any]] = { # te vae dit speed attn cache "reference": dict(te = "none", vae = "none", dit = "none", speed = "off", attn = "native", cache = "off"), "compile": dict(te = "none", vae = "none", dit = "none", speed = "default", attn = "native", cache = "off"), "cudnn": dict(te = "none", vae = "none", dit = "none", speed = "default", attn = "auto", cache = "off"), "fbcache": dict(te = "none", vae = "none", dit = "none", speed = "default", attn = "auto", cache = "auto"), "ditquant": dict(te = "none", vae = "none", dit = "auto", speed = "default", attn = "auto", cache = "auto"), "shipped": dict(te = "auto", vae = "auto", dit = "auto", speed = "default", attn = "auto", cache = "auto"), "speedmax": dict(te = "auto", vae = "auto", dit = "auto", speed = "max", attn = "auto", cache = "auto"), "flash4": dict(te = "auto", vae = "auto", dit = "auto", speed = "default", attn = "flash4", cache = "auto"), # diagnostics: isolate whether the DiT-quant + compile crash needs FBCache. "diag_ditq_default_nocache": dict( te = "none", vae = "none", dit = "auto", speed = "default", attn = "native", cache = "off" ), "diag_ditq_max_nocache": dict( te = "none", vae = "none", dit = "auto", speed = "max", attn = "native", cache = "off" ), "diag_ditq_nocompile": dict( te = "none", vae = "none", dit = "auto", speed = "eager", attn = "native", cache = "off" ), # which quant scheme survives torch.compile? (fp8/mslk fails "fake tensors"; test int8/mxfp8) "diag_ditint8_compile": dict( te = "none", vae = "none", dit = "int8", speed = "default", attn = "native", cache = "off" ), "diag_ditmxfp8_compile": dict( te = "none", vae = "none", dit = "mxfp8", speed = "default", attn = "native", cache = "off" ), "diag_ditint8_fbcache": dict( te = "none", vae = "none", dit = "int8", speed = "default", attn = "native", cache = "auto" ), # isolate the quant x FBCache over-caching interaction at production size: "ditfp8_nocache": dict( te = "none", vae = "none", dit = "auto", speed = "default", attn = "auto", cache = "off" ), "te_fbcache": dict( te = "auto", vae = "none", dit = "none", speed = "default", attn = "auto", cache = "auto" ), # Companion-quant accuracy isolation vs the bit-exact reference (uncached, so the cache # cannot mask it): TE-only and VAE-only on top of the trim+cudnn+compile stack. With the # compile rounding fixed (emulate_precision_casts), the companions are the next-largest # divergence source, and only one of them should pay for it. "diag_te_nocache": dict( te = "auto", vae = "none", dit = "none", speed = "default", attn = "auto", cache = "off" ), "diag_vae_nocache": dict( te = "none", vae = "auto", dit = "none", speed = "default", attn = "auto", cache = "off" ), "ditfp8_fbcache": dict( te = "none", vae = "none", dit = "auto", speed = "default", attn = "auto", cache = "auto" ), "ditint8_fbcache_prod": dict( te = "none", vae = "none", dit = "int8", speed = "default", attn = "auto", cache = "auto" ), # mixed-fp8 vs int8 head-to-head (Phase 3): the DiT-quant accuracy comparison on Wan/Hunyuan. # fp8 here goes through the production quantize_transformer family exclude (input embedders # kept bf16), so it is only non-black if the mixed-fp8 wiring is live. cache on AND off. "ditfp8mixed_nocache": dict( te = "none", vae = "none", dit = "fp8", speed = "default", attn = "native", cache = "off" ), "ditfp8mixed_fbcache": dict( te = "none", vae = "none", dit = "fp8", speed = "default", attn = "auto", cache = "auto" ), "ditint8_nocache": dict( te = "none", vae = "none", dit = "int8", speed = "default", attn = "native", cache = "off" ), # Hunyuan padded-text trim isolation: "cudnn" above is the same stack WITH the trim # (it auto-engages under any active speed tier), so trim_off isolates its win. The # trim key defaults True everywhere else; only this row forces it off. "trim_off": dict( te = "none", vae = "none", dit = "none", speed = "default", attn = "auto", cache = "off", trim = False, ), # Compile isolation at matched attention/trim: eager tier (channels_last + cudnn # benchmark, NO compile) vs the "cudnn" row (default tier = regional compile). "eager_trim": dict(te = "none", vae = "none", dit = "none", speed = "eager", attn = "auto", cache = "off"), # int8 DiT baseline at the production attention stack (cudnn + trim), cache off, # directly comparable to the "cudnn" dense row. "int8_cudnn": dict( te = "none", vae = "none", dit = "int8", speed = "default", attn = "auto", cache = "off" ), # The full companion-quant stack WITHOUT step caching: te/vae auto + dit auto (dense # -fit skip on Hunyuan) + compile + cudnn + trim. The stacked-best candidate default # when the family's auto cache policy stays off. "shipped_nocache": dict( te = "auto", vae = "auto", dit = "auto", speed = "default", attn = "auto", cache = "off" ), # Compile-parity isolation: the same compiled stack as "cudnn" but with inductor's # emulate_precision_casts turned back OFF after the speed layer set it, so the row # measures the numeric + speed effect of the round-2 parity flag per family. "epc_off": dict( te = "none", vae = "none", dit = "none", speed = "default", attn = "auto", cache = "off", epc = False, ), # Explicit cache modes at the dense compiled stack (bypass the family AUTO policy), # for FBCache-vs-MagCache head-to-head rows at identical settings. "fbcache_explicit": dict( te = "none", vae = "none", dit = "none", speed = "default", attn = "auto", cache = "fbcache" ), "magcache_explicit": dict( te = "none", vae = "none", dit = "none", speed = "default", attn = "auto", cache = "magcache" ), } def _ref_cache_path(out, *, family, seed, steps, num_frames, width, height): """Reference-frame cache keyed by every parameter that changes the reference clip. The persisted reference (for the "run reference once, score other configs in parallel processes" workflow) shares the default --out dir across runs, so a single unkeyed ref_frames.npz would let a later reference-less run of a different family / seed / steps / frames / resolution score LPIPS against the wrong baseline. Key by all of them so a reference-less run only reuses a reference computed for the same parameters.""" from pathlib import Path return Path(out) / ( f"ref_frames_{family}_seed{seed}_st{steps}_f{num_frames}_{width}x{height}.npz" ) def _build_pipe(repo: str, force_fp32_vae: bool): import torch diffusers = _import_diffusers() # Wan-style VAEs decode in fp32 for numerical stability (the loader pins this via # vae_force_fp32). A scalar bf16 torch_dtype truncates the fp32-stored VAE weights at # load, and a later .to(float32) only widens the already-lossy values (banding), so the # bench would measure a decode path production never runs. Pin the VAE fp32 per-component # exactly like the production loader (video.py: {"vae": fp32, "default": bf16}). torch_dtype = torch.bfloat16 if force_fp32_vae: torch_dtype = {"vae": torch.float32, "default": torch.bfloat16} pipe = diffusers.DiffusionPipeline.from_pretrained(repo, torch_dtype = torch_dtype) # Stays on CPU: the caller applies the configured levers FIRST and only then places on # CUDA, mirroring the loader (video.py quantizes before apply_memory_plan). Placing the # dense pipeline first would OOM configs whose quantized form fits but whose dense form # does not, and would record a dense load peak for a quantized row. if force_fp32_vae and getattr(pipe, "vae", None) is not None: pipe.vae.to(torch.float32) # belt-and-suspenders; a no-op on the primary path above return pipe class _SecondExpertView: """Present ``pipe.transformer_2`` as ``.transformer`` so the single-DiT lever functions run on the second expert of a dual-expert MoE (Wan2.2-A14B) unforked -- mirrors the loader's _SecondDiTView (video.py). Attribute reads delegate to the real pipe except ``transformer``, and a ``transformer`` write (e.g. torch.compile reassigning it) is routed to ``transformer_2``.""" def __init__(self, pipe): object.__setattr__(self, "_pipe", pipe) @property def transformer(self): return self._pipe.transformer_2 def __getattr__(self, name): return getattr(self._pipe, name) def __setattr__(self, name, value): setattr(self._pipe, "transformer_2" if name == "transformer" else name, value) def _apply_levers( pipe, cfg: dict, *, fam_name: str, fam_obj, force_fp32_vae: bool, default_steps: int, cache_threshold: Optional[float] = None, cache_quality: Optional[str] = None, logger = None, ) -> dict: """Apply the configured levers with the loader's own argument values, in the loader's order: quant (dit -> te -> vae) THEN optimisation layers (cache -> attention -> speed). For a dual-expert MoE (pipe.transformer_2 present) every DiT-touching lever is applied to BOTH experts via _SecondExpertView, exactly like the loader, so A14B latency + accuracy are real.""" from core.inference.diffusion_precision import quantize_text_encoders from core.inference.diffusion_vae_quant import quantize_vae from core.inference.diffusion_transformer_quant import ( quantize_transformer, is_int8_memory_fallback, ) from core.inference.diffusion_speed import apply_speed_optims, snapshot_backend_flags from core.inference.diffusion_attention import ( select_attention_backend, apply_attention_backend, install_hunyuan_attention_trim, ) from core.inference.diffusion_cache import ( apply_step_cache, auto_cache_mode, auto_cache_quality, normalize_cache_quality, FBCACHE_MIN_STEPS, ) tgt = _target() engaged = { "dit": None, "te": None, "vae": None, "attn": None, "cache": None, "speed_optims": {}, } # DiT-touching levers run per expert: [pipe] for a single-DiT family, plus a second-expert view # for a dual-expert MoE. Each view exposes the expert as ``.transformer``. views = [pipe] if getattr(pipe, "transformer_2", None) is not None: views.append(_SecondExpertView(pipe)) # DiT quant (pipeline kind, resident): mutates each expert's transformer in place. Mirror the # loader's dense-fit skip: for an AUTO request on an int8-fallback family (HunyuanVideo-1.5, where # fp8 is black-framed so auto lands on int8, a memory-only lever ~7% slower AND less accurate than # dense+compile), run the dense DiT instead when it fits resident -- and the benchmark always # loads resident (no offload). Explicit int8/fp8 configs are honored (the whole point of the # sweep). Without this the "shipped"/"ditquant" auto rows would measure int8 where the loader runs # dense, overstating the shipped cost on Hunyuan. dense_fit_skip = cfg["dit"] == "auto" and is_int8_memory_fallback(tgt, fam_name) if cfg["dit"] not in ("none", "off") and not dense_fit_skip: schemes = [ quantize_transformer(v, tgt, mode = cfg["dit"], family = fam_name, logger = logger) for v in views ] # All-or-none across experts, mirroring the production loader (video.py): the first # expert is mutated in place, so a second-expert miss cannot fall back to dense and # the loader fails that load. Rejecting here keeps the benchmark from publishing a # dit_scheme + timings for a mixed quantized/dense pipeline users cannot actually load. n_engaged = sum(1 for s in schemes if s is not None) if 0 < n_engaged < len(views): raise RuntimeError( f"dit quant '{cfg['dit']}' engaged on only {n_engaged}/{len(views)} experts; " "production rejects this partial state, so the row would be unloadable" ) engaged["dit"] = schemes[0] engaged["dit_experts"] = schemes _empty() dit_quant_active = engaged["dit"] is not None # TE quant (once; text encoders are shared, not per-expert). if cfg["te"] not in ("none", "off"): engaged["te"] = quantize_text_encoders( pipe, tgt, mode = cfg["te"], family = fam_name, offload_active = False, logger = logger ) _empty() # VAE quant (once; Wan force_fp32 pins dense inside quantize_vae regardless). if cfg["vae"] not in ("none", "off"): engaged["vae"] = quantize_vae( pipe, tgt, mode = cfg["vae"], family = fam_name, offload_active = False, force_fp32 = force_fp32_vae, logger = logger, ) _empty() # ── optimisation layers ── # Process-wide flags (cudnn.benchmark, TF32/fp16 accumulation, # emulate_precision_casts): snapshotted here and restored by _run_config after # the row completes, like the production unload path. Without the restore a # speed-enabled row would poison a later reference/off row in the same process. engaged["_flags_snapshot"] = snapshot_backend_flags() speed = cfg["speed"] speed_active = speed != "off" # A quantized DiT must be compiled (eager dynamic quant ~30x slower), matching the loader. if dit_quant_active and speed == "off": speed = "default" speed_active = True # Step cache FIRST (compile keys fullgraph off an active cache); per expert. cache_active = False if cfg["cache"] in ("auto", "fbcache", "magcache"): # Per-family auto mode, exactly like the loader (video.py): MagCache for the # HunyuanVideo-1.5 families, FBCache elsewhere. An explicit "fbcache"/"magcache" # config value bypasses the auto policy (the head-to-head rows). if cfg["cache"] == "auto": cache_request = ( auto_cache_mode(fam_name) if default_steps >= FBCACHE_MIN_STEPS else None ) else: cache_request = cfg["cache"] if cache_request is not None: # Quality preset resolution, exactly like the loader (video.py): an unset # request takes the family's measured auto default. Expert names zip with # the views (the loader's expert-view iteration contract) so a dual-expert # MoE resolves per-expert MagCache curves. quality = normalize_cache_quality(cache_quality) or auto_cache_quality(fam_name) experts = ("transformer", "transformer_2") for v, expert in zip(views, experts): engaged["cache"] = apply_step_cache( v, mode = cache_request, threshold = cache_threshold, quant_active = dit_quant_active, family = fam_name, steps = default_steps, quality = quality, expert = expert, logger = logger, ) cache_active = engaged["cache"] not in (None, "off") # HunyuanVideo-1.5 joint-attention trim (per expert), BEFORE the backend set so the requested # kernel pins onto the new processors -- exactly the loader's order. Drops the ~99% zero-padded # text tokens so the fused SDPA kernel runs (~18x/DiT-forward, cosine ~1.0). A speed lever, so # gated on an active tier like the loader; no-op for every non-Hunyuan family. trim_engaged = False if speed_active and cfg.get("trim", True): for v in views: trim_engaged = install_hunyuan_attention_trim(v, fam_obj, logger = logger) or trim_engaged engaged["attn_trim"] = trim_engaged # Attention (per expert). backend = select_attention_backend(tgt, cfg["attn"], speed_active = speed_active) for v in views: engaged["attn"] = apply_attention_backend(v, backend, logger = logger) # Speed profile (per expert; compiles each denoiser). if speed != "off": for v in views: engaged["speed_optims"] = apply_speed_optims( v, tgt, is_gguf = False, family = fam_obj, speed_mode = speed, cache_active = cache_active, offload_active = False, logger = logger, ) engaged["_effective_speed"] = speed # emulate_precision_casts A/B: the speed layer sets the flag True inside # _compile_repeated_blocks; compile is lazy (first forward), so flipping it back off # here gives the pre-round-2 inductor numerics for the whole run ("epc_off" config). if not cfg.get("epc", True): try: import torch torch._inductor.config.emulate_precision_casts = False engaged["epc"] = False except Exception: pass return engaged def _timed_video( pipe, *, steps, width, height, num_frames, guidance, seed, cache_mode, dit_quant_active, default_steps, guidance_via_guider = False, cache_threshold = None, cache_quality = None, family = None, logger = None, ): """One clip generation. Re-checks the step cache per generation (maybe_toggle_step_cache) exactly like the loader, then times total + per-step. Returns (output, total_s, [per_step_ms]).""" import torch from core.inference.diffusion_cache import ( _disengage_step_cache, apply_step_cache, auto_cache_mode, auto_cache_quality, maybe_toggle_step_cache, normalize_cache_quality, ) # Toggle on EVERY expert view, exactly like the loader's per-view recheck # (video.py iterates _views_for): toggling only the primary pipe would disable # FBCache on pipe.transformer while transformer_2 stays cached, measuring a # mixed cache state production never runs on a dual-expert MoE. views = [pipe] if getattr(pipe, "transformer_2", None) is not None: views.append(_SecondExpertView(pipe)) if cache_mode == "auto": for v, expert in zip(views, ("transformer", "transformer_2")): try: maybe_toggle_step_cache( v, steps = steps, quant_active = dit_quant_active, threshold = cache_threshold, mode = auto_cache_mode(family), family = family, quality = normalize_cache_quality(cache_quality) or auto_cache_quality(family), expert = expert, logger = logger, ) except Exception: pass elif cache_mode == "magcache": # An explicit magcache row never toggles off, but production re-engages it # when the actual step count differs from the configured one (video.py: the # marker carries "#s{steps}"; endswith, not substring, so "#s5" cannot match # inside "#s50"). _apply_levers installed the cache at default_steps, so a # --steps override would otherwise time a stale curve/skip budget users # never run. for v, expert in zip(views, ("transformer", "transformer_2")): transformer = getattr(v, "transformer", None) marker = getattr(transformer, "_unsloth_step_cache", None) if not marker or str(marker).endswith(f"#s{int(steps)}"): continue try: _disengage_step_cache( transformer, reason = f"explicit magcache re-interpolating for {steps} steps", logger = logger, ) apply_step_cache( v, mode = "magcache", threshold = cache_threshold, quant_active = dit_quant_active, family = family, steps = steps, quality = normalize_cache_quality(cache_quality) or auto_cache_quality(family), expert = expert, logger = logger, ) except Exception: pass # Clear step-cache residuals before EVERY generation, mirroring production # (VideoBackend.generate calls _reset_step_cache before each run): diffusers keys # the residuals on the long-lived transformer, so the measured iterations after a # warmup at identical shape/seed would otherwise start against the previous clip's # cache state -- a behavior users never get. Best-effort, uncached is a no-op. for name in ("transformer", "transformer_2"): module = getattr(pipe, name, None) reset = getattr(module, "_reset_stateful_cache", None) or getattr( module, "reset_stateful_hooks", None ) if callable(reset): try: reset() except Exception: pass g = torch.Generator(device = "cuda").manual_seed(seed) step_ts: list[float] = [] last = [0.0] def _cb(pp, i, t, kw): torch.cuda.synchronize() now = time.perf_counter() if last[0]: step_ts.append((now - last[0]) * 1000.0) last[0] = now return kw kwargs = dict( prompt = PROMPT, width = width, height = height, num_frames = num_frames, num_inference_steps = steps, generator = g, ) restore_step = None if guidance_via_guider: # HunyuanVideo-1.5: CFG lives on a guider component and __call__ takes no # guidance_scale / callback_on_step_end (the loader writes the scale onto pipe.guider). guider = getattr(pipe, "guider", None) if guider is not None and hasattr(guider, "guidance_scale"): try: guider.guidance_scale = guidance except Exception: pass # __call__ ignores callback_on_step_end here, so the _cb step timer never fires and # step_ts would stay empty -> per_step_ms published as 0.0 for every Hunyuan row despite a # non-zero total latency. Time the denoise via a scheduler.step wrapper instead (the same # inter-step delta the callback records elsewhere), restored after generation. sched = getattr(pipe, "scheduler", None) orig_step = getattr(sched, "step", None) if callable(orig_step): def _timed_step(*a, **k): torch.cuda.synchronize() now = time.perf_counter() if last[0]: step_ts.append((now - last[0]) * 1000.0) last[0] = now return orig_step(*a, **k) sched.step = _timed_step restore_step = lambda: setattr(sched, "step", orig_step) # noqa: E731 else: kwargs["guidance_scale"] = guidance kwargs["callback_on_step_end"] = _cb _sync() t0 = time.perf_counter() try: out = pipe(**kwargs) finally: if restore_step is not None: restore_step() _sync() return out, (time.perf_counter() - t0), step_ts def _run_config( name: str, cfg: dict, *, family: str, steps: int, width: int, height: int, num_frames: int, seed: int, iters: int, out: Path, cache_threshold: Optional[float] = None, cache_quality: Optional[str] = None, logger = None, ): import numpy as np from core.inference.video_families import detect_video_family spec = _FAMILIES[family] repo = spec["repo"] force_fp32 = spec.get("vae_force_fp32", False) guidance = spec.get("guidance", 5.0) fam_obj = detect_video_family(repo) default_steps = getattr(fam_obj, "default_steps", 50) gvg = bool(getattr(fam_obj, "guidance_via_guider", False)) _empty() _reset_peak() # Fresh dynamo state per config so a prior config's compiled graphs cannot leak into this one # (each config builds a fresh pipe; without this, later configs in a multi-config run can be # measured against a dirty compile cache). try: import torch torch._dynamo.reset() except Exception: pass # Loader order (video.py): build on CPU, apply quant/optim levers, THEN place on CUDA. # Placing dense first would OOM configs whose quantized form fits but dense does not, and # load_peak_gb would record the dense placement instead of the measured config. pipe = _build_pipe(repo, force_fp32) engaged = _apply_levers( pipe, cfg, fam_name = family, fam_obj = fam_obj, force_fp32_vae = force_fp32, default_steps = default_steps, cache_threshold = cache_threshold, cache_quality = cache_quality, logger = logger, ) pipe = pipe.to("cuda") _sync() load_peak = _peak_gb() _empty() weights_gb = _alloc_gb() dit_active = engaged["dit"] is not None cache_mode = cfg["cache"] # warmup (pays the one-time compile / autotune) warmup_t0 = time.perf_counter() _timed_video( pipe, steps = steps, width = width, height = height, num_frames = num_frames, guidance = guidance, seed = seed, cache_mode = cache_mode, dit_quant_active = dit_active, default_steps = default_steps, guidance_via_guider = gvg, cache_threshold = cache_threshold, cache_quality = cache_quality, family = family, logger = logger, ) warmup_s = time.perf_counter() - warmup_t0 _reset_peak() dts, steps_ms = [], [] last_out = None for _ in range(iters): last_out, dt, st = _timed_video( pipe, steps = steps, width = width, height = height, num_frames = num_frames, guidance = guidance, seed = seed, cache_mode = cache_mode, dit_quant_active = dit_active, default_steps = default_steps, guidance_via_guider = gvg, cache_threshold = cache_threshold, cache_quality = cache_quality, family = family, logger = logger, ) dts.append(dt) steps_ms.append(_median(st) if st else 0.0) gen_peak = _peak_gb() arrs = _frames_to_arrays(last_out) # save a mid frame for eyeballing try: if arrs: from PIL import Image Image.fromarray(arrs[len(arrs) // 2]).save(out / f"vid_{family}_{name}.png") except Exception: pass # Persist reference frames so parallel per-config processes can score LPIPS against them. if name == "reference" and arrs: try: import numpy as _np _np.savez_compressed( _ref_cache_path( out, family = family, seed = seed, steps = steps, num_frames = num_frames, width = width, height = height, ), *arrs, ) except Exception: pass # Persist EVERY config's frames too, so a row generated before the reference exists # (parallel per-GPU processes) can be LPIPS-rescored offline instead of publishing null. if arrs: try: import numpy as _np _np.savez_compressed(out / f"frames_{family}_{name}.npz", *arrs) except Exception: pass # Report the GENERATION-time cache state, not the load-time one: the per-generation # recheck in _timed_video can toggle an auto cache off (steps below the threshold) or # re-size an explicit magcache, and production status reports the post-toggle state. # The marker is "{mode}@{threshold}[#s{steps}]", written by apply_step_cache and # cleared by _disengage_step_cache, so its mode prefix IS the live state. cache_marker = getattr(getattr(pipe, "transformer", None), "_unsloth_step_cache", None) row = { "config": name, "family": family, "levers": cfg, "dit_scheme": engaged["dit"] or "dense", "dit_experts": engaged.get("dit_experts"), "te_scheme": engaged["te"] or "dense", "vae_scheme": engaged["vae"] or "dense", "attn": engaged["attn"] or "native", "cache": str(cache_marker).split("@")[0] if cache_marker else "off", "cache_at_load": engaged["cache"] or "off", "effective_speed": engaged.get("_effective_speed"), "speed_optims": engaged["speed_optims"], "attn_trim": engaged.get("attn_trim", False), "cache_threshold": cache_threshold, "cache_quality": cache_quality, "cache_marker": cache_marker, "load_peak_gb": round(load_peak, 2), "weights_gb": round(weights_gb, 2), "gen_peak_gb": round(gen_peak, 2), "warmup_s": round(warmup_s, 3), "gen_latency_s": round(_median(dts), 3), "per_step_ms": round(_median(steps_ms), 1), "n_frames": len(arrs), "mean_luma": _mean_luma(arrs), } del pipe _empty() # Restore the process-wide backend flags this config's speed layer mutated # (snapshot captured in _apply_levers before any mutation), mirroring the # production unload path, so a later config in the same --configs run starts # from clean flags instead of inheriting cudnn.benchmark / TF32 / # emulate_precision_casts from this row. An exception skips this, but the main # loop does not continue past a failed config, so no row can run with leaked # flags. try: from core.inference.diffusion_speed import restore_backend_flags restore_backend_flags(engaged.get("_flags_snapshot")) except Exception: pass return row, arrs def main(argv = None) -> int: ap = argparse.ArgumentParser(description = __doc__) ap.add_argument("--family", default = "wan2.2-ti2v-5b", choices = sorted(_FAMILIES)) ap.add_argument( "--configs", default = ",".join(_CONFIGS), help = "comma list from: " + ",".join(_CONFIGS) ) ap.add_argument("--steps", type = int, default = 30) ap.add_argument("--num-frames", type = int, default = 25) ap.add_argument("--width", type = int, default = 512) ap.add_argument("--height", type = int, default = 320) ap.add_argument("--seed", type = int, default = 42) ap.add_argument("--iters", type = int, default = 3) ap.add_argument("--out", default = "outputs/video_speedmem") ap.add_argument( "--cache-threshold", type = float, default = None, help = "FBCache residual-diff threshold override (None -> the production default)", ) ap.add_argument( "--cache-quality", default = None, choices = ("quality", "balanced", "fast"), help = "Step-cache quality preset (None -> the family's production auto default)", ) args = ap.parse_args(argv) import logging logging.basicConfig(level = logging.INFO, format = "%(message)s") logger = logging.getLogger("videobench") out = Path(args.out) out.mkdir(parents = True, exist_ok = True) names = [c.strip() for c in args.configs.split(",") if c.strip()] for n in names: if n not in _CONFIGS: raise SystemExit(f"unknown config '{n}'; choose from {list(_CONFIGS)}") print( f"== video speed+mem bench: family={args.family} configs={names} " f"steps={args.steps} frames={args.num_frames} {args.width}x{args.height} ==", flush = True, ) rows = [] ref_arrs = None # If not (re)computing the reference in this run, load persisted reference frames for LPIPS. if "reference" not in names: ref_npz = _ref_cache_path( out, family = args.family, seed = args.seed, steps = args.steps, num_frames = args.num_frames, width = args.width, height = args.height, ) if ref_npz.exists(): try: import numpy as _np with _np.load(ref_npz) as z: ref_arrs = [z[k] for k in z.files] except Exception: ref_arrs = None row_arrs: list = [] for n in names: row, arrs = _run_config( n, _CONFIGS[n], family = args.family, steps = args.steps, width = args.width, height = args.height, num_frames = args.num_frames, seed = args.seed, iters = args.iters, out = out, cache_threshold = args.cache_threshold, cache_quality = args.cache_quality, logger = logger, ) if n == "reference": ref_arrs = arrs row["lpips_vs_reference"] = _mean_lpips(ref_arrs, arrs) if ref_arrs is not None else None rows.append(row) row_arrs.append(arrs) print( f" [{n}] {json.dumps({k: row[k] for k in ('dit_scheme','te_scheme','attn','cache','effective_speed','weights_gb','gen_peak_gb','gen_latency_s','per_step_ms','lpips_vs_reference')})}", flush = True, ) # --configs accepts an arbitrary order, so rows finalized BEFORE the reference clip was # generated (e.g. --configs shipped,reference) scored lpips_vs_reference as None; rescore # them now that the reference frames exist instead of silently publishing null. if ref_arrs is not None: for r, arrs in zip(rows, row_arrs): if r["lpips_vs_reference"] is None: r["lpips_vs_reference"] = _mean_lpips(ref_arrs, arrs) # speedups relative to reference (if present) ref_lat = next((r["gen_latency_s"] for r in rows if r["config"] == "reference"), None) for r in rows: r["speedup_vs_reference"] = ( round(ref_lat / r["gen_latency_s"], 3) if ref_lat and r["gen_latency_s"] else None ) dest = out / f"video_{args.family}_{'-'.join(names)}.json" with open(dest, "w", encoding = "utf-8") as fh: json.dump(rows, fh, indent = 2) print(f"wrote {dest}", flush = True) return 0 if __name__ == "__main__": raise SystemExit(main())