"""A LoRA that matches by NAME but not by SHAPE must say so. The case that motivated this: MiniMax H3 ships in two forms. The full one derives adaLN's input from a `time_embedder` (time_embed_dim 2688); the pruned "curve" one replaces that with a small precomputed basis, so every `adaln_proj.linear` is narrower on its input side. A turbo LoRA trained on the full model matches every key of a curve-form checkpoint and then fails to merge into 51 of them — previously visible only as generic per-key warnings from comfy, while FunPack's own status line still read like a clean load. """ import sys import types import pytest import torch sys.path.insert(0, ".") @pytest.fixture def mm(): import model_management return model_management class _Adapter: """Stands in for comfy's LoRAAdapter: weights = (up[out, rank], down[rank, in], ...).""" def __init__(self, out_dim, in_dim, rank=8): self.weights = (torch.zeros(out_dim, rank), torch.zeros(rank, in_dim), 1.0, None, None) class _Model: def __init__(self, shapes, use_curves=False): self._sd = {k: torch.zeros(*v) for k, v in shapes.items()} diffusion = types.SimpleNamespace(use_adaln_curves=use_curves) inner = types.SimpleNamespace(state_dict=lambda: self._sd, diffusion_model=diffusion) self.model = inner def test_a_fitting_lora_reports_nothing(mm): model = _Model({"diffusion_model.blocks.0.attn.qkv.weight": (768, 768)}) patches = {"diffusion_model.blocks.0.attn.qkv.weight": _Adapter(768, 768)} assert mm._mismatched_lora_keys(model, patches) == [] def test_a_narrower_input_is_caught(mm): """The exact curve-form shape: the weight's input side is the small basis, the LoRA's is the full time-embedding width.""" model = _Model({"diffusion_model.blocks.0.adaln_proj.linear.weight": (16128, 64)}) patches = {"diffusion_model.blocks.0.adaln_proj.linear.weight": _Adapter(16128, 2688)} bad = mm._mismatched_lora_keys(model, patches) assert len(bad) == 1 key, got, want = bad[0] assert got == (16128, 2688) and want == (16128, 64) def test_keys_the_model_does_not_have_are_not_flagged(mm): model = _Model({"diffusion_model.blocks.0.attn.qkv.weight": (768, 768)}) patches = {"diffusion_model.blocks.99.nonexistent.weight": _Adapter(10, 10)} assert mm._mismatched_lora_keys(model, patches) == [] def test_non_lora_adapters_are_left_to_comfy(mm): """LoKr/LoHa/diff entries have no up/down pair to measure — guessing at them would produce false alarms on formats this check does not model.""" model = _Model({"diffusion_model.blocks.0.attn.qkv.weight": (768, 768)}) patches = {"diffusion_model.blocks.0.attn.qkv.weight": types.SimpleNamespace(weights=None)} assert mm._mismatched_lora_keys(model, patches) == [] def test_curve_note_names_the_cause_and_the_remedy(mm): model = _Model({"diffusion_model.blocks.0.adaln_proj.linear.weight": (16128, 64)}, use_curves=True) bad = [("diffusion_model.blocks.0.adaln_proj.linear.weight", (16128, 2688), (16128, 64))] note = mm._adaln_curve_note(model, bad) assert note and "curve-form" in note assert "cannot be projected" in note # The rest of the LoRA is fine, and saying so stops it reading as a failed load. assert "rest of the LoRA still applies" in note def test_no_curve_note_on_a_full_width_checkpoint(mm): """Same mismatched key, but this checkpoint is not curve-form — so that diagnosis would be a wrong explanation, which is worse than a generic one.""" model = _Model({"diffusion_model.blocks.0.adaln_proj.linear.weight": (16128, 2688)}, use_curves=False) bad = [("diffusion_model.blocks.0.adaln_proj.linear.weight", (16128, 99), (16128, 2688))] assert mm._adaln_curve_note(model, bad) is None def test_no_curve_note_when_the_mismatch_is_elsewhere(mm): model = _Model({"diffusion_model.blocks.0.attn.qkv.weight": (768, 768)}, use_curves=True) bad = [("diffusion_model.blocks.0.attn.qkv.weight", (768, 512), (768, 768))] assert mm._adaln_curve_note(model, bad) is None def test_a_broken_model_never_breaks_the_load(mm): """Diagnostics are not worth failing a generation over.""" class Exploding: model = types.SimpleNamespace( state_dict=lambda: (_ for _ in ()).throw(RuntimeError("boom"))) assert mm._mismatched_lora_keys(Exploding(), {"k": _Adapter(1, 1)}) == [] # --- dropping, not just reporting ------------------------------------------------------ # Reporting a mismatch and then handing it to comfy anyway means comfy attempts it: it # materialises the full lora_A @ lora_B delta and only then discovers it cannot be reshaped # into the weight. On a curve-form H3 checkpoint that is a 96768x2688 tensor per block, 51 # times, while dynamic VRAM staging is streaming the model in. def _patched_resolve(mm, monkeypatch, model, patches): monkeypatch.setattr(mm.comfy.lora, "load_lora", lambda *a, **kw: dict(patches), raising=False) monkeypatch.setattr(mm.comfy.lora, "model_lora_keys_unet", lambda *a, **kw: {}, raising=False) monkeypatch.setattr(mm.comfy.lora, "model_lora_keys_clip", lambda *a, **kw: {}, raising=False) monkeypatch.setattr(mm.comfy.lora_convert, "convert_lora", lambda sd: sd, raising=False) return mm.resolve_lora_patches(model, {"x": torch.zeros(1)}) def test_mismatched_patches_are_removed_from_what_comfy_gets(mm, monkeypatch): model = _Model({"diffusion_model.blocks.0.adaln_proj.linear.weight": (64, 8), "diffusion_model.blocks.0.attn.qkv_proj.weight": (64, 64)}, use_curves=True) patches = { "diffusion_model.blocks.0.adaln_proj.linear.weight": _Adapter(64, 2688), # cannot fit "diffusion_model.blocks.0.attn.qkv_proj.weight": _Adapter(64, 64), # fits } out, note = _patched_resolve(mm, monkeypatch, model, patches) assert "diffusion_model.blocks.0.adaln_proj.linear.weight" not in out assert "diffusion_model.blocks.0.attn.qkv_proj.weight" in out # the rest still applies assert "DROPPED" in note def test_a_fitting_lora_is_handed_over_whole(mm, monkeypatch): """Dropping is driven by the shape check — not applied to everything on the way past.""" model = _Model({"diffusion_model.blocks.0.attn.qkv_proj.weight": (64, 64)}) patches = {"diffusion_model.blocks.0.attn.qkv_proj.weight": _Adapter(64, 64)} out, note = _patched_resolve(mm, monkeypatch, model, patches) assert len(out) == 1 and "DROPPED" not in note class _MidAdapter(_Adapter): """A locon entry: comfy rebuilds mat2 from the mid weight, so the plain pair is not what the merge will require.""" def __init__(self, out_dim, in_dim, rank=8): super().__init__(out_dim, in_dim, rank) self.weights = (self.weights[0], self.weights[1], 1.0, torch.zeros(rank, rank, 1, 1), None, None) class _ReshapeAdapter(_Adapter): """A `reshape` entry: comfy PADS the target weight before merging, so the weight in the model is not the shape the delta has to fit.""" def __init__(self, out_dim, in_dim, target, rank=8): super().__init__(out_dim, in_dim, rank) self.weights = (self.weights[0], self.weights[1], 1.0, None, None, target) def test_a_locon_mid_entry_is_left_to_comfy(mm): """Dropping it here would lose an adapter that merges correctly — the check has no model of what mat2 becomes, so it must abstain rather than guess.""" model = _Model({"w": (768, 768)}) assert mm._mismatched_lora_keys(model, {"w": _MidAdapter(768, 512)}) == [] def test_a_padded_target_is_left_to_comfy(mm): model = _Model({"w": (768, 768)}) assert mm._mismatched_lora_keys(model, {"w": _ReshapeAdapter(1024, 768, (1024, 768))}) == [] def test_a_delta_that_fits_by_count_but_not_by_shape_is_dropped(mm): """The dangerous case, and the reason the count alone is not the test. comfy's reshape(weight.shape) SUCCEEDS here and adds a transposed delta without erroring — the weight is corrupted silently and the run dies later with an all-NaN latent.""" model = _Model({"w": (768, 768)}) # 589824 elements, same as 384x1536 assert mm._mismatched_lora_keys(model, {"w": _Adapter(384, 1536)}) != [] def test_a_conv_weight_is_measured_the_way_comfy_flattens_it(mm): model = _Model({"w": (64, 32, 3, 3)}) # comfy flattens to (64, 288) assert mm._mismatched_lora_keys(model, {"w": _Adapter(64, 288)}) == [] assert mm._mismatched_lora_keys(model, {"w": _Adapter(64, 32)}) != [] def test_the_dropped_message_carries_both_shapes(mm, monkeypatch, capsys): """"trained against a different variant" is not actionable on its own. The numbers say whether every layer is off by the same factor or only some are.""" model = _Model({"diffusion_model.blocks.0.attn.qkv_proj.weight": (16128, 5376)}) patches = {"diffusion_model.blocks.0.attn.qkv_proj.weight": _Adapter(5376, 5376)} bad = mm._mismatched_lora_keys(model, patches) assert len(bad) == 1 key, dims, want = bad[0] assert dims == (5376, 5376) assert want == (16128, 5376) def _status_line(slot, name, fragment): return f"lora_{slot}: {name} applied=+1.000 source=base mode=keys=400 fmt=as-is{fragment}" def test_the_stack_verdict_blames_the_checkpoint_when_nothing_fits(mm): """The question the per-file lines cannot answer. One LoRA dropping everything is that file; every LoRA dropping everything is the one thing they share — the model.""" lines = [_status_line(0, "a.safetensors", " | 400/400 DROPPED (shape mismatch)"), _status_line(1, "b.safetensors", " | 128/128 DROPPED (shape mismatch)")] verdict = mm.FunPackLoraLoader._stack_shape_verdict(lines) assert "all 2 LoRAs" in verdict assert "checkpoint" in verdict def test_a_partial_drop_does_not_blame_the_checkpoint(mm): lines = [_status_line(0, "a.safetensors", " | 51/400 DROPPED (shape mismatch)"), _status_line(1, "b.safetensors", "")] verdict = mm.FunPackLoraLoader._stack_shape_verdict(lines) assert "1 of 2 LoRAs" in verdict assert "checkpoint" not in verdict def test_one_lora_dropping_everything_is_not_a_stack_verdict(mm): """A single file says nothing about what every file has in common.""" lines = [_status_line(0, "a.safetensors", " | 400/400 DROPPED (shape mismatch)")] assert "checkpoint" not in (mm.FunPackLoraLoader._stack_shape_verdict(lines) or "") def test_a_clean_stack_says_nothing(mm): lines = [_status_line(0, "a.safetensors", ""), _status_line(1, "b.safetensors", "")] assert mm.FunPackLoraLoader._stack_shape_verdict(lines) is None def test_the_log_line_names_the_file_and_the_share(mm, monkeypatch, capsys): """Which file, and how much of it. "51 dropped" out of 400 is a variant mismatch; out of 51 it is the whole LoRA, and the difference decides where to look next.""" import comfy.lora model = _Model({"w": (768, 768)}) monkeypatch.setattr(comfy.lora, "model_lora_keys_unet", lambda *a, **k: {}) monkeypatch.setattr(comfy.lora, "load_lora", lambda *a, **k: {"w": _Adapter(384, 1536)}) monkeypatch.setattr(mm._log, "_last_by_key", {}, raising=False) _, note = mm.resolve_lora_patches(model, {}, name="turbo.safetensors") out = capsys.readouterr().out assert "turbo.safetensors" in out assert "1 of 1" in out assert "DROPPED before the merge" in out assert "1/1 DROPPED" in note # ── a file that matches nothing has to say what it looked like ────────────── # # "matched no weights" is not actionable. The two sets of NAMES are: they say whether the # file is for another architecture (nothing alike) or for the same one under different # module names (a rename rule away). Every entry in LORA_KEY_PREFIXES was born from one of # these, and each was found by hand because the log did not print this. def test_a_lora_that_matches_nothing_prints_both_namings(mm, monkeypatch, capsys): monkeypatch.setattr(mm.comfy.lora, "model_lora_keys_unet", lambda model, out: {"diffusion_model.blocks.0.attn.qkv_proj.weight": "blocks.0.attn.qkv_proj.weight"}, raising=False) monkeypatch.setattr(mm.comfy.lora, "model_lora_keys_clip", lambda *a, **kw: {}, raising=False) monkeypatch.setattr(mm.comfy.lora_convert, "convert_lora", lambda sd: sd, raising=False) monkeypatch.setattr(mm.comfy.lora, "load_lora", lambda *a, **kw: {}, raising=False) mm._log.reset() lora = {"base_model.model.dit.blocks.0.self_attn.qkv.lora_A.weight": torch.zeros(1)} patches, status = mm.resolve_lora_patches(_Model({}), lora, name="adapter.safetensors") assert patches == {} and status == "MATCHED NOTHING" out = capsys.readouterr().out assert "doing nothing at all" in out assert "self_attn.qkv" in out # what the file has assert "attn.qkv_proj" in out # what the model wants def test_the_sample_says_how_many_it_did_not_show(mm): """Three names is enough to compare two namings by eye; the count says it is a sample.""" line = mm._sample_keys({f"k{i}": i for i in range(10)}) assert line.count(",") == 2 and "(+7 more)" in line def test_an_empty_side_reads_as_empty_rather_than_blank(mm): assert mm._sample_keys({}) == "(none)" def test_the_curve_note_leads_with_the_answer_and_prints_whole(mm, monkeypatch, capsys): """It is trimmed to its first sentence like everything else, so the verdict — how many dropped, and that the rest still applies — has to BE the first sentence. It also asks for its full text: this only prints when a LoRA is half-landing, which is the one moment the whole explanation is worth its line.""" import funpack_log as L note = mm._adaln_curve_note( _Model({"diffusion_model.blocks.0.adaln_proj.linear.weight": (16128, 64)}, use_curves=True), [("diffusion_model.blocks.0.adaln_proj.linear.weight", (16128, 2688), (16128, 64))]) first = L.compact(note) assert "DROPPED" in first and "rest of the LoRA still applies" in first import inspect assert "full=True" in inspect.getsource(mm.resolve_lora_patches) # ── a wrapper head: rename it, do not strip it ────────────────────────────── # # A PEFT adapter trained against a wrapper that calls the DiT something else carries comfy's # module paths EXACTLY, under a different head — `base_model.model.dit.` where comfy's # key map holds `diffusion_model.`. Stripping leaves `blocks.0...`, comfy's map has no # bare form to land on, and the whole file silently applies nothing. Measured: 532 keys, zero # applied, no error. Renaming the head is the fix, and it is scored against the real key map # alongside every strip so it wins only where it actually lands. def test_a_wrapper_head_is_renamed_to_comfys_not_removed(mm): assert ("base_model.model.dit.", "diffusion_model.") in mm.LORA_KEY_RENAMES def test_renaming_a_head_yields_the_keys_comfys_map_holds(mm): renamed = mm._rekey( {"base_model.model.dit.audio_patch_proj.lora_A.weight": 1, "base_model.model.dit.blocks.0.adaln_proj.linear.lora_B.weight": 2}, "base_model.model.dit.", "diffusion_model.") assert set(renamed) == {"diffusion_model.audio_patch_proj.lora_A.weight", "diffusion_model.blocks.0.adaln_proj.linear.lora_B.weight"} def test_the_empty_replacement_is_a_plain_strip(mm): stripped = mm._rekey({"base_model.model.dit.blocks.0.attn.qkv_proj.lora_A.weight": 1}, "base_model.model.dit.") assert set(stripped) == {"blocks.0.attn.qkv_proj.lora_A.weight"} def test_a_head_the_file_does_not_carry_is_declined_not_applied(mm): """None, so the candidate is skipped rather than scored as a no-op variant of itself.""" assert mm._rekey({"diffusion_model.x.lora_A.weight": 1}, "base_model.model.") is None def test_a_rename_wins_over_a_strip_when_it_matches_more(mm, monkeypatch): """Both are scored against the real key map and the best is kept, so a rename that lands beats a strip that does not without either being ordered ahead of the other.""" key_map = {"diffusion_model.audio_patch_proj": "audio_patch_proj.weight"} monkeypatch.setattr(mm, "_count_lora_matches", lambda lora, km: sum(1 for k in lora if k.rsplit(".lora_", 1)[0] in km)) best, note = mm._best_lora_keying( {"base_model.model.dit.audio_patch_proj.lora_A.weight": 1}, key_map) assert note == "base_model.model.dit. -> diffusion_model." assert "diffusion_model.audio_patch_proj.lora_A.weight" in best def test_a_file_that_needs_nothing_is_left_as_is(mm, monkeypatch): monkeypatch.setattr(mm, "_count_lora_matches", lambda lora, km: len(lora)) lora = {"diffusion_model.x.lora_A.weight": 1} best, note = mm._best_lora_keying(lora, {}) assert best is lora and note == "as-is"