digital-garbage
b5594b00a7
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.<path>`, 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.
114 lines
4.5 KiB
Python
114 lines
4.5 KiB
Python
"""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.<path>. 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
|