From 5f6ed35929ae8eb9f0d207ce692c09dfcd018015 Mon Sep 17 00:00:00 2001 From: digital-garbage <133010426+digital-garbage@users.noreply.github.com> Date: Thu, 5 Nov 2026 15:58:29 +0200 Subject: [PATCH] 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. --- core/brownian.py | 26 +++ core/tests/test_brownian.py | 13 ++ modules/loaders/lora/keys.py | 13 +- .../loaders/tests/test_diffusers_h3_lora.py | 26 +++ modules/sampling/tests/test_combinations.py | 203 +++++++++++++++--- 5 files changed, 253 insertions(+), 28 deletions(-) diff --git a/core/brownian.py b/core/brownian.py index be65036a..407356b8 100644 --- a/core/brownian.py +++ b/core/brownian.py @@ -15,6 +15,29 @@ import torch _installed = False +def _internals_present(ksampling) -> bool: + """The replacement reads a few private attributes of ComfyUI's BatchedBrownianTree and torchsde's interval. A version + that renamed any of them gets ComfyUI's own sampler left as it is, instead of failing at the first noise draw.""" + try: + tree = ksampling.BatchedBrownianTree(torch.zeros(1, 1), torch.tensor(0.1), torch.tensor(0.9), 1, cpu=True) + span = tree.trees[0]._interval + float(span._start), float(span._end) + tree.sort, tree.cpu_tree, tree.sign, tree.batched + tree.trees[0](0.2, 0.3) # takes plain numbers + return True + except Exception: # noqa: BLE001 -- any mismatch means: leave the stock sampler alone + return False + + +def _log_skipped(): + try: + from . import log + log.warn("FunUI Brownian", "this ComfyUI or torchsde has a different Brownian tree than expected, so the stock sampler is left as it is " + "(SDE samplers may print torchsde's start warnings; the pictures are not affected)") + except Exception: # noqa: BLE001 -- a log line is never worth failing the install + pass + + def install(): """Swap ComfyUI's BrownianTreeNoiseSampler for the widened one. Safe to call twice.""" global _installed @@ -23,6 +46,9 @@ def install(): _installed = True return original = ksampling.BrownianTreeNoiseSampler + if not _internals_present(ksampling): + _log_skipped() + return class BrownianTreeNoiseSampler(original): _funui_widened = True diff --git a/core/tests/test_brownian.py b/core/tests/test_brownian.py index 5a8b1e52..1504c9fe 100644 --- a/core/tests/test_brownian.py +++ b/core/tests/test_brownian.py @@ -107,3 +107,16 @@ def test_the_noise_dtype_is_the_stock_dtype_for_a_reduced_precision_latent(patch want = stock(x, sig.min(), sig.max(), seed=7, cpu=True)(sig[0], sig[1]) got = patched(x, sig.min(), sig.max(), seed=7, cpu=True)(sig[0], sig[1]) assert got.dtype == want.dtype + + +def test_a_torchsde_with_other_internals_leaves_the_stock_sampler_in_place(monkeypatch): + original = ksampling.BrownianTreeNoiseSampler + mod = _brownian() + mod._installed = False + monkeypatch.setattr(mod, "_internals_present", lambda k: False) + mod.install() + assert ksampling.BrownianTreeNoiseSampler is original and not mod._installed + + +def test_the_real_internals_are_recognised(): + assert _brownian()._internals_present(ksampling) is True diff --git a/modules/loaders/lora/keys.py b/modules/loaders/lora/keys.py index e3a80018..899cef52 100644 --- a/modules/loaders/lora/keys.py +++ b/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) diff --git a/modules/loaders/tests/test_diffusers_h3_lora.py b/modules/loaders/tests/test_diffusers_h3_lora.py index 37736da6..32ffea74 100644 --- a/modules/loaders/tests/test_diffusers_h3_lora.py +++ b/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 diff --git a/modules/sampling/tests/test_combinations.py b/modules/sampling/tests/test_combinations.py index e8513529..d8f4a606 100644 --- a/modules/sampling/tests/test_combinations.py +++ b/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))