diff --git a/conftest.py b/conftest.py index 59de0475..e839377c 100644 --- a/conftest.py +++ b/conftest.py @@ -248,9 +248,12 @@ _TINY_LTX = [] @pytest.fixture -def tiny_ltx(comfyui): +def tiny_ltx(comfyui, monkeypatch): """One model for the whole run (building it costs seconds): tests clone its patcher and never - leave anything installed on the model itself.""" + leave anything installed on the model itself. The held-back list is emptied, so a feature is + tested on LTX whether or not it is shown there yet (`ltx_held_back` puts it back).""" + from modules.models import ltx + monkeypatch.setattr(ltx, "HELD_BACK", ()) if not _TINY_LTX: _TINY_LTX.append(TinyLTX()) return _TINY_LTX[0] diff --git a/core/traits.py b/core/traits.py index 039c319b..c1ced890 100644 --- a/core/traits.py +++ b/core/traits.py @@ -43,14 +43,26 @@ def split(specs: Iterable[ModuleSpec], available: Sequence[str]) -> Tuple[List[M genuinely cannot work on. """ have: Set[str] = set(available or ()) + held = held_back(have) compatible, incompatible = [], [] for spec in specs: - (compatible if set(spec.requires) <= have else incompatible).append(spec) + (compatible if set(spec.requires) <= have and spec.id not in held else incompatible).append(spec) return compatible, incompatible +#: A model may announce `hides:` for a module that runs on it but has not been proven on it yet. +#: Hidden like an incompatible one, and one line to take back when it has been. +HOLDS = "hides:" + + +def held_back(available: Iterable[str]) -> Set[str]: + return {t[len(HOLDS):] for t in (available or ()) if str(t).startswith(HOLDS)} + + def missing_for(spec: ModuleSpec, available: Sequence[str]) -> List[str]: """Which traits a module wanted and did not get. For the modules dump.""" + if spec.id in held_back(available): + return ["a proven run on this model (held back for now)"] return sorted(set(spec.requires) - set(available or ())) diff --git a/modules/models/ltx/__init__.py b/modules/models/ltx/__init__.py index 3b73bf7f..26634313 100644 --- a/modules/models/ltx/__init__.py +++ b/modules/models/ltx/__init__.py @@ -63,11 +63,20 @@ def detect(keys) -> bool: return any(all(f"{prefix}{key}" in keyset for key in _SIGNATURE_KEYS) for prefix in _PREFIXES) +#: Enhancements that run on LTX but have not been proven on a real run there: held back (hidden) until they have. +#: Proven on LTX so far: Noise merging and the Attention explorer. Take an id off this list once it has been. +HELD_BACK = ("conditioning_identity_transfer", "guide_frames", "camera_move", "shot_memory", "audio_coupling", "decisiveness", + "dynashift", "momentum", "output_guidance", "reins", "sharpness", "trajectory_guidance", "velocity_bias", "alg", + "attention_temperature", "block_repeat", "context_windows", "explore_first_step", "identity_bank", "prompt_nudge", + "score_slider", "temporal_style") + + def probe_traits(keys) -> list: """What `traits()` would say, read before anything loads.""" if not detect(keys): return [] - return ["audio_stream", "temporal_latent", "temporal_compression", "ltx_av", "block_scope", "block_rows"] + return ["audio_stream", "temporal_latent", "temporal_compression", "ltx_av", "block_scope", "block_rows", + *(_traits.HOLDS + i for i in HELD_BACK)] def probe_blocks(keys): @@ -86,7 +95,7 @@ def traits(model): # returns (video, audio), where the hooks that edit rows expect one packed tensor. # block_rows: the hooks that read or push the video rows of a block's output (REINS, Block repeat) # handle the pair too (see those modules). - return ["audio_stream", "ltx_av", "block_scope", "block_rows"] + return ["audio_stream", "ltx_av", "block_scope", "block_rows", *(_traits.HOLDS + i for i in HELD_BACK)] def frames_for(length: int) -> int: diff --git a/modules/models/ltx/tests/test_ltx.py b/modules/models/ltx/tests/test_ltx.py index ad26e7ef..bb29a217 100644 --- a/modules/models/ltx/tests/test_ltx.py +++ b/modules/models/ltx/tests/test_ltx.py @@ -291,3 +291,17 @@ def test_the_h3_block_count_ignores_its_token_refiner(comfyui): h3 = import_module("modules.models.minimax_h3") keys = {f"blocks.{i}.attn.q.weight" for i in range(50)} | {f"token_refiner.blocks.{i}.x" for i in range(2)} assert h3.probe_blocks(keys) == 50 + + +def test_unproven_enhancements_are_held_back_but_the_proven_ones_stay(comfyui): + from modules import models # noqa: F401 + from modules.models import ltx + from core import traits + traits_ = ltx.probe_traits({"audio_adaln_single.linear.weight"}) + held = traits.held_back(traits_) + assert {"reins", "block_repeat", "camera_move"} <= held + assert not held & {"noise_merge", "attention_bandit", "taste"} + spec = type("S", (), {"id": "reins", "requires": []})() + assert traits.missing_for(spec, traits_) == ["a proven run on this model (held back for now)"] + other = type("S", (), {"id": "noise_merge", "requires": []})() + assert traits.split([spec, other], traits_) == ([other], [spec])