fix: DynaShift and taste guidance read the step from ComfyUI's positional call; they never steered

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

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

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}

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}

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

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)

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

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")

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]

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")

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: