feat(loaders): raven_lora on our own diffusion loader

FunPack owns its loading pipeline, so the causal lane must not need 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. Same checkpoint either way — it is a LoRA, there is no
separate RAVEN model file, and every message now says so.

The weight load itself is DELEGATED to the RAVEN package rather than
reimplemented, and not for convenience: the adapter is an FP32 activation
residual registered on the base Linear leaves and must be attached before the
ModelPatcher is built. ModelPatcher.model_size() memoises, so an adapter
attached afterwards is invisible to Comfy's memory ledger and never moved by
partial CPU offload — the model would report itself smaller than it is and the
residual would end up on the wrong device.

Selecting a LoRA when the causal class is unavailable RAISES with the fix in
the message, rather than quietly loading a plain model: a model without the
adapter would leave the causal lane reading an attention pattern nothing was
trained for, which reads as a quality problem with no visible cause.

Strength is fixed at 1.0 — the adapter is what teaches the chunked attention
pattern, so "partly on" is not a useful state.

Tests: 1730 passed (+12).

loaders.py

@@ -73,6 +73,11 @@ 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.
 
@@ -128,6 +133,44 @@ 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."""
 
@@ -166,6 +209,24 @@ 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."}),
                 "fp16_accumulation": ("BOOLEAN", {
                     "default": False,
                     "tooltip": "Accumulate fp16 matmuls in fp16. Faster on recent NVIDIA cards; "
@@ -228,7 +289,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):
+                   sla=False, raven_lora="None"):
         notes = [f"FunPack Diffusion Model Loader | {model_name}"]
 
         accum = set_fp16_accumulation(fp16_accumulation)
@@ -265,8 +326,13 @@ class FunPackDiffusionModelLoader:
             notes.append(gguf_note)
         else:
             state_dict, metadata = comfy.utils.load_torch_file(path, return_metadata=True)
-        model = comfy.sd.load_diffusion_model_state_dict(
-            state_dict, model_options=model_options, metadata=metadata)
+        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)
         if model is None:
             raise RuntimeError(
                 f"ERROR: could not detect a diffusion model in {model_name}. "

movie_editor/frontend/engine_settings.js

@@ -333,7 +333,7 @@
       detail: "H3 runs video and audio on different flow schedules but only one sigma grid reaches the sampler, so audio drifts further the bigger the step. This corrects it for one multiply per step. Exact with FunPack Distilled Flow, Hybrid Euler 2S and euler; leave it off with res_multistep / dpmpp_2m and the rest of the multistep family, which already absorb most of the error. No effect when both sigma shifts are equal, or on two-evals-per-step samplers." },
     { name: "h3_causal_chunks", label: "Remembering chunks (H3)", kind: "bool", default: false,
       hint: "Generates the clip in time chunks that each remember the ones before them, instead of denoising the whole clip at once — so a long clip stays one continuous shot instead of separate scenes stitched together.",
-      detail: "Needs the ComfyUI-MiniMax-H3-RAVEN-Streaming pack installed and its RAVEN LoRA loaded; the chunked attention pattern is what that LoRA was trained to read. Without either it says so and the run continues normally. The wired sampler and CFG are not used on this path. There is no per-chunk redo: a chunk is committed the moment it finishes and the next is already built on it, so disliking one chunk means a new seed. EXPERIMENTAL." },
+      detail: "Requires the model to be loaded with a RAVEN streaming LoRA — pick one in Models, on the diffusion model's raven_lora setting. It is a LoRA on the ordinary H3 checkpoint, not a different model file, but it is not optional here: the chunked attention pattern is what it was trained to read. Without it this says so and the run continues normally. The wired sampler and CFG are not used on this path. There is no per-chunk redo: a chunk is committed the moment it finishes and the next is already built on it, so disliking one chunk means a new seed. EXPERIMENTAL." },
     { name: "h3_causal_step_rule", label: "Chunk step rule", kind: "combo", choices: ["consistency", "euler", "euler_ancestral"], default: "consistency",
       dependsOn: "h3_causal_chunks", dependsValue: true,
       hint: "How each chunk reaches its next noise level. 'consistency' is what the RAVEN LoRA was distilled for. 'euler' is the ordinary step used everywhere else, so a chunked run can be compared against a normal one without this being a second thing that changed." },

raven_causal.py

@@ -293,8 +293,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 — load it through RAVEN Model Loader, which also "
-                      "attaches the RAVEN LoRA the causal attention pattern was trained for.")
+                      "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.")
     blocks = getattr(diffusion, "blocks", None)
     if not blocks:
         return None, "the loaded model exposes no DiT blocks to cache"

samplers.py

@@ -3179,7 +3179,7 @@ class FunPackLTXAVSceneChainSampler:
                 }),
                 "h3_causal_chunks": ("BOOLEAN", {
                     "default": False,
-                    "tooltip": "EXPERIMENTAL, MiniMax H3 only. Generates the clip in time CHUNKS that each remember the ones before them, instead of denoising the whole clip at once. A chunk is finished, written into a key/value cache as settled context, and the next chunk is generated attending to it — so a long clip stays one continuous shot with real memory of what was already drawn, rather than a chain of separate scenes stitched together. Needs the ComfyUI-MiniMax-H3-RAVEN-Streaming pack installed AND its RAVEN LoRA loaded: the chunked attention pattern is what that LoRA was trained to read, and the base H3 weights have never seen a key/value cache. Without either, this says so and the run continues normally. One consequence you cannot opt out of: there is no per-chunk redo — a chunk enters the cache the moment it finishes and the next is already built on it, so disliking chunk 3 of 12 means a new seed.",
+                    "tooltip": "EXPERIMENTAL, MiniMax H3 only. Generates the clip in time CHUNKS that each remember the ones before them, instead of denoising the whole clip at once. A chunk is finished, written into a key/value cache as settled context, and the next chunk is generated attending to it — so a long clip stays one continuous shot with real memory of what was already drawn, rather than a chain of separate scenes stitched together. Requires the model to be loaded with a RAVEN streaming LoRA: set `raven_lora` on FunPack Diffusion Model Loader. That is a LoRA on the ordinary H3 checkpoint, not a different model file — but it is not optional here, because the chunked attention pattern is what it was trained to read and the base H3 weights have never seen a key/value cache. Without it this says so and the run continues normally. One consequence you cannot opt out of: there is no per-chunk redo — a chunk enters the cache the moment it finishes and the next is already built on it, so disliking chunk 3 of 12 means a new seed.",
                 }),
                 "h3_causal_step_rule": (list(_CAUSAL_STEP_RULES), {
                     "default": "consistency",

tests/test_raven_causal.py

@@ -524,7 +524,10 @@ def test_a_stock_bidirectional_model_is_refused_by_name(monkeypatch):
     monkeypatch.setitem(sys.modules, "raven_streaming.cache", fake_cache)
     session, reason = rc.build_session(object(), [], {})
     assert session is None
-    assert "chunk-causal DiT" in reason and "RAVEN Model Loader" in reason
+    # 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
 
 
 # ── the Chain Sampler dispatch ──────────────────────────────────────────────

tests/test_raven_loader_option.py

@@ -0,0 +1,130 @@
+"""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