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