diff --git a/h3_causal.py b/h3_causal.py index 820112e7..7fbbbed1 100644 --- a/h3_causal.py +++ b/h3_causal.py @@ -172,18 +172,28 @@ class ChunkPlan: would be a second source of truth for the clip's shape — and the first symptom of it drifting is a clip that looks subtly wrong rather than an exception. - Two kinds of selection: + Three kinds of selection, and they are three because THE PROMPT AND THE ANCHOR ARE NOT + THE SAME KIND OF CONTEXT: - `prefix` is everything before the target streams — the text, the keyframe cond rows and - the reference blocks. It is cached ONCE, as cache chunk 0, and every media chunk reads it. - That is why i2v and r2v survive here: the conditioning rows keep the layout upstream built - for them instead of being re-packed into a chunk that has no place for them. + `text` is the prompt span. Cache chunk 0, and pinned for the whole clip by any sink at + all — every moment of the video is the same prompt. + + `cond` is the keyframe pins and the reference blocks. Cache chunk 1, so it is pinned by + `sink >= 2` and EVICTED by `sink == 1`. In a dense forward the anchor is one pin among + hundreds of frames' worth of rows; pinned into a 3-block context it is a third of + everything the model can see, and the model does what it was trained to do with + conditioning rows — compose from them. Releasing it lets the shot progress. Keeping the + conditioning rows in the layout upstream built for them is still what makes i2v and r2v + work here at all; this only decides how long they stay visible. `chunk(i)` is media chunk i: its audio rows followed by its video rows, in the packing - order the target segments already use. Media chunk i is CACHE chunk i + 1, because the - prefix took index 0. + order the target segments already use. Media chunk i is CACHE chunk i + 2, because the + text and the conditioning took 0 and 1. """ + TEXT_CACHE_INDEX = 0 + COND_CACHE_INDEX = 1 + def __init__(self, layout, bounds, audio_t): self.layout = layout self.bounds = list(bounds) @@ -211,8 +221,19 @@ class ChunkPlan: f"the soundtrack does not cover the picture — {self.audio_t} audio latents " f"against {latent_t} video latents leaves chunk {empty[0]} silent. H3 runs " f"audio at 40/s against a fixed 24 fps; this clip is not on that clock.") + prefix = [(a, b, kind) for a, b, kind in segments if b <= audio_seg[0]] self.prefix_rows = torch.arange(0, audio_seg[0], dtype=torch.long) - self.prefix_runs = [(a, b, kind) for a, b, kind in segments if b <= audio_seg[0]] + self.prefix_runs = prefix + text = [r for r in prefix if r[2] == "text"] + cond = [r for r in prefix if r[2] != "text"] + self.text_runs = text + self.text_rows = torch.arange(text[0][0], text[-1][1], dtype=torch.long) if text \ + else torch.zeros(0, dtype=torch.long) + # Local indices: the conditioning is its own sequence now, not a tail of the text. + base = cond[0][0] if cond else 0 + self.cond_runs = [(a - base, b - base, kind) for a, b, kind in cond] + self.cond_rows = torch.arange(base, cond[-1][1], dtype=torch.long) if cond \ + else torch.zeros(0, dtype=torch.long) @property def n_chunks(self): @@ -220,8 +241,13 @@ class ChunkPlan: @staticmethod def cache_index(index): - """Media chunk i lives at cache index i + 1; the prefix holds 0.""" - return int(index) + 1 + """Media chunk i lives at cache index i + 2; the text holds 0 and the conditioning 1. + + The conditioning slot is reserved whether or not this clip has any — an index that + moved with the presence of an anchor would make `sink` mean a different thing on a + t2v run than on an i2v one. + """ + return int(index) + 2 def chunk(self, index): """(row indices, local runs) for one media chunk.""" @@ -569,35 +595,16 @@ def _causal_classes(): return self.token_refiner(self.condition_proj(states), transformer_options=transformer_options) - def prefill_text(self, context, plan, cache, *, video_sigma=0.0, audio_sigma=0.0, - transformer_options={}, minimax_payload=None): - """Run the prompt and the conditioning rows once, and cache them as chunk 0. + def _prefill_span(self, plan, rows, runs, cache, chunk_index, t_v, t_a, dtype, + device, payload, transformer_options, text_states=None): + """Run one prefix span and write its K/V. Returns nothing — only the cache matters. - The prefix is the text, the keyframe cond rows and the reference blocks — the part - of the sequence that does not change from chunk to chunk. Caching it once is not - only cheaper: every later chunk assumes the text keys it attends are the ones the - model saw, and recomputing them per chunk would quietly make that false. - - Cached at sigma 0 by default, so the prefix is CLEAN context. The cond rows pin - near 1 in upstream's own table regardless, and text is never noised, so this is the - state those rows would be in at the end of any schedule. + `chunk_index` is both where this span is written AND what it is allowed to read, + so the conditioning written at 1 attends the text already at 0. """ - payload = minimax_payload or {} - rows = plan.prefix_rows - if rows.numel() == 0: - raise CacheError("this layout has no prefix rows to prefill") - device = comfy.model_management.get_torch_device() - if hasattr(context, "device"): - device = context.device - dtype = context.dtype - t_v = float(1.0 - max(float(video_sigma), 1e-6)) - t_a = float(1.0 - max(float(audio_sigma), 1e-6)) - runs = plan.prefix_runs seg_t, t_row, unique_t = self._modulation({k for _, _, k in runs}, t_v, t_a, payload) - text_states = self._text_states(context, transformer_options) cond_video = self._cond_video_rows(payload, device) cond_audio = self._cond_audio_rows(payload, device) - h = torch.empty(int(rows.shape[0]), self.hidden_size, dtype=dtype, device=device) voff = aoff = 0 for a, b, kind in runs: @@ -616,14 +623,48 @@ def _causal_classes(): "carries no cond_audio_latents to fill them") h[a:b] = self.audio_patch_proj(cond_audio[aoff:aoff + n]).to(dtype) aoff += n - rope_freqs = rope_rotation_table( self.rope_freqs(plan.layout.position_ids[rows], device), dtype) self._run_blocks(h, self._t_emb(unique_t, device, dtype), self._mod_segments(runs, seg_t, t_row, payload.get("text_token_tags")), - rope_freqs, transformer_options, cache, 0, True) - cache.finish_chunk(0) + rope_freqs, transformer_options, cache, chunk_index, True) + cache.finish_chunk(chunk_index) + + def prefill_text(self, context, plan, cache, *, video_sigma=0.0, audio_sigma=0.0, + transformer_options={}, minimax_payload=None): + """Cache the prompt, then the conditioning, as two separate chunks. + + Two forwards rather than one because the prompt and the anchor are not the same + kind of context and must be evictable separately — see ChunkPlan. The conditioning + pass reads the text already in the cache, so the anchor is still represented WITH + the prompt rather than in isolation. + + Cached at sigma 0 by default, so the prefix is CLEAN context. The cond rows pin + near 1 in upstream's own table regardless, and text is never noised, so this is + the state those rows would be in at the end of any schedule. + """ + payload = minimax_payload or {} + if plan.text_rows.numel() == 0: + raise CacheError("this layout has no prompt rows to prefill") + device = context.device if hasattr(context, "device") \ + else comfy.model_management.get_torch_device() + dtype = context.dtype + t_v = float(1.0 - max(float(video_sigma), 1e-6)) + t_a = float(1.0 - max(float(audio_sigma), 1e-6)) + text_states = self._text_states(context, transformer_options) + + self._prefill_span(plan, plan.text_rows, plan.text_runs, cache, + plan.TEXT_CACHE_INDEX, t_v, t_a, dtype, device, payload, + transformer_options, text_states=text_states) + if plan.cond_rows.numel(): + self._prefill_span(plan, plan.cond_rows, plan.cond_runs, cache, + plan.COND_CACHE_INDEX, t_v, t_a, dtype, device, payload, + transformer_options) + else: + # Reserve the slot even with nothing in it, so `sink` counts the same on a + # t2v run as on an i2v one. + cache.finish_chunk(plan.COND_CACHE_INDEX) def forward_chunk(self, video_latent, audio_latent, plan, index, cache, *, video_sigma, audio_sigma, context=None, update_cache=False, diff --git a/movie_editor/frontend/engine_settings.js b/movie_editor/frontend/engine_settings.js index bf6426b0..629c016e 100644 --- a/movie_editor/frontend/engine_settings.js +++ b/movie_editor/frontend/engine_settings.js @@ -339,7 +339,8 @@ hint: "How each chunk reaches its next noise level. 'consistency' is what the RAVEN LoRA was distilled for. 'euler' is the ordinary step used everywhere else, so a chunked run can be compared against a normal one without this being a second thing that changed." }, { name: "h3_causal_sink", label: "Chunks pinned from the start", kind: "int", default: 2, min: 1, max: 64, dependsOn: "h3_causal_chunks", dependsValue: true, - hint: "Chunks that stay visible to every later chunk and are never forgotten. Chunk 0 is the prompt, so 2 means the prompt plus the opening shot — which is what holds a character together across a long clip. Higher costs memory." }, + hint: "Chunks that stay visible to every later chunk and are never forgotten, counted from the front: 1 is the prompt, 2 adds the anchor image and any reference, 3 adds the opening shot. Higher costs memory.", + detail: "The prompt is every moment of the clip, so it is always pinned. The anchor is not: in a dense render it is one pin among hundreds of frames, but pinned into a three-block context it is a third of everything the model can see — and the model then does what it was trained to do with a conditioning row, which is compose from it. If every chunk restarts from your input image instead of continuing the shot, drop this to 1." }, { name: "h3_causal_window", label: "Recent chunks remembered", kind: "int", default: 2, min: 0, max: 64, dependsOn: "h3_causal_chunks", dependsValue: true, hint: "Short-term memory: how many of the most recent chunks stay visible, carrying motion and continuity from one moment to the next. 0 makes the clip lean entirely on its opening." }, diff --git a/samplers.py b/samplers.py index 6995853f..1b7b8bc4 100644 --- a/samplers.py +++ b/samplers.py @@ -3187,7 +3187,7 @@ class FunPackLTXAVSceneChainSampler: }), "h3_causal_sink": ("INT", { "default": 2, "min": 1, "max": 64, - "tooltip": "EXPERIMENTAL, MiniMax H3 only. How many chunks are pinned in the cache from the very start and never evicted. Chunk 0 is the prompt, so 2 means the prompt plus the opening shot stay visible to every later chunk — which is what keeps a character consistent across a long clip. Higher costs more memory and more attention per chunk.", + "tooltip": "EXPERIMENTAL, MiniMax H3 only. How many chunks are pinned in the cache from the very start and never evicted, counted from the front: 1 is the prompt, 2 adds the anchor image and any reference, 3 adds the opening shot. The prompt is every moment of the clip so it is always pinned. The ANCHOR is not — in a dense render it is one pin among hundreds of frames, but pinned into a three-block context it is a third of everything the model can see, and the model then does what it was trained to do with a conditioning row: compose from it. If every chunk restarts from your input image instead of continuing, drop this to 1. Higher costs more memory and more attention per chunk.", }), "h3_causal_window": ("INT", { "default": 2, "min": 0, "max": 64, diff --git a/tests/_h3_causal_upstream_check.py b/tests/_h3_causal_upstream_check.py index b3cc224a..533db06e 100644 --- a/tests/_h3_causal_upstream_check.py +++ b/tests/_h3_causal_upstream_check.py @@ -185,6 +185,34 @@ def _forward_phase(): rows, _ = plan.chunk(i) assert cache._store[(0, plan.cache_index(i))][0].shape[-2] == rows.shape[0] + # An i2v clip: the anchor is prefilled as its OWN cache chunk, reading the text already + # cached at 0, so it can be evicted by `sink` without taking the prompt with it. + pinned_layout = PackedLayout(text_len, latent_t, lat_h, lat_w, audio_t, + keyframes=[{"resolved_frame_index": 0}], frame_count=39) + pinned_plan = hc.build_plan(pinned_layout, latent_t, audio_t) + assert pinned_plan.cond_rows.numel(), "the keyframe produced no conditioning rows" + pinned_payload = {"layout": pinned_layout, "seed": 0, + "cond_video_latents": [torch.randn(1, 24, 1, lat_h, lat_w)]} + pinned_cache = hc.ChunkKVCache(len(causal.blocks), sink=2, window=2, + device=torch.device("cpu"), offload=False) + with torch.no_grad(): + causal.prefill_text(context, pinned_plan, pinned_cache, + minimax_payload=pinned_payload) + assert pinned_cache._store[(0, 0)][0].shape[-2] == text_len + assert (pinned_cache._store[(0, 1)][0].shape[-2] + == pinned_plan.cond_rows.shape[0]), "the anchor did not cache on its own" + v_start, v_stop, a_start, a_stop = pinned_plan.bounds[0] + out_v, out_a = causal.forward_chunk( + video[:, :, v_start:v_stop], audio[..., a_start:a_stop], pinned_plan, 0, + pinned_cache, video_sigma=0.8, audio_sigma=0.4, context=context, + minimax_payload=pinned_payload) + assert torch.isfinite(out_v).all() and torch.isfinite(out_a).all() + # sink 1 releases the anchor and keeps the prompt; sink 2 keeps both + pinned_cache.sink = 1 + assert pinned_cache.retained_indices(6) == [0], pinned_cache.retained_indices(6) + pinned_cache.sink = 2 + assert pinned_cache.retained_indices(6) == [0, 1], pinned_cache.retained_indices(6) + # A ComfyUI without the newer `time_shift_slope` must still build the lane. Importing it # like the rest made one missing symbol refuse the whole feature on 0.34.0, and the # refusal then blamed a setting — this is that regression. diff --git a/tests/test_h3_causal_plan.py b/tests/test_h3_causal_plan.py index b1b3e0c1..c5da36f0 100644 --- a/tests/test_h3_causal_plan.py +++ b/tests/test_h3_causal_plan.py @@ -98,10 +98,35 @@ def test_the_tail_chunk_owns_whatever_audio_is_left(): assert plan.bounds[-1][3] == AUDIO_T -def test_media_chunk_zero_is_cache_chunk_one(): - """The prefix holds cache index 0, so `sink=2` pins the prompt plus the opening shot.""" - assert hc.ChunkPlan.cache_index(0) == 1 - assert hc.ChunkPlan.cache_index(5) == 6 +def test_media_chunk_zero_is_cache_chunk_two(): + """The prompt holds 0 and the conditioning 1, so the media chunks start at 2.""" + assert hc.ChunkPlan.cache_index(0) == 2 + assert hc.ChunkPlan.cache_index(5) == 7 + + +def test_the_prompt_and_the_conditioning_are_separate_cache_chunks(): + """They are not the same kind of context. The prompt is every moment of the clip; the + anchor is one moment of it, and pinning it into a 3-block context makes the model compose + from it instead of continuing.""" + plan = _plan((("text", TEXT_LEN), ("cond", 6), ("ref_img", 12))) + assert [k for _, _, k in plan.text_runs] == ["text"] + assert [k for _, _, k in plan.cond_runs] == ["cond", "ref_img"] + assert plan.text_rows.tolist() + plan.cond_rows.tolist() == plan.prefix_rows.tolist() + + +def test_the_conditioning_runs_are_local_to_their_own_sequence(): + """They are prefilled as a sequence of their own, so a run still starting at the text's + row count would write past the end of the buffer.""" + plan = _plan((("text", TEXT_LEN), ("cond", 6))) + assert plan.cond_runs == [(0, 6, "cond")] + + +def test_a_clip_with_no_conditioning_still_reserves_the_slot(): + """An index that moved with the presence of an anchor would make `sink` mean a different + thing on a t2v run than on an i2v one.""" + plan = _plan() + assert plan.cond_rows.numel() == 0 + assert hc.ChunkPlan.cache_index(0) == 2 # ── conditioning rows keep upstream's layout ──────────────────────────────── diff --git a/tests/test_h3_causal_session.py b/tests/test_h3_causal_session.py index 87327d50..d5848e7d 100644 --- a/tests/test_h3_causal_session.py +++ b/tests/test_h3_causal_session.py @@ -231,10 +231,11 @@ def test_the_commit_of_one_chunk_precedes_the_first_forward_of_the_next(): def test_each_chunk_writes_to_its_own_cache_index_above_the_prefix(): - diffusion, session, _, _ = _run() + """The prompt holds 0 and the conditioning 1, so nothing a media chunk writes may land + on either — overwriting the prompt's K/V would silently unprompt the rest of the clip.""" + _, session, _, _ = _run() written = {key[1] for key in session["cache"]._store} - assert 0 not in written or session["cache"].committed_chunks - assert written <= set(range(1, session["plan"].n_chunks + 1)) + assert written <= set(range(2, session["plan"].n_chunks + 2)) def test_the_rollout_returns_a_clip_of_the_shape_it_was_given():