Ship a modular workflow that carries keyframes and references in one run

combined_layout.py

@@ -0,0 +1,260 @@
+"""Packed layout for a MiniMax-H3 request carrying BOTH keyframes and references.
+
+Shipped with the checkpoint so `ModularPipeline.from_pretrained(..., trust_remote_code=True)` can serve a request
+that has keyframes *and* references, which the stock blocks cannot express.
+
+diffusers ships two builders and dispatches either/or: `references` wins and keyframes are dropped. ComfyUI packs
+both (`PackedLayout(..., keyframes=..., refs=...)` in `comfy/ldm/minimax/model.py`), so the layout exists — this is
+a port of it, not an invention.
+
+Row order:  [ text | keyframe cond | reference blocks | target audio | target video ]
+
+The coupling that makes naive composition wrong: references pack between the text and the targets, so the *target
+timeline* starts after their spans, and the keyframe anchors — which live on the target timeline — shift by exactly
+that amount. ComfyUI does this with a pre-pass (`cursor = text_len + sum(_ref_t_span(blk))`); so does this.
+
+Correctness is pinned by `validate_against_stock()`: with no keyframes this must reproduce diffusers' `ref2va`
+layout bit for bit, and with no references its `fl2va` layout, both including `position_ids` in float64.
+"""
+
+from __future__ import annotations
+
+import numpy as np
+import torch
+from diffusers.modular_pipelines.minimax_h3.before_denoise import (
+    _ROPE_FRAME_RESCALE,
+    _ROPE_FRAMES_PER_LATENT,
+    _fill_audio_positions,
+    _frame_position_grid,
+    _temporal_position_grid,
+    MiniMaxH3PrepareLayoutStep,
+    MiniMaxH3Ref2VAPrepareLayoutStep,
+)
+
+
+def _target_span_sum(num_latent_frames: int) -> float:
+    """Rotary time the generated frames span, by numpy pairwise summation — the order the keyframe anchor uses."""
+    spans = np.ones(num_latent_frames, dtype=np.float64) * _ROPE_FRAME_RESCALE
+    for offset in range(len(_ROPE_FRAMES_PER_LATENT)):
+        spans[offset :: len(_ROPE_FRAMES_PER_LATENT)] *= _ROPE_FRAMES_PER_LATENT[offset]
+    return float(spans.sum())
+
+
+def _video_span_sequential(num_latent_frames: int) -> float:
+    """The same series summed sequentially — the order a reference block advances the clock by. The two differ in
+    the last ulp from 16 latent frames on, and the reference implementation keeps both, one per call site."""
+    return sum(
+        _ROPE_FRAME_RESCALE * _ROPE_FRAMES_PER_LATENT[index % len(_ROPE_FRAMES_PER_LATENT)]
+        for index in range(num_latent_frames)
+    )
+
+
+def _reference_span(reference, visual_geometry, audio_row_counts, audio_channels) -> float:
+    """`_ref_t_span` from ComfyUI: the time axis a reference block occupies ahead of the target streams."""
+    if reference.kind == "image":
+        next(visual_geometry)
+        return 1.0
+    if reference.kind == "audio":
+        return float(next(audio_row_counts) // audio_channels)
+    if reference.kind == "video":
+        audio_latents = (next(audio_row_counts) // audio_channels) if reference.has_audio else 0
+        frames, _, _ = next(visual_geometry)
+        return max(float(audio_latents), _video_span_sequential(frames))
+    raise ValueError(f"A reference must be an 'image', a 'video' or an 'audio', got {reference.kind!r}.")
+
+
+def build_combined_packed_sequence(
+    text_token_tags: torch.Tensor,
+    references: list,
+    condition_latents: list[torch.Tensor],
+    audio_condition_latents: list[torch.Tensor],
+    num_latent_frames: int,
+    latent_height: int,
+    latent_width: int,
+    num_audio_latents: int,
+    patch_size: tuple[int, int, int],
+    audio_channels: int,
+    audio_tag: int,
+    video_tag: int,
+    keyframe_anchors: tuple[str, ...] = (),
+):
+    """`condition_latents` is keyframe latents first (one per `keyframe_anchors` entry), then the reference latents."""
+    _, patch_h, patch_w = patch_size
+    num_keyframes = len(keyframe_anchors)
+    keyframe_latents, reference_latents = condition_latents[:num_keyframes], condition_latents[num_keyframes:]
+
+    num_text_tokens = text_token_tags.shape[0]
+    rows_per_frame = (latent_height // patch_h) * (latent_width // patch_w)
+    num_keyframe_rows = num_keyframes * rows_per_frame
+    num_reference_video_rows = sum(
+        frames * (height // patch_h) * (width // patch_w)
+        for frames, height, width in (tuple(latents.shape[2:5]) for latents in reference_latents)
+    )
+    num_reference_audio_rows = sum(rows.shape[0] for rows in audio_condition_latents)
+    num_target_video_rows = num_latent_frames * rows_per_frame
+    num_target_audio_rows = num_audio_latents * audio_channels
+    sequence_length = (
+        num_text_tokens
+        + num_keyframe_rows
+        + num_reference_video_rows
+        + num_reference_audio_rows
+        + num_target_audio_rows
+        + num_target_video_rows
+    )
+
+    position_ids = torch.zeros(sequence_length, 3, dtype=torch.float64)
+    position_ids[:num_text_tokens, 0] = torch.arange(num_text_tokens, dtype=torch.float64)
+    target_frame_grid, target_width_grid = _frame_position_grid(latent_height, latent_width, patch_h, patch_w)
+
+    # Pre-pass: how far the references push the target timeline out. The keyframe anchors ride on that timeline.
+    span_geometry = iter(tuple(latents.shape[2:5]) for latents in reference_latents)
+    span_audio = iter(rows.shape[0] for rows in audio_condition_latents)
+    reference_span = sum(
+        _reference_span(reference, span_geometry, span_audio, audio_channels) for reference in references
+    )
+    target_origin = float(num_text_tokens) + reference_span
+
+    video_indices, audio_indices = [], []
+
+    # 1. Keyframe conditioning rows, immediately after the text, on the target spatial grid.
+    cursor = num_text_tokens
+    for index, anchor in enumerate(keyframe_anchors):
+        if anchor == "first":
+            anchor_time = target_origin
+        elif anchor == "last":
+            anchor_time = target_origin + _target_span_sum(num_latent_frames) - _ROPE_FRAME_RESCALE
+        else:
+            raise ValueError(f"A keyframe anchor must be 'first' or 'last', got {anchor!r}.")
+        frames = keyframe_latents[index].shape[2]
+        rows = slice(cursor, cursor + frames * rows_per_frame)
+        cursor = rows.stop
+        video_indices.append(torch.arange(rows.start, rows.stop))
+        position_ids[rows, 0] = anchor_time
+        position_ids[rows, 1:] = target_frame_grid.repeat(frames, 1)
+
+    # 2. Reference blocks, on their own clock starting where the text ends — exactly the stock `ref2va` walk.
+    visual_geometry = iter(tuple(latents.shape[2:5]) for latents in reference_latents)
+    audio_row_counts = iter(rows.shape[0] for rows in audio_condition_latents)
+    rotary_time = float(num_text_tokens)
+    for reference in references:
+        if reference.kind == "image":
+            frames, height, width = next(visual_geometry)
+            rows = slice(cursor, cursor + frames * (height // patch_h) * (width // patch_w))
+            cursor = rows.stop
+            video_indices.append(torch.arange(rows.start, rows.stop))
+            frame_grid, _ = _frame_position_grid(height, width, patch_h, patch_w)
+            position_ids[rows, 0] = rotary_time
+            position_ids[rows, 1:] = frame_grid
+            rotary_time += 1.0
+        elif reference.kind == "audio":
+            num_rows = next(audio_row_counts)
+            latents = num_rows // audio_channels
+            rows = slice(cursor, cursor + num_rows)
+            cursor = rows.stop
+            audio_indices.append(torch.arange(rows.start, rows.stop))
+            _fill_audio_positions(position_ids, rows, latents, rotary_time, target_width_grid, audio_channels)
+            rotary_time += float(latents)
+        elif reference.kind == "video":
+            num_rows = next(audio_row_counts) if reference.has_audio else 0
+            latents = num_rows // audio_channels
+            frames, height, width = next(visual_geometry)
+            audio_rows = slice(cursor, cursor + num_rows)
+            video_rows = slice(audio_rows.stop, audio_rows.stop + frames * (height // patch_h) * (width // patch_w))
+            cursor = video_rows.stop
+            audio_indices.append(torch.arange(audio_rows.start, audio_rows.stop))
+            video_indices.append(torch.arange(video_rows.start, video_rows.stop))
+            frame_grid, width_grid = _frame_position_grid(height, width, patch_h, patch_w)
+            _fill_audio_positions(position_ids, audio_rows, latents, rotary_time, width_grid, audio_channels)
+            frame_time = _temporal_position_grid(frames, rotary_time)
+            position_ids[video_rows, 0] = frame_time.repeat_interleave(frame_grid.shape[0])
+            position_ids[video_rows, 1:] = frame_grid.repeat(frames, 1)
+            rotary_time += max(float(latents), _video_span_sequential(frames))
+        else:
+            raise ValueError(f"A reference must be an 'image', a 'video' or an 'audio', got {reference.kind!r}.")
+
+    # 3. The generated rows, on the timeline the references left behind — the same origin the keyframes anchored to.
+    audio_start = cursor
+    video_start = audio_start + num_target_audio_rows
+    _fill_audio_positions(
+        position_ids,
+        slice(audio_start, video_start),
+        num_audio_latents,
+        target_origin,
+        target_width_grid,
+        audio_channels,
+    )
+    frame_time = _temporal_position_grid(num_latent_frames, target_origin)
+    position_ids[video_start:, 0] = frame_time.repeat_interleave(target_frame_grid.shape[0])
+    position_ids[video_start:, 1:] = target_frame_grid.repeat(num_latent_frames, 1)
+
+    video_indices = torch.cat(video_indices + [torch.arange(video_start, sequence_length)])
+    audio_indices = torch.cat(audio_indices + [torch.arange(audio_start, video_start)])
+    text_indices = torch.arange(num_text_tokens)
+
+    token_tags = torch.empty(sequence_length, dtype=torch.long)
+    token_tags[text_indices] = text_token_tags.to(torch.long)
+    token_tags[audio_indices] = audio_tag
+    token_tags[video_indices] = video_tag
+
+    return (
+        position_ids,
+        token_tags,
+        video_indices,
+        audio_indices,
+        text_indices,
+        num_keyframe_rows + num_reference_video_rows,
+        num_reference_audio_rows,
+    )
+
+
+def validate_against_stock(verbose: bool = True) -> dict:
+    """Both degenerate cases must reproduce the shipped builders exactly."""
+
+    class _Ref:
+        def __init__(self, kind, has_audio=False):
+            self.kind, self.has_audio = kind, has_audio
+
+    geometry = dict(num_latent_frames=8, latent_height=34, latent_width=60, num_audio_latents=200,
+                    patch_size=(1, 2, 2), audio_channels=2, audio_tag=2, video_tag=0)
+    tags = torch.randint(0, 2, (57,))
+    report = {}
+
+    # (a) references only -> the stock ref2va layout
+    refs = [_Ref("image"), _Ref("video", has_audio=True), _Ref("audio")]
+    ref_latents = [torch.zeros(1, 16, 1, 32, 32), torch.zeros(1, 16, 3, 34, 60)]
+    ref_audio = [torch.zeros(60, 8), torch.zeros(40, 8)]
+    mine = build_combined_packed_sequence(tags, refs, ref_latents, ref_audio, **geometry)
+    stock = MiniMaxH3Ref2VAPrepareLayoutStep.build_ref2va_packed_sequence(
+        tags, refs, ref_latents, ref_audio, **geometry)
+    report["refs_only"] = _compare(mine, stock)
+
+    # (b) keyframes only -> the stock fl2va layout
+    anchors = ("first", "last")
+    kf_latents = [torch.zeros(1, 16, 1, 34, 60), torch.zeros(1, 16, 1, 34, 60)]
+    mine = build_combined_packed_sequence(tags, [], kf_latents, [], keyframe_anchors=anchors, **geometry)
+    stock = MiniMaxH3PrepareLayoutStep.build_packed_sequence(tags, keyframe_anchors=anchors, **geometry)
+    report["keyframes_only"] = _compare(mine, stock)
+
+    if verbose:
+        for case, result in report.items():
+            print(f"{case}: {result}")
+    return report
+
+
+def _compare(mine, stock) -> dict:
+    names = ["position_ids", "token_tags", "video_indices", "audio_indices", "text_indices",
+             "num_condition_video_rows", "num_condition_audio_rows"]
+    out = {}
+    for name, a, b in zip(names, mine, stock):
+        if isinstance(a, torch.Tensor):
+            out[name] = "identical" if a.shape == b.shape and torch.equal(a, b) else f"DIFFERS {tuple(a.shape)} vs {tuple(b.shape)}"
+        else:
+            out[name] = "identical" if a == b else f"DIFFERS {a} vs {b}"
+    return out
+
+
+if __name__ == "__main__":
+    report = validate_against_stock()
+    bad = {c: {k: v for k, v in r.items() if v != "identical"} for c, r in report.items()}
+    bad = {c: v for c, v in bad.items() if v}
+    print("\nVALIDATION", "PASSED" if not bad else f"FAILED: {bad}")