digital-garbage
d80c17bfaa
Validated on the rental and it does not earn its place: 3.5 minutes against 1 for the same clip, worse picture, garbage audio. Some of that was mine — no weight prefetch across the 50 blocks, a host-offloaded K/V cache, and an extra commit forward per chunk — but fixing those only buys parity on speed, and the technique's real win is unbounded length, which H3 already covers natively at the lengths in use here. Gone: h3_causal.py, the four h3_causal_* sampler widgets and their Engine Settings rows, `chunk_causal` on the diffusion loader, and five test files. KEPT, because it is not about this at all: the LoRA key RENAME table. Stripping a wrapper head leaves `blocks.0...`, comfy's key map has no bare form, and the file silently applies nothing — 532 keys, zero applied, no error. That was breaking any PEFT-style adapter named this way, not just the one that found it. Its tests moved to test_lora_shape_mismatch.py alongside the MATCHED NOTHING diagnostics. Two INPUT_TYPES widgets go with this (h3_causal_* on the sampler, chunk_causal on the loader). A saved workflow carrying them loads fine — both were optional and are matched by name — but the values are ignored from here. Tests: 1700 python, 68 JS.
358 lines
17 KiB
Python
358 lines
17 KiB
Python
"""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.<path>` where comfy's
|
|
# key map holds `diffusion_model.<path>`. 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"
|