From f99944e5e4bc762f794df81e1afd1db0d3d03dd9 Mon Sep 17 00:00:00 2001 From: digital-garbage <133010426+digital-garbage@users.noreply.github.com> Date: Sun, 27 Sep 2026 18:53:51 +0300 Subject: [PATCH] fix: DynaShift and taste guidance read the step from ComfyUI's positional call; they never steered Co-Authored-By: Claude Opus 5.5 --- core/streams.py | 17 +++++++++++++++++ core/tests/test_streams.py | 14 ++++++++++++++ modules/models/minimax_h3/__init__.py | 4 +++- modules/models/minimax_h3/rows.py | 11 +++++++++++ modules/sampling/dynashift/__init__.py | 7 ++++--- .../sampling/dynashift/tests/test_dynashift.py | 6 +++--- modules/sampling/output_guidance/__init__.py | 3 ++- .../tests/test_output_guidance.py | 6 +++--- .../sampling/trajectory_guidance/__init__.py | 3 ++- .../tests/test_trajectory_guidance.py | 2 +- 10 files changed, 60 insertions(+), 13 deletions(-) create mode 100644 core/tests/test_streams.py diff --git a/core/streams.py b/core/streams.py index 9d1c98f1..c55f852d 100644 --- a/core/streams.py +++ b/core/streams.py @@ -14,3 +14,20 @@ def video_of(x): if getattr(x, "dim", None) and not getattr(x, "is_nested", False) and x.dim() == 5: return x, lambda video: video return None + + +MODEL_ARGS = ("c_concat", "c_crossattn", "control", "transformer_options") + + +def model_args(args, kwargs) -> dict: + """An APPLY_MODEL wrapper's call, by name. + + ComfyUI hands c_concat, c_crossattn, control and transformer_options on + POSITIONALLY (model_base.apply_model), and the model's extra conds by name. + Reading `kwargs.get("transformer_options")` alone finds nothing on a real + run -- a step gate built on it never opens, silently. A wrapper that calls + on with changed values may pass everything by name; this reads both shapes. + """ + named = dict(zip(MODEL_ARGS, args)) + named.update(kwargs) + return named diff --git a/core/tests/test_streams.py b/core/tests/test_streams.py new file mode 100644 index 00000000..d3204f83 --- /dev/null +++ b/core/tests/test_streams.py @@ -0,0 +1,14 @@ +"""model_args: an APPLY_MODEL wrapper reads its call the same whichever way it came.""" + +from core.streams import model_args + + +def test_positional_as_comfyui_calls_it(): + # model_base.apply_model: execute(x, t, c_concat, c_crossattn, control, transformer_options, **kw) + named = model_args((None, "c", None, {"sigmas": 1}), {"minimax_payload": {}}) + assert named["transformer_options"] == {"sigmas": 1} + assert named["c_crossattn"] == "c" and "minimax_payload" in named + + +def test_by_name_as_an_outer_wrapper_may_call_on(): + assert model_args((), {"transformer_options": {"a": 1}})["transformer_options"] == {"a": 1} diff --git a/modules/models/minimax_h3/__init__.py b/modules/models/minimax_h3/__init__.py index 619312f2..bd60ab19 100644 --- a/modules/models/minimax_h3/__init__.py +++ b/modules/models/minimax_h3/__init__.py @@ -33,6 +33,7 @@ from .anchor import rescale_pins as _rescale_pins from .latent_upscaler import load_upscaler as _load_upscaler from . import prompt_layout as _prompt_layout from .rows import target_rows as _target_rows +from .rows import text_rows as _text_rows from .rope import query_rotation as _query_rotation from .streams import video_stream as _video_stream @@ -195,4 +196,5 @@ PROVIDES = {"empty_latent": empty_latent, "decode": decode, "detect": detect, "prompt_rows": _prompt_layout.prompt_rows, "video_time_rows": _prompt_layout.video_time_rows, "query_rotation": _query_rotation, - "video_stream": _video_stream} + "video_stream": _video_stream, + "text_rows": _text_rows} diff --git a/modules/models/minimax_h3/rows.py b/modules/models/minimax_h3/rows.py index bd3a1876..730c5dd0 100644 --- a/modules/models/minimax_h3/rows.py +++ b/modules/models/minimax_h3/rows.py @@ -33,3 +33,14 @@ def target_rows(mod_segments, seq_len, device, stream): mask = torch.zeros(seq_len, dtype=torch.bool, device=device) mask[a:b] = True return mask + + +def text_rows(model_kwargs, length): + """[length] bool mask of the WORD rows in the text conditioning the model is + called with, or None when it can't be told. H3's text span carries picture + tokens too (reference images, tag 0); words are tag 1.""" + payload = (model_kwargs or {}).get("minimax_payload") + tags = payload.get("text_token_tags") if isinstance(payload, dict) else None + if tags is None or not hasattr(tags, "view") or tags.numel() != length: + return None + return tags.view(-1) == 1 diff --git a/modules/sampling/dynashift/__init__.py b/modules/sampling/dynashift/__init__.py index 2042aa1c..acd087b3 100644 --- a/modules/sampling/dynashift/__init__.py +++ b/modules/sampling/dynashift/__init__.py @@ -186,7 +186,8 @@ def install(patcher, values, key): def apply_model(executor, x, t, *args, **kwargs): out = executor(x, t, *args, **kwargs) - to = kwargs.get("transformer_options") + named = streams.model_args(args, kwargs) + to = named.get("transformer_options") split = streams.video_of(out) if split is None: _say("off this run: could not find the picture in this model's latent") @@ -194,13 +195,13 @@ def install(patcher, values, key): video, rebuild = split if dit_hooks.last_step(to): captured["latent"] = video[0].detach().half() - cond = _pooled(kwargs.get("c_crossattn")) + cond = _pooled(named.get("c_crossattn")) if cond is not None: captured["cond"] = cond amount = strength * dit_hooks.late_half(to) if amount <= 0.0: return out - shifted = shift(video, live["bank"], amount, threshold, kwargs.get("c_crossattn")) + shifted = shift(video, live["bank"], amount, threshold, named.get("c_crossattn")) return out if shifted is None else rebuild(shifted) patcher.add_wrapper_with_key(WrappersMP.APPLY_MODEL, key, apply_model) diff --git a/modules/sampling/dynashift/tests/test_dynashift.py b/modules/sampling/dynashift/tests/test_dynashift.py index dc509a69..7871c6e5 100644 --- a/modules/sampling/dynashift/tests/test_dynashift.py +++ b/modules/sampling/dynashift/tests/test_dynashift.py @@ -71,7 +71,7 @@ def test_it_banks_the_last_step_and_steers_once_rated(tiny_h3, monkeypatch): monkeypatch.setattr(store, "current_prompt_id", lambda i=i: f"run-{i}") wrap, outer, status = load() for o in outer: - o(lambda: wrap(lambda *a, **k: x0, x0, None, transformer_options=_to(3))) + o(lambda: wrap(lambda *a, **k: x0, x0, None, None, None, None, _to(3))) assert store.rate(f"run-{i}", rating)["recorded"] == ["dynashift"] wrap, outer, _status = load() @@ -80,8 +80,8 @@ def test_it_banks_the_last_step_and_steers_once_rated(tiny_h3, monkeypatch): for o in outer: o(lambda: None) # a run starts: the bank is read assert any("2 disliked, 0 liked banked" in e["message"] for e in log.history()) - early = wrap(lambda *a, **k: x0, x0, None, transformer_options=_to(1)) + early = wrap(lambda *a, **k: x0, x0, None, None, None, None, _to(1)) assert early is x0 # first half: untouched - late = wrap(lambda *a, **k: x0, x0, None, transformer_options=_to(3)) + late = wrap(lambda *a, **k: x0, x0, None, None, None, None, _to(3)) assert not torch.allclose(late.tensors[0], bad) assert late.tensors[1] is audio # sound untouched diff --git a/modules/sampling/output_guidance/__init__.py b/modules/sampling/output_guidance/__init__.py index 45ff3776..2f7674dc 100644 --- a/modules/sampling/output_guidance/__init__.py +++ b/modules/sampling/output_guidance/__init__.py @@ -72,7 +72,8 @@ def install(patcher, values, key): def apply_model(executor, x, t, *args, **kwargs): out = executor(x, t, *args, **kwargs) - to = kwargs.get("transformer_options") + named = streams.model_args(args, kwargs) + to = named.get("transformer_options") split = streams.video_of(out) if split is None: _say("off this run: could not find the picture in this model's latent") diff --git a/modules/sampling/output_guidance/tests/test_output_guidance.py b/modules/sampling/output_guidance/tests/test_output_guidance.py index b84a181b..4e498411 100644 --- a/modules/sampling/output_guidance/tests/test_output_guidance.py +++ b/modules/sampling/output_guidance/tests/test_output_guidance.py @@ -58,13 +58,13 @@ def test_late_guidance_banks_the_final_picture_and_steers_only_late(tiny_h3, mon x0 = _x0() wrap, outer = _load(tiny_h3, "output_guidance") for o in outer: - o(lambda: wrap(lambda *a, **k: x0, x0, None, transformer_options=_to(3))) + o(lambda: wrap(lambda *a, **k: x0, x0, None, None, None, None, _to(3))) assert store.rate("run-1", "liked")["recorded"] == ["x0_final"] _teach("x0_final", ["final"]) wrap, outer = _load(tiny_h3, "output_guidance") _start(outer) - assert wrap(lambda *a, **k: x0, x0, None, transformer_options=_to(1)) is x0 - late = wrap(lambda *a, **k: x0, x0, None, transformer_options=_to(3)) + assert wrap(lambda *a, **k: x0, x0, None, None, None, None, _to(1)) is x0 + late = wrap(lambda *a, **k: x0, x0, None, None, None, None, _to(3)) assert not torch.allclose(late.tensors[0], x0.tensors[0]) assert late.tensors[1] is x0.tensors[1] diff --git a/modules/sampling/trajectory_guidance/__init__.py b/modules/sampling/trajectory_guidance/__init__.py index 93ec8c31..6f1b97df 100644 --- a/modules/sampling/trajectory_guidance/__init__.py +++ b/modules/sampling/trajectory_guidance/__init__.py @@ -78,7 +78,8 @@ def install(patcher, values, key): def apply_model(executor, x, t, *args, **kwargs): out = executor(x, t, *args, **kwargs) - q = quarter(kwargs.get("transformer_options")) + named = streams.model_args(args, kwargs) + q = quarter(named.get("transformer_options")) split = streams.video_of(out) if q is None or split is None: _say("off this run: could not place this step on the schedule or find the picture") diff --git a/modules/sampling/trajectory_guidance/tests/test_trajectory_guidance.py b/modules/sampling/trajectory_guidance/tests/test_trajectory_guidance.py index 137e4bea..ab960ede 100644 --- a/modules/sampling/trajectory_guidance/tests/test_trajectory_guidance.py +++ b/modules/sampling/trajectory_guidance/tests/test_trajectory_guidance.py @@ -60,7 +60,7 @@ def test_early_guidance_steers_in_the_first_quarter_and_banks_every_quarter(tiny wrap, outer = _load(tiny_h3, "trajectory_guidance") def run(): - return [wrap(lambda *a, **k: x0, x0, None, transformer_options=_to(i)) for i in range(4)] + return [wrap(lambda *a, **k: x0, x0, None, None, None, None, _to(i)) for i in range(4)] results = {} for o in outer: