def build_ref2va_packed_sequence(
text_token_tags: torch.Tensor,
references: list[MiniMaxH3PreparedReference],
num_latent_frames: int,
latent_height: int,
latent_width: int,
num_audio_latents: int,
patch_size: tuple[int, int, int],
) -> MiniMaxH3PackedLayout:
"""Build `[text | ordered references | target audio | target video]`."""
if patch_size != (1, 2, 2):
raise ValueError(f"MiniMax-H3 Ref2VA requires patch_size=(1, 2, 2), got {patch_size}.")
if text_token_tags.ndim != 1:
raise ValueError(f"text_token_tags must be one-dimensional, got {tuple(text_token_tags.shape)}.")
valid_text_tags = (text_token_tags == MINIMAX_H3_TEXT_TAG) | (text_token_tags == MINIMAX_H3_VIDEO_TAG)
if not bool(valid_text_tags.all()):
raise ValueError("text_token_tags may contain only text and vision tags.")
if not references:
raise ValueError("Ref2VA requires at least one prepared reference.")
patch_t, patch_h, patch_w = patch_size
target_geometry = (num_latent_frames, latent_height, latent_width)
if any(value <= 0 for value in target_geometry) or num_audio_latents <= 0:
raise ValueError("Ref2VA target latent geometry must be positive.")
if any(value % patch for value, patch in zip(target_geometry, patch_size, strict=True)):
raise ValueError(f"Target geometry {target_geometry} is not divisible by patch {patch_size}.")
visual_row_counts = [
0 if reference.media_type == "audio" else _num_video_rows(reference, patch_size) for reference in references
]
for reference in references:
if reference.media_type == "audio" and not reference.has_audio:
raise ValueError("An audio reference must carry decoded waveform latents.")
if reference.has_audio and reference.num_audio_latents <= 0:
raise ValueError("An audio-bearing reference has no resolved audio latents.")
num_text_tokens = int(text_token_tags.shape[0])
num_target_video_rows = (num_latent_frames // patch_t) * (latent_height // patch_h) * (latent_width // patch_w)
num_target_audio_rows = num_audio_latents * MINIMAX_H3_AUDIO_CHANNELS
num_reference_video_rows = sum(visual_row_counts)
num_reference_audio_rows = sum(reference.num_audio_latents * MINIMAX_H3_AUDIO_CHANNELS for reference in references
if reference.has_audio)
sequence_length = (num_text_tokens + 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)
video_indices: list[torch.Tensor] = []
audio_indices: list[torch.Tensor] = []
cursor = num_text_tokens
rotary_time = float(num_text_tokens)
for reference, visual_row_count in zip(references, visual_row_counts, strict=True):
if reference.media_type == "image":
rows = slice(cursor, cursor + visual_row_count)
cursor = rows.stop
video_indices.append(torch.arange(rows.start, rows.stop))
frame_grid, _ = _frame_position_grid(
reference.latent_height,
reference.latent_width,
patch_h,
patch_w,
)
position_ids[rows, 0] = rotary_time
position_ids[rows, 1:] = frame_grid
rotary_time += 1.0
elif reference.media_type == "audio":
count = reference.num_audio_latents * MINIMAX_H3_AUDIO_CHANNELS
rows = slice(cursor, cursor + count)
cursor = rows.stop
audio_indices.append(torch.arange(rows.start, rows.stop))
_fill_audio_positions(position_ids, rows, reference.num_audio_latents, rotary_time, target_width_grid)
rotary_time += float(reference.num_audio_latents)
elif reference.media_type == "video":
audio_count = reference.num_audio_latents * MINIMAX_H3_AUDIO_CHANNELS if reference.has_audio else 0
audio_rows = slice(cursor, cursor + audio_count)
video_rows = slice(audio_rows.stop, audio_rows.stop + visual_row_count)
cursor = video_rows.stop
if audio_count:
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(
reference.latent_height,
reference.latent_width,
patch_h,
patch_w,
)
if audio_count:
_fill_audio_positions(
position_ids,
audio_rows,
reference.num_audio_latents,
rotary_time,
width_grid,
)
frame_time = temporal_position_grid(reference.num_latent_frames, rotary_time)
rows_per_frame = frame_grid.shape[0]
position_ids[video_rows, 0] = frame_time.repeat_interleave(rows_per_frame)
position_ids[video_rows, 1:] = frame_grid.repeat(reference.num_latent_frames // patch_t, 1)
rotary_time += max(
float(reference.num_audio_latents if reference.has_audio else 0),
_reference_temporal_span(reference.num_latent_frames),
)
else:
raise ValueError(f"Unsupported prepared reference type: {reference.media_type!r}.")
audio_start = cursor
video_start = audio_start + num_target_audio_rows
_fill_audio_positions(position_ids, slice(audio_start, video_start), num_audio_latents, rotary_time,
target_width_grid)
frame_time = temporal_position_grid(num_latent_frames, rotary_time)
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 // patch_t, 1)
video_indices.append(torch.arange(video_start, sequence_length))
audio_indices.append(torch.arange(audio_start, video_start))
packed_video_indices = torch.cat(video_indices)
packed_audio_indices = torch.cat(audio_indices)
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[packed_audio_indices] = MINIMAX_H3_AUDIO_TAG
token_tags[packed_video_indices] = MINIMAX_H3_VIDEO_TAG
return MiniMaxH3PackedLayout(
sequence_length=sequence_length,
position_ids=position_ids,
token_tags=token_tags,
video_indices=packed_video_indices,
audio_indices=packed_audio_indices,
text_indices=text_indices,
num_condition_video_rows=num_reference_video_rows,
num_condition_audio_rows=num_reference_audio_rows,
num_video_latent_frames=num_latent_frames,
latent_height=latent_height,
latent_width=latent_width,
num_audio_latents=num_audio_latents,
)