"""FunPack NVFP4 on-the-fly model patching (EXPERIMENTAL). Blackwell's 5th-gen tensor cores run NVFP4 (e2m1 with per-16-element FP8-E4M3 block scales + a per-tensor FP32 scale) at roughly 2x FP8 / 4x BF16 GEMM throughput. ComfyUI core already ships the whole runtime for it (comfy.quant_ops layouts + MixedPrecisionOps dispatch, dynamic activation quantization, LoRA fallback, dequant emulation on non-FP4 GPUs) - but only for checkpoints that were quantized ahead of time. Stock loaders offer fp8 at best, and on sm_120 the fp8 path is mostly storage, not tensor-core speedup. This module closes that gap: it quantizes an existing bf16/fp16 LTX-AV checkpoint to NVFP4 *at load time* as a pure state-dict transform - selected 2D Linear weights are re-serialized in core's exact quantized-checkpoint format (weight/weight_scale/ weight_scale_2 + per-layer `comfy_quant` metadata), then handed to the normal `load_diffusion_model_state_dict` path. Core's own detection routes the model through MixedPrecisionOps; every layer we did NOT tag simply stays at compute dtype. No custom forward, no monkeypatching, no new inference code paths to maintain. Zone-aware precision map (the part generic quantizers can't do): LTX-AV is a joint audio-video transformer and its zones are not equally robust. By default we quantize only the video-branch Linears (attn1/attn2/ff) of the middle blocks, keeping at full precision: - the audio branch (audio_attn1/audio_attn2/audio_ff) - joint attention is the part that corrupts first (see the self-consistency post-block injection failure), - the cross-modal bridges (audio_to_video_attn / video_to_audio_attn), - the first/last N blocks (standard quantization practice; they carry the most input/output-adjacent signal), - everything that is not a plain 2D matmul (norms, ada modulation, embeddings). Scopes widen from there for experimentation. W4A4 without calibration WILL cost some fidelity - this is an experimental speed/VRAM lever, off unless you wire the node in. Requires: ComfyUI with comfy.quant_ops (comfy_kitchen), torch cu13+ on Blackwell (sm_120+) for real FP4 GEMMs. On anything older the model still loads and runs through core's dequant emulation (no speedup, same quality) - useful for A/B sanity checks. """ from __future__ import annotations import json import logging import os import re import torch # --------------------------------------------------------------------------- # Static input-scale calibration # --------------------------------------------------------------------------- # Core quantizes each layer's INPUT dynamically (amax reduction + quantize) on every # call unless the layer has a static `input_scale` parameter. Live probe on sm_120: # the dynamic path costs ~0.30 ms/layer-call and a provided scale cuts it to ~0.11 ms # (amax dominates) - across ~800 quantized layers that is a top-line cost of every # model call. Calibration is capture-then-freeze: a "calibrate" run records each # quantized layer's input amax via forward-pre-hooks into a per-checkpoint sidecar; # later loads inject `input_scale = amax * MARGIN / (448 * 6)` into the state dict, # where core's loader picks it up natively. MARGIN leaves headroom because NVFP4 # clips activations above the scale; per the disposable-keys philosophy the sidecar # is per-machine and rebuilt in one run whenever absent. CALIBRATION_MARGIN = 1.25 _F8F4_MAX = 448.0 * 6.0 # F8_E4M3_MAX * F4_E2M1_MAX INPUT_SCALE_MODES = ["dynamic", "auto-calibrate"] def _scales_path(checkpoint_path): base = os.path.join(os.path.dirname(os.path.abspath(__file__)), "refinements", "nvfp4_scales") return os.path.join(base, os.path.basename(checkpoint_path) + ".json") def load_calibration(checkpoint_path): """{module_path: input_amax} recorded by a previous calibrate run, or None.""" try: p = _scales_path(checkpoint_path) if not os.path.exists(p): return None with open(p, "r", encoding="utf-8") as f: data = json.load(f) amax = data.get("amax") or {} return {str(k): float(v) for k, v in amax.items()} or None except Exception as e: # noqa: BLE001 logging.warning(f"[FunPack NVFP4] could not read calibration sidecar: {e}") return None def inject_input_scales(sd, target_weight_keys, calibration, margin=CALIBRATION_MARGIN): """Add `.input_scale` next to each quantized weight whose module path has a calibrated amax. Keys are matched by suffix so any checkpoint prefix ('model.diffusion_model.' or none) works. Returns the number injected.""" injected = 0 for k in target_weight_keys: prefix = k[: -len("weight")] # '...transformer_blocks.3.attn1.to_q.' module_path = prefix.rstrip(".") for cal_key, amax in calibration.items(): if module_path == cal_key or module_path.endswith("." + cal_key): sd[prefix + "input_scale"] = torch.tensor( float(amax) * margin / _F8F4_MAX, dtype=torch.float32) injected += 1 break return injected def install_calibration_capture(model, checkpoint_path): """Forward-pre-hooks on every quantized Linear recording the running input amax. The sidecar is written on the first full sweep and refreshed periodically; hooks live for the model instance's lifetime (a calibration run is expected to be slightly SLOWER - it pays the amax it is trying to eliminate, once).""" store = {"amax": {}, "writes": 0, "calls": 0} path = _scales_path(checkpoint_path) def _flush(): try: os.makedirs(os.path.dirname(path), exist_ok=True) tmp = path + ".tmp" with open(tmp, "w", encoding="utf-8") as f: json.dump({"amax": store["amax"], "margin": CALIBRATION_MARGIN}, f) os.replace(tmp, path) store["writes"] += 1 except Exception as e: # noqa: BLE001 logging.warning(f"[FunPack NVFP4] calibration flush failed: {e}") hooked = 0 try: for name, m in model.model.diffusion_model.named_modules(): if getattr(m, "quant_format", None) != "nvfp4": continue def _hook(_module, args, _name=name): try: x = args[0] if args else None if isinstance(x, torch.Tensor): a = float(x.detach().abs().amax()) if a > store["amax"].get(_name, 0.0): store["amax"][_name] = a store["calls"] += 1 if store["calls"] % 4096 == 0: _flush() except Exception: pass m.register_forward_pre_hook(_hook) hooked += 1 except Exception as e: # noqa: BLE001 logging.warning(f"[FunPack NVFP4] calibration capture install failed: {e}") return 0 if hooked: # First flush after one full model call so a single generation is enough. def _first_flush_hook(_module, _args): if store["writes"] == 0 and len(store["amax"]) >= hooked: _flush() try: last = [m for _, m in model.model.diffusion_model.named_modules() if getattr(m, "quant_format", None) == "nvfp4"][-1] last.register_forward_pre_hook(_first_flush_hook) except Exception: pass return hooked QUANTIZE_SCOPES = [ "video blocks", # video-branch Linears in middle blocks (default, safest) "video + cross-modal", # + audio_to_video_attn / video_to_audio_attn bridges "all blocks", # every Linear inside transformer_blocks (incl. audio branch) "all 2D layers", # + eligible Linears outside the blocks (most aggressive) ] # Video-branch submodules inside a BasicAVTransformerBlock (comfy/ldm/lightricks/av_model.py). _VIDEO_MODULES = ("attn1.", "attn2.", "ff.") _CROSS_MODULES = ("audio_to_video_attn.", "video_to_audio_attn.") # Matches "...transformer_blocks.." with any leading prefix # ("model.diffusion_model." or none - load_diffusion_model_state_dict strips prefixes # from ALL keys uniformly, so our metadata keys ride along). _BLOCK_RE = re.compile(r"(?:^|\.)transformer_blocks\.(\d+)\.(.*)$") # Outside the blocks ("all 2D layers" scope) never touch input/output-adjacent or # modulation weights even when they are 2D Linears. _GENERIC_SKIP = ("proj_out", "adaln", "embed", "emb.", "norm", "scale_shift", "patchify") # Below this, quantization overhead outweighs any GEMM win and relative error is worst. _MIN_DIM = 512 def _block_ids(sd) -> list[int]: ids = set() for k in sd: m = _BLOCK_RE.search(k) if m: ids.add(int(m.group(1))) return sorted(ids) def _selected(key_body: str, scope: str, keep: set[int]) -> bool: """Decide whether the Linear at `key_body` (state-dict key minus '.weight') is quantized under `scope`, given the high-precision `keep` block set.""" m = _BLOCK_RE.search(key_body + ".") if not m: if scope != "all 2D layers": return False low = key_body.lower() return not any(s in low for s in _GENERIC_SKIP) idx, rest = int(m.group(1)), m.group(2) if idx in keep: return False if scope in ("all blocks", "all 2D layers"): return True if any(rest.startswith(p) for p in _VIDEO_MODULES): return True if scope == "video + cross-modal" and any(rest.startswith(p) for p in _CROSS_MODULES): return True return False def _eligible(key: str, w) -> bool: return ( key.endswith(".weight") and isinstance(w, torch.Tensor) and w.dim() == 2 and w.shape[0] >= _MIN_DIM and w.shape[1] >= _MIN_DIM and w.shape[1] % 16 == 0 # NVFP4 blocks run along in_features ) def quantize_state_dict_nvfp4(sd, scope="video blocks", keep_first_blocks=2, keep_last_blocks=2, device=None, progress_callback=None): """Return (new_sd, quantized_layer_names, kept_block_ids). Pure transform: selected 2D weights are replaced by NVFP4 QuantizedTensor serialization (weight uint8-packed, weight_scale fp8-e4m3 block scales, weight_scale_2 fp32 tensor scale) plus a per-layer `comfy_quant` metadata key - byte-for-byte the format `comfy.ops._load_quantized_module` consumes. Unselected keys pass through untouched. """ import comfy.quant_ops as quant_ops if not quant_ops._CK_AVAILABLE: raise RuntimeError( "comfy_kitchen is unavailable in this ComfyUI - NVFP4 quantization needs it. " "Update ComfyUI / reinstall requirements." ) if scope not in QUANTIZE_SCOPES: raise ValueError(f"Unknown quantize scope: {scope}") for k in sd: if k.endswith(".comfy_quant") or k.endswith("scaled_fp8"): raise ValueError( "This checkpoint is already quantized (comfy_quant / scaled_fp8 metadata " "found). NVFP4 patching needs the original bf16/fp16 weights." ) ids = _block_ids(sd) keep: set[int] = set() if ids: if keep_first_blocks > 0: keep |= set(ids[:keep_first_blocks]) if keep_last_blocks > 0: keep |= set(ids[-keep_last_blocks:]) targets = [k for k, w in sd.items() if _eligible(k, w) and _selected(k[: -len(".weight")], scope, keep)] meta = torch.tensor(list(json.dumps({"format": "nvfp4"}).encode("utf-8")), dtype=torch.uint8) out = {} done = 0 for k, w in sd.items(): if k not in targets: out[k] = w continue wq = w if wq.dtype not in (torch.float32, torch.bfloat16): # fp16 and fp8-stored checkpoints alike: quantize from bf16. wq = wq.to(torch.bfloat16) if device is not None: wq = wq.to(device) qt = quant_ops.QuantizedTensor.from_float(wq, "TensorCoreNVFP4Layout", scale="recalculate") for kk, vv in qt.state_dict(k).items(): out[kk] = vv.to("cpu") out[k[: -len("weight")] + "comfy_quant"] = meta.clone() del qt, wq done += 1 if progress_callback is not None: progress_callback(done, len(targets)) if done % 100 == 0: logging.info(f"[FunPack NVFP4] quantized {done}/{len(targets)} layers…") return out, targets, keep def log_nvfp4_diagnostics(model, device=None): """Print everything needed to tell WHY NVFP4 is (or is not) fast on this machine. 'Loaded fine' and 'runs fast' are different claims. The three ways the fast path silently degrades, in the order they bite: 1. comfy_kitchen's CUDA backend is disabled when torch is built for CUDA < 13 - quantized matmuls then fall through to the pure-PyTorch eager backend, which is CORRECT but no faster than bf16 (the exact "speed didn't change" symptom). 2. The GPU lacks FP4 tensor cores (pre-Blackwell) - core flags every nvfp4 layer _full_precision_mm and dequantizes (visible in the layer census below). 3. Everything is native but the quantized scope is too small a slice of the total step time to move the needle. A timed 4096x4096 GEMM probe settles it empirically: on a native FP4 path the quantized matmul must beat bf16 clearly; parity or worse means emulation.""" lines = [] try: import comfy_kitchen as ck for name, info in ck.list_backends().items(): if info.get("available"): caps = info.get("capabilities") or [] lines.append(f" ck backend '{name}': AVAILABLE, scaled_mm_nvfp4=" f"{'yes' if 'scaled_mm_nvfp4' in caps else 'NO'}") else: lines.append(f" ck backend '{name}': unavailable ({info.get('unavailable_reason')})") except Exception as e: # noqa: BLE001 lines.append(f" comfy_kitchen introspection failed: {e}") cap = torch.cuda.get_device_capability(device) if torch.cuda.is_available() else None lines.append(f" torch {torch.__version__} | cuda build {torch.version.cuda} | device capability {cap}") if torch.version.cuda is not None: try: if tuple(map(int, str(torch.version.cuda).split(".")))[0] < 13: lines.append(" ⚠ torch is built for CUDA < 13 -> comfy_kitchen DISABLES its CUDA " "backend -> nvfp4 matmuls run on the eager (pure PyTorch) backend: " "correct but NO speedup. Install a cu13x torch build to unlock FP4 GEMMs.") except Exception: # noqa: BLE001 pass n_q = n_emul = 0 try: for m in model.model.diffusion_model.modules(): if getattr(m, "quant_format", None) == "nvfp4": n_q += 1 if getattr(m, "_full_precision_mm", False): n_emul += 1 lines.append(f" nvfp4 layers in loaded model: {n_q}" + (f" — {n_emul} flagged FULL-PRECISION EMULATION (no FP4 compute on this GPU)" if n_emul else " — all on the quantized matmul path")) except Exception as e: # noqa: BLE001 lines.append(f" layer census failed: {e}") if torch.cuda.is_available(): try: import time import comfy.quant_ops as quant_ops d = torch.device(device) if device is not None else torch.device("cuda") w = torch.randn(4096, 4096, dtype=torch.bfloat16, device=d) qw = quant_ops.QuantizedTensor.from_float(w, "TensorCoreNVFP4Layout", scale="recalculate") def _timed(fn, n=30): for _ in range(5): fn() torch.cuda.synchronize(d) t0 = time.perf_counter() for _ in range(n): fn() torch.cuda.synchronize(d) return (time.perf_counter() - t0) / n # Decomposed probe at two token counts: M=1024 (small scene chunk) and # M=16384 (realistic video token count). Times the input quantization and # the FP4 GEMM SEPARATELY — "no speedup" has two very different fixes # depending on which one is at fault: # * GEMM alone beats bf16, quantization dominates -> per-call overhead # problem (x1232 layers/step); needs static input scales or fusion. # * GEMM alone loses to bf16 -> the FP4 kernel isn't engaging at all. results = {} for m in (1024, 16384): x = torch.randn(m, 4096, dtype=torch.bfloat16, device=d) t_bf16 = _timed(lambda: torch.nn.functional.linear(x, w)) t_quant = _timed(lambda: quant_ops.QuantizedTensor.from_float(x, "TensorCoreNVFP4Layout")) # Same quantization with a PREcomputed static scale: isolates the amax # reduction (which per-layer input_scale calibration could remove) from # tensor-construction/kernel-launch overhead (which it can't). s = (x.abs().amax() / (448.0 * 6.0)).float() t_quant_s = _timed(lambda: quant_ops.QuantizedTensor.from_float(x, "TensorCoreNVFP4Layout", scale=s)) qx = quant_ops.QuantizedTensor.from_float(x, "TensorCoreNVFP4Layout") t_gemm = _timed(lambda: torch.nn.functional.linear(qx, qw)) total = t_quant + t_gemm results[m] = (t_bf16, t_quant, t_quant_s, t_gemm) lines.append( f" probe 4096x4096 @ M={m}: bf16 {t_bf16 * 1e3:.3f} ms | " f"input-quant {t_quant * 1e3:.3f} ms (static-scale {t_quant_s * 1e3:.3f} ms) " f"+ fp4-gemm {t_gemm * 1e3:.3f} ms = {total * 1e3:.3f} ms " f"({t_bf16 / max(total, 1e-9):.2f}x net, {t_bf16 / max(t_gemm, 1e-9):.2f}x gemm-only)") del x, qx # Verdict from the LARGE-M run (realistic video token counts); small-M numbers # are launch-overhead dominated by nature and must not drive the conclusion. t_bf16, t_quant, t_quant_s, t_gemm = results[16384] total = t_quant + t_gemm if t_gemm < t_bf16 * 0.8 and total < t_bf16 * 0.8: lines.append(" verdict: FAST PATH ACTIVE — FP4 GEMMs win at real token counts. " "If end-to-end speed still looks flat, the input-quant overhead " f"(~{t_quant * 1e3:.2f} ms/layer-call, largely size-independent) is " "cancelling the win on SMALL layers: prefer scope 'video blocks' " "(skips the audio branch, whose matmuls are too small to profit).") if t_quant_s < t_quant * 0.6: lines.append(" note: static input scales would reclaim most of the input-quant " "cost (amax dominates) — calibration support is worth building.") else: lines.append(" note: static input scales would NOT help much " "(overhead is construction/launch, not amax).") elif t_gemm < t_bf16 * 0.8: lines.append(" verdict: FP4 GEMM fast but INPUT QUANTIZATION dominates even at " "large M — overhead-bound; scope down to the largest layers only.") else: lines.append(" verdict: FP4 GEMM not faster than bf16 even at M=16384 — the " "kernel is not engaging; report these numbers upstream.") except Exception as e: # noqa: BLE001 lines.append(f" GEMM probe failed: {e}") else: lines.append(" GEMM probe skipped (no CUDA device)") print("[FunPack NVFP4] diagnostics:\n" + "\n".join(lines)) return {"nvfp4_layers": n_q, "emulated_layers": n_emul} class FunPackNVFP4ModelLoader: """EXPERIMENTAL: load a diffusion model with on-the-fly NVFP4 quantization. Drop-in replacement for 'Load Diffusion Model' aimed at LTX-AV on Blackwell (sm_120+): selected transformer Linears run real FP4 tensor-core GEMMs (~2x fp8 throughput) and weight VRAM drops to ~1/4 of bf16. Audio branch, cross-modal bridges, first/last blocks, and all non-matmul weights stay at full precision by default. On GPUs without FP4 compute the model still works via emulation. """ CATEGORY = "FunPack/Experimental" RETURN_TYPES = ("MODEL",) RETURN_NAMES = ("model",) FUNCTION = "load" DESCRIPTION = __doc__ @classmethod def INPUT_TYPES(cls): import folder_paths return { "required": { "unet_name": (folder_paths.get_filename_list("diffusion_models"), {"tooltip": "Original bf16/fp16 checkpoint - NOT a pre-quantized one."}), "quantize_scope": (QUANTIZE_SCOPES, { "default": "video blocks", "tooltip": "Which Linears go NVFP4. 'video blocks' keeps the whole audio " "path + cross-modal bridges at full precision (safest). Widen " "only if quality holds.", }), "keep_first_blocks": ("INT", {"default": 2, "min": 0, "max": 16, "tooltip": "First N transformer blocks stay full precision."}), "keep_last_blocks": ("INT", {"default": 2, "min": 0, "max": 16, "tooltip": "Last N transformer blocks stay full precision."}), "input_scales": (INPUT_SCALE_MODES, { "default": "dynamic", "tooltip": "auto-calibrate: first run records each layer's input range " "(slightly slower) into a per-checkpoint sidecar; every later " "load injects static input scales, removing the per-call amax " "cost (~0.2 ms x every quantized layer x every step). Reload " "the model after the calibration run to activate.", }), } } def load(self, unet_name, quantize_scope, keep_first_blocks, keep_last_blocks, input_scales="dynamic"): import comfy.model_management as mm import comfy.sd import comfy.utils import folder_paths path = folder_paths.get_full_path_or_raise("diffusion_models", unet_name) # return_metadata is NOT optional here: LTX-AV variants carry their architecture # config in the safetensors metadata, and detection without it builds the WRONG # model (6 vs 9 ada params, connector dims) -> size-mismatch wall on load. sd, metadata = comfy.utils.load_torch_file(path, return_metadata=True) if mm.supports_nvfp4_compute(mm.get_torch_device()): logging.info("[FunPack NVFP4] FP4 tensor-core compute available - native GEMM path.") else: logging.warning( "[FunPack NVFP4] This GPU has no FP4 compute (needs Blackwell, sm_100+). " "The model will load and run through dequant emulation: correct output, " "no speedup. Use this only for A/B quality checks." ) device = mm.get_torch_device() if torch.cuda.is_available() else None pbar = comfy.utils.ProgressBar(1) def _tick(done, total): if total: pbar.total = total pbar.update_absolute(done) sd, targets, keep = quantize_state_dict_nvfp4( sd, scope=quantize_scope, keep_first_blocks=keep_first_blocks, keep_last_blocks=keep_last_blocks, device=device, progress_callback=_tick, ) if not targets: raise RuntimeError( "No layers matched the NVFP4 quantization scope - is this an LTX-AV " "diffusion model checkpoint? (Expected transformer_blocks.* keys.)" ) logging.info( f"[FunPack NVFP4] quantized {len(targets)} Linear layers " f"(scope='{quantize_scope}', full-precision blocks: {sorted(keep) or 'none'})." ) calibration = load_calibration(path) if input_scales == "auto-calibrate" else None if calibration: injected = inject_input_scales(sd, targets, calibration) logging.info( f"[FunPack NVFP4] static input scales ACTIVE on {injected}/{len(targets)} layers " f"(calibrated {len(calibration)}, margin {CALIBRATION_MARGIN}x) — per-call input " "amax eliminated. Delete the sidecar under refinements/nvfp4_scales/ to recalibrate." ) model = comfy.sd.load_diffusion_model_state_dict(sd, metadata=metadata) if model is None: raise RuntimeError("ComfyUI could not detect the model type after NVFP4 patching.") if input_scales == "auto-calibrate" and not calibration: hooked = install_calibration_capture(model, path) logging.info( f"[FunPack NVFP4] CALIBRATING input scales on {hooked} layers this session " "(slightly slower). Run one full generation, then reload the model (change any " "widget or restart) — static scales activate on the next load." ) try: log_nvfp4_diagnostics(model, device) except Exception as e: # noqa: BLE001 logging.warning(f"[FunPack NVFP4] diagnostics failed: {e}") return (model,)