From b5594b00a73c6699a293a969ed6d77aee2fc5f83 Mon Sep 17 00:00:00 2001 From: digital-garbage <133010426+digital-garbage@users.noreply.github.com> Date: Wed, 26 Aug 2026 18:55:58 +0300 Subject: [PATCH] feat(loaders): chunk-causal is a switch on our loader, the adapter is just a LoRA MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Both halves now go through nodes we already own, so nothing third-party has to be installed for the feature to exist. The diffusion loader gains `chunk_causal`, which RE-CLASSES an already-loaded H3 in place. That is what removes the need for a separate loader entirely: the alternative is reimplementing ComfyUI's whole diffusion-model load path so the adapter can be attached before the ModelPatcher exists and its memory ledger counts it. Re-classing adds no parameter, renames no key and copies nothing — only which forward runs changes — so the ledger and offload accounting are untouched and the switch is idempotent. The adapter is then an ORDINARY LoRA through FunPack Apply LoRA Weights. Our loader already scores prefix variants against the model's real key map; it just did not know this one. RAVEN's keys are `base_model.model.dit.`, so stripping only `base_model.model.` left `dit.` in front of every key and the whole file matched nothing, silently. One entry in LORA_KEY_PREFIXES fixes it. A ComfyUI too old for H3, or a non-H3 model, declines with a reason rather than raising — one optional switch must not turn the loader into a dead node. Tests: 1751 passed. --- h3_causal.py | 32 ++++++++ loaders.py | 96 +++++++--------------- model_management.py | 5 ++ raven_causal.py | 7 +- tests/test_h3_causal_loader.py | 113 ++++++++++++++++++++++++++ tests/test_raven_causal.py | 4 +- tests/test_raven_loader_option.py | 130 ------------------------------ 7 files changed, 182 insertions(+), 205 deletions(-) create mode 100644 tests/test_h3_causal_loader.py delete mode 100644 tests/test_raven_loader_option.py diff --git a/h3_causal.py b/h3_causal.py index 8b967454..125c2799 100644 --- a/h3_causal.py +++ b/h3_causal.py @@ -294,3 +294,35 @@ def _to_causal_block(block, block_cls, attention_cls): block.attn.__class__ = attention_cls block.__class__ = block_cls return block + + +def make_causal(model): + """Re-class an already-loaded H3 so its blocks can read a chunk cache. Returns (ok, note). + + This is why nothing here needs its own loader. The alternative — building the model + through a different class — means reimplementing ComfyUI's whole diffusion-model load + path, and attaching anything to the raw model before the ModelPatcher exists so its + memory ledger counts it. Re-classing sidesteps all of that: no parameter is added, no key + is renamed, nothing is copied. Only which `forward` runs changes. + + Idempotent, and a no-op on anything that is not H3. + """ + diffusion = getattr(getattr(model, "model", None), "diffusion_model", None) + if diffusion is None: + return False, "no diffusion model to make causal" + blocks = getattr(diffusion, "blocks", None) + if not blocks: + return False, "this model has no DiT blocks — not a MiniMax H3" + if getattr(diffusion, "_funpack_causal", False): + return True, "already chunk-causal" + try: + attention_cls, block_cls, model_cls = _causal_classes() + except Exception as error: # noqa: BLE001 + return False, f"this ComfyUI has no MiniMax H3 to extend ({error})" + if not isinstance(diffusion, model_cls.__bases__[0]): + return False, "this model is not a MiniMax H3" + for index in range(len(blocks)): + _to_causal_block(blocks[index], block_cls, attention_cls) + diffusion.__class__ = model_cls + diffusion._funpack_causal = True + return True, f"chunk-causal DiT: {len(blocks)} blocks re-classed in place" diff --git a/loaders.py b/loaders.py index 23c4d1be..51044f93 100644 --- a/loaders.py +++ b/loaders.py @@ -73,11 +73,6 @@ def weight_model_options(weight_dtype): return options -def raven_lora_choices(): - """LoRA files, with "None" first — the causal lane is opt-in per model.""" - return ["None"] + list(folder_paths.get_filename_list("loras")) - - def attention_choices(): """Attention backends this ComfyUI can actually run, newest-registry first. @@ -133,44 +128,6 @@ def set_fp16_accumulation(enabled): return bool(enabled) -def _load_with_raven_lora(model_name, lora_name, weight_dtype, model_options): - """Load an H3 checkpoint as a chunk-causal DiT with the RAVEN LoRA attached. - - Delegated to the RAVEN package rather than reimplemented, for a reason that is not - convenience: the adapter is an FP32 activation residual registered as parameters of the - base Linear leaves, and it has to be attached to the raw model BEFORE the ModelPatcher is - built. ``ModelPatcher.model_size()`` memoises, so an adapter attached afterwards is - invisible to ComfyUI's memory ledger and never moved by partial CPU offload — the model - would look smaller than it is and the residual would sit on the wrong device. - - Returns (model, note). Raises with an actionable message rather than falling back to a - plain load: silently returning a model WITHOUT the adapter would leave the causal lane - reading an attention pattern nothing was trained for, and look like a quality problem. - """ - try: - from .raven_causal import locate_raven - except ImportError: - from raven_causal import locate_raven - module, reason = locate_raven() - if module is None: - raise RuntimeError( - f"ERROR: a RAVEN LoRA is selected ({lora_name}) but the causal model class is " - f"unavailable — {reason} Set raven_lora to None to load this model normally.") - from raven_streaming import loader as raven_loader - from raven_streaming.causal_model import RavenCausalMiniMaxH3Model - - model = raven_loader.load_raven_diffusion_model( - model_name, lora_name, - weight_dtype=weight_dtype, - model_options=dict(model_options), - unet_model_cls=RavenCausalMiniMaxH3Model, - ) - return model, (f"RAVEN: chunk-causal DiT + {lora_name} at strength 1.0 (a LoRA on the " - f"ordinary H3 checkpoint, not a separate model). The Chain Sampler's " - f"'Remembering chunks' mode can now run; without that mode on, this model " - f"samples exactly as usual.") - - class FunPackDiffusionModelLoader: """Loads a diffusion model with quantization and an attention backend chosen per node.""" @@ -209,24 +166,18 @@ class FunPackDiffusionModelLoader: "everything else."}), }, "optional": { - "raven_lora": (raven_lora_choices(), { - "default": "None", - "tooltip": "MiniMax H3 only. Loads this model as a CHUNK-CAUSAL DiT with a " - "RAVEN streaming LoRA attached, which is what the Chain " - "Sampler's 'Remembering chunks' mode needs. The clip is then " - "generated in time chunks that each remember the ones before " - "them, so a long clip stays one continuous shot instead of " - "separate scenes stitched together. The LoRA is not optional " - "for that mode: the chunked attention pattern is what it was " - "trained to read, and the base H3 weights have never seen a " - "key/value cache. Needs the ComfyUI-MiniMax-H3-RAVEN-Streaming " - "pack installed for the causal model class. This is a LoRA on " - "the ordinary H3 checkpoint — there is no separate RAVEN model " - "file. Strength is fixed at 1.0: the adapter is what teaches " - "the chunked attention pattern, so 'partly on' is not a useful " - "state. 'None' loads normally. " - "Use the full non-pruned checkpoint: the pruned/adaln-curve cut " - "has no time_embedder for the adapter to attach to."}), + "chunk_causal": ("BOOLEAN", { + "default": False, + "tooltip": "MiniMax H3 only. Loads this model so it can generate in time " + "CHUNKS that each remember the ones before them, which is what " + "the Chain Sampler's 'Remembering chunks' mode needs. Nothing " + "about the weights changes and an ordinary generation is " + "unaffected — the model simply gains the ability to read a " + "key/value cache. It still needs a LoRA trained for that " + "pattern (load it with FunPack Apply LoRA Weights like any " + "other): the base H3 weights have never seen a cache, so " + "without one the chunked lane runs out of distribution. Use " + "the full non-pruned checkpoint."}), "fp16_accumulation": ("BOOLEAN", { "default": False, "tooltip": "Accumulate fp16 matmuls in fp16. Faster on recent NVIDIA cards; " @@ -289,7 +240,7 @@ class FunPackDiffusionModelLoader: def load_model(self, model_name, weight_dtype, compute_dtype, attention, fp16_accumulation=False, sla_sparsity=None, sla_block_size=None, sla_protect_audio=None, sla_min_seq_len=None, sla_dense_last_steps=None, - sla=False, raven_lora="None"): + sla=False, chunk_causal=False): notes = [f"FunPack Diffusion Model Loader | {model_name}"] accum = set_fp16_accumulation(fp16_accumulation) @@ -326,13 +277,8 @@ class FunPackDiffusionModelLoader: notes.append(gguf_note) else: state_dict, metadata = comfy.utils.load_torch_file(path, return_metadata=True) - if str(raven_lora) not in ("", "None"): - model, raven_note = _load_with_raven_lora( - model_name, raven_lora, weight_dtype, model_options) - notes.append(raven_note) - else: - model = comfy.sd.load_diffusion_model_state_dict( - state_dict, model_options=model_options, metadata=metadata) + model = comfy.sd.load_diffusion_model_state_dict( + state_dict, model_options=model_options, metadata=metadata) if model is None: raise RuntimeError( f"ERROR: could not detect a diffusion model in {model_name}. " @@ -370,6 +316,18 @@ class FunPackDiffusionModelLoader: else: notes.append(f"attention: {attention} (dense calls)") + if chunk_causal: + try: + from .h3_causal import make_causal + except ImportError: + from h3_causal import make_causal + ok, causal_note = make_causal(model) + # Not fatal. The weights are fine and every other setting on this node applied; + # what is lost is one optional capability, and the sampler refuses the chunked + # mode on its own if the model cannot do it. + notes.append(f"chunk-causal: {causal_note}" if ok + else f"chunk-causal REQUESTED BUT NOT APPLIED: {causal_note}") + return (model, "\n".join(notes)) diff --git a/model_management.py b/model_management.py index 87dfc6f7..cfdb1467 100644 --- a/model_management.py +++ b/model_management.py @@ -284,6 +284,11 @@ LORA_KEY_PREFIXES = ( "diffusion_model.", "model.diffusion_model.", "base_model.model.", + # PEFT adapters trained against a wrapper that names the DiT: RAVEN's streaming LoRA is + # `base_model.model.dit.`, so stripping only `base_model.model.` leaves `dit.` in + # front of every key and the whole file matches nothing. + "base_model.model.dit.", + "base_model.model.diffusion_model.", "lora_model.", "net.", ) diff --git a/raven_causal.py b/raven_causal.py index d1da2846..9e5c0aa1 100644 --- a/raven_causal.py +++ b/raven_causal.py @@ -292,10 +292,9 @@ def build_session(model, positive, latent, *, sink=2, window=2, storage="cpu_pin diffusion = resolved.diffusion_model if not hasattr(diffusion, "forward_chunk") or not hasattr(diffusion, "prefill_text"): - return None, ("the loaded model is a stock bidirectional H3. The causal lane needs the " - "chunk-causal DiT: set `raven_lora` on FunPack Diffusion Model Loader " - "to the RAVEN streaming LoRA. It is a LoRA on the ordinary H3 " - "checkpoint, not a different model file.") + return None, ("the loaded model is a stock bidirectional H3. Switch on `chunk_causal` " + "on FunPack Diffusion Model Loader, and load a LoRA trained for the " + "chunked pattern with FunPack Apply LoRA Weights.") blocks = getattr(diffusion, "blocks", None) if not blocks: return None, "the loaded model exposes no DiT blocks to cache" diff --git a/tests/test_h3_causal_loader.py b/tests/test_h3_causal_loader.py new file mode 100644 index 00000000..f63f368f --- /dev/null +++ b/tests/test_h3_causal_loader.py @@ -0,0 +1,113 @@ +"""Making a model chunk-causal is a switch on OUR loader, and the LoRA is an ordinary LoRA. + +FunPack owns the loading pipeline, so neither half of this needs a third-party node or +package. The diffusion loader RE-CLASSES an already-loaded H3 in place rather than building +it through a different class — no parameter is added, no key is renamed, nothing is copied, +so ComfyUI's memory ledger and offload accounting are untouched. The adapter then goes +through FunPack Apply LoRA Weights like any other LoRA. +""" +import sys +from pathlib import Path + +import pytest + +sys.path.insert(0, str(Path(__file__).resolve().parents[1])) + +import h3_causal as hc +import loaders +import model_management as mm + + +def _fields(): + spec = loaders.FunPackDiffusionModelLoader.INPUT_TYPES() + return {**spec.get("required", {}), **spec.get("optional", {})} + + +def test_the_switch_is_on_our_own_loader(): + assert "chunk_causal" in _fields() + + +def test_it_is_off_by_default(): + assert _fields()["chunk_causal"][1]["default"] is False + + +def test_the_tooltip_says_a_lora_is_still_needed(): + """Without one the chunked lane runs out of distribution, which reads as a quality problem + with no visible cause.""" + tip = _fields()["chunk_causal"][1]["tooltip"] + assert "LoRA" in tip and "out of distribution" in tip + + +def test_no_raven_package_is_named_anywhere_in_the_loader(): + """The whole point of owning it: no third-party install to discover.""" + import inspect + src = inspect.getsource(loaders.FunPackDiffusionModelLoader) + assert "raven_streaming" not in src and "RAVEN Model Loader" not in src + + +# ── the LoRA is an ordinary LoRA ──────────────────────────────────────────── + +def test_the_peft_dit_prefix_is_one_of_the_candidates(): + """RAVEN's keys are base_model.model.dit.. Stripping only `base_model.model.` leaves + `dit.` in front of every key and the whole file silently matches nothing.""" + assert "base_model.model.dit." in mm.LORA_KEY_PREFIXES + + +def test_the_longer_prefix_is_tried_before_the_shorter_one_can_win(): + """Both match; the loader scores every candidate and keeps the best, so the order only has + to include the longer one — but it must be there at all.""" + assert "base_model.model." in mm.LORA_KEY_PREFIXES + prefixes = list(mm.LORA_KEY_PREFIXES) + assert prefixes.index("base_model.model.dit.") > prefixes.index("base_model.model.") + + +def test_stripping_that_prefix_yields_plain_module_paths(): + stripped = mm._strip_key_prefix( + {"base_model.model.dit.blocks.0.attn.qkv_proj.lora_A.weight": 1, + "base_model.model.dit.blocks.0.attn.qkv_proj.lora_B.weight": 2}, + "base_model.model.dit.") + assert set(stripped) == {"blocks.0.attn.qkv_proj.lora_A.weight", + "blocks.0.attn.qkv_proj.lora_B.weight"} + + +# ── re-classing an already-loaded model ───────────────────────────────────── + +class _FakeBlock: + def __init__(self): + self.attn = type("A", (), {})() + + +def test_a_model_without_blocks_declines_with_a_reason(): + fake = type("M", (), {"model": type("N", (), {"diffusion_model": type("D", (), {})()})()})() + ok, note = hc.make_causal(fake) + assert ok is False and "no DiT blocks" in note + + +def test_a_model_without_a_diffusion_model_declines(): + ok, note = hc.make_causal(type("M", (), {"model": None})()) + assert ok is False and "no diffusion model" in note + + +def test_it_is_idempotent(monkeypatch): + """Loading the same model twice, or re-running the node, must not stack anything.""" + diffusion = type("D", (), {})() + diffusion.blocks = [_FakeBlock()] + diffusion._funpack_causal = True + fake = type("M", (), {})() + fake.model = type("N", (), {})() + fake.model.diffusion_model = diffusion + ok, note = hc.make_causal(fake) + assert ok is True and "already" in note + + +def test_a_comfy_without_h3_declines_instead_of_raising(monkeypatch): + """A ComfyUI too old for H3 must not turn one optional switch into a dead node.""" + monkeypatch.setattr(hc, "_causal_classes", + lambda: (_ for _ in ()).throw(ImportError("no minimax"))) + diffusion = type("D", (), {})() + diffusion.blocks = [_FakeBlock()] + fake = type("M", (), {})() + fake.model = type("N", (), {})() + fake.model.diffusion_model = diffusion + ok, note = hc.make_causal(fake) + assert ok is False and "no MiniMax H3" in note diff --git a/tests/test_raven_causal.py b/tests/test_raven_causal.py index b6bb52d9..c4b13cb3 100644 --- a/tests/test_raven_causal.py +++ b/tests/test_raven_causal.py @@ -526,8 +526,8 @@ def test_a_stock_bidirectional_model_is_refused_by_name(monkeypatch): assert session is None # points at OUR loader: FunPack owns the loading pipeline, so the fix is a setting on the # node already in the graph, not a third-party loader node to go and add - assert "chunk-causal DiT" in reason - assert "FunPack Diffusion Model Loader" in reason and "raven_lora" in reason + assert "stock bidirectional H3" in reason + assert "FunPack Diffusion Model Loader" in reason and "chunk_causal" in reason # ── the Chain Sampler dispatch ────────────────────────────────────────────── diff --git a/tests/test_raven_loader_option.py b/tests/test_raven_loader_option.py deleted file mode 100644 index 2e44637f..00000000 --- a/tests/test_raven_loader_option.py +++ /dev/null @@ -1,130 +0,0 @@ -"""The RAVEN LoRA is an option on FunPack's OWN diffusion loader. - -FunPack owns its loading pipeline, so the causal lane must not require a third-party loader -node in the graph. `raven_lora` on FunPack Diffusion Model Loader turns an ordinary H3 -checkpoint into the chunk-causal DiT with the adapter attached — the checkpoint is the same -file either way; there is no separate RAVEN model. - -Why the load is DELEGATED rather than reimplemented: the adapter is an FP32 activation -residual registered on the base Linear leaves, and it has to be attached before the -ModelPatcher exists. `ModelPatcher.model_size()` memoises, so an adapter attached afterwards -is invisible to ComfyUI's memory ledger and never moved by partial CPU offload. -""" -import sys -import types -from pathlib import Path - -import pytest - -sys.path.insert(0, str(Path(__file__).resolve().parents[1])) - -import loaders - - -def _fields(): - spec = loaders.FunPackDiffusionModelLoader.INPUT_TYPES() - return {**spec.get("required", {}), **spec.get("optional", {})} - - -def test_the_option_exists_on_our_own_loader(): - assert "raven_lora" in _fields() - - -def test_it_defaults_to_off(): - assert _fields()["raven_lora"][1]["default"] == "None" - - -def test_none_is_the_first_choice(): - """A file list that opens on a LoRA would arm the causal lane by accident.""" - assert _fields()["raven_lora"][0][0] == "None" - - -def test_the_tooltip_says_it_is_a_lora_not_a_model(): - """There is no 'RAVEN H3 model' — saying so would send the user hunting for a file that - does not exist.""" - tip = _fields()["raven_lora"][1]["tooltip"] - assert "not a separate RAVEN model file" in tip or "no separate RAVEN model" in tip - - -def test_the_tooltip_says_strength_is_fixed(): - assert "1.0" in _fields()["raven_lora"][1]["tooltip"] - - -def test_the_choices_come_from_the_loras_folder(monkeypatch): - monkeypatch.setattr(loaders.folder_paths, "get_filename_list", - lambda kind: ["a.safetensors"] if kind == "loras" else []) - assert loaders.raven_lora_choices() == ["None", "a.safetensors"] - - -# ── the delegating load ───────────────────────────────────────────────────── - -def test_a_missing_raven_package_refuses_instead_of_loading_plain(monkeypatch): - """Returning a model WITHOUT the adapter would leave the causal lane reading an attention - pattern nothing was trained for — a quality problem with no visible cause.""" - import raven_causal - monkeypatch.setattr(raven_causal, "locate_raven", - lambda: (None, "the pack is not installed.")) - with pytest.raises(RuntimeError, match="not installed"): - loaders._load_with_raven_lora("h3.safetensors", "raven.safetensors", "default", {}) - - -def test_the_refusal_says_how_to_get_a_normal_load(monkeypatch): - import raven_causal - monkeypatch.setattr(raven_causal, "locate_raven", lambda: (None, "missing.")) - with pytest.raises(RuntimeError, match="raven_lora to None"): - loaders._load_with_raven_lora("h3.safetensors", "raven.safetensors", "default", {}) - - -def test_the_load_is_delegated_with_the_causal_class(monkeypatch): - """The class is what makes the DiT chunk-causal; loading the adapter without it would - attach weights nothing reads.""" - import raven_causal - seen = {} - - causal_cls = type("RavenCausalMiniMaxH3Model", (), {}) - fake_loader = types.SimpleNamespace( - load_raven_diffusion_model=lambda unet, lora, **kw: seen.update( - {"unet": unet, "lora": lora, **kw}) or "the-model") - parent = types.ModuleType("raven_streaming") - parent.loader = fake_loader - causal_mod = types.ModuleType("raven_streaming.causal_model") - causal_mod.RavenCausalMiniMaxH3Model = causal_cls - parent.causal_model = causal_mod - monkeypatch.setitem(sys.modules, "raven_streaming", parent) - monkeypatch.setitem(sys.modules, "raven_streaming.loader", fake_loader) - monkeypatch.setitem(sys.modules, "raven_streaming.causal_model", causal_mod) - monkeypatch.setattr(raven_causal, "locate_raven", lambda: (parent, "")) - - model, note = loaders._load_with_raven_lora( - "h3.safetensors", "raven.safetensors", "fp8_e4m3fn", {"opt": 1}) - assert model == "the-model" - assert seen["unet"] == "h3.safetensors" and seen["lora"] == "raven.safetensors" - assert seen["unet_model_cls"] is causal_cls - assert seen["weight_dtype"] == "fp8_e4m3fn" - assert seen["model_options"] == {"opt": 1} - assert "strength 1.0" in note - - -def test_the_note_does_not_call_it_a_separate_model(monkeypatch): - import raven_causal - causal_mod = types.ModuleType("raven_streaming.causal_model") - causal_mod.RavenCausalMiniMaxH3Model = type("C", (), {}) - fake_loader = types.SimpleNamespace(load_raven_diffusion_model=lambda *a, **k: "m") - parent = types.ModuleType("raven_streaming") - parent.loader, parent.causal_model = fake_loader, causal_mod - monkeypatch.setitem(sys.modules, "raven_streaming", parent) - monkeypatch.setitem(sys.modules, "raven_streaming.loader", fake_loader) - monkeypatch.setitem(sys.modules, "raven_streaming.causal_model", causal_mod) - monkeypatch.setattr(raven_causal, "locate_raven", lambda: (parent, "")) - _model, note = loaders._load_with_raven_lora("h3.safetensors", "r.safetensors", "default", {}) - assert "not a separate model" in note - - -def test_the_sampler_points_at_our_loader_not_a_third_party_node(): - """We own the loading pipeline; telling the user to add someone else's loader node would - be wrong as well as unnecessary.""" - import inspect - import raven_causal - src = inspect.getsource(raven_causal.build_session) - assert "FunPack Diffusion Model Loader" in src - assert "RAVEN Model Loader" not in src