Add video speed/memory lever benchmark; extend quant bench with DiT mode
video_speedmem_bench.py drives the real video-loader lever functions (quantize_transformer / quantize_text_encoders / quantize_vae / apply_speed_optims / apply_attention_backend / apply_step_cache) with the loader's own defaults, so each measured config reflects a real load. It decomposes the video speed/memory stack (compile, cuDNN attention, First-Block-Cache, DiT/TE/VAE quant) with per-step latency, peak resident GB, and per-frame LPIPS vs a bit-exact reference. This is the harness that surfaced and validated the Wan fp8 black-frame fix. quant_speedmem_bench.py gains the DiT-quant mode (dense vs fp8/int8/mxfp8 speed, peak memory, and LPIPS vs the dense render) plus the shared LPIPS(AlexNet) helper.
This commit is contained in:
Daniel Han
parent
cbf24dd847
commit
5d501f5086
2 changed files with 669 additions and 2 deletions
|
|
@ -105,6 +105,60 @@ 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 images (net kept 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 round(float(fn(_t(ref_arr), _t(arr)).item()), 4)
|
||||
except Exception:
|
||||
return None
|
||||
|
||||
|
||||
def _timed_generate(pipe, *, steps, res, seed):
|
||||
"""One generation returning (PIL image, total seconds, [per-step ms]). Per-step wall-clock
|
||||
via callback_on_step_end (each step synchronised)."""
|
||||
import time as _time
|
||||
|
||||
import torch
|
||||
|
||||
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
|
||||
|
||||
_sync()
|
||||
t0 = _time.perf_counter()
|
||||
img = pipe(
|
||||
prompt=PROMPT, width=res, height=res, num_inference_steps=steps, generator=g,
|
||||
callback_on_step_end=_cb,
|
||||
).images[0]
|
||||
_sync()
|
||||
return img, (_time.perf_counter() - t0), step_ts
|
||||
|
||||
|
||||
def _import_diffusers():
|
||||
"""diffusers with the bnb quantiser disabled (we quant via torchao / layerwise only)."""
|
||||
import torch # noqa: F401 (torch/torchao first so their extensions register)
|
||||
|
|
@ -567,11 +621,107 @@ def measure_e2e(
|
|||
return rows
|
||||
|
||||
|
||||
# ── mode: dit (transformer quant, the real speed lever) ───────────────────────
|
||||
def _compile_blocks(transformer) -> bool:
|
||||
"""Regional block compile (the real feature path); torchao dynamic quant is ~30x slower eager,
|
||||
so both dense and quant variants are compiled for a fair speedup comparison."""
|
||||
fn = getattr(transformer, "compile_repeated_blocks", None)
|
||||
if not callable(fn):
|
||||
return False
|
||||
for kw in ({"dynamic": True}, {}):
|
||||
try:
|
||||
fn(**kw)
|
||||
return True
|
||||
except Exception:
|
||||
continue
|
||||
return False
|
||||
|
||||
|
||||
def _dit_run(
|
||||
repo: str, family: str, *, dit_quant: str, steps: int, res: int, seed: int, iters: int,
|
||||
compile_blocks: bool = True, logger=None,
|
||||
):
|
||||
"""Load the full pipeline dense, quantise ONLY the transformer (TE + VAE stay dense to isolate
|
||||
the DiT), regional-compile it (the real feature path), then measure per-step + total latency and
|
||||
peak resident memory. Returns row + image."""
|
||||
import torch
|
||||
|
||||
from core.inference.diffusion_transformer_quant import quantize_transformer
|
||||
|
||||
diffusers = _import_diffusers()
|
||||
_empty(); _reset_peak()
|
||||
pipe = diffusers.DiffusionPipeline.from_pretrained(repo, torch_dtype=torch.bfloat16).to("cuda")
|
||||
load_peak = _peak_gb()
|
||||
engaged = None
|
||||
if dit_quant and dit_quant != "none":
|
||||
engaged = quantize_transformer(pipe, _target(), mode=dit_quant, family=family, logger=logger)
|
||||
_empty()
|
||||
weights_gb = _alloc_gb()
|
||||
compiled = _compile_blocks(getattr(pipe, "transformer", None)) if compile_blocks else False
|
||||
|
||||
img, _, _ = _timed_generate(pipe, steps=steps, res=res, seed=seed) # warmup (triggers compile)
|
||||
_reset_peak()
|
||||
dts, steps_ms, last_img = [], [], img
|
||||
for _ in range(iters):
|
||||
last_img, dt, st = _timed_generate(pipe, steps=steps, res=res, seed=seed)
|
||||
dts.append(dt)
|
||||
steps_ms.append(_median(st) if st else 0.0)
|
||||
gen_peak = _peak_gb()
|
||||
del pipe
|
||||
_empty()
|
||||
return {
|
||||
"family": family,
|
||||
"dit_quant": dit_quant,
|
||||
"dit_scheme": engaged or "dense",
|
||||
"compiled": compiled,
|
||||
"load_peak_gb": round(load_peak, 2),
|
||||
"weights_gb": round(weights_gb, 2),
|
||||
"gen_peak_gb": round(gen_peak, 2),
|
||||
"gen_latency_s": round(_median(dts), 3),
|
||||
"per_step_ms": round(_median(steps_ms), 1),
|
||||
}, last_img
|
||||
|
||||
|
||||
def measure_dit(family: str, *, schemes, steps: int, res: int, seed: int, iters: int, out: Path, logger=None):
|
||||
"""Dense reference + each DiT scheme (auto/fp8/int8/mxfp8), reporting speedup, peak-memory drop,
|
||||
and LPIPS(AlexNet) vs the dense render (the whole-image accuracy metric)."""
|
||||
import numpy as np
|
||||
|
||||
repo = _FAMILIES[family]["repo"]
|
||||
dense_row, dense_img = _dit_run(
|
||||
repo, family, dit_quant="none", steps=steps, res=res, seed=seed, iters=iters, logger=logger
|
||||
)
|
||||
try:
|
||||
dense_img.save(out / f"dit_{family}_dense.png")
|
||||
except Exception:
|
||||
pass
|
||||
ref_arr = np.array(dense_img)
|
||||
dense_row["lpips_vs_dense"] = 0.0
|
||||
dense_row["speedup_vs_dense"] = 1.0
|
||||
rows = [dense_row]
|
||||
print(f" dit dense: {json.dumps(dense_row)}", flush=True)
|
||||
base_lat = dense_row["gen_latency_s"] or 1.0
|
||||
for scheme in schemes:
|
||||
row, img = _dit_run(
|
||||
repo, family, dit_quant=scheme, steps=steps, res=res, seed=seed, iters=iters, logger=logger
|
||||
)
|
||||
row["lpips_vs_dense"] = _lpips_alex(ref_arr, np.array(img))
|
||||
row["speedup_vs_dense"] = round(base_lat / row["gen_latency_s"], 3) if row["gen_latency_s"] else None
|
||||
try:
|
||||
img.save(out / f"dit_{family}_{scheme}.png")
|
||||
except Exception:
|
||||
pass
|
||||
rows.append(row)
|
||||
print(f" dit {scheme:5s}: {json.dumps(row)}", flush=True)
|
||||
return rows
|
||||
|
||||
|
||||
# ── main ──────────────────────────────────────────────────────────────────────
|
||||
def main(argv=None) -> int:
|
||||
ap = argparse.ArgumentParser(description=__doc__)
|
||||
ap.add_argument("--family", required=True, choices=sorted(_FAMILIES))
|
||||
ap.add_argument("--mode", required=True, choices=("te", "vae", "e2e", "teacc"))
|
||||
ap.add_argument("--mode", required=True, choices=("te", "vae", "e2e", "teacc", "dit"))
|
||||
ap.add_argument("--dit-schemes", default="auto", help="dit mode: comma list e.g. auto,fp8,int8,mxfp8")
|
||||
ap.add_argument("--warmup", type=int, default=2)
|
||||
ap.add_argument("--iters", type=int, default=5)
|
||||
ap.add_argument("--steps", type=int, default=20, help="e2e denoise steps")
|
||||
|
|
@ -592,7 +742,13 @@ def main(argv=None) -> int:
|
|||
out.mkdir(parents=True, exist_ok=True)
|
||||
|
||||
print(f"== speed+mem bench: family={args.family} mode={args.mode} ==", flush=True)
|
||||
if args.mode == "teacc":
|
||||
if args.mode == "dit":
|
||||
schemes = [s.strip() for s in args.dit_schemes.split(",") if s.strip()]
|
||||
rows = measure_dit(
|
||||
args.family, schemes=schemes, steps=args.steps, res=args.res, seed=args.seed,
|
||||
iters=args.e2e_iters, out=out, logger=logger
|
||||
)
|
||||
elif args.mode == "teacc":
|
||||
rows = measure_te_accuracy(args.family, logger=logger)
|
||||
elif args.mode == "te":
|
||||
rows = measure_te(args.family, warmup=args.warmup, iters=args.iters, scheme=args.te_scheme, logger=logger)
|
||||
|
|
|
|||
511
scripts/video_speedmem_bench.py
Normal file
511
scripts/video_speedmem_bench.py
Normal file
|
|
@ -0,0 +1,511 @@
|
|||
# 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")
|
||||
# The video base repos live in the workspace BACKUP cache; honor an override but default to it.
|
||||
os.environ.setdefault("HF_HOME", "/mnt/disks/unslothai/ubuntu/workspace_81/BACKUP_05/temp/hf_cache")
|
||||
|
||||
_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,
|
||||
},
|
||||
}
|
||||
|
||||
|
||||
# ── 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_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"),
|
||||
"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"),
|
||||
}
|
||||
|
||||
|
||||
def _build_pipe(repo: str, force_fp32_vae: bool):
|
||||
import torch
|
||||
|
||||
diffusers = _import_diffusers()
|
||||
pipe = diffusers.DiffusionPipeline.from_pretrained(repo, torch_dtype=torch.bfloat16)
|
||||
pipe = pipe.to("cuda")
|
||||
# Wan-style VAEs decode in fp32 for numerical stability (the loader pins this via
|
||||
# vae_force_fp32); loading bf16 bands every clip, so mirror the loader.
|
||||
if force_fp32_vae and getattr(pipe, "vae", None) is not None:
|
||||
pipe.vae.to(torch.float32)
|
||||
return pipe
|
||||
|
||||
|
||||
def _apply_levers(pipe, cfg: dict, *, fam_name: str, fam_obj, force_fp32_vae: bool, default_steps: int,
|
||||
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)."""
|
||||
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
|
||||
from core.inference.diffusion_speed import apply_speed_optims, snapshot_backend_flags
|
||||
from core.inference.diffusion_attention import select_attention_backend, apply_attention_backend
|
||||
from core.inference.diffusion_cache import (
|
||||
apply_step_cache,
|
||||
TC_FBCACHE,
|
||||
FBCACHE_MIN_STEPS,
|
||||
)
|
||||
|
||||
tgt = _target()
|
||||
engaged = {"dit": None, "te": None, "vae": None, "attn": None, "cache": None, "speed_optims": {}}
|
||||
|
||||
# DiT quant (pipeline kind, resident): mutates pipe.transformer in place.
|
||||
if cfg["dit"] not in ("none", "off"):
|
||||
engaged["dit"] = quantize_transformer(
|
||||
pipe, tgt, mode=cfg["dit"], family=fam_name, logger=logger
|
||||
)
|
||||
_empty()
|
||||
dit_quant_active = engaged["dit"] is not None
|
||||
|
||||
# TE quant.
|
||||
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 (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 ──
|
||||
snapshot_backend_flags() # process-wide flags; benchmark process is short-lived so no restore.
|
||||
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).
|
||||
cache_active = False
|
||||
if cfg["cache"] == "auto":
|
||||
cache_request = TC_FBCACHE if default_steps >= FBCACHE_MIN_STEPS else None
|
||||
if cache_request is not None:
|
||||
engaged["cache"] = apply_step_cache(
|
||||
pipe, mode=cache_request, threshold=None,
|
||||
quant_active=dit_quant_active, logger=logger,
|
||||
)
|
||||
cache_active = engaged["cache"] not in (None, "off")
|
||||
|
||||
# Attention.
|
||||
backend = select_attention_backend(tgt, cfg["attn"], speed_active=speed_active)
|
||||
engaged["attn"] = apply_attention_backend(pipe, backend, logger=logger)
|
||||
|
||||
# Speed profile.
|
||||
if speed != "off":
|
||||
engaged["speed_optims"] = apply_speed_optims(
|
||||
pipe, tgt, is_gguf=False, family=fam_obj, speed_mode=speed,
|
||||
cache_active=cache_active, offload_active=False, logger=logger,
|
||||
)
|
||||
engaged["_effective_speed"] = speed
|
||||
return engaged
|
||||
|
||||
|
||||
def _timed_video(pipe, *, steps, width, height, num_frames, guidance, seed, cache_mode,
|
||||
dit_quant_active, default_steps, logger=None):
|
||||
"""One clip generation. Re-checks FBCache 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 maybe_toggle_step_cache, FBCACHE_MIN_STEPS
|
||||
|
||||
if cache_mode == "auto":
|
||||
try:
|
||||
maybe_toggle_step_cache(
|
||||
pipe, steps=steps, quant_active=dit_quant_active, threshold=None, logger=logger
|
||||
)
|
||||
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
|
||||
|
||||
_sync()
|
||||
t0 = time.perf_counter()
|
||||
out = pipe(
|
||||
prompt=PROMPT,
|
||||
width=width,
|
||||
height=height,
|
||||
num_frames=num_frames,
|
||||
num_inference_steps=steps,
|
||||
guidance_scale=guidance,
|
||||
generator=g,
|
||||
callback_on_step_end=_cb,
|
||||
)
|
||||
_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, 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)
|
||||
|
||||
_empty(); _reset_peak()
|
||||
pipe = _build_pipe(repo, force_fp32)
|
||||
load_peak = _peak_gb()
|
||||
engaged = _apply_levers(
|
||||
pipe, cfg, fam_name=family, fam_obj=fam_obj, force_fp32_vae=force_fp32,
|
||||
default_steps=default_steps, logger=logger,
|
||||
)
|
||||
_empty()
|
||||
weights_gb = _alloc_gb()
|
||||
dit_active = engaged["dit"] is not None
|
||||
cache_mode = cfg["cache"]
|
||||
|
||||
# warmup (pays the one-time compile / autotune)
|
||||
_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,
|
||||
logger=logger,
|
||||
)
|
||||
_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,
|
||||
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(out / "ref_frames.npz", *arrs)
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
row = {
|
||||
"config": name,
|
||||
"family": family,
|
||||
"levers": cfg,
|
||||
"dit_scheme": engaged["dit"] or "dense",
|
||||
"te_scheme": engaged["te"] or "dense",
|
||||
"vae_scheme": engaged["vae"] or "dense",
|
||||
"attn": engaged["attn"] or "native",
|
||||
"cache": engaged["cache"] or "off",
|
||||
"effective_speed": engaged.get("_effective_speed"),
|
||||
"speed_optims": engaged["speed_optims"],
|
||||
"load_peak_gb": round(load_peak, 2),
|
||||
"weights_gb": round(weights_gb, 2),
|
||||
"gen_peak_gb": round(gen_peak, 2),
|
||||
"gen_latency_s": round(_median(dts), 3),
|
||||
"per_step_ms": round(_median(steps_ms), 1),
|
||||
"n_frames": len(arrs),
|
||||
}
|
||||
del pipe
|
||||
_empty()
|
||||
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")
|
||||
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 = out / "ref_frames.npz"
|
||||
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
|
||||
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, 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)
|
||||
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)
|
||||
|
||||
# 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())
|
||||
Loading…
Add table
Add a link
Reference in a new issue