feat(h3): the causal lane runs on our own model, with no third-party pack
The two entry points the rollout was waiting on — prefill_text and forward_chunk — are written, and with them the whole lane is ours. raven_causal.py is gone: its pure parts (step rules, the audio sigma remap, the rollout loop) moved into h3_causal.py, and its bridge to ComfyUI-MiniMax-H3-RAVEN-Streaming is deleted along with the probe that hunted for that package in custom_nodes. Nothing has to be installed for the mode to exist. How a chunk is cut: ONE full-clip PackedLayout is built by ComfyUI's own extra_conds, and every chunk is a SELECTION of its rows. So a chunk's RoPE angles, modality tags and timestep classes are the ones a dense forward would have given those same rows, and the keyframe pins and reference blocks keep the rows upstream packed for them — which is what makes i2v and r2v work here rather than be refused. Those rows are the prefix: cached once as chunk 0, so media chunk i is cache chunk i + 1 and a sink of 2 pins the prompt plus the opening shot. A chunk is committed CLEAN — one extra forward at sigma 0 after its last step. Caching the last denoising step's K/V would hand every later chunk a context that still carries that step's noise. Two failures that used to be a crash fifty blocks deep are now refusals by name: a clip whose soundtrack does not cover its picture, and a clip too short to have anything for a cache to remember. Tests: 1800 python (+1 net; the RAVEN bridge's tests are gone, the plan, session and real-model checks are new), 68 JS. The new upstream check drives the REAL DiT at toy width through both lanes and confirms the safety property holds: with no cache passed, the causal model is bit-identical to stock H3. Not yet run on the rental.
This commit is contained in:
digital-garbage
+562
-16
@@ -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)
|
||||
|
||||
-413
@@ -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,
|
||||
)
|
||||
+5
-5
@@ -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 "
|
||||
|
||||
@@ -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])
|
||||
@@ -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 == []
|
||||
@@ -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
|
||||
|
||||
|
||||
@@ -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:]
|
||||
@@ -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 == []
|
||||
Reference in New Issue
Block a user