feat: H3 attention decoupling -- damp video<->audio joint self-attention, no S^2 mask

H3 has no cross-attn to mask (unlike LTXAV's Bounded Attention), so a dense
attention mask would be GB-scale over its ~37k-token packed sequence. Uses
comfy's own optimized_attention_override extension point instead: runs the
same joint attention as separate video-query/audio-query passes against a K
where the other modality's rows are scaled down, at the same total FLOPs as
one pass. New h3_av_decouple sampler input + Engine Settings entry, off by
default (0 installs nothing).

Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>

h3_repr_steering.py

@@ -78,6 +78,24 @@ def state_path(refinement_key):
     return refinement_state_path(refinement_key, "repr_steer", prefix="refine_v2", extension="pt")
 
 
+def _mask_from_mod_segments(mod_segments, seq_len, device, tag):
+    """Shared by video_mask_from_mod_segments/audio_mask_from_mod_segments -- see those for
+    what `tag` means. Returns None (not an all-False mask) when nothing matched, so callers
+    no-op on a shape this has never seen rather than silently mask everything out."""
+    mask = torch.zeros(seq_len, dtype=torch.bool, device=device)
+    found = False
+    for a, b, row in mod_segments or ():
+        if torch.is_tensor(row):
+            tags = (row % 3) == tag
+            if tags.any():
+                mask[a:b][tags] = True
+                found = True
+        elif int(row) % 3 == tag:
+            mask[a:b] = True
+            found = True
+    return mask if found else None
+
+
 def video_mask_from_mod_segments(mod_segments, seq_len, device):
     """mod_segments (from the H3 block hook's own args) -> a [seq_len] bool mask, True on
     VIDEO rows. Each entry is (a, b, row) where row is either a scalar `t_row*3 + tag` or,
@@ -85,18 +103,13 @@ def video_mask_from_mod_segments(mod_segments, seq_len, device):
     construction (`seg_tag` in model.py: video/cond/ref_img all tag 0). Returns None if
     nothing matches, which happens on a shape this has never seen -- callers no-op rather
     than guess."""
-    mask = torch.zeros(seq_len, dtype=torch.bool, device=device)
-    found = False
-    for a, b, row in mod_segments or ():
-        if torch.is_tensor(row):
-            tags = (row % 3) == 0
-            if tags.any():
-                mask[a:b][tags] = True
-                found = True
-        elif int(row) % 3 == 0:
-            mask[a:b] = True
-            found = True
-    return mask if found else None
+    return _mask_from_mod_segments(mod_segments, seq_len, device, 0)
+
+
+def audio_mask_from_mod_segments(mod_segments, seq_len, device):
+    """Same as video_mask_from_mod_segments but for AUDIO rows (tag 2 -- audio/cond_audio/
+    ref_audio in model.py's `seg_tag`)."""
+    return _mask_from_mod_segments(mod_segments, seq_len, device, 2)
 
 
 def capture(hidden_state, video_mask):

movie_editor/backend/settings_card.py

@@ -268,6 +268,7 @@ _EDITOR_DEFAULTS = {
     "embed_guidance_source": "relative",
     "embed_guidance_strength": 0.02,
     "frame_overlap": 16,
+    "h3_av_decouple": 0.0,
     "h3_block_repeat": "",
     "h3_block_repeat_times": 1,
     "h3_block_repeat_span_loop": False,

movie_editor/frontend/engine_settings.js

@@ -334,6 +334,9 @@
     { name: "h3_repr_steering_block", label: "REINS steer at block(s)", kind: "text", default: "25", dependsOn: "h3_repr_steering",
       hint: "TEST-ONLY: which block(s) actually get steered, instead of always block 25 — every block 0-49 is captured read-only regardless. A single block, a range, or a comma list: 25 | 31-40 | 4,5,6 — same syntax as Repeat block(s). Each named block steers with its OWN learned direction from its OWN liked/disliked history at that block; a block without enough data just captures, same as any unnamed candidate block.",
       detail: "Blocks 40-49 carry the most residual movement (48+49 alone about a third of it), so a steering push has the most to act on there. Blocks 0-1 are structurally distinct rather than just 'quiet early blocks' — 0 is the raw input projection, 1 does the first cross-patch mixing — expect steering there to behave differently in kind, not just 'weaker', than the middle of the stack." },
+    { name: "h3_av_decouple", label: "Video/audio attention decoupling (H3, experimental)", kind: "float", default: 0.0, min: 0.0, max: 1.0, step: 0.01,
+      hint: "EXPERIMENTAL, unvalidated, H3 only. Weakens how much video and audio attend to each other, for the failure pattern where one comes out clean and the other doesn't. 0 = off. Not rating-driven — a manual dial, same as Repeat block(s).",
+      detail: "H3 has no separate cross-attention to mask — video, text and audio share one joint self-attention. Instead of building a mask over that whole sequence (unaffordable at H3's token count), this runs the same attention as separate video-query and audio-query passes against a version of the keys where the other modality is scaled down by this amount. Same total cost as the ordinary single pass. 1.0 cuts cross-modal attention entirely, which likely also kills legitimate audio/video sync — search down from a small value, not up to it." },
     { name: "identity_transfer_enabled", label: "Best-FaceID compatibility", kind: "bool", default: false,
       hint: "Feeds the identity pin image the way Best-FaceID identity LoRAs expect it. Needs an Identity pin set.",
       detail: "Replaces Continuity's identity-pin guide with separate, non-rendered reference tokens plus an optional ArcFace projector below. Load the LoRA itself in Models. No effect without an identity pin." },
@@ -520,7 +523,7 @@
   const CHAIN_VIEW_KNOBS = {
     chain_continuity: ["carry_i2v_guides", "carry_overlap_through_anchor"],
     chain_timing: ["frame_overlap", "transition_duration", "use_same_seed", "cut_opening_frames"],
-    chain_guidance: ["cfg", "embed_guidance", "embed_guidance_source", "embed_guidance_strength", "score_slider", "score_slider_strength", "taste_nearest_prompt", "output_guidance", "output_guidance_strength", "trajectory_guidance", "trajectory_guidance_strength", "dynashift", "dynashift_strength", "dynashift_threshold", "h3_repr_steering", "h3_repr_steering_strength", "h3_repr_steering_block"],
+    chain_guidance: ["cfg", "embed_guidance", "embed_guidance_source", "embed_guidance_strength", "score_slider", "score_slider_strength", "taste_nearest_prompt", "output_guidance", "output_guidance_strength", "trajectory_guidance", "trajectory_guidance_strength", "dynashift", "dynashift_strength", "dynashift_threshold", "h3_repr_steering", "h3_repr_steering_strength", "h3_repr_steering_block", "h3_av_decouple"],
     chain_decode: ["decode_noise_scale", "decode_timestep", "decode_tile_size"],
     chain_experimental: ["context_windows", "context_window_length", "context_window_overlap", "context_window_schedule", "context_window_fuse", "context_window_freenoise", "context_window_retain_first", "joyai_memory", "joyai_memory_size", "joyai_fix_frames", "joyai_frame_select", "joyai_memory_strength", "joyai_audio_memory", "v2a_grad_scale", "alg_blur_guides", "alg_guide_blur_strength", "alg_guide_blur_sigma_threshold", "bounded_attention_enabled", "h3_block_repeat", "h3_block_repeat_span_loop", "h3_block_repeat_video_only", "h3_block_repeat_times", "h3_block_repeat_last_steps", "identity_transfer_enabled", "source_id", "phase_scale", "id_strength", "arcface_mode", "debug_log"],
   };

samplers.py

@@ -2806,6 +2806,10 @@ class FunPackLTXAVSceneChainSampler:
                     "default": 0.05, "min": 0.0, "max": 2.0, "step": 0.01,
                     "tooltip": "How much of the learned direction to add, as a fraction of that block's own hidden-state norm. Unlike the attention-bias channel, there is no architectural ceiling here -- push too far and coherence breaks with no warning. Start low; the way to tell it is doing anything at all is to push this deliberately high on a same-seed test and see if the video visibly moves, then back off.",
                 }),
+                "h3_av_decouple": ("FLOAT", {
+                    "default": 0.0, "min": 0.0, "max": 1.0, "step": 0.01,
+                    "tooltip": "EXPERIMENTAL, unvalidated (H3 only). Damps how much video and audio attend to EACH OTHER inside H3's joint self-attention, without touching any hidden state -- for the failure pattern where one comes out clean and the other doesn't. 0 = off (installs nothing, zero cost). Higher = weaker cross-modal influence; 1.0 fully cuts it, which likely also kills legitimate AV sync, so this is a knob to search down from, not up to.",
+                }),
                 "alg_guide_blur_strength": ("FLOAT", {
                     "default": 2.0, "min": 1.0, "max": 4.0, "step": 0.1,
                     "tooltip": "Downsample factor for the guide-frame blur (alg_blur_guides). Higher = blurrier guide/JoyAI frames during the affected steps. Independent of the sampler's anchor alg_strength.",
@@ -3703,6 +3707,7 @@ class FunPackLTXAVSceneChainSampler:
                       bounded_attention_enabled=False,
                       h3_repr_steering=False, h3_repr_steering_strength=0.05,
                       h3_repr_steering_block="25", h3_repr_capture_slot="",
+                      h3_av_decouple=0.0,
                       h3_block_repeat="", h3_block_repeat_times=1,
                       h3_block_repeat_video_only=False, h3_block_repeat_span_loop=False,
                       h3_block_repeat_last_steps=0,
@@ -3795,6 +3800,7 @@ class FunPackLTXAVSceneChainSampler:
             model = self._install_h3_repr_steering(
                 model, refinement_key, h3_repr_steering_strength, _repr_capture,
                 steer_block=h3_repr_steering_block)
+        model = self._install_h3_av_decouple(model, h3_av_decouple)
 
         try:
             sampled = comfy.sample.sample_custom(
@@ -6817,6 +6823,102 @@ class FunPackLTXAVSceneChainSampler:
                         "not applied and not capturing this run")
             return model
 
+    def _install_h3_av_decouple(self, model, strength):
+        """EXPERIMENTAL, unvalidated (H3 only). Damps cross-modal attention instead of
+        touching any hidden state: H3 packs video/text/audio into ONE joint self-attention
+        (no separate cross-attn to mask, unlike LTXAV -- see project_minimax_h3 memory,
+        which is why Bounded Attention was ruled out there as GB-scale on a ~37k-token
+        sequence). This does not build that S×S mask at all. It uses comfy's own
+        `optimized_attention_override` extension point (comfy/ldm/modules/attention.py's
+        `wrap_attn` -- already wired into every backend, nothing to monkeypatch) to run the
+        SAME joint attention as separate video-query and audio-query passes against a K
+        where the OTHER modality's rows are scaled down by `strength`; text/ref/cond rows
+        stay on one ordinary pass, untouched. Splitting the query dimension like this costs
+        nothing extra -- same total Q*K*V work as the single joint call, just partitioned --
+        and the only new memory is two full-length clones of K (audio and video zeroed out
+        by contrast, at most `strength`), not a dense attention matrix.
+
+        `strength` in [0, 1]; 0 or blank installs nothing (exact pass-through, not just a
+        no-op override) so an unset knob costs literally zero.
+        """
+        try:
+            _strength = float(strength or 0.0)
+        except (TypeError, ValueError):
+            _strength = 0.0
+        if _strength <= 0.0:
+            return model
+        _strength = min(_strength, 1.0)
+        try:
+            try:
+                from . import h3_repr_steering as _rs
+            except ImportError:
+                import h3_repr_steering as _rs
+
+            patched = model.clone()
+            to = patched.model_options.get("transformer_options", {}).copy()
+            patches_replace = dict(to.get("patches_replace", {}))
+            dit_patches = dict(patches_replace.get("dit", {}))
+
+            # mod_segments is only ever handed to blocks directly (never threaded onto
+            # transformer_options by core), so the only way to see it here is to catch it
+            # off a block hook -- block 0 always runs first and mod_segments doesn't change
+            # block to block, so one hook is enough. Chains through whatever's already
+            # hooked there (h3_repr_steering hooks all 50 blocks) so this never disables it.
+            _seg_holder = {"mod_segments": None}
+            _inner0 = dit_patches.get(("double_block", 0))
+
+            def _capture_hook(args, extra):
+                _seg_holder["mod_segments"] = args.get("mod_segments")
+                return extra["original_block"](args) if _inner0 is None else _inner0(args, extra)
+
+            dit_patches[("double_block", 0)] = _capture_hook
+            patches_replace["dit"] = dit_patches
+            to["patches_replace"] = patches_replace
+
+            _mask_cache = {}  # seq_len -> (video_mask, audio_mask); one live resolution per gen
+
+            def _override(func, q, k, v, heads, mask=None, skip_reshape=True,
+                          transformer_options=None, **kwargs):
+                seg = _seg_holder["mod_segments"]
+                if seg is None or mask is not None or not skip_reshape:
+                    return func(q, k, v, heads, mask=mask, skip_reshape=skip_reshape, **kwargs)
+                seq_len = q.shape[-2]
+                cached = _mask_cache.get(seq_len)
+                if cached is None:
+                    cached = (_rs.video_mask_from_mod_segments(seg, seq_len, q.device),
+                              _rs.audio_mask_from_mod_segments(seg, seq_len, q.device))
+                    _mask_cache[seq_len] = cached
+                vmask, amask = cached
+                if vmask is None or amask is None or k.shape[-2] != seq_len:
+                    return func(q, k, v, heads, mask=mask, skip_reshape=skip_reshape, **kwargs)
+                other = ~(vmask | amask)
+                out = torch.empty(q.shape[0], seq_len, heads * q.shape[-1],
+                                  dtype=q.dtype, device=q.device)
+                if amask.any():
+                    k_damped = k.clone()
+                    k_damped[:, :, vmask, :] *= (1.0 - _strength)
+                    out[:, amask, :] = func(q[:, :, amask, :], k_damped, v, heads,
+                                            mask=None, skip_reshape=skip_reshape, **kwargs)
+                if vmask.any():
+                    k_damped = k.clone()
+                    k_damped[:, :, amask, :] *= (1.0 - _strength)
+                    out[:, vmask, :] = func(q[:, :, vmask, :], k_damped, v, heads,
+                                            mask=None, skip_reshape=skip_reshape, **kwargs)
+                if other.any():
+                    out[:, other, :] = func(q[:, :, other, :], k, v, heads,
+                                            mask=None, skip_reshape=skip_reshape, **kwargs)
+                return out
+
+            to["optimized_attention_override"] = _override
+            patched.model_options["transformer_options"] = to
+            print(f"[FunPackSceneChain] H3 attention decoupling: damping cross-modal "
+                  f"video<->audio attention by {_strength:g}.")
+            return patched
+        except Exception as _e:  # noqa: BLE001
+            _log.failed("FunPackSceneChain", "H3 attention decoupling", _e,
+                        "not applied this run")
+            return model
+
     @staticmethod
     def _parse_block_spec(spec, n_blocks=50):
         """"40" | "38-42" | "10,40,44" | "" -> a set of block indices. Anything unparseable
@@ -7462,6 +7564,7 @@ class FunPackLTXAVSceneChainSampler:
                dynashift=False, dynashift_strength=0.3, dynashift_threshold=0.6,
                h3_repr_steering=False, h3_repr_steering_strength=0.05,
                h3_repr_steering_block="25", h3_repr_capture_slot="",
+               h3_av_decouple=0.0,
                h3_block_repeat="", h3_block_repeat_times=1,
                h3_block_repeat_video_only=False, h3_block_repeat_span_loop=False,
                h3_block_repeat_last_steps=0,
@@ -8360,6 +8463,7 @@ class FunPackLTXAVSceneChainSampler:
                     h3_repr_steering_strength=h3_repr_steering_strength,
                     h3_repr_steering_block=h3_repr_steering_block,
                     h3_repr_capture_slot=h3_repr_capture_slot,
+                    h3_av_decouple=h3_av_decouple,
                     h3_block_repeat=h3_block_repeat,
                     h3_block_repeat_times=h3_block_repeat_times,
                     h3_block_repeat_video_only=h3_block_repeat_video_only,

tests/test_h3_av_decouple.py

@@ -0,0 +1,155 @@
+"""H3 attention decoupling (_install_h3_av_decouple): damps cross-modal video<->audio
+attention via comfy's optimized_attention_override extension point, without ever
+materializing an S x S mask. Exercised at the install/override level with fake tensors --
+no real attention backend or GPU needed."""
+import sys
+from pathlib import Path
+
+import torch
+
+sys.path.insert(0, str(Path(__file__).resolve().parents[1]))
+sys.path.insert(0, ".")
+import _comfy_stubs  # noqa: E402, F401
+from samplers import FunPackLTXAVSceneChainSampler as S  # noqa: E402
+
+
+class _FakeModel:
+    def __init__(self):
+        self.model_options = {}
+
+    def clone(self):
+        m = _FakeModel()
+        m.model_options = dict(self.model_options)
+        return m
+
+
+# rows: 0-1 video, 2 text, 3-4 audio, 5 video -- exercises a modality repeating
+# non-contiguously, same as a real packed layout can produce.
+_MOD_SEGMENTS = [(0, 2, 0), (2, 3, 1), (3, 5, 2), (5, 6, 0)]
+_VIDEO_IDX = [0, 1, 5]
+_AUDIO_IDX = [3, 4]
+_OTHER_IDX = [2]
+
+
+def _fire_capture_hook(patched, seq_len=6, mod_segments=_MOD_SEGMENTS):
+    """Runs block 0's hook once, the way the real DiT forward loop would, so the override
+    has a mod_segments layout to work with."""
+    dit = patched.model_options["transformer_options"]["patches_replace"]["dit"]
+    hook = dit[("double_block", 0)]
+    img = torch.ones(seq_len, 3)
+    args = {"img": img, "mod_segments": mod_segments}
+    extra = {"original_block": lambda a: {"img": a["img"]}}
+    hook(args, extra)
+
+
+def _qkv(seq_len=6, heads=1, head_dim=2):
+    q = torch.arange(seq_len * head_dim, dtype=torch.float32).view(1, heads, seq_len, head_dim).clone()
+    k = torch.ones(1, heads, seq_len, head_dim) * 2.0
+    v = torch.zeros(1, heads, seq_len, head_dim)
+    return q, k, v
+
+
+def _fake_func_factory(calls):
+    def fake_func(q, k, v, heads, mask=None, skip_reshape=True, **kw):
+        calls.append({"q_len": q.shape[-2], "k_sum": float(k.sum()), "mask": mask})
+        return torch.full((q.shape[0], q.shape[-2], heads * q.shape[-1]), float(k.sum()))
+    return fake_func
+
+
+def test_zero_strength_is_a_true_noop():
+    model = _FakeModel()
+    node = S()
+    assert node._install_h3_av_decouple(model, 0.0) is model
+    assert node._install_h3_av_decouple(model, "") is model
+    assert node._install_h3_av_decouple(model, "not-a-number") is model
+
+
+def test_enabled_installs_override_and_block0_hook():
+    node = S()
+    patched = node._install_h3_av_decouple(_FakeModel(), 0.5)
+    to = patched.model_options["transformer_options"]
+    assert callable(to["optimized_attention_override"])
+    assert ("double_block", 0) in to["patches_replace"]["dit"]
+
+
+def test_before_mod_segments_captured_override_is_a_passthrough():
+    node = S()
+    patched = node._install_h3_av_decouple(_FakeModel(), 0.5)
+    override = patched.model_options["transformer_options"]["optimized_attention_override"]
+    calls = []
+    q, k, v = _qkv()
+    out = override(_fake_func_factory(calls), q, k, v, heads=1, mask=None, skip_reshape=True)
+    assert len(calls) == 1, "no mod_segments yet -> single unmodified call, not split"
+    assert calls[0]["k_sum"] == float(k.sum())
+    assert torch.allclose(out, torch.full_like(out, float(k.sum())))
+
+
+def test_an_explicit_mask_is_never_overridden():
+    node = S()
+    patched = node._install_h3_av_decouple(_FakeModel(), 0.5)
+    _fire_capture_hook(patched)
+    override = patched.model_options["transformer_options"]["optimized_attention_override"]
+    calls = []
+    q, k, v = _qkv()
+    real_mask = torch.zeros(1, 1, 6, 6, dtype=torch.bool)
+    override(_fake_func_factory(calls), q, k, v, heads=1, mask=real_mask, skip_reshape=True)
+    assert len(calls) == 1 and calls[0]["mask"] is real_mask
+
+
+def test_splits_into_three_groups_with_correctly_damped_keys():
+    node = S()
+    patched = node._install_h3_av_decouple(_FakeModel(), 0.5)
+    _fire_capture_hook(patched)
+    override = patched.model_options["transformer_options"]["optimized_attention_override"]
+    calls = []
+    q, k, v = _qkv(seq_len=6, heads=1, head_dim=2)
+    out = override(_fake_func_factory(calls), q, k, v, heads=1, mask=None, skip_reshape=True)
+
+    assert len(calls) == 3, "video-query, audio-query and other-query passes"
+    by_len = {c["q_len"]: c for c in calls}
+    assert set(by_len) == {len(_VIDEO_IDX), len(_AUDIO_IDX), len(_OTHER_IDX)}
+
+    full_k_sum = float(k.sum())  # 6 rows * 2 dims * 2.0 = 24
+    # audio-query pass: video rows (3 of them) halved in K -> k_sum drops by half of their share
+    per_row = 2 * 2.0  # head_dim * key value
+    audio_call = by_len[len(_AUDIO_IDX)]
+    assert audio_call["k_sum"] == full_k_sum - len(_VIDEO_IDX) * per_row * 0.5
+    video_call = by_len[len(_VIDEO_IDX)]
+    assert video_call["k_sum"] == full_k_sum - len(_AUDIO_IDX) * per_row * 0.5
+    other_call = by_len[len(_OTHER_IDX)]
+    assert other_call["k_sum"] == full_k_sum, "text/other rows ride the ordinary, undamped K"
+
+    # scatter landed each group's output back at its own rows
+    assert torch.allclose(out[0, _AUDIO_IDX], torch.full((len(_AUDIO_IDX), 2), audio_call["k_sum"]))
+    assert torch.allclose(out[0, _VIDEO_IDX], torch.full((len(_VIDEO_IDX), 2), video_call["k_sum"]))
+    assert torch.allclose(out[0, _OTHER_IDX], torch.full((len(_OTHER_IDX), 2), other_call["k_sum"]))
+
+
+def test_chains_through_an_already_installed_block0_hook():
+    """If h3_repr_steering (or anything else) already hooked block 0, av_decouple must not
+    silently replace it -- it should still see mod_segments AND still get the other hook's
+    modified output."""
+    node = S()
+    model = _FakeModel()
+    to = {"patches_replace": {"dit": {
+        ("double_block", 0): lambda args, extra: {"img": args["img"] * 0.0}
+    }}}
+    model.model_options["transformer_options"] = to
+    patched = node._install_h3_av_decouple(model, 0.5)
+    dit = patched.model_options["transformer_options"]["patches_replace"]["dit"]
+    img = torch.ones(6, 3)
+    out = dit[("double_block", 0)]({"img": img, "mod_segments": _MOD_SEGMENTS},
+                                   {"original_block": lambda a: {"img": a["img"]}})
+    assert torch.allclose(out["img"], torch.zeros(6, 3)), "existing hook's effect must survive"
+
+    override = patched.model_options["transformer_options"]["optimized_attention_override"]
+    calls = []
+    q, k, v = _qkv()
+    override(_fake_func_factory(calls), q, k, v, heads=1, mask=None, skip_reshape=True)
+    assert len(calls) == 3, "mod_segments was captured despite chaining through the other hook"
+
+
+if __name__ == "__main__":
+    test_zero_strength_is_a_true_noop()
+    test_enabled_installs_override_and_block0_hook()
+    print("ok (run via pytest for the rest)")