diff --git a/h3_causal.py b/h3_causal.py index 125c2799..9575d408 100644 --- a/h3_causal.py +++ b/h3_causal.py @@ -20,9 +20,18 @@ here can degrade an ordinary generation. GEOMETRY IS BORROWED, NOT REDERIVED ----------------------------------- A chunk's RoPE positions must be the ones the dense layout would have given those same rows, -or the model is being told the clip is a different shape than it is. So the FULL clip's grid -is built with upstream's own `_frame_grid` / `_video_grid` / `_audio_grid` and then SLICED per -chunk. An upstream change to the grid is inherited rather than silently disagreed with. +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 ------------------------ @@ -154,20 +163,219 @@ def chunk_bounds(latent_t, audio_t): return bounds -def clip_grids(text_len, latent_t, latent_h, latent_w, audio_t): - """The whole clip's video and audio position grids, and the rows per latent frame. +class ChunkPlan: + """Which packed rows each chunk owns, cut out of ONE full-clip PackedLayout. - Built with upstream's own helpers and sliced per chunk afterwards, so a chunk's RoPE - angles are exactly the ones the dense layout would have given those rows. Re-deriving the - geometry here instead would be a second source of truth for the clip's shape. + 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. + + Two kinds of selection: + + `prefix` is everything before the target streams — the text, the keyframe cond rows and + the reference blocks. It is cached ONCE, as cache chunk 0, and every media chunk reads it. + That is why i2v and r2v survive here: the conditioning rows keep the layout upstream built + for them instead of being re-packed into a chunk that has no place for them. + + `chunk(i)` is media chunk i: its audio rows followed by its video rows, in the packing + order the target segments already use. Media chunk i is CACHE chunk i + 1, because the + prefix took index 0. """ - from comfy.ldm.minimax.model import _audio_grid, _frame_grid, _video_grid - frame, w_grid = _frame_grid(int(latent_h), int(latent_w)) - video = _video_grid(int(latent_t), frame, float(text_len)) - audio = _audio_grid(float(text_len), int(audio_t), - float(w_grid[0]), float(w_grid[-1])) - return video, audio, int(frame.shape[0]) + def __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.") + self.prefix_rows = torch.arange(0, audio_seg[0], dtype=torch.long) + self.prefix_runs = [(a, b, kind) for a, b, kind in segments if b <= audio_seg[0]] + + @property + def n_chunks(self): + return len(self.bounds) + + @staticmethod + def cache_index(index): + """Media chunk i lives at cache index i + 1; the prefix holds 0.""" + return int(index) + 1 + + 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, 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. + + `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): + 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) + + 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): @@ -190,8 +398,11 @@ def _causal_classes(): """ import comfy.model_management import comfy.quant_ops - from comfy.ldm.minimax.model import (Attention, DiTBlock, MiniMaxH3Model, - _mod_gate, _mod_scale_shift, optimized_attention) + from comfy.ldm.minimax.model import (AUDIO_COND_TIMESTEP, VISUAL_COND_TIMESTEP, Attention, + DiTBlock, MiniMaxH3Model, _mod_gate, _mod_scale_shift, + optimized_attention, pack_audio, patchify_video, + rope_rotation_table, time_shift_slope, unpack_audio, + unpatchify_video) class CausalAttention(Attention): """Upstream attention, with the retained K/V prepended to this chunk's own. @@ -282,6 +493,183 @@ def _causal_classes(): 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_text(self, context, plan, cache, *, video_sigma=0.0, audio_sigma=0.0, + transformer_options={}, minimax_payload=None): + """Run the prompt and the conditioning rows once, and cache them as chunk 0. + + The prefix is the text, the keyframe cond rows and the reference blocks — the part + of the sequence that does not change from chunk to chunk. Caching it once is not + only cheaper: every later chunk assumes the text keys it attends are the ones the + model saw, and recomputing them per chunk would quietly make that false. + + Cached at sigma 0 by default, so the prefix is CLEAN context. The cond rows pin + near 1 in upstream's own table regardless, and text is never noised, so this is the + state those rows would be in at the end of any schedule. + """ + payload = minimax_payload or {} + rows = plan.prefix_rows + if rows.numel() == 0: + raise CacheError("this layout has no prefix rows to prefill") + device = comfy.model_management.get_torch_device() + if hasattr(context, "device"): + device = context.device + dtype = context.dtype + t_v = float(1.0 - max(float(video_sigma), 1e-6)) + t_a = float(1.0 - max(float(audio_sigma), 1e-6)) + runs = plan.prefix_runs + seg_t, t_row, unique_t = self._modulation({k for _, _, k in runs}, t_v, t_a, payload) + text_states = self._text_states(context, transformer_options) + cond_video = self._cond_video_rows(payload, device) + cond_audio = self._cond_audio_rows(payload, device) + + h = torch.empty(int(rows.shape[0]), self.hidden_size, dtype=dtype, device=device) + voff = aoff = 0 + for a, b, kind in runs: + 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, 0, True) + cache.finish_chunk(0) + + 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. + """ + 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) + 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) + return (video_out.to(video_latent.dtype), + (slope.to(audio_out.dtype) * audio_out).to(audio_latent.dtype)) + return CausalAttention, CausalDiTBlock, CausalMiniMaxH3Model @@ -326,3 +714,161 @@ def make_causal(model): 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 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"): + return None, ("the loaded model is a stock bidirectional H3. Switch on `chunk_causal` " + "on FunPack Diffusion Model Loader, and load a LoRA trained for the " + "chunked pattern with FunPack LoRA Loader — it applies on its own, no " + "stack node needed.") + 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) + 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, 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 {}) + + # LOAD FIRST. 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). This lane calls the DiT + # directly, so without this the weights may still be offloaded AND every object patch + # FunPack installed (the AdaLN modality gains, the token-refiner edit) would silently not + # be there. Once, for the whole rollout: the loop must not fight Comfy's offload decision + # chunk by chunk. + patcher = session.get("patcher") + if patcher is not None: + try: + import comfy.model_management + comfy.model_management.load_models_gpu([patcher], memory_required=0, + force_full_load=False) + except Exception as error: # noqa: BLE001 + print(f"[FunPack H3] causal lane could not pre-load the model ({error}); the " + f"rollout continues, but object patches may not be installed.") + + 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, cancel=cancel) diff --git a/raven_causal.py b/raven_causal.py deleted file mode 100644 index acdeedfd..00000000 --- a/raven_causal.py +++ /dev/null @@ -1,413 +0,0 @@ -"""Chunk-causal rollout for MiniMax H3, on FunPack's own schedule and step rule. - -WHAT THIS IS ------------- -RAVEN's contribution is not a sampler. It is a way of *shaping the sequence*: instead of -denoising the whole clip at once, the clip is cut into time chunks, each chunk is carried to -completion, committed into a KV cache as clean context, and the next chunk is generated -attending to that cache. The model gets real memory of what it already drew, so a long clip -stays one continuous shot instead of a stitched chain. - -That shaping is separable from the four-step consistency sampler RAVEN ships with it, and this -module keeps them apart so each can be tested on its own: - -* the CHUNKING and the cache come from the RAVEN package (``prefill_text`` / ``forward_chunk``); -* the SCHEDULE is FunPack's — any scheduler, any step count, including bong_tangent; -* the STEP RULE is selectable — RAVEN's consistency transition, or plain flow euler, or - ancestral euler. - -WHY THE FIRST CHUNK IS DENSE ----------------------------- -RAVEN's own sampler refuses keyframe pins (``fl2va``) and reference blocks (``ref2va``): its -causal layout models text and target rows only, so an i2v anchor has nowhere to go. Rather -than extend that layout, chunk 0 is sampled through FunPack's ORDINARY dense H3 path — where -pins, references, region locks, AdaLN gains and everything else already work — and only THEN -committed into the cache as clean context. Chunks 1..N continue causally from it. - -So the anchor lands where an anchor belongs (the opening frames), every conditioning feature -FunPack has keeps working, and the causal lane is only ever asked to do the thing it was -trained for: continue from cached context. - -REQUIREMENTS ------------- -The RAVEN LoRA is not optional. The chunk-causal attention pattern is what it was trained to -read; the base H3 weights have never seen a KV cache. Without the LoRA this produces output, -and the output is out of distribution. The caller is told, once, rather than discovering it. -""" - -from __future__ import annotations - -import importlib -import os -import sys - -import torch - -#: video latents per chunk, and the clock span that many latents cover. Re-derived here so a -#: probe can check them against the installed package instead of trusting a remembered number. -VIDEO_LATENTS_PER_CHUNK = 5 -CHUNK_T_SPAN = 85.0 / 3.0 - -#: how a chunk's next-sigma is reached from its prediction -STEP_RULES = ("consistency", "euler", "euler_ancestral") - -_PROBE = {"state": None} - - -def locate_raven(): - """Import ``raven_streaming``, from ComfyUI's custom_nodes if it is not already on the path. - - Returns (module, reason). The module is None when it cannot be had, and `reason` says why - in the terms the user can act on — this is an optional capability, not a dependency. - """ - if _PROBE["state"] is not None: - return _PROBE["state"] - try: - module = importlib.import_module("raven_streaming") - _PROBE["state"] = (module, "") - return _PROBE["state"] - except ImportError: - pass - # Not on sys.path: ComfyUI does not add a custom node's directory for other packs to - # import from. Find it the way a user would describe it — the sibling directory. - here = os.path.dirname(os.path.abspath(__file__)) - custom_nodes = os.path.dirname(here) - for name in sorted(os.listdir(custom_nodes)) if os.path.isdir(custom_nodes) else []: - candidate = os.path.join(custom_nodes, name) - if not os.path.isdir(os.path.join(candidate, "raven_streaming")): - continue - if candidate not in sys.path: - sys.path.append(candidate) - try: - module = importlib.import_module("raven_streaming") - _PROBE["state"] = (module, "") - return _PROBE["state"] - except Exception as error: # found, but will not import - _PROBE["state"] = (None, f"found {name} but it did not import ({error})") - return _PROBE["state"] - _PROBE["state"] = ( - None, - "ComfyUI-MiniMax-H3-RAVEN-Streaming is not installed. The causal rollout borrows its " - "chunk cache; install it into custom_nodes and restart ComfyUI.", - ) - return _PROBE["state"] - - -#: pixel frames one full chunk covers: FRAME_PER_TOKEN over 5 latents is 1+4+4+4+4 -FRAMES_PER_CHUNK = 17 -#: audio latents per pixel frame — 40 latents/s against the model's fixed 24 fps -AUDIO_LATENTS_PER_FRAME = 40.0 / 24.0 - - -def chunk_bounds(latent_t, audio_t): - """Plan the chunk cut: (video_start, video_stop, audio_start, audio_stop) per chunk. - - Video cuts every ``VIDEO_LATENTS_PER_CHUNK`` latents. Audio is NOT cut by a remembered - 29/28/28 cadence — it is derived from the shared clock, so a chunk owns the audio latents - whose time falls inside its video span and the cadence comes out as a consequence. - - The LIVE rollout uses the RAVEN package's own layout object, because that is what its - ``forward_chunk`` indexes. This exists to plan and report a run before the package is - touched (how many chunks, how long), and to be checkable against theirs. - """ - latent_t, audio_t = int(latent_t), int(audio_t) - if latent_t <= 0: - return [] - bounds = [] - audio_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: - frames_done = FRAMES_PER_CHUNK * (index + 1) - audio_stop = min(audio_t, int(round(frames_done * AUDIO_LATENTS_PER_FRAME))) - audio_stop = max(audio_stop, audio_cursor) - bounds.append((video_start, video_stop, audio_cursor, audio_stop)) - audio_cursor = audio_stop - return bounds - - -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, 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. - - `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): - 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) - - 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 build_session(model, positive, latent, *, sink=2, window=2, storage="cpu_pinned", - device=None, compute_dtype=None): - """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, not for a log parser. - - The pieces are the RAVEN package's own — its contracts parse the sockets, its layout cuts - the chunks, its cache holds the K/V, its causal model does the per-chunk forward. FunPack - supplies only the schedule and the step rule, which is the whole point of the split. - """ - module, reason = locate_raven() - if module is None: - return None, reason - try: - from raven_streaming import cache as cache_mod - from raven_streaming import contracts - except Exception as error: # installed but not importable - return None, f"the RAVEN package is present but its modules did not load ({error})" - try: - resolved = contracts.resolve_model(model) - conditioning = contracts.parse_conditioning(positive) - request = contracts.parse_latent(latent, warn_experimental=False) - layout = request.layout(conditioning.text_len, warn_experimental=False) - except Exception as error: - return None, f"this run is not one the causal lane can take: {error}" - - diffusion = resolved.diffusion_model - if not hasattr(diffusion, "forward_chunk") or not hasattr(diffusion, "prefill_text"): - return None, ("the loaded model is a stock bidirectional H3. Switch on `chunk_causal` " - "on FunPack Diffusion Model Loader, and load a LoRA trained for the " - "chunked pattern with FunPack LoRA Loader — it applies on its own, no " - "stack node needed.") - blocks = getattr(diffusion, "blocks", None) - if not blocks: - return None, "the loaded model exposes no DiT blocks to cache" - - if device is None: - device = resolved.load_device if resolved.load_device is not None else request.device - if compute_dtype is None: - # One dtype for the prefill AND every chunk: the cache is filled in one and read in - # the other otherwise, and the attention module refuses that mid-rollout. - compute_dtype = getattr(diffusion, "dtype", None) or request.dtype - - kv = cache_mod.ChunkKVCache(len(blocks), sink=int(sink), - window=None if window is None else int(window), - storage=str(storage)) - try: - from raven_streaming import contracts as _c - video_channels = int(_c.VIDEO_LATENT_CHANNELS) - audio_channels = int(_c.AUDIO_LATENT_CHANNELS) - except Exception: - video_channels, audio_channels = 24, 32 - return { - "module": module, - "model": diffusion, - "video_channels": video_channels, - "audio_channels": audio_channels, - "shift_video": float(getattr(diffusion, "sigma_shift_video", 12.0)), - "shift_audio": float(getattr(diffusion, "sigma_shift_audio", 3.0)), - "patcher": resolved.patcher, - "conditioning": conditioning, - "request": request, - "layout": layout, - "cache": kv, - "device": torch.device(device), - "compute_dtype": compute_dtype, - }, "" - - -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, cancel=None, transformer_options=None): - """Drive one causal rollout over an assembled session. Returns (video, audio). - - The text is written into the cache as chunk 0, alone and once. Folding it into the first - media chunk would let text rows attend media rows, and every later chunk assumes the text - keys it cached are the ones the model saw. - """ - layout = session["layout"] - model = session["model"] - cache = session["cache"] - # LOAD FIRST. 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). This lane calls the DiT - # directly, so without this the weights may still be offloaded AND every object patch - # FunPack installed (the AdaLN modality gains, the token-refiner edit) would silently not - # be there. Once, for the whole rollout: the loop must not fight Comfy's offload decision - # chunk by chunk. - patcher = session.get("patcher") - if patcher is not None: - try: - import comfy.model_management - comfy.model_management.load_models_gpu([patcher], memory_required=0, - force_full_load=False) - except Exception as error: # noqa: BLE001 - print(f"[FunPack H3] causal lane could not pre-load the model ({error}); " - f"the rollout continues, but object patches may not be installed.") - device = session["device"] - compute_dtype = session["compute_dtype"] - conditioning = session["conditioning"] - options = transformer_options or {} - - model.prefill_text( - conditioning.cross_attn.to(device=device, dtype=compute_dtype), - cache=cache, - transformer_options=options, - text_token_tags=(None if conditioning.token_tags is None - else conditioning.token_tags.to(device)), - compute_dtype=compute_dtype, - ) - - 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(layout.video_latent_shape(session["video_channels"])) - audio_noise = draw_noise(layout.audio_latent_shape(session["audio_channels"])) - - chunks = [(c.video_start, c.video_stop, c.audio_start, c.audio_stop) for c in layout.chunks] - shift_video, shift_audio = 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 further the step, the worse it is. - return model.forward_chunk( - video_latent=video_xt, audio_latent=audio_xt, layout=layout, chunk_index=index, - cache=cache, role="noise", video_sigma=float(sigma), - audio_sigma=float(audio_sigma_for(sigma, shift_video, shift_audio)), - update_cache=False, transformer_options=options, compute_dtype=compute_dtype, - ) - - def commit(video_x0, audio_x0, index): - model.forward_chunk( - video_latent=video_x0, audio_latent=audio_x0, layout=layout, chunk_index=index, - cache=cache, role="clean", - video_eps=draw_noise(video_x0.shape), audio_eps=draw_noise(audio_x0.shape), - update_cache=True, transformer_options=options, compute_dtype=compute_dtype, - ) - - return causal_rollout( - chunks=chunks, 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, cancel=cancel, - ) diff --git a/samplers.py b/samplers.py index a7a76825..8e915dc2 100644 --- a/samplers.py +++ b/samplers.py @@ -2841,10 +2841,10 @@ def _normalize_video_denoised(denoised, video_mask, sigma, ref, strength, # the loop can never drift apart. Imported defensively: a broken optional module must not stop # the node pack from registering. try: - from .raven_causal import STEP_RULES as _CAUSAL_STEP_RULES + from .h3_causal import STEP_RULES as _CAUSAL_STEP_RULES except ImportError: try: - from raven_causal import STEP_RULES as _CAUSAL_STEP_RULES + from h3_causal import STEP_RULES as _CAUSAL_STEP_RULES except ImportError: _CAUSAL_STEP_RULES = ("consistency", "euler", "euler_ancestral") @@ -6889,10 +6889,10 @@ class FunPackLTXAVSceneChainSampler: if not getattr(self, "_h3_causal_chunks", False): return None try: - from . import raven_causal as causal + from . import h3_causal as causal from . import minimax_h3 as h3mod except ImportError: - import raven_causal as causal + import h3_causal as causal import minimax_h3 as h3mod if not h3mod.is_h3_model(model): _log.note_on_change( @@ -6914,7 +6914,7 @@ class FunPackLTXAVSceneChainSampler: schedule = [float(v) for v in sigmas] _log.note_on_change( "h3:causal_on", "FunPackSceneChain", - f"H3 causal chunks: {len(session['layout'].chunks)} chunks, " + f"H3 causal chunks: {session['plan'].n_chunks} chunks, " f"{len(schedule) - 1} steps each, step rule " f"'{getattr(self, '_h3_causal_step_rule', 'consistency')}', sink " f"{getattr(self, '_h3_causal_sink', 2)} / window " diff --git a/tests/_h3_causal_upstream_check.py b/tests/_h3_causal_upstream_check.py new file mode 100644 index 00000000..073be44b --- /dev/null +++ b/tests/_h3_causal_upstream_check.py @@ -0,0 +1,198 @@ +"""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] + + # 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_raven_causal.py b/tests/test_causal_rollout.py similarity index 62% rename from tests/test_raven_causal.py rename to tests/test_causal_rollout.py index c4b13cb3..ecde7912 100644 --- a/tests/test_raven_causal.py +++ b/tests/test_causal_rollout.py @@ -1,15 +1,17 @@ -"""Chunk-causal rollout: RAVEN's sequence shaping, on FunPack's schedule and step rule. +"""Chunk-causal rollout: the cut, the step rules, and the order the loop does things in. -RAVEN's contribution is not a sampler. It is a way of shaping the sequence — cut the clip into +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 it ships with, and these tests -pin the separation: the chunking is theirs, the schedule is ours, and the step rule is a knob -so a chunked run can be compared against an unchunked one without two variables moving. +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 -import types from pathlib import Path import pytest @@ -17,7 +19,7 @@ import torch sys.path.insert(0, str(Path(__file__).resolve().parents[1])) -import raven_causal as rc +import h3_causal as rc # ── the chunk cut ─────────────────────────────────────────────────────────── @@ -78,7 +80,7 @@ def test_an_empty_clip_plans_nothing(): # ── the step rules ────────────────────────────────────────────────────────── def test_consistency_jumps_to_x0_and_renoises_with_fresh_noise(): - """RAVEN's own transition, and the one its LoRA is distilled for.""" + """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) @@ -166,27 +168,38 @@ def test_every_advertised_rule_actually_runs(): assert out.shape == (4,) -# ── the package probe ─────────────────────────────────────────────────────── +# ── the audio stream's own sigma ──────────────────────────────────────────── -def test_a_missing_package_names_what_to_install(monkeypatch): - monkeypatch.setitem(rc._PROBE, "state", None) - monkeypatch.setattr(rc.importlib, "import_module", - lambda name: (_ for _ in ()).throw(ImportError(name))) - monkeypatch.setattr(rc.os.path, "isdir", lambda p: False) - module, reason = rc.locate_raven() - assert module is None - assert "RAVEN-Streaming" in reason and "custom_nodes" in reason +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_the_probe_is_only_paid_once(monkeypatch): - monkeypatch.setitem(rc._PROBE, "state", ("sentinel", "")) - assert rc.locate_raven() == ("sentinel", "") +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, no -# RAVEN package and no weights. What is being tested is the ORDER of things: which chunks are +# 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. @@ -360,174 +373,8 @@ def test_the_chunk_callback_reports_whether_it_was_generated(): assert all(not known for _i, known in events[1:]) -# ── the bridge to the RAVEN package ───────────────────────────────────────── -# -# build_session assembles their contracts/layout/cache/causal model; run_session drives it -# with OUR schedule and step rule. Tested against a fake package, because the real one needs -# weights and a GPU — what is being pinned is the wiring, not the numerics. - -class FakeLayout: - def __init__(self, latent_t=12, audio_t=100, latent_h=4, latent_w=4): - self.latent_t, self.audio_t = latent_t, audio_t - self.latent_h, self.latent_w = latent_h, latent_w - bounds = rc.chunk_bounds(latent_t, audio_t) - self.chunks = [types.SimpleNamespace(index=i, video_start=a, video_stop=b, - audio_start=c, audio_stop=d) - for i, (a, b, c, d) in enumerate(bounds)] - - def video_latent_shape(self, dim=24): - return (1, dim, self.latent_t, self.latent_h, self.latent_w) - - def audio_latent_shape(self, dim=32): - return (1, dim, 2, self.audio_t) - - -class FakeCausalModel: - dtype = torch.float32 - sigma_shift_video = 12.0 - sigma_shift_audio = 3.0 - - def __init__(self): - self.blocks = [object()] * 4 - self.prefills = 0 - self.forwards = [] - self.commits = [] - - def prefill_text(self, context, **kw): - self.prefills += 1 - return int(context.shape[1]) - - def forward_chunk(self, *, video_latent, audio_latent, chunk_index, role, - video_sigma=None, audio_sigma=None, update_cache=None, **kw): - if role == "clean": - self.commits.append(chunk_index) - else: - self.forwards.append((chunk_index, video_sigma, audio_sigma)) - return [torch.zeros_like(video_latent), torch.zeros_like(audio_latent)] - - -def _session(model=None, layout=None): - model = model or FakeCausalModel() - return { - "module": None, "model": model, "patcher": None, - "conditioning": types.SimpleNamespace( - cross_attn=torch.zeros(1, 7, 16), token_tags=None, text_len=7), - "request": None, "layout": layout or FakeLayout(), - "cache": object(), "device": torch.device("cpu"), - "compute_dtype": torch.float32, - "video_channels": 24, "audio_channels": 32, - "shift_video": 12.0, "shift_audio": 3.0, - } - - -def test_the_text_is_prefilled_exactly_once(): - """Text is cache chunk 0 and is written alone: folding it into the first media chunk - would let text rows attend media rows, and every later chunk assumes otherwise.""" - model = FakeCausalModel() - rc.run_session(_session(model), sigmas=(1.0, 0.5, 0.0)) - assert model.prefills == 1 - - -def test_the_audio_stream_gets_its_own_sigma(): - """H3 denoises audio on a shift of 3 against the video's 12. Handing the audio rows the - video's sigma is the error h3_audio_clock exists to correct.""" - model = FakeCausalModel() - rc.run_session(_session(model), sigmas=(1.0, 0.83, 0.0)) - video_sigmas = [v for _i, v, _a in model.forwards] - audio_sigmas = [a for _i, _v, a in model.forwards] - assert 0.83 in video_sigmas - assert not any(abs(a - 0.83) < 1e-9 for a in audio_sigmas) - assert any(abs(a - rc.audio_sigma_for(0.83, 12.0, 3.0)) < 1e-9 for a in audio_sigmas) - - -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 - - -def test_every_chunk_is_committed_through_the_clean_role(): - model = FakeCausalModel() - session = _session(model) - rc.run_session(session, sigmas=(1.0, 0.5, 0.0)) - assert model.commits == [c.index for c in session["layout"].chunks] - - -def test_the_noise_shapes_come_from_the_layout(): - model = FakeCausalModel() - layout = FakeLayout(latent_t=7, audio_t=37) - video, audio = rc.run_session(_session(model, layout), sigmas=(1.0, 0.0)) - assert tuple(video.shape) == layout.video_latent_shape(24) - assert tuple(audio.shape) == layout.audio_latent_shape(32) - - -def test_the_same_seed_gives_the_same_clip(): - a = rc.run_session(_session(), sigmas=(1.0, 0.5, 0.0), seed=7)[0] - b = rc.run_session(_session(), sigmas=(1.0, 0.5, 0.0), seed=7)[0] - assert torch.allclose(a, b) - - -def test_a_different_seed_gives_a_different_clip(): - a = rc.run_session(_session(), sigmas=(1.0, 0.5, 0.0), seed=7)[0] - b = rc.run_session(_session(), sigmas=(1.0, 0.5, 0.0), seed=8)[0] - assert not torch.allclose(a, b) - - -def test_our_schedule_reaches_the_model_unchanged(): - model = FakeCausalModel() - rc.run_session(_session(model), sigmas=(1.0, 0.94, 0.83, 0.72, 0.55, 0.30, 0.10, 0.0)) - first = [v for i, v, _a in model.forwards if i == 0] - assert first == [1.0, 0.94, 0.83, 0.72, 0.55, 0.30, 0.10] - - -def test_a_stock_bidirectional_model_is_refused_by_name(monkeypatch): - """The causal lane needs the chunk-causal DiT; a stock H3 has no forward_chunk and would - otherwise fail somewhere deep in the rollout instead of at the door.""" - monkeypatch.setitem(rc._PROBE, "state", (types.SimpleNamespace(), "")) - fake_contracts = types.SimpleNamespace( - resolve_model=lambda m, **k: types.SimpleNamespace( - diffusion_model=types.SimpleNamespace(), patcher=None, load_device="cpu"), - parse_conditioning=lambda p: types.SimpleNamespace(text_len=7, token_tags=None), - parse_latent=lambda l, **k: types.SimpleNamespace( - layout=lambda *a, **k: FakeLayout(), device="cpu", dtype=torch.float32), - ) - # The parent module has to be registered too, and the submodules attached to it as - # ATTRIBUTES: `from raven_streaming import cache` resolves through the parent, so a - # sys.modules entry alone leaves the import reading nothing. Same trap the suite's own - # comfy stubs were built to avoid. - fake_cache = types.SimpleNamespace(ChunkKVCache=lambda *a, **k: object()) - parent = types.ModuleType("raven_streaming") - parent.contracts, parent.cache = fake_contracts, fake_cache - monkeypatch.setitem(sys.modules, "raven_streaming", parent) - monkeypatch.setitem(sys.modules, "raven_streaming.contracts", fake_contracts) - monkeypatch.setitem(sys.modules, "raven_streaming.cache", fake_cache) - session, reason = rc.build_session(object(), [], {}) - assert session is None - # points at OUR loader: FunPack owns the loading pipeline, so the fix is a setting on the - # node already in the graph, not a third-party loader node to go and add - assert "stock bidirectional H3" in reason - assert "FunPack Diffusion Model Loader" in reason and "chunk_causal" in reason +class _FakePlan: + n_chunks = 3 # ── the Chain Sampler dispatch ────────────────────────────────────────────── @@ -576,18 +423,18 @@ def test_a_non_h3_model_falls_through_and_says_so(monkeypatch, capsys): assert "not a MiniMax H3 model" in capsys.readouterr().out -def test_a_missing_raven_package_falls_through_and_says_why(monkeypatch, capsys): +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, "the pack is not installed.")) + 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 "not installed" in out and "samples normally" in out + assert "stock bidirectional" in out and "samples normally" in out def test_a_rollout_failure_is_raised_not_swallowed(monkeypatch): @@ -599,7 +446,7 @@ def test_a_rollout_failure_is_raised_not_swallowed(monkeypatch): import minimax_h3 monkeypatch.setattr(minimax_h3, "is_h3_model", lambda m: True) monkeypatch.setattr(rc, "build_session", - lambda *a, **k: ({"layout": FakeLayout()}, "")) + 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, @@ -616,12 +463,11 @@ def test_the_run_says_the_sampler_and_cfg_are_not_used(monkeypatch, capsys): lambda tag, who, msg: print(msg), raising=False) import minimax_h3 monkeypatch.setattr(minimax_h3, "is_h3_model", lambda m: True) - layout = FakeLayout() # 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: ({"layout": layout}, "")) + 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, @@ -630,78 +476,3 @@ def test_the_run_says_the_sampler_and_cfg_are_not_used(monkeypatch, capsys): out = capsys.readouterr().out assert "CFG" in out and "not used" in out.replace("NOT used", "not used") assert "no per-chunk redo" in out.lower() - - -# ── 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. - -def test_the_model_is_loaded_before_the_rollout(monkeypatch): - import comfy.model_management - calls = [] - monkeypatch.setattr(comfy.model_management, "load_models_gpu", - lambda models, **kw: calls.append((models, kw)), raising=False) - model = FakeCausalModel() - session = _session(model) - session["patcher"] = "the-patcher" - rc.run_session(session, sigmas=(1.0, 0.0)) - assert calls and calls[0][0] == ["the-patcher"] - - -def test_the_load_happens_before_the_first_forward(monkeypatch): - import comfy.model_management - order = [] - monkeypatch.setattr(comfy.model_management, "load_models_gpu", - lambda models, **kw: order.append("load"), raising=False) - model = FakeCausalModel() - - real_forward = model.forward_chunk - def spy(**kw): - order.append("forward") - return real_forward(**kw) - model.forward_chunk = spy - - session = _session(model) - session["patcher"] = "the-patcher" - rc.run_session(session, sigmas=(1.0, 0.0)) - assert order[0] == "load" - - -def test_the_model_is_loaded_once_not_per_chunk(monkeypatch): - """Loading per chunk would fight Comfy's own offload decision all the way down the clip.""" - import comfy.model_management - calls = [] - monkeypatch.setattr(comfy.model_management, "load_models_gpu", - lambda models, **kw: calls.append(models), raising=False) - session = _session() - session["patcher"] = "the-patcher" - rc.run_session(session, sigmas=(1.0, 0.5, 0.0)) - assert len(calls) == 1 - - -def test_a_failed_load_does_not_abort_the_run(monkeypatch, capsys): - """It degrades to whatever Comfy already had resident, and says the patches may be - missing — rather than turning a recoverable placement problem into a dead run.""" - import comfy.model_management - monkeypatch.setattr( - comfy.model_management, "load_models_gpu", - lambda models, **kw: (_ for _ in ()).throw(RuntimeError("out of memory")), - raising=False) - session = _session() - session["patcher"] = "the-patcher" - video, _audio = rc.run_session(session, sigmas=(1.0, 0.0)) - assert video is not None - assert "object patches may not be installed" in capsys.readouterr().out - - -def test_no_patcher_means_no_load_attempt(monkeypatch): - import comfy.model_management - calls = [] - monkeypatch.setattr(comfy.model_management, "load_models_gpu", - lambda models, **kw: calls.append(models), raising=False) - rc.run_session(_session(), sigmas=(1.0, 0.0)) # _session has no patcher - assert calls == [] diff --git a/tests/test_h3_causal_loader.py b/tests/test_h3_causal_loader.py index 301c8264..2b3ba958 100644 --- a/tests/test_h3_causal_loader.py +++ b/tests/test_h3_causal_loader.py @@ -127,8 +127,8 @@ def test_the_loader_message_names_the_lora_loader(): def test_the_sampler_message_names_the_lora_loader(): import inspect - import raven_causal - src = inspect.getsource(raven_causal.build_session) + import h3_causal + src = inspect.getsource(h3_causal.build_session) assert "FunPack LoRA Loader" in src assert "Apply LoRA Weights" not in src diff --git a/tests/test_h3_causal_plan.py b/tests/test_h3_causal_plan.py new file mode 100644 index 00000000..b1b3e0c1 --- /dev/null +++ b/tests/test_h3_causal_plan.py @@ -0,0 +1,156 @@ +"""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_one(): + """The prefix holds cache index 0, so `sink=2` pins the prompt plus the opening shot.""" + assert hc.ChunkPlan.cache_index(0) == 1 + assert hc.ChunkPlan.cache_index(5) == 6 + + +# ── 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 new file mode 100644 index 00000000..ce6ab337 --- /dev/null +++ b/tests/test_h3_causal_session.py @@ -0,0 +1,333 @@ +"""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_bidirectional_model_is_refused_by_name(): + class Stock: + blocks = [object()] + + session, reason = _build(diffusion=Stock()) + assert session is None + assert "chunk_causal" in reason and "stock bidirectional" 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(): + diffusion, session, _, _ = _run() + written = {key[1] for key in session["cache"]._store} + assert 0 not in written or session["cache"].committed_chunks + assert written <= set(range(1, session["plan"].n_chunks + 1)) + + +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. + +def _loading_session(diffusion=None): + diffusion = diffusion or _Diffusion() + session, reason = _build(diffusion=diffusion) + assert session is not None, reason + return diffusion, session + + +def test_the_model_is_loaded_before_the_rollout(monkeypatch): + import comfy.model_management + calls = [] + monkeypatch.setattr(comfy.model_management, "load_models_gpu", + lambda models, **kw: calls.append((models, kw)), raising=False) + _, session = _loading_session() + hc.run_session(session, sigmas=[1.0, 0.0]) + assert calls and calls[0][0] == [session["patcher"]] + + +def test_the_load_happens_before_the_first_forward(monkeypatch): + import comfy.model_management + order = [] + monkeypatch.setattr(comfy.model_management, "load_models_gpu", + lambda models, **kw: order.append("load"), raising=False) + diffusion, session = _loading_session() + 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_not_per_chunk(monkeypatch): + """Loading per chunk would fight Comfy's own offload decision all the way down the clip.""" + import comfy.model_management + calls = [] + monkeypatch.setattr(comfy.model_management, "load_models_gpu", + lambda models, **kw: calls.append(models), raising=False) + _, session = _loading_session() + hc.run_session(session, sigmas=[1.0, 0.5, 0.0]) + assert len(calls) == 1 + + +def test_a_failed_load_does_not_abort_the_run(monkeypatch, capsys): + """It degrades to whatever Comfy already had resident, and says the patches may be + missing — rather than turning a recoverable placement problem into a dead run.""" + import comfy.model_management + monkeypatch.setattr( + comfy.model_management, "load_models_gpu", + lambda models, **kw: (_ for _ in ()).throw(RuntimeError("out of memory")), + raising=False) + _, session = _loading_session() + video, _audio = hc.run_session(session, sigmas=[1.0, 0.0]) + assert video is not None + assert "object patches may not be installed" in capsys.readouterr().out + + +def test_no_patcher_means_no_load_attempt(monkeypatch): + import comfy.model_management + calls = [] + monkeypatch.setattr(comfy.model_management, "load_models_gpu", + lambda models, **kw: calls.append(models), raising=False) + _, session = _loading_session() + session["patcher"] = None + hc.run_session(session, sigmas=[1.0, 0.0]) + assert calls == []