From d80c17bfaadc59517e6fd8efdf73ca853f5dc460 Mon Sep 17 00:00:00 2001 From: digital-garbage <133010426+digital-garbage@users.noreply.github.com> Date: Wed, 26 Aug 2026 23:39:32 +0300 Subject: [PATCH] maint: remove the chunk-causal lane MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Validated on the rental and it does not earn its place: 3.5 minutes against 1 for the same clip, worse picture, garbage audio. Some of that was mine — no weight prefetch across the 50 blocks, a host-offloaded K/V cache, and an extra commit forward per chunk — but fixing those only buys parity on speed, and the technique's real win is unbounded length, which H3 already covers natively at the lengths in use here. Gone: h3_causal.py, the four h3_causal_* sampler widgets and their Engine Settings rows, `chunk_causal` on the diffusion loader, and five test files. KEPT, because it is not about this at all: the LoRA key RENAME table. Stripping a wrapper head leaves `blocks.0...`, comfy's key map has no bare form, and the file silently applies nothing — 532 keys, zero applied, no error. That was breaking any PEFT-style adapter named this way, not just the one that found it. Its tests moved to test_lora_shape_mismatch.py alongside the MATCHED NOTHING diagnostics. Two INPUT_TYPES widgets go with this (h3_causal_* on the sampler, chunk_causal on the loader). A saved workflow carrying them loads fine — both were optional and are matched by name — but the values are ignored from here. Tests: 1700 python, 68 JS. --- h3_causal.py | 995 ------------------ loaders.py | 30 +- model_management.py | 12 +- movie_editor/frontend/engine_defaults.test.js | 4 +- movie_editor/frontend/engine_settings.js | 15 +- movie_editor/frontend/pipeline_caps.js | 8 - samplers.py | 140 +-- tests/_h3_causal_upstream_check.py | 250 ----- tests/test_causal_rollout.py | 504 --------- tests/test_funpack_log.py | 10 +- tests/test_h3_causal.py | 210 ---- tests/test_h3_causal_loader.py | 208 ---- tests/test_h3_causal_plan.py | 181 ---- tests/test_h3_causal_session.py | 403 ------- tests/test_lora_shape_mismatch.py | 55 +- 15 files changed, 74 insertions(+), 2951 deletions(-) delete mode 100644 h3_causal.py delete mode 100644 tests/_h3_causal_upstream_check.py delete mode 100644 tests/test_causal_rollout.py delete mode 100644 tests/test_h3_causal.py delete mode 100644 tests/test_h3_causal_loader.py delete mode 100644 tests/test_h3_causal_plan.py delete mode 100644 tests/test_h3_causal_session.py diff --git a/h3_causal.py b/h3_causal.py deleted file mode 100644 index ba80372b..00000000 --- a/h3_causal.py +++ /dev/null @@ -1,995 +0,0 @@ -"""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 ONE full-clip -`PackedLayout` is built by ComfyUI's own `extra_conds` and every chunk is a SELECTION of its -rows. An upstream change to the layout is inherited rather than silently disagreed with, and -the keyframe pins and reference blocks keep the rows upstream packed for them — which is what -makes i2v and r2v work on this lane instead of being refused by it. - -WHAT IS OURS AND WHAT IS THE MODEL'S ------------------------------------- -The chunk cut, the cache, the schedule and the step rule are all here: any scheduler, any step -count, and a step rule that can be set to plain flow euler so a chunked run can be compared -against a dense one without two variables moving. Nothing in this file needs a third-party -package installed. - -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 - - -class ChunkPlan: - """Which packed rows each chunk owns, cut out of ONE full-clip PackedLayout. - - The layout is upstream's own, built for the WHOLE clip. Every chunk is then a selection of - its rows, so a chunk's RoPE angles, modality tags and timestep classes are exactly the ones - the dense forward would have given those same rows. Re-deriving the geometry per chunk - 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. - - Three kinds of selection, and they are three because THE PROMPT AND THE ANCHOR ARE NOT - THE SAME KIND OF CONTEXT: - - `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 + 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) - self.audio_t = int(audio_t) - segments = list(layout.segments) - audio_seg = next((s for s in segments if s[2] == "audio"), None) - video_seg = next((s for s in segments if s[2] == "video"), None) - if audio_seg is None or video_seg is None: - raise CacheError("this layout has no target audio/video segment to chunk") - latent_t = self.bounds[-1][1] if self.bounds else 0 - if latent_t <= 0: - raise CacheError("a clip with no video latents has nothing to chunk") - # Rows per latent frame, ASKED OF THE LAYOUT rather than re-derived from the patch - # size. The layout already answered this question when it packed the video segment, - # and a second answer is a second thing that can drift. - self.frame_rows = (video_seg[1] - video_seg[0]) // latent_t - self._audio_base, self._video_base = audio_seg[0], video_seg[0] - # Every chunk must own rows of BOTH streams: the final layer reads one contiguous - # span of each, so a chunk with no audio is a StopIteration fifty blocks deep rather - # than an answer. It only happens when the two streams disagree about the clip's - # length, so say that instead. - empty = [i for i, (vs, ve, a, b) in enumerate(self.bounds) if ve <= vs or b <= a] - if empty: - raise CacheError( - 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 = 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): - return len(self.bounds) - - @staticmethod - def cache_index(index): - """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.""" - v_start, v_stop, a_start, a_stop = self.bounds[int(index)] - audio = [self._audio_base + r for r in chunk_rows(self.audio_t, a_start, a_stop)] - video = range(self._video_base + v_start * self.frame_rows, - self._video_base + v_stop * self.frame_rows) - rows = torch.tensor(list(audio) + list(video), dtype=torch.long) - n_audio = len(audio) - runs = [(0, n_audio, "audio"), (n_audio, rows.shape[0], "video")] - return rows, [r for r in runs if r[1] > r[0]] - - -def build_plan(layout, latent_t, audio_t): - """A ChunkPlan for a clip, with the chunk cuts derived from the shared clock.""" - return ChunkPlan(layout, chunk_bounds(latent_t, audio_t), audio_t) - - -#: how a chunk's next-sigma is reached from its prediction -STEP_RULES = ("consistency", "euler", "euler_ancestral") - - -def step(rule, x_t, prediction, sigma, sigma_next, noise, eta=1.0): - """Advance one step, from a model prediction to the next sigma. - - `prediction` is x0. H3's head emits data-ward velocity and the caller converts with - ``x0 = x_t + sigma * v`` before getting here — no negation; that is the model's own - convention, and getting it wrong inverts the whole clip rather than degrading it. - - consistency - RAVEN's transition: jump to x0 and re-noise to the next sigma with FRESH noise, never - the noise already in x_t. This is what the RAVEN LoRA is distilled for. - euler - The ordinary flow step. Deterministic, and the same rule FunPack's other H3 sampling - uses — so a chunked run can be compared against an unchunked one without the step rule - being a second variable. - euler_ancestral - Euler with `eta` of the step's noise put back, in the RECTIFIED-FLOW form. The VP - formula is wrong on a flow model and was a real bug here before. - """ - sigma, sigma_next = float(sigma), float(sigma_next) - if rule == "consistency": - return (1.0 - sigma_next) * prediction + sigma_next * noise - if rule not in ("euler", "euler_ancestral"): - raise ValueError(f"unknown step rule {rule!r}; expected one of {STEP_RULES}") - if sigma <= 0: - return prediction - derivative = (x_t - prediction) / sigma # dx/dsigma on the flow ODE - if rule == "euler" or sigma_next <= 0 or eta <= 0: - return x_t + (sigma_next - sigma) * derivative - # rectified-flow ancestral: shorten the step, then put back what was taken out - down_ratio = 1.0 + (sigma_next / sigma - 1.0) * float(eta) - sigma_down = sigma_next * down_ratio - stepped = x_t + (sigma_down - sigma) * derivative - renoise = max(0.0, sigma_next ** 2 - sigma_down ** 2) ** 0.5 - return stepped + noise * renoise - - -def audio_sigma_for(sigma, shift_video, shift_audio): - """The audio stream's own sigma for a video sigma, through the flow time-shift. - - H3 runs the two streams on DIFFERENT shifted grids — 12 for video, 3 for audio — and the - sampler only ever sees one array. Reusing the video sigma for the audio rows tells the - model the soundtrack is at a noise level it is not at, and the error grows with the step - size, which is why it bites hardest on the few-step schedules this lane is built for. - - Same remap ``h3_audio_clock`` uses: undo the video shift to recover the unshifted time, - then apply the audio shift to it. - """ - sigma = float(sigma) - if sigma <= 0.0 or float(shift_video) == float(shift_audio): - return sigma - unshifted = sigma / (float(shift_video) - (float(shift_video) - 1.0) * sigma) - return (float(shift_audio) * unshifted - / (1.0 + (float(shift_audio) - 1.0) * unshifted)) - - -def causal_rollout(*, chunks, sigmas, forward, commit, draw_noise, - video_noise, audio_noise, step_rule="consistency", eta=1.0, - known_chunks=0, known_video=None, known_audio=None, - on_chunk=None, on_step=None, cancel=None): - """Generate a clip chunk by chunk, each one attending to the cache of the ones before it. - - Everything the model does arrives as a callable, so the loop is testable without ComfyUI, - without the RAVEN package and without weights — and so the same loop can be driven by a - different backend later without being rewritten. - - forward(video_xt, audio_xt, index, sigma) - -> (video_velocity, audio_velocity) for that chunk at that sigma. Cache read-only. - commit(video_x0, audio_x0, index) - writes the finished chunk into the KV cache as clean context. Called ONCE per chunk, - after its last step, and for the pre-known chunks too — a chunk that is not committed - is a chunk the rest of the clip cannot see. - draw_noise(shape) - fresh noise. Called in a fixed order so a seeded run reproduces exactly. - on_step(chunk_index, done, total) - after every model call, with the count of chunk-steps finished out of the clip's - total. This lane does not go through comfy's sampler, so nothing else ticks a - progress bar — and a long generation showing no progress is indistinguishable from - a hung one. - - `known_chunks` leading chunks are taken from `known_video`/`known_audio` instead of being - sampled. That is how an i2v anchor survives: chunk 0 comes from FunPack's ORDINARY dense - path — where pins, references and region locks all work — and is only then committed here, - so the causal lane is never asked to model conditioning rows it has no layout for. - - `sigmas` is FunPack's own schedule, descending and ending at 0. Any scheduler, any step - count; the RAVEN adapter's distilled 4 is a default, not a constraint. - """ - if step_rule not in STEP_RULES: - raise ValueError(f"unknown step rule {step_rule!r}; expected one of {STEP_RULES}") - sigmas = [float(s) for s in sigmas] - if len(sigmas) < 2: - raise ValueError("a schedule needs at least one step (two sigmas)") - known = int(known_chunks) - if known and (known_video is None or known_audio is None): - raise ValueError("known_chunks was given without the latents to take them from") - - video_out = video_noise.clone() - audio_out = audio_noise.clone() - - for index, (v_start, v_stop, a_start, a_stop) in enumerate(chunks): - if cancel is not None: - cancel(index) - if index < known: - # Already generated elsewhere. It still has to be COMMITTED, or every later - # chunk continues from a clip that appears to start after it. - video_x0 = known_video[:, :, v_start:v_stop] - audio_x0 = known_audio[..., a_start:a_stop] - video_out[:, :, v_start:v_stop] = video_x0 - audio_out[..., a_start:a_stop] = audio_x0 - commit(video_x0, audio_x0, index) - if on_chunk is not None: - on_chunk(index, video_x0, audio_x0, True) - continue - - video_xt = video_noise[:, :, v_start:v_stop].clone() - audio_xt = audio_noise[..., a_start:a_stop].clone() - for position in range(len(sigmas) - 1): - if cancel is not None: - cancel(index) - sigma, sigma_next = sigmas[position], sigmas[position + 1] - video_v, audio_v = forward(video_xt, audio_xt, index, sigma) - # H3's head is data-ward velocity; x0 = x_t + sigma * v, no negation. - video_x0 = video_xt + sigma * video_v - audio_x0 = audio_xt + sigma * audio_v - # Both draws happen every step even when they are multiplied by zero at the end: - # skipping one would shift the noise stream for every later chunk and a seeded - # run would stop reproducing. - video_eps = draw_noise(video_xt.shape) - audio_eps = draw_noise(audio_xt.shape) - video_xt = step(step_rule, video_xt, video_x0, sigma, sigma_next, video_eps, eta) - audio_xt = step(step_rule, audio_xt, audio_x0, sigma, sigma_next, audio_eps, eta) - if on_step is not None: - on_step(index, index * (len(sigmas) - 1) + position + 1, - len(chunks) * (len(sigmas) - 1)) - - video_out[:, :, v_start:v_stop] = video_xt - audio_out[..., a_start:a_stop] = audio_xt - commit(video_xt, audio_xt, index) - if on_chunk is not None: - on_chunk(index, video_xt, audio_xt, False) - - return video_out, audio_out - - -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.ldm.minimax.model as upstream - import comfy.model_management - import comfy.quant_ops - # The names without which there is no lane at all. A missing one is named by the - # ImportError, which is how a ComfyUI too old for this gets diagnosed rather than guessed. - from comfy.ldm.minimax.model import (Attention, DiTBlock, MiniMaxH3Model, _mod_gate, - _mod_scale_shift, optimized_attention, pack_audio, - patchify_video, rope_rotation_table, unpack_audio, - unpatchify_video) - - # And the names that ARRIVED LATER. Importing these the same way made one missing symbol - # refuse the whole feature on an older ComfyUI — which is what happened, and the refusal - # then blamed a setting. Each falls back to what that ComfyUI's own dense forward does, - # so the chunked lane never applies a correction the dense path is not applying either. - VISUAL_COND_TIMESTEP = getattr(upstream, "VISUAL_COND_TIMESTEP", 0.999) - AUDIO_COND_TIMESTEP = getattr(upstream, "AUDIO_COND_TIMESTEP", 1.0) - time_shift_slope = getattr(upstream, "time_shift_slope", None) - - 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) - - # ── the two entry points the rollout drives ───────────────────────────────── - - def _modulation(self, kinds, t_v, t_a, payload): - """seg_t / t_row / t_emb for the kinds actually present in this sequence. - - Upstream's own analytic table: text and video ride the video's time, audio rides - the audio's, and the condition rows pin near 1 because that is the noise level - they were augmented to. Only the kinds present get a timestep ROW, so a chunk with - no cond rows does not carry an unused one into the AdaLN table. - """ - vis_aug = float(payload.get("visual_cond_noise_aug", VISUAL_COND_TIMESTEP)) - aud_aug = float(payload.get("audio_cond_noise_aug", AUDIO_COND_TIMESTEP)) - seg_t = {"text": t_v, "video": t_v, "audio": t_a, - "cond": max(t_v, vis_aug), "ref_img": max(t_v, vis_aug), - "ref_audio": max(t_a, aud_aug)} - unique_t = sorted({seg_t[k] for k in kinds}) - return seg_t, {t: i for i, t in enumerate(unique_t)}, unique_t - - def _t_emb(self, unique_t, device, dtype): - t_vals = torch.tensor(unique_t, dtype=torch.float32, device=device) - if self.use_adaln_curves: - table = comfy.model_management.cast_to(self.adaln_t_table, device=device) - pos = t_vals.clamp(0.0, 1.0) * (table.shape[0] - 1) - i0 = pos.floor().long().clamp(max=table.shape[0] - 2) - return torch.lerp(table[i0], table[i0 + 1], (pos - i0).unsqueeze(1)) - return self.time_embedder(t_vals).to(dtype) - - def _mod_segments(self, runs, seg_t, t_row, text_tags): - """Local (start, stop, adaln row) runs, splitting the text span by its tag runs. - - The text span mixes modality tags — a reference image's vision pads sit inside the - text and carry the VIDEO tag — so it cannot be treated as one uniform run without - modulating the picture rows as if they were words. - """ - seg_tag = {"text": TEXT_TAG, "video": VIDEO_TAG, "audio": AUDIO_TAG, - "cond": VIDEO_TAG, "ref_img": VIDEO_TAG, "ref_audio": AUDIO_TAG} - out = [] - for a, b, kind in runs: - base = t_row[seg_t[kind]] * 3 - if kind == "text" and text_tags is not None: - tags = text_tags.view(-1).tolist() - run_start = 0 - for i in range(1, b - a + 1): - if i == b - a or tags[i] != tags[run_start]: - out.append((a + run_start, a + i, base + int(tags[run_start]))) - run_start = i - else: - out.append((a, b, base + seg_tag[kind])) - return out - - def _run_blocks(self, h, t_emb, mod_segments, rope_freqs, transformer_options, - cache, chunk_index, update_cache): - for layer, block in enumerate(self.blocks): - h = block(h, t_emb, mod_segments, rope_freqs, - transformer_options=transformer_options, cache=cache, - layer_idx=layer, chunk_index=chunk_index, - update_cache=update_cache) - return h - - def _text_states(self, context, transformer_options): - states = context[0] - if states.shape[-1] == self.hidden_size: - return states - return self.token_refiner(self.condition_proj(states), - transformer_options=transformer_options) - - 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. - - `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. - """ - seg_t, t_row, unique_t = self._modulation({k for _, _, k in runs}, t_v, t_a, payload) - 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: - n = b - a - if kind == "text": - h[a:b] = text_states - elif kind in ("cond", "ref_img"): - if cond_video is None: - raise CacheError(f"the layout has {kind} rows but the payload carries " - f"no cond_video_latents to fill them") - h[a:b] = self.video_patch_proj(cond_video[voff:voff + n]).to(dtype) - voff += n - else: # ref_audio - if cond_audio is None: - raise CacheError("the layout has ref_audio rows but the payload " - "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, 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, - transformer_options={}, minimax_payload=None, - sigma_shift_video=None, sigma_shift_audio=None): - """One media chunk's forward against the cache. Returns (video_v, audio_v). - - DATA-WARD velocity, which is what `x0 = x_t + sigma * v` wants — upstream's - `_forward` negates on the way out because ComfyUI's flow sampler expects the other - sign, and following it here would invert the whole clip rather than degrade it. - - The audio velocity carries the same `time_shift_slope` factor upstream applies, so - the caller can integrate BOTH streams against the video sigma and still be on the - audio's own shifted grid — on a ComfyUI that applies it. 0.34.0 has no such - function and its dense forward returns the audio velocity unscaled, so neither - does this: the chunked lane must not correct for something the dense path on the - same install is leaving alone. - """ - payload = minimax_payload or {} - rows, runs = plan.chunk(index) - device = video_latent.device - dtype = context.dtype if context is not None else video_latent.dtype - latent_t = video_latent.shape[2] - lat_h, lat_w = video_latent.shape[3], video_latent.shape[4] - sigma_v = max(float(video_sigma), 1e-6) - t_v = float(1.0 - sigma_v) - t_a = float(1.0 - max(float(audio_sigma), 1e-6)) - seg_t, t_row, unique_t = self._modulation({k for _, _, k in runs}, t_v, t_a, payload) - - video_embed = self.video_patch_proj( - patchify_video(video_latent.to(torch.float32), self.patch_size)).to(dtype) - audio_embed = self.audio_patch_proj( - pack_audio(audio_latent.to(torch.float32))).to(dtype) - h = torch.empty(int(rows.shape[0]), self.hidden_size, dtype=dtype, device=device) - for a, b, kind in runs: - h[a:b] = video_embed if kind == "video" else audio_embed - - rope_freqs = rope_rotation_table( - self.rope_freqs(plan.layout.position_ids[rows], device), dtype) - t_emb = self._t_emb(unique_t, device, dtype) - h = self._run_blocks(h, t_emb, - self._mod_segments(runs, seg_t, t_row, None), - rope_freqs, transformer_options, cache, - plan.cache_index(index), update_cache) - - video_seg = next((a, b, t_row[seg_t["video"]]) for a, b, k in runs if k == "video") - audio_seg = next((a, b, t_row[seg_t["audio"]]) for a, b, k in runs if k == "audio") - v, a = self.final_layer(h, t_emb, video_seg, audio_seg) - video_out = unpatchify_video(v, latent_t, lat_h // 2, lat_w // 2, - self.latents_dim, self.patch_size) - audio_out = unpack_audio(a) - if time_shift_slope is not None: - shift_v = float(self.sigma_shift_video if sigma_shift_video is None - else sigma_shift_video) - shift_a = float(self.sigma_shift_audio if sigma_shift_audio is None - else sigma_shift_audio) - slope = time_shift_slope(torch.tensor(sigma_v), shift_v, shift_a) - audio_out = slope.to(audio_out.dtype) * audio_out - return (video_out.to(video_latent.dtype), audio_out.to(audio_latent.dtype)) - - 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 - - -def make_causal(model): - """Re-class an already-loaded H3 so its blocks can read a chunk cache. Returns (ok, note). - - This is why nothing here needs its own loader. The alternative — building the model - through a different class — means reimplementing ComfyUI's whole diffusion-model load - path, and attaching anything to the raw model before the ModelPatcher exists so its - memory ledger counts it. Re-classing sidesteps all of that: no parameter is added, no key - is renamed, nothing is copied. Only which `forward` runs changes. - - Idempotent, and a no-op on anything that is not H3. - """ - diffusion = getattr(getattr(model, "model", None), "diffusion_model", None) - if diffusion is None: - return False, "no diffusion model to make causal" - blocks = getattr(diffusion, "blocks", None) - if not blocks: - return False, "this model has no DiT blocks — not a MiniMax H3" - if getattr(diffusion, "_funpack_causal", False): - return True, "already chunk-causal" - try: - attention_cls, block_cls, model_cls = _causal_classes() - except Exception as error: # noqa: BLE001 - return False, f"this ComfyUI has no MiniMax H3 to extend ({error})" - if not isinstance(diffusion, model_cls.__bases__[0]): - return False, "this model is not a MiniMax H3" - for index in range(len(blocks)): - _to_causal_block(blocks[index], block_cls, attention_cls) - diffusion.__class__ = model_cls - diffusion._funpack_causal = True - return True, f"chunk-causal DiT: {len(blocks)} blocks re-classed in place" - - -# ── driving a whole clip ──────────────────────────────────────────────────────────── - - -def _unbind_av(samples): - """(video, audio) out of the nested container H3 latents travel in.""" - if hasattr(samples, "unbind"): - parts = list(samples.unbind()) - if len(parts) == 2: - return parts[0], parts[1] - raise CacheError("this latent is not an H3 video+audio pair") - - -def _ensure_loaded(patcher, noise_shape, conds): - """Put the weights on the card the way the dense path does, and say so if it cannot. - - THIS MUST HAPPEN BEFORE ANY FORWARD, AND BEFORE THE CONDITIONING IS PREPARED. The - ordinary path reaches the model through comfy.sample.sample_custom, whose prepare_sampling - calls load_models_gpu — and load_models_gpu is also what applies a patcher's OBJECT - PATCHES (partially_load -> patch_model), which is how FunPack's AdaLN modality gains and - token-refiner edit get installed at all. This lane calls the DiT directly, so without this - the weights may still be offloaded AND every object patch would silently not be there. - - Preparing the conditioning is itself a partial forward — extra_conds runs H3's token - refiner — so doing it before the load is what produced `mat1 is on cuda:0, different from - other tensors on cpu`. ComfyUI's own prepare_sampling is used rather than a bare - load_models_gpu so the memory estimate, the wrappers and the placement are the ones the - dense path would have made for this clip. - - Returns None when the model is where it should be, or the reason it is not — the caller - decides whether that is fatal. It is at build time (nothing has run yet, so refusing is - clean) and it is not at run time (the rollout is already under way). - """ - if patcher is None: - return None - try: - import comfy.sampler_helpers as helpers - prepare = helpers.prepare_sampling - convert = helpers.convert_cond - except (ImportError, AttributeError): - # A ComfyUI that does not expose the helper still has to get its weights placed. - try: - import comfy.model_management - comfy.model_management.load_models_gpu([patcher], force_full_load=False) - return None - except Exception as error: # noqa: BLE001 - return str(error) - try: - # convert_cond FIRST. `estimate_memory` calls `extra_conds_shapes(**cond)`, so it - # wants ComfyUI's INTERNAL cond dicts, not the [tensor, meta] pairs a node is handed. - # Passing the node-level list raised, the load never happened, and the conditioning - # then failed on a model that was still half on the CPU — one bug reported twice. - prepared = {k: convert(v) for k, v in (conds or {}).items()} - prepare(patcher, tuple(noise_shape), prepared, - getattr(patcher, "model_options", None)) - return None - except Exception as error: # noqa: BLE001 - return str(error) - - -def build_session(model, positive, latent, *, sink=2, window=2, device=None, offload=True): - """Assemble everything one causal run needs, or explain why it cannot be assembled. - - Returns (session, reason). `session` is None when the run cannot be built, and `reason` is - written for the person who has to fix it rather than for a log parser. - - The PAYLOAD is built by ComfyUI's own `extra_conds`, not by a parser of our own. That is - deliberate: the keyframe rows, the reference blocks, the condition noise augmentation and - the packed layout are then bit-for-bit the ones the dense path would have used, so an - upstream change to any of them is inherited instead of silently disagreed with. - """ - diffusion = getattr(getattr(model, "model", None), "diffusion_model", None) - if diffusion is None: - return None, "this model has no diffusion model to run." - if not hasattr(diffusion, "forward_chunk"): - # Do it here rather than telling the user to go and switch something on. Re-classing - # adds no parameter, renames no key and costs nothing, and with no cache passed the - # model is bit-identical to stock H3 — so there is nothing to protect them from and - # no reason a mode that is already switched ON should refuse itself over it. - ok, note = make_causal(model) - if not ok: - return None, (f"{note} (the diffusion model is " - f"{type(diffusion).__name__}).") - # Say that it happened HERE. Reaching this point means the loader's `chunk_causal` - # did not take — either it is off, or its value never arrived — and knowing which - # end did the work is the difference between a fixed setting and a mystery. - print("[FunPack H3] the chunk-causal DiT was installed by the Chain Sampler, not by " - "the loader: FunPack Diffusion Model Loader's `chunk_causal` is off or its " - f"value did not reach the node. Nothing is lost — {note}.") - blocks = getattr(diffusion, "blocks", None) - if not blocks: - return None, "the loaded model exposes no DiT blocks to cache." - try: - video, audio = _unbind_av(latent["samples"]) - except Exception as error: # noqa: BLE001 - return None, f"this run is not one the causal lane can take: {error}" - if video.shape[0] != 1: - return None, "MiniMax H3 samples one clip at a time; this batch is larger." - - if device is None: - device = getattr(model, "load_device", None) or video.device - device = torch.device(device) - # Before extra_conds, which runs the token refiner: a partial forward that needs the - # weights where the tensors it is handed already are. Refused rather than warned about, - # because everything after this point assumes the model is loaded and the failure would - # otherwise surface as a device mismatch three lines later. - failure = _ensure_loaded(model, video.shape, {"positive": positive}) - if failure is not None: - return None, f"the model could not be placed for the causal lane: {failure}" - try: - base = model.model - kwargs = dict(positive[0][1]) - kwargs["cross_attn"] = positive[0][0].to(device) - kwargs["device"] = device - kwargs["latent_shapes"] = [tuple(video.shape), tuple(audio.shape)] - kwargs.setdefault("seed", 0) - conds = base.extra_conds(**kwargs) - payload = conds["minimax_payload"].cond - context = conds["c_crossattn"].cond - except Exception as error: # noqa: BLE001 - return None, f"the conditioning could not be prepared for the causal lane: {error}" - - layout = payload.get("layout") - if layout is None: - return None, ("this run produced no packed layout — the causal lane needs the same " - "one the dense path builds.") - try: - plan = build_plan(layout, int(video.shape[2]), int(audio.shape[-1])) - except CacheError as error: - return None, str(error) - if plan.n_chunks < 2: - return None, (f"this clip is {plan.n_chunks} chunk long — there is nothing for a " - f"chunk cache to remember.") - - cache = ChunkKVCache(len(blocks), sink=int(sink), - window=None if window is None else int(window), - device=device, offload=bool(offload)) - return { - "model": diffusion, - "patcher": model, - "plan": plan, - "cache": cache, - "context": context, - "payload": payload, - "device": device, - "video_shape": tuple(video.shape), - "audio_shape": tuple(audio.shape), - "shift_video": float(getattr(diffusion, "sigma_shift_video", 12.0)), - "shift_audio": float(getattr(diffusion, "sigma_shift_audio", 3.0)), - }, "" - - -def run_session(session, *, sigmas, step_rule="consistency", eta=1.0, seed=0, - known_chunks=0, known_video=None, known_audio=None, - on_chunk=None, on_step=None, cancel=None, transformer_options=None): - """Drive one causal rollout over an assembled session. Returns (video, audio).""" - model, cache, plan = session["model"], session["cache"], session["plan"] - device, payload = session["device"], session["payload"] - options = dict(transformer_options or {}) - - # Again, once for the whole rollout: a session can be built and run much later, and - # another model may have evicted this one in between. The loop must not fight Comfy's - # offload decision chunk by chunk, so this is the only place it happens. - failure = _ensure_loaded(session.get("patcher"), session["video_shape"], {}) - if failure is not None: - print(f"[FunPack H3] causal lane could not re-place the model ({failure}); the " - f"rollout continues on whatever is already resident.") - - cache.clear() - with torch.no_grad(): - model.prefill_text(session["context"], plan, cache, - transformer_options=options, minimax_payload=payload) - - generator = torch.Generator(device="cpu").manual_seed(int(seed)) - - def draw_noise(shape): - return torch.randn(tuple(shape), generator=generator, dtype=torch.float32).to(device) - - video_noise = draw_noise(session["video_shape"]) - audio_noise = draw_noise(session["audio_shape"]) - shift_v, shift_a = session["shift_video"], session["shift_audio"] - - def forward(video_xt, audio_xt, index, sigma): - # The two streams do NOT share a sigma. H3 denoises audio on its own shifted grid - # (3 against the video's 12), and handing the audio rows the video's sigma is exactly - # the error h3_audio_clock exists to correct — the longer the step, the worse it is. - with torch.no_grad(): - return model.forward_chunk( - video_xt, audio_xt, plan, index, cache, - video_sigma=float(sigma), - audio_sigma=float(audio_sigma_for(sigma, shift_v, shift_a)), - context=session["context"], update_cache=False, - transformer_options=options, minimax_payload=payload, - sigma_shift_video=shift_v, sigma_shift_audio=shift_a) - - def commit(video_x0, audio_x0, index): - # One extra forward per chunk, at sigma 0, purely to write CLEAN K/V. Caching the K/V - # of the last denoising step instead would hand every later chunk a context that still - # carries that step's noise. - with torch.no_grad(): - model.forward_chunk( - video_x0, audio_x0, plan, index, cache, - video_sigma=0.0, audio_sigma=0.0, context=session["context"], - update_cache=True, transformer_options=options, minimax_payload=payload, - sigma_shift_video=shift_v, sigma_shift_audio=shift_a) - cache.finish_chunk(plan.cache_index(index)) - - return causal_rollout( - chunks=plan.bounds, sigmas=sigmas, forward=forward, commit=commit, - draw_noise=draw_noise, video_noise=video_noise, audio_noise=audio_noise, - step_rule=step_rule, eta=eta, known_chunks=known_chunks, - known_video=known_video, known_audio=known_audio, on_chunk=on_chunk, - on_step=on_step, cancel=cancel) diff --git a/loaders.py b/loaders.py index d098fcd9..a0027eba 100644 --- a/loaders.py +++ b/loaders.py @@ -166,18 +166,6 @@ class FunPackDiffusionModelLoader: "everything else."}), }, "optional": { - "chunk_causal": ("BOOLEAN", { - "default": False, - "tooltip": "MiniMax H3 only. Loads this model so it can generate in time " - "CHUNKS that each remember the ones before them, which is what " - "the Chain Sampler's 'Remembering chunks' mode needs. Nothing " - "about the weights changes and an ordinary generation is " - "unaffected — the model simply gains the ability to read a " - "key/value cache. It still needs a LoRA trained for that " - "pattern — load it with FunPack LoRA Loader like any other " - "LoRA: the base H3 weights have never seen a cache, so " - "without one the chunked lane runs out of distribution. Use " - "the full non-pruned checkpoint."}), "fp16_accumulation": ("BOOLEAN", { "default": False, "tooltip": "Accumulate fp16 matmuls in fp16. Faster on recent NVIDIA cards; " @@ -240,7 +228,7 @@ class FunPackDiffusionModelLoader: def load_model(self, model_name, weight_dtype, compute_dtype, attention, fp16_accumulation=False, sla_sparsity=None, sla_block_size=None, sla_protect_audio=None, sla_min_seq_len=None, sla_dense_last_steps=None, - sla=False, chunk_causal=False): + sla=False): notes = [f"FunPack Diffusion Model Loader | {model_name}"] accum = set_fp16_accumulation(fp16_accumulation) @@ -316,22 +304,6 @@ class FunPackDiffusionModelLoader: else: notes.append(f"attention: {attention} (dense calls)") - if chunk_causal: - try: - from .h3_causal import make_causal - except ImportError: - from h3_causal import make_causal - ok, causal_note = make_causal(model) - # Not fatal. The weights are fine and every other setting on this node applied; - # what is lost is one optional capability, and the sampler refuses the chunked - # mode on its own if the model cannot do it. - notes.append(f"chunk-causal: {causal_note}" if ok - else f"chunk-causal REQUESTED BUT NOT APPLIED: {causal_note}") - # ON THE CONSOLE, not only in `status`. The Editor's fixed graph does not show - # this node's status output anywhere, so a refusal here used to be invisible — - # and the sampler's own refusal then had to GUESS why, which it got wrong. - print(f"[FunPackDiffusionModelLoader] {model_name}: {notes[-1]}") - return (model, "\n".join(notes)) diff --git a/model_management.py b/model_management.py index e77b9de4..da208cc0 100644 --- a/model_management.py +++ b/model_management.py @@ -284,9 +284,9 @@ LORA_KEY_PREFIXES = ( "diffusion_model.", "model.diffusion_model.", "base_model.model.", - # PEFT adapters trained against a wrapper that names the DiT: RAVEN's streaming LoRA is - # `base_model.model.dit.`, so stripping only `base_model.model.` leaves `dit.` in - # front of every key and the whole file matches nothing. + # PEFT adapters trained against a wrapper that names the DiT: `base_model.model.dit. + # `, where stripping only `base_model.model.` leaves `dit.` in front of every key + # and the whole file matches nothing. "base_model.model.dit.", "base_model.model.diffusion_model.", "lora_model.", @@ -297,9 +297,9 @@ LORA_KEY_PREFIXES = ( #: Stripping is not enough for these: comfy's key map has no bare form, only #: `diffusion_model.`, so the head has to be REPLACED rather than removed. LORA_KEY_RENAMES = ( - # RAVEN's streaming LoRA for MiniMax H3: `base_model.model.dit.blocks.0.adaln_proj.linear` - # is comfy's `diffusion_model.blocks.0.adaln_proj.linear`, module for module. Stripping - # left `blocks.0...`, which matches nothing, and the file silently did nothing at all. + # Measured on a MiniMax H3 adapter: `base_model.model.dit.blocks.0.adaln_proj.linear` is + # comfy's `diffusion_model.blocks.0.adaln_proj.linear`, module for module. Stripping left + # `blocks.0...`, which matched nothing, and 532 keys silently did nothing at all. ("base_model.model.dit.", "diffusion_model."), ("base_model.model.", "diffusion_model."), ("dit.", "diffusion_model."), diff --git a/movie_editor/frontend/engine_defaults.test.js b/movie_editor/frontend/engine_defaults.test.js index ca78b050..99b8cd4b 100644 --- a/movie_editor/frontend/engine_defaults.test.js +++ b/movie_editor/frontend/engine_defaults.test.js @@ -88,6 +88,6 @@ test("a store without flushSave does not throw", () => { test("booleans clear when set back to their default", () => { const s = fakeStore(); - writeSamplerInput(s, { name: "h3_causal_chunks", default: false }, false, true); - assert.deepStrictEqual(s.calls[0], ["unset", "h3_causal_chunks"]); + writeSamplerInput(s, { name: "segmented_detailing", default: false }, false, true); + assert.deepStrictEqual(s.calls[0], ["unset", "segmented_detailing"]); }); diff --git a/movie_editor/frontend/engine_settings.js b/movie_editor/frontend/engine_settings.js index 629c016e..8323a694 100644 --- a/movie_editor/frontend/engine_settings.js +++ b/movie_editor/frontend/engine_settings.js @@ -331,19 +331,6 @@ { name: "h3_audio_clock", label: "H3 audio clock (few-step audio)", kind: "bool", default: false, hint: "MiniMax H3 only: removes audio distortion on few-step turbo schedules. Free.", detail: "H3 runs video and audio on different flow schedules but only one sigma grid reaches the sampler, so audio drifts further the bigger the step. This corrects it for one multiply per step. Exact with FunPack Distilled Flow, Hybrid Euler 2S and euler; leave it off with res_multistep / dpmpp_2m and the rest of the multistep family, which already absorb most of the error. No effect when both sigma shifts are equal, or on two-evals-per-step samplers." }, - { name: "h3_causal_chunks", label: "Remembering chunks (H3)", kind: "bool", default: false, - hint: "Generates the clip in time chunks that each remember the ones before them, instead of denoising the whole clip at once — so a long clip stays one continuous shot instead of separate scenes stitched together.", - detail: "Needs two things in Models: chunk_causal switched on for the diffusion model, and a chunk-causal LoRA (the RAVEN streaming one) added to the LoRA list. It is a LoRA on the ordinary H3 checkpoint, not a different model file, but it is not optional here: the chunked attention pattern is what it was trained to read. Without it this says so and the run continues normally. The wired sampler and CFG are not used on this path. There is no per-chunk redo: a chunk is committed the moment it finishes and the next is already built on it, so disliking one chunk means a new seed. EXPERIMENTAL." }, - { name: "h3_causal_step_rule", label: "Chunk step rule", kind: "combo", choices: ["consistency", "euler", "euler_ancestral"], default: "consistency", - dependsOn: "h3_causal_chunks", dependsValue: true, - 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, 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." }, { name: "h3_gain_mode", label: "H3 render gains", kind: "combo", choices: ["learned", "manual"], default: "learned", hint: "MiniMax H3 only. 'learned' takes the five render strengths from your ratings and ignores the dials below. 'manual' uses the dials as typed.", detail: "Learned needs the Refiner's 'Rating-learned render gains (H3)' left on — with it off, or before anything is rated, learned mode renders at the model's trained strengths and the dials still do nothing. Switch to manual to drive them yourself." }, @@ -564,7 +551,7 @@ 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", "dynashift", "dynashift_strength", "dynashift_threshold"], 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", "plateau_cache", "plateau_cache_threshold", "h3_audio_clock", "h3_causal_chunks", "h3_causal_step_rule", "h3_causal_sink", "h3_causal_window", "h3_gain_mode", "h3_gain_video", "h3_gain_prompt", "h3_gain_audio", "h3_prompt_scale", "h3_taste_bias", "segmented_detailing", "detail_targets", "detail_strength", "detail_threshold", "detail_max_area", "detail_mode", "detail_denoise", "mid_scene_guide", "mid_scene_guide_strength", "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", "identity_transfer_enabled", "source_id", "phase_scale", "id_strength", "arcface_mode", "debug_log"], + chain_experimental: ["context_windows", "context_window_length", "context_window_overlap", "context_window_schedule", "context_window_fuse", "context_window_freenoise", "context_window_retain_first", "plateau_cache", "plateau_cache_threshold", "h3_audio_clock", "h3_gain_mode", "h3_gain_video", "h3_gain_prompt", "h3_gain_audio", "h3_prompt_scale", "h3_taste_bias", "segmented_detailing", "detail_targets", "detail_strength", "detail_threshold", "detail_max_area", "detail_mode", "detail_denoise", "mid_scene_guide", "mid_scene_guide_strength", "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", "identity_transfer_enabled", "source_id", "phase_scale", "id_strength", "arcface_mode", "debug_log"], }; function countChainView(p, id, st) { diff --git a/movie_editor/frontend/pipeline_caps.js b/movie_editor/frontend/pipeline_caps.js index bcd569c5..4d83d3b8 100644 --- a/movie_editor/frontend/pipeline_caps.js +++ b/movie_editor/frontend/pipeline_caps.js @@ -171,10 +171,6 @@ // Left on against an LTX pipeline they are just as silently inert, so they get a chip // by the same rule — the user should not have to spend a generation to find out. // key -> why it cannot run off H3. - const H3_CAUSAL_REASON = - "The chunk cache is an H3 lane: it needs the chunk-causal DiT and the RAVEN LoRA, both " - + "of which are MiniMax H3 only. LTX has no equivalent."; - const H3_RENDER_GAIN_REASON = "The render gains scale H3's per-modality AdaLN gates and its token refiner. LTX has " + "neither: one modulation path for all streams, and the prompt reaches it through " @@ -187,10 +183,6 @@ // The render gains scale H3's per-modality AdaLN gates and its token refiner. LTX has // neither — one modulation path, cross-attention instead of a packed text span — so // there is nothing for these to reach on any other family. - h3_causal_chunks: H3_CAUSAL_REASON, - h3_causal_step_rule: H3_CAUSAL_REASON, - h3_causal_sink: H3_CAUSAL_REASON, - h3_causal_window: H3_CAUSAL_REASON, h3_gain_mode: H3_RENDER_GAIN_REASON, h3_gain_video: H3_RENDER_GAIN_REASON, h3_gain_prompt: H3_RENDER_GAIN_REASON, diff --git a/samplers.py b/samplers.py index ea136d12..ae34b626 100644 --- a/samplers.py +++ b/samplers.py @@ -2838,17 +2838,6 @@ def _normalize_video_denoised(denoised, video_mask, sigma, ref, strength, return denoised return denoised -# The causal lane's step rules, read from the module that implements them so the widget and -# the loop can never drift apart. Imported defensively: a broken optional module must not stop -# the node pack from registering. -try: - from .h3_causal import STEP_RULES as _CAUSAL_STEP_RULES -except ImportError: - try: - from h3_causal import STEP_RULES as _CAUSAL_STEP_RULES - except ImportError: - _CAUSAL_STEP_RULES = ("consistency", "euler", "euler_ancestral") - class FunPackLTXAVSceneChainSampler: @classmethod @@ -3178,22 +3167,6 @@ class FunPackLTXAVSceneChainSampler: "default": 1.0, "min": 0.0, "max": 2.0, "step": 0.05, "tooltip": "EXPERIMENTAL, MiniMax H3 only. Per-block write gain on the AUDIO rows (see h3_gain_video for the mechanism). This is the one modality whose gate can be moved without touching the other two at the point of writing — though the streams still share one attention pass, so a changed audio row is read by the video on the NEXT block. Below 1.0 the soundtrack is built more conservatively. Free. 1.0 = untouched.", }), - "h3_causal_chunks": ("BOOLEAN", { - "default": False, - "tooltip": "EXPERIMENTAL, MiniMax H3 only. Generates the clip in time CHUNKS that each remember the ones before them, instead of denoising the whole clip at once. A chunk is finished, written into a key/value cache as settled context, and the next chunk is generated attending to it — so a long clip stays one continuous shot with real memory of what was already drawn, rather than a chain of separate scenes stitched together. Two things are needed: switch on `chunk_causal` on FunPack Diffusion Model Loader, and load a chunk-causal LoRA (the RAVEN streaming one) with FunPack LoRA Loader like any other LoRA. That is a LoRA on the ordinary H3 checkpoint, not a different model file — but it is not optional here, because the chunked attention pattern is what it was trained to read and the base H3 weights have never seen a key/value cache. Without it this says so and the run continues normally. One consequence you cannot opt out of: there is no per-chunk redo — a chunk enters the cache the moment it finishes and the next is already built on it, so disliking chunk 3 of 12 means a new seed.", - }), - "h3_causal_step_rule": (list(_CAUSAL_STEP_RULES), { - "default": "consistency", - "tooltip": "EXPERIMENTAL, MiniMax H3 only. How each chunk's step reaches the next noise level, when h3_causal_chunks is on. 'consistency' is RAVEN's own transition and what its LoRA was distilled for — jump to the clean prediction and re-noise with fresh noise. 'euler' is the ordinary flow step this sampler uses everywhere else, so a chunked run can be compared against an unchunked one without the step rule being a second thing that changed. 'euler_ancestral' is euler with noise put back, in the rectified-flow form. Start on consistency; the other two are for finding out what the chunking alone is worth.", - }), - "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, 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, - "tooltip": "EXPERIMENTAL, MiniMax H3 only. How many of the MOST RECENT chunks stay in the cache besides the pinned ones. This is short-term memory: it carries motion and continuity from one moment to the next. 0 keeps only the pinned chunks, which makes the clip lean entirely on its opening. Higher is smoother and costs more per chunk.", - }), "h3_taste_bias": ("FLOAT", { "default": 0.0, "min": -0.30, "max": 0.30, "step": 0.01, "tooltip": "EXPERIMENTAL, MiniMax H3 only. Pushes the prompt toward what you have RATED WELL, in the model's own space. The Refiner already learns a 'liked direction' from your ratings and steers the conditioning along it — but that edit then goes through H3's token refiner, whose two attention blocks mix it across tokens and whose final norm scales its magnitude away, so what the model reads is not what was learned. This adds the same direction AFTER the refiner, where it lands as sent. Measured as a fraction of a typical prompt row's magnitude, so it means the same thing on every prompt and every checkpoint. Negative pushes AWAY from what you liked. Needs a refinement key with at least 3 liked runs; without one it does nothing and says so. Free. IGNORED unless h3_gain_mode is 'manual' — in 'learned' mode this is learned from your ratings like the other gains. 0.0 = off.", @@ -4022,16 +3995,11 @@ class FunPackLTXAVSceneChainSampler: model = self._install_h3_token_refiner(model, positive) try: - # The causal lane replaces the sampler entirely rather than wrapping it, so it is - # tried here and falls through (loudly) when it cannot run. - sampled = self._run_h3_causal(model, positive, samples, sigmas, seed, - pbar=pbar, step_offset=step_offset) - if sampled is None: - sampled = comfy.sample.sample_custom( - model, noise, float(cfg), sampler, sigmas, positive, negative, samples, - noise_mask=latent.get("noise_mask"), seed=int(seed), - callback=_progress_cb if pbar is not None else None, - ) + sampled = comfy.sample.sample_custom( + model, noise, float(cfg), sampler, sigmas, positive, negative, samples, + noise_mask=latent.get("noise_mask"), seed=int(seed), + callback=_progress_cb if pbar is not None else None, + ) finally: self._remove_bounded_attention(_ba_handles) self._assert_finite_sample(sampled) @@ -6750,8 +6718,6 @@ class FunPackLTXAVSceneChainSampler: #: what Studio tags onto entry 0 when it has learned gains for this key H3_GAINS_META = "funpack_h3_gains" - #: how a chunk reaches its next sigma when the causal lane is running - H3_CAUSAL_STEP_RULES = _CAUSAL_STEP_RULES #: and the learned taste direction that rides with them H3_TASTE_DIR_META = "funpack_h3_taste_dir" #: the value of each render gain that means "untouched". Not all of them are 1.0: @@ -6873,96 +6839,6 @@ class FunPackLTXAVSceneChainSampler: return None return (projected.float() / norm) * float(strength) - def _run_h3_causal(self, model, positive, samples, sigmas, seed, - pbar=None, step_offset=0): - """Generate the clip in remembering chunks instead of denoising it all at once. - - Returns the finished latent, or None to fall through to ordinary sampling — every - reason for falling through is printed, because a silently ignored mode looks exactly - like a mode that ran and did nothing. - - What this deliberately does NOT use: the wired SAMPLER and the CFG. The chunk loop is - not a k-diffusion sampler and cannot host one, and H3 generates at CFG 1 regardless. - Both are stated on the console rather than dropped quietly. - """ - if not getattr(self, "_h3_causal_chunks", False): - return None - try: - from . import h3_causal as causal - from . import minimax_h3 as h3mod - except ImportError: - import h3_causal as causal - import minimax_h3 as h3mod - if not h3mod.is_h3_model(model): - _log.feature("FunPackSceneChain", "Chunk cache", False, - "not a MiniMax H3 model. The chunk cache is an H3 lane; this run " - "samples normally.") - return None - session, reason = causal.build_session( - model, positive, {"samples": samples}, - sink=int(getattr(self, "_h3_causal_sink", 2)), - window=int(getattr(self, "_h3_causal_window", 2)), - ) - if session is None: - _log.feature("FunPackSceneChain", "Chunk cache", False, - f"{reason} This run samples normally.") - return None - schedule = [float(v) for v in sigmas] - _log.feature( - "FunPackSceneChain", "Chunk cache", True, - f"{session['plan'].n_chunks} chunks, {len(schedule) - 1} steps each, " - f"{getattr(self, '_h3_causal_step_rule', 'consistency')}, sink " - f"{getattr(self, '_h3_causal_sink', 2)} / window " - f"{getattr(self, '_h3_causal_window', 2)}. EXPERIMENTAL, and it needs a " - f"chunk-causal LoRA loaded with FunPack LoRA Loader — the chunked attention " - f"pattern is what that LoRA was trained to read, and the base H3 weights have " - f"never seen a key/value cache. The wired sampler and CFG are NOT used on this " - f"path: the chunk loop is not a k-diffusion sampler, and H3 generates at CFG 1 " - f"anyway. There is no per-chunk redo.") - # This lane never reaches comfy's sampler, so nothing else ticks the progress bar or - # checks for a cancel. A long generation showing no progress is indistinguishable - # from a hung one, and an interrupt that does nothing is worse. - base_label = self._current_phase_label() - scene_steps = max(1, len(schedule) - 1) - - def _on_step(index, done, total): - self._set_phase(f"{base_label}{' · ' if base_label else ''}" - f"chunk {int(index) + 1}/{session['plan'].n_chunks}") - if pbar is not None and total: - # Mapped onto the budget this scene was allotted, not counted raw: the bar was - # sized in scenes x steps, and a chunked scene runs that many steps PER CHUNK. - pbar.update_absolute(int(step_offset) - + int(round(scene_steps * float(done) / float(total)))) - - def _cancel(_index): - comfy.model_management.throw_exception_if_processing_interrupted() - - try: - video, audio = causal.run_session( - session, sigmas=schedule, - step_rule=str(getattr(self, "_h3_causal_step_rule", "consistency")), - seed=int(seed), on_step=_on_step, cancel=_cancel, - ) - except Exception as error: # noqa: BLE001 - _log.failed("FunPackSceneChain", "H3 causal rollout", error, - "this scene is NOT generated — nothing is silently substituted") - raise - finally: - self._set_phase(base_label) - return self._pack_h3_av(samples, video, audio) - - @staticmethod - def _pack_h3_av(template, video, audio): - """Put the two streams back into the nested container the rest of the chain expects. - - `template` is the incoming latent, used only for its dtype: the rollout works in fp32 - and handing the decoder a different dtype than it was given is a difference nothing - downstream would report. - """ - import comfy.nested_tensor - dtype = getattr(template, "dtype", None) or torch.float32 - return comfy.nested_tensor.NestedTensor((video.to(dtype), audio.to(dtype))) - def _install_h3_token_refiner(self, model, positive): """Edit the token refiner's OUTPUT: prompt loudness, and the learned taste push. @@ -7351,8 +7227,6 @@ class FunPackLTXAVSceneChainSampler: identity_transfer_enabled=False, identity_projector="None", source_id=2.0, phase_scale=1.0, id_strength=1.0, arcface_mode="auto_adjust", debug_log=False, carry_overlap_through_anchor=False, - h3_causal_chunks=False, h3_causal_step_rule="consistency", - h3_causal_sink=2, h3_causal_window=2, plateau_cache=False, plateau_cache_threshold=0.975, taste_nearest_prompt=False, segmented_detailing=False, detail_targets="hands", @@ -7388,10 +7262,6 @@ class FunPackLTXAVSceneChainSampler: self._h3_gain_audio = max(0.0, min(2.0, float(h3_gain_audio))) self._h3_prompt_scale = max(0.0, min(2.0, float(h3_prompt_scale))) self._h3_taste_bias = max(-0.30, min(0.30, float(h3_taste_bias))) - self._h3_causal_chunks = bool(h3_causal_chunks) - self._h3_causal_step_rule = str(h3_causal_step_rule) - self._h3_causal_sink = max(1, int(h3_causal_sink)) - self._h3_causal_window = max(0, int(h3_causal_window)) self._h3_gain_mode = str(h3_gain_mode or "learned").strip().lower() # The gate every rating-driven wrapper shares. On H3 it is read off the schedule's # own base grid; on LTX it stays the absolute-sigma gate it was validated with. diff --git a/tests/_h3_causal_upstream_check.py b/tests/_h3_causal_upstream_check.py deleted file mode 100644 index 533db06e..00000000 --- a/tests/_h3_causal_upstream_check.py +++ /dev/null @@ -1,250 +0,0 @@ -"""Chunk plan and modulation, checked against the REAL ComfyUI MiniMax H3. - -Run in a clean subprocess by `test_h3_causal_plan.py`, because the suite stubs `comfy` and -the real package cannot be imported alongside the stub. This is the test that catches a -ComfyUI update moving the packed layout out from under us. -""" -import sys - -TEXT_LEN, LATENT_T, LAT_H, LAT_W, AUDIO_T = 120, 57, 32, 48, 320 - - -def main(funpack, comfy_root): - sys.path.insert(0, comfy_root) - sys.path.insert(0, funpack) - import torch - from comfy.ldm.minimax.model import (VISUAL_COND_TIMESTEP, PackedLayout, _frame_grid, - pack_audio) - - import h3_causal as hc - - frame_rows = _frame_grid(LAT_H, LAT_W)[0].shape[0] - - def plan_for(**kw): - layout = PackedLayout(TEXT_LEN, LATENT_T, LAT_H, LAT_W, AUDIO_T, **kw) - return hc.build_plan(layout, LATENT_T, AUDIO_T) - - def tiles(plan): - rows = torch.cat([plan.prefix_rows] + [plan.chunk(i)[0] for i in range(plan.n_chunks)]) - return sorted(rows.tolist()) == list(range(plan.layout.seq_len)) - - # 192 pixel frames is 57 video latents: eleven 17-frame chunks and a 5-frame tail. - plan = plan_for() - assert plan.n_chunks == 12, plan.n_chunks - assert tiles(plan), "plain clip does not tile" - assert [k for _, _, k in plan.prefix_runs] == ["text"] - - # a chunk's audio rows are two spans, and `pack_audio` on the chunk's own slice produces - # them in exactly that order — which is what lets forward_chunk pack the slice directly - rows, (audio_run, video_run) = plan.chunk(2) - audio = rows[audio_run[0]:audio_run[1]].tolist() - half = len(audio) // 2 - assert audio[half] - audio[0] == AUDIO_T, "the two stereo spans are not one clip apart" - v_start, v_stop, a_start, a_stop = plan.bounds[2] - packed = pack_audio(torch.zeros(1, 32, 2, a_stop - a_start)) - assert packed.shape[0] == len(audio), (packed.shape, len(audio)) - assert video_run[1] - video_run[0] == (v_stop - v_start) * frame_rows - - # conditioning rows keep the layout upstream packed for them - refs = plan_for(refs=[{"kind": "image", "latent_h": LAT_H, "latent_w": LAT_W}]) - assert [k for _, _, k in refs.prefix_runs] == ["text", "ref_img"] - assert tiles(refs), "ref clip does not tile" - for i in range(refs.n_chunks): - assert {k for _, _, k in refs.chunk(i)[1]} == {"audio", "video"} - - pins = plan_for(keyframes=[{"resolved_frame_index": 0}], frame_count=192) - assert [k for _, _, k in pins.prefix_runs] == ["text", "cond"] - assert tiles(pins), "pinned clip does not tile" - - # modulation: only the kinds present get a timestep row, and identical times share one - model_cls = hc._causal_classes()[2] - seg_t, t_row, unique = model_cls._modulation(None, {"audio", "video"}, 0.8, 0.6, {}) - assert len(unique) == 2 and t_row[seg_t["video"]] != t_row[seg_t["audio"]] - _, t_row, unique = model_cls._modulation(None, {"audio", "video"}, 0.5, 0.5, {}) - assert len(unique) == 1 and len(t_row) == 1 - seg_t, _, _ = model_cls._modulation(None, {"video", "cond"}, 0.1, 0.1, {}) - assert abs(seg_t["cond"] - VISUAL_COND_TIMESTEP) < 1e-9 - - # the text span splits by tag runs; everything else is one run per kind - seg_t, t_row, _ = model_cls._modulation(None, {"text"}, 0.8, 0.6, {}) - out = model_cls._mod_segments(None, [(0, 6, "text")], seg_t, t_row, - torch.tensor([1, 1, 0, 0, 0, 1])) - assert [(a, b, row % 3) for a, b, row in out] == [(0, 2, 1), (2, 5, 0), (5, 6, 1)] - - # `_mod_scale_shift` writes each run in place, so a gap would leave rows unmodulated - rows, runs = plan.chunk(1) - seg_t, t_row, _ = model_cls._modulation(None, {k for _, _, k in runs}, 0.8, 0.6, {}) - cursor = 0 - for a, b, _row in model_cls._mod_segments(None, runs, seg_t, t_row, None): - assert a == cursor - cursor = b - assert cursor == rows.shape[0] - - _forward_phase() - print("OK") - - -def _install_rope_fallback(): - """A pure-torch stand-in for the fused RMSNorm+rope kernel, for hosts whose comfy_kitchen - build lacks it. Both lanes use the same one, so the dense-vs-causal comparison stays fair. - """ - import torch - - import comfy.quant_ops - - if hasattr(comfy.quant_ops.ck, "rms_rope_split_half_"): - return - - def rms_rope(q, k, table, qw, kw, epsilon=1e-5, rot_dim=None): - half = table.shape[-3] - rot = rot_dim or half * 2 - cos = table[..., 0, 0].squeeze(0).squeeze(-2)[:, None, :] - sin = table[..., 1, 0].squeeze(0).squeeze(-2)[:, None, :] - for tensor, weight in ((q, qw), (k, kw)): - x = tensor[0] - scale = torch.rsqrt(x.float().pow(2).mean(-1, keepdim=True) + epsilon) - x.copy_(x * scale.to(x.dtype) * weight) - lo, hi = x[..., :rot // 2].clone(), x[..., rot // 2:rot].clone() - x[..., :rot // 2] = lo * cos - hi * sin - x[..., rot // 2:rot] = lo * sin + hi * cos - return q, k - - comfy.quant_ops.ck.rms_rope_split_half_ = rms_rope - - -def _forward_phase(): - """The real DiT, at toy width, driven through both lanes.""" - import copy - - import torch - - import comfy.ops - from comfy.ldm.minimax.model import MiniMaxH3Model, PackedLayout - - import h3_causal as hc - - _install_rope_fallback() - text_len, latent_t, lat_h, lat_w, audio_t = 7, 12, 8, 8, 65 - torch.manual_seed(0) - stock = MiniMaxH3Model( - hidden_size=192, num_layers=2, token_refiner_num_layers=1, num_attention_heads=2, - attention_head_dim=96, ffn_hidden_size=192, text_dim=64, timestep_input_dim=16, - time_embed_hidden_size=192, time_embed_dim=32, rope_inv_freq_len=16, - dtype=torch.float32, device=torch.device("cpu"), - operations=comfy.ops.disable_weight_init) - for _name, param in stock.named_parameters(): - param.data.normal_(0, 0.02) - stock.rope.inv_freq.copy_(1.0 / (10000 ** (torch.arange(16).float() / 16))) - stock.eval() - - causal = copy.deepcopy(stock) - keys = set(causal.state_dict().keys()) - ok, note = hc.make_causal(type("Patcher", (), {"model": type("M", (), {})()})()) - assert not ok, note # nothing to re-class on a bare object - attention_cls, block_cls, model_cls = hc._causal_classes() - for i in range(len(causal.blocks)): - hc._to_causal_block(causal.blocks[i], block_cls, attention_cls) - causal.__class__ = model_cls - # Re-classing must not rename a single weight: wrapping used to, and unpatching crashed. - assert set(causal.state_dict().keys()) == keys, "re-classing renamed state-dict keys" - - video = torch.randn(1, 24, latent_t, lat_h, lat_w) - audio = torch.randn(1, 32, 2, audio_t) - context = torch.randn(1, text_len, 64) - layout = PackedLayout(text_len, latent_t, lat_h, lat_w, audio_t) - payload = {"layout": layout, "seed": 0} - - # THE safety property: with no cache passed, the causal model is the stock one. - with torch.no_grad(): - a = stock._forward([video, audio], torch.tensor([800.0]), context, - minimax_payload=payload) - b = causal._forward([video, audio], torch.tensor([800.0]), context, - minimax_payload=payload) - assert torch.equal(a[0], b[0]) and torch.equal(a[1], b[1]), "the dense lane changed" - - plan = hc.build_plan(layout, latent_t, audio_t) - cache = hc.ChunkKVCache(len(causal.blocks), sink=2, window=2, - device=torch.device("cpu"), offload=False) - with torch.no_grad(): - causal.prefill_text(context, plan, cache, minimax_payload=payload) - assert cache._store[(0, 0)][0].shape[-2] == text_len, "the prefix cached the wrong rows" - for i in range(plan.n_chunks): - v_start, v_stop, a_start, a_stop = plan.bounds[i] - chunk_v = video[:, :, v_start:v_stop] - chunk_a = audio[..., a_start:a_stop] - out_v, out_a = causal.forward_chunk( - chunk_v, chunk_a, plan, i, cache, video_sigma=0.8, audio_sigma=0.4, - context=context, minimax_payload=payload) - assert out_v.shape == chunk_v.shape, (out_v.shape, chunk_v.shape) - assert out_a.shape == chunk_a.shape, (out_a.shape, chunk_a.shape) - assert torch.isfinite(out_v).all() and torch.isfinite(out_a).all() - causal.forward_chunk(chunk_v, chunk_a, plan, i, cache, video_sigma=0.0, - audio_sigma=0.0, context=context, update_cache=True, - minimax_payload=payload) - cache.finish_chunk(plan.cache_index(i)) - 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. - from comfy.ldm.minimax.model import time_shift_slope as _slope - del sys.modules["comfy.ldm.minimax.model"].time_shift_slope - try: - old_attn, old_block, older = hc._causal_classes() - assert hasattr(older, "forward_chunk"), "the lane refused itself without the slope" - stripped = copy.deepcopy(stock) - for i in range(len(stripped.blocks)): - hc._to_causal_block(stripped.blocks[i], old_block, old_attn) - stripped.__class__ = older - older_cache = hc.ChunkKVCache(len(stripped.blocks), sink=2, window=2, - device=torch.device("cpu"), offload=False) - with torch.no_grad(): - stripped.prefill_text(context, plan, older_cache, minimax_payload=payload) - v_start, v_stop, a_start, a_stop = plan.bounds[0] - out_v, out_a = stripped.forward_chunk( - video[:, :, v_start:v_stop], audio[..., a_start:a_stop], plan, 0, older_cache, - video_sigma=0.8, audio_sigma=0.4, context=context, minimax_payload=payload) - assert torch.isfinite(out_v).all() and torch.isfinite(out_a).all() - finally: - sys.modules["comfy.ldm.minimax.model"].time_shift_slope = _slope - - # a clip whose streams disagree about its length is refused by name, not deep in a block - try: - hc.build_plan(PackedLayout(text_len, latent_t, lat_h, lat_w, 20), latent_t, 20) - except hc.CacheError as error: - assert "does not cover the picture" in str(error), error - else: - raise AssertionError("a silent chunk was not refused") - - -if __name__ == "__main__": - main(sys.argv[1], sys.argv[2]) diff --git a/tests/test_causal_rollout.py b/tests/test_causal_rollout.py deleted file mode 100644 index 6af64d30..00000000 --- a/tests/test_causal_rollout.py +++ /dev/null @@ -1,504 +0,0 @@ -"""Chunk-causal rollout: the cut, the step rules, and the order the loop does things in. - -Chunk-causal generation is not a sampler but a way of shaping the sequence — cut the clip into -time chunks, carry each to completion, commit it into a KV cache as clean context, generate the -next one attending to that cache. The model gets real memory of what it already drew. - -That shaping is separable from the 4-step consistency sampler the technique is usually shipped -with, and these tests pin the separation: the schedule is ours, and the step rule is a knob so -a chunked run can be compared against a dense one without two variables moving. - -Everything the model does arrives as a callable, so all of this runs with no ComfyUI and no -weights. The wiring to a real model is in `test_h3_causal_session.py`. -""" -import sys -from pathlib import Path - -import pytest -import torch - -sys.path.insert(0, str(Path(__file__).resolve().parents[1])) - -import h3_causal as rc - - -# ── the chunk cut ─────────────────────────────────────────────────────────── - -def test_video_cuts_every_five_latents(): - bounds = rc.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_the_tail_chunk_is_shorter_not_dropped(): - """H3's grid is 5k+2 latents, so there is ALWAYS a 2-latent tail. Dropping it would cut - the last five frames off every clip.""" - bounds = rc.chunk_bounds(latent_t=17, audio_t=140) - assert bounds[-1][0], bounds[-1][1] == (15, 17) - assert bounds[-1][1] - bounds[-1][0] == 2 - - -def test_every_video_latent_lands_in_exactly_one_chunk(): - bounds = rc.chunk_bounds(latent_t=22, audio_t=180) - covered = [] - for start, stop, _a, _b in bounds: - covered.extend(range(start, stop)) - assert covered == list(range(22)) - - -def test_every_audio_latent_lands_in_exactly_one_chunk(): - bounds = rc.chunk_bounds(latent_t=22, audio_t=180) - covered = [] - for _s, _e, start, stop in bounds: - covered.extend(range(start, stop)) - assert covered == list(range(180)) - - -def test_audio_boundaries_come_from_the_clock_not_a_cadence(): - """40 audio latents/s against a fixed 24 fps, 17 frames per chunk. The 28/29 alternation - is a CONSEQUENCE of that ratio — hard-coding it would silently drift on a changed grid.""" - bounds = rc.chunk_bounds(latent_t=17, audio_t=int(round((17 * 3 + 5) / 24 * 40))) - assert bounds[0][3] == round(17 * 40 / 24) - assert bounds[1][3] == round(34 * 40 / 24) - - -def test_audio_boundaries_never_go_backwards(): - for latent_t in (2, 7, 12, 22, 57): - audio_t = int(round((17 * ((latent_t - 2) // 5) + 5) / 24 * 40)) - bounds = rc.chunk_bounds(latent_t, audio_t) - cursors = [b[2] for b in bounds] + [bounds[-1][3]] - assert cursors == sorted(cursors) - - -def test_a_single_chunk_clip_is_one_chunk(): - assert len(rc.chunk_bounds(latent_t=2, audio_t=8)) == 1 - - -def test_an_empty_clip_plans_nothing(): - assert rc.chunk_bounds(latent_t=0, audio_t=0) == [] - - -# ── the step rules ────────────────────────────────────────────────────────── - -def test_consistency_jumps_to_x0_and_renoises_with_fresh_noise(): - """The distilled transition a chunk-causal LoRA is trained for.""" - x0 = torch.full((4,), 2.0) - noise = torch.full((4,), 1.0) - out = rc.step("consistency", torch.zeros(4), x0, sigma=0.8, sigma_next=0.5, noise=noise) - assert torch.allclose(out, 0.5 * x0 + 0.5 * noise) - - -def test_consistency_at_the_last_step_is_exactly_x0(): - x0 = torch.full((4,), 3.0) - out = rc.step("consistency", torch.ones(4), x0, sigma=0.3, sigma_next=0.0, - noise=torch.full((4,), 9.0)) - assert torch.allclose(out, x0) - - -def test_consistency_ignores_the_noise_already_in_x_t(): - """Fresh noise every step is the point — reusing x_t's noise would correlate the steps.""" - x0 = torch.full((4,), 2.0) - a = rc.step("consistency", torch.zeros(4), x0, 0.8, 0.5, torch.full((4,), 1.0)) - b = rc.step("consistency", torch.full((4,), 99.0), x0, 0.8, 0.5, torch.full((4,), 1.0)) - assert torch.allclose(a, b) - - -def test_euler_is_the_ordinary_flow_step(): - x_t = torch.full((4,), 1.0) - x0 = torch.full((4,), 0.0) - out = rc.step("euler", x_t, x0, sigma=1.0, sigma_next=0.5, noise=torch.zeros(4)) - # derivative = (x_t - x0)/sigma = 1; x + (0.5 - 1.0) * 1 = 0.5 - assert torch.allclose(out, torch.full((4,), 0.5)) - - -def test_euler_reaches_x0_exactly_at_sigma_zero(): - x_t = torch.full((4,), 1.0) - x0 = torch.full((4,), 0.25) - out = rc.step("euler", x_t, x0, sigma=0.7, sigma_next=0.0, noise=torch.zeros(4)) - assert torch.allclose(out, x0, atol=1e-6) - - -def test_euler_ignores_the_noise_it_is_handed(): - x_t = torch.full((4,), 1.0) - x0 = torch.zeros(4) - a = rc.step("euler", x_t, x0, 1.0, 0.5, torch.zeros(4)) - b = rc.step("euler", x_t, x0, 1.0, 0.5, torch.full((4,), 50.0)) - assert torch.allclose(a, b) - - -def test_ancestral_with_zero_eta_is_euler(): - x_t, x0 = torch.full((4,), 1.0), torch.zeros(4) - a = rc.step("euler_ancestral", x_t, x0, 1.0, 0.5, torch.full((4,), 3.0), eta=0.0) - b = rc.step("euler", x_t, x0, 1.0, 0.5, torch.zeros(4)) - assert torch.allclose(a, b) - - -def test_ancestral_adds_noise_when_eta_is_positive(): - x_t, x0 = torch.full((4,), 1.0), torch.zeros(4) - a = rc.step("euler_ancestral", x_t, x0, 1.0, 0.5, torch.full((4,), 3.0), eta=1.0) - b = rc.step("euler", x_t, x0, 1.0, 0.5, torch.zeros(4)) - assert not torch.allclose(a, b) - - -def test_ancestral_uses_the_rectified_flow_form_not_the_vp_one(): - """The VP formula is wrong on a flow model — it was a real bug here before. At eta=1 the - step goes all the way down to sigma_down = sigma_next**2/sigma, not to sigma_next.""" - x_t, x0 = torch.full((4,), 1.0), torch.zeros(4) - sigma, sigma_next, eta = 1.0, 0.5, 1.0 - out = rc.step("euler_ancestral", x_t, x0, sigma, sigma_next, torch.zeros(4), eta=eta) - down_ratio = 1.0 + (sigma_next / sigma - 1.0) * eta - expected = x_t + (sigma_next * down_ratio - sigma) * ((x_t - x0) / sigma) - assert torch.allclose(out, expected) - - -def test_ancestral_at_the_last_step_does_not_renoise(): - """Noise added at sigma_next == 0 would land in the output.""" - x_t, x0 = torch.full((4,), 1.0), torch.full((4,), 0.25) - out = rc.step("euler_ancestral", x_t, x0, 0.7, 0.0, torch.full((4,), 99.0), eta=1.0) - assert torch.allclose(out, x0, atol=1e-6) - - -def test_an_unknown_rule_is_refused_not_guessed(): - with pytest.raises(ValueError, match="unknown step rule"): - rc.step("dpmpp_2m", torch.zeros(4), torch.zeros(4), 1.0, 0.5, torch.zeros(4)) - - -def test_every_advertised_rule_actually_runs(): - for rule in rc.STEP_RULES: - out = rc.step(rule, torch.ones(4), torch.zeros(4), 1.0, 0.5, torch.zeros(4)) - assert out.shape == (4,) - - -# ── the audio stream's own sigma ──────────────────────────────────────────── - -def test_the_audio_remap_matches_comfys_own_time_shift(): - """Pinned against upstream's formula so a drift in ours is caught here rather than as - quiet audio distortion. - - Transcribed rather than imported: the suite stubs `comfy`, so importing the real module - is not possible here. This is `time_shift_sigma` from comfy/ldm/minimax/model.py verbatim - — invert to the base grid, re-apply the other shift. Verified equal to the live function - across this range when run against the real checkout. - """ - def upstream(sigma, from_shift, to_shift): - base = sigma / (from_shift + sigma * (1.0 - from_shift)) - return to_shift * base / (1.0 + (to_shift - 1.0) * base) - - for sigma in (1.0, 0.94, 0.83, 0.72, 0.55, 0.30, 0.10, 0.01): - assert abs(rc.audio_sigma_for(sigma, 12.0, 3.0) - - upstream(sigma, 12.0, 3.0)) < 1e-9 - - -def test_equal_shifts_leave_the_sigma_alone(): - assert rc.audio_sigma_for(0.7, 12.0, 12.0) == 0.7 - - -def test_sigma_zero_stays_zero(): - """A remapped zero would leave the last step short of clean.""" - assert rc.audio_sigma_for(0.0, 12.0, 3.0) == 0.0 - - -# ── the rollout ───────────────────────────────────────────────────────────── -# -# Everything the model does arrives as a callable, so the loop runs here with no ComfyUI and no weights. What is being tested is the ORDER of things: which chunks are -# sampled, which are taken as given, what gets committed to the cache, and that a seeded run -# draws its noise in a fixed sequence. - -class Recorder: - def __init__(self, latent_t=12, audio_t=100, channels=2): - self.chunks = rc.chunk_bounds(latent_t, audio_t) - self.video_noise = torch.zeros(1, channels, latent_t, 2, 2) - self.audio_noise = torch.zeros(1, channels, 2, audio_t) - self.forwards = [] - self.commits = [] - self.draws = 0 - - def forward(self, video_xt, audio_xt, index, sigma): - self.forwards.append((index, round(float(sigma), 6))) - return torch.zeros_like(video_xt), torch.zeros_like(audio_xt) - - def commit(self, video_x0, audio_x0, index): - self.commits.append((index, tuple(video_x0.shape), tuple(audio_x0.shape))) - - def draw(self, shape): - self.draws += 1 - return torch.zeros(shape) - - def run(self, sigmas=(1.0, 0.5, 0.0), **kw): - return rc.causal_rollout( - chunks=self.chunks, sigmas=sigmas, forward=self.forward, commit=self.commit, - draw_noise=self.draw, video_noise=self.video_noise, - audio_noise=self.audio_noise, **kw) - - -def test_every_chunk_is_committed_to_the_cache(): - """A chunk that is not committed is a chunk the rest of the clip cannot see — which is - the whole reason to chunk in the first place.""" - r = Recorder() - r.run() - assert [c[0] for c in r.commits] == list(range(len(r.chunks))) - - -def test_a_chunk_is_committed_once_not_once_per_step(): - r = Recorder() - r.run(sigmas=(1.0, 0.8, 0.5, 0.2, 0.0)) - assert len(r.commits) == len(r.chunks) - - -def test_chunks_are_generated_in_order(): - """Chunk i+1 attends to chunk i's cache, so out-of-order would read an empty cache.""" - r = Recorder() - r.run() - assert [i for i, _s in r.forwards] == sorted(i for i, _s in r.forwards) - - -def test_the_schedule_is_ours_not_a_fixed_four(): - """Any scheduler, any step count — the distilled 4 is a default, not a constraint.""" - r = Recorder() - r.run(sigmas=(1.0, 0.94, 0.83, 0.72, 0.55, 0.30, 0.10, 0.0)) - per_chunk = [s for i, s in r.forwards if i == 0] - assert per_chunk == [1.0, 0.94, 0.83, 0.72, 0.55, 0.30, 0.10] - - -def test_every_step_sees_its_own_sigma(): - r = Recorder() - r.run(sigmas=(1.0, 0.5, 0.0)) - assert [s for i, s in r.forwards if i == 1] == [1.0, 0.5] - - -def test_noise_is_drawn_for_both_streams_every_step(): - """Even at sigma_next == 0, where it is multiplied away: skipping the draw would shift - the stream for every later chunk and a seeded run would stop reproducing.""" - r = Recorder() - r.run(sigmas=(1.0, 0.5, 0.0)) - assert r.draws == len(r.chunks) * 2 * 2 - - -def test_the_output_covers_the_whole_clip(): - r = Recorder() - video, audio = r.run() - assert video.shape == r.video_noise.shape - assert audio.shape == r.audio_noise.shape - - -def test_the_initial_noise_is_not_mutated(): - r = Recorder() - before = r.video_noise.clone() - r.run() - assert torch.allclose(r.video_noise, before) - - -# ── the pre-known first chunk: how an i2v anchor survives ─────────────────── - -def test_a_known_chunk_is_not_resampled(): - """Chunk 0 comes from FunPack's ordinary dense path, where pins and references work. The - causal lane is never asked to model conditioning rows it has no layout for.""" - r = Recorder() - r.run(known_chunks=1, - known_video=torch.full_like(r.video_noise, 7.0), - known_audio=torch.full_like(r.audio_noise, 7.0)) - assert 0 not in [i for i, _s in r.forwards] - - -def test_a_known_chunk_is_still_committed(): - """The anchor chunk must reach the cache or the rest of the clip continues from a clip - that appears to start after it.""" - r = Recorder() - r.run(known_chunks=1, - known_video=torch.full_like(r.video_noise, 7.0), - known_audio=torch.full_like(r.audio_noise, 7.0)) - assert r.commits[0][0] == 0 - - -def test_the_known_chunk_reaches_the_output_verbatim(): - r = Recorder() - known_v = torch.full_like(r.video_noise, 7.0) - video, _audio = r.run(known_chunks=1, known_video=known_v, - known_audio=torch.full_like(r.audio_noise, 7.0)) - v_start, v_stop = r.chunks[0][0], r.chunks[0][1] - assert torch.allclose(video[:, :, v_start:v_stop], known_v[:, :, v_start:v_stop]) - - -def test_later_chunks_are_still_sampled_after_a_known_one(): - r = Recorder() - r.run(known_chunks=1, - known_video=torch.full_like(r.video_noise, 7.0), - known_audio=torch.full_like(r.audio_noise, 7.0)) - assert sorted({i for i, _s in r.forwards}) == list(range(1, len(r.chunks))) - - -def test_known_chunks_without_latents_is_refused(): - """Silently sampling over the anchor would look like the anchor was ignored.""" - r = Recorder() - with pytest.raises(ValueError, match="known_chunks"): - r.run(known_chunks=1) - - -# ── refusals and cancellation ─────────────────────────────────────────────── - -def test_an_unknown_step_rule_is_refused_before_any_forward(): - r = Recorder() - with pytest.raises(ValueError, match="unknown step rule"): - r.run(step_rule="dpmpp_3m") - assert r.forwards == [] - - -def test_a_schedule_with_no_steps_is_refused(): - r = Recorder() - with pytest.raises(ValueError, match="at least one step"): - r.run(sigmas=(1.0,)) - - -def test_cancellation_is_checked_before_each_chunk_and_each_step(): - """Per chunk is not often enough: a chunk is a full denoise, so an interrupt could sit - unanswered for every step of it while the user watches nothing happen.""" - r = Recorder() - seen = [] - - def cancel(index): - seen.append(index) - if len(seen) == 3: - raise KeyboardInterrupt - with pytest.raises(KeyboardInterrupt): - r.run(cancel=cancel) - assert seen == [0, 0, 0] # chunk 0, then once per step inside it - - -def test_every_model_call_reports_progress(): - """This lane never reaches comfy's sampler, so nothing else ticks a progress bar — and a - long generation showing none is indistinguishable from a hung one.""" - r = Recorder() - ticks = [] - r.run(sigmas=(1.0, 0.5, 0.0), on_step=lambda i, done, total: ticks.append((i, done, total))) - steps = 2 - assert len(ticks) == len(r.chunks) * steps - assert [t[1] for t in ticks] == list(range(1, len(ticks) + 1)) - assert {t[2] for t in ticks} == {len(r.chunks) * steps} # a stable denominator - - -def test_progress_counts_the_whole_clip_not_the_chunk(): - """A bar that restarts at every chunk says nothing about how far along the clip is.""" - r = Recorder() - ticks = [] - r.run(sigmas=(1.0, 0.0), on_step=lambda i, done, total: ticks.append(done)) - assert ticks == list(range(1, len(r.chunks) + 1)) - - -def test_the_chunk_callback_reports_whether_it_was_generated(): - """A preview lane needs to know which chunks are new; a pre-known one has nothing to show.""" - r = Recorder() - events = [] - r.run(known_chunks=1, - known_video=torch.full_like(r.video_noise, 7.0), - known_audio=torch.full_like(r.audio_noise, 7.0), - on_chunk=lambda i, v, a, was_known: events.append((i, was_known))) - assert events[0] == (0, True) - assert all(not known for _i, known in events[1:]) - - -class _FakePlan: - n_chunks = 3 - - -# ── the Chain Sampler dispatch ────────────────────────────────────────────── -# -# The causal lane REPLACES the sampler rather than wrapping it, so the thing that matters is -# that it falls through — loudly — whenever it cannot run. A mode that is silently ignored -# looks exactly like a mode that ran and did nothing. - -def _sampler(**attrs): - import samplers - node = samplers.FunPackLTXAVSceneChainSampler.__new__( - samplers.FunPackLTXAVSceneChainSampler) - for k, v in attrs.items(): - setattr(node, k, v) - return node - - -def test_the_mode_is_off_by_default(): - import samplers - spec = samplers.FunPackLTXAVSceneChainSampler.INPUT_TYPES() - fields = {**spec.get("required", {}), **spec.get("optional", {})} - assert fields["h3_causal_chunks"][1]["default"] is False - - -def test_the_widget_offers_exactly_the_rules_the_loop_implements(): - """A widget listing a rule the loop refuses would be a crash three minutes into a run.""" - import samplers - spec = samplers.FunPackLTXAVSceneChainSampler.INPUT_TYPES() - fields = {**spec.get("required", {}), **spec.get("optional", {})} - assert tuple(fields["h3_causal_step_rule"][0]) == rc.STEP_RULES - - -def test_off_means_the_normal_sampler_runs(): - node = _sampler(_h3_causal_chunks=False) - assert node._run_h3_causal(object(), [], object(), [1.0, 0.0], 0) is None - - -def test_a_non_h3_model_falls_through_and_says_so(monkeypatch, capsys): - import samplers - monkeypatch.setattr(samplers._log, "note_on_change", - lambda tag, who, msg: print(msg), raising=False) - import minimax_h3 - monkeypatch.setattr(minimax_h3, "is_h3_model", lambda m: False) - node = _sampler(_h3_causal_chunks=True) - assert node._run_h3_causal(object(), [], object(), [1.0, 0.0], 0) is None - assert "Chunk cache: Inactive | Not a MiniMax H3 model" in capsys.readouterr().out - - -def test_an_unbuildable_causal_lane_falls_through_and_says_why(monkeypatch, capsys): - import samplers - monkeypatch.setattr(samplers._log, "note_on_change", - lambda tag, who, msg: print(msg), raising=False) - import minimax_h3 - monkeypatch.setattr(minimax_h3, "is_h3_model", lambda m: True) - monkeypatch.setattr(rc, "build_session", - lambda *a, **k: (None, "this model is a stock bidirectional H3.")) - node = _sampler(_h3_causal_chunks=True, _h3_causal_sink=2, _h3_causal_window=2) - assert node._run_h3_causal(object(), [], object(), [1.0, 0.0], 0) is None - out = capsys.readouterr().out - assert "Chunk cache: Inactive" in out and "stock bidirectional" in out - - -def test_a_rollout_failure_is_raised_not_swallowed(monkeypatch): - """Falling back to ordinary sampling after a partial rollout would return a clip the user - did not ask for, under a mode they think ran.""" - import samplers - monkeypatch.setattr(samplers._log, "note_on_change", lambda *a, **k: None, raising=False) - monkeypatch.setattr(samplers._log, "failed", lambda *a, **k: None, raising=False) - import minimax_h3 - monkeypatch.setattr(minimax_h3, "is_h3_model", lambda m: True) - monkeypatch.setattr(rc, "build_session", - lambda *a, **k: ({"plan": _FakePlan()}, "")) - monkeypatch.setattr(rc, "run_session", - lambda *a, **k: (_ for _ in ()).throw(RuntimeError("cache overflow"))) - node = _sampler(_h3_causal_chunks=True, _h3_causal_sink=2, _h3_causal_window=2, - _h3_causal_step_rule="consistency") - with pytest.raises(RuntimeError, match="cache overflow"): - node._run_h3_causal(object(), [], object(), [1.0, 0.0], 0) - - -def test_the_run_says_the_sampler_and_cfg_are_not_used(monkeypatch, capsys): - """Both are wired on the node and ignored on this path. Dropping them quietly would be - the kind of silent no-op this project logs against.""" - import samplers - monkeypatch.setattr(samplers._log, "note_on_change", - lambda tag, who, msg: print(msg), raising=False) - import minimax_h3 - monkeypatch.setattr(minimax_h3, "is_h3_model", lambda m: True) - # the suite's comfy stub has NestedTensor as a bare object; the packer needs to construct one - import comfy.nested_tensor - monkeypatch.setattr(comfy.nested_tensor, "NestedTensor", lambda pair: tuple(pair), - raising=False) - monkeypatch.setattr(rc, "build_session", lambda *a, **k: ({"plan": _FakePlan()}, "")) - monkeypatch.setattr(rc, "run_session", - lambda *a, **k: (torch.zeros(1, 24, 12, 4, 4), torch.zeros(1, 32, 2, 100))) - node = _sampler(_h3_causal_chunks=True, _h3_causal_sink=2, _h3_causal_window=2, - _h3_causal_step_rule="consistency") - node._run_h3_causal(object(), [], torch.zeros(1), [1.0, 0.5, 0.0], 0) - assert "Chunk cache: Active" in capsys.readouterr().out - # The console line is one clause by design; the caveats live in the full text, which - # FUNPACK_LOG=verbose prints. What must not happen is them being dropped from the source. - import inspect - src = inspect.getsource(samplers.FunPackLTXAVSceneChainSampler._run_h3_causal) - assert "CFG" in src and "NOT used" in src - assert "no per-chunk redo" in src.lower() diff --git a/tests/test_funpack_log.py b/tests/test_funpack_log.py index 39a0f674..74f25fe0 100644 --- a/tests/test_funpack_log.py +++ b/tests/test_funpack_log.py @@ -177,7 +177,7 @@ def test_verbose_prints_the_whole_thing(capsys): def test_an_em_dash_is_not_a_cut(capsys): """It introduces the reason as often as the essay, and a renderer that guesses wrong eats the one part of the line worth reading.""" - L.feature("FunPack", "Chunk cache", False, "not a MiniMax H3 model — an H3-only lane") + L.feature("FunPack", "Region lock", False, "not a MiniMax H3 model — an H3-only lane") assert "an H3-only lane" in capsys.readouterr().out @@ -191,10 +191,10 @@ def test_a_feature_repeats_only_when_its_state_changes(capsys): """A standing condition is news once. Twenty identical lines are what makes a log stop being read at all.""" for _ in range(5): - L.feature("FunPack", "Chunk cache", False, "not a MiniMax H3 model") - assert capsys.readouterr().out.count("Chunk cache") == 1 - L.feature("FunPack", "Chunk cache", True, "22 chunks") - assert "Chunk cache: Active" in capsys.readouterr().out + L.feature("FunPack", "Region lock", False, "not a MiniMax H3 model") + assert capsys.readouterr().out.count("Region lock") == 1 + L.feature("FunPack", "Region lock", True, "384 of 1024 patches") + assert "Region lock: Active" in capsys.readouterr().out def test_a_skip_and_an_inactive_feature_read_the_same(capsys): diff --git a/tests/test_h3_causal.py b/tests/test_h3_causal.py deleted file mode 100644 index febff3a7..00000000 --- a/tests/test_h3_causal.py +++ /dev/null @@ -1,210 +0,0 @@ -"""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) diff --git a/tests/test_h3_causal_loader.py b/tests/test_h3_causal_loader.py deleted file mode 100644 index 169d9dca..00000000 --- a/tests/test_h3_causal_loader.py +++ /dev/null @@ -1,208 +0,0 @@ -"""Making a model chunk-causal is a switch on OUR loader, and the LoRA is an ordinary LoRA. - -FunPack owns the loading pipeline, so neither half of this needs a third-party node or -package. The diffusion loader RE-CLASSES an already-loaded H3 in place rather than building -it through a different class — no parameter is added, no key is renamed, nothing is copied, -so ComfyUI's memory ledger and offload accounting are untouched. The adapter then goes -through FunPack Apply LoRA Weights like any other LoRA. -""" -import sys -from pathlib import Path - -import pytest - -sys.path.insert(0, str(Path(__file__).resolve().parents[1])) - -import h3_causal as hc -import loaders -import model_management as mm - - -def _fields(): - spec = loaders.FunPackDiffusionModelLoader.INPUT_TYPES() - return {**spec.get("required", {}), **spec.get("optional", {})} - - -def test_the_switch_is_on_our_own_loader(): - assert "chunk_causal" in _fields() - - -def test_it_is_off_by_default(): - assert _fields()["chunk_causal"][1]["default"] is False - - -def test_the_tooltip_says_a_lora_is_still_needed(): - """Without one the chunked lane runs out of distribution, which reads as a quality problem - with no visible cause.""" - tip = _fields()["chunk_causal"][1]["tooltip"] - assert "LoRA" in tip and "out of distribution" in tip - - -def test_no_raven_package_is_named_anywhere_in_the_loader(): - """The whole point of owning it: no third-party install to discover.""" - import inspect - src = inspect.getsource(loaders.FunPackDiffusionModelLoader) - assert "raven_streaming" not in src and "RAVEN Model Loader" not in src - - -# ── the LoRA is an ordinary LoRA ──────────────────────────────────────────── - -def test_the_peft_dit_prefix_is_one_of_the_candidates(): - """RAVEN's keys are base_model.model.dit.. Stripping only `base_model.model.` leaves - `dit.` in front of every key and the whole file silently matches nothing.""" - assert "base_model.model.dit." in mm.LORA_KEY_PREFIXES - - -def test_the_longer_prefix_is_tried_before_the_shorter_one_can_win(): - """Both match; the loader scores every candidate and keeps the best, so the order only has - to include the longer one — but it must be there at all.""" - assert "base_model.model." in mm.LORA_KEY_PREFIXES - prefixes = list(mm.LORA_KEY_PREFIXES) - assert prefixes.index("base_model.model.dit.") > prefixes.index("base_model.model.") - - -def test_stripping_that_prefix_yields_plain_module_paths(): - stripped = mm._rekey( - {"base_model.model.dit.blocks.0.attn.qkv_proj.lora_A.weight": 1, - "base_model.model.dit.blocks.0.attn.qkv_proj.lora_B.weight": 2}, - "base_model.model.dit.") - assert set(stripped) == {"blocks.0.attn.qkv_proj.lora_A.weight", - "blocks.0.attn.qkv_proj.lora_B.weight"} - - -# ── and why stripping was not enough ──────────────────────────────────────── -# -# Measured on the rental: the RAVEN LoRA matched NOTHING through every strip. Its keys are -# comfy's module paths exactly, under a different head — `base_model.model.dit.` for -# comfy's `diffusion_model.` — and comfy's key map has no bare form to land on, so -# stripping produced `blocks.0...` and matched none of it. The file did nothing at all. - -def test_the_raven_head_is_renamed_to_comfys_not_removed(): - assert ("base_model.model.dit.", "diffusion_model.") in mm.LORA_KEY_RENAMES - - -def test_renaming_that_head_yields_the_keys_comfys_map_holds(): - renamed = mm._rekey( - {"base_model.model.dit.audio_patch_proj.lora_A.weight": 1, - "base_model.model.dit.blocks.0.adaln_proj.linear.lora_B.weight": 2}, - "base_model.model.dit.", "diffusion_model.") - assert set(renamed) == {"diffusion_model.audio_patch_proj.lora_A.weight", - "diffusion_model.blocks.0.adaln_proj.linear.lora_B.weight"} - - -def test_a_rename_wins_over_a_strip_when_it_matches_more(monkeypatch): - """Both are scored against the real key map and the best is kept, so a rename that lands - beats a strip that does not without either being ordered ahead of the other.""" - key_map = {"diffusion_model.audio_patch_proj": "audio_patch_proj.weight"} - monkeypatch.setattr(mm, "_count_lora_matches", - lambda lora, km: sum(1 for k in lora - if k.rsplit(".lora_", 1)[0] in km)) - lora = {"base_model.model.dit.audio_patch_proj.lora_A.weight": 1} - best, note = mm._best_lora_keying(lora, key_map) - assert note == "base_model.model.dit. -> diffusion_model." - assert "diffusion_model.audio_patch_proj.lora_A.weight" in best - - -# ── re-classing an already-loaded model ───────────────────────────────────── - -class _FakeBlock: - def __init__(self): - self.attn = type("A", (), {})() - - -def test_a_model_without_blocks_declines_with_a_reason(): - fake = type("M", (), {"model": type("N", (), {"diffusion_model": type("D", (), {})()})()})() - ok, note = hc.make_causal(fake) - assert ok is False and "no DiT blocks" in note - - -def test_a_model_without_a_diffusion_model_declines(): - ok, note = hc.make_causal(type("M", (), {"model": None})()) - assert ok is False and "no diffusion model" in note - - -def test_it_is_idempotent(monkeypatch): - """Loading the same model twice, or re-running the node, must not stack anything.""" - diffusion = type("D", (), {})() - diffusion.blocks = [_FakeBlock()] - diffusion._funpack_causal = True - fake = type("M", (), {})() - fake.model = type("N", (), {})() - fake.model.diffusion_model = diffusion - ok, note = hc.make_causal(fake) - assert ok is True and "already" in note - - -def test_a_comfy_without_h3_declines_instead_of_raising(monkeypatch): - """A ComfyUI too old for H3 must not turn one optional switch into a dead node.""" - monkeypatch.setattr(hc, "_causal_classes", - lambda: (_ for _ in ()).throw(ImportError("no minimax"))) - diffusion = type("D", (), {})() - diffusion.blocks = [_FakeBlock()] - fake = type("M", (), {})() - fake.model = type("N", (), {})() - fake.model.diffusion_model = diffusion - ok, note = hc.make_causal(fake) - assert ok is False and "no MiniMax H3" in note - - -# ── which node the user is actually sent to ───────────────────────────────── -# -# FunPack LoRA Loader loads AND applies, from its own file list. FunPack Apply LoRA Weights is -# the optional stack producer for prompt-specific trained strengths. Naming the second one as -# the way in sends the user to add a node they do not need. - -def test_the_loader_message_names_the_lora_loader(): - tip = _fields()["chunk_causal"][1]["tooltip"] - assert "FunPack LoRA Loader" in tip - assert "Apply LoRA Weights" not in tip - - -def test_the_sampler_message_names_the_lora_loader(): - import inspect - import samplers - src = inspect.getsource(samplers.FunPackLTXAVSceneChainSampler._run_h3_causal) - assert "FunPack LoRA Loader" in src - assert "Apply LoRA Weights" not in src - - -def test_the_lora_loader_does_not_advertise_itself_as_stack_only(): - """Its description used to read as though a stack node were required.""" - import model_management - text = model_management.FunPackLoraLoader.DESCRIPTION - assert "only needed" in text or "own list" in text - - -def test_the_stack_input_is_optional(): - """If it were required, 'the loader applies them on its own' would be false.""" - import model_management - spec = model_management.FunPackLoraLoader.INPUT_TYPES() - assert "lora_stack" in spec.get("optional", {}) - assert "lora_stack" not in spec.get("required", {}) - - -def test_the_file_list_is_the_required_way_in(): - import model_management - spec = model_management.FunPackLoraLoader.INPUT_TYPES() - assert "lora_list" in spec.get("required", {}) - - -def test_no_surface_still_points_at_the_removed_raven_lora_widget(): - """It was a combo on the diffusion loader and is now a boolean plus an ordinary LoRA. - A tooltip naming the old widget sends the user hunting for a setting that is not there.""" - root = Path(__file__).resolve().parents[1] - for name in ("samplers.py", "loaders.py", - "movie_editor/frontend/engine_settings.js", - "movie_editor/frontend/pipeline_caps.js"): - text = (root / name).read_text(encoding="utf-8", errors="replace") - assert "raven_lora" not in text, name - - -def test_the_lora_stack_says_what_it_matched_on_the_console(): - """`keys=N fmt=...` only ever went into the node's `status` output, which the Editor's - fixed graph shows nowhere — so a LoRA that matched nothing looked exactly like one that - applied, and the chunk-causal lane had no way to tell the user which it was.""" - import inspect - src = inspect.getsource(mm.FunPackLoraLoader.load_loras) - assert '_log.feature(' in src - assert 'no LoRA is selected' in src diff --git a/tests/test_h3_causal_plan.py b/tests/test_h3_causal_plan.py deleted file mode 100644 index c5da36f0..00000000 --- a/tests/test_h3_causal_plan.py +++ /dev/null @@ -1,181 +0,0 @@ -"""The chunk plan: which packed rows each chunk owns. - -The plan is the join between our rollout and upstream's packed sequence, so what these tests -guard is agreement rather than any behaviour of our own. A plan that tiles the sequence wrongly -does not raise — it hands the model rows from the wrong moment of the clip, and the only -symptom is a video that is subtly incoherent. - -The suite stubs `comfy`, so the plan is exercised here against a fake layout with upstream's -segment structure, and against the REAL PackedLayout in a clean subprocess at the bottom. -""" -import os -import subprocess -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 # noqa: E402 - -TEXT_LEN, LATENT_T, AUDIO_T, FRAME_ROWS = 120, 57, 320, 6 - - -class _Layout: - """Upstream's segment structure, without upstream. Order is the contract: the prefix - first, then the target audio, then the target video.""" - - def __init__(self, prefix=(("text", TEXT_LEN),)): - segments, row = [], 0 - for kind, n in prefix: - segments.append((row, row + n, kind)) - row += n - segments.append((row, row + AUDIO_T * 2, "audio")) - row += AUDIO_T * 2 - segments.append((row, row + LATENT_T * FRAME_ROWS, "video")) - self.seq_len = row + LATENT_T * FRAME_ROWS - self.segments = segments - - -def _plan(prefix=(("text", TEXT_LEN),)): - return hc.build_plan(_Layout(prefix), LATENT_T, AUDIO_T) - - -def _all_rows(plan): - return torch.cat([plan.prefix_rows] + [plan.chunk(i)[0] for i in range(plan.n_chunks)]) - - -# ── the cut ───────────────────────────────────────────────────────────────── - -def test_a_real_clip_cuts_into_the_expected_number_of_chunks(): - """192 pixel frames is 57 video latents: eleven 17-frame chunks and a 5-frame tail.""" - assert _plan().n_chunks == 12 - - -def test_the_chunks_and_the_prefix_tile_the_whole_sequence_exactly(): - """No row generated twice, none missed. An overlap would denoise a moment twice from two - different contexts; a gap would leave noise in the finished clip.""" - plan = _plan() - assert sorted(_all_rows(plan).tolist()) == list(range(plan.layout.seq_len)) - - -def test_the_prefix_is_everything_before_the_target_streams(): - plan = _plan() - assert [k for _, _, k in plan.prefix_runs] == ["text"] - assert plan.prefix_rows.tolist() == list(range(TEXT_LEN)) - - -def test_a_chunk_packs_its_audio_before_its_video(): - """Upstream packs the target audio segment before the target video one, and the final - layer reads both back by contiguous slice.""" - assert [k for _, _, k in _plan().chunk(3)[1]] == ["audio", "video"] - - -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 — so a chunk owns a slice of each, not one contiguous run.""" - rows, (audio_run, _) = _plan().chunk(2) - audio = rows[audio_run[0]:audio_run[1]].tolist() - half = len(audio) // 2 - left, right = audio[:half], audio[half:] - assert right[0] - left[0] == AUDIO_T - assert [r - left[0] for r in left] == [r - right[0] for r in right] - - -def test_the_video_rows_of_a_chunk_are_its_own_latent_frames(): - plan = _plan() - v_start, v_stop, _, _ = plan.bounds[4] - _, (_, video_run) = plan.chunk(4) - assert video_run[1] - video_run[0] == (v_stop - v_start) * FRAME_ROWS - - -def test_the_tail_chunk_owns_whatever_audio_is_left(): - """The clip's audio must be covered exactly — a soundtrack that stops short of the - picture is the failure this rules out.""" - plan = _plan() - assert plan.bounds[-1][3] == AUDIO_T - - -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 ──────────────────────────────── - -def test_reference_rows_land_in_the_prefix_not_in_a_media_chunk(): - """This is what makes r2v work here: the reference keeps the rows upstream packed for it - instead of being re-packed into a chunk that has no place for them.""" - plan = _plan((("text", TEXT_LEN), ("ref_img", 90))) - assert [k for _, _, k in plan.prefix_runs] == ["text", "ref_img"] - for i in range(plan.n_chunks): - assert {k for _, _, k in plan.chunk(i)[1]} == {"audio", "video"} - - -def test_the_tiling_still_holds_with_conditioning_rows_present(): - plan = _plan((("text", TEXT_LEN), ("cond", 6), ("ref_audio", 40))) - assert sorted(_all_rows(plan).tolist()) == list(range(plan.layout.seq_len)) - - -def test_a_clip_whose_streams_disagree_about_its_length_is_refused(): - """A chunk with no audio rows is a StopIteration fifty blocks deep, because the final - layer reads one contiguous span of each stream. Say what is actually wrong instead.""" - layout = _Layout() - with pytest.raises(hc.CacheError, match="does not cover the picture"): - hc.ChunkPlan(layout, hc.chunk_bounds(LATENT_T, 12), 12) - - -def test_a_layout_with_no_target_streams_is_refused(): - class Bare: - segments = [(0, 10, "text")] - - with pytest.raises(hc.CacheError): - hc.ChunkPlan(Bare(), [(0, 1, 0, 1)], 1) - - -# ── against the real ComfyUI, if one is here ──────────────────────────────── - -REAL_COMFY = Path.home() / "Documents" / "ComfyUI" - - -@pytest.mark.skipif(not (REAL_COMFY / "comfy" / "ldm" / "minimax" / "model.py").exists(), - reason="no local ComfyUI with MiniMax H3 to check against") -def test_the_plan_and_both_lanes_agree_with_the_real_model(): - """The one that catches a ComfyUI update moving the layout out from under us. It also - drives the REAL DiT at toy width through both lanes, which is where the dense-lane safety - property is actually checked: with no cache passed, the causal model IS the stock one.""" - root = Path(__file__).resolve().parents[1] - out = subprocess.run( - [sys.executable, str(Path(__file__).parent / "_h3_causal_upstream_check.py"), - str(root), str(REAL_COMFY)], - capture_output=True, text=True, cwd=str(REAL_COMFY), - env=dict(os.environ, PYTHONPATH="")) - assert "OK" in out.stdout, out.stderr[-3000:] diff --git a/tests/test_h3_causal_session.py b/tests/test_h3_causal_session.py deleted file mode 100644 index 5c9bfcde..00000000 --- a/tests/test_h3_causal_session.py +++ /dev/null @@ -1,403 +0,0 @@ -"""Assembling and driving a causal run: what it refuses, and what order it does things in. - -The rollout arithmetic is tested in `test_causal_rollout.py`. What is tested here is the -wiring around it — which cache index each chunk writes to, that the prefix is cached exactly -once, and that a chunk's context is committed CLEAN rather than at the last noisy step. -""" -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 # noqa: E402 - -TEXT_LEN, LATENT_T, AUDIO_T, FRAME_ROWS = 40, 12, 64, 6 - - -class _Layout: - def __init__(self): - self.segments = [(0, TEXT_LEN, "text"), - (TEXT_LEN, TEXT_LEN + AUDIO_T * 2, "audio"), - (TEXT_LEN + AUDIO_T * 2, - TEXT_LEN + AUDIO_T * 2 + LATENT_T * FRAME_ROWS, "video")] - self.seq_len = self.segments[-1][1] - self.position_ids = torch.zeros(self.seq_len, 3, dtype=torch.float64) - - -class _Diffusion: - """Records what the rollout asked of it, and nothing else.""" - - sigma_shift_video, sigma_shift_audio = 12.0, 3.0 - - def __init__(self, layers=2): - self.blocks = [object()] * layers - self.calls = [] - - def prefill_text(self, context, plan, cache, **kw): - self.calls.append(("prefill", kw.get("video_sigma", 0.0))) - cache.finish_chunk(0) - - def forward_chunk(self, video, audio, plan, index, cache, *, video_sigma, audio_sigma, - update_cache=False, **kw): - self.calls.append(("commit" if update_cache else "forward", index, - round(float(video_sigma), 6), round(float(audio_sigma), 6))) - if update_cache: - for layer in range(len(self.blocks)): - cache.write(layer, plan.cache_index(index), torch.zeros(1, 1, 1, 1), - torch.zeros(1, 1, 1, 1)) - return torch.zeros_like(video), torch.zeros_like(audio) - - -class _Samples: - def __init__(self, video, audio): - self._parts = (video, audio) - - def unbind(self): - return self._parts - - -class _Patcher: - load_device = torch.device("cpu") - - def __init__(self, diffusion, payload=None, fail=None): - self.model = _Base(diffusion, payload, fail) - self.model.diffusion_model = diffusion - - -class _Base: - def __init__(self, diffusion, payload, fail): - self.diffusion_model = diffusion - self._payload = payload - self._fail = fail - - def extra_conds(self, **kwargs): - if self._fail: - raise RuntimeError(self._fail) - payload = dict(self._payload or {}) - payload.setdefault("layout", _Layout()) - return {"minimax_payload": _Cond(payload), - "c_crossattn": _Cond(torch.zeros(1, TEXT_LEN, 8))} - - -class _Cond: - def __init__(self, cond): - self.cond = cond - - -def _latent(t=LATENT_T, batch=1): - return {"samples": _Samples(torch.zeros(batch, 24, t, 8, 8), - torch.zeros(1, 32, 2, AUDIO_T))} - - -def _positive(): - return [[torch.zeros(1, TEXT_LEN, 8), {}]] - - -def _build(**kw): - diffusion = kw.pop("diffusion", None) or _Diffusion() - return hc.build_session(_Patcher(diffusion, kw.pop("payload", None), kw.pop("fail", None)), - _positive(), kw.pop("latent", None) or _latent(), **kw) - - -# ── what it refuses, and how it says so ───────────────────────────────────── - -def test_a_stock_model_is_made_causal_rather_than_refused(monkeypatch): - """Re-classing costs nothing and changes no weight, so a mode that is already switched ON - must not refuse itself over a switch on another node. It does the work instead.""" - class Stock: - blocks = [object()] - - made = [] - monkeypatch.setattr(hc, "make_causal", - lambda model: (made.append(model) or (True, "re-classed"))) - session, reason = _build(diffusion=Stock()) - assert made and session is not None, reason - - -def test_a_model_that_cannot_be_made_causal_says_what_it_actually_is(): - """The old message asserted the user had not switched something on, which was a guess — - and a wrong one sends them looking in the wrong place.""" - class NotH3: - blocks = [object()] - - session, reason = _build(diffusion=NotH3()) - assert session is None - assert "NotH3" in reason - - -def test_a_model_with_no_blocks_is_refused(): - class Empty: - blocks = [] - - def forward_chunk(self): - pass - - session, reason = _build(diffusion=Empty()) - assert session is None and "no DiT blocks" in reason - - -def test_a_batch_larger_than_one_is_refused(): - session, reason = _build(latent=_latent(batch=2)) - assert session is None and "one clip at a time" in reason - - -def test_a_latent_that_is_not_a_video_audio_pair_is_refused(): - session, reason = _build(latent={"samples": torch.zeros(1, 24, 4, 8, 8)}) - assert session is None and "causal lane can take" in reason - - -def test_a_clip_too_short_to_chunk_is_refused_rather_than_run(): - """One chunk means there is nothing for a cache to remember — running it would pay the - lane's overhead for exactly none of its benefit.""" - session, reason = _build(latent=_latent(t=4)) - assert session is None and "nothing for a chunk cache to remember" in reason - - -def test_a_conditioning_the_model_cannot_prepare_is_reported_not_raised(): - session, reason = _build(fail="no keyframes here") - assert session is None and "no keyframes here" in reason - - -def test_a_run_with_no_packed_layout_is_refused(): - session, reason = _build(payload={"layout": None}) - assert session is None and "no packed layout" in reason - - -def test_a_buildable_run_reports_no_reason(): - session, reason = _build() - assert session is not None and reason == "" - assert session["plan"].n_chunks == 3 - assert session["cache"].num_layers == 2 - - -def test_the_cache_takes_the_sink_and_window_it_is_given(): - session, _ = _build(sink=1, window=None) - assert session["cache"].sink == 1 and session["cache"].window is None - - -# ── the order things happen in ────────────────────────────────────────────── - -def _run(**kw): - diffusion = _Diffusion() - session, reason = _build(diffusion=diffusion) - assert session is not None, reason - kw.setdefault("sigmas", [1.0, 0.5, 0.0]) - video, audio = hc.run_session(session, **kw) - return diffusion, session, video, audio - - -def test_the_prefix_is_cached_once_for_the_whole_rollout(): - """Recomputing it per chunk would quietly make every later chunk's assumption false: that - the text keys it attends are the ones the model saw.""" - diffusion, _, _, _ = _run() - assert [c for c in diffusion.calls if c[0] == "prefill"] == [("prefill", 0.0)] - - -def test_every_chunk_is_forwarded_once_per_step_and_committed_once(): - diffusion, session, _, _ = _run() - chunks = session["plan"].n_chunks - forwards = [c for c in diffusion.calls if c[0] == "forward"] - commits = [c for c in diffusion.calls if c[0] == "commit"] - assert len(forwards) == chunks * 2 # two steps in the schedule - assert [c[1] for c in commits] == list(range(chunks)) - - -def test_a_chunk_is_committed_clean_not_at_its_last_noisy_step(): - """Caching the K/V of the last denoising step would hand every later chunk a context that - still carries that step's noise.""" - diffusion, _, _, _ = _run() - assert all(c[2] == 0.0 and c[3] == 0.0 for c in diffusion.calls if c[0] == "commit") - - -def test_the_audio_rows_are_given_their_own_sigma(): - """H3 denoises audio on its own shifted grid. Handing it the video's sigma is exactly the - error `h3_audio_clock` exists to correct.""" - diffusion, _, _, _ = _run(sigmas=[0.8, 0.0]) - forward = next(c for c in diffusion.calls if c[0] == "forward") - assert forward[3] != forward[2] - - -def test_the_commit_of_one_chunk_precedes_the_first_forward_of_the_next(): - """A chunk that is not committed before the next one starts is a chunk the rest of the - clip cannot see.""" - diffusion, _, _, _ = _run() - order = [c for c in diffusion.calls if c[0] in ("forward", "commit")] - for position, call in enumerate(order): - if call[0] == "forward" and call[1] > 0: - assert ("commit", call[1] - 1, 0.0, 0.0) in order[:position] - - -def test_each_chunk_writes_to_its_own_cache_index_above_the_prefix(): - """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 written <= set(range(2, session["plan"].n_chunks + 2)) - - -def test_the_rollout_returns_a_clip_of_the_shape_it_was_given(): - _, session, video, audio = _run() - assert tuple(video.shape) == session["video_shape"] - assert tuple(audio.shape) == session["audio_shape"] - - -def test_a_seeded_run_reproduces_exactly(): - """The noise is drawn in a fixed order so a repeated seed gives the same clip; a skipped - draw anywhere would shift the stream for every later chunk.""" - first = _run(seed=7)[2] - second = _run(seed=7)[2] - assert torch.equal(first, second) - - -def test_a_different_seed_gives_a_different_clip(): - assert not torch.equal(_run(seed=7)[2], _run(seed=8)[2]) - - -def test_the_cache_starts_empty_even_if_the_session_is_reused(): - """A second scene through the same session must not inherit the first one's memory.""" - diffusion = _Diffusion() - session, _ = _build(diffusion=diffusion) - hc.run_session(session, sigmas=[1.0, 0.0]) - before = len(session["cache"]._store) - hc.run_session(session, sigmas=[1.0, 0.0]) - assert len(session["cache"]._store) == before - - -def test_an_unknown_step_rule_is_refused_rather_than_silently_defaulted(): - with pytest.raises(ValueError): - _run(step_rule="whatever") - - -# ── loading, and what loading also installs ───────────────────────────────── -# -# The ordinary path reaches the model through comfy.sample.sample_custom, whose -# prepare_sampling calls load_models_gpu. That call is ALSO what applies a patcher's object -# patches (partially_load -> patch_model), which is how FunPack's AdaLN modality gains and -# token-refiner edit get installed at all. This lane calls the DiT directly, so skipping the -# load would leave the weights offloaded AND both of those silently absent. - -@pytest.fixture(autouse=True) -def helpers(monkeypatch): - """A live comfy.sampler_helpers whose prepare_sampling is swappable PER TEST. - - Installed once and patched with monkeypatch afterwards: `install_module` merges onto a - module object the whole session shares, so a test that installs a raising stub and does - not restore it poisons every test that runs after it. - """ - from tests._comfy_stubs import install_module - - module = install_module( - "comfy.sampler_helpers", - prepare_sampling=lambda *a, **kw: None, - convert_cond=lambda cond: [dict(c[1], cross_attn=c[0]) for c in cond]) - - def swap(fn): - monkeypatch.setattr(module, "prepare_sampling", fn) - - monkeypatch.setattr(module, "prepare_sampling", lambda *a, **kw: None) - return swap - - -@pytest.fixture -def loads(helpers): - """Records every prepare_sampling call, the way ComfyUI's own would be made.""" - calls = [] - helpers(lambda *a, **kw: calls.append(a)) - return calls - - -def test_the_model_is_loaded_before_the_conditioning_is_prepared(helpers): - """extra_conds runs H3's token refiner — a partial forward. Preparing it before the load - is what produced `mat1 is on cuda:0, different from other tensors on cpu`.""" - order = [] - - class Watched(_Base): - def extra_conds(self, **kwargs): - order.append("conds") - return super().extra_conds(**kwargs) - - diffusion = _Diffusion() - patcher = _Patcher(diffusion) - patcher.model = Watched(diffusion, None, None) - patcher.model.diffusion_model = diffusion - helpers(lambda *a, **kw: order.append("load")) - session, reason = hc.build_session(patcher, _positive(), _latent()) - assert session is not None, reason - assert order[:2] == ["load", "conds"] - - -def test_the_load_is_given_the_clips_own_shape(loads): - """prepare_sampling estimates memory from the noise shape; handing it the wrong one gets - a placement decision made for a clip that is not this one.""" - _build() - assert loads and loads[0][1] == (1, 24, LATENT_T, 8, 8) - - -def test_the_conditioning_is_converted_before_the_memory_estimate(loads): - """estimate_memory calls extra_conds_shapes(**cond), so it wants ComfyUI's INTERNAL cond - dicts — handing it the [tensor, meta] pairs a node receives raises, the load silently does - not happen, and the failure resurfaces as a device mismatch three lines later.""" - _build() - assert isinstance(loads[0][2]["positive"][0], dict) - - -def test_the_load_happens_before_the_first_forward(helpers): - order = [] - diffusion = _Diffusion() - session, reason = _build(diffusion=diffusion) - assert session is not None, reason - helpers(lambda *a, **kw: order.append("load")) - real = diffusion.forward_chunk - - def spy(*a, **kw): - order.append("forward") - return real(*a, **kw) - - diffusion.forward_chunk = spy - hc.run_session(session, sigmas=[1.0, 0.0]) - assert order[0] == "load" - - -def test_the_model_is_loaded_once_per_rollout_not_per_chunk(loads): - """Loading per chunk would fight Comfy's own offload decision all the way down the clip.""" - session, reason = _build() - assert session is not None, reason - loads.clear() - hc.run_session(session, sigmas=[1.0, 0.5, 0.0]) - assert len(loads) == 1 - - -def _boom(*a, **kw): - raise RuntimeError("out of memory") - - -def test_a_load_that_fails_at_build_time_refuses_the_lane(helpers): - """Everything after the load assumes the model is resident, so a warning here surfaces - three lines later as a device mismatch — one bug reported twice.""" - helpers(_boom) - session, reason = _build() - assert session is None - assert "could not be placed" in reason and "out of memory" in reason - - -def test_a_load_that_fails_mid_rollout_does_not_abort_it(helpers, capsys): - """By then the rollout is under way. It degrades to whatever Comfy already had resident - and says so, rather than turning a recoverable placement problem into a dead run.""" - session, reason = _build() - assert session is not None, reason - helpers(_boom) - video, _audio = hc.run_session(session, sigmas=[1.0, 0.0]) - assert video is not None - assert "could not re-place the model" in capsys.readouterr().out - - -def test_no_patcher_means_no_load_attempt(loads): - session, reason = _build() - assert session is not None, reason - session["patcher"] = None - loads.clear() - hc.run_session(session, sigmas=[1.0, 0.0]) - assert loads == [] diff --git a/tests/test_lora_shape_mismatch.py b/tests/test_lora_shape_mismatch.py index 8b0cf60b..df3a3aa3 100644 --- a/tests/test_lora_shape_mismatch.py +++ b/tests/test_lora_shape_mismatch.py @@ -267,7 +267,7 @@ def test_a_lora_that_matches_nothing_prints_both_namings(mm, monkeypatch, capsys mm._log.reset() lora = {"base_model.model.dit.blocks.0.self_attn.qkv.lora_A.weight": torch.zeros(1)} - patches, status = mm.resolve_lora_patches(_Model({}), lora, name="raven.safetensors") + patches, status = mm.resolve_lora_patches(_Model({}), lora, name="adapter.safetensors") assert patches == {} and status == "MATCHED NOTHING" out = capsys.readouterr().out @@ -302,3 +302,56 @@ def test_the_curve_note_leads_with_the_answer_and_prints_whole(mm, monkeypatch, import inspect assert "full=True" in inspect.getsource(mm.resolve_lora_patches) + + +# ── a wrapper head: rename it, do not strip it ────────────────────────────── +# +# A PEFT adapter trained against a wrapper that calls the DiT something else carries comfy's +# module paths EXACTLY, under a different head — `base_model.model.dit.` where comfy's +# key map holds `diffusion_model.`. Stripping leaves `blocks.0...`, comfy's map has no +# bare form to land on, and the whole file silently applies nothing. Measured: 532 keys, zero +# applied, no error. Renaming the head is the fix, and it is scored against the real key map +# alongside every strip so it wins only where it actually lands. + +def test_a_wrapper_head_is_renamed_to_comfys_not_removed(mm): + assert ("base_model.model.dit.", "diffusion_model.") in mm.LORA_KEY_RENAMES + + +def test_renaming_a_head_yields_the_keys_comfys_map_holds(mm): + renamed = mm._rekey( + {"base_model.model.dit.audio_patch_proj.lora_A.weight": 1, + "base_model.model.dit.blocks.0.adaln_proj.linear.lora_B.weight": 2}, + "base_model.model.dit.", "diffusion_model.") + assert set(renamed) == {"diffusion_model.audio_patch_proj.lora_A.weight", + "diffusion_model.blocks.0.adaln_proj.linear.lora_B.weight"} + + +def test_the_empty_replacement_is_a_plain_strip(mm): + stripped = mm._rekey({"base_model.model.dit.blocks.0.attn.qkv_proj.lora_A.weight": 1}, + "base_model.model.dit.") + assert set(stripped) == {"blocks.0.attn.qkv_proj.lora_A.weight"} + + +def test_a_head_the_file_does_not_carry_is_declined_not_applied(mm): + """None, so the candidate is skipped rather than scored as a no-op variant of itself.""" + assert mm._rekey({"diffusion_model.x.lora_A.weight": 1}, "base_model.model.") is None + + +def test_a_rename_wins_over_a_strip_when_it_matches_more(mm, monkeypatch): + """Both are scored against the real key map and the best is kept, so a rename that lands + beats a strip that does not without either being ordered ahead of the other.""" + key_map = {"diffusion_model.audio_patch_proj": "audio_patch_proj.weight"} + monkeypatch.setattr(mm, "_count_lora_matches", + lambda lora, km: sum(1 for k in lora + if k.rsplit(".lora_", 1)[0] in km)) + best, note = mm._best_lora_keying( + {"base_model.model.dit.audio_patch_proj.lora_A.weight": 1}, key_map) + assert note == "base_model.model.dit. -> diffusion_model." + assert "diffusion_model.audio_patch_proj.lora_A.weight" in best + + +def test_a_file_that_needs_nothing_is_left_as_is(mm, monkeypatch): + monkeypatch.setattr(mm, "_count_lora_matches", lambda lora, km: len(lora)) + lora = {"diffusion_model.x.lora_A.weight": 1} + best, note = mm._best_lora_keying(lora, {}) + assert best is lora and note == "as-is"