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 <noreply@anthropic.com>
2026-07-02 06:16:08 +03:00
parent c8710597d0
commit 3dd2d33f1d
3 changed files with 423 additions and 0 deletions
@@ -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",
 }

@@ -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.<idx>.<rest>" 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,)
@@ -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})"