nodes/olivv-cs-ComfyUI-FunPack
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

feat(h3): the anchor is its own cache chunk, so it can be released

Browse Source 
Every chunk restarting from the input image on the rental. The prompt and the
anchor shared one cache chunk, so pinning the prompt pinned the anchor with it —
for the whole clip, in a context of three visible blocks. In a dense render the
anchor is one pin among hundreds of frames' worth of rows; there it was a third
of everything the model could see, and the model did what it was trained to do
with a conditioning row: compose from it.

They are now two chunks. The prompt is 0 and always pinned — it is every moment
of the clip. The conditioning is 1, so `sink` decides how long it stays: 1
releases the anchor after the opening, 2 keeps it (today's behaviour), 3 also
pins the opening shot. Media chunks moved to i + 2, and the conditioning slot is
reserved even on a t2v run so `sink` counts the same either way.

The prefill is two forwards, and the conditioning pass reads the text already in
the cache — so the anchor is still represented WITH the prompt rather than in
isolation.

Note the widget's meaning shifted by one: sink 2 was prompt+anchor+opening shot
and is now prompt+anchor.

Tests: 1806 python, 68 JS. The upstream check now prefills a real keyframe and
confirms sink 1 releases it while keeping the prompt.
This commit is contained in:
digital-garbage
2026-08-26 22:49:42 +03:00
parent 8e808c62e5
commit 985cf8048d
6 changed files with 142 additions and 46 deletions
Download Patch File Download Diff File Expand all files Collapse all files
h3_causal.py
+78 -37
View File
@@ -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,
movie_editor/frontend/engine_settings.js
+2 -1
View File
@@ -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." },
samplers.py
+1 -1
View File
@@ -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,
tests/_h3_causal_upstream_check.py
+28
View File
@@ -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.
tests/test_h3_causal_plan.py
+29 -4
View File
@@ -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 ────────────────────────────────
tests/test_h3_causal_session.py
+4 -3
View File
@@ -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():