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