Matrix: every modifier is listed or exempt with a reason, must report active, and must wire in; crossed with samplers and Two stages

Brownian leaves the stock sampler alone when torchsde internals differ. LoRA: unwrapped refiner layout, absent blocks, LoKr.

modules/loaders/lora/keys.py

@@ -46,7 +46,7 @@ def from_diffusers(lora):
         return lora, {}
     out, extra, sizes = {}, {}, {}
     for k, v in lora.items():
-        m = _PART.match(_WRAPPER.sub("", k))                      # whatever training wrapper sits in front of the block names
+        m = _PART.match(_WRAPPER.sub("", k, count=1))                      # whatever training wrapper sits in front of the block names
         if not m:
             out[k] = v
             continue
@@ -83,8 +83,17 @@ def _rekey(lora, prefix, replacement):
     return {(replacement + k[len(prefix):] if k.startswith(prefix) else k): v for k, v in lora.items()}
 
 
+def _absent(sd, key):
+    """True when a patch aims at a weight of the diffusion model that this model does not have (a block number past its
+    depth, a renamed layer): ComfyUI skips such a key without a word, so it must not be counted as applied."""
+    name = key[0] if isinstance(key, tuple) else key
+    return str(name).startswith("diffusion_model.") and name not in sd
+
+
 def _misfit(sd, key, adapter):
     """True when a plain LoRA pair cannot make the weight it names (both dimensions, not the count)."""
+    if getattr(adapter, "name", "lora") != "lora":                 # LoKr, LoHa, ... keep other factors in those slots: not an up/down pair
+        return False
     w = getattr(adapter, "weights", None)
     window = key[1] if isinstance(key, tuple) else None            # (dim, start, length): a LoRA aimed at some rows of a fused weight (Q, K, V of qkv_proj)
     target = sd.get(key[0] if isinstance(key, tuple) else key)
@@ -120,7 +129,7 @@ def match(model, clip, lora):
             if len(got) > len(best):
                 best, how = got, f"{prefix} -> {replacement or '(none)'}"
     sd = model.model.state_dict()
-    bad = [k for k, a in best.items() if _misfit(sd, k, a)]
+    bad = [k for k, a in best.items() if _absent(sd, k) or _misfit(sd, k, a)]
     for k in bad:
         best.pop(k)
     return best, how, len(bad)

modules/loaders/tests/test_diffusers_h3_lora.py

@@ -134,3 +134,29 @@ def test_a_qkv_pair_of_the_wrong_width_is_dropped_not_merged_scrambled(patcher):
     lora[k] = torch.randn(96, RANK)                                   # 96 rows where this model has 64 per part
     patches, how, dropped = match(patcher, None, lora)
     assert dropped >= 1 and len(patches) == 5
+
+
+def test_a_refiner_layout_file_without_a_wrapper_reaches_the_refiner_blocks(patcher):
+    from modules.loaders.lora.keys import match
+    lora, want = _file("token_refiner.refiner_blocks.0", torch.Generator().manual_seed(0))
+    lora = {k[len("transformer."):]: v for k, v in lora.items()}
+    patches, how, dropped = match(patcher, None, lora)
+    assert dropped == 0 and len(patches) == 6, (how, len(patches))
+    state = patcher.model.state_dict()
+    assert all((k[0] if isinstance(k, tuple) else k) in state for k in patches)
+
+
+def test_a_block_the_model_does_not_have_is_dropped_and_counted_not_reported_as_applied(patcher):
+    from modules.loaders.lora.keys import match
+    lora, _ = _file("transformer_blocks.9", torch.Generator().manual_seed(0))            # the tiny model has 2 blocks
+    patches, how, dropped = match(patcher, None, lora)
+    assert not patches and dropped >= 1
+
+
+def test_a_dense_lokr_on_a_native_weight_is_not_mistaken_for_a_misfit_pair(patcher):
+    from modules.loaders.lora.keys import match
+    gen = torch.Generator().manual_seed(0)
+    lora = {"diffusion_model.blocks.1.attn.out_proj.lokr_w1": torch.randn(8, 8, generator=gen),
+            "diffusion_model.blocks.1.attn.out_proj.lokr_w2": torch.randn(8, 8, generator=gen)}
+    patches, how, dropped = match(patcher, None, lora)
+    assert len(patches) == 1 and dropped == 0

modules/sampling/tests/test_combinations.py

@@ -1,6 +1,15 @@
 """Combination rule (Dex): a feature is tested alone AND with the others, in every pair and all together, on the real
 MiniMaxH3 class (tiny, CPU) through ComfyUI's own guider, with and without Two stages. A new feature goes into FEATURES.
 
+Three things keep it honest:
+  * the registry guard: every modifier a module provides must be in FEATURES (run here), in LTX_ONLY (run in
+    test_combinations_ltx.py) or in EXEMPT with the reason it cannot run on a Mac. A new module fails until it is listed;
+  * a feature counts as installed only when its own status line says it is active (not "needs ...", "failed ...");
+  * a feature alone must change what the model is wired with (wrappers, hooks, options), and the ones that act without
+    ratings must change the picture.
+The runs cross the features with the sampler (euler and the SDE ones that use the Brownian noise) and Two stages on/off.
+CFG stays 1 and Two stages keeps its default resize: Dex's decision, not a gap.
+
 Not covered here and said so: triples (the all-on case catches what only shows with many on), and anything that needs a
 GPU or real weights."""
 
@@ -24,10 +33,31 @@ FEATURES = {
     "late_branch": {"block": 2},
     "decisiveness": {},
     "shot_memory": {},
-    "block_map": {},
     "turbo_sampler": {},
     "camera_move": {"pan_x": 0.3, "zoom": 1.2},
     "explore_first_step": {},
+    "alg": {},
+    "block_influence": {},
+    "dynashift": {},
+    "momentum": {},
+    "output_guidance": {},
+    "prompt_nudge": {},
+    "score_slider": {},
+    "shadow_negative": {},
+    "trajectory_guidance": {},
+    "velocity_bias": {},
+    "video_detail": {},
+}
+#: Modifiers that need the LTX audio+video or spatial-latent model: run in test_combinations_ltx.py.
+LTX_ONLY = {"audio_coupling", "context_windows", "identity_bank", "sharpen", "temporal_style"}
+#: Modifiers that cannot run here, and why. Said loudly: a feature in this list is not proven on this machine.
+EXEMPT = {
+    "ptd": "needs a PTD adapter file in ComfyUI/models/ptd (none on this Mac)",
+    "block_repeat": "repeats a block range of the real model; the tiny model has too few blocks (its own tests run it)",
+    "phrase_probe": "an analysis tool run on request, it installs no hook (its own tests cover it)",
+    "taste": "always present: it is the rating store every other feature reads",
+    "two_stage": "a dimension of every run here, not a feature in the subsets",
+    "block_map": "follows a saved map, there is no switch to put in a subset: every run here has a map saved, test_block_map_is_wired_only_when_a_map_is_saved and the bare run below cover no map",
 }
 SUBSETS = [()] + [(f,) for f in FEATURES] + list(itertools.combinations(FEATURES, 2)) + [tuple(FEATURES)]
 
@@ -62,46 +92,167 @@ def _stage_lora(model, stage):
     return kept
 
 
-def _run(model, subset, staged, stage=None):
-    import comfy.samplers
-    from comfy.nested_tensor import NestedTensor
+STAGES = {                                   # one stage, and Two stages with its default resize (CFG is 1: H3 and LTX run at 1)
+    "one-stage": None,
+    "two-stages": {"enabled": True, "size": 0.5, "steps": 2, "back_up": 0.5},
+}
+SAMPLERS = ["euler", "dpmpp_2m_sde", "dpmpp_2m_sde_gpu", "dpmpp_sde"]      # the SDE ones draw Brownian noise
+
+
+def _status_line(status, name):
+    for line in str(status).splitlines():
+        if line.strip().startswith(name + ":"):
+            return line.strip()
+    return None
+
+
+NOT_ACTIVE = ("needs ", "failed", "inactive", "unavailable", "switched off")
+
+
+def _fingerprint(patched):
+    """What the model is wired with: enough to tell a patcher that a feature touched from one it did not."""
+    to = patched.model_options.get("transformer_options", {})
+    return (sorted(map(str, patched.wrappers)), sorted(str(k) for k, v in patched.wrappers.items() if v),
+            sorted(sum(len(w) for w in v.values()) for v in patched.wrappers.values()),
+            sorted(map(str, patched.callbacks)), sorted(str(k) for k in to), sorted(map(str, patched.model_options)),
+            sorted(map(str, getattr(patched, "attachments", {}))), sorted(map(str, patched.object_patches)))
+
+
+def _load(model, subset, stage_settings, lora_stage=None):
     from modules.sampling.modifiers.nodes import FunUILoadModifiers
     settings = {"taste": {"key": "fox"}}
     for f in subset:
         settings[f] = {"enabled": True, **FEATURES[f]}
-    if staged:
-        settings["two_stage"] = {"enabled": True, "size": 0.5, "steps": 2, "back_up": 0.5}
-    patched, status = FunUILoadModifiers.execute(_stage_lora(model, stage) if stage else model.clone(), settings).result
-    missing = [f for f in subset if f not in status]
-    assert not missing, f"never installed (a test that installs nothing proves nothing): {missing}: {status}"
+    if stage_settings:
+        settings["two_stage"] = dict(stage_settings)
+    patched, status = FunUILoadModifiers.execute(_stage_lora(model, lora_stage) if lora_stage else model.clone(), settings).result
+    for f in subset:
+        line = _status_line(status, f)
+        assert line, f"never installed (a test that installs nothing proves nothing): {f}: {status}"
+        assert not any(bad in line.lower() for bad in NOT_ACTIVE), f"{f} did not switch on: {line}"
+    return patched, status
+
+
+def _run(model, subset, staged, stage=None, sampler="euler", cfg=1.0, variant=None):
+    import comfy.samplers
+    from comfy.nested_tensor import NestedTensor
+    stage_settings = variant if variant is not None else (STAGES["two-stages"] if staged else None)
+    patched, _status = _load(model, subset, stage_settings, stage)
     guider = comfy.samplers.CFGGuider(patched)
     cond = [[torch.randn(1, 5, 48), {}]]
-    guider.set_conds(cond, cond)
-    guider.set_cfg(1.0)
+    guider.set_conds(cond, [[torch.randn(1, 5, 48), {}]] if cfg != 1.0 else cond)
+    guider.set_cfg(cfg)
     latent = NestedTensor((torch.zeros(1, 24, 2, 8, 8), torch.zeros(1, 32, 2, 3)))
     noise = NestedTensor((torch.randn(1, 24, 2, 8, 8), torch.randn(1, 32, 2, 3)))
-    return guider.sample(noise, latent, comfy.samplers.sampler_object("euler"), torch.tensor(SIGMAS),
+    return guider.sample(noise, latent, comfy.samplers.sampler_object(sampler), torch.tensor(SIGMAS),
                          seed=1, disable_pbar=True)
 
 
+def _check(out, label):
+    video, audio = out.unbind()
+    assert torch.isfinite(video).all() and torch.isfinite(audio).all(), label
+    assert video.shape == (1, 24, 2, 8, 8) and audio.shape == (1, 32, 2, 3), label
+    assert video.dtype == audio.dtype == torch.float32 and video.device.type == "cpu", label       # the input's dtype and device come back
+
+
+@pytest.fixture
+def fresh(model, tmp_path, monkeypatch):
+    from core import dit_hooks
+    from modules.system.taste import store
+    monkeypatch.setattr(store, "ROOT", tmp_path / "taste")
+    _save_a_map(tmp_path, monkeypatch, dit_hooks.block_count(model))
+    return model
+
+
+def test_every_modifier_a_module_provides_is_in_the_matrix_or_exempt_with_a_reason():
+    import importlib
+    import pkgutil
+
+    import modules
+    found = set()
+    for info in pkgutil.walk_packages(modules.__path__, "modules."):
+        if not info.ispkg or info.name.rsplit(".", 1)[-1] in ("tests", "_core"):
+            continue
+        try:
+            mod = importlib.import_module(info.name)
+        except Exception:                                   # noqa: BLE001 -- a missing dependency is not this test's business
+            continue
+        if hasattr(mod, "install") and hasattr(mod, "ID") and getattr(mod, "PROVIDES", {}).get("modifier") is mod.install:
+            found.add(mod.ID)
+    listed = set(FEATURES) | LTX_ONLY | set(EXEMPT)
+    assert found <= listed, f"add these to FEATURES, LTX_ONLY or EXEMPT (with the reason): {sorted(found - listed)}"
+    assert listed - found <= {"phrase_probe"}, f"listed but no module provides it: {sorted(listed - found)}"
+    assert all(EXEMPT[k] for k in EXEMPT)
+
+
+#: Installed, but wires nothing until its own switch (recording) is on at run time.
+NOT_WIRED_UNTIL_USED = {"block_influence": "measures per-block movement only when recording is on"}
+
+
+@pytest.mark.parametrize("feature", [f for f in FEATURES if f not in NOT_WIRED_UNTIL_USED])
+def test_a_feature_alone_is_really_wired_in(fresh, feature):
+    base, _ = _load(fresh, (), None)
+    patched, status = _load(fresh, (feature,), None)
+    assert _fingerprint(patched) != _fingerprint(base), f"{feature} switched on and changed nothing the model is wired with: {_status_line(status, feature)}"
+
+
+#: Features that act on this tiny model without ratings: they must change the picture, not just install.
+ACTS = ("momentum", "video_detail", "noise_merge", "sharpness", "late_branch", "decisiveness", "turbo_sampler", "camera_move")
+
+
+@pytest.mark.parametrize("feature", ACTS)
+def test_a_feature_that_needs_no_ratings_changes_the_picture(fresh, feature):
+    torch.manual_seed(0)
+    base = _run(fresh, (), False)
+    torch.manual_seed(0)
+    out = _run(fresh, (feature,), False)
+    assert any(not torch.equal(a, b) for a, b in zip(out.unbind(), base.unbind())), f"{feature} ran and changed nothing"
+
+
 @pytest.mark.parametrize("staged", [False, True], ids=["one-stage", "two-stages"])
 @pytest.mark.parametrize("subset", SUBSETS, ids=_name)
-def test_every_combination_finishes_with_a_finite_picture_and_sound(model, tmp_path, monkeypatch, subset, staged):
-    from core import dit_hooks
-    from modules.system.taste import store
-    monkeypatch.setattr(store, "ROOT", tmp_path / "taste")
-    _save_a_map(tmp_path, monkeypatch, dit_hooks.block_count(model))
-    video, audio = _run(model, subset, staged).unbind()
-    assert torch.isfinite(video).all() and torch.isfinite(audio).all(), _name(subset)
-    assert video.shape == (1, 24, 2, 8, 8)
+def test_every_combination_finishes_with_a_finite_picture_and_sound(fresh, subset, staged):
+    _check(_run(fresh, subset, staged), _name(subset))
 
 
-@pytest.mark.parametrize("stage", ["First", "Second"])
+ALONE_AND_ALL = [()] + [(f,) for f in FEATURES] + [tuple(FEATURES)]
+
+
+@pytest.mark.parametrize("variant", STAGES, ids=list(STAGES))
+@pytest.mark.parametrize("sampler", SAMPLERS)
+@pytest.mark.parametrize("subset", [(), tuple(FEATURES)] + [(f,) for f in ("reins", "stas", "noise_merge", "turbo_sampler", "camera_move")], ids=_name)
+def test_features_cross_the_sampler_and_the_two_stages_setting(fresh, subset, sampler, variant):
+    import warnings
+    with warnings.catch_warnings():
+        warnings.filterwarnings("error", message="Should have")             # torchsde's own warning about the Brownian start
+        out = _run(fresh, subset, False, sampler=sampler, variant=STAGES[variant])
+    _check(out, f"{_name(subset)} {sampler} {variant}")
+
+
+@pytest.mark.parametrize("stage", ["First", "Second", "Both"])
 @pytest.mark.parametrize("subset", [(), tuple(FEATURES)], ids=_name)
-def test_a_lora_kept_for_one_stage_runs_with_every_feature(model, tmp_path, monkeypatch, subset, stage):
+def test_a_lora_kept_for_one_stage_runs_with_every_feature(fresh, subset, stage):
+    _check(_run(fresh, subset, True, stage), (_name(subset), stage))
+
+
+def test_block_map_is_wired_only_when_a_map_is_saved(model, tmp_path, monkeypatch):
     from core import dit_hooks
-    from modules.system.taste import store
-    monkeypatch.setattr(store, "ROOT", tmp_path / "taste")
+    from modules.sampling.block_map import measure
+    monkeypatch.setattr(measure, "FOLDER", tmp_path / "maps")
+    monkeypatch.setattr(measure, "SWITCH", tmp_path / "block_map.enabled")
+    without, _ = _load(model, (), None)
     _save_a_map(tmp_path, monkeypatch, dit_hooks.block_count(model))
-    video, audio = _run(model, subset, True, stage).unbind()
-    assert torch.isfinite(video).all() and torch.isfinite(audio).all(), (_name(subset), stage)
+    with_map, status = _load(model, (), None)
+    assert with_map.model_options["transformer_options"].get("funui_block_map") is not None
+    assert without.model_options["transformer_options"].get("funui_block_map") is None, status
+
+
+@pytest.mark.parametrize("sampler", SAMPLERS)
+@pytest.mark.parametrize("subset", [(), tuple(FEATURES)], ids=_name)
+def test_no_saved_map_is_a_plain_run_with_every_feature(model, tmp_path, monkeypatch, subset, sampler):
+    from modules.sampling.block_map import measure
+    from modules.system.taste import store
+    monkeypatch.setattr(measure, "FOLDER", tmp_path / "maps")
+    monkeypatch.setattr(measure, "SWITCH", tmp_path / "block_map.enabled")
+    monkeypatch.setattr(store, "ROOT", tmp_path / "taste")
+    _check(_run(model, subset, True, sampler=sampler), (_name(subset), sampler))