feat(h3): edit the token refiner's output, and h3_prompt_scale

Qwen text reaches the DiT through condition_proj + a 2-block token
refiner ending in an RMSNorm. Everything steered today happens BEFORE
that, so the refiner mixes the edit across tokens and its final norm
renormalizes the magnitude away — what lands is not what was learned.

TokenRefinerEdit patches the refiner's output instead: a scale and/or a
bias over a row range, landing unchanged in the space the 50 blocks read.
project_into_refiner_space maps a conditioning-space direction across with
condition_proj's weight only (a direction is a difference of points, so
the bias would move the origin).

Ships h3_prompt_scale: how loudly the prompt is read through attention,
distinct from h3_gain_prompt (how much each block writes back into those
rows). Confined to the prompt rows via minimax_token_tags, so a wired
reference's vision block is left alone. Free, 1.0 = model not cloned.

The bias path is the mechanism for a learned direction and is not wired to
one yet: a single [hidden] vector added to every prompt row is
prompt-length-independent, which is the shape a learned direction needs
here since Qwen does not pad.

minimax_h3.py

@@ -1184,3 +1184,103 @@ def apply_adaln_gains(model, gains):
     named = ", ".join(f"{n}={v:g}" for n, v in sorted((gains or {}).items()) if float(v) != 1.0)
     return patched, (f"AdaLN modality gain on {count} blocks ({named}). Keyframe pins and "
                      f"reference images ride the VIDEO tag, so the video gain moves them too.")
+
+
+# ── token-refiner output ─────────────────────────────────────────────────────
+#
+# `context` (Qwen, [1, L, 5120]) reaches the DiT through
+# ``condition_proj`` (Linear 5120 -> 5376) then ``token_refiner`` (2 blocks + RMSNorm).
+# Everything FunPack steers today happens BEFORE that, on the Qwen tensor, which means the
+# refiner then processes the edit — its two attention blocks mix it across tokens and its
+# final RMSNorm renormalizes the magnitude away. The direction that lands is not the
+# direction that was learned.
+#
+# Editing the refiner's OUTPUT lands unchanged, and lands in the space the 50 blocks
+# actually consume. Two edits are supported:
+#
+#   scale — multiply the rows. The blocks read text through attention, so magnitude is how
+#           loudly the prompt competes with the picture for attention. Distinct from
+#           h3_gain_prompt, which controls how much each block WRITES BACK into these rows.
+#   bias  — add a direction. A single [hidden] vector added to every prompt row is
+#           PROMPT-LENGTH INDEPENDENT, which per-position edits are not (Qwen does not pad,
+#           so H3's sequence length changes with every prompt edit). That is the shape a
+#           learned direction has to have here to transfer between prompts.
+#
+# Both are confined to a row range: the `<Picture N>` label and vision-block rows sit at the
+# head of this span and are not ours to steer.
+
+
+class TokenRefinerEdit(torch.nn.Module):
+    """Wraps ``token_refiner`` and edits its output over a row range."""
+
+    def __init__(self, inner, scale=1.0, bias=None, row_start=0, row_end=None):
+        super().__init__()
+        self.inner = inner
+        self.scale = float(scale)
+        self.bias = bias
+        self.row_start = int(row_start)
+        self.row_end = row_end
+
+    def forward(self, x, *args, **kwargs):
+        out = self.inner(x, *args, **kwargs)
+        try:
+            rows = out.shape[-2]
+            start = max(0, min(int(self.row_start), rows))
+            end = rows if self.row_end is None else max(start, min(int(self.row_end), rows))
+            if end <= start:
+                return out
+            out = out.clone()
+            span = out[..., start:end, :]
+            if self.scale != 1.0:
+                span = span * self.scale
+            if self.bias is not None:
+                span = span + self.bias.to(device=span.device, dtype=span.dtype)
+            out[..., start:end, :] = span
+        except Exception:      # a refiner shape we do not recognise: leave it alone
+            return self.inner(x, *args, **kwargs)
+        return out
+
+
+def project_into_refiner_space(model, direction):
+    """Map a conditioning-space (Qwen) direction into refiner space with ``condition_proj``.
+
+    The bias, not just the weight, would move the origin — a DIRECTION is a difference of two
+    points, so only the linear part carries over. Returns None when the projection cannot be
+    read or the width does not match.
+    """
+    try:
+        proj = model.get_model_object("diffusion_model.condition_proj")
+        weight = proj.weight
+    except Exception:
+        return None
+    try:
+        vector = direction.reshape(-1).to(dtype=weight.dtype, device=weight.device)
+        if vector.shape[0] != weight.shape[1]:
+            return None
+        return torch.mv(weight, vector)
+    except Exception:
+        return None
+
+
+def apply_token_refiner_edit(model, scale=1.0, bias=None, row_start=0, row_end=None):
+    """Attach a scale and/or bias to the token refiner's output. Returns (model, note).
+
+    ``add_object_patch`` again, so ComfyUI restores it on unpatch and nothing outlives the
+    run. A scale of 1.0 with no bias returns the model unchanged rather than cloning it.
+    """
+    if float(scale) == 1.0 and bias is None:
+        return model, None
+    key = "diffusion_model.token_refiner"
+    try:
+        inner = model.get_model_object(key)
+    except Exception as error:
+        return model, f"token-refiner edit skipped — no token_refiner to patch ({error})."
+    patched = model.clone()
+    patched.add_object_patch(key, TokenRefinerEdit(inner, scale, bias, row_start, row_end))
+    parts = []
+    if float(scale) != 1.0:
+        parts.append(f"scale={float(scale):g}")
+    if bias is not None:
+        parts.append("bias")
+    where = f"rows {row_start}:{row_end if row_end is not None else 'end'}"
+    return patched, f"token-refiner edit ({', '.join(parts)}) on {where}."

samplers.py

@@ -3165,6 +3165,10 @@ class FunPackLTXAVSceneChainSampler:
                     "default": 1.0, "min": 0.0, "max": 2.0, "step": 0.05,
                     "tooltip": "EXPERIMENTAL, MiniMax H3 only. Per-block write gain on the AUDIO rows (see h3_gain_video for the mechanism). This is the one modality whose gate can be moved without touching the other two at the point of writing — though the streams still share one attention pass, so a changed audio row is read by the video on the NEXT block. Below 1.0 the soundtrack is built more conservatively. Free. 1.0 = untouched.",
                 }),
+                "h3_prompt_scale": ("FLOAT", {
+                    "default": 1.0, "min": 0.0, "max": 2.0, "step": 0.05,
+                    "tooltip": "EXPERIMENTAL, MiniMax H3 only. How LOUD the prompt is when the picture reads it. H3 puts the Qwen text through condition_proj + a 2-block token refiner (ending in an RMSNorm) before the DiT sees it, and the blocks then read the text through attention — so the magnitude of those refined rows is how strongly the prompt competes for attention against the picture and the reference. This multiplies them: above 1.0 the prompt is harder to ignore, below 1.0 it recedes and the reference or the anchor gets more say. Applied AFTER the refiner's final norm, so it lands as set rather than being renormalized away. 1.0 = untouched, and the model is not cloned. Free — one multiply on the text rows, no extra model calls. Confined to the PROMPT rows: the <Picture N> label and vision block sit ahead of them and are left alone (read from minimax_token_tags; if those cannot be read the whole text span is scaled and the console says so). DIFFERENT from h3_gain_prompt, which controls how much each block WRITES BACK into the text rows — this controls how loudly they are READ.",
+                }),
                 "alg_anchor": ("BOOLEAN", {
                     "default": False,
                     "tooltip": "EXPERIMENTAL: run ALG's i2v anchor blur (arXiv:2506.08456) on WHATEVER sampler is wired — a stock KSampler with any sampler_name, Hybrid Euler 2S, a two-evals-per-step sampler like heun, anything. The blur de-statics an anchored scene by hiding the anchor's high-frequency detail during the near-pure-noise steps, so the model cannot shortcut to a video that just matches the still. It is the same guidance as the FunPack Distilled Flow sampler's own alg_enabled, and this switch drives that one too when Distilled Flow is the wired sampler, so there is one control wherever you are. The swap is decided by the step's sigma alone, which is an argument of every model call, so it does not need to run inside a sampler's loop. No effect on a scene with no i2v anchor.",
@@ -3978,6 +3982,7 @@ class FunPackLTXAVSceneChainSampler:
         # key, no rating and no Studio state, so it works on a graph with the Refiner absent
         # entirely. 1.0 on all three does not even clone the model.
         model = self._install_h3_adaln_gains(model)
+        model = self._install_h3_prompt_scale(model, positive)
 
         try:
             sampled = comfy.sample.sample_custom(
@@ -6684,6 +6689,56 @@ class FunPackLTXAVSceneChainSampler:
         w_idx = idx % w
         return (w_idx >= (w // 2)).long()
 
+    def _install_h3_prompt_scale(self, model, positive):
+        """Scale the token refiner's OUTPUT over the prompt rows.
+
+        Sampler-side and self-contained like the AdaLN gains: no key, no rating, no learned
+        direction. The row range comes from the conditioning's own `minimax_token_tags`,
+        which H3 carries on every run — Studio is not involved.
+        """
+        scale = float(getattr(self, "_h3_prompt_scale", 1.0))
+        if scale == 1.0:
+            return model
+        try:
+            from . import minimax_h3 as h3mod
+            from . import h3_token_weights as _tw
+        except ImportError:
+            import minimax_h3 as h3mod
+            import h3_token_weights as _tw
+        if not h3mod.is_h3_model(model):
+            _log.note_on_change(
+                "h3:prompt_scale_family",
+                "[FunPackSceneChain] h3_prompt_scale is set but this is not a MiniMax H3 "
+                "model — there is no token refiner to edit. Not applied.")
+            return model
+        start = 0
+        try:
+            entry = positive[0] if isinstance(positive, list) and positive else None
+            meta = entry[1] if isinstance(entry, (list, tuple)) and len(entry) >= 2 else {}
+            cond = entry[0] if entry is not None else None
+            cond_len = int(cond.shape[1]) if hasattr(cond, "shape") and cond.dim() >= 2 else 0
+            region = _tw.prompt_region(meta.get("minimax_token_tags"), cond_len)
+            if region:
+                start = int(region[0])
+            elif cond_len:
+                _log.note_on_change(
+                    "h3:prompt_scale_span",
+                    "[FunPackSceneChain] h3_prompt_scale: the prompt/reference boundary could "
+                    "not be read from the token tags, so the WHOLE text span is scaled — "
+                    "including a reference's vision block, if one is wired.")
+        except Exception as error:  # noqa: BLE001
+            _log.failed("FunPackSceneChain", "h3_prompt_scale row range", error,
+                        "the whole text span is scaled instead of just the prompt")
+        try:
+            patched, note = h3mod.apply_token_refiner_edit(model, scale=scale, row_start=start)
+        except Exception as error:  # noqa: BLE001
+            _log.failed("FunPackSceneChain", "h3_prompt_scale", error,
+                        "the prompt is read at its trained strength")
+            return model
+        if note:
+            print(f"[FunPackSceneChain] {note}")
+        return patched
+
     def _install_h3_adaln_gains(self, model):
         """Scale every DiT block's AdaLN gates per modality, from the three sampler widgets.
 
@@ -7051,6 +7106,7 @@ class FunPackLTXAVSceneChainSampler:
                second_pass_sampler=None,
                h3_audio_clock=False,
                h3_gain_video=1.0, h3_gain_prompt=1.0, h3_gain_audio=1.0,
+               h3_prompt_scale=1.0,
                audio_vae=None, h3_keyframes=None,
                unique_id=None, prompt=None):
         if not isinstance(positive, list) or not positive:
@@ -7069,6 +7125,7 @@ class FunPackLTXAVSceneChainSampler:
         self._h3_gain_video = max(0.0, min(2.0, float(h3_gain_video)))
         self._h3_gain_prompt = max(0.0, min(2.0, float(h3_gain_prompt)))
         self._h3_gain_audio = max(0.0, min(2.0, float(h3_gain_audio)))
+        self._h3_prompt_scale = max(0.0, min(2.0, float(h3_prompt_scale)))
         # The gate every rating-driven wrapper shares. On H3 it is read off the schedule's
         # own base grid; on LTX it stays the absolute-sigma gate it was validated with.
         _steer_ramp = _make_steer_ramp(sigmas, self._is_h3)

tests/test_h3_token_refiner_edit.py

@@ -0,0 +1,181 @@
+"""Editing the token refiner's OUTPUT rather than the Qwen tensor before it.
+
+H3 takes `context` [1, L, 5120] through condition_proj (-> 5376) and a 2-block token
+refiner that ends in an RMSNorm. Everything FunPack steers today happens BEFORE that, so
+the refiner processes the edit: its attention mixes it across tokens and its final norm
+renormalizes the magnitude away. An edit applied AFTER lands as set, in the space the 50
+blocks consume.
+
+Two edits: `scale` (how loudly the prompt is read through attention) and `bias` (a single
+[hidden] vector added to every prompt row — prompt-LENGTH-INDEPENDENT, unlike anything
+per-position, because Qwen does not pad and H3's sequence length moves with every edit).
+"""
+import inspect
+import sys
+from pathlib import Path
+
+import pytest
+import torch
+
+sys.path.insert(0, str(Path(__file__).resolve().parents[1]))
+
+import minimax_h3 as h3
+
+HID = 6
+
+
+class FakeRefiner(torch.nn.Module):
+    def forward(self, x, *a, **k):
+        return torch.ones(x.shape[-2], HID)
+
+
+class FakePatcher:
+    def __init__(self, proj_weight=None):
+        self.objects = {"diffusion_model.token_refiner": FakeRefiner()}
+        if proj_weight is not None:
+            self.objects["diffusion_model.condition_proj"] = \
+                torch.nn.Linear(proj_weight.shape[1], proj_weight.shape[0])
+            with torch.no_grad():
+                self.objects["diffusion_model.condition_proj"].weight.copy_(proj_weight)
+        self.patched = {}
+
+    def get_model_object(self, name):
+        if name in self.patched:
+            return self.patched[name]
+        if name not in self.objects:
+            raise KeyError(name)
+        return self.objects[name]
+
+    def clone(self):
+        other = FakePatcher.__new__(FakePatcher)
+        other.objects, other.patched = self.objects, dict(self.patched)
+        return other
+
+    def add_object_patch(self, name, obj):
+        self.patched[name] = obj
+
+
+def run(edit, rows=8):
+    return edit(torch.zeros(1, rows, 4))
+
+
+# ── the edit ────────────────────────────────────────────────────────────────
+
+def test_scale_multiplies_the_rows():
+    out = run(h3.TokenRefinerEdit(FakeRefiner(), scale=0.5))
+    assert torch.allclose(out, torch.full((8, HID), 0.5))
+
+
+def test_rows_before_the_start_are_untouched():
+    """The <Picture N> label and vision block sit at the head of this span."""
+    out = run(h3.TokenRefinerEdit(FakeRefiner(), scale=0.0, row_start=3))
+    assert torch.all(out[:3] == 1.0)
+    assert torch.all(out[3:] == 0.0)
+
+
+def test_a_row_end_bounds_the_other_side():
+    out = run(h3.TokenRefinerEdit(FakeRefiner(), scale=2.0, row_start=2, row_end=5))
+    assert torch.all(out[:2] == 1.0) and torch.all(out[5:] == 1.0)
+    assert torch.all(out[2:5] == 2.0)
+
+
+def test_a_bias_is_added_to_every_row_in_range():
+    bias = torch.arange(HID, dtype=torch.float32)
+    out = run(h3.TokenRefinerEdit(FakeRefiner(), bias=bias, row_start=1))
+    assert torch.allclose(out[1], 1.0 + bias)
+    assert torch.allclose(out[5], 1.0 + bias)          # same vector on every row
+    assert torch.all(out[0] == 1.0)
+
+
+def test_the_bias_is_length_independent():
+    """The property that makes it a usable learning target: one vector, any prompt length."""
+    bias = torch.full((HID,), 0.25)
+    short = run(h3.TokenRefinerEdit(FakeRefiner(), bias=bias), rows=4)
+    long = run(h3.TokenRefinerEdit(FakeRefiner(), bias=bias), rows=40)
+    assert torch.allclose(short[0], long[0])
+
+
+def test_scale_and_bias_compose_scale_first():
+    bias = torch.full((HID,), 1.0)
+    out = run(h3.TokenRefinerEdit(FakeRefiner(), scale=2.0, bias=bias))
+    assert torch.allclose(out, torch.full((8, HID), 3.0))   # 1*2 + 1
+
+
+def test_an_out_of_range_start_is_a_no_op_not_a_crash():
+    out = run(h3.TokenRefinerEdit(FakeRefiner(), scale=0.0, row_start=99))
+    assert torch.all(out == 1.0)
+
+
+def test_the_inner_output_is_not_mutated_in_place():
+    inner = FakeRefiner()
+    edit = h3.TokenRefinerEdit(inner, scale=0.0)
+    run(edit)
+    assert torch.all(inner(torch.zeros(1, 8, 4)) == 1.0)
+
+
+# ── attaching it ────────────────────────────────────────────────────────────
+
+def test_no_edit_returns_the_model_unpatched():
+    model = FakePatcher()
+    out, note = h3.apply_token_refiner_edit(model, scale=1.0, bias=None)
+    assert out is model and note is None and model.patched == {}
+
+
+def test_it_patches_the_refiner_object():
+    model = FakePatcher()
+    out, note = h3.apply_token_refiner_edit(model, scale=1.2, row_start=4)
+    assert isinstance(out.patched["diffusion_model.token_refiner"], h3.TokenRefinerEdit)
+    assert "scale=1.2" in note and "rows 4:end" in note
+
+
+def test_object_patches_so_nothing_outlives_the_run():
+    src = inspect.getsource(h3.apply_token_refiner_edit)
+    assert "add_object_patch" in src and "model.clone()" in src
+
+
+def test_a_model_without_a_refiner_declines_with_a_reason():
+    class Bare:
+        def get_model_object(self, name):
+            raise KeyError(name)
+    out, note = h3.apply_token_refiner_edit(Bare(), scale=0.5)
+    assert isinstance(out, Bare) and "no token_refiner" in note
+
+
+# ── projecting a conditioning-space direction ───────────────────────────────
+
+def test_a_direction_is_projected_by_the_weight_only():
+    """A direction is a difference of two points, so the bias would move the origin."""
+    weight = torch.randn(HID, 4)
+    model = FakePatcher(proj_weight=weight)
+    direction = torch.randn(4)
+    out = h3.project_into_refiner_space(model, direction)
+    assert torch.allclose(out, weight @ direction, atol=1e-5)
+
+
+def test_a_mismatched_width_is_refused():
+    model = FakePatcher(proj_weight=torch.randn(HID, 4))
+    assert h3.project_into_refiner_space(model, torch.randn(9)) is None
+
+
+def test_no_projection_available_returns_none():
+    assert h3.project_into_refiner_space(FakePatcher(), torch.randn(4)) is None
+
+
+# ── independence from the refinement path ───────────────────────────────────
+
+def test_the_installer_reads_no_key_or_rating():
+    import samplers
+    fn = samplers.FunPackLTXAVSceneChainSampler._install_h3_prompt_scale
+    src = inspect.getsource(fn)
+    for line in (inspect.getdoc(fn) or "").splitlines():
+        src = src.replace(line, "")
+    for forbidden in ("refinement_key", "rating", "value_fn", "global_state", "phrase_memory"):
+        assert forbidden not in src
+
+
+def test_the_widget_defaults_to_untouched():
+    import samplers
+    spec = samplers.FunPackLTXAVSceneChainSampler.INPUT_TYPES()
+    fields = {**spec.get("required", {}), **spec.get("optional", {})}
+    assert fields["h3_prompt_scale"][1]["default"] == 1.0
+    assert "h3_gain_prompt" in fields["h3_prompt_scale"][1]["tooltip"]   # says how they differ