combined_layout.py

"""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}")