diff --git a/scripts/fp8_layer_ablation.py b/scripts/fp8_layer_ablation.py index 3862e251db..d4498040b3 100644 --- a/scripts/fp8_layer_ablation.py +++ b/scripts/fp8_layer_ablation.py @@ -67,7 +67,6 @@ from core.inference.diffusion_transformer_quant import ( # noqa: E402 # ── cuda helpers ──────────────────────────────────────────────────────────────── def _empty() -> None: import torch - gc.collect() if torch.cuda.is_available(): torch.cuda.empty_cache() @@ -121,8 +120,9 @@ class _Stop(Exception): pass -def _capture_forward_tuple(pipe, *, steps, width, height, num_frames, guidance, seed, - capture_call, gvg): +def _capture_forward_tuple( + pipe, *, steps, width, height, num_frames, guidance, seed, capture_call, gvg +): """Run a bf16 generation and grab the ``capture_call``-th transformer forward's (args, kwargs) to CPU, then abort. This is a REAL mid-schedule DiT input.""" import torch @@ -138,11 +138,15 @@ def _capture_forward_tuple(pipe, *, steps, width, height, num_frames, guidance, raise _Stop() return None - handle = pipe.transformer.register_forward_pre_hook(_pre, with_kwargs=True) - g = torch.Generator(device="cuda").manual_seed(seed) + handle = pipe.transformer.register_forward_pre_hook(_pre, with_kwargs = True) + g = torch.Generator(device = "cuda").manual_seed(seed) kwargs = dict( - prompt=PROMPT, width=width, height=height, num_frames=num_frames, - num_inference_steps=steps, generator=g, + prompt = PROMPT, + width = width, + height = height, + num_frames = num_frames, + num_inference_steps = steps, + generator = g, ) if gvg: guider = getattr(pipe, "guider", None) @@ -192,8 +196,12 @@ def _score(ref, cand) -> dict: a = canf[mask] b = reff[mask] cos = float(torch.nn.functional.cosine_similarity(a.unsqueeze(0), b.unsqueeze(0)).item()) - rel = float((torch.linalg.vector_norm(a - b) / (torch.linalg.vector_norm(b) + 1e-12)).item()) - norm_ratio = float((torch.linalg.vector_norm(a) / (torch.linalg.vector_norm(b) + 1e-12)).item()) + rel = float( + (torch.linalg.vector_norm(a - b) / (torch.linalg.vector_norm(b) + 1e-12)).item() + ) + norm_ratio = float( + (torch.linalg.vector_norm(a) / (torch.linalg.vector_norm(b) + 1e-12)).item() + ) else: cos, rel, norm_ratio = 0.0, float("inf"), 0.0 return { @@ -201,7 +209,9 @@ def _score(ref, cand) -> dict: "relL2": round(rel, 5), "frac_nonfinite": round(frac_nonfinite, 6), "norm_ratio": round(norm_ratio, 5), - "cand_amax": round(float(canf[fin_c].abs().max().item()) if fin_c.any() else float("inf"), 3), + "cand_amax": round( + float(canf[fin_c].abs().max().item()) if fin_c.any() else float("inf"), 3 + ), } @@ -217,12 +227,19 @@ def _verdict(sc: dict) -> str: # ── quant helpers ────────────────────────────────────────────────────────────── -def _fp8_filter(min_features: int, exclude_tokens=(), only_tokens=()): +def _fp8_filter( + min_features: int, + exclude_tokens = (), + only_tokens = (), +): """A production fp8 filter (require_bf16) optionally narrowed: - exclude_tokens: skip linears whose fqn contains any token (keep them bf16) - only_tokens : quantise ONLY linears whose fqn contains any token.""" - base = make_filter_fn(min_features, exclude_name_tokens=exclude_tokens, - require_bf16=(TQ_FP8 in _REQUIRE_BF16_SCHEMES)) + base = make_filter_fn( + min_features, + exclude_name_tokens = exclude_tokens, + require_bf16 = (TQ_FP8 in _REQUIRE_BF16_SCHEMES), + ) def filt(module, fqn: str = "") -> bool: if not base(module, fqn): @@ -237,16 +254,24 @@ def _fp8_filter(min_features: int, exclude_tokens=(), only_tokens=()): return filt -def _ablate(dense_cpu, tup, *, min_features, fast_accum=None, exclude_tokens=(), only_tokens=(), - instrument=False): +def _ablate( + dense_cpu, + tup, + *, + min_features, + fast_accum = None, + exclude_tokens = (), + only_tokens = (), + instrument = False, +): """Deepcopy the dense DiT -> GPU -> fp8-quantise the selected subset -> one forward. Returns (cpu_output_tensor, finiteness_records|None).""" import torch from torchao.quantization import quantize_ m = copy.deepcopy(dense_cpu).to("cuda") - filt = _fp8_filter(min_features, exclude_tokens=exclude_tokens, only_tokens=only_tokens) - quantize_(m, _make_quant_config(TQ_FP8, fast_accum=fast_accum), filter_fn=filt) + filt = _fp8_filter(min_features, exclude_tokens = exclude_tokens, only_tokens = only_tokens) + quantize_(m, _make_quant_config(TQ_FP8, fast_accum = fast_accum), filter_fn = filt) records = None hooks = [] @@ -255,7 +280,11 @@ def _ablate(dense_cpu, tup, *, min_features, fast_accum=None, exclude_tokens=(), def mk(nm): def hook(mod, inp, out): - o = out if torch.is_tensor(out) else (out[0] if isinstance(out, (list, tuple)) and out else None) + o = ( + out + if torch.is_tensor(out) + else (out[0] if isinstance(out, (list, tuple)) and out else None) + ) if o is None: return of = o.detach().float() @@ -265,15 +294,20 @@ def _ablate(dense_cpu, tup, *, min_features, fast_accum=None, exclude_tokens=(), xin = inp[0].detach().float() in_amax = float(xin.abs().max().item()) if xin.dim() >= 2: - ra = xin.reshape(-1, xin.shape[-1]).abs().amax(dim=-1) + ra = xin.reshape(-1, xin.shape[-1]).abs().amax(dim = -1) in_frac_zero_rows = round(float((ra < 1e-6).float().mean().item()), 4) - records.append({ - "name": nm, - "out_finite": bool(torch.isfinite(of).all().item()), - "out_amax": float(of.abs().max().item()) if torch.isfinite(of).any() else float("inf"), - "in_amax": in_amax, - "in_frac_zero_rows": in_frac_zero_rows, - }) + records.append( + { + "name": nm, + "out_finite": bool(torch.isfinite(of).all().item()), + "out_amax": float(of.abs().max().item()) + if torch.isfinite(of).any() + else float("inf"), + "in_amax": in_amax, + "in_frac_zero_rows": in_frac_zero_rows, + } + ) + return hook for nm, mod in m.named_modules(): @@ -332,7 +366,7 @@ def _zero_row_diag(tup: dict) -> dict: for key, val in tup.get("kwargs", {}).items(): if torch.is_tensor(val) and val.dim() >= 2 and val.is_floating_point(): x = val.float().reshape(-1, val.shape[-1]) - row_amax = x.abs().amax(dim=-1) + row_amax = x.abs().amax(dim = -1) n = int(row_amax.numel()) n_zero = int((row_amax < 1e-6).sum().item()) diag[key] = { @@ -360,16 +394,19 @@ def _bf16_reference_and_stats(dense_gpu, tup, *, min_features): return x = inp[0].detach().float() x2 = x.reshape(-1, x.shape[-1]).abs() # [tokens, channels] - amax_tok = x2.amax(dim=-1) - med_tok = x2.median(dim=-1).values + amax_tok = x2.amax(dim = -1) + med_tok = x2.median(dim = -1).values spread = float((amax_tok / (med_tok + 1e-9)).mean().item()) - stats.append({ - "name": nm, - "in_amax": float(x2.max().item()), - "spread": round(spread, 2), - "in_features": int(mod.in_features), - "out_features": int(mod.out_features), - }) + stats.append( + { + "name": nm, + "in_amax": float(x2.max().item()), + "spread": round(spread, 2), + "in_features": int(mod.in_features), + "out_features": int(mod.out_features), + } + ) + return hook for nm, mod in dense_gpu.named_modules(): @@ -380,31 +417,35 @@ def _bf16_reference_and_stats(dense_gpu, tup, *, min_features): ref = ref.detach().to("cpu") for h in hooks: h.remove() - stats.sort(key=lambda s: s["spread"], reverse=True) + stats.sort(key = lambda s: s["spread"], reverse = True) return ref, stats -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("--phase", type=int, default=0, choices=[0, 1, 2]) - ap.add_argument("--only", default="", help="phase 2: comma tokens -- fp8 ONLY these linears") - ap.add_argument("--exclude", default="", help="phase 2: comma tokens -- fp8 all EXCEPT these") - ap.add_argument("--steps", type=int, default=20) - 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("--capture-call", type=int, default=8, - help="which transformer forward call to capture (~mid schedule)") - ap.add_argument("--min-features", type=int, default=DEFAULT_MIN_LINEAR_FEATURES) - ap.add_argument("--out", default="outputs/fp8_ablation") +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("--phase", type = int, default = 0, choices = [0, 1, 2]) + ap.add_argument("--only", default = "", help = "phase 2: comma tokens -- fp8 ONLY these linears") + ap.add_argument("--exclude", default = "", help = "phase 2: comma tokens -- fp8 all EXCEPT these") + ap.add_argument("--steps", type = int, default = 20) + 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( + "--capture-call", + type = int, + default = 8, + help = "which transformer forward call to capture (~mid schedule)", + ) + ap.add_argument("--min-features", type = int, default = DEFAULT_MIN_LINEAR_FEATURES) + ap.add_argument("--out", default = "outputs/fp8_ablation") args = ap.parse_args(argv) import torch out = Path(args.out) - out.mkdir(parents=True, exist_ok=True) + out.mkdir(parents = True, exist_ok = True) spec = _FAMILIES[args.family] repo = spec["repo"] @@ -416,79 +457,111 @@ def main(argv=None) -> int: fam_obj = detect_video_family(repo) gvg = bool(getattr(fam_obj, "guidance_via_guider", False)) - print(f"== fp8 layer ablation: family={args.family} repo={repo} phase={args.phase} ==", flush=True) + print( + f"== fp8 layer ablation: family={args.family} repo={repo} phase={args.phase} ==", flush = True + ) t0 = time.perf_counter() pipe = _build_pipe(repo, force_fp32) - print(f"[load] pipe built in {time.perf_counter()-t0:.1f}s", flush=True) + print(f"[load] pipe built in {time.perf_counter()-t0:.1f}s", flush = True) tup = _capture_forward_tuple( - pipe, steps=args.steps, width=args.width, height=args.height, - num_frames=args.num_frames, guidance=guidance, seed=args.seed, - capture_call=args.capture_call, gvg=gvg, + pipe, + steps = args.steps, + width = args.width, + height = args.height, + num_frames = args.num_frames, + guidance = guidance, + seed = args.seed, + capture_call = args.capture_call, + gvg = gvg, ) - ks = {k: (tuple(v.shape) if torch.is_tensor(v) else type(v).__name__) - for k, v in tup.get("kwargs", {}).items()} - print(f"[capture] call #{args.capture_call} kwargs={ks}", flush=True) + ks = { + k: (tuple(v.shape) if torch.is_tensor(v) else type(v).__name__) + for k, v in tup.get("kwargs", {}).items() + } + print(f"[capture] call #{args.capture_call} kwargs={ks}", flush = True) # Keep the dense bf16 DiT (still on GPU) as reference source; copy to CPU as the ablation seed. dense_gpu = pipe.transformer - ref, stats = _bf16_reference_and_stats(dense_gpu, tup, min_features=args.min_features) - print(f"[bf16 ref] out shape={tuple(ref.shape)} amax={float(ref.float().abs().max()):.3f} " - f"n_quantized_linears={len(stats)}", flush=True) + ref, stats = _bf16_reference_and_stats(dense_gpu, tup, min_features = args.min_features) + print( + f"[bf16 ref] out shape={tuple(ref.shape)} amax={float(ref.float().abs().max()):.3f} " + f"n_quantized_linears={len(stats)}", + flush = True, + ) dense_cpu = copy.deepcopy(dense_gpu).to("cpu") # Free the pipeline (VAE/text encoders) to leave the GPU for one transformer at a time. del pipe, dense_gpu _empty() - report: dict[str, Any] = {"family": args.family, "repo": repo, "capture_call": args.capture_call, - "input_shapes": ks, "ref_amax": round(float(ref.float().abs().max()), 3), - "zero_row_diag": _zero_row_diag(tup)} - print(f"[zero-row diag] {report['zero_row_diag']}", flush=True) + report: dict[str, Any] = { + "family": args.family, + "repo": repo, + "capture_call": args.capture_call, + "input_shapes": ks, + "ref_amax": round(float(ref.float().abs().max()), 3), + "zero_row_diag": _zero_row_diag(tup), + } + print(f"[zero-row diag] {report['zero_row_diag']}", flush = True) if args.phase == 0: # ── Phase 0a: proxy validity -- fp8-ALL (production accum = None auto-detect) ── - out_all, _ = _ablate(dense_cpu, tup, min_features=args.min_features, fast_accum=None) + out_all, _ = _ablate(dense_cpu, tup, min_features = args.min_features, fast_accum = None) sc_all = _score(ref, out_all) report["fp8_all_autoaccum"] = {**sc_all, "verdict": _verdict(sc_all)} - print(f"[0a fp8-ALL auto-accum] {sc_all} -> {_verdict(sc_all)}", flush=True) + print(f"[0a fp8-ALL auto-accum] {sc_all} -> {_verdict(sc_all)}", flush = True) # ── Phase 0b-ii: accumulate flip ── for fa in (False, True): - o, _ = _ablate(dense_cpu, tup, min_features=args.min_features, fast_accum=fa) + o, _ = _ablate(dense_cpu, tup, min_features = args.min_features, fast_accum = fa) sc = _score(ref, o) report[f"fp8_all_fast_accum_{fa}"] = {**sc, "verdict": _verdict(sc)} - print(f"[0b accum={fa}] {sc} -> {_verdict(sc)}", flush=True) + print(f"[0b accum={fa}] {sc} -> {_verdict(sc)}", flush = True) # ── Phase 0b-i: per-Linear finiteness (instrumented fp8-ALL forward) ── - _, records = _ablate(dense_cpu, tup, min_features=args.min_features, fast_accum=None, - instrument=True) + _, records = _ablate( + dense_cpu, tup, min_features = args.min_features, fast_accum = None, instrument = True + ) if records: nonfinite = [r for r in records if not r["out_finite"]] report["first_nonfinite_linears"] = nonfinite[:10] report["n_nonfinite_linears"] = len(nonfinite) - finite_sorted = sorted([r for r in records if r["out_finite"]], - key=lambda r: r["out_amax"], reverse=True) + finite_sorted = sorted( + [r for r in records if r["out_finite"]], key = lambda r: r["out_amax"], reverse = True + ) report["top_out_amax_linears"] = finite_sorted[:10] - print(f"[0b finiteness] {len(nonfinite)}/{len(records)} linears non-finite; " - f"first={[r['name'] for r in nonfinite[:5]]}", flush=True) - print(f"[0b top out_amax] {[(r['name'], round(r['out_amax'],1)) for r in finite_sorted[:5]]}", - flush=True) + print( + f"[0b finiteness] {len(nonfinite)}/{len(records)} linears non-finite; " + f"first={[r['name'] for r in nonfinite[:5]]}", + flush = True, + ) + print( + f"[0b top out_amax] {[(r['name'], round(r['out_amax'],1)) for r in finite_sorted[:5]]}", + flush = True, + ) report["top_spread_linears"] = stats[:20] - print(f"[stats] top-spread linears: " - f"{[(s['name'], s['spread']) for s in stats[:8]]}", flush=True) + print( + f"[stats] top-spread linears: " f"{[(s['name'], s['spread']) for s in stats[:8]]}", + flush = True, + ) va = report["fp8_all_autoaccum"]["verdict"] if va.startswith("BROKEN"): - print("\n[GATE] fp8-ALL reproduces the black signal in one forward -> proxy VALID; " - "proceed to Phase 1 bucket ablation.", flush=True) + print( + "\n[GATE] fp8-ALL reproduces the black signal in one forward -> proxy VALID; " + "proceed to Phase 1 bucket ablation.", + flush = True, + ) elif va == "CLEAN": - print("\n[GATE] fp8-ALL is CLEAN in one forward but full gen is black -> failure is " - "multi-step/cache; single-forward bisection cannot localise -> STOP, keep int8.", - flush=True) + print( + "\n[GATE] fp8-ALL is CLEAN in one forward but full gen is black -> failure is " + "multi-step/cache; single-forward bisection cannot localise -> STOP, keep int8.", + flush = True, + ) else: - print(f"\n[GATE] fp8-ALL verdict={va}: borderline; inspect scores.", flush=True) + print(f"\n[GATE] fp8-ALL verdict={va}: borderline; inspect scores.", flush = True) elif args.phase == 1: # ── Phase 1: bucket ablation (necessity + sufficiency) ── kind = _dit_kind(args.family) @@ -496,53 +569,68 @@ def main(argv=None) -> int: results: dict[str, Any] = {} # baseline: fp8-ALL (should be BROKEN, matching Phase 0) - o, _ = _ablate(dense_cpu, tup, min_features=args.min_features) + o, _ = _ablate(dense_cpu, tup, min_features = args.min_features) results["fp8_all"] = {**_score(ref, o)} results["fp8_all"]["verdict"] = _verdict(results["fp8_all"]) - print(f"[1 baseline fp8-ALL] {results['fp8_all']['verdict']} {results['fp8_all']}", flush=True) + print( + f"[1 baseline fp8-ALL] {results['fp8_all']['verdict']} {results['fp8_all']}", flush = True + ) for bname, toks in buckets.items(): # NECESSITY: fp8 everything EXCEPT this bucket -- if CLEAN, this bucket is the culprit. - o, _ = _ablate(dense_cpu, tup, min_features=args.min_features, exclude_tokens=toks) + o, _ = _ablate(dense_cpu, tup, min_features = args.min_features, exclude_tokens = toks) sc_ex = _score(ref, o) # SUFFICIENCY: fp8 ONLY this bucket -- if BROKEN, this bucket alone reproduces damage. - o2, _ = _ablate(dense_cpu, tup, min_features=args.min_features, only_tokens=toks) + o2, _ = _ablate(dense_cpu, tup, min_features = args.min_features, only_tokens = toks) sc_only = _score(ref, o2) results[bname] = { "exclude": {**sc_ex, "verdict": _verdict(sc_ex)}, "only": {**sc_only, "verdict": _verdict(sc_only)}, } - print(f"[1 {bname}] EXCLUDE->{_verdict(sc_ex)} (cos {sc_ex['cosine']}, " - f"nf {sc_ex['frac_nonfinite']}) | ONLY->{_verdict(sc_only)} " - f"(cos {sc_only['cosine']}, nf {sc_only['frac_nonfinite']})", flush=True) + print( + f"[1 {bname}] EXCLUDE->{_verdict(sc_ex)} (cos {sc_ex['cosine']}, " + f"nf {sc_ex['frac_nonfinite']}) | ONLY->{_verdict(sc_only)} " + f"(cos {sc_only['cosine']}, nf {sc_only['frac_nonfinite']})", + flush = True, + ) report["phase1"] = results else: # ── Phase 2: instrument a specific only/exclude set, find first overflow ── only = tuple(t.strip() for t in args.only.split(",") if t.strip()) exclude = tuple(t.strip() for t in args.exclude.split(",") if t.strip()) - o, records = _ablate(dense_cpu, tup, min_features=args.min_features, - only_tokens=only, exclude_tokens=exclude, instrument=True) + o, records = _ablate( + dense_cpu, + tup, + min_features = args.min_features, + only_tokens = only, + exclude_tokens = exclude, + instrument = True, + ) sc = _score(ref, o) report["phase2"] = {"only": only, "exclude": exclude, **sc, "verdict": _verdict(sc)} - print(f"[2 only={only} exclude={exclude}] {sc} -> {_verdict(sc)}", flush=True) + print(f"[2 only={only} exclude={exclude}] {sc} -> {_verdict(sc)}", flush = True) if records: # first non-finite linears in execution order, with their input amax: an in_amax of # ~0 at the first-broken layer confirms a zero-amax (padding) input row -> scale 0 -> inf. nonfinite = [r for r in records if not r["out_finite"]] report["phase2_first_nonfinite"] = nonfinite[:15] report["phase2_n_nonfinite"] = len(nonfinite) - print(f"[2 finiteness] {len(nonfinite)}/{len(records)} quantised-path linears non-finite", - flush=True) + print( + f"[2 finiteness] {len(nonfinite)}/{len(records)} quantised-path linears non-finite", + flush = True, + ) for r in nonfinite[:8]: - print(f" first-inf {r['name']} in_amax={r['in_amax']} " - f"in_frac_zero_rows={r['in_frac_zero_rows']} out_amax={r['out_amax']}", - flush=True) + print( + f" first-inf {r['name']} in_amax={r['in_amax']} " + f"in_frac_zero_rows={r['in_frac_zero_rows']} out_amax={r['out_amax']}", + flush = True, + ) dest = out / f"phase{args.phase}_{args.family}.json" - with open(dest, "w", encoding="utf-8") as fh: - json.dump(report, fh, indent=2) - print(f"wrote {dest}", flush=True) + with open(dest, "w", encoding = "utf-8") as fh: + json.dump(report, fh, indent = 2) + print(f"wrote {dest}", flush = True) return 0 diff --git a/scripts/hunyuan_attn_diag.py b/scripts/hunyuan_attn_diag.py index 4a7210cfb7..eb8273204a 100644 --- a/scripts/hunyuan_attn_diag.py +++ b/scripts/hunyuan_attn_diag.py @@ -10,6 +10,7 @@ so we know whether the win is the N-reduction (trim) or the mask-elimination (nu Run: CUDA_VISIBLE_DEVICES=3 python scripts/hunyuan_attn_diag.py [--repo ...] [--frames 121] """ + from __future__ import annotations import argparse @@ -62,8 +63,8 @@ def _block_pre_hook(module, args, kwargs): CAP["dtype"] = hs.dtype if amask is not None: m = amask.bool() - CAP["text_valid_per_batch"] = m.sum(dim=1).tolist() - CAP["text_cols_valid_any"] = int(m.any(dim=0).sum()) # what our global-trim would keep + CAP["text_valid_per_batch"] = m.sum(dim = 1).tolist() + CAP["text_cols_valid_any"] = int(m.any(dim = 0).sum()) # what our global-trim would keep raise _StopCapture @@ -72,14 +73,16 @@ def _model_pre_hook(module, args, kwargs): def g(name): return kwargs.get(name) - for key, mkey in (("encoder_hidden_states", "encoder_attention_mask"), - ("encoder_hidden_states_2", "encoder_attention_mask_2")): + for key, mkey in ( + ("encoder_hidden_states", "encoder_attention_mask"), + ("encoder_hidden_states_2", "encoder_attention_mask_2"), + ): s = g(key) m = g(mkey) if s is not None: CAP.setdefault("streams", {})[key] = { "len": int(s.shape[1]), - "valid": (m.bool().sum(dim=1).tolist() if m is not None else None), + "valid": (m.bool().sum(dim = 1).tolist() if m is not None else None), } ie = g("image_embeds") if ie is not None: @@ -88,44 +91,52 @@ def _model_pre_hook(module, args, kwargs): return None -def _time_sdpa(q, k, v, mask, iters=30): +def _time_sdpa( + q, + k, + v, + mask, + iters = 30, +): # q,k,v: [B, H, N, D] torch.cuda.synchronize() for _ in range(3): # warmup - F.scaled_dot_product_attention(q, k, v, attn_mask=mask) + F.scaled_dot_product_attention(q, k, v, attn_mask = mask) torch.cuda.synchronize() t0 = time.perf_counter() for _ in range(iters): - F.scaled_dot_product_attention(q, k, v, attn_mask=mask) + F.scaled_dot_product_attention(q, k, v, attn_mask = mask) torch.cuda.synchronize() return (time.perf_counter() - t0) / iters * 1e3 # ms def main(): ap = argparse.ArgumentParser() - ap.add_argument("--repo", default="hunyuanvideo-community/HunyuanVideo-1.5-Diffusers-480p_t2v") - ap.add_argument("--frames", type=int, default=121) - ap.add_argument("--width", type=int, default=832) - ap.add_argument("--height", type=int, default=480) + ap.add_argument("--repo", default = "hunyuanvideo-community/HunyuanVideo-1.5-Diffusers-480p_t2v") + ap.add_argument("--frames", type = int, default = 121) + ap.add_argument("--width", type = int, default = 832) + ap.add_argument("--height", type = int, default = 480) args = ap.parse_args() diffusers = _import_diffusers() dev = "cuda:0" - print(f"loading {args.repo} ...", flush=True) - pipe = diffusers.DiffusionPipeline.from_pretrained(args.repo, torch_dtype=torch.bfloat16) + print(f"loading {args.repo} ...", flush = True) + pipe = diffusers.DiffusionPipeline.from_pretrained(args.repo, torch_dtype = torch.bfloat16) pipe = pipe.to(dev) - pipe.transformer.register_forward_pre_hook(_model_pre_hook, with_kwargs=True) - pipe.transformer.transformer_blocks[0].register_forward_pre_hook(_block_pre_hook, with_kwargs=True) + pipe.transformer.register_forward_pre_hook(_model_pre_hook, with_kwargs = True) + pipe.transformer.transformer_blocks[0].register_forward_pre_hook( + _block_pre_hook, with_kwargs = True + ) - print("running 1 capture step ...", flush=True) + print("running 1 capture step ...", flush = True) try: pipe( - prompt="a cat playing piano", - num_frames=args.frames, - width=args.width, - height=args.height, - num_inference_steps=1, + prompt = "a cat playing piano", + num_frames = args.frames, + width = args.width, + height = args.height, + num_inference_steps = 1, ) except _StopCapture: pass @@ -133,14 +144,24 @@ def main(): # the StopCapture may surface wrapped; if we captured, continue if "n_video" not in CAP: raise - print(f"(generation aborted after capture: {type(exc).__name__})", flush=True) + print(f"(generation aborted after capture: {type(exc).__name__})", flush = True) - print("\n===== CAPTURED SHAPES =====", flush=True) - for kk in ("batch", "n_video", "n_text", "heads", "dim_head", "dtype", - "text_valid_per_batch", "text_cols_valid_any", "image_embeds_len", - "image_is_t2v", "streams"): + print("\n===== CAPTURED SHAPES =====", flush = True) + for kk in ( + "batch", + "n_video", + "n_text", + "heads", + "dim_head", + "dtype", + "text_valid_per_batch", + "text_cols_valid_any", + "image_embeds_len", + "image_is_t2v", + "streams", + ): if kk in CAP: - print(f" {kk}: {CAP[kk]}", flush=True) + print(f" {kk}: {CAP[kk]}", flush = True) B = CAP["batch"] H = CAP["heads"] @@ -152,42 +173,62 @@ def main(): keep_text = CAP.get("text_cols_valid_any", n_text) N_trim = n_video + keep_text dtype = CAP["dtype"] - print(f"\n joint N = {N} (video {n_video} + text {n_text}); " - f"trimmed N = {N_trim} (text kept {keep_text})", flush=True) + print( + f"\n joint N = {N} (video {n_video} + text {n_text}); " + f"trimmed N = {N_trim} (text kept {keep_text})", + flush = True, + ) def mk(n): - return torch.randn(B, H, n, D, device=dev, dtype=dtype) + return torch.randn(B, H, n, D, device = dev, dtype = dtype) # (a) dense mask over full N (current). Build [B,1,N,N] bool (mostly True). - print("\n===== SDPA TIMING (ms/call, real shapes) =====", flush=True) + print("\n===== SDPA TIMING (ms/call, real shapes) =====", flush = True) q, k, v = mk(N), mk(N), mk(N) - dense = torch.ones(B, 1, N, N, dtype=torch.bool, device=dev) + dense = torch.ones(B, 1, N, N, dtype = torch.bool, device = dev) # emulate text padding: last (n_text - keep_text) columns invalid if n_text - keep_text > 0: - dense[:, :, :, n_video + keep_text:] = False - dense[:, :, n_video + keep_text:, :] = False + dense[:, :, :, n_video + keep_text :] = False + dense[:, :, n_video + keep_text :, :] = False t_dense = _time_sdpa(q, k, v, dense) mask_gb = dense.numel() / 1e9 - print(f" (a) dense [B,1,N,N] mask N={N:>6} : {t_dense:7.3f} ms (mask {mask_gb:.2f} GB)", flush=True) + print( + f" (a) dense [B,1,N,N] mask N={N:>6} : {t_dense:7.3f} ms (mask {mask_gb:.2f} GB)", + flush = True, + ) # (b) no mask over full N (upper bound of flash path if all valid) t_none = _time_sdpa(q, k, v, None) - print(f" (b) attn_mask=None N={N:>6} : {t_none:7.3f} ms ({t_dense/t_none:.2f}x vs a)", flush=True) + print( + f" (b) attn_mask=None N={N:>6} : {t_none:7.3f} ms ({t_dense/t_none:.2f}x vs a)", + flush = True, + ) # (c) trimmed N, dense all-True mask (text padding removed but mask still built) qt, kt, vt = mk(N_trim), mk(N_trim), mk(N_trim) - dense_t = torch.ones(B, 1, N_trim, N_trim, dtype=torch.bool, device=dev) + dense_t = torch.ones(B, 1, N_trim, N_trim, dtype = torch.bool, device = dev) t_dense_trim = _time_sdpa(qt, kt, vt, dense_t) - print(f" (c) dense mask @trimmed N={N_trim:>6} : {t_dense_trim:7.3f} ms ({t_dense/t_dense_trim:.2f}x vs a)", flush=True) + print( + f" (c) dense mask @trimmed N={N_trim:>6} : {t_dense_trim:7.3f} ms ({t_dense/t_dense_trim:.2f}x vs a)", + flush = True, + ) # (d) trimmed N, no mask (trim + null: the full proposed fast path) t_none_trim = _time_sdpa(qt, kt, vt, None) - print(f" (d) no mask @trimmed N={N_trim:>6} : {t_none_trim:7.3f} ms ({t_dense/t_none_trim:.2f}x vs a)", flush=True) + print( + f" (d) no mask @trimmed N={N_trim:>6} : {t_none_trim:7.3f} ms ({t_dense/t_none_trim:.2f}x vs a)", + flush = True, + ) - print("\n Interpretation:", flush=True) - print(f" trim-only ceiling (a->c): {(1-t_dense_trim/t_dense)*100:5.1f}% attn saving", flush=True) - print(f" null-only ceiling (a->b): {(1-t_none/t_dense)*100:5.1f}% attn saving", flush=True) - print(f" trim+null (a->d): {(1-t_none_trim/t_dense)*100:5.1f}% attn saving", flush=True) + print("\n Interpretation:", flush = True) + print( + f" trim-only ceiling (a->c): {(1-t_dense_trim/t_dense)*100:5.1f}% attn saving", + flush = True, + ) + print(f" null-only ceiling (a->b): {(1-t_none/t_dense)*100:5.1f}% attn saving", flush = True) + print( + f" trim+null (a->d): {(1-t_none_trim/t_dense)*100:5.1f}% attn saving", flush = True + ) if __name__ == "__main__": diff --git a/scripts/hunyuan_int8_profile.py b/scripts/hunyuan_int8_profile.py index b840007d79..5853af1567 100644 --- a/scripts/hunyuan_int8_profile.py +++ b/scripts/hunyuan_int8_profile.py @@ -38,7 +38,12 @@ def _dynamo_counts(): """(recompiles, graph_breaks, unique_compiles) from dynamo counters, best-effort.""" try: from torch._dynamo.utils import counters - rc = sum(v for k, v in counters.get("recompiles", {}).items()) if "recompiles" in counters else 0 + + rc = ( + sum(v for k, v in counters.get("recompiles", {}).items()) + if "recompiles" in counters + else 0 + ) gb = sum(counters.get("graph_break", {}).values()) # total frames compiled stats = counters.get("stats", {}) @@ -48,8 +53,9 @@ def _dynamo_counts(): return -1, -1, -1 -def _profile_mode(mode: str, *, steps: int, width: int, height: int, num_frames: int, guidance: float, - seed: int): +def _profile_mode( + mode: str, *, steps: int, width: int, height: int, num_frames: int, guidance: float, seed: int +): import torch from core.inference.video_families import detect_video_family from core.inference.diffusion_cache import maybe_toggle_step_cache # noqa: F401 @@ -67,39 +73,59 @@ def _profile_mode(mode: str, *, steps: int, width: int, height: int, num_frames: default_steps = getattr(fam_obj, "default_steps", 50) pipe = _build_pipe(_REPO, False) - cfg = dict(te="none", vae="none", dit=("int8" if mode == "int8" else "none"), - speed="default", attn="native", cache="off") - engaged = _apply_levers(pipe, cfg, fam_name="hunyuanvideo-1.5", fam_obj=fam_obj, - force_fp32_vae=False, default_steps=default_steps) - print(f"[{mode}] dit_scheme={engaged['dit'] or 'dense'} speed={engaged.get('_effective_speed')}", - flush=True) + cfg = dict( + te = "none", + vae = "none", + dit = ("int8" if mode == "int8" else "none"), + speed = "default", + attn = "native", + cache = "off", + ) + engaged = _apply_levers( + pipe, + cfg, + fam_name = "hunyuanvideo-1.5", + fam_obj = fam_obj, + force_fp32_vae = False, + default_steps = default_steps, + ) + print( + f"[{mode}] dit_scheme={engaged['dit'] or 'dense'} speed={engaged.get('_effective_speed')}", + flush = True, + ) # per-forward GPU timing via cuda events on the transformer. fwd_ms: list[float] = [] starts: list = [] def _pre(mod, args, kwargs): - ev = torch.cuda.Event(enable_timing=True) + ev = torch.cuda.Event(enable_timing = True) ev.record() starts.append(ev) return None def _post(mod, args, output): - end = torch.cuda.Event(enable_timing=True) + end = torch.cuda.Event(enable_timing = True) end.record() torch.cuda.synchronize() if starts: fwd_ms.append(starts[-1].elapsed_time(end)) - h1 = pipe.transformer.register_forward_pre_hook(_pre, with_kwargs=True) + h1 = pipe.transformer.register_forward_pre_hook(_pre, with_kwargs = True) h2 = pipe.transformer.register_forward_hook(_post) def _gen(tag): fwd_ms.clear() starts.clear() - g = torch.Generator(device="cuda").manual_seed(seed) - kwargs = dict(prompt=PROMPT, width=width, height=height, num_frames=num_frames, - num_inference_steps=steps, generator=g) + g = torch.Generator(device = "cuda").manual_seed(seed) + kwargs = dict( + prompt = PROMPT, + width = width, + height = height, + num_frames = num_frames, + num_inference_steps = steps, + generator = g, + ) if gvg: guider = getattr(pipe, "guider", None) if guider is not None and hasattr(guider, "guidance_scale"): @@ -118,36 +144,47 @@ def _profile_mode(mode: str, *, steps: int, width: int, height: int, num_frames: p10 = srt[max(0, n // 10)] if n else 0.0 p90 = srt[min(n - 1, (9 * n) // 10)] if n else 0.0 slow = sum(1 for x in fwd_ms if x > 2.0 * med) if med else 0 - print(f"[{mode}:{tag}] total={total:.2f}s n_fwd={n} med={med:.1f}ms " - f"p10={p10:.1f} p90={p90:.1f} max={max(fwd_ms) if fwd_ms else 0:.1f} " - f"slow(>2x med)={slow} | recompiles={rc} graph_breaks={gb} unique_graphs={uc}", - flush=True) + print( + f"[{mode}:{tag}] total={total:.2f}s n_fwd={n} med={med:.1f}ms " + f"p10={p10:.1f} p90={p90:.1f} max={max(fwd_ms) if fwd_ms else 0:.1f} " + f"slow(>2x med)={slow} | recompiles={rc} graph_breaks={gb} unique_graphs={uc}", + flush = True, + ) return total _gen("warmup") # pays compile _gen("timed") _gen("timed2") - h1.remove(); h2.remove() + h1.remove() + h2.remove() del pipe import gc + gc.collect() torch.cuda.empty_cache() -def main(argv=None): +def main(argv = None): ap = argparse.ArgumentParser() - ap.add_argument("--modes", default="dense,int8") - 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("--guidance", type=float, default=6.0) - ap.add_argument("--seed", type=int, default=42) + ap.add_argument("--modes", default = "dense,int8") + 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("--guidance", type = float, default = 6.0) + ap.add_argument("--seed", type = int, default = 42) args = ap.parse_args(argv) for mode in [m.strip() for m in args.modes.split(",") if m.strip()]: - _profile_mode(mode, steps=args.steps, width=args.width, height=args.height, - num_frames=args.num_frames, guidance=args.guidance, seed=args.seed) + _profile_mode( + mode, + steps = args.steps, + width = args.width, + height = args.height, + num_frames = args.num_frames, + guidance = args.guidance, + seed = args.seed, + ) return 0 diff --git a/scripts/hunyuan_trim_e2e.py b/scripts/hunyuan_trim_e2e.py index 2992e584ca..8899e9853c 100644 --- a/scripts/hunyuan_trim_e2e.py +++ b/scripts/hunyuan_trim_e2e.py @@ -7,6 +7,7 @@ wall-clock. Stock is only run at MODEST settings (a full 121-frame stock gen is Run: CUDA_VISIBLE_DEVICES=3 python scripts/hunyuan_trim_e2e.py """ + from __future__ import annotations import argparse @@ -23,7 +24,7 @@ import torch _REPO_ROOT = Path(__file__).resolve().parent.parent sys.path.insert(0, str(_REPO_ROOT / "studio" / "backend")) OUT = _REPO_ROOT / "outputs" / "video_speedmem" -OUT.mkdir(parents=True, exist_ok=True) +OUT.mkdir(parents = True, exist_ok = True) def _import_diffusers(): @@ -36,12 +37,18 @@ def _import_diffusers(): def _gen(pipe, seed, frames, steps, w, h): - g = torch.Generator(device="cuda").manual_seed(seed) + g = torch.Generator(device = "cuda").manual_seed(seed) torch.cuda.synchronize() t0 = time.perf_counter() - out = pipe(prompt="a cat playing piano on a stage, cinematic", - num_frames=frames, width=w, height=h, - num_inference_steps=steps, generator=g, output_type="np") + out = pipe( + prompt = "a cat playing piano on a stage, cinematic", + num_frames = frames, + width = w, + height = h, + num_inference_steps = steps, + generator = g, + output_type = "np", + ) torch.cuda.synchronize() wall = (time.perf_counter() - t0) * 1e3 frames_np = out.frames[0] # [F,H,W,C] in [0,1] @@ -54,7 +61,12 @@ def _luma(frames): return float((0.299 * r + 0.587 * gg + 0.114 * b).mean()) -def _lpips_mean(loss_fn, a, b, stride=4): +def _lpips_mean( + loss_fn, + a, + b, + stride = 4, +): vals = [] for i in range(0, len(a), stride): ta = torch.from_numpy(a[i]).permute(2, 0, 1).unsqueeze(0).float() * 2 - 1 @@ -66,14 +78,14 @@ def _lpips_mean(loss_fn, a, b, stride=4): def main(): ap = argparse.ArgumentParser() - ap.add_argument("--repo", default="hunyuanvideo-community/HunyuanVideo-1.5-Diffusers-480p_t2v") - ap.add_argument("--cmp-frames", type=int, default=25) - ap.add_argument("--cmp-steps", type=int, default=50) - ap.add_argument("--cmp-w", type=int, default=832) - ap.add_argument("--cmp-h", type=int, default=480) - ap.add_argument("--full-frames", type=int, default=121) - ap.add_argument("--full-steps", type=int, default=50) - ap.add_argument("--seed", type=int, default=1234) + ap.add_argument("--repo", default = "hunyuanvideo-community/HunyuanVideo-1.5-Diffusers-480p_t2v") + ap.add_argument("--cmp-frames", type = int, default = 25) + ap.add_argument("--cmp-steps", type = int, default = 50) + ap.add_argument("--cmp-w", type = int, default = 832) + ap.add_argument("--cmp-h", type = int, default = 480) + ap.add_argument("--full-frames", type = int, default = 121) + ap.add_argument("--full-steps", type = int, default = 50) + ap.add_argument("--seed", type = int, default = 1234) args = ap.parse_args() diffusers = _import_diffusers() @@ -81,43 +93,61 @@ def main(): from core.inference.video_families import detect_video_family import lpips - print(f"loading {args.repo} ...", flush=True) - pipe = diffusers.DiffusionPipeline.from_pretrained(args.repo, torch_dtype=torch.bfloat16).to("cuda") + print(f"loading {args.repo} ...", flush = True) + pipe = diffusers.DiffusionPipeline.from_pretrained(args.repo, torch_dtype = torch.bfloat16).to( + "cuda" + ) fam = detect_video_family(args.repo) or detect_video_family("hunyuanvideo-1.5") - loss_fn = lpips.LPIPS(net="alex").cuda() + loss_fn = lpips.LPIPS(net = "alex").cuda() fr, st, w, hh = args.cmp_frames, args.cmp_steps, args.cmp_w, args.cmp_h # ---- STOCK x2 (nondeterminism floor) ---- - print(f"\n[stock#1] gen {fr}f/{st}steps {w}x{hh} ...", flush=True) + print(f"\n[stock#1] gen {fr}f/{st}steps {w}x{hh} ...", flush = True) stock1, w_stock = _gen(pipe, args.seed, fr, st, w, hh) - print(f"[stock#1] wall={w_stock:.0f} ms luma={_luma(stock1):.4f} frames={stock1.shape}", flush=True) - print(f"[stock#2] gen (same seed, measures run-to-run nondeterminism) ...", flush=True) + print( + f"[stock#1] wall={w_stock:.0f} ms luma={_luma(stock1):.4f} frames={stock1.shape}", + flush = True, + ) + print(f"[stock#2] gen (same seed, measures run-to-run nondeterminism) ...", flush = True) stock2, _ = _gen(pipe, args.seed, fr, st, w, hh) # ---- TRIM (same seed) ---- - engaged = install_hunyuan_attention_trim(pipe, fam, logger=None) - print(f"\ninstall_hunyuan_attention_trim engaged={engaged}", flush=True) + engaged = install_hunyuan_attention_trim(pipe, fam, logger = None) + print(f"\ninstall_hunyuan_attention_trim engaged={engaged}", flush = True) trim, w_trim = _gen(pipe, args.seed, fr, st, w, hh) - print(f"[trim ] wall={w_trim:.0f} ms luma={_luma(trim):.4f} ({w_stock/w_trim:.2f}x vs stock)", flush=True) + print( + f"[trim ] wall={w_trim:.0f} ms luma={_luma(trim):.4f} ({w_stock/w_trim:.2f}x vs stock)", + flush = True, + ) floor = _lpips_mean(loss_fn, stock1, stock2) lp = _lpips_mean(loss_fn, stock1, trim) - print(f"\nLPIPS(stock#1, stock#2) = {floor:.5f} <- nondeterminism floor", flush=True) - print(f"LPIPS(stock#1, trim ) = {lp:.5f} <- trim vs stock", flush=True) - print(f" => trim is {'WITHIN' if lp <= floor * 1.5 + 0.002 else 'ABOVE'} the nondeterminism floor", flush=True) + print(f"\nLPIPS(stock#1, stock#2) = {floor:.5f} <- nondeterminism floor", flush = True) + print(f"LPIPS(stock#1, trim ) = {lp:.5f} <- trim vs stock", flush = True) + print( + f" => trim is {'WITHIN' if lp <= floor * 1.5 + 0.002 else 'ABOVE'} the nondeterminism floor", + flush = True, + ) # ---- FULL-RES TRIM (wall-clock + black-frame guard; stock at full-res is ~19 min, skipped) ---- - print(f"\n[trim-full] gen {args.full_frames}f/{args.full_steps}steps {args.cmp_w}x{args.cmp_h} ...", flush=True) + print( + f"\n[trim-full] gen {args.full_frames}f/{args.full_steps}steps {args.cmp_w}x{args.cmp_h} ...", + flush = True, + ) full, w_full = _gen(pipe, args.seed, args.full_frames, args.full_steps, args.cmp_w, args.cmp_h) - print(f"[trim-full] wall={w_full/1000:.1f} s luma={_luma(full):.4f} frames={full.shape}", flush=True) + print( + f"[trim-full] wall={w_full/1000:.1f} s luma={_luma(full):.4f} frames={full.shape}", + flush = True, + ) try: from PIL import Image + Image.fromarray((trim[0] * 255).astype("uint8")).save(OUT / "vid_hunyuan_trim_cmp.png") Image.fromarray((full[0] * 255).astype("uint8")).save(OUT / "vid_hunyuan_trim_full.png") - print(f"\nsaved sample frames to {OUT}", flush=True) + print(f"\nsaved sample frames to {OUT}", flush = True) except Exception as exc: # noqa: BLE001 - print(f"(png save skipped: {exc})", flush=True) + print(f"(png save skipped: {exc})", flush = True) if __name__ == "__main__": diff --git a/scripts/hunyuan_trim_fp32ref.py b/scripts/hunyuan_trim_fp32ref.py index 87372e656b..cedf3a26af 100644 --- a/scripts/hunyuan_trim_fp32ref.py +++ b/scripts/hunyuan_trim_fp32ref.py @@ -8,6 +8,7 @@ quality loss. If trim is clearly farther from fp32 than stock, it is a real regr Run: CUDA_VISIBLE_DEVICES=1 python scripts/hunyuan_trim_fp32ref.py """ + from __future__ import annotations import argparse @@ -35,14 +36,25 @@ def _import_diffusers(): def _gen(pipe, seed, fr, st, w, h): - g = torch.Generator(device="cuda").manual_seed(seed) - out = pipe(prompt="a cat playing piano on a stage, cinematic", - num_frames=fr, width=w, height=h, num_inference_steps=st, - generator=g, output_type="np") + g = torch.Generator(device = "cuda").manual_seed(seed) + out = pipe( + prompt = "a cat playing piano on a stage, cinematic", + num_frames = fr, + width = w, + height = h, + num_inference_steps = st, + generator = g, + output_type = "np", + ) return np.asarray(out.frames[0]) -def _lpips_mean(loss_fn, a, b, stride=3): +def _lpips_mean( + loss_fn, + a, + b, + stride = 3, +): vals = [] for i in range(0, len(a), stride): ta = torch.from_numpy(a[i]).permute(2, 0, 1).unsqueeze(0).float().cuda() * 2 - 1 @@ -54,12 +66,12 @@ def _lpips_mean(loss_fn, a, b, stride=3): def main(): ap = argparse.ArgumentParser() - ap.add_argument("--repo", default="hunyuanvideo-community/HunyuanVideo-1.5-Diffusers-480p_t2v") - ap.add_argument("--frames", type=int, default=25) - ap.add_argument("--steps", type=int, default=30) - ap.add_argument("--w", type=int, default=832) - ap.add_argument("--h", type=int, default=480) - ap.add_argument("--seed", type=int, default=1234) + ap.add_argument("--repo", default = "hunyuanvideo-community/HunyuanVideo-1.5-Diffusers-480p_t2v") + ap.add_argument("--frames", type = int, default = 25) + ap.add_argument("--steps", type = int, default = 30) + ap.add_argument("--w", type = int, default = 832) + ap.add_argument("--h", type = int, default = 480) + ap.add_argument("--seed", type = int, default = 1234) args = ap.parse_args() diffusers = _import_diffusers() @@ -68,37 +80,45 @@ def main(): import lpips fam = detect_video_family(args.repo) or detect_video_family("hunyuanvideo-1.5") - loss_fn = lpips.LPIPS(net="alex").cuda() + loss_fn = lpips.LPIPS(net = "alex").cuda() fr, st, w, h, seed = args.frames, args.steps, args.w, args.h, args.seed # ---- bf16 stock + bf16 trim on one pipe ---- - print(f"loading bf16 {args.repo} ...", flush=True) - pipe = diffusers.DiffusionPipeline.from_pretrained(args.repo, torch_dtype=torch.bfloat16).to("cuda") - print(f"[bf16 stock] gen {fr}f/{st}steps ...", flush=True) + print(f"loading bf16 {args.repo} ...", flush = True) + pipe = diffusers.DiffusionPipeline.from_pretrained(args.repo, torch_dtype = torch.bfloat16).to( + "cuda" + ) + print(f"[bf16 stock] gen {fr}f/{st}steps ...", flush = True) stock = _gen(pipe, seed, fr, st, w, h) - install_hunyuan_attention_trim(pipe, fam, logger=None) - print("[bf16 trim ] gen ...", flush=True) + install_hunyuan_attention_trim(pipe, fam, logger = None) + print("[bf16 trim ] gen ...", flush = True) trim = _gen(pipe, seed, fr, st, w, h) del pipe - gc.collect(); torch.cuda.empty_cache() + gc.collect() + torch.cuda.empty_cache() # ---- fp32 reference (stock masked attention, upcast) ---- - print("loading fp32 reference ...", flush=True) - pipe32 = diffusers.DiffusionPipeline.from_pretrained(args.repo, torch_dtype=torch.float32).to("cuda") - print(f"[fp32 gold ] gen {fr}f/{st}steps ...", flush=True) + print("loading fp32 reference ...", flush = True) + pipe32 = diffusers.DiffusionPipeline.from_pretrained(args.repo, torch_dtype = torch.float32).to( + "cuda" + ) + print(f"[fp32 gold ] gen {fr}f/{st}steps ...", flush = True) gold = _gen(pipe32, seed, fr, st, w, h) d_stock = _lpips_mean(loss_fn, gold, stock) d_trim = _lpips_mean(loss_fn, gold, trim) d_st = _lpips_mean(loss_fn, stock, trim) - print("\n===== ACCURACY vs fp32 reference =====", flush=True) - print(f" LPIPS(fp32, bf16-stock) = {d_stock:.5f}", flush=True) - print(f" LPIPS(fp32, bf16-trim ) = {d_trim:.5f}", flush=True) - print(f" LPIPS(bf16-stock, trim) = {d_st:.5f}", flush=True) + print("\n===== ACCURACY vs fp32 reference =====", flush = True) + print(f" LPIPS(fp32, bf16-stock) = {d_stock:.5f}", flush = True) + print(f" LPIPS(fp32, bf16-trim ) = {d_trim:.5f}", flush = True) + print(f" LPIPS(bf16-stock, trim) = {d_st:.5f}", flush = True) if d_stock is not None and d_trim is not None: - verdict = "NOT less accurate (trim ~= stock vs fp32)" if d_trim <= d_stock * 1.25 + 0.01 \ + verdict = ( + "NOT less accurate (trim ~= stock vs fp32)" + if d_trim <= d_stock * 1.25 + 0.01 else "LESS accurate (trim farther from fp32 than stock)" - print(f"\n VERDICT: {verdict}", flush=True) + ) + print(f"\n VERDICT: {verdict}", flush = True) if __name__ == "__main__": diff --git a/scripts/hunyuan_trim_validate.py b/scripts/hunyuan_trim_validate.py index 97936ce071..4c2a364e2a 100644 --- a/scripts/hunyuan_trim_validate.py +++ b/scripts/hunyuan_trim_validate.py @@ -7,6 +7,7 @@ real production shape (default 121 frames / 480p). Run: CUDA_VISIBLE_DEVICES=3 python scripts/hunyuan_trim_validate.py """ + from __future__ import annotations import argparse @@ -45,14 +46,18 @@ def _capture_hook(module, args, kwargs): raise _Stop -def _forward(transformer, no_grad=True): +def _forward(transformer, no_grad = True): ctx = torch.no_grad() if no_grad else torch.enable_grad() with ctx: out = transformer(*CAP["args"], **CAP["kwargs"]) return out[0] if isinstance(out, tuple) else out.sample -def _median_ms(fn, iters=8, warmup=2): +def _median_ms( + fn, + iters = 8, + warmup = 2, +): for _ in range(warmup): fn() torch.cuda.synchronize() @@ -70,7 +75,7 @@ def _median_ms(fn, iters=8, warmup=2): def _compare(a, b): a = a.float().flatten() b = b.float().flatten() - cos = torch.nn.functional.cosine_similarity(a, b, dim=0).item() + cos = torch.nn.functional.cosine_similarity(a, b, dim = 0).item() max_abs = (a - b).abs().max().item() denom = a.abs().max().item() or 1.0 return cos, max_abs, max_abs / denom @@ -78,11 +83,11 @@ def _compare(a, b): def main(): ap = argparse.ArgumentParser() - ap.add_argument("--repo", default="hunyuanvideo-community/HunyuanVideo-1.5-Diffusers-480p_t2v") - ap.add_argument("--frames", type=int, default=121) - ap.add_argument("--width", type=int, default=832) - ap.add_argument("--height", type=int, default=480) - ap.add_argument("--compile", action="store_true", help="also test regional compile of blocks") + ap.add_argument("--repo", default = "hunyuanvideo-community/HunyuanVideo-1.5-Diffusers-480p_t2v") + ap.add_argument("--frames", type = int, default = 121) + ap.add_argument("--width", type = int, default = 832) + ap.add_argument("--height", type = int, default = 480) + ap.add_argument("--compile", action = "store_true", help = "also test regional compile of blocks") args = ap.parse_args() diffusers = _import_diffusers() @@ -90,16 +95,26 @@ def main(): from core.inference.video_families import detect_video_family dev = "cuda:0" - print(f"loading {args.repo} ...", flush=True) - pipe = diffusers.DiffusionPipeline.from_pretrained(args.repo, torch_dtype=torch.bfloat16).to(dev) + print(f"loading {args.repo} ...", flush = True) + pipe = diffusers.DiffusionPipeline.from_pretrained(args.repo, torch_dtype = torch.bfloat16).to( + dev + ) fam = detect_video_family(args.repo) or detect_video_family("hunyuanvideo-1.5") - print(f"family: {getattr(fam, 'name', None)} / transformer_class={getattr(fam, 'transformer_class', None)}", flush=True) + print( + f"family: {getattr(fam, 'name', None)} / transformer_class={getattr(fam, 'transformer_class', None)}", + flush = True, + ) - h = pipe.transformer.register_forward_pre_hook(_capture_hook, with_kwargs=True) - print("capturing one real forward input ...", flush=True) + h = pipe.transformer.register_forward_pre_hook(_capture_hook, with_kwargs = True) + print("capturing one real forward input ...", flush = True) try: - pipe(prompt="a cat playing piano", num_frames=args.frames, - width=args.width, height=args.height, num_inference_steps=1) + pipe( + prompt = "a cat playing piano", + num_frames = args.frames, + width = args.width, + height = args.height, + num_inference_steps = 1, + ) except _Stop: pass h.remove() @@ -108,46 +123,58 @@ def main(): ehs = k.get("encoder_hidden_states") m = k.get("encoder_attention_mask") ie = k.get("image_embeds") - print(f"\ncaptured: encoder_hidden_states={tuple(ehs.shape)} mask_valid={m.bool().sum(1).tolist()}" - f" image_embeds={tuple(ie.shape) if ie is not None else None}" - f" image_all_zero={bool(torch.all(ie==0).item()) if ie is not None else None}", flush=True) + print( + f"\ncaptured: encoder_hidden_states={tuple(ehs.shape)} mask_valid={m.bool().sum(1).tolist()}" + f" image_embeds={tuple(ie.shape) if ie is not None else None}" + f" image_all_zero={bool(torch.all(ie==0).item()) if ie is not None else None}", + flush = True, + ) # ---- STOCK forward (reference) + timing ---- transformer = pipe.transformer out_stock = _forward(transformer).detach().clone() t_stock = _median_ms(lambda: _forward(transformer)) - print(f"\nSTOCK forward: {t_stock:8.2f} ms out={tuple(out_stock.shape)}", flush=True) + print(f"\nSTOCK forward: {t_stock:8.2f} ms out={tuple(out_stock.shape)}", flush = True) # ---- install trim, re-run same inputs ---- - engaged = install_hunyuan_attention_trim(pipe, fam, logger=None) - print(f"install_hunyuan_attention_trim engaged = {engaged}", flush=True) + engaged = install_hunyuan_attention_trim(pipe, fam, logger = None) + print(f"install_hunyuan_attention_trim engaged = {engaged}", flush = True) out_trim = _forward(transformer).detach().clone() t_trim = _median_ms(lambda: _forward(transformer)) cos, max_abs, rel = _compare(out_stock, out_trim) - print(f"TRIM forward: {t_trim:8.2f} ms ({t_stock/t_trim:.2f}x faster)", flush=True) - print(f"\nACCURACY stock-vs-trim: cosine={cos:.8f} max_abs={max_abs:.4e} rel_max={rel:.4e}", flush=True) + print(f"TRIM forward: {t_trim:8.2f} ms ({t_stock/t_trim:.2f}x faster)", flush = True) + print( + f"\nACCURACY stock-vs-trim: cosine={cos:.8f} max_abs={max_abs:.4e} rel_max={rel:.4e}", + flush = True, + ) finite = bool(torch.isfinite(out_trim).all().item()) - print(f"trim output finite: {finite}", flush=True) + print(f"trim output finite: {finite}", flush = True) if args.compile: - print("\ncompiling blocks (compile_repeated_blocks, mode=default, dynamic=True) ...", flush=True) + print( + "\ncompiling blocks (compile_repeated_blocks, mode=default, dynamic=True) ...", + flush = True, + ) try: for _a in ("recompile_limit", "cache_size_limit"): if hasattr(torch._dynamo.config, _a): setattr(torch._dynamo.config, _a, 64) - transformer.compile_repeated_blocks(fullgraph=False, dynamic=True) + transformer.compile_repeated_blocks(fullgraph = False, dynamic = True) out_c = _forward(transformer).detach().clone() # triggers compile - t_c = _median_ms(lambda: _forward(transformer), iters=5, warmup=1) + t_c = _median_ms(lambda: _forward(transformer), iters = 5, warmup = 1) cos_c, ma_c, rel_c = _compare(out_stock, out_c) cnt = torch._dynamo.utils.counters - print(f"TRIM+COMPILE forward: {t_c:8.2f} ms ({t_stock/t_c:.2f}x vs stock)", flush=True) - print(f" accuracy vs stock: cosine={cos_c:.8f} max_abs={ma_c:.4e}", flush=True) - print(f" dynamo recompiles={sum(cnt['recompiles'].values()) if 'recompiles' in cnt else '?'}" - f" graph_breaks={sum(cnt['graph_break'].values()) if 'graph_break' in cnt else 0}", flush=True) + print(f"TRIM+COMPILE forward: {t_c:8.2f} ms ({t_stock/t_c:.2f}x vs stock)", flush = True) + print(f" accuracy vs stock: cosine={cos_c:.8f} max_abs={ma_c:.4e}", flush = True) + print( + f" dynamo recompiles={sum(cnt['recompiles'].values()) if 'recompiles' in cnt else '?'}" + f" graph_breaks={sum(cnt['graph_break'].values()) if 'graph_break' in cnt else 0}", + flush = True, + ) except Exception as exc: # noqa: BLE001 import traceback - print(f"COMPILE FAILED: {type(exc).__name__}: {exc}", flush=True) + print(f"COMPILE FAILED: {type(exc).__name__}: {exc}", flush = True) traceback.print_exc() diff --git a/scripts/quant_accuracy_sweep.py b/scripts/quant_accuracy_sweep.py index 2416597b80..2a3f1484de 100644 --- a/scripts/quant_accuracy_sweep.py +++ b/scripts/quant_accuracy_sweep.py @@ -71,14 +71,17 @@ _VAE_FAMILIES: dict[str, dict[str, Any]] = { def _check_deps() -> None: import importlib.util as ilu - - missing = [m for m in ("torch", "torchao", "diffusers", "lpips", "numpy", "PIL") if not ilu.find_spec(m)] + missing = [ + m + for m in ("torch", "torchao", "diffusers", "lpips", "numpy", "PIL") + if not ilu.find_spec(m) + ] if missing: print( "missing deps: " + ", ".join(missing) + "\n" " uv pip install torch torchao diffusers lpips numpy pillow", - file=sys.stderr, - flush=True, + file = sys.stderr, + flush = True, ) raise SystemExit(2) @@ -102,7 +105,7 @@ def _load_vae(repo: str, subfolder: str, device: str): import torch diffusers = _import_diffusers() - vae = diffusers.AutoModel.from_pretrained(repo, subfolder=subfolder, torch_dtype=torch.bfloat16) + vae = diffusers.AutoModel.from_pretrained(repo, subfolder = subfolder, torch_dtype = torch.bfloat16) vae = vae.to(device).eval() return vae @@ -148,11 +151,13 @@ def _ref_images(args: argparse.Namespace, size: int) -> list: from PIL import Image ref_dir = Path(args.ref_image_dir) - files = sorted(glob.glob(str(ref_dir / "*.jpg")) + glob.glob(str(ref_dir / "*.png")))[: args.num_samples] + files = sorted(glob.glob(str(ref_dir / "*.jpg")) + glob.glob(str(ref_dir / "*.png")))[ + : args.num_samples + ] imgs = [] for f in files: im = Image.open(f).convert("RGB").resize((size, size), Image.BICUBIC) - imgs.append(np.asarray(im, dtype=np.uint8)) + imgs.append(np.asarray(im, dtype = np.uint8)) return imgs @@ -189,11 +194,11 @@ def _make_latents(vae: Any, args: argparse.Namespace, device: str): x = torch.from_numpy(arr).float().permute(2, 0, 1).unsqueeze(0).div(127.5).sub(1.0) if is_3d: x = x.unsqueeze(2).repeat(1, 1, args.enc_frames, 1, 1) # static clip [1,3,T,H,W] - x = x.to(device=device, dtype=torch.bfloat16) + x = x.to(device = device, dtype = torch.bfloat16) lat.append(_encode_latent(vae, x)) if lat: return lat, is_3d - print(" (no ref images found; falling back to random latents)", flush=True) + print(" (no ref images found; falling back to random latents)", flush = True) for seed in range(args.num_samples): g = torch.Generator().manual_seed(1000 + seed) shape = ( @@ -201,8 +206,8 @@ def _make_latents(vae: Any, args: argparse.Namespace, device: str): if is_3d else (1, channels, args.latent_hw, args.latent_hw) ) - z = torch.randn(shape, generator=g, dtype=torch.float32) - lat.append(z.to(device=device, dtype=torch.bfloat16)) + z = torch.randn(shape, generator = g, dtype = torch.float32) + lat.append(z.to(device = device, dtype = torch.bfloat16)) return lat, is_3d @@ -215,7 +220,7 @@ def _decode(vae: Any, z: Any): try: out = vae.decode(z) except TypeError: - out = vae.decode(z, return_dict=True) + out = vae.decode(z, return_dict = True) sample = out.sample if hasattr(out, "sample") else out[0] sample = sample.float().clamp(-1, 1) # [B,C,H,W] (image) or [B,C,T,H,W] (video). Emit one frame per temporal slot. @@ -228,9 +233,11 @@ def _decode(vae: Any, z: Any): frames.append(sample[0]) imgs = [] for f in frames: - arr = ((f.permute(1, 2, 0).cpu().numpy() + 1.0) * 127.5).round().clip(0, 255).astype(np.uint8) + arr = ( + ((f.permute(1, 2, 0).cpu().numpy() + 1.0) * 127.5).round().clip(0, 255).astype(np.uint8) + ) if arr.shape[2] == 1: - arr = np.repeat(arr, 3, axis=2) + arr = np.repeat(arr, 3, axis = 2) imgs.append(arr) return imgs # list of HxWx3 uint8 @@ -247,13 +254,21 @@ class _Lpips: self.torch = torch self.device = device - self.fn = lpips.LPIPS(net="alex", verbose=False).to(device).eval() + self.fn = lpips.LPIPS(net = "alex", verbose = False).to(device).eval() def __call__(self, a: Any, b: Any) -> float: t = self.torch def to_t(x): - return t.from_numpy(x).float().permute(2, 0, 1).unsqueeze(0).div(127.5).sub(1.0).to(self.device) + return ( + t.from_numpy(x) + .float() + .permute(2, 0, 1) + .unsqueeze(0) + .div(127.5) + .sub(1.0) + .to(self.device) + ) with t.no_grad(): return float(self.fn(to_t(a), to_t(b)).item()) @@ -299,12 +314,18 @@ def _apply_fp8_dynamic_no1x1(vae_q: Any) -> None: if w is None or w.dim() < 2 or w.shape[0] % 16 or w.shape[1] % 16: return False ks = getattr(module, "kernel_size", None) - if isinstance(ks, tuple) and all(k == 1 for k in ks): # pointwise conv -> torchao kernel fails + if isinstance(ks, tuple) and all( + k == 1 for k in ks + ): # pointwise conv -> torchao kernel fails return False name = fqn.lower() if fqn else "" return not any(tok in name for tok in _VAE_KEEP_DENSE_TOKENS) - quantize_(vae_q, Float8DynamicActivationFloat8WeightConfig(granularity=PerTensor()), filter_fn=filter_fn) + quantize_( + vae_q, + Float8DynamicActivationFloat8WeightConfig(granularity = PerTensor()), + filter_fn = filter_fn, + ) def _sweep_vae(args: argparse.Namespace, lp: "_Lpips", out_dir: Path) -> list[dict]: @@ -318,7 +339,6 @@ def _sweep_vae(args: argparse.Namespace, lp: "_Lpips", out_dir: Path) -> list[di class _Target: def __init__(self): import torch - self.device = "cuda" self.dtype = torch.bfloat16 @@ -326,30 +346,33 @@ def _sweep_vae(args: argparse.Namespace, lp: "_Lpips", out_dir: Path) -> list[di rows: list[dict] = [] for family in args.family: repo = _VAE_FAMILIES[family]["repo"] - print(f"\n=== VAE {family} ({repo}) ===", flush=True) + print(f"\n=== VAE {family} ({repo}) ===", flush = True) t0 = time.time() try: vae = _load_vae(repo, "vae", "cuda") except Exception as exc: # noqa: BLE001 - print(f" load FAILED: {type(exc).__name__}: {str(exc)[:200]}", flush=True) + print(f" load FAILED: {type(exc).__name__}: {str(exc)[:200]}", flush = True) rows.append({"family": family, "scheme": "-", "error": f"load: {exc}"}) continue ch, is_3d = _latent_spec(vae) cls = type(vae).__name__ - print(f" {cls} latent_ch={ch} {'3D' if is_3d else '2D'} loaded {time.time()-t0:.0f}s", flush=True) + print( + f" {cls} latent_ch={ch} {'3D' if is_3d else '2D'} loaded {time.time()-t0:.0f}s", + flush = True, + ) latents, _ = _make_latents(vae, args, "cuda") ref_by_sample = [_decode(vae, z) for z in latents] fam_dir = out_dir / family - fam_dir.mkdir(parents=True, exist_ok=True) + fam_dir.mkdir(parents = True, exist_ok = True) Image.fromarray(ref_by_sample[0][0]).save(fam_dir / "dense_s0.png") for scheme in args.scheme: try: vae_q = copy.deepcopy(vae) except Exception as exc: # noqa: BLE001 - print(f" [{scheme}] deepcopy FAILED: {exc}", flush=True) + print(f" [{scheme}] deepcopy FAILED: {exc}", flush = True) continue pipe = type("P", (), {"vae": vae_q})() # "fp8_dynamic_no1x1" is a diagnostic that bypasses quantize_vae to apply the @@ -359,17 +382,19 @@ def _sweep_vae(args: argparse.Namespace, lp: "_Lpips", out_dir: Path) -> list[di _apply_fp8_dynamic_no1x1(vae_q) engaged = scheme except Exception as exc: # noqa: BLE001 - print(f" [{scheme}] apply FAILED: {str(exc)[:120]}", flush=True) + print(f" [{scheme}] apply FAILED: {str(exc)[:120]}", flush = True) del vae_q _empty_cache() continue else: engaged = vq.quantize_vae( - pipe, target, mode=scheme, family=family, offload_active=False, force_fp32=False + pipe, target, mode = scheme, family = family, offload_active = False, force_fp32 = False ) if engaged != scheme: - print(f" [{scheme}] NOT engaged (returned {engaged}); skipping", flush=True) - rows.append({"family": family, "vae_class": cls, "scheme": scheme, "verdict": "NOT_ENGAGED"}) + print(f" [{scheme}] NOT engaged (returned {engaged}); skipping", flush = True) + rows.append( + {"family": family, "vae_class": cls, "scheme": scheme, "verdict": "NOT_ENGAGED"} + ) del vae_q _empty_cache() continue @@ -377,9 +402,15 @@ def _sweep_vae(args: argparse.Namespace, lp: "_Lpips", out_dir: Path) -> list[di q_by_sample = [_decode(vae_q, z) for z in latents] except Exception as exc: # noqa: BLE001 — the shipped caster produced a VAE that crashes at decode emsg = f"{type(exc).__name__}: {str(exc)[:100]}" - print(f" [{scheme}] DECODE CRASH: {emsg}", flush=True) + print(f" [{scheme}] DECODE CRASH: {emsg}", flush = True) rows.append( - {"family": family, "vae_class": cls, "scheme": scheme, "verdict": "CRASH", "error": emsg} + { + "family": family, + "vae_class": cls, + "scheme": scheme, + "verdict": "CRASH", + "error": emsg, + } ) del vae_q _empty_cache() @@ -402,7 +433,7 @@ def _sweep_vae(args: argparse.Namespace, lp: "_Lpips", out_dir: Path) -> list[di rows.append(row) print( f" [{scheme}] LPIPS={m['lpips']} PSNR={m['psnr']} SSIM={m['ssim']} -> {verdict}", - flush=True, + flush = True, ) del vae_q _empty_cache() @@ -414,7 +445,6 @@ def _sweep_vae(args: argparse.Namespace, lp: "_Lpips", out_dir: Path) -> list[di def _empty_cache() -> None: try: import torch - torch.cuda.empty_cache() except Exception: pass @@ -448,14 +478,16 @@ def _apply_auto(pipe: Any, family: str, components: list[str]) -> dict[str, Opti tgt = _Target() engaged: dict[str, Optional[str]] = {} if "transformer" in components: - engaged["transformer"] = tq.quantize_transformer(pipe, tgt, mode="auto", family=family) + engaged["transformer"] = tq.quantize_transformer(pipe, tgt, mode = "auto", family = family) if "text_encoder" in components: try: - engaged["text_encoder"] = dp.quantize_text_encoders(pipe, tgt, mode="auto", family=family) + engaged["text_encoder"] = dp.quantize_text_encoders( + pipe, tgt, mode = "auto", family = family + ) except Exception as exc: # noqa: BLE001 engaged["text_encoder"] = f"err:{type(exc).__name__}" if "vae" in components: - engaged["vae"] = vq.quantize_vae(pipe, tgt, mode="auto", family=family) + engaged["vae"] = vq.quantize_vae(pipe, tgt, mode = "auto", family = family) return engaged @@ -466,7 +498,7 @@ def _sweep_e2e(args: argparse.Namespace, lp: "_Lpips", out_dir: Path) -> list[di rows: list[dict] = [] for family in args.family: model = args.e2e_model or _VAE_FAMILIES.get(family, {}).get("repo") - print(f"\n=== E2E {family} ({model}) ===", flush=True) + print(f"\n=== E2E {family} ({model}) ===", flush = True) prompts = args.prompts or _E2E_PROMPTS seeds = args.seeds @@ -474,13 +506,13 @@ def _sweep_e2e(args: argparse.Namespace, lp: "_Lpips", out_dir: Path) -> list[di imgs = [] for pi, prompt in enumerate(prompts): for seed in seeds: - g = torch.Generator(device="cuda").manual_seed(seed) + g = torch.Generator(device = "cuda").manual_seed(seed) kw = dict( - prompt=prompt, - num_inference_steps=args.steps, - generator=g, - height=args.height, - width=args.width, + prompt = prompt, + num_inference_steps = args.steps, + generator = g, + height = args.height, + width = args.width, ) if args.guidance is not None: kw["guidance_scale"] = args.guidance @@ -490,10 +522,10 @@ def _sweep_e2e(args: argparse.Namespace, lp: "_Lpips", out_dir: Path) -> list[di try: pipe = diffusers.AutoPipelineForText2Image.from_pretrained( - model, torch_dtype=torch.bfloat16 + model, torch_dtype = torch.bfloat16 ).to("cuda") except Exception as exc: # noqa: BLE001 - print(f" pipe load FAILED: {type(exc).__name__}: {str(exc)[:200]}", flush=True) + print(f" pipe load FAILED: {type(exc).__name__}: {str(exc)[:200]}", flush = True) rows.append({"family": family, "error": f"load: {exc}"}) continue ref = _gen(pipe) @@ -501,10 +533,10 @@ def _sweep_e2e(args: argparse.Namespace, lp: "_Lpips", out_dir: Path) -> list[di _empty_cache() pipe2 = diffusers.AutoPipelineForText2Image.from_pretrained( - model, torch_dtype=torch.bfloat16 + model, torch_dtype = torch.bfloat16 ).to("cuda") engaged = _apply_auto(pipe2, family, args.e2e_components) - print(f" engaged: {engaged}", flush=True) + print(f" engaged: {engaged}", flush = True) q = _gen(pipe2) del pipe2 _empty_cache() @@ -513,7 +545,7 @@ def _sweep_e2e(args: argparse.Namespace, lp: "_Lpips", out_dir: Path) -> list[di ls = [] fam_dir = out_dir / f"e2e_{family}" - fam_dir.mkdir(parents=True, exist_ok=True) + fam_dir.mkdir(parents = True, exist_ok = True) for (pi, seed, a), (_, _, b) in zip(ref, q): aa, bb = np.asarray(a.convert("RGB")), np.asarray(b.convert("RGB")) ls.append(lp(aa, bb)) @@ -524,7 +556,7 @@ def _sweep_e2e(args: argparse.Namespace, lp: "_Lpips", out_dir: Path) -> list[di rows.append( {"family": family, "engaged": engaged, "mean_lpips": mean_l, "verdict": verdict} ) - print(f" mean LPIPS={mean_l} -> {verdict}", flush=True) + print(f" mean LPIPS={mean_l} -> {verdict}", flush = True) return rows @@ -532,28 +564,33 @@ def _sweep_e2e(args: argparse.Namespace, lp: "_Lpips", out_dir: Path) -> list[di def _write(out_dir: Path, mode: str, rows: list[dict]) -> None: - (out_dir / f"{mode}_results.json").write_text(json.dumps(rows, indent=2)) - print(f"\nwrote {out_dir / f'{mode}_results.json'}", flush=True) - print(f"\n=== {mode.upper()} RESULTS ===", flush=True) + (out_dir / f"{mode}_results.json").write_text(json.dumps(rows, indent = 2)) + print(f"\nwrote {out_dir / f'{mode}_results.json'}", flush = True) + print(f"\n=== {mode.upper()} RESULTS ===", flush = True) if mode == "vae": - print(f" bars: LPIPS <= {LPIPS_BAR}, SSIM >= {SSIM_BAR}", flush=True) - print(f" {'family':<20}{'scheme':<14}{'LPIPS':>9}{'PSNR':>9}{'SSIM':>9} verdict", flush=True) + print(f" bars: LPIPS <= {LPIPS_BAR}, SSIM >= {SSIM_BAR}", flush = True) + print( + f" {'family':<20}{'scheme':<14}{'LPIPS':>9}{'PSNR':>9}{'SSIM':>9} verdict", flush = True + ) for r in rows: if "error" in r: - print(f" {r['family']:<20}{'(error)':<14} {r['error'][:60]}", flush=True) + print(f" {r['family']:<20}{'(error)':<14} {r['error'][:60]}", flush = True) continue print( f" {r['family']:<20}{r['scheme']:<14}{_f(r.get('lpips')):>9}" f"{_f(r.get('psnr')):>9}{_f(r.get('ssim')):>9} {r.get('verdict')}", - flush=True, + flush = True, ) else: - print(f" bar: mean LPIPS <= {E2E_LPIPS_BAR}", flush=True) + print(f" bar: mean LPIPS <= {E2E_LPIPS_BAR}", flush = True) for r in rows: if "error" in r: - print(f" {r['family']}: (error) {r['error'][:80]}", flush=True) + print(f" {r['family']}: (error) {r['error'][:80]}", flush = True) continue - print(f" {r['family']:<20} mean_lpips={r.get('mean_lpips')} {r.get('verdict')} {r.get('engaged')}", flush=True) + print( + f" {r['family']:<20} mean_lpips={r.get('mean_lpips')} {r.get('verdict')} {r.get('engaged')}", + flush = True, + ) def _f(v: Any) -> str: @@ -565,36 +602,49 @@ def _f(v: Any) -> str: def _build_parser() -> argparse.ArgumentParser: p = argparse.ArgumentParser( - description="Decoded-image accuracy sweep for the auto VAE / end-to-end quantisation.", - formatter_class=argparse.ArgumentDefaultsHelpFormatter, + description = "Decoded-image accuracy sweep for the auto VAE / end-to-end quantisation.", + formatter_class = argparse.ArgumentDefaultsHelpFormatter, ) - p.add_argument("--mode", choices=["vae", "e2e"], default="vae") - p.add_argument("--family", nargs="+", default=list(_VAE_FAMILIES.keys())) - p.add_argument("--scheme", nargs="+", default=["fp8_dynamic", "fp8_dynamic_no1x1", "fp8"]) - p.add_argument("--num-samples", type=int, default=5, help="latents to average over") - p.add_argument("--latents", choices=["encode", "random"], default="encode", - help="encode natural photos (in-distribution) or seeded N(0,1) latents") - p.add_argument("--ref-image-dir", default="outputs/quant_accuracy/_refs", - help="natural photos to encode for the round-trip") - p.add_argument("--enc-hw", type=int, default=512, help="2D encode pixel H=W") - p.add_argument("--enc-hw-3d", type=int, default=256, help="3D encode pixel H=W") - p.add_argument("--enc-frames", type=int, default=9, help="3D encode pixel frame count") - p.add_argument("--latent-hw", type=int, default=64, help="2D random-latent H=W (x8 -> 512px)") - p.add_argument("--latent-hw-3d", type=int, default=32, help="3D random-latent H=W") - p.add_argument("--latent-t", type=int, default=3, help="3D random-latent temporal length") - p.add_argument("--lpips-device", default="cpu", help="device for the LPIPS net (keep off the measured GPU)") - p.add_argument("--out-dir", default="outputs/quant_accuracy") + p.add_argument("--mode", choices = ["vae", "e2e"], default = "vae") + p.add_argument("--family", nargs = "+", default = list(_VAE_FAMILIES.keys())) + p.add_argument("--scheme", nargs = "+", default = ["fp8_dynamic", "fp8_dynamic_no1x1", "fp8"]) + p.add_argument("--num-samples", type = int, default = 5, help = "latents to average over") + p.add_argument( + "--latents", + choices = ["encode", "random"], + default = "encode", + help = "encode natural photos (in-distribution) or seeded N(0,1) latents", + ) + p.add_argument( + "--ref-image-dir", + default = "outputs/quant_accuracy/_refs", + help = "natural photos to encode for the round-trip", + ) + p.add_argument("--enc-hw", type = int, default = 512, help = "2D encode pixel H=W") + p.add_argument("--enc-hw-3d", type = int, default = 256, help = "3D encode pixel H=W") + p.add_argument("--enc-frames", type = int, default = 9, help = "3D encode pixel frame count") + p.add_argument("--latent-hw", type = int, default = 64, help = "2D random-latent H=W (x8 -> 512px)") + p.add_argument("--latent-hw-3d", type = int, default = 32, help = "3D random-latent H=W") + p.add_argument("--latent-t", type = int, default = 3, help = "3D random-latent temporal length") + p.add_argument( + "--lpips-device", default = "cpu", help = "device for the LPIPS net (keep off the measured GPU)" + ) + p.add_argument("--out-dir", default = "outputs/quant_accuracy") # e2e-only - p.add_argument("--e2e-model", default=None, help="full model repo for --mode e2e") - p.add_argument("--e2e-components", nargs="+", default=["transformer", "text_encoder", "vae"], - choices=["transformer", "text_encoder", "vae"], - help="which components to auto-quantise for the e2e (isolate the VAE with: --e2e-components vae)") - p.add_argument("--prompts", nargs="*", default=None) - p.add_argument("--seeds", nargs="*", type=int, default=[12345]) - p.add_argument("--steps", type=int, default=8) - p.add_argument("--guidance", type=float, default=None) - p.add_argument("--height", type=int, default=1024) - p.add_argument("--width", type=int, default=1024) + p.add_argument("--e2e-model", default = None, help = "full model repo for --mode e2e") + p.add_argument( + "--e2e-components", + nargs = "+", + default = ["transformer", "text_encoder", "vae"], + choices = ["transformer", "text_encoder", "vae"], + help = "which components to auto-quantise for the e2e (isolate the VAE with: --e2e-components vae)", + ) + p.add_argument("--prompts", nargs = "*", default = None) + p.add_argument("--seeds", nargs = "*", type = int, default = [12345]) + p.add_argument("--steps", type = int, default = 8) + p.add_argument("--guidance", type = float, default = None) + p.add_argument("--height", type = int, default = 1024) + p.add_argument("--width", type = int, default = 1024) return p @@ -602,7 +652,7 @@ def main(argv: Optional[list[str]] = None) -> int: args = _build_parser().parse_args(argv) _check_deps() out_dir = Path(args.out_dir).resolve() - out_dir.mkdir(parents=True, exist_ok=True) + out_dir.mkdir(parents = True, exist_ok = True) lp = _Lpips(args.lpips_device) if args.mode == "vae": rows = _sweep_vae(args, lp, out_dir) diff --git a/scripts/quant_speedmem_bench.py b/scripts/quant_speedmem_bench.py index 98bd82b735..229ad05c9b 100644 --- a/scripts/quant_speedmem_bench.py +++ b/scripts/quant_speedmem_bench.py @@ -69,33 +69,28 @@ _TE_ATTRS = ("text_encoder", "text_encoder_2", "text_encoder_3") # ── cuda memory / timing helpers (lifted from diffusion_bench.py) ────────────── 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() @@ -116,12 +111,11 @@ def _lpips_alex(ref_arr, arr): fn = _LP.get("fn") if fn is None: - fn = lpips.LPIPS(net="alex", verbose=False).eval() + 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(): @@ -137,7 +131,7 @@ def _timed_generate(pipe, *, steps, res, seed): import torch - g = torch.Generator(device="cuda").manual_seed(seed) + g = torch.Generator(device = "cuda").manual_seed(seed) step_ts: list[float] = [] last = [0.0] @@ -152,8 +146,12 @@ def _timed_generate(pipe, *, steps, res, seed): _sync() t0 = _time.perf_counter() img = pipe( - prompt=PROMPT, width=res, height=res, num_inference_steps=steps, generator=g, - callback_on_step_end=_cb, + 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 @@ -177,15 +175,14 @@ def _target(): import torch - return types.SimpleNamespace(device="cuda", dtype=torch.bfloat16) + return types.SimpleNamespace(device = "cuda", dtype = torch.bfloat16) # ── model_index component class resolution ──────────────────────────────────── def _model_index(repo: str) -> dict: from huggingface_hub import hf_hub_download - path = hf_hub_download(repo, "model_index.json") - with open(path, "r", encoding="utf-8") as fh: + with open(path, "r", encoding = "utf-8") as fh: return json.load(fh) @@ -203,7 +200,7 @@ def _load_named(repo: str, name: str): lib, cls_name = spec module = importlib.import_module(lib) klass = getattr(module, cls_name) - return klass.from_pretrained(repo, subfolder=name, torch_dtype=torch.bfloat16) + return klass.from_pretrained(repo, subfolder = name, torch_dtype = torch.bfloat16) def _load_text_encoders(repo: str, device: str): @@ -244,7 +241,7 @@ def _load_tokenizers(repo: str): ("text_encoder_3", "tokenizer_3"), ): try: - toks[te_attr] = AutoTokenizer.from_pretrained(repo, subfolder=tk_sub) + toks[te_attr] = AutoTokenizer.from_pretrained(repo, subfolder = tk_sub) except Exception: toks[te_attr] = None return toks @@ -256,15 +253,15 @@ def _encoder_hidden(te, ids, mask): with torch.inference_mode(): try: - out = te(input_ids=ids, attention_mask=mask, output_hidden_states=True) + out = te(input_ids = ids, attention_mask = mask, output_hidden_states = True) except TypeError: - out = te(ids, output_hidden_states=True) + out = te(ids, output_hidden_states = True) hs = getattr(out, "last_hidden_state", None) if hs is None: hidden = getattr(out, "hidden_states", None) hs = hidden[-1] if hidden else (out[0] if isinstance(out, (tuple, list)) else out) m = mask.unsqueeze(-1).to(hs.dtype) - v = (hs * m).sum(1) / m.sum(1).clamp(min=1) + v = (hs * m).sum(1) / m.sum(1).clamp(min = 1) return v.float().flatten() @@ -280,7 +277,7 @@ def _te_hidden_refs(bag, toks, device): continue vecs = [] for p in _PROMPT_SUITE: - enc = tok(p, return_tensors="pt", padding="max_length", truncation=True, max_length=64) + enc = tok(p, return_tensors = "pt", padding = "max_length", truncation = True, max_length = 64) ids = enc["input_ids"].to(device) mask = enc.get("attention_mask") mask = mask.to(device) if mask is not None else torch.ones_like(ids) @@ -289,7 +286,12 @@ def _te_hidden_refs(bag, toks, device): return refs -def measure_te_accuracy(family: str, *, schemes=("fp8", "fp8_dynamic"), logger=None) -> list[dict]: +def measure_te_accuracy( + family: str, + *, + schemes = ("fp8", "fp8_dynamic"), + logger = None, +) -> list[dict]: """Hidden-state cosine / relL2 of each TE scheme vs the dense bf16 encoder (per encoder), on real prompts. Bar (PR#150): cosine >= 0.99 and min_cosine >= 0.98.""" import torch @@ -310,14 +312,16 @@ def measure_te_accuracy(family: str, *, schemes=("fp8", "fp8_dynamic"), logger=N rows: list[dict] = [] for scheme in schemes: bag = _load_text_encoders(repo, device) - engaged = quantize_text_encoders(bag, _target(), mode=scheme, family=family, logger=logger) + engaged = quantize_text_encoders(bag, _target(), mode = scheme, family = family, logger = logger) cur = _te_hidden_refs(bag, toks, device) for attr, ref_vecs in refs.items(): q_vecs = cur.get(attr, []) if not q_vecs: continue - cosines = [F.cosine_similarity(r, q, dim=0).item() for r, q in zip(ref_vecs, q_vecs)] - rell2 = [((q - r).norm() / r.norm().clamp(min=1e-8)).item() for r, q in zip(ref_vecs, q_vecs)] + cosines = [F.cosine_similarity(r, q, dim = 0).item() for r, q in zip(ref_vecs, q_vecs)] + rell2 = [ + ((q - r).norm() / r.norm().clamp(min = 1e-8)).item() for r, q in zip(ref_vecs, q_vecs) + ] mean_cos = sum(cosines) / len(cosines) min_cos = min(cosines) rows.append( @@ -340,7 +344,6 @@ def _encode_once(bag) -> None: """One forward through every present text encoder on a fixed short token batch. Uses a length within each encoder's max positions (CLIP caps at 77) so position embeddings never overflow.""" import torch - with torch.inference_mode(): for attr in _TE_ATTRS: te = getattr(bag, attr, None) @@ -350,10 +353,10 @@ def _encode_once(bag) -> None: vocab = int(getattr(cfg, "vocab_size", 30000) or 30000) maxpos = int(getattr(cfg, "max_position_embeddings", 64) or 64) length = max(8, min(64, maxpos)) - ids = torch.randint(1, min(vocab, 30000), (1, length), device="cuda") + ids = torch.randint(1, min(vocab, 30000), (1, length), device = "cuda") mask = torch.ones_like(ids) try: - te(input_ids=ids, attention_mask=mask) + te(input_ids = ids, attention_mask = mask) except TypeError: te(ids) @@ -363,7 +366,7 @@ def _load_vae(repo: str, device: str): import torch diffusers = _import_diffusers() - vae = diffusers.AutoModel.from_pretrained(repo, subfolder="vae", torch_dtype=torch.bfloat16) + vae = diffusers.AutoModel.from_pretrained(repo, subfolder = "vae", torch_dtype = torch.bfloat16) return vae.to(device).eval() @@ -395,12 +398,11 @@ def _latent_spec(vae) -> tuple[int, bool]: def _decode_once(vae, z) -> None: import torch - with torch.inference_mode(): try: out = vae.decode(z) except TypeError: - out = vae.decode(z, return_dict=True) + out = vae.decode(z, return_dict = True) _ = out.sample if hasattr(out, "sample") else out[0] @@ -411,8 +413,8 @@ def _make_latent(vae, device: str): channels, is_3d = _latent_spec(vae) g = torch.Generator().manual_seed(1234) shape = (1, channels, 3, 32, 32) if is_3d else (1, channels, 64, 64) - z = torch.randn(shape, generator=g, dtype=torch.float32) - return z.to(device=device, dtype=torch.bfloat16) + z = torch.randn(shape, generator = g, dtype = torch.float32) + return z.to(device = device, dtype = torch.bfloat16) # ── measurement primitives ──────────────────────────────────────────────────── @@ -431,7 +433,14 @@ def _time_median(fn, *, warmup: int, iters: int) -> float: # ── mode: te ────────────────────────────────────────────────────────────────── -def measure_te(family: str, *, warmup: int, iters: int, scheme: str = "auto", logger=None) -> list[dict]: +def measure_te( + family: str, + *, + warmup: int, + iters: int, + scheme: str = "auto", + logger = None, +) -> list[dict]: from core.inference.diffusion_precision import quantize_text_encoders repo = _FAMILIES[family]["repo"] @@ -439,20 +448,22 @@ def measure_te(family: str, *, warmup: int, iters: int, scheme: str = "auto", lo rows: list[dict] = [] # dense - _empty(); _reset_peak() + _empty() + _reset_peak() bag = _load_text_encoders(repo, device) _sync() mem_dense = _alloc_gb() _reset_peak() - lat_dense = _time_median(lambda: _encode_once(bag), warmup=warmup, iters=iters) + lat_dense = _time_median(lambda: _encode_once(bag), warmup = warmup, iters = iters) peak_dense = _peak_gb() # quant in place (scheme="auto" resolves the ladder; else force an explicit scheme to compare) - engaged = quantize_text_encoders(bag, _target(), mode=scheme, family=family, logger=logger) - _empty(); _sync() + engaged = quantize_text_encoders(bag, _target(), mode = scheme, family = family, logger = logger) + _empty() + _sync() mem_quant = _alloc_gb() _reset_peak() - lat_quant = _time_median(lambda: _encode_once(bag), warmup=warmup, iters=iters) + lat_quant = _time_median(lambda: _encode_once(bag), warmup = warmup, iters = iters) peak_quant = _peak_gb() del bag @@ -470,38 +481,55 @@ def measure_te(family: str, *, warmup: int, iters: int, scheme: str = "auto", lo "peak_quant_gb": round(peak_quant, 3), "lat_dense_ms": round(lat_dense, 2), "lat_quant_ms": round(lat_quant, 2), - "lat_delta_pct": round((lat_quant - lat_dense) / lat_dense * 100.0, 1) if lat_dense else None, + "lat_delta_pct": round((lat_quant - lat_dense) / lat_dense * 100.0, 1) + if lat_dense + else None, } ) return rows # ── mode: vae ─────────────────────────────────────────────────────────────── -def _measure_vae_scheme(family: str, repo: str, mode: str, *, warmup: int, iters: int, logger=None) -> dict: +def _measure_vae_scheme( + family: str, + repo: str, + mode: str, + *, + warmup: int, + iters: int, + logger = None, +) -> dict: import types from core.inference.diffusion_vae_quant import quantize_vae device = "cuda" - _empty(); _reset_peak() + _empty() + _reset_peak() vae = _load_vae(repo, device) z = _make_latent(vae, device) _sync() mem_dense = _alloc_gb() _reset_peak() - lat_dense = _time_median(lambda: _decode_once(vae, z), warmup=warmup, iters=iters) + lat_dense = _time_median(lambda: _decode_once(vae, z), warmup = warmup, iters = iters) peak_dense = _peak_gb() # quantize_vae reads pipe.vae, so hand it a bag exposing .vae (it mutates that module in place). # Pass the family's force_fp32 (Wan) so the real dense-only behaviour is reflected. force_fp32 = bool(_FAMILIES.get(family, {}).get("vae_force_fp32", False)) engaged = quantize_vae( - types.SimpleNamespace(vae=vae), _target(), mode=mode, family=family, force_fp32=force_fp32, logger=logger + types.SimpleNamespace(vae = vae), + _target(), + mode = mode, + family = family, + force_fp32 = force_fp32, + logger = logger, ) - _empty(); _sync() + _empty() + _sync() mem_quant = _alloc_gb() _reset_peak() - lat_quant = _time_median(lambda: _decode_once(vae, z), warmup=warmup, iters=iters) + lat_quant = _time_median(lambda: _decode_once(vae, z), warmup = warmup, iters = iters) peak_quant = _peak_gb() del vae, z @@ -519,40 +547,65 @@ def _measure_vae_scheme(family: str, repo: str, mode: str, *, warmup: int, iters "peak_quant_gb": round(peak_quant, 3), "lat_dense_ms": round(lat_dense, 2), "lat_quant_ms": round(lat_quant, 2), - "lat_delta_pct": round((lat_quant - lat_dense) / lat_dense * 100.0, 1) if lat_dense else None, + "lat_delta_pct": round((lat_quant - lat_dense) / lat_dense * 100.0, 1) + if lat_dense + else None, } -def measure_vae(family: str, *, warmup: int, iters: int, logger=None) -> list[dict]: +def measure_vae( + family: str, + *, + warmup: int, + iters: int, + logger = None, +) -> list[dict]: repo = _FAMILIES[family]["repo"] - rows = [_measure_vae_scheme(family, repo, "auto", warmup=warmup, iters=iters, logger=logger)] + rows = [_measure_vae_scheme(family, repo, "auto", warmup = warmup, iters = iters, logger = logger)] if family == "flux.2": # the one image family where the explicit fp8_dynamic conv opt-in is measured in-bar. - rows.append(_measure_vae_scheme(family, repo, "fp8_dynamic", warmup=warmup, iters=iters, logger=logger)) + rows.append( + _measure_vae_scheme( + family, repo, "fp8_dynamic", warmup = warmup, iters = iters, logger = logger + ) + ) return rows # ── mode: e2e (qwen-image) ──────────────────────────────────────────────────── -def _e2e_run(repo: str, *, quant: bool, steps: int, res: int, seed: int, iters: int, family: str, logger=None): +def _e2e_run( + repo: str, + *, + quant: bool, + steps: int, + res: int, + seed: int, + iters: int, + family: str, + logger = None, +): import torch from core.inference.diffusion_precision import quantize_text_encoders from core.inference.diffusion_vae_quant import quantize_vae diffusers = _import_diffusers() - _empty(); _reset_peak() - pipe = diffusers.DiffusionPipeline.from_pretrained(repo, torch_dtype=torch.bfloat16) + _empty() + _reset_peak() + pipe = diffusers.DiffusionPipeline.from_pretrained(repo, torch_dtype = torch.bfloat16) pipe = pipe.to("cuda") load_peak = _peak_gb() te_scheme = vae_scheme = None if quant: - te_scheme = quantize_text_encoders(pipe, _target(), mode="auto", family=family, logger=logger) - vae_scheme = quantize_vae(pipe, _target(), mode="auto", family=family, logger=logger) + te_scheme = quantize_text_encoders( + pipe, _target(), mode = "auto", family = family, logger = logger + ) + vae_scheme = quantize_vae(pipe, _target(), mode = "auto", family = family, logger = logger) _empty() weights_gb = _alloc_gb() def _gen(): - g = torch.Generator(device="cuda").manual_seed(seed) + g = torch.Generator(device = "cuda").manual_seed(seed) step_ts: list[float] = [] last = [0.0] @@ -567,12 +620,12 @@ def _e2e_run(repo: str, *, quant: bool, steps: int, res: int, seed: int, iters: _sync() t0 = time.perf_counter() img = pipe( - prompt=PROMPT, - width=res, - height=res, - num_inference_steps=steps, - generator=g, - callback_on_step_end=_cb, + 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 @@ -600,7 +653,15 @@ def _e2e_run(repo: str, *, quant: bool, steps: int, res: int, seed: int, iters: def measure_e2e( - family: str, *, steps: int, res: int, seed: int, iters: int, out: Path, variant: str = "both", logger=None + family: str, + *, + steps: int, + res: int, + seed: int, + iters: int, + out: Path, + variant: str = "both", + logger = None, ) -> list[dict]: repo = _FAMILIES[family]["repo"] # "both" runs dense then auto in one process (fast, but the 2nd run is on a hotter GPU / a @@ -610,14 +671,21 @@ def measure_e2e( rows = [] for quant in variants: row, img = _e2e_run( - repo, quant=quant, steps=steps, res=res, seed=seed, iters=iters, family=family, logger=logger + repo, + quant = quant, + steps = steps, + res = res, + seed = seed, + iters = iters, + family = family, + logger = logger, ) try: img.save(out / f"e2e_{family}_{row['variant']}.png") except Exception: pass rows.append(row) - print(f" e2e {row['variant']:5s}: {json.dumps(row)}", flush=True) + print(f" e2e {row['variant']:5s}: {json.dumps(row)}", flush = True) return rows @@ -638,8 +706,16 @@ def _compile_blocks(transformer) -> bool: def _dit_run( - repo: str, family: str, *, dit_quant: str, steps: int, res: int, seed: int, iters: int, - compile_blocks: bool = True, logger=None, + 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 @@ -649,21 +725,24 @@ def _dit_run( 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") + _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) + 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) + 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) + 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() @@ -682,14 +761,24 @@ def _dit_run( }, last_img -def measure_dit(family: str, *, schemes, steps: int, res: int, seed: int, iters: int, out: Path, logger=None): +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 + 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") @@ -699,78 +788,105 @@ def measure_dit(family: str, *, schemes, steps: int, res: int, seed: int, iters: 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) + 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 + 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 + 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) + 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", "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") - ap.add_argument("--res", type=int, default=1024, help="e2e image size") - ap.add_argument("--seed", type=int, default=42) - ap.add_argument("--e2e-iters", type=int, default=3) - ap.add_argument("--variant", choices=("both", "dense", "auto"), default="both", help="e2e variant(s)") - ap.add_argument("--te-scheme", default="auto", help="te mode: auto | fp8_dynamic | fp8 | int8") - ap.add_argument("--out", default="outputs/quant_speedmem") +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", "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") + ap.add_argument("--res", type = int, default = 1024, help = "e2e image size") + ap.add_argument("--seed", type = int, default = 42) + ap.add_argument("--e2e-iters", type = int, default = 3) + ap.add_argument( + "--variant", choices = ("both", "dense", "auto"), default = "both", help = "e2e variant(s)" + ) + ap.add_argument("--te-scheme", default = "auto", help = "te mode: auto | fp8_dynamic | fp8 | int8") + ap.add_argument("--out", default = "outputs/quant_speedmem") args = ap.parse_args(argv) import logging - logging.basicConfig(level=logging.INFO, format="%(message)s") + logging.basicConfig(level = logging.INFO, format = "%(message)s") logger = logging.getLogger("speedmem") out = Path(args.out) - out.mkdir(parents=True, exist_ok=True) + out.mkdir(parents = True, exist_ok = True) - print(f"== speed+mem bench: family={args.family} mode={args.mode} ==", flush=True) + print(f"== speed+mem bench: family={args.family} mode={args.mode} ==", flush = True) 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 + 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) + 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) + rows = measure_te( + args.family, warmup = args.warmup, iters = args.iters, scheme = args.te_scheme, logger = logger + ) elif args.mode == "vae": - rows = measure_vae(args.family, warmup=args.warmup, iters=args.iters, logger=logger) + rows = measure_vae(args.family, warmup = args.warmup, iters = args.iters, logger = logger) else: rows = measure_e2e( - args.family, steps=args.steps, res=args.res, seed=args.seed, iters=args.e2e_iters, - out=out, variant=args.variant, logger=logger + args.family, + steps = args.steps, + res = args.res, + seed = args.seed, + iters = args.e2e_iters, + out = out, + variant = args.variant, + logger = logger, ) for r in rows: - print(" " + json.dumps(r), flush=True) + print(" " + json.dumps(r), flush = True) suffix = "" if args.mode == "e2e" and args.variant != "both": suffix = f"_{args.variant}" elif args.mode == "te" and args.te_scheme != "auto": suffix = f"_{args.te_scheme}" dest = out / f"{args.mode}_{args.family}{suffix}.json" - with open(dest, "w", encoding="utf-8") as fh: - json.dump(rows, fh, indent=2) - print(f"wrote {dest}", flush=True) + with open(dest, "w", encoding = "utf-8") as fh: + json.dump(rows, fh, indent = 2) + print(f"wrote {dest}", flush = True) return 0 diff --git a/scripts/sdpa_mask_backend_probe.py b/scripts/sdpa_mask_backend_probe.py index 3dd83edeaa..4f2a030869 100644 --- a/scripts/sdpa_mask_backend_probe.py +++ b/scripts/sdpa_mask_backend_probe.py @@ -4,6 +4,7 @@ """Which SDPA backends tolerate a dense bool attn_mask, and at what cost, at Hunyuan's real joint shape (B=1, H=16, N=50345, D=128, bf16)? Decides whether nulling the all-True mask is the real win on the PRODUCTION cuDNN path (not just the native math fallback).""" + import time import torch import torch.nn.functional as F @@ -14,10 +15,10 @@ dev, dt = "cuda:0", torch.bfloat16 def mk(): - return torch.randn(B, H, N, D, device=dev, dtype=dt) + return torch.randn(B, H, N, D, device = dev, dtype = dt) -def timed(fn, iters=20): +def timed(fn, iters = 20): torch.cuda.synchronize() for _ in range(3): try: @@ -33,7 +34,7 @@ def timed(fn, iters=20): q, k, v = mk(), mk(), mk() -dense = torch.ones(B, 1, N, N, dtype=torch.bool, device=dev) +dense = torch.ones(B, 1, N, N, dtype = torch.bool, device = dev) backends = { "default(dispatch)": None, @@ -46,17 +47,18 @@ backends = { print(f"shape B={B} H={H} N={N} D={D} {dt}\n") print(f"{'backend':<20}{'mask=dense(ms)':>18}{'mask=None(ms)':>18}") for name, bk in backends.items(): + def run_dense(): if bk is None: - return F.scaled_dot_product_attention(q, k, v, attn_mask=dense) + return F.scaled_dot_product_attention(q, k, v, attn_mask = dense) with sdpa_kernel(bk): - return F.scaled_dot_product_attention(q, k, v, attn_mask=dense) + return F.scaled_dot_product_attention(q, k, v, attn_mask = dense) def run_none(): if bk is None: - return F.scaled_dot_product_attention(q, k, v, attn_mask=None) + return F.scaled_dot_product_attention(q, k, v, attn_mask = None) with sdpa_kernel(bk): - return F.scaled_dot_product_attention(q, k, v, attn_mask=None) + return F.scaled_dot_product_attention(q, k, v, attn_mask = None) dms = timed(run_dense) nms = timed(run_none) diff --git a/scripts/video_speedmem_bench.py b/scripts/video_speedmem_bench.py index 5a9e33051e..511a5ee8eb 100644 --- a/scripts/video_speedmem_bench.py +++ b/scripts/video_speedmem_bench.py @@ -84,33 +84,28 @@ _FAMILIES: dict[str, dict[str, Any]] = { # ── 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() @@ -131,12 +126,11 @@ def _lpips_alex(ref_arr, arr): fn = _LP.get("fn") if fn is None: - fn = lpips.LPIPS(net="alex", verbose=False).eval() + 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(): @@ -210,11 +204,10 @@ 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, + device = "cuda", + dtype = torch.bfloat16, + supports_default_torch_compile = True, ) @@ -224,33 +217,59 @@ def _target(): # `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"), + "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"), + "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"), + "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"), + "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" + ), # 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"), + "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" + ), } @@ -258,7 +277,7 @@ 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 = 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. @@ -287,8 +306,16 @@ class _SecondExpertView: 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, - logger=None) -> dict: +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). For a dual-expert MoE (pipe.transformer_2 present) every DiT-touching lever is applied to BOTH experts @@ -305,7 +332,14 @@ def _apply_levers(pipe, cfg: dict, *, fam_name: str, fam_obj, force_fp32_vae: bo ) tgt = _target() - engaged = {"dit": None, "te": None, "vae": None, "attn": None, "cache": None, "speed_optims": {}} + 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``. @@ -316,7 +350,7 @@ def _apply_levers(pipe, cfg: dict, *, fam_name: str, fam_obj, force_fp32_vae: bo # DiT quant (pipeline kind, resident): mutates each expert's transformer in place. if cfg["dit"] not in ("none", "off"): schemes = [ - quantize_transformer(v, tgt, mode=cfg["dit"], family=fam_name, logger=logger) + quantize_transformer(v, tgt, mode = cfg["dit"], family = fam_name, logger = logger) for v in views ] engaged["dit"] = schemes[0] @@ -327,15 +361,20 @@ def _apply_levers(pipe, cfg: dict, *, fam_name: str, fam_obj, force_fp32_vae: bo # 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 + 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, + pipe, + tgt, + mode = cfg["vae"], + family = fam_name, + offload_active = False, + force_fp32 = force_fp32_vae, + logger = logger, ) _empty() @@ -356,29 +395,51 @@ def _apply_levers(pipe, cfg: dict, *, fam_name: str, fam_obj, force_fp32_vae: bo if cache_request is not None: for v in views: engaged["cache"] = apply_step_cache( - v, mode=cache_request, threshold=None, - quant_active=dit_quant_active, logger=logger, + v, + mode = cache_request, + threshold = None, + quant_active = dit_quant_active, + logger = logger, ) cache_active = engaged["cache"] not in (None, "off") # Attention (per expert). - backend = select_attention_backend(tgt, cfg["attn"], speed_active=speed_active) + backend = select_attention_backend(tgt, cfg["attn"], speed_active = speed_active) for v in views: - engaged["attn"] = apply_attention_backend(v, backend, logger=logger) + 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, + v, + 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, guidance_via_guider=False, logger=None): +def _timed_video( + pipe, + *, + steps, + width, + height, + num_frames, + guidance, + seed, + cache_mode, + dit_quant_active, + default_steps, + guidance_via_guider = False, + 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 @@ -388,12 +449,12 @@ def _timed_video(pipe, *, steps, width, height, num_frames, guidance, seed, cach if cache_mode == "auto": try: maybe_toggle_step_cache( - pipe, steps=steps, quant_active=dit_quant_active, threshold=None, logger=logger + pipe, steps = steps, quant_active = dit_quant_active, threshold = None, logger = logger ) except Exception: pass - g = torch.Generator(device="cuda").manual_seed(seed) + g = torch.Generator(device = "cuda").manual_seed(seed) step_ts: list[float] = [] last = [0.0] @@ -406,8 +467,12 @@ def _timed_video(pipe, *, steps, width, height, num_frames, guidance, seed, cach return kw kwargs = dict( - prompt=PROMPT, width=width, height=height, num_frames=num_frames, - num_inference_steps=steps, generator=g, + prompt = PROMPT, + width = width, + height = height, + num_frames = num_frames, + num_inference_steps = steps, + generator = g, ) if guidance_via_guider: # HunyuanVideo-1.5: CFG lives on a guider component and __call__ takes no @@ -429,8 +494,20 @@ def _timed_video(pipe, *, steps, width, height, num_frames, guidance, seed, cach 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): +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 @@ -443,21 +520,26 @@ def _run_config(name: str, cfg: dict, *, family: str, steps: int, width: int, he default_steps = getattr(fam_obj, "default_steps", 50) gvg = bool(getattr(fam_obj, "guidance_via_guider", False)) - _empty(); _reset_peak() + _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 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, + 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() @@ -466,18 +548,36 @@ def _run_config(name: str, cfg: dict, *, family: str, steps: int, width: int, he # 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, - guidance_via_guider=gvg, logger=logger, + 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, + 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, - guidance_via_guider=gvg, logger=logger, + 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, + logger = logger, ) dts.append(dt) steps_ms.append(_median(st) if st else 0.0) @@ -487,7 +587,6 @@ def _run_config(name: str, cfg: dict, *, family: str, steps: int, width: int, he try: if arrs: from PIL import Image - Image.fromarray(arrs[len(arrs) // 2]).save(out / f"vid_{family}_{name}.png") except Exception: pass @@ -496,7 +595,6 @@ def _run_config(name: str, cfg: dict, *, family: str, steps: int, width: int, he if name == "reference" and arrs: try: import numpy as _np - _np.savez_compressed(out / "ref_frames.npz", *arrs) except Exception: pass @@ -526,35 +624,39 @@ def _run_config(name: str, cfg: dict, *, family: str, steps: int, width: int, he 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") +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") + logging.basicConfig(level = logging.INFO, format = "%(message)s") logger = logging.getLogger("videobench") out = Path(args.out) - out.mkdir(parents=True, exist_ok=True) + 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) + 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 @@ -564,22 +666,32 @@ def main(argv=None) -> int: 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, + 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) + 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) @@ -589,9 +701,9 @@ def main(argv=None) -> int: ) 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) + with open(dest, "w", encoding = "utf-8") as fh: + json.dump(rows, fh, indent = 2) + print(f"wrote {dest}", flush = True) return 0 diff --git a/studio/backend/core/inference/diffusion_attention.py b/studio/backend/core/inference/diffusion_attention.py index 30849ada65..e70a542e5d 100644 --- a/studio/backend/core/inference/diffusion_attention.py +++ b/studio/backend/core/inference/diffusion_attention.py @@ -437,9 +437,9 @@ def _null_mask_processor_cls(): self, attn, hidden_states, - encoder_hidden_states=None, - attention_mask=None, - image_rotary_emb=None, + encoder_hidden_states = None, + attention_mask = None, + image_rotary_emb = None, ): # Fast path only when the pre-hook removed all padding (attn_mask redundant); a # constant python bool so torch.compile const-folds the branch (no graph break). @@ -447,9 +447,9 @@ def _null_mask_processor_cls(): return super().__call__( attn, hidden_states, - encoder_hidden_states=encoder_hidden_states, - attention_mask=attention_mask, - image_rotary_emb=image_rotary_emb, + encoder_hidden_states = encoder_hidden_states, + attention_mask = attention_mask, + image_rotary_emb = image_rotary_emb, ) # Null path = the stock body with the mask block removed and attn_mask=None. @@ -466,9 +466,8 @@ def _null_mask_processor_cls(): if image_rotary_emb is not None: from diffusers.models.embeddings import apply_rotary_emb - - query = apply_rotary_emb(query, image_rotary_emb, sequence_dim=1) - key = apply_rotary_emb(key, image_rotary_emb, sequence_dim=1) + query = apply_rotary_emb(query, image_rotary_emb, sequence_dim = 1) + key = apply_rotary_emb(key, image_rotary_emb, sequence_dim = 1) if encoder_hidden_states is not None: encoder_query = attn.add_q_proj(encoder_hidden_states) @@ -484,19 +483,19 @@ def _null_mask_processor_cls(): if attn.norm_added_k is not None: encoder_key = attn.norm_added_k(encoder_key) - query = torch.cat([query, encoder_query], dim=1) - key = torch.cat([key, encoder_key], dim=1) - value = torch.cat([value, encoder_value], dim=1) + query = torch.cat([query, encoder_query], dim = 1) + key = torch.cat([key, encoder_key], dim = 1) + value = torch.cat([value, encoder_value], dim = 1) hidden_states = dispatch_attention_fn( query, key, value, - attn_mask=None, - dropout_p=0.0, - is_causal=False, - backend=self._attention_backend, - parallel_config=self._parallel_config, + attn_mask = None, + dropout_p = 0.0, + is_causal = False, + backend = self._attention_backend, + parallel_config = self._parallel_config, ) hidden_states = hidden_states.flatten(2, 3) @@ -530,7 +529,7 @@ def _trim_stream(states, mask): if states is None or mask is None or mask.dim() != 2: return states, mask, True # nothing to mask -> treat as no-padding mb = mask.bool() - keep = mb.any(dim=0) # column valid for at least one batch element + keep = mb.any(dim = 0) # column valid for at least one batch element if not bool(keep.all()): states = states[:, keep] mask = mask[:, keep] @@ -642,7 +641,12 @@ def _install_null_processors(dit: Any, logger: Any) -> bool: return installed > 0 -def install_hunyuan_attention_trim(pipe: Any, family: Any, *, logger: Any = None) -> bool: +def install_hunyuan_attention_trim( + pipe: Any, + family: Any, + *, + logger: Any = None, +) -> bool: """HunyuanVideo-1.5 only: make the joint attention skip padded text tokens (see module note). Installs a null-mask processor on every denoiser DiT block plus an eager pre-hook that trims @@ -662,7 +666,7 @@ def install_hunyuan_attention_trim(pipe: Any, family: Any, *, logger: Any = None continue if getattr(dit, "_unsloth_trim_hook", None) is None: try: - handle = dit.register_forward_pre_hook(_hunyuan_trim_pre_hook, with_kwargs=True) + handle = dit.register_forward_pre_hook(_hunyuan_trim_pre_hook, with_kwargs = True) dit._unsloth_trim_hook = handle except Exception as exc: # noqa: BLE001 — optimisation only _warn(logger, "hunyuan_attn_trim", exc) diff --git a/studio/backend/core/inference/diffusion_prequant.py b/studio/backend/core/inference/diffusion_prequant.py index ecce54efab..e7ff6de402 100644 --- a/studio/backend/core/inference/diffusion_prequant.py +++ b/studio/backend/core/inference/diffusion_prequant.py @@ -319,6 +319,7 @@ def _validate_checkpoint( ckpt_excludes = meta.get("exclude_name_tokens") if ckpt_excludes is not None: from .diffusion_transformer_quant import exclude_tokens_for_scheme + # The exclude set is derived from scheme AND family (fp8 keeps a family's padded-conditioning # embedder bf16); use the checkpoint's recorded family so an artifact baked under an older # token list (e.g. a Wan fp8 checkpoint built before the condition_embedder exclude) is diff --git a/studio/backend/core/inference/diffusion_vae_quant.py b/studio/backend/core/inference/diffusion_vae_quant.py index 089a8b3f04..aeb7823012 100644 --- a/studio/backend/core/inference/diffusion_vae_quant.py +++ b/studio/backend/core/inference/diffusion_vae_quant.py @@ -73,12 +73,12 @@ _VAE_AUTO_LADDER = (VAE_QUANT_FP8,) # ltx-2 SSIM 0.942. FLUX.2 (klein/dev) and Hunyuan-1.5 pass fp8_dynamic, so they are not denied. # The vae_force_fp32 video families (Wan) are gated separately at the loader (never quantise). _VAE_FAMILY_SCHEME_DENY: dict[str, frozenset[str]] = { - "sdxl": frozenset({VAE_QUANT_FP8, VAE_QUANT_FP8_DYNAMIC}), - "flux.1": frozenset({VAE_QUANT_FP8_DYNAMIC}), - "flux.1-kontext": frozenset({VAE_QUANT_FP8_DYNAMIC}), - "qwen-image": frozenset({VAE_QUANT_FP8_DYNAMIC}), - "qwen-image-edit": frozenset({VAE_QUANT_FP8_DYNAMIC}), - "ltx-2": frozenset({VAE_QUANT_FP8_DYNAMIC}), + "sdxl": frozenset({VAE_QUANT_FP8, VAE_QUANT_FP8_DYNAMIC}), + "flux.1": frozenset({VAE_QUANT_FP8_DYNAMIC}), + "flux.1-kontext": frozenset({VAE_QUANT_FP8_DYNAMIC}), + "qwen-image": frozenset({VAE_QUANT_FP8_DYNAMIC}), + "qwen-image-edit": frozenset({VAE_QUANT_FP8_DYNAMIC}), + "ltx-2": frozenset({VAE_QUANT_FP8_DYNAMIC}), } # Cache of device -> bool for the fp8_dynamic conv smoke probe (run once per device). @@ -255,7 +255,10 @@ def quantize_vae( # leaves it dense; the big video Conv3d VAEs clear the floor and quantise. Explicit # requests below skip this gate (the user asked for it directly). if _vae_param_bytes(vae) < _VAE_AUTO_MIN_BYTES: - _note(logger, "vae auto: VAE under the ~1GB size floor; staying dense (quant saves ~nothing)") + _note( + logger, + "vae auto: VAE under the ~1GB size floor; staying dense (quant saves ~nothing)", + ) return None mode = select_vae_quant_scheme( target, diff --git a/studio/backend/core/inference/video.py b/studio/backend/core/inference/video.py index 5b36ab677f..9febd1b62e 100644 --- a/studio/backend/core/inference/video.py +++ b/studio/backend/core/inference/video.py @@ -1353,7 +1353,7 @@ class VideoBackend: # attention so it runs the fused (cuDNN/flash) SDPA kernel instead of the dense-mask # fallback (~18x/DiT-forward at 121 frames, cosine ~1.0). Must precede the backend set # so the requested kernel pins onto the new processors. No-op for every other family. - trim = install_hunyuan_attention_trim(view, fam, logger=logger) + trim = install_hunyuan_attention_trim(view, fam, logger = logger) engaged = apply_attention_backend( view, select_attention_backend( diff --git a/studio/backend/tests/test_diffusion_attention.py b/studio/backend/tests/test_diffusion_attention.py index 13a183d251..8ec592634a 100644 --- a/studio/backend/tests/test_diffusion_attention.py +++ b/studio/backend/tests/test_diffusion_attention.py @@ -462,7 +462,7 @@ def test_trim_stream_layout_agnostic_drops_only_global_padding(): def test_trim_stream_full_mask_is_noop(): states = torch.ones(1, 4, 2) - mask = torch.ones(1, 4, dtype=torch.long) + mask = torch.ones(1, 4, dtype = torch.long) out_s, out_m, all_valid = att._trim_stream(states, mask) assert out_s.shape == (1, 4, 2) and all_valid is True @@ -483,9 +483,9 @@ def test_trim_stream_mixed_batch_not_all_valid(): assert all_valid is False -def _fake_dit(n_blocks=2): - blocks = [types.SimpleNamespace(attn=types.SimpleNamespace()) for _ in range(n_blocks)] - return types.SimpleNamespace(transformer_blocks=blocks) +def _fake_dit(n_blocks = 2): + blocks = [types.SimpleNamespace(attn = types.SimpleNamespace()) for _ in range(n_blocks)] + return types.SimpleNamespace(transformer_blocks = blocks) def test_trim_pre_hook_empties_t2v_image_and_trims_and_flags(): @@ -508,7 +508,7 @@ def test_trim_stream_all_invalid_yields_empty_but_valid(): # A fully-padded secondary stream (e.g. unused byt5 in t2v) trims to 0 length and reports # all_valid True (vacuous) so it does NOT drop the fast path -- it just contributes no tokens. states = torch.ones(1, 5, 2) - mask = torch.zeros(1, 5, dtype=torch.long) + mask = torch.zeros(1, 5, dtype = torch.long) out_s, out_m, all_valid = att._trim_stream(states, mask) assert out_s.shape == (1, 0, 2) and all_valid is True @@ -522,7 +522,7 @@ def test_trim_pre_hook_byt5_all_invalid_keeps_fast_path(): "encoder_hidden_states": torch.arange(4.0).reshape(1, 4, 1), "encoder_attention_mask": torch.tensor([[1, 1, 1, 0]]), "encoder_hidden_states_2": torch.ones(1, 6, 1), - "encoder_attention_mask_2": torch.zeros(1, 6, dtype=torch.long), # all padding + "encoder_attention_mask_2": torch.zeros(1, 6, dtype = torch.long), # all padding } _, out = att._hunyuan_trim_pre_hook(dit, (), kwargs) assert out["encoder_hidden_states"].shape == (1, 3, 1) @@ -538,7 +538,7 @@ def test_trim_pre_hook_empty_primary_reverts_and_disables(): kwargs = { "image_embeds": torch.zeros(1, 5, 3), "encoder_hidden_states": mllm, - "encoder_attention_mask": torch.zeros(1, 4, dtype=torch.long), # 0 valid + "encoder_attention_mask": torch.zeros(1, 4, dtype = torch.long), # 0 valid } _, out = att._hunyuan_trim_pre_hook(dit, (), kwargs) assert out["encoder_hidden_states"] is mllm # reverted (not emptied) @@ -590,14 +590,14 @@ def test_trim_pre_hook_absent_stream_not_written_back(): def test_install_trim_noop_for_non_hunyuan_family(): - fam = types.SimpleNamespace(transformer_class="WanTransformer3DModel") - pipe = types.SimpleNamespace(transformer=types.SimpleNamespace()) + fam = types.SimpleNamespace(transformer_class = "WanTransformer3DModel") + pipe = types.SimpleNamespace(transformer = types.SimpleNamespace()) assert att.install_hunyuan_attention_trim(pipe, fam) is False def test_install_trim_noop_when_transformer_class_mismatch(): # Family claims Hunyuan but the loaded module isn't -> no processors touched, no diffusers # import; returns False rather than swapping an unknown attention processor. - fam = types.SimpleNamespace(transformer_class="HunyuanVideo15Transformer3DModel") - pipe = types.SimpleNamespace(transformer=types.SimpleNamespace()) # class name mismatch + fam = types.SimpleNamespace(transformer_class = "HunyuanVideo15Transformer3DModel") + pipe = types.SimpleNamespace(transformer = types.SimpleNamespace()) # class name mismatch assert att.install_hunyuan_attention_trim(pipe, fam) is False diff --git a/studio/backend/tests/test_diffusion_transformer_quant.py b/studio/backend/tests/test_diffusion_transformer_quant.py index 4d33fc7937..ff29c7f4e6 100644 --- a/studio/backend/tests/test_diffusion_transformer_quant.py +++ b/studio/backend/tests/test_diffusion_transformer_quant.py @@ -608,7 +608,10 @@ def test_family_deny_auto_skips_fp8_for_hunyuan(monkeypatch): _stub_torch(monkeypatch, cc = (10, 0)) _allow(monkeypatch, {TQ_FP8, TQ_NVFP4, TQ_MXFP8, TQ_INT8}) assert select_transformer_quant_scheme(_target(), "auto", family = "hunyuanvideo-1.5") == TQ_INT8 - assert select_transformer_quant_scheme(_target(), "auto", family = "hunyuanvideo-1.5-720p") == TQ_INT8 + assert ( + select_transformer_quant_scheme(_target(), "auto", family = "hunyuanvideo-1.5-720p") + == TQ_INT8 + ) # ltx-2 keeps the ladder head (fp8) -- it is not a black-frame family. assert select_transformer_quant_scheme(_target(), "auto", family = "ltx-2") == TQ_FP8 @@ -690,7 +693,12 @@ def test_quantize_transformer_fp8_wan_excludes_condition_embedder(monkeypatch): torch.bfloat16 = "bfloat16" class _Linear: - def __init__(self, inf, outf, dtype = "bfloat16"): + def __init__( + self, + inf, + outf, + dtype = "bfloat16", + ): self.in_features, self.out_features = inf, outf self.weight = types.SimpleNamespace(dtype = dtype) @@ -711,4 +719,6 @@ def test_quantize_transformer_fp8_wan_excludes_condition_embedder(monkeypatch): assert filt(big, "condition_embedder.time_embedder.linear_1") is False assert filt(big, "blocks.0.attn1.to_q") is True assert filt(big, "blocks.0.ffn.net.0.proj") is True - assert filt(big, "blocks.0.attn2.to_k") is True # cross-attn K/V stay fp8 (embedder bias rescues rows) + assert ( + filt(big, "blocks.0.attn2.to_k") is True + ) # cross-attn K/V stay fp8 (embedder bias rescues rows) diff --git a/studio/backend/tests/test_diffusion_vae_quant.py b/studio/backend/tests/test_diffusion_vae_quant.py index d3748d3ad2..08c057f053 100644 --- a/studio/backend/tests/test_diffusion_vae_quant.py +++ b/studio/backend/tests/test_diffusion_vae_quant.py @@ -250,7 +250,11 @@ def test_fp8_dynamic_conv_filter(monkeypatch): ff = captured["filter_fn"] nn = torch.nn - def _mod(cls, shape, kernel_size = None): + def _mod( + cls, + shape, + kernel_size = None, + ): m = cls() m.weight = _Weight(shape) if kernel_size is not None: @@ -264,7 +268,10 @@ def test_fp8_dynamic_conv_filter(monkeypatch): assert ff(_mod(nn.Linear, (512, 512)), "decoder.mid_block.attentions.0.to_q") is True # POINTWISE (1x1 / 1x1x1) convs excluded even at %16 channels: torchao 0.17's fp8 conv # kernel rejects them ("Activation and filter channels must match") -> crash at decode. - assert ff(_mod(nn.Conv2d, (128, 128, 1, 1), (1, 1)), "decoder.mid_block.attentions.0.proj_conv") is False + assert ( + ff(_mod(nn.Conv2d, (128, 128, 1, 1), (1, 1)), "decoder.mid_block.attentions.0.proj_conv") + is False + ) assert ff(_mod(nn.Conv3d, (64, 64, 1, 1, 1), (1, 1, 1)), "decoder.time_mix.conv") is False # Channels not a multiple of 16 excluded (torchao would skip them regardless): the RGB # in/out head (C=3) and any off-16 dim. @@ -388,7 +395,9 @@ def test_quantize_vae_auto_resolves_and_applies(monkeypatch): _stub_torch(monkeypatch, cc = (10, 0)) _stub_capability(monkeypatch, (10, 0)) _allow_vae(monkeypatch, {VAE_QUANT_FP8_DYNAMIC, VAE_QUANT_FP8}) - monkeypatch.setattr(vq, "_cast_vae_fp8_dynamic", lambda v, t: pytest.fail("auto must not use fp8_dynamic")) + monkeypatch.setattr( + vq, "_cast_vae_fp8_dynamic", lambda v, t: pytest.fail("auto must not use fp8_dynamic") + ) calls: list = [] monkeypatch.setattr(vq, "_cast_vae_fp8", lambda v, t: calls.append(v)) vae = object() @@ -406,7 +415,9 @@ def test_quantize_vae_auto_size_gate_skips_small(monkeypatch): _allow_vae(monkeypatch, {VAE_QUANT_FP8_DYNAMIC, VAE_QUANT_FP8}) monkeypatch.setattr(vq, "_vae_param_bytes", lambda v: 200_000_000) # ~0.2 GB image VAE monkeypatch.setattr(vq, "_cast_vae_fp8", lambda v, t: pytest.fail("small VAE must stay dense")) - monkeypatch.setattr(vq, "_cast_vae_fp8_dynamic", lambda v, t: pytest.fail("small VAE must stay dense")) + monkeypatch.setattr( + vq, "_cast_vae_fp8_dynamic", lambda v, t: pytest.fail("small VAE must stay dense") + ) pipe = types.SimpleNamespace(vae = object()) assert quantize_vae(pipe, _target(), mode = "auto", family = "flux.1") is None @@ -442,13 +453,18 @@ def test_quantize_vae_explicit_fp8_dynamic_probe_gates(monkeypatch): _stub_torch(monkeypatch, cc = (10, 0)) _allow_vae(monkeypatch, {VAE_QUANT_FP8_DYNAMIC, VAE_QUANT_FP8}) monkeypatch.setattr(vq, "_vae_fp8_dynamic_probe", lambda device: False) - monkeypatch.setattr(vq, "_cast_vae_fp8_dynamic", lambda v, t: pytest.fail("probe failed: must not cast")) + monkeypatch.setattr( + vq, "_cast_vae_fp8_dynamic", lambda v, t: pytest.fail("probe failed: must not cast") + ) pipe = types.SimpleNamespace(vae = object()) assert quantize_vae(pipe, _target(), mode = "fp8_dynamic", family = "flux.2-klein") is None monkeypatch.setattr(vq, "_vae_fp8_dynamic_probe", lambda device: True) calls: list = [] monkeypatch.setattr(vq, "_cast_vae_fp8_dynamic", lambda v, t: calls.append(v)) - assert quantize_vae(pipe, _target(), mode = "fp8_dynamic", family = "flux.2-klein") == VAE_QUANT_FP8_DYNAMIC + assert ( + quantize_vae(pipe, _target(), mode = "fp8_dynamic", family = "flux.2-klein") + == VAE_QUANT_FP8_DYNAMIC + ) assert len(calls) == 1 @@ -467,5 +483,11 @@ def test_real_family_deny_list_policy(monkeypatch): assert select_vae_quant_scheme(_target(), "auto", family = "ltx-2") == VAE_QUANT_FP8 assert select_vae_quant_scheme(_target(), "fp8_dynamic", family = "flux.1") is None # FLUX.2 / Hunyuan keep fp8_dynamic available as an explicit opt-in (measured in-bar). - assert select_vae_quant_scheme(_target(), "fp8_dynamic", family = "flux.2-klein") == VAE_QUANT_FP8_DYNAMIC - assert select_vae_quant_scheme(_target(), "fp8_dynamic", family = "hunyuanvideo-1.5") == VAE_QUANT_FP8_DYNAMIC + assert ( + select_vae_quant_scheme(_target(), "fp8_dynamic", family = "flux.2-klein") + == VAE_QUANT_FP8_DYNAMIC + ) + assert ( + select_vae_quant_scheme(_target(), "fp8_dynamic", family = "hunyuanvideo-1.5") + == VAE_QUANT_FP8_DYNAMIC + )