From 3dd2d33f1d39df1e67f99f0e11c6e711640e2cd7 Mon Sep 17 00:00:00 2001 From: digital-garbage <133010426+digital-garbage@users.noreply.github.com> Date: Thu, 2 Jul 2026 06:16:08 +0300 Subject: [PATCH] model: experimental NVFP4 on-the-fly patching for LTX-AV on Blackwell FunPackNVFP4ModelLoader quantizes an existing bf16/fp16 checkpoint to NVFP4 at load time as a pure state-dict transform: selected 2D Linears are re-serialized in ComfyUI core's quantized-checkpoint format (packed fp4 weights + fp8-e4m3 block scales + fp32 tensor scale + per-layer comfy_quant metadata) and handed to the stock loading path. Core's MixedPrecisionOps then does everything: real FP4 tensor-core GEMMs on sm_100+ (sm_120 incl.), dynamic activation scales (no calibration pass), per-layer LoRA fallback, and dequant emulation on older GPUs. No custom forward, no monkeypatching. Zone-aware precision map for the joint AV transformer: default scope quantizes only video-branch Linears (attn1/attn2/ff) of the middle blocks; the audio branch, the audio<->video cross-modal bridges, the first/last N blocks, and all non-matmul weights stay full precision. Scopes widen to +cross-modal, all blocks, and all 2D layers for experimentation. Tests drive the REAL comfy.quant_ops/comfy.ops via comfy_kitchen's eager backend on CPU: scope selection, serialized-format contract, prefix handling, refusal of pre-quantized checkpoints, and a load+forward roundtrip through MixedPrecisionOps on both the NVFP4 and emulation paths. The Movie Editor Models modal picks the node up automatically as a Unet/Diffusion Model candidate (object_info-driven). Co-Authored-By: Claude Fable 5 --- __init__.py | 7 ++ nvfp4_patch.py | 253 ++++++++++++++++++++++++++++++++++++++ tests/test_nvfp4_patch.py | 163 ++++++++++++++++++++++++ 3 files changed, 423 insertions(+) create mode 100644 nvfp4_patch.py create mode 100644 tests/test_nvfp4_patch.py diff --git a/__init__.py b/__init__.py index c3741389..8104589e 100644 --- a/__init__.py +++ b/__init__.py @@ -20,6 +20,7 @@ if __package__: FunPackVideoStitch, ) from .model_management import FunPackApplyLoraWeights, FunPackLoraLoader + from .nvfp4_patch import FunPackNVFP4ModelLoader from .samplers import FunPackHybridEuler2SSampler, FunPackDistilledFlowSampler, FunPackLTXAVSceneChainSampler from .templates import FunPackRefinementKeyLoader try: @@ -62,6 +63,10 @@ else: except Exception: FunPackApplyLoraWeights = None FunPackLoraLoader = None + try: + from nvfp4_patch import FunPackNVFP4ModelLoader + except Exception: + FunPackNVFP4ModelLoader = None try: from samplers import FunPackHybridEuler2SSampler, FunPackDistilledFlowSampler, FunPackLTXAVSceneChainSampler except Exception: @@ -96,6 +101,7 @@ NODE_CLASS_MAPPINGS = { "FunPackLTXAVSceneChainSampler": FunPackLTXAVSceneChainSampler, "FunPackApplyLoraWeights": FunPackApplyLoraWeights, "FunPackLoraLoader": FunPackLoraLoader, + "FunPackNVFP4ModelLoader": FunPackNVFP4ModelLoader, "FunPackRefinementKeyLoader": FunPackRefinementKeyLoader, } NODE_CLASS_MAPPINGS = {name: cls for name, cls in NODE_CLASS_MAPPINGS.items() if cls is not None} @@ -121,6 +127,7 @@ NODE_DISPLAY_NAME_MAPPINGS = { "FunPackLTXAVSceneChainSampler": "FunPack LTXAV Scene Chain Sampler", "FunPackApplyLoraWeights": "FunPack Apply LoRA Weights", "FunPackLoraLoader": "FunPack LoRA Loader", + "FunPackNVFP4ModelLoader": "FunPack NVFP4 Model Loader (Experimental)", "FunPackRefinementKeyLoader": "FunPack Refinement Key Loader", } diff --git a/nvfp4_patch.py b/nvfp4_patch.py new file mode 100644 index 00000000..e28de5ce --- /dev/null +++ b/nvfp4_patch.py @@ -0,0 +1,253 @@ +"""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 re + +import torch + +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 + + +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."}), + } + } + + def load(self, unet_name, quantize_scope, keep_first_blocks, keep_last_blocks): + 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) + sd = comfy.utils.load_torch_file(path) + + 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'})." + ) + + model = comfy.sd.load_diffusion_model_state_dict(sd) + if model is None: + raise RuntimeError("ComfyUI could not detect the model type after NVFP4 patching.") + return (model,) diff --git a/tests/test_nvfp4_patch.py b/tests/test_nvfp4_patch.py new file mode 100644 index 00000000..0f7a03a1 --- /dev/null +++ b/tests/test_nvfp4_patch.py @@ -0,0 +1,163 @@ +"""NVFP4 on-the-fly patching: transform format + scope selection + core-load roundtrip. + +Uses the REAL comfy.quant_ops / comfy.ops from the local ComfyUI install (comfy_kitchen's +eager backend quantizes and matmuls NVFP4 on CPU), so these tests prove byte-for-byte +compatibility with the format `comfy.ops._load_quantized_module` consumes - not just our +own invariants. Skipped wholesale if ComfyUI or comfy_kitchen is unavailable. +""" +import os +import sys +from pathlib import Path + +import pytest +import torch + +sys.path.insert(0, str(Path(__file__).resolve().parents[1])) + +_COMFY_ROOT = os.environ.get("COMFYUI_ROOT", str(Path.home() / "Documents" / "ComfyUI")) +if _COMFY_ROOT not in sys.path and os.path.isdir(_COMFY_ROOT): + sys.path.insert(0, _COMFY_ROOT) + +quant_ops = pytest.importorskip("comfy.quant_ops") +if not quant_ops._CK_AVAILABLE: + pytest.skip("comfy_kitchen unavailable", allow_module_level=True) + +import comfy.ops # noqa: E402 + +from nvfp4_patch import ( # noqa: E402 + QUANTIZE_SCOPES, + _selected, + quantize_state_dict_nvfp4, +) + +D = 512 # >= _MIN_DIM and a multiple of 16 + + +def _w(out_f=D, in_f=D): + return torch.randn(out_f, in_f, dtype=torch.bfloat16) * 0.02 + + +def _fake_ltxav_sd(n_blocks=4, prefix=""): + """Minimal LTX-AV-shaped state dict: video attn/ff, audio branch, cross-modal + bridges, plus non-block layers that must never be touched by block scopes.""" + sd = {} + for i in range(n_blocks): + b = f"{prefix}transformer_blocks.{i}." + for mod in ("attn1", "attn2", "audio_attn1", "audio_to_video_attn", "video_to_audio_attn"): + for lin in ("to_q", "to_k", "to_v", "to_out.0"): + sd[f"{b}{mod}.{lin}.weight"] = _w() + sd[f"{b}ff.net.0.proj.weight"] = _w(2 * D, D) + sd[f"{b}ff.net.2.weight"] = _w(D, 2 * D) + sd[f"{b}audio_ff.net.0.proj.weight"] = _w(2 * D, D) + sd[f"{b}scale_shift_table"] = torch.randn(6, D, dtype=torch.bfloat16) # 2D but not .weight + sd[f"{b}attn1.norm_q.weight"] = torch.randn(D, dtype=torch.bfloat16) # 1D + sd[f"{prefix}patchify_proj.weight"] = _w() + sd[f"{prefix}proj_out.weight"] = _w() + sd[f"{prefix}adaln_single.linear.weight"] = _w(6 * D, D) + return sd + + +def _quant_layers(out_sd): + return {k[: -len(".comfy_quant")] for k in out_sd if k.endswith(".comfy_quant")} + + +def test_video_scope_protects_audio_cross_edges_and_nonblock(): + sd = _fake_ltxav_sd() + out, targets, keep = quantize_state_dict_nvfp4(sd, "video blocks", 1, 1) + q = _quant_layers(out) + assert keep == {0, 3} + # middle-block video linears quantized + assert "transformer_blocks.1.attn1.to_q" in q + assert "transformer_blocks.2.ff.net.2" in q + # first/last blocks untouched + assert not any(l.startswith("transformer_blocks.0.") or l.startswith("transformer_blocks.3.") for l in q) + # audio branch, cross-modal bridges, non-block layers untouched + assert not any("audio" in l for l in q) + assert not any(l.split(".")[-2:] == ["patchify_proj"] or "proj_out" in l or "adaln" in l for l in q) + # untouched keys pass through identically + assert torch.equal(out["transformer_blocks.0.attn1.to_q.weight"], + sd["transformer_blocks.0.attn1.to_q.weight"]) + assert len(targets) == len(q) + + +def test_scopes_widen_monotonically(): + sd = _fake_ltxav_sd() + counts = [] + for scope in QUANTIZE_SCOPES: + _, targets, _ = quantize_state_dict_nvfp4(sd, scope, 1, 1) + counts.append(len(targets)) + assert counts == sorted(counts) and counts[0] < counts[-1] + # cross-modal scope adds exactly the two bridges of the two middle blocks (4 linears each) + assert counts[1] - counts[0] == 2 * 2 * 4 + + +def test_all_2d_scope_still_skips_sensitive_generic_layers(): + sd = _fake_ltxav_sd() + out, _, _ = quantize_state_dict_nvfp4(sd, "all 2D layers", 0, 0) + q = _quant_layers(out) + for name in ("proj_out", "adaln_single.linear", "patchify_proj"): + assert not any(name in l for l in q) + assert any("audio_attn1" in l for l in q) # audio IS included in the widest scopes + + +def test_serialized_format_matches_core_loader_contract(): + sd = _fake_ltxav_sd(n_blocks=3) + out, _, _ = quantize_state_dict_nvfp4(sd, "video blocks", 1, 1) + layer = "transformer_blocks.1.attn1.to_q" + assert out[f"{layer}.weight"].dtype == torch.uint8 # packed fp4 pairs + assert out[f"{layer}.weight"].shape[1] == D // 2 + assert out[f"{layer}.weight_scale"].dtype == torch.float8_e4m3fn + assert out[f"{layer}.weight_scale_2"].dtype == torch.float32 + import json + conf = json.loads(bytes(out[f"{layer}.comfy_quant"].tolist())) + assert conf == {"format": "nvfp4"} + + +def test_prefix_and_refusals(): + sd = _fake_ltxav_sd(prefix="model.diffusion_model.") + out, targets, _ = quantize_state_dict_nvfp4(sd, "video blocks", 1, 1) + assert any(k.startswith("model.diffusion_model.") and k.endswith(".comfy_quant") for k in out) + assert targets + # already-quantized checkpoints are refused + with pytest.raises(ValueError): + quantize_state_dict_nvfp4(out, "video blocks", 1, 1) + with pytest.raises(ValueError): + quantize_state_dict_nvfp4({"scaled_fp8": torch.zeros(2)}, "video blocks", 1, 1) + + +def test_selector_never_matches_bias_or_small(): + sd = {"transformer_blocks.1.attn1.to_q.bias": torch.randn(D, dtype=torch.bfloat16), + "transformer_blocks.1.attn1.small.weight": _w(64, 64)} + _, targets, _ = quantize_state_dict_nvfp4(sd, "all blocks", 0, 0) + assert targets == [] + assert _selected("transformer_blocks.1.attn1.to_q", "video blocks", set()) + assert not _selected("transformer_blocks.1.audio_attn1.to_q", "video blocks", set()) + + +@pytest.mark.parametrize("disabled", [[], ["nvfp4"]]) +def test_core_mixed_precision_linear_loads_and_forwards(disabled): + """The emitted keys load through comfy.ops.mixed_precision_ops Linear and produce + output close to the bf16 reference - both on the eager NVFP4 matmul path and on + the dequant-emulation path used by GPUs without FP4 compute.""" + torch.manual_seed(7) + w = _w() + sd = {"transformer_blocks.1.attn1.to_q.weight": w, + "transformer_blocks.1.attn1.to_q.bias": torch.zeros(D, dtype=torch.bfloat16)} + out, targets, _ = quantize_state_dict_nvfp4(sd, "video blocks", 0, 0) + assert targets == ["transformer_blocks.1.attn1.to_q.weight"] + + ops = comfy.ops.mixed_precision_ops({}, torch.bfloat16, disabled=disabled) + lin = ops.Linear(D, D, bias=True, device="cpu", dtype=torch.bfloat16) + layer_sd = {k.replace("transformer_blocks.1.attn1.to_q.", ""): v for k, v in out.items()} + missing, unexpected = lin.load_state_dict(layer_sd, strict=False) + assert not missing and not unexpected + assert isinstance(lin.weight, quant_ops.QuantizedTensor) + assert lin.quant_format == "nvfp4" + + x = torch.randn(2, 8, D, dtype=torch.bfloat16) * 0.5 + with torch.no_grad(): + y = lin.forward(x) + y_ref = torch.nn.functional.linear(x, w) + assert y.shape == y_ref.shape + rel = (y.float() - y_ref.float()).abs().mean() / (y_ref.float().abs().mean() + 1e-8) + assert rel < 0.25, f"relative error too high: {rel:.3f} (disabled={disabled})"