diff --git a/h3_causal.py b/h3_causal.py new file mode 100644 index 00000000..8b967454 --- /dev/null +++ b/h3_causal.py @@ -0,0 +1,296 @@ +"""FunPack's own chunk-causal MiniMax H3: a KV cache, and a DiT that reads it. + +WHY THIS EXISTS +--------------- +H3's stock DiT does ONE dense forward over the whole packed sequence. Chunk-causal generation +needs something it does not have: a forward over a slice of the clip that attends to a cache +of the slices already finished. That is the machinery here. + +Every class is a SUBCLASS of its ComfyUI counterpart, with the same module names and the same +state dict, so: + +* the weights load with no remapping and no conversion; +* a LoRA trained for the chunk-causal pattern attaches through ComfyUI's ordinary LoRA path + (its adapter already reads PEFT `lora_A`/`lora_B` keys); +* with no cache passed, every forward defers to its parent and is bit-identical to stock H3. + +That last point is the safety property. `cache=None` is the dense lane, unchanged, so nothing +here can degrade an ordinary generation. + +GEOMETRY IS BORROWED, NOT REDERIVED +----------------------------------- +A chunk's RoPE positions must be the ones the dense layout would have given those same rows, +or the model is being told the clip is a different shape than it is. So the FULL clip's grid +is built with upstream's own `_frame_grid` / `_video_grid` / `_audio_grid` and then SLICED per +chunk. An upstream change to the grid is inherited rather than silently disagreed with. + +THE LORA IS NOT OPTIONAL +------------------------ +Chunk-causal attention is a pattern the base H3 weights have never seen — they were trained +with every row attending every other row. Running this without an adapter trained for it +produces output, and the output is out of distribution. The loader says so. +""" + +from __future__ import annotations + +import torch + +#: video latents per chunk, and the pixel frames that many latents cover (1+4+4+4+4) +VIDEO_LATENTS_PER_CHUNK = 5 +FRAMES_PER_CHUNK = 17 +#: audio latents per pixel frame — 40/s against the model's fixed 24 fps +AUDIO_LATENTS_PER_FRAME = 40.0 / 24.0 + +#: AdaLN modality tags, from the dense model's own seg_tag +VIDEO_TAG, TEXT_TAG, AUDIO_TAG = 0, 1, 2 + + +class CacheError(RuntimeError): + """The cache was asked for something it cannot honestly answer.""" + + +class ChunkKVCache: + """Per-layer K/V of the chunks already finished, with sink + window eviction. + + Two kinds of memory, and they do different jobs: + + `sink` chunks are pinned from the very start and never evicted. Chunk 0 is the text, so a + sink of 2 means the prompt plus the opening shot stay visible to every later chunk — this + is what holds a character together across a long clip. + + `window` is the most recent chunks besides those, which is short-term continuity: motion + and lighting carrying from one moment to the next. + + Anything in between is dropped, which is what keeps memory flat instead of growing with + clip length. `window=None` keeps everything. + """ + + def __init__(self, num_layers, sink=2, window=2, device=None, offload=True): + if int(num_layers) <= 0: + raise CacheError(f"num_layers must be positive, got {num_layers!r}") + if int(sink) < 0: + raise CacheError(f"sink must be >= 0, got {sink!r}") + if window is not None and int(window) < 0: + raise CacheError(f"window must be None or >= 0, got {window!r}") + self.num_layers = int(num_layers) + self.sink = int(sink) + self.window = None if window is None else int(window) + self.device = device + # The retained K/V of a long clip is measured in tens of GiB on the card. Holding it + # on the host and copying one chunk back per layer trades bandwidth for the ability to + # run at all. + self.offload = bool(offload) + self.committed_chunks = 0 + self._store = {} # (layer, chunk_index) -> (k, v) + + def _keep(self, chunk_index, next_index): + if chunk_index < self.sink: + return True + if self.window is None: + return True + return chunk_index > next_index - 1 - self.window + + def retained_indices(self, next_index): + """Which committed chunks are visible to the chunk about to be generated.""" + return [i for i in range(self.committed_chunks) if self._keep(i, next_index)] + + def read(self, layer, next_index): + """The retained K and V for one layer, in chunk order, or (None, None).""" + parts = [self._store.get((layer, i)) for i in self.retained_indices(next_index)] + parts = [p for p in parts if p is not None] + if not parts: + return None, None + keys = torch.cat([p[0] for p in parts], dim=-2) + values = torch.cat([p[1] for p in parts], dim=-2) + if self.device is not None: + keys, values = keys.to(self.device), values.to(self.device) + return keys, values + + def write(self, layer, chunk_index, keys, values): + """Store one layer's K/V for a finished chunk.""" + if self.offload: + keys, values = keys.detach().to("cpu"), values.detach().to("cpu") + else: + keys, values = keys.detach(), values.detach() + self._store[(layer, int(chunk_index))] = (keys, values) + + def finish_chunk(self, chunk_index): + """Mark a chunk committed, and drop whatever is now out of both sink and window.""" + self.committed_chunks = max(self.committed_chunks, int(chunk_index) + 1) + self.prune() + + def prune(self): + keep = set(self.retained_indices(self.committed_chunks)) + for key in [k for k in self._store if k[1] not in keep]: + del self._store[key] + + def clear(self): + self._store.clear() + self.committed_chunks = 0 + + +def chunk_bounds(latent_t, audio_t): + """(video_start, video_stop, audio_start, audio_stop) per chunk. + + Video cuts every ``VIDEO_LATENTS_PER_CHUNK`` latents. Audio is derived from the SHARED + CLOCK — 40 latents per second against a fixed 24 fps, 17 pixel frames per chunk — so the + familiar 28/29/28/28 cadence comes out as a consequence and stays correct if the grid ever + moves, instead of being a remembered pattern that silently drifts. + """ + latent_t, audio_t = int(latent_t), int(audio_t) + if latent_t <= 0: + return [] + bounds, cursor = [], 0 + starts = list(range(0, latent_t, VIDEO_LATENTS_PER_CHUNK)) + for index, video_start in enumerate(starts): + video_stop = min(latent_t, video_start + VIDEO_LATENTS_PER_CHUNK) + if index == len(starts) - 1: + audio_stop = audio_t # the tail owns whatever is left + else: + done = FRAMES_PER_CHUNK * (index + 1) + audio_stop = max(cursor, min(audio_t, int(round(done * AUDIO_LATENTS_PER_FRAME)))) + bounds.append((video_start, video_stop, cursor, audio_stop)) + cursor = audio_stop + return bounds + + +def clip_grids(text_len, latent_t, latent_h, latent_w, audio_t): + """The whole clip's video and audio position grids, and the rows per latent frame. + + Built with upstream's own helpers and sliced per chunk afterwards, so a chunk's RoPE + angles are exactly the ones the dense layout would have given those rows. Re-deriving the + geometry here instead would be a second source of truth for the clip's shape. + """ + from comfy.ldm.minimax.model import _audio_grid, _frame_grid, _video_grid + + frame, w_grid = _frame_grid(int(latent_h), int(latent_w)) + video = _video_grid(int(latent_t), frame, float(text_len)) + audio = _audio_grid(float(text_len), int(audio_t), + float(w_grid[0]), float(w_grid[-1])) + return video, audio, int(frame.shape[0]) + + +def chunk_rows(audio_t, audio_start, audio_stop): + """Which packed audio ROWS a chunk owns. + + `pack_audio` is channel-major over the WHOLE clip — all of the left channel, then all of + the right — so a chunk's rows are two separate spans, not one. Treating them as one + contiguous slice silently hands the model the wrong half of the stereo field. + """ + left = list(range(int(audio_start), int(audio_stop))) + right = [int(audio_t) + i for i in left] + return left + right + + +def _causal_classes(): + """Build the three subclasses against whatever H3 this ComfyUI has. + + Deferred into a function because these subclass ``comfy.ldm.minimax.model``, which is not + importable on a ComfyUI without H3 — and a node pack must still register there. + """ + import comfy.model_management + import comfy.quant_ops + from comfy.ldm.minimax.model import (Attention, DiTBlock, MiniMaxH3Model, + _mod_gate, _mod_scale_shift, optimized_attention) + + class CausalAttention(Attention): + """Upstream attention, with the retained K/V prepended to this chunk's own. + + Q stays current-only: this chunk is what is being generated. K and V are the cache + followed by this chunk, so the chunk sees all of itself and all of what was retained. + There is no mask — that visibility IS the causal pattern, expressed by what is in the + buffer rather than by hiding parts of a bigger one. + + Cached keys keep the absolute positions they were rotated with, so nothing is ever + re-based. That is what makes the cache reusable across steps. + """ + + def forward(self, x, rope_freqs=None, transformer_options={}, + cache=None, layer_idx=0, chunk_index=0, update_cache=False): + if cache is None: + # the dense lane, untouched + return super().forward(x, rope_freqs=rope_freqs, + transformer_options=transformer_options) + s = x.shape[0] + q, k, v = self.qkv_proj(x).split(self.heads * self.head_dim, dim=-1) + v = v.view(s, self.heads, self.head_dim) + if rope_freqs is not None: + q = q.view(1, s, self.heads, self.head_dim) + k = k.view(1, s, self.heads, self.head_dim) + qw = comfy.model_management.cast_to(self.q_norm.weight, device=x.device) + kw = comfy.model_management.cast_to(self.k_norm.weight, device=x.device) + rot = rope_freqs.shape[-3] * 2 + if comfy.model_management.in_training: + q, k = comfy.quant_ops.ck.rms_rope_split_half( + q, k, rope_freqs, qw, kw, epsilon=self.q_norm.eps, rot_dim=rot) + else: + comfy.quant_ops.ck.rms_rope_split_half_( + q, k, rope_freqs, qw, kw, epsilon=self.q_norm.eps, rot_dim=rot) + q, k = q[0], k[0] + else: + q = self.q_norm(q.view(s, self.heads, self.head_dim)) + k = self.k_norm(k.view(s, self.heads, self.head_dim)) + q = q.transpose(0, 1).unsqueeze(0) + k = k.transpose(0, 1).unsqueeze(0) + v = v.transpose(0, 1).unsqueeze(0) + # Written BEFORE the cache is prepended: what this chunk contributes to later + # chunks is its own K/V, never the concatenation it happened to attend over. + if update_cache: + cache.write(layer_idx, chunk_index, k, v) + past_k, past_v = cache.read(layer_idx, chunk_index) + if past_k is not None: + k = torch.cat((past_k.to(dtype=k.dtype, device=k.device), k), dim=-2) + v = torch.cat((past_v.to(dtype=v.dtype, device=v.device), v), dim=-2) + out = optimized_attention(q, k, v, self.heads, mask=None, skip_reshape=True, + transformer_options=transformer_options) + return self.out_proj(out.squeeze(0)) + + class CausalDiTBlock(DiTBlock): + """Upstream block, with the cache threaded to its attention. + + A verbatim fork of the parent body — the norm / AdaLN / gate / MLP order is the + checkpoint's contract and must not drift — differing only in the kwargs passed on. + """ + + def forward(self, x, t_emb, mod_segments, rope_freqs, transformer_options={}, + cache=None, layer_idx=0, chunk_index=0, update_cache=False): + if cache is None: + return super().forward(x, t_emb, mod_segments, rope_freqs, + transformer_options=transformer_options) + shift_msa, scale_msa, gate_msa, shift_mlp, scale_mlp, gate_mlp = self.adaln_proj(t_emb) + h = _mod_scale_shift(self.norm1(x), shift_msa, scale_msa, mod_segments) + attended = self.attn(h, rope_freqs=rope_freqs, + transformer_options=transformer_options, + cache=cache, layer_idx=layer_idx, + chunk_index=chunk_index, update_cache=update_cache) + x = _mod_gate(x, gate_msa, attended, mod_segments) + h = _mod_scale_shift(self.norm2(x), shift_mlp, scale_mlp, mod_segments) + return _mod_gate(x, gate_mlp, self.mlp(h), mod_segments) + + class CausalMiniMaxH3Model(MiniMaxH3Model): + """Upstream H3 whose blocks can read a chunk cache. Same state dict, key for key. + + The blocks are rebuilt as causal ones after construction, ONE AT A TIME so peak memory + stays at a single extra block rather than a second copy of all fifty. Each rebuilt + block adopts the original's submodules by reference, so no weight is copied and no key + is renamed — the hard rule this project already learned about wrapper modules. + """ + + def __init__(self, *args, **kwargs): + super().__init__(*args, **kwargs) + for index in range(len(self.blocks)): + self.blocks[index] = _to_causal_block(self.blocks[index], + CausalDiTBlock, CausalAttention) + + return CausalAttention, CausalDiTBlock, CausalMiniMaxH3Model + + +def _to_causal_block(block, block_cls, attention_cls): + """Re-class one block and its attention in place, adopting every submodule by reference. + + NOT a wrapper module. Wrapping renames the state-dict keys of everything inside it, which + has broken unpatching here before; re-classing changes only which `forward` runs. + """ + block.attn.__class__ = attention_cls + block.__class__ = block_cls + return block diff --git a/tests/test_h3_causal.py b/tests/test_h3_causal.py new file mode 100644 index 00000000..febff3a7 --- /dev/null +++ b/tests/test_h3_causal.py @@ -0,0 +1,210 @@ +"""FunPack's own chunk-causal H3: the cache, the chunk geometry, and the re-classed DiT. + +H3's stock DiT does one dense forward over the whole packed sequence. Chunk-causal generation +needs a forward over a SLICE that attends to a cache of the slices already finished. This is +that machinery, owned rather than borrowed — no third-party package has to be installed for +the feature to exist. + +The safety property everything rests on: with no cache passed, every forward defers to its +parent, so an ordinary generation cannot be degraded by any of this. +""" +import sys +from pathlib import Path + +import pytest +import torch + +sys.path.insert(0, str(Path(__file__).resolve().parents[1])) + +import h3_causal as hc + + +# ── the cache: what a chunk is allowed to remember ────────────────────────── + +def _cache(**kw): + kw.setdefault("num_layers", 2) + return hc.ChunkKVCache(**kw) + + +def _kv(rows, value): + return torch.full((1, 2, rows, 4), float(value)), torch.full((1, 2, rows, 4), float(value)) + + +def test_an_empty_cache_has_nothing_to_read(): + assert _cache().read(0, 0) == (None, None) + + +def test_a_committed_chunk_is_readable(): + cache = _cache() + k, v = _kv(3, 1.0) + cache.write(0, 0, k, v) + cache.finish_chunk(0) + got_k, got_v = cache.read(0, 1) + assert got_k.shape[-2] == 3 and got_v.shape[-2] == 3 + + +def test_layers_do_not_share_their_cache(): + """Layer 7 reading layer 3's keys would be silent nonsense, not a crash.""" + cache = _cache(num_layers=4) + cache.write(0, 0, *_kv(3, 1.0)) + cache.finish_chunk(0) + assert cache.read(1, 1) == (None, None) + + +def test_chunks_are_read_back_in_order(): + cache = _cache() + for index, value in enumerate((1.0, 2.0)): + cache.write(0, index, *_kv(2, value)) + cache.finish_chunk(index) + keys, _v = cache.read(0, 2) + assert keys[0, 0, 0, 0] == 1.0 and keys[0, 0, 2, 0] == 2.0 + + +def test_sink_chunks_are_never_evicted(): + """The prompt and the opening shot are what hold a character together over a long clip.""" + cache = _cache(sink=1, window=1) + for index in range(6): + cache.write(0, index, *_kv(2, index)) + cache.finish_chunk(index) + assert 0 in cache.retained_indices(6) + + +def test_the_window_keeps_only_the_most_recent(): + cache = _cache(sink=1, window=2) + for index in range(6): + cache.write(0, index, *_kv(2, index)) + cache.finish_chunk(index) + assert cache.retained_indices(6) == [0, 4, 5] + + +def test_a_zero_window_leans_entirely_on_the_sink(): + cache = _cache(sink=2, window=0) + for index in range(5): + cache.write(0, index, *_kv(2, index)) + cache.finish_chunk(index) + assert cache.retained_indices(5) == [0, 1] + + +def test_no_window_remembers_everything(): + cache = _cache(sink=1, window=None) + for index in range(5): + cache.write(0, index, *_kv(2, index)) + cache.finish_chunk(index) + assert cache.retained_indices(5) == [0, 1, 2, 3, 4] + + +def test_evicted_chunks_are_actually_dropped(): + """Retention that only hides rows would grow memory with clip length, which is the thing + chunking exists to avoid.""" + cache = _cache(sink=1, window=1) + for index in range(8): + cache.write(0, index, *_kv(64, index)) + cache.finish_chunk(index) + assert len(cache._store) <= 3 + + +def test_a_negative_sink_is_refused(): + with pytest.raises(hc.CacheError): + _cache(sink=-1) + + +def test_zero_layers_is_refused(): + with pytest.raises(hc.CacheError): + hc.ChunkKVCache(0) + + +def test_clearing_forgets_everything(): + cache = _cache() + cache.write(0, 0, *_kv(2, 1.0)) + cache.finish_chunk(0) + cache.clear() + assert cache.committed_chunks == 0 and cache.read(0, 1) == (None, None) + + +def test_offloaded_entries_leave_the_compute_device(): + """The retained K/V of a long clip is tens of GiB; holding it on the card is what makes a + long clip impossible rather than merely slow.""" + cache = _cache(offload=True) + cache.write(0, 0, *_kv(2, 1.0)) + assert cache._store[(0, 0)][0].device.type == "cpu" + + +# ── the chunk cut ─────────────────────────────────────────────────────────── + +def test_video_cuts_every_five_latents(): + bounds = hc.chunk_bounds(latent_t=12, audio_t=100) + assert [(a, b) for a, b, _c, _d in bounds] == [(0, 5), (5, 10), (10, 12)] + + +def test_every_latent_lands_in_exactly_one_chunk(): + bounds = hc.chunk_bounds(latent_t=22, audio_t=180) + video, audio = [], [] + for vs, ve, as_, ae in bounds: + video.extend(range(vs, ve)) + audio.extend(range(as_, ae)) + assert video == list(range(22)) and audio == list(range(180)) + + +def test_the_audio_cadence_falls_out_of_the_clock(): + """40 latents/s against a fixed 24 fps, 17 frames per chunk. The 28/29 alternation is a + CONSEQUENCE — hard-coding it would drift silently if the grid ever moved.""" + frames = 17 * 11 + 5 + bounds = hc.chunk_bounds(latent_t=5 * 11 + 2, audio_t=int(round(frames / 24 * 40))) + per_chunk = [ae - as_ for _a, _b, as_, ae in bounds] + assert per_chunk[:4] == [28, 29, 28, 28] + + +def test_the_tail_chunk_is_kept_not_dropped(): + """H3's grid is 5k+2 latents, so there is ALWAYS a two-latent tail. Dropping it would cut + the last five frames off every clip.""" + bounds = hc.chunk_bounds(latent_t=17, audio_t=140) + assert bounds[-1][1] - bounds[-1][0] == 2 + + +# ── stereo rows: the trap in packed audio ─────────────────────────────────── + +def test_a_chunks_audio_rows_are_two_spans_not_one(): + """pack_audio is channel-major over the WHOLE clip — all of the left channel, then all of + the right. Treating a chunk's rows as one contiguous slice hands the model the wrong half + of the stereo field, quietly.""" + rows = hc.chunk_rows(audio_t=100, audio_start=28, audio_stop=57) + assert rows[:29] == list(range(28, 57)) + assert rows[29:] == list(range(128, 157)) + + +def test_the_row_count_is_two_per_latent(): + rows = hc.chunk_rows(audio_t=100, audio_start=0, audio_stop=28) + assert len(rows) == 56 + + +def test_the_first_chunk_starts_at_both_channel_origins(): + rows = hc.chunk_rows(audio_t=50, audio_start=0, audio_stop=4) + assert rows == [0, 1, 2, 3, 50, 51, 52, 53] + + +# ── re-classing, not wrapping ─────────────────────────────────────────────── + +def test_a_block_is_re_classed_in_place_and_keeps_its_submodules(): + """Wrapping renames the state-dict keys of everything inside, which has broken unpatching + here before. Re-classing changes only which forward runs.""" + class FakeAttention: + pass + + class FakeBlock: + def __init__(self): + self.attn = FakeAttention() + self.norm1 = "the-real-norm" + + class CausalAttn(FakeAttention): + pass + + class CausalBlock(FakeBlock): + def __init__(self): + raise AssertionError("re-classing must not re-run __init__") + + block = FakeBlock() + attn, norm = block.attn, block.norm1 + out = hc._to_causal_block(block, CausalBlock, CausalAttn) + assert out is block + assert out.attn is attn and out.norm1 is norm + assert isinstance(out, CausalBlock) and isinstance(out.attn, CausalAttn)