diff --git a/h3_repr_steering.py b/h3_repr_steering.py index 80dc960c..e4a7f393 100644 --- a/h3_repr_steering.py +++ b/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): diff --git a/movie_editor/backend/settings_card.py b/movie_editor/backend/settings_card.py index 8e191e89..e089da98 100644 --- a/movie_editor/backend/settings_card.py +++ b/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, diff --git a/movie_editor/frontend/engine_settings.js b/movie_editor/frontend/engine_settings.js index bfdeeac8..8f690674 100644 --- a/movie_editor/frontend/engine_settings.js +++ b/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"], }; diff --git a/samplers.py b/samplers.py index cc8f5876..caadca26 100644 --- a/samplers.py +++ b/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, diff --git a/tests/test_h3_av_decouple.py b/tests/test_h3_av_decouple.py new file mode 100644 index 00000000..519baaa8 --- /dev/null +++ b/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)")