905 lines
39 KiB
Python
905 lines
39 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
|
|
"""MiniMax H3 text-conditioning contract."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import math
|
|
from collections.abc import Mapping, Sequence
|
|
from dataclasses import dataclass, fields
|
|
from typing import Any
|
|
|
|
import torch
|
|
|
|
MINIMAX_H3_TEXT_CONDITIONING_SCHEMA = "minimax_h3.text_conditioning/v1"
|
|
MINIMAX_H3_TEXT_HIDDEN_SIZE = 5120
|
|
MINIMAX_H3_PRESENTATION_TASK_KEY = "_minimax_h3_presentation_task"
|
|
MINIMAX_H3_CONDITION_LABELS_KEY = "_minimax_h3_condition_labels"
|
|
|
|
|
|
def _validate_contiguous_strided_layout(name: str, tensor: torch.Tensor) -> None:
|
|
if tensor.layout != torch.strided:
|
|
raise ValueError(
|
|
f"{MINIMAX_H3_TEXT_CONDITIONING_SCHEMA}: {name} must use contiguous strided layout, got {tensor.layout}"
|
|
)
|
|
if not tensor.is_contiguous():
|
|
raise ValueError(
|
|
f"{MINIMAX_H3_TEXT_CONDITIONING_SCHEMA}: {name} must use contiguous "
|
|
f"strided layout, got stride={tuple(tensor.stride())}"
|
|
)
|
|
|
|
|
|
@dataclass(frozen=True)
|
|
class MiniMaxH3TextConditioning:
|
|
"""``minimax_h3.text_conditioning/v1`` semantic payload."""
|
|
|
|
hidden_states: torch.Tensor
|
|
token_tags: torch.Tensor
|
|
|
|
@classmethod
|
|
def from_payload(
|
|
cls,
|
|
payload: Mapping[str, Any],
|
|
) -> MiniMaxH3TextConditioning:
|
|
"""Validate the semantic payload consumed by the diffusion stage."""
|
|
hidden_states = payload.get("hidden_states")
|
|
token_tags = payload.get("token_tags")
|
|
if not isinstance(hidden_states, torch.Tensor) or not isinstance(token_tags, torch.Tensor):
|
|
raise ValueError(
|
|
f"{MINIMAX_H3_TEXT_CONDITIONING_SCHEMA}: conditioning requires hidden_states and token_tags tensors"
|
|
)
|
|
if hidden_states.ndim != 2 or hidden_states.shape[-1] != MINIMAX_H3_TEXT_HIDDEN_SIZE:
|
|
raise ValueError(
|
|
f"{MINIMAX_H3_TEXT_CONDITIONING_SCHEMA}: hidden_states must have shape "
|
|
f"[tokens, {MINIMAX_H3_TEXT_HIDDEN_SIZE}], got {tuple(hidden_states.shape)}"
|
|
)
|
|
if hidden_states.dtype != torch.bfloat16:
|
|
raise ValueError(
|
|
f"{MINIMAX_H3_TEXT_CONDITIONING_SCHEMA}: hidden_states must have dtype "
|
|
f"torch.bfloat16, got {hidden_states.dtype}"
|
|
)
|
|
_validate_contiguous_strided_layout("hidden_states", hidden_states)
|
|
if token_tags.ndim != 1 or token_tags.shape[0] != hidden_states.shape[0]:
|
|
raise ValueError(
|
|
f"{MINIMAX_H3_TEXT_CONDITIONING_SCHEMA}: token_tags must align with hidden_states, got "
|
|
f"token_tags={tuple(token_tags.shape)} and hidden_states={tuple(hidden_states.shape)}"
|
|
)
|
|
if token_tags.dtype != torch.int64:
|
|
raise ValueError(
|
|
f"{MINIMAX_H3_TEXT_CONDITIONING_SCHEMA}: token_tags must have dtype torch.int64, got {token_tags.dtype}"
|
|
)
|
|
_validate_contiguous_strided_layout("token_tags", token_tags)
|
|
if not torch.all((token_tags == 0) | (token_tags == 1)):
|
|
raise ValueError(
|
|
f"{MINIMAX_H3_TEXT_CONDITIONING_SCHEMA}: text-encoder token_tags must contain only 0 and 1"
|
|
)
|
|
return cls(hidden_states=hidden_states, token_tags=token_tags)
|
|
|
|
@classmethod
|
|
def from_omni_payload(
|
|
cls,
|
|
payload: Mapping[str, Any],
|
|
) -> MiniMaxH3TextConditioning:
|
|
"""Validate and adapt the existing ``OmniPayload`` stage-wire view."""
|
|
hidden_states = payload.get("hidden_states")
|
|
if not isinstance(hidden_states, Mapping):
|
|
raise ValueError(f"{MINIMAX_H3_TEXT_CONDITIONING_SCHEMA}: text encoder returned no hidden_states payload")
|
|
hidden = hidden_states.get("output")
|
|
if not isinstance(hidden, torch.Tensor):
|
|
raise ValueError(
|
|
f"{MINIMAX_H3_TEXT_CONDITIONING_SCHEMA}: text encoder returned no hidden_states.output tensor"
|
|
)
|
|
|
|
meta = payload.get("meta")
|
|
if not isinstance(meta, Mapping):
|
|
raise ValueError(f"{MINIMAX_H3_TEXT_CONDITIONING_SCHEMA}: text encoder returned no conditioning metadata")
|
|
token_role_ids = meta.get("token_role_ids")
|
|
if not isinstance(token_role_ids, torch.Tensor):
|
|
raise ValueError(f"{MINIMAX_H3_TEXT_CONDITIONING_SCHEMA}: text encoder returned no token_role_ids tensor")
|
|
if token_role_ids.ndim != 2 or token_role_ids.shape[-1] != 1:
|
|
raise ValueError(
|
|
f"{MINIMAX_H3_TEXT_CONDITIONING_SCHEMA}: stage-wire token_role_ids must have shape "
|
|
f"[tokens, 1], got {tuple(token_role_ids.shape)}"
|
|
)
|
|
if token_role_ids.dtype != torch.int64:
|
|
raise ValueError(
|
|
f"{MINIMAX_H3_TEXT_CONDITIONING_SCHEMA}: stage-wire token_role_ids must have dtype "
|
|
f"torch.int64, got {token_role_ids.dtype}"
|
|
)
|
|
_validate_contiguous_strided_layout("stage-wire token_role_ids", token_role_ids)
|
|
|
|
return cls.from_payload(
|
|
{
|
|
"hidden_states": hidden,
|
|
"token_tags": token_role_ids.squeeze(-1),
|
|
}
|
|
)
|
|
|
|
def to_payload(self) -> dict[str, torch.Tensor]:
|
|
return {
|
|
"hidden_states": self.hidden_states,
|
|
"token_tags": self.token_tags,
|
|
}
|
|
|
|
|
|
MINIMAX_H3_ENCODER_REQUEST_KEY = "minimax_h3_encoder_request"
|
|
MINIMAX_H3_ENCODER_LAYOUT_KEY = "minimax_h3_encoder_layout"
|
|
STAGE_SCHEMA_VERSION = 2
|
|
_ENCODER_WIRE_SCHEMA_ID = 1
|
|
|
|
_TASK_TO_CODE = {"t2va": 1, "fl2va": 2, "ref2va": 3}
|
|
_CODE_TO_TASK = {value: key for key, value in _TASK_TO_CODE.items()}
|
|
_REF_KIND_TO_CODE = {"image": 1, "video": 2, "video_audio": 3, "audio": 4}
|
|
_CODE_TO_REF_KIND = {value: key for key, value in _REF_KIND_TO_CODE.items()}
|
|
_INTEGER_DTYPES = frozenset(
|
|
{
|
|
torch.int8,
|
|
torch.int16,
|
|
torch.int32,
|
|
torch.int64,
|
|
torch.uint8,
|
|
}
|
|
)
|
|
_VIDEO_CONDITION_WIDTH = 96
|
|
_AUDIO_CONDITION_WIDTH = 32
|
|
_AUDIO_CONDITION_CHANNELS = 2
|
|
|
|
|
|
def _validate_edit_mask_structure(
|
|
mask: torch.Tensor,
|
|
*,
|
|
name: str,
|
|
shape: tuple[int, ...],
|
|
) -> None:
|
|
"""Validate only the internal tensor contract after encoder-boundary parsing."""
|
|
if mask.dtype != torch.float32 or tuple(mask.shape) != shape:
|
|
raise ValueError(f"MiniMax H3 {name} edit mask must be FP32 with shape {shape}")
|
|
_validate_contiguous_strided_layout(f"{name} edit mask", mask)
|
|
|
|
|
|
def _validate_edit_tensors(
|
|
clean_rows: torch.Tensor | None,
|
|
raw_mask: torch.Tensor | None,
|
|
*,
|
|
name: str,
|
|
width: int,
|
|
mask_shape: tuple[int, ...],
|
|
) -> int:
|
|
values = (clean_rows, raw_mask)
|
|
if all(value is None for value in values):
|
|
return 0
|
|
if clean_rows is None or raw_mask is None:
|
|
raise ValueError(f"MiniMax H3 {name} edit requires clean rows and a raw mask tensor")
|
|
if clean_rows.dtype != torch.float32 or clean_rows.ndim != 2 or clean_rows.shape[1] != width:
|
|
raise ValueError(f"MiniMax H3 {name} edit rows must be FP32 with shape [rows, {width}]")
|
|
row_count = int(clean_rows.shape[0])
|
|
if row_count <= 0:
|
|
raise ValueError(f"MiniMax H3 {name} edit rows must not be empty")
|
|
_validate_edit_mask_structure(raw_mask, name=name, shape=mask_shape)
|
|
if raw_mask.device != clean_rows.device:
|
|
raise ValueError(f"MiniMax H3 {name} edit tensors must share one device")
|
|
_validate_contiguous_strided_layout(f"{name} edit rows", clean_rows)
|
|
return row_count
|
|
|
|
|
|
def _packed_vector_rows(value: torch.Tensor, *, width: int) -> torch.Tensor:
|
|
flat = value.reshape(-1)
|
|
padding = (-int(flat.numel())) % width
|
|
if padding:
|
|
flat = torch.cat((flat, flat.new_zeros(padding)))
|
|
return flat.reshape(-1, width)
|
|
|
|
|
|
def _pack_condition_slot(
|
|
condition: torch.Tensor | None,
|
|
clean_rows: torch.Tensor | None,
|
|
raw_mask: torch.Tensor | None,
|
|
*,
|
|
width: int,
|
|
) -> torch.Tensor:
|
|
if clean_rows is None:
|
|
return condition if condition is not None else torch.empty((0,), dtype=torch.float32)
|
|
if raw_mask is None:
|
|
raise ValueError("MiniMax H3 edit rows require a raw mask tensor")
|
|
values = [clean_rows, _packed_vector_rows(raw_mask, width=width)]
|
|
if condition is not None:
|
|
values.insert(0, condition)
|
|
return torch.cat(values).contiguous()
|
|
|
|
|
|
def _unpack_condition_slot(
|
|
value: torch.Tensor,
|
|
*,
|
|
condition_rows: int,
|
|
edit_rows: int,
|
|
mask_shape: tuple[int, ...],
|
|
width: int,
|
|
name: str,
|
|
) -> tuple[torch.Tensor | None, torch.Tensor | None, torch.Tensor | None]:
|
|
mask_values = math.prod(mask_shape) if mask_shape else int(edit_rows > 0)
|
|
mask_wire_rows = (mask_values + width - 1) // width
|
|
expected_rows = condition_rows + edit_rows + mask_wire_rows
|
|
if expected_rows == 0:
|
|
if value.dtype != torch.float32 or tuple(value.shape) != (0,):
|
|
raise ValueError(f"MiniMax H3 empty {name} condition slot must be FP32 with shape [0]")
|
|
return None, None, None
|
|
if value.dtype != torch.float32 or value.ndim != 2 or tuple(value.shape) != (expected_rows, width):
|
|
raise ValueError(
|
|
f"MiniMax H3 packed {name} condition must be FP32 with shape "
|
|
f"[{expected_rows}, {width}], got {value.dtype} {tuple(value.shape)}"
|
|
)
|
|
cursor = 0
|
|
condition = value[:condition_rows].contiguous() if condition_rows else None
|
|
cursor += condition_rows
|
|
if edit_rows == 0:
|
|
return condition, None, None
|
|
clean_rows = value[cursor : cursor + edit_rows].contiguous()
|
|
cursor += edit_rows
|
|
raw_mask = value[cursor : cursor + mask_wire_rows].reshape(-1)[:mask_values].clone()
|
|
if mask_shape:
|
|
raw_mask = raw_mask.reshape(mask_shape)
|
|
else:
|
|
raw_mask = raw_mask.reshape(())
|
|
return condition, clean_rows, raw_mask
|
|
|
|
|
|
def _wire_layout(payload: Mapping[str, Any]) -> torch.Tensor:
|
|
private_metadata = payload.get("kv_metadata")
|
|
if not isinstance(private_metadata, Mapping):
|
|
raise ValueError("MiniMax H3 wire payload requires private layout metadata")
|
|
layout = private_metadata.get(MINIMAX_H3_ENCODER_LAYOUT_KEY)
|
|
if isinstance(layout, torch.Tensor):
|
|
return layout
|
|
raise ValueError("MiniMax H3 wire payload requires one packed encoder layout tensor")
|
|
|
|
|
|
def _validate_condition_tensors(
|
|
visual: torch.Tensor | None,
|
|
visual_shapes: Sequence[Sequence[int]],
|
|
audio: torch.Tensor | None,
|
|
audio_lengths: Sequence[int],
|
|
) -> None:
|
|
if (visual is None) != (not visual_shapes):
|
|
raise ValueError("MiniMax H3 visual condition and shapes must be present together")
|
|
if (audio is None) != (not audio_lengths):
|
|
raise ValueError("MiniMax H3 audio condition and lengths must be present together")
|
|
|
|
parsed_shapes = [tuple(int(item) for item in shape) for shape in visual_shapes]
|
|
if any(
|
|
len(shape) != 3 or any(item <= 0 for item in shape) or shape[1] % 2 or shape[2] % 2 for shape in parsed_shapes
|
|
):
|
|
raise ValueError("MiniMax H3 visual condition shapes must be positive [T, H, W] triplets with even H/W")
|
|
if visual is not None:
|
|
expected_rows = sum(t * (h // 2) * (w // 2) for t, h, w in parsed_shapes)
|
|
if (
|
|
visual.dtype != torch.float32
|
|
or visual.ndim != 2
|
|
or tuple(visual.shape)
|
|
!= (
|
|
expected_rows,
|
|
_VIDEO_CONDITION_WIDTH,
|
|
)
|
|
):
|
|
raise ValueError(
|
|
"MiniMax H3 visual condition must be FP32 with shape "
|
|
f"[{expected_rows}, {_VIDEO_CONDITION_WIDTH}], got {visual.dtype} {tuple(visual.shape)}"
|
|
)
|
|
|
|
parsed_lengths = [int(length) for length in audio_lengths]
|
|
if any(length <= 0 for length in parsed_lengths):
|
|
raise ValueError("MiniMax H3 audio condition lengths must be positive")
|
|
if audio is not None:
|
|
expected_rows = _AUDIO_CONDITION_CHANNELS * sum(parsed_lengths)
|
|
if (
|
|
audio.dtype != torch.float32
|
|
or audio.ndim != 2
|
|
or tuple(audio.shape)
|
|
!= (
|
|
expected_rows,
|
|
_AUDIO_CONDITION_WIDTH,
|
|
)
|
|
):
|
|
raise ValueError(
|
|
"MiniMax H3 audio condition must be FP32 with shape "
|
|
f"[{expected_rows}, {_AUDIO_CONDITION_WIDTH}], got {audio.dtype} {tuple(audio.shape)}"
|
|
)
|
|
|
|
|
|
def _ref_block_rows(ref_blocks: Sequence[Mapping[str, Any]]) -> list[list[int]]:
|
|
rows: list[list[int]] = []
|
|
for block in ref_blocks:
|
|
kind = str(block.get("kind") or "")
|
|
kind_code = _REF_KIND_TO_CODE.get(kind)
|
|
if kind_code is None:
|
|
raise ValueError(f"unsupported MiniMax H3 ref kind {kind!r}")
|
|
rows.append(
|
|
[
|
|
kind_code,
|
|
int(block.get("ref_audio_t", 0)),
|
|
int(block.get("latent_t", 0)),
|
|
int(block.get("latent_h", 0)),
|
|
int(block.get("latent_w", 0)),
|
|
]
|
|
)
|
|
return rows
|
|
|
|
|
|
def _decode_ref_blocks(value: torch.Tensor) -> tuple[dict[str, Any], ...]:
|
|
if value.ndim != 2 or value.shape[1] != 5:
|
|
raise ValueError("MiniMax H3 ref_blocks metadata must have shape [N, 5]")
|
|
blocks: list[dict[str, Any]] = []
|
|
for kind_code, ref_audio_t, ref_t, ref_h, ref_w in value.detach().cpu().tolist():
|
|
kind = _CODE_TO_REF_KIND.get(int(kind_code))
|
|
if kind is None:
|
|
raise ValueError(f"unsupported MiniMax H3 ref kind code {kind_code!r}")
|
|
block: dict[str, Any] = {"kind": kind}
|
|
for key, item in (
|
|
("ref_audio_t", ref_audio_t),
|
|
("latent_t", ref_t),
|
|
("latent_h", ref_h),
|
|
("latent_w", ref_w),
|
|
):
|
|
if item:
|
|
block[key] = int(item)
|
|
blocks.append(block)
|
|
return tuple(blocks)
|
|
|
|
|
|
@dataclass(frozen=True)
|
|
class MiniMaxH3EncoderMediaInput:
|
|
task: str
|
|
height: int
|
|
width: int
|
|
num_frames: int
|
|
latent_t: int
|
|
audio_t: int
|
|
images: tuple[torch.Tensor, ...] = ()
|
|
videos: tuple[torch.Tensor, ...] = ()
|
|
video_audios: tuple[tuple[torch.Tensor, int] | None, ...] = ()
|
|
audios: tuple[tuple[torch.Tensor, int], ...] = ()
|
|
keyframe_frame_indices: tuple[int, ...] = ()
|
|
audio_mode: str = "native"
|
|
video_edit: torch.Tensor | None = None
|
|
video_edit_mask: torch.Tensor | None = None
|
|
audio_edit: tuple[torch.Tensor, int] | None = None
|
|
audio_edit_mask: torch.Tensor | None = None
|
|
|
|
@classmethod
|
|
def from_mm_tensors(
|
|
cls,
|
|
values: Sequence[torch.Tensor],
|
|
metadata: Mapping[str, Any],
|
|
) -> MiniMaxH3EncoderMediaInput:
|
|
if not isinstance(values, Sequence) or isinstance(values, (str, bytes)):
|
|
raise ValueError("MiniMax H3 encoder media must be a tensor sequence")
|
|
tensors = list(values)
|
|
if any(not isinstance(item, torch.Tensor) for item in tensors):
|
|
raise ValueError("MiniMax H3 encoder media input must contain only tensors")
|
|
task = str(metadata["task"])
|
|
if task not in _TASK_TO_CODE:
|
|
raise ValueError(f"unsupported MiniMax H3 task {task!r}")
|
|
height = int(metadata["height"])
|
|
width = int(metadata["width"])
|
|
num_frames = int(metadata["num_frames"])
|
|
latent_t = int(metadata["latent_t"])
|
|
audio_t = int(metadata["audio_t"])
|
|
image_count = int(metadata.get("image_count", 0))
|
|
video_count = int(metadata.get("video_count", 0))
|
|
audio_count = int(metadata.get("audio_count", 0))
|
|
if min(height, width, num_frames, latent_t, audio_t) <= 0:
|
|
raise ValueError("MiniMax H3 encoder media dimensions must be positive")
|
|
if min(image_count, video_count, audio_count) < 0:
|
|
raise ValueError("MiniMax H3 encoder media counts must be non-negative")
|
|
|
|
audio_flags = tuple(bool(value) for value in metadata.get("video_audio_flags", ()))
|
|
video_audio_sample_rates = tuple(int(value) for value in metadata.get("video_audio_sample_rates", ()))
|
|
audio_sample_rates = tuple(int(value) for value in metadata.get("audio_sample_rates", ()))
|
|
if len(audio_flags) != video_count or len(video_audio_sample_rates) != video_count:
|
|
raise ValueError("MiniMax H3 video audio flags must align with videos")
|
|
if len(audio_sample_rates) != audio_count:
|
|
raise ValueError("MiniMax H3 audio sample rates must align with audio inputs")
|
|
if any((rate > 0) != has_audio for rate, has_audio in zip(video_audio_sample_rates, audio_flags, strict=True)):
|
|
raise ValueError("MiniMax H3 embedded audio sample rates must align with audio flags")
|
|
if any(rate <= 0 for rate in audio_sample_rates):
|
|
raise ValueError("MiniMax H3 audio sample rates must be positive")
|
|
|
|
cursor = 0
|
|
images = tuple(tensors[cursor : cursor + image_count])
|
|
cursor += image_count
|
|
videos = tuple(tensors[cursor : cursor + video_count])
|
|
cursor += video_count
|
|
if len(images) != image_count or any(image.ndim != 3 or image.shape[-1] != 3 for image in images):
|
|
raise ValueError("MiniMax H3 encoder images must have shape [H, W, 3]")
|
|
if len(videos) != video_count or any(video.ndim != 4 or video.shape[-1] != 3 for video in videos):
|
|
raise ValueError("MiniMax H3 encoder videos must have shape [T, H, W, 3]")
|
|
|
|
video_audios: list[tuple[torch.Tensor, int] | None] = []
|
|
for has_audio, sample_rate in zip(audio_flags, video_audio_sample_rates, strict=True):
|
|
if not has_audio:
|
|
video_audios.append(None)
|
|
continue
|
|
if cursor >= len(tensors):
|
|
raise ValueError("MiniMax H3 encoder media input is truncated")
|
|
waveform = tensors[cursor]
|
|
cursor += 1
|
|
if waveform.ndim not in (1, 2):
|
|
raise ValueError("MiniMax H3 embedded audio must have shape [samples] or [channels, samples]")
|
|
video_audios.append((waveform, sample_rate))
|
|
|
|
audios: list[tuple[torch.Tensor, int]] = []
|
|
for sample_rate in audio_sample_rates:
|
|
if cursor >= len(tensors):
|
|
raise ValueError("MiniMax H3 encoder media input is truncated")
|
|
waveform = tensors[cursor]
|
|
cursor += 1
|
|
if waveform.ndim not in (1, 2):
|
|
raise ValueError("MiniMax H3 audio must have shape [samples] or [channels, samples]")
|
|
audios.append((waveform, sample_rate))
|
|
|
|
video_edit = None
|
|
video_edit_mask = None
|
|
if bool(metadata.get("has_video_edit", False)):
|
|
if cursor + 2 > len(tensors):
|
|
raise ValueError("MiniMax H3 encoder video edit input is truncated")
|
|
video_edit, video_edit_mask = tensors[cursor : cursor + 2]
|
|
cursor += 2
|
|
if tuple(video_edit.shape) != (num_frames, height, width, 3):
|
|
raise ValueError(
|
|
"MiniMax H3 encoder video edit must have shape "
|
|
f"[{num_frames}, {height}, {width}, 3], got {tuple(video_edit.shape)}"
|
|
)
|
|
_validate_edit_mask_structure(
|
|
video_edit_mask,
|
|
name="video",
|
|
shape=(latent_t, height // 16, width // 16),
|
|
)
|
|
|
|
audio_edit = None
|
|
audio_edit_mask = None
|
|
audio_edit_sample_rate = int(metadata.get("audio_edit_sample_rate", 0))
|
|
if audio_edit_sample_rate:
|
|
if cursor + 2 > len(tensors):
|
|
raise ValueError("MiniMax H3 encoder audio edit input is truncated")
|
|
waveform, audio_edit_mask = tensors[cursor : cursor + 2]
|
|
cursor += 2
|
|
if waveform.ndim not in (1, 2) or audio_edit_sample_rate <= 0:
|
|
raise ValueError("MiniMax H3 encoder audio edit requires a waveform and positive sample rate")
|
|
_validate_edit_mask_structure(audio_edit_mask, name="audio", shape=(2, audio_t))
|
|
audio_edit = (waveform, audio_edit_sample_rate)
|
|
if cursor != len(tensors):
|
|
raise ValueError(f"MiniMax H3 encoder media input has {len(tensors) - cursor} trailing tensors")
|
|
return cls(
|
|
audio_mode=str(metadata.get("audio_mode", "native")),
|
|
task=task,
|
|
height=height,
|
|
width=width,
|
|
num_frames=num_frames,
|
|
latent_t=latent_t,
|
|
audio_t=audio_t,
|
|
images=images,
|
|
videos=videos,
|
|
video_audios=tuple(video_audios),
|
|
audios=tuple(audios),
|
|
keyframe_frame_indices=tuple(int(value) for value in metadata.get("keyframe_frame_indices", ())),
|
|
video_edit=video_edit,
|
|
video_edit_mask=video_edit_mask,
|
|
audio_edit=audio_edit,
|
|
audio_edit_mask=audio_edit_mask,
|
|
)
|
|
|
|
def to_mm_tensors(self) -> list[torch.Tensor]:
|
|
task_code = _TASK_TO_CODE.get(self.task)
|
|
if task_code is None:
|
|
raise ValueError(f"unsupported MiniMax H3 task {self.task!r}")
|
|
if min(self.height, self.width, self.num_frames, self.latent_t, self.audio_t) <= 0:
|
|
raise ValueError("MiniMax H3 encoder media dimensions must be positive")
|
|
if len(self.video_audios) != len(self.videos):
|
|
raise ValueError("MiniMax H3 video audio slots must align with videos")
|
|
if any(image.ndim != 3 or image.shape[-1] != 3 for image in self.images):
|
|
raise ValueError("MiniMax H3 encoder images must have shape [H, W, 3]")
|
|
if any(video.ndim != 4 or video.shape[-1] != 3 for video in self.videos):
|
|
raise ValueError("MiniMax H3 encoder videos must have shape [T, H, W, 3]")
|
|
audio_items = [item for item in self.video_audios if item is not None] + list(self.audios)
|
|
if any(waveform.ndim not in (1, 2) or int(sample_rate) <= 0 for waveform, sample_rate in audio_items):
|
|
raise ValueError("MiniMax H3 encoder audio must have a waveform and positive sample rate")
|
|
video_edit_values = (self.video_edit, self.video_edit_mask)
|
|
if any(value is not None for value in video_edit_values):
|
|
if self.video_edit is None or self.video_edit_mask is None:
|
|
raise ValueError("MiniMax H3 encoder video edit requires a source and mask")
|
|
if tuple(self.video_edit.shape) != (self.num_frames, self.height, self.width, 3):
|
|
raise ValueError(
|
|
"MiniMax H3 encoder video edit must have shape "
|
|
f"[{self.num_frames}, {self.height}, {self.width}, 3], got {tuple(self.video_edit.shape)}"
|
|
)
|
|
_validate_edit_mask_structure(
|
|
self.video_edit_mask,
|
|
name="video",
|
|
shape=(self.latent_t, self.height // 16, self.width // 16),
|
|
)
|
|
audio_edit_values = (self.audio_edit, self.audio_edit_mask)
|
|
if any(value is not None for value in audio_edit_values):
|
|
if self.audio_edit is None or self.audio_edit_mask is None:
|
|
raise ValueError("MiniMax H3 encoder audio edit requires a source and mask")
|
|
waveform, sample_rate = self.audio_edit
|
|
if waveform.ndim not in (1, 2) or int(sample_rate) <= 0:
|
|
raise ValueError("MiniMax H3 encoder audio edit requires a waveform and positive sample rate")
|
|
_validate_edit_mask_structure(self.audio_edit_mask, name="audio", shape=(2, self.audio_t))
|
|
tensors = [*self.images, *self.videos]
|
|
for item in self.video_audios:
|
|
if item is None:
|
|
continue
|
|
waveform, _sample_rate = item
|
|
tensors.append(waveform)
|
|
tensors.extend(waveform for waveform, _sample_rate in self.audios)
|
|
if self.video_edit is not None:
|
|
if self.video_edit_mask is None:
|
|
raise ValueError("MiniMax H3 encoder video edit requires a mask")
|
|
tensors.extend((self.video_edit, self.video_edit_mask))
|
|
if self.audio_edit is not None:
|
|
if self.audio_edit_mask is None:
|
|
raise ValueError("MiniMax H3 encoder audio edit requires a mask")
|
|
tensors.extend((self.audio_edit[0], self.audio_edit_mask))
|
|
return tensors
|
|
|
|
def to_metadata(self) -> dict[str, Any]:
|
|
metadata = {
|
|
"audio_mode": self.audio_mode,
|
|
"task": self.task,
|
|
"height": self.height,
|
|
"width": self.width,
|
|
"num_frames": self.num_frames,
|
|
"latent_t": self.latent_t,
|
|
"audio_t": self.audio_t,
|
|
"image_count": len(self.images),
|
|
"video_count": len(self.videos),
|
|
"audio_count": len(self.audios),
|
|
"video_audio_flags": [item is not None for item in self.video_audios],
|
|
"video_audio_sample_rates": [item[1] if item is not None else 0 for item in self.video_audios],
|
|
"audio_sample_rates": [sample_rate for _waveform, sample_rate in self.audios],
|
|
"keyframe_frame_indices": list(self.keyframe_frame_indices),
|
|
}
|
|
if self.video_edit is not None:
|
|
metadata["has_video_edit"] = True
|
|
if self.audio_edit is not None:
|
|
metadata["audio_edit_sample_rate"] = int(self.audio_edit[1])
|
|
return metadata
|
|
|
|
|
|
@dataclass(frozen=True)
|
|
class MiniMaxH3EncoderMediaConditioning:
|
|
task: str
|
|
height: int
|
|
width: int
|
|
num_frames: int
|
|
latent_t: int
|
|
audio_t: int
|
|
visual_condition: torch.Tensor | None = None
|
|
visual_condition_shapes: tuple[tuple[int, int, int], ...] = ()
|
|
audio_condition: torch.Tensor | None = None
|
|
audio_condition_lengths: tuple[int, ...] = ()
|
|
ref_blocks: tuple[dict[str, Any], ...] = ()
|
|
keyframe_frame_indices: tuple[int, ...] = ()
|
|
video_edit_clean_rows: torch.Tensor | None = None
|
|
video_edit_mask: torch.Tensor | None = None
|
|
audio_edit_clean_rows: torch.Tensor | None = None
|
|
audio_edit_mask: torch.Tensor | None = None
|
|
audio_edit_source_t: int = 0
|
|
|
|
def to_omni_components(
|
|
self,
|
|
) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
|
|
task_code = _TASK_TO_CODE.get(self.task)
|
|
if task_code is None:
|
|
raise ValueError(f"unsupported MiniMax H3 task {self.task!r}")
|
|
if min(self.height, self.width, self.num_frames, self.latent_t, self.audio_t) <= 0:
|
|
raise ValueError("MiniMax H3 encoder media dimensions must be positive")
|
|
_validate_condition_tensors(
|
|
self.visual_condition,
|
|
self.visual_condition_shapes,
|
|
self.audio_condition,
|
|
self.audio_condition_lengths,
|
|
)
|
|
video_mask_shape = (self.latent_t, self.height // 16, self.width // 16)
|
|
audio_mask_shape = (2, self.audio_t)
|
|
video_edit_rows = _validate_edit_tensors(
|
|
self.video_edit_clean_rows,
|
|
self.video_edit_mask,
|
|
name="video",
|
|
width=_VIDEO_CONDITION_WIDTH,
|
|
mask_shape=video_mask_shape,
|
|
)
|
|
audio_edit_rows = _validate_edit_tensors(
|
|
self.audio_edit_clean_rows,
|
|
self.audio_edit_mask,
|
|
name="audio",
|
|
width=_AUDIO_CONDITION_WIDTH,
|
|
mask_shape=audio_mask_shape,
|
|
)
|
|
expected_video_rows = self.latent_t * (self.height // 32) * (self.width // 32)
|
|
if video_edit_rows not in (0, expected_video_rows):
|
|
raise ValueError(f"MiniMax H3 video edit has {video_edit_rows} rows, expected {expected_video_rows}")
|
|
expected_audio_rows = self.audio_t * _AUDIO_CONDITION_CHANNELS
|
|
if audio_edit_rows not in (0, expected_audio_rows):
|
|
raise ValueError(f"MiniMax H3 audio edit has {audio_edit_rows} rows, expected {expected_audio_rows}")
|
|
if audio_edit_rows == 0:
|
|
if self.audio_edit_source_t != 0:
|
|
raise ValueError("MiniMax H3 audio_edit_source_t requires audio edit rows")
|
|
elif not 0 < self.audio_edit_source_t <= self.audio_t:
|
|
raise ValueError(f"MiniMax H3 audio_edit_source_t must be in [1, {self.audio_t}]")
|
|
ref_rows = _ref_block_rows(self.ref_blocks)
|
|
layout = [
|
|
_ENCODER_WIRE_SCHEMA_ID,
|
|
STAGE_SCHEMA_VERSION,
|
|
task_code,
|
|
self.height,
|
|
self.width,
|
|
self.num_frames,
|
|
self.latent_t,
|
|
self.audio_t,
|
|
len(self.visual_condition_shapes),
|
|
len(self.audio_condition_lengths),
|
|
len(ref_rows),
|
|
len(self.keyframe_frame_indices),
|
|
video_edit_rows,
|
|
audio_edit_rows,
|
|
self.audio_edit_source_t,
|
|
]
|
|
layout.extend(item for shape in self.visual_condition_shapes for item in shape)
|
|
layout.extend(self.audio_condition_lengths)
|
|
layout.extend(item for row in ref_rows for item in row)
|
|
layout.extend(self.keyframe_frame_indices)
|
|
return (
|
|
_pack_condition_slot(
|
|
self.visual_condition,
|
|
self.video_edit_clean_rows,
|
|
self.video_edit_mask,
|
|
width=_VIDEO_CONDITION_WIDTH,
|
|
),
|
|
_pack_condition_slot(
|
|
self.audio_condition,
|
|
self.audio_edit_clean_rows,
|
|
self.audio_edit_mask,
|
|
width=_AUDIO_CONDITION_WIDTH,
|
|
),
|
|
torch.tensor(layout, dtype=torch.int64),
|
|
)
|
|
|
|
|
|
@dataclass(frozen=True)
|
|
class MiniMaxH3EncoderConditioning:
|
|
hidden_states: torch.Tensor
|
|
token_tags: torch.Tensor
|
|
task: str
|
|
height: int
|
|
width: int
|
|
num_frames: int
|
|
latent_t: int
|
|
audio_t: int
|
|
visual_condition: torch.Tensor | None = None
|
|
visual_condition_shapes: tuple[tuple[int, int, int], ...] = ()
|
|
audio_condition: torch.Tensor | None = None
|
|
audio_condition_lengths: tuple[int, ...] = ()
|
|
ref_blocks: tuple[dict[str, Any], ...] = ()
|
|
keyframe_frame_indices: tuple[int, ...] = ()
|
|
video_edit_clean_rows: torch.Tensor | None = None
|
|
video_edit_mask: torch.Tensor | None = None
|
|
audio_edit_clean_rows: torch.Tensor | None = None
|
|
audio_edit_mask: torch.Tensor | None = None
|
|
audio_edit_source_t: int = 0
|
|
|
|
@classmethod
|
|
def from_components(
|
|
cls,
|
|
text: MiniMaxH3TextConditioning,
|
|
media: MiniMaxH3EncoderMediaConditioning,
|
|
) -> MiniMaxH3EncoderConditioning:
|
|
"""Combine local encoder results without a stage-wire round trip."""
|
|
text = MiniMaxH3TextConditioning.from_payload(text.to_payload())
|
|
_validate_condition_tensors(
|
|
media.visual_condition,
|
|
media.visual_condition_shapes,
|
|
media.audio_condition,
|
|
media.audio_condition_lengths,
|
|
)
|
|
return cls(
|
|
hidden_states=text.hidden_states,
|
|
token_tags=text.token_tags,
|
|
**{item.name: getattr(media, item.name) for item in fields(media)},
|
|
)
|
|
|
|
@classmethod
|
|
def from_omni_payload(cls, payload: Mapping[str, Any]) -> MiniMaxH3EncoderConditioning:
|
|
hidden_payload = payload.get("hidden_states")
|
|
embed_payload = payload.get("embed")
|
|
meta_payload = payload.get("meta")
|
|
hidden_states = hidden_payload.get("output") if isinstance(hidden_payload, Mapping) else None
|
|
token_tags = meta_payload.get("token_role_ids") if isinstance(meta_payload, Mapping) else None
|
|
visual = embed_payload.get("embedding") if isinstance(embed_payload, Mapping) else None
|
|
audio = embed_payload.get("speech_feat") if isinstance(embed_payload, Mapping) else None
|
|
if not all(isinstance(item, torch.Tensor) for item in (hidden_states, token_tags, visual, audio)):
|
|
raise ValueError("MiniMax H3 encoder wire requires text, visual and audio tensors")
|
|
assert isinstance(hidden_states, torch.Tensor)
|
|
assert isinstance(token_tags, torch.Tensor)
|
|
assert isinstance(visual, torch.Tensor)
|
|
assert isinstance(audio, torch.Tensor)
|
|
if token_tags.ndim == 2 and token_tags.shape[-1] == 1:
|
|
token_tags = token_tags.squeeze(-1)
|
|
text = MiniMaxH3TextConditioning.from_payload({"hidden_states": hidden_states, "token_tags": token_tags})
|
|
for name, value in (("visual", visual), ("audio", audio)):
|
|
if value.numel() == 0 and (value.dtype != torch.float32 or tuple(value.shape) != (0,)):
|
|
raise ValueError(f"MiniMax H3 empty {name} condition slot must be FP32 with shape [0]")
|
|
if not isinstance(meta_payload, Mapping):
|
|
raise ValueError("MiniMax H3 wire payload requires a meta mapping")
|
|
layout = _wire_layout(payload)
|
|
if layout.dtype not in _INTEGER_DTYPES or layout.ndim != 1:
|
|
raise ValueError("MiniMax H3 encoder layout must be a one-dimensional integer tensor")
|
|
values = [int(item) for item in layout.detach().cpu().tolist()]
|
|
if len(values) < 2:
|
|
raise ValueError("MiniMax H3 encoder layout header is truncated")
|
|
schema_id, version = values[:2]
|
|
if schema_id != _ENCODER_WIRE_SCHEMA_ID or version != STAGE_SCHEMA_VERSION:
|
|
raise ValueError(f"unsupported MiniMax H3 encoder wire schema {schema_id}:{version}")
|
|
if len(values) < 15:
|
|
raise ValueError("MiniMax H3 encoder layout header is truncated")
|
|
(
|
|
_schema_id,
|
|
_version,
|
|
task_code,
|
|
height,
|
|
width,
|
|
num_frames,
|
|
latent_t,
|
|
audio_t,
|
|
visual_count,
|
|
audio_count,
|
|
ref_count,
|
|
keyframe_count,
|
|
video_edit_rows,
|
|
audio_edit_rows,
|
|
audio_edit_source_t,
|
|
) = values[:15]
|
|
task = _CODE_TO_TASK.get(task_code)
|
|
if task is None:
|
|
raise ValueError(f"unsupported MiniMax H3 wire task code {task_code!r}")
|
|
if min(height, width, num_frames, latent_t, audio_t) <= 0:
|
|
raise ValueError("MiniMax H3 encoder wire dimensions must be positive")
|
|
if min(visual_count, audio_count, ref_count, keyframe_count) < 0:
|
|
raise ValueError("MiniMax H3 encoder layout counts must be non-negative")
|
|
if min(video_edit_rows, audio_edit_rows, audio_edit_source_t) < 0:
|
|
raise ValueError("MiniMax H3 encoder edit layout values must be non-negative")
|
|
expected_size = 15 + 3 * visual_count + audio_count + 5 * ref_count + keyframe_count
|
|
if len(values) != expected_size:
|
|
raise ValueError(f"MiniMax H3 encoder layout has {len(values)} integers, expected {expected_size}")
|
|
cursor = 15
|
|
visual_shapes = tuple(
|
|
tuple(values[cursor + 3 * index : cursor + 3 * (index + 1)]) for index in range(visual_count)
|
|
)
|
|
cursor += 3 * visual_count
|
|
audio_lengths = tuple(values[cursor : cursor + audio_count])
|
|
cursor += audio_count
|
|
ref_rows = values[cursor : cursor + 5 * ref_count]
|
|
cursor += 5 * ref_count
|
|
ref_blocks_tensor = torch.tensor(ref_rows, dtype=torch.int64).reshape(-1, 5)
|
|
keyframe_frame_indices = tuple(values[cursor : cursor + keyframe_count])
|
|
video_mask_shape = (latent_t, height // 16, width // 16) if video_edit_rows else ()
|
|
audio_mask_shape = (2, audio_t) if audio_edit_rows else ()
|
|
visual_condition_rows = sum(t * (h // 2) * (w // 2) for t, h, w in visual_shapes)
|
|
audio_condition_rows = _AUDIO_CONDITION_CHANNELS * sum(audio_lengths)
|
|
visual_condition, video_edit_clean_rows, video_edit_mask = _unpack_condition_slot(
|
|
visual,
|
|
condition_rows=visual_condition_rows,
|
|
edit_rows=video_edit_rows,
|
|
mask_shape=video_mask_shape,
|
|
width=_VIDEO_CONDITION_WIDTH,
|
|
name="visual",
|
|
)
|
|
audio_condition, audio_edit_clean_rows, audio_edit_mask = _unpack_condition_slot(
|
|
audio,
|
|
condition_rows=audio_condition_rows,
|
|
edit_rows=audio_edit_rows,
|
|
mask_shape=audio_mask_shape,
|
|
width=_AUDIO_CONDITION_WIDTH,
|
|
name="audio",
|
|
)
|
|
_validate_condition_tensors(
|
|
visual_condition,
|
|
visual_shapes,
|
|
audio_condition,
|
|
audio_lengths,
|
|
)
|
|
_validate_edit_tensors(
|
|
video_edit_clean_rows,
|
|
video_edit_mask,
|
|
name="video",
|
|
width=_VIDEO_CONDITION_WIDTH,
|
|
mask_shape=(latent_t, height // 16, width // 16),
|
|
)
|
|
_validate_edit_tensors(
|
|
audio_edit_clean_rows,
|
|
audio_edit_mask,
|
|
name="audio",
|
|
width=_AUDIO_CONDITION_WIDTH,
|
|
mask_shape=(2, audio_t),
|
|
)
|
|
expected_video_edit_rows = latent_t * (height // 32) * (width // 32)
|
|
if video_edit_rows not in (0, expected_video_edit_rows):
|
|
raise ValueError(f"MiniMax H3 video edit has {video_edit_rows} rows, expected {expected_video_edit_rows}")
|
|
if audio_edit_rows not in (0, audio_t * _AUDIO_CONDITION_CHANNELS):
|
|
raise ValueError(f"MiniMax H3 audio edit has {audio_edit_rows} rows, expected {audio_t * 2}")
|
|
if (audio_edit_rows == 0 and audio_edit_source_t != 0) or (
|
|
audio_edit_rows and not 0 < audio_edit_source_t <= audio_t
|
|
):
|
|
raise ValueError("MiniMax H3 encoder audio edit source length is inconsistent")
|
|
return cls(
|
|
hidden_states=text.hidden_states,
|
|
token_tags=text.token_tags,
|
|
task=task,
|
|
height=height,
|
|
width=width,
|
|
num_frames=num_frames,
|
|
latent_t=latent_t,
|
|
audio_t=audio_t,
|
|
visual_condition=visual_condition,
|
|
visual_condition_shapes=visual_shapes,
|
|
audio_condition=audio_condition,
|
|
audio_condition_lengths=audio_lengths,
|
|
ref_blocks=_decode_ref_blocks(ref_blocks_tensor),
|
|
keyframe_frame_indices=keyframe_frame_indices,
|
|
video_edit_clean_rows=video_edit_clean_rows,
|
|
video_edit_mask=video_edit_mask,
|
|
audio_edit_clean_rows=audio_edit_clean_rows,
|
|
audio_edit_mask=audio_edit_mask,
|
|
audio_edit_source_t=audio_edit_source_t,
|
|
)
|
|
|
|
def to_omni_payload(self) -> dict[str, Any]:
|
|
text = MiniMaxH3TextConditioning.from_payload(
|
|
{
|
|
"hidden_states": self.hidden_states,
|
|
"token_tags": self.token_tags,
|
|
}
|
|
)
|
|
visual, audio, layout = MiniMaxH3EncoderMediaConditioning(
|
|
task=self.task,
|
|
height=self.height,
|
|
width=self.width,
|
|
num_frames=self.num_frames,
|
|
latent_t=self.latent_t,
|
|
audio_t=self.audio_t,
|
|
visual_condition=self.visual_condition,
|
|
visual_condition_shapes=self.visual_condition_shapes,
|
|
audio_condition=self.audio_condition,
|
|
audio_condition_lengths=self.audio_condition_lengths,
|
|
ref_blocks=self.ref_blocks,
|
|
keyframe_frame_indices=self.keyframe_frame_indices,
|
|
video_edit_clean_rows=self.video_edit_clean_rows,
|
|
video_edit_mask=self.video_edit_mask,
|
|
audio_edit_clean_rows=self.audio_edit_clean_rows,
|
|
audio_edit_mask=self.audio_edit_mask,
|
|
audio_edit_source_t=self.audio_edit_source_t,
|
|
).to_omni_components()
|
|
return {
|
|
"hidden_states": {"output": text.hidden_states},
|
|
"embed": {
|
|
"embedding": visual,
|
|
"speech_feat": audio,
|
|
},
|
|
"meta": {"token_role_ids": text.token_tags},
|
|
"kv_metadata": {MINIMAX_H3_ENCODER_LAYOUT_KEY: layout},
|
|
}
|
|
|
|
|
|
__all__ = [
|
|
"MINIMAX_H3_CONDITION_LABELS_KEY",
|
|
"MINIMAX_H3_ENCODER_LAYOUT_KEY",
|
|
"MINIMAX_H3_ENCODER_REQUEST_KEY",
|
|
"MINIMAX_H3_PRESENTATION_TASK_KEY",
|
|
"MINIMAX_H3_TEXT_CONDITIONING_SCHEMA",
|
|
"MINIMAX_H3_TEXT_HIDDEN_SIZE",
|
|
"STAGE_SCHEMA_VERSION",
|
|
"MiniMaxH3EncoderConditioning",
|
|
"MiniMaxH3EncoderMediaConditioning",
|
|
"MiniMaxH3EncoderMediaInput",
|
|
"MiniMaxH3TextConditioning",
|
|
]
|