823 lines
31 KiB
Python
823 lines
31 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
|
|
"""MiniMax H3 remote-code VAE adapters and exact latent contracts."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import importlib
|
|
import json
|
|
from collections.abc import Iterator
|
|
from contextlib import AbstractContextManager, contextmanager, nullcontext
|
|
from pathlib import Path
|
|
from typing import Any, Protocol
|
|
|
|
import torch
|
|
import torch.distributed as dist
|
|
import torch.nn as nn
|
|
from PIL import Image
|
|
from transformers.dynamic_module_utils import get_class_from_dynamic_module
|
|
from vllm.logger import init_logger
|
|
|
|
from vllm_omni.diffusion.distributed.autoencoders.distributed_vae_executor import (
|
|
DistributedVaeMixin,
|
|
)
|
|
from vllm_omni.diffusion.distributed.parallel_state import (
|
|
get_data_parallel_world_size,
|
|
get_world_group,
|
|
model_parallel_is_initialized,
|
|
)
|
|
from vllm_omni.diffusion.offloader.module_residency import (
|
|
BoundedAllocatorCache,
|
|
PinnedModuleStager,
|
|
)
|
|
|
|
from .ops import install_h3_vae_optimizations
|
|
from .packed_tokens import minimax_h3_patchify_video_latent
|
|
from .temporal_chunks import (
|
|
MiniMaxH3TemporalChunkCompatibilityError,
|
|
MiniMaxH3TemporalDecodePlan,
|
|
decode_minimax_h3_temporal_chunks_compat,
|
|
prepare_minimax_h3_temporal_chunks_compat,
|
|
)
|
|
|
|
MINIMAX_H3_KEYFRAME_ENCODE_SEED = 42
|
|
MINIMAX_H3_AUDIO_SAMPLE_RATE = 32000
|
|
MINIMAX_H3_AUDIO_CHANNELS = 2
|
|
|
|
|
|
logger = init_logger(__name__)
|
|
|
|
|
|
class MiniMaxH3VideoChunkCallback(Protocol):
|
|
"""Consume one ordered, temporally committed H3 video chunk.
|
|
|
|
``frames`` is an owned, contiguous float32 RGB tensor in ``BCTHW`` layout
|
|
and the ``[0, 1]`` value range. The callback runs synchronously on the
|
|
producer's current device stream. Before reading it on another stream, the
|
|
consumer must establish ordering from the producer stream (for example via
|
|
an event) and retain the allocation for that stream's lifetime.
|
|
"""
|
|
|
|
def __call__(
|
|
self,
|
|
frames: torch.Tensor,
|
|
*,
|
|
chunk_index: int,
|
|
total_chunks: int,
|
|
frame_start: int,
|
|
is_final: bool,
|
|
) -> None: ...
|
|
|
|
|
|
class MiniMaxH3ChunkedDecodeUnsupportedError(RuntimeError):
|
|
"""Raised when loaded MiniMax H3 code cannot honor the chunk contract."""
|
|
|
|
|
|
class MiniMaxH3ChunkCallbackPeerError(RuntimeError):
|
|
"""Raised on peer VAE ranks when the output-rank callback fails."""
|
|
|
|
|
|
def _load_component_config(component_path: str) -> dict[str, Any]:
|
|
config_path = Path(component_path) / "config.json"
|
|
config = json.loads(config_path.read_text(encoding="utf-8"))
|
|
channels = int(config["latent_channels"])
|
|
for key in ("latents_mean", "latents_std"):
|
|
values = config.get(key)
|
|
if not isinstance(values, list) or len(values) != channels:
|
|
raise ValueError(f"{config_path}: {key} must contain {channels} values")
|
|
return config
|
|
|
|
|
|
def _load_remote_component(
|
|
component_path: str,
|
|
config: dict[str, Any],
|
|
) -> nn.Module:
|
|
auto_map = config.get("auto_map") or {}
|
|
class_reference = auto_map.get("AutoModel")
|
|
if not isinstance(class_reference, str):
|
|
raise ValueError(f"{component_path}/config.json must define auto_map.AutoModel")
|
|
component_cls = get_class_from_dynamic_module(
|
|
class_reference,
|
|
component_path,
|
|
)
|
|
# Build on the host regardless of the ambient default device. Online
|
|
# quantization wraps pipeline construction in a `with torch.device(<accel>)`
|
|
# block for the DiT's quantized linears, and the checkpoint's own VAE code
|
|
# builds constants with ops that have no accelerator kernel (BigVGAN's
|
|
# anti-aliasing filters call torch.kaiser_window). Callers place the module
|
|
# explicitly right after this returns, so nothing depends on the context.
|
|
with torch.device("cpu"):
|
|
return component_cls.from_pretrained(component_path)
|
|
|
|
|
|
class _AudioVAEDeterminismContext(AbstractContextManager):
|
|
def __enter__(self):
|
|
backends = torch.backends
|
|
self._saved = (
|
|
backends.cuda.matmul.allow_tf32,
|
|
backends.cudnn.allow_tf32,
|
|
backends.cudnn.benchmark,
|
|
backends.cudnn.deterministic,
|
|
backends.cudnn.enabled,
|
|
backends.cuda.flash_sdp_enabled(),
|
|
backends.cuda.mem_efficient_sdp_enabled(),
|
|
backends.cuda.math_sdp_enabled(),
|
|
)
|
|
backends.cuda.matmul.allow_tf32 = False
|
|
backends.cudnn.allow_tf32 = False
|
|
backends.cudnn.benchmark = False
|
|
backends.cudnn.deterministic = True
|
|
backends.cudnn.enabled = False
|
|
backends.cuda.enable_flash_sdp(False)
|
|
backends.cuda.enable_mem_efficient_sdp(False)
|
|
backends.cuda.enable_math_sdp(True)
|
|
return self
|
|
|
|
def __exit__(self, exc_type, exc, traceback):
|
|
backends = torch.backends
|
|
(
|
|
backends.cuda.matmul.allow_tf32,
|
|
backends.cudnn.allow_tf32,
|
|
backends.cudnn.benchmark,
|
|
backends.cudnn.deterministic,
|
|
backends.cudnn.enabled,
|
|
flash,
|
|
memory_efficient,
|
|
math_sdp,
|
|
) = self._saved
|
|
backends.cuda.enable_flash_sdp(flash)
|
|
backends.cuda.enable_mem_efficient_sdp(memory_efficient)
|
|
backends.cuda.enable_math_sdp(math_sdp)
|
|
return False
|
|
|
|
|
|
class MiniMaxH3VideoVAE(nn.Module, DistributedVaeMixin):
|
|
"""Adapter around the checkpoint's native parallel-tiled video VAE."""
|
|
|
|
def __init__(
|
|
self,
|
|
component_path: str,
|
|
*,
|
|
device: torch.device,
|
|
load_device: torch.device | None = None,
|
|
) -> None:
|
|
super().__init__()
|
|
self._device_target = device
|
|
self.config_dict = _load_component_config(component_path)
|
|
self.remote = _load_remote_component(
|
|
component_path,
|
|
self.config_dict,
|
|
)
|
|
# Match the reference loader contract before installing inference-only
|
|
# decoder fast paths. Keyframe encoding remains FP32; decoder Linear
|
|
# weights may be materialized in FP16 because reference decode casts
|
|
# those same tensors through CUDA autocast on every tile.
|
|
initial_device = load_device or device
|
|
self.remote.eval().to(device=initial_device, dtype=torch.float32)
|
|
decoder = getattr(self.remote.model, "decoder", None)
|
|
if decoder is not None:
|
|
install_h3_vae_optimizations(
|
|
decoder,
|
|
device=device,
|
|
)
|
|
self._stager = None
|
|
if initial_device.type == "cpu" and device.type not in ("cpu", "meta"):
|
|
self._stager = PinnedModuleStager(
|
|
self.remote,
|
|
device,
|
|
pin_memory=True,
|
|
)
|
|
self.model = self.remote.model
|
|
self.use_tiling = True
|
|
self.use_slicing = False
|
|
self.parallel_size = 1
|
|
self.device_module = torch.get_device_module()
|
|
|
|
def load_to_device(self) -> None:
|
|
if self._stager is not None:
|
|
self._stager.load()
|
|
else:
|
|
self.remote.to(self._device_target)
|
|
|
|
def set_omni_component_cache(self, cache: BoundedAllocatorCache | None) -> None:
|
|
self._omni_component_cache = cache
|
|
if self._stager is not None:
|
|
self._stager.set_cache_retention(cache)
|
|
|
|
def offload_to_cpu(self) -> None:
|
|
if self._stager is not None:
|
|
self._stager.offload()
|
|
else:
|
|
self.remote.to("cpu")
|
|
cache = getattr(self, "_omni_component_cache", None)
|
|
if cache is None:
|
|
torch.accelerator.empty_cache()
|
|
else:
|
|
cache.release_if_needed()
|
|
|
|
def set_parallel_size(
|
|
self,
|
|
parallel_size: int,
|
|
mode: str = "tile",
|
|
) -> None:
|
|
if mode != "tile":
|
|
raise ValueError(f"MiniMax H3 VAE supports its native tile parallel mode only, got {mode!r}")
|
|
group = get_world_group().device_group
|
|
world_size = dist.get_world_size(group)
|
|
rank = dist.get_rank(group)
|
|
parallel_size = int(parallel_size)
|
|
if parallel_size not in (1, world_size):
|
|
raise ValueError(
|
|
"MiniMax H3 native VAE patch parallelism currently requires "
|
|
"vae_patch_parallel_size=1 or the full DiT group size "
|
|
f"({world_size}), got {parallel_size}"
|
|
)
|
|
self.parallel_size = parallel_size
|
|
enabled = parallel_size > 1
|
|
|
|
state = self._native_parallel_state()
|
|
state.clear()
|
|
state.update(
|
|
group_size=parallel_size,
|
|
group_rank=rank if enabled else 0,
|
|
local_process_group=group if enabled else None,
|
|
sp_size=parallel_size,
|
|
sp_rank=rank if enabled else 0,
|
|
sp_enabled=enabled,
|
|
sp_process_group=group if enabled else None,
|
|
tp_size=1,
|
|
tp_rank=0,
|
|
)
|
|
self.model.parallel_tiling = enabled
|
|
|
|
def _native_parallel_state(self) -> dict[str, Any]:
|
|
"""Return the checkpoint's own mutable parallel-state dict."""
|
|
|
|
package = self.remote.__class__.__module__.rsplit(".", 1)[0]
|
|
parallel_module = importlib.import_module(f"{package}.parallel")
|
|
return parallel_module.get_parallel_state()
|
|
|
|
def _decoder_tile_count(self, latent: torch.Tensor) -> int:
|
|
"""Number of decoder tiles the checkpoint will split ``latent`` into.
|
|
|
|
Mirrors the checkpoint's ``decode_tiled``: the grid is computed from the
|
|
pixel-space dimensions, so it is a pure function of the latent shape and
|
|
resolves identically on every rank.
|
|
"""
|
|
|
|
ratio = int(self.model.vae_ratio)
|
|
rows, _, _ = self.model.split_tiles(int(latent.shape[-2]) * ratio, True)
|
|
cols, _, _ = self.model.split_tiles(int(latent.shape[-1]) * ratio, True)
|
|
return len(rows) * len(cols)
|
|
|
|
@contextmanager
|
|
def _rank_local_tiling(self) -> Iterator[None]:
|
|
"""Run one decode with tiling kept on this rank, then restore the group.
|
|
|
|
Used only when there are fewer tiles than ranks. Every rank then decodes
|
|
every tile, which is slower than sharing the work but is correct and
|
|
involves no collective.
|
|
"""
|
|
|
|
state = self._native_parallel_state()
|
|
saved_state = dict(state)
|
|
saved_tiling = self.model.parallel_tiling
|
|
state.update(
|
|
group_size=1,
|
|
group_rank=0,
|
|
local_process_group=None,
|
|
sp_size=1,
|
|
sp_rank=0,
|
|
sp_enabled=False,
|
|
sp_process_group=None,
|
|
tp_size=1,
|
|
tp_rank=0,
|
|
)
|
|
self.model.parallel_tiling = False
|
|
try:
|
|
yield
|
|
finally:
|
|
state.clear()
|
|
state.update(saved_state)
|
|
self.model.parallel_tiling = saved_tiling
|
|
|
|
def is_distributed_enabled(self) -> bool:
|
|
return self.parallel_size > 1 and dist.is_initialized()
|
|
|
|
def _denormalize_latent(self, latent: torch.Tensor) -> torch.Tensor:
|
|
channels = int(self.config_dict["latent_channels"])
|
|
mean = torch.tensor(
|
|
self.config_dict["latents_mean"],
|
|
device=latent.device,
|
|
dtype=latent.dtype,
|
|
).view(1, channels, 1, 1, 1)
|
|
std = torch.tensor(
|
|
self.config_dict["latents_std"],
|
|
device=latent.device,
|
|
dtype=latent.dtype,
|
|
).view(1, channels, 1, 1, 1)
|
|
return latent * std + mean
|
|
|
|
def _decode_tiling_context(self, latent: torch.Tensor) -> AbstractContextManager:
|
|
# The checkpoint hands rank r the tiles ``range(r, num_tiles, sp_size)``
|
|
# and then rejects an empty share inside the gather. A rank with no
|
|
# tiles raises and leaves the collective while the others block in it
|
|
# forever, so too few tiles hangs the whole stage rather than failing
|
|
# it. Tile count depends only on the latent shape, so every rank takes
|
|
# this branch together.
|
|
num_tiles = self._decoder_tile_count(latent)
|
|
if self.parallel_size > 1 and num_tiles < self.parallel_size:
|
|
logger.warning_once(
|
|
"MiniMax-H3 VAE decode splits into %d tile(s) but the tile group has "
|
|
"%d ranks; decoding rank-locally for this shape instead, which is "
|
|
"slower but avoids ranks without tiles hanging the collective.",
|
|
num_tiles,
|
|
self.parallel_size,
|
|
)
|
|
return self._rank_local_tiling()
|
|
return nullcontext()
|
|
|
|
def _normalize_decoded_frames(self, decoded: torch.Tensor) -> torch.Tensor:
|
|
frames = self.model.processor.revert_tensor(decoded)
|
|
if frames.ndim == 4:
|
|
frames = frames.unsqueeze(0).transpose(1, 2)
|
|
if frames.ndim != 5:
|
|
raise ValueError(f"unexpected decoded video shape {tuple(frames.shape)}")
|
|
return frames.float()
|
|
|
|
def _chunk_process_group(self) -> dist.ProcessGroup | None:
|
|
if not self.is_distributed_enabled():
|
|
return None
|
|
if model_parallel_is_initialized() and get_data_parallel_world_size() > 1:
|
|
raise MiniMaxH3ChunkedDecodeUnsupportedError(
|
|
"MiniMax-H3 temporal chunks do not yet support combining data parallelism with VAE patch parallelism"
|
|
)
|
|
group = self._native_parallel_state().get("sp_process_group")
|
|
if group is None or dist.get_world_size(group) != self.parallel_size:
|
|
raise RuntimeError("MiniMax-H3 VAE chunk decode has an invalid spatial-parallel group")
|
|
return group
|
|
|
|
@staticmethod
|
|
def _is_chunk_output_rank(
|
|
chunk_callback: MiniMaxH3VideoChunkCallback | None,
|
|
*,
|
|
group: dist.ProcessGroup | None,
|
|
device: torch.device,
|
|
) -> bool:
|
|
if group is None:
|
|
return chunk_callback is not None
|
|
|
|
group_rank = dist.get_rank(group)
|
|
callback_owners = torch.tensor(
|
|
[int(chunk_callback is not None)],
|
|
dtype=torch.int32,
|
|
device=device,
|
|
)
|
|
dist.all_reduce(callback_owners, group=group)
|
|
valid_owner = int(callback_owners.item()) == 1 and ((group_rank == 0) == (chunk_callback is not None))
|
|
if not valid_owner:
|
|
raise ValueError("MiniMax-H3 chunk decode requires a callback on exactly VAE group rank 0")
|
|
return group_rank == 0
|
|
|
|
@staticmethod
|
|
def _synchronize_chunk_callback_error(
|
|
callback_error: BaseException | None,
|
|
*,
|
|
group: dist.ProcessGroup | None,
|
|
device: torch.device,
|
|
) -> None:
|
|
if group is None:
|
|
if callback_error is not None:
|
|
raise callback_error
|
|
return
|
|
|
|
is_output_rank = dist.get_rank(group) == 0
|
|
failed = torch.tensor(
|
|
[int(is_output_rank and callback_error is not None)],
|
|
dtype=torch.int32,
|
|
device=device,
|
|
)
|
|
dist.broadcast(
|
|
failed,
|
|
src=dist.get_global_rank(group, 0),
|
|
group=group,
|
|
)
|
|
if int(failed.item()) == 0:
|
|
return
|
|
if is_output_rank:
|
|
assert callback_error is not None
|
|
raise callback_error
|
|
raise MiniMaxH3ChunkCallbackPeerError("MiniMax-H3 video chunk callback failed on this VAE group's output rank")
|
|
|
|
@staticmethod
|
|
def _synchronize_chunk_preflight_error(
|
|
preflight_error: Exception | None,
|
|
temporal_dtype: torch.dtype | None,
|
|
*,
|
|
group: dist.ProcessGroup | None,
|
|
device: torch.device,
|
|
) -> None:
|
|
if group is None:
|
|
if preflight_error is not None:
|
|
raise preflight_error
|
|
return
|
|
failed = torch.tensor(
|
|
[int(preflight_error is not None)],
|
|
dtype=torch.int32,
|
|
device=device,
|
|
)
|
|
dist.all_reduce(failed, op=dist.ReduceOp.MAX, group=group)
|
|
if int(failed.item()) == 0:
|
|
dtype_codes = {
|
|
None: 0,
|
|
torch.float16: 1,
|
|
torch.bfloat16: 2,
|
|
torch.float32: 3,
|
|
}
|
|
code = torch.tensor(
|
|
[dtype_codes.get(temporal_dtype, -1)],
|
|
dtype=torch.int32,
|
|
device=device,
|
|
)
|
|
lowest = code.clone()
|
|
highest = code.clone()
|
|
dist.all_reduce(lowest, op=dist.ReduceOp.MIN, group=group)
|
|
dist.all_reduce(highest, op=dist.ReduceOp.MAX, group=group)
|
|
if int(lowest.item()) == int(highest.item()) and int(lowest.item()) >= 0:
|
|
return
|
|
raise MiniMaxH3ChunkedDecodeUnsupportedError(
|
|
"MiniMax-H3 VAE ranks selected incompatible temporal concat dtypes"
|
|
)
|
|
if preflight_error is not None:
|
|
raise preflight_error
|
|
raise MiniMaxH3ChunkedDecodeUnsupportedError("A peer VAE rank rejected the MiniMax-H3 temporal chunk preflight")
|
|
|
|
@torch.inference_mode()
|
|
def encode_image(self, image: Image.Image) -> torch.Tensor:
|
|
previous_parallel = self.model.parallel_tiling
|
|
self.model.parallel_tiling = False
|
|
parameter = next(self.parameters())
|
|
previous_dtype = parameter.dtype
|
|
if previous_dtype != torch.float32:
|
|
self.to(torch.float32)
|
|
devices = [parameter.device] if parameter.device.type != "cpu" else []
|
|
try:
|
|
with torch.random.fork_rng(devices=devices, device_type=parameter.device.type):
|
|
torch.default_generator.manual_seed(MINIMAX_H3_KEYFRAME_ENCODE_SEED)
|
|
for device in devices:
|
|
with self.device_module.device(device):
|
|
self.device_module.manual_seed(MINIMAX_H3_KEYFRAME_ENCODE_SEED)
|
|
latent = self.model.encode_images(
|
|
image,
|
|
use_fp16_latent=True,
|
|
)[0]
|
|
finally:
|
|
self.model.parallel_tiling = previous_parallel
|
|
if previous_dtype != torch.float32:
|
|
self.to(previous_dtype)
|
|
|
|
# Match the reference contract exactly: normalization and patchify
|
|
# happen on CPU in FP32 after the sampled encode. The condition noise
|
|
# path is sensitive enough that doing these elementwise operations on
|
|
# CUDA can noticeably change the final conditioned video.
|
|
latent = latent.float().cpu()
|
|
if latent.ndim == 4:
|
|
latent = latent[None]
|
|
channels = int(self.config_dict["latent_channels"])
|
|
mean = torch.tensor(
|
|
self.config_dict["latents_mean"],
|
|
).view(1, channels, 1, 1, 1)
|
|
std = torch.tensor(
|
|
self.config_dict["latents_std"],
|
|
).view(1, channels, 1, 1, 1)
|
|
return minimax_h3_patchify_video_latent(
|
|
(latent - mean) / std,
|
|
patch_size=(1, 2, 2),
|
|
).float()
|
|
|
|
@torch.inference_mode()
|
|
def encode_video(
|
|
self,
|
|
frames: Any,
|
|
) -> tuple[torch.Tensor, tuple[int, int, int]]:
|
|
parameter = next(self.parameters())
|
|
previous_dtype = parameter.dtype
|
|
if previous_dtype != torch.float32:
|
|
self.to(torch.float32)
|
|
devices = [parameter.device] if parameter.device.type != "cpu" else []
|
|
try:
|
|
with torch.random.fork_rng(devices=devices, device_type=parameter.device.type):
|
|
torch.default_generator.manual_seed(MINIMAX_H3_KEYFRAME_ENCODE_SEED)
|
|
for device in devices:
|
|
with self.device_module.device(device):
|
|
self.device_module.manual_seed(MINIMAX_H3_KEYFRAME_ENCODE_SEED)
|
|
latent = self.model.encode_videos(
|
|
frames,
|
|
use_fp16_latent=True,
|
|
)[0]
|
|
finally:
|
|
if previous_dtype != torch.float32:
|
|
self.to(previous_dtype)
|
|
|
|
latent = latent.float().cpu()
|
|
if latent.ndim == 4:
|
|
latent = latent[None]
|
|
channels = int(self.config_dict["latent_channels"])
|
|
if latent.ndim != 5 or int(latent.shape[1]) != channels:
|
|
raise ValueError(f"unexpected reference video latent shape {tuple(latent.shape)}")
|
|
shape = (
|
|
int(latent.shape[2]),
|
|
int(latent.shape[3]),
|
|
int(latent.shape[4]),
|
|
)
|
|
mean = torch.tensor(
|
|
self.config_dict["latents_mean"],
|
|
).view(1, channels, 1, 1, 1)
|
|
std = torch.tensor(
|
|
self.config_dict["latents_std"],
|
|
).view(1, channels, 1, 1, 1)
|
|
rows = minimax_h3_patchify_video_latent(
|
|
(latent - mean) / std,
|
|
patch_size=(1, 2, 2),
|
|
).float()
|
|
return rows, shape
|
|
|
|
@torch.inference_mode()
|
|
def decode_latent(self, latent: torch.Tensor) -> torch.Tensor:
|
|
with self._decode_tiling_context(latent):
|
|
decoded = self.model.decode_base(self._denormalize_latent(latent))
|
|
return self._normalize_decoded_frames(decoded)
|
|
|
|
@torch.inference_mode()
|
|
def decode_latent_with_chunks(
|
|
self,
|
|
latent: torch.Tensor,
|
|
chunk_callback: MiniMaxH3VideoChunkCallback | None,
|
|
) -> torch.Tensor:
|
|
"""Decode a video while publishing committed temporal chunks.
|
|
|
|
Every rank collaborating in the configured H3 VAE tile group must call
|
|
this method, with ``chunk_callback`` supplied only on group rank 0. A
|
|
callback failure is deferred until the compatibility decoder finishes its
|
|
remaining collectives, then propagated to every group rank. With
|
|
``parallel_size=1`` the runtime designates each request's output owner
|
|
by supplying a callback there and ``None`` on replica peers.
|
|
|
|
This is a model-local producer seam before the typed media boundary in
|
|
RFC #6541 / PR #6615. Representation conversion, transport,
|
|
presentation policy, and encoding remain runtime-owned. Chunks retain
|
|
the VAE canvas dimensions; a runtime consumer must apply the same
|
|
requested-size crop as the complete H3 pipeline output.
|
|
"""
|
|
|
|
group = self._chunk_process_group()
|
|
plan: MiniMaxH3TemporalDecodePlan | None = None
|
|
preflight_error: Exception | None = None
|
|
if latent.ndim != 5:
|
|
preflight_error = MiniMaxH3ChunkedDecodeUnsupportedError(
|
|
f"MiniMax-H3 video latent must be rank 5, got {tuple(latent.shape)}"
|
|
)
|
|
else:
|
|
try:
|
|
denormalized = self._denormalize_latent(latent)
|
|
plan = prepare_minimax_h3_temporal_chunks_compat(
|
|
self.model,
|
|
denormalized,
|
|
validate_sources=not getattr(
|
|
self,
|
|
"_minimax_h3_chunk_sources_validated",
|
|
False,
|
|
),
|
|
)
|
|
self._minimax_h3_chunk_sources_validated = True
|
|
del denormalized
|
|
except MiniMaxH3TemporalChunkCompatibilityError as exc:
|
|
preflight_error = MiniMaxH3ChunkedDecodeUnsupportedError(str(exc))
|
|
except Exception as exc: # noqa: BLE001
|
|
# Every VAE rank must reach the preflight reduction even when
|
|
# one node's private planner or environment fails locally.
|
|
exc.__traceback__ = None
|
|
preflight_error = exc
|
|
self._synchronize_chunk_preflight_error(
|
|
preflight_error,
|
|
plan.temporal_dtype if plan is not None else None,
|
|
group=group,
|
|
device=latent.device,
|
|
)
|
|
assert plan is not None
|
|
is_output_rank = self._is_chunk_output_rank(
|
|
chunk_callback,
|
|
group=group,
|
|
device=latent.device,
|
|
)
|
|
callback_error: BaseException | None = None
|
|
next_chunk_index = 0
|
|
expected_total_chunks: int | None = None
|
|
next_frame_start = 0
|
|
saw_final = False
|
|
|
|
def publish(
|
|
decoded_chunk: torch.Tensor,
|
|
*,
|
|
chunk_index: int,
|
|
total_chunks: int,
|
|
frame_start: int,
|
|
is_final: bool,
|
|
) -> None:
|
|
nonlocal callback_error, expected_total_chunks
|
|
nonlocal next_chunk_index, next_frame_start, saw_final
|
|
if not is_output_rank or callback_error is not None:
|
|
return
|
|
try:
|
|
if expected_total_chunks is None:
|
|
expected_total_chunks = total_chunks
|
|
if (
|
|
saw_final
|
|
or total_chunks <= 0
|
|
or total_chunks != expected_total_chunks
|
|
or chunk_index != next_chunk_index
|
|
or chunk_index >= total_chunks
|
|
or frame_start != next_frame_start
|
|
or is_final != (chunk_index + 1 == total_chunks)
|
|
):
|
|
raise RuntimeError(
|
|
"MiniMax-H3 temporal backend emitted discontinuous chunk metadata: "
|
|
f"index={chunk_index}/{next_chunk_index}, "
|
|
f"total={total_chunks}/{expected_total_chunks}, "
|
|
f"frame_start={frame_start}/{next_frame_start}, "
|
|
f"final={is_final}, after_final={saw_final}"
|
|
)
|
|
frames = self._normalize_decoded_frames(decoded_chunk)
|
|
frames = frames.clone(memory_format=torch.contiguous_format)
|
|
frame_count = int(frames.shape[2])
|
|
if frame_count <= 0:
|
|
raise RuntimeError("MiniMax-H3 temporal backend emitted an empty chunk")
|
|
next_chunk_index += 1
|
|
next_frame_start += frame_count
|
|
saw_final = is_final
|
|
assert chunk_callback is not None
|
|
chunk_callback(
|
|
frames,
|
|
chunk_index=chunk_index,
|
|
total_chunks=total_chunks,
|
|
frame_start=frame_start,
|
|
is_final=is_final,
|
|
)
|
|
except BaseException as exc: # noqa: BLE001
|
|
# Keep a failed GPU chunk and its callback locals out of the
|
|
# traceback while peer ranks finish the remaining collectives.
|
|
exc.__traceback__ = None
|
|
callback_error = exc
|
|
|
|
with self._decode_tiling_context(latent):
|
|
decoded = decode_minimax_h3_temporal_chunks_compat(
|
|
self.model,
|
|
plan,
|
|
publish,
|
|
)
|
|
frames = self._normalize_decoded_frames(decoded)
|
|
if (
|
|
is_output_rank
|
|
and callback_error is None
|
|
and (not saw_final or next_chunk_index != expected_total_chunks or next_frame_start != int(frames.shape[2]))
|
|
):
|
|
callback_error = RuntimeError(
|
|
"MiniMax-H3 temporal backend did not reconstruct the full decode: "
|
|
f"final={saw_final}, chunks={next_chunk_index}/{expected_total_chunks}, "
|
|
f"frames={next_frame_start}/{frames.shape[2]}"
|
|
)
|
|
self._synchronize_chunk_callback_error(
|
|
callback_error,
|
|
group=group,
|
|
device=frames.device,
|
|
)
|
|
return frames
|
|
|
|
|
|
class MiniMaxH3AudioVAE(nn.Module):
|
|
def __init__(
|
|
self,
|
|
component_path: str,
|
|
*,
|
|
device: torch.device,
|
|
load_device: torch.device | None = None,
|
|
) -> None:
|
|
super().__init__()
|
|
self._device_target = device
|
|
self.config_dict = _load_component_config(component_path)
|
|
self.remote = _load_remote_component(
|
|
component_path,
|
|
self.config_dict,
|
|
)
|
|
# The checkpoint's audio VAE contract is FP32 for both reference
|
|
# encoding and waveform decoding.
|
|
initial_device = load_device or device
|
|
self.remote.eval().to(device=initial_device, dtype=torch.float32)
|
|
self._stager = None
|
|
if initial_device.type == "cpu" and device.type not in ("cpu", "meta"):
|
|
self._stager = PinnedModuleStager(
|
|
self.remote,
|
|
device,
|
|
pin_memory=True,
|
|
)
|
|
self.model = self.remote.model
|
|
self.sample_rate = int(self.config_dict["sample_rate"])
|
|
|
|
def load_to_device(self) -> None:
|
|
if self._stager is not None:
|
|
self._stager.load()
|
|
else:
|
|
self.remote.to(self._device_target)
|
|
|
|
def set_omni_component_cache(self, cache: BoundedAllocatorCache | None) -> None:
|
|
self._omni_component_cache = cache
|
|
if self._stager is not None:
|
|
self._stager.set_cache_retention(cache)
|
|
|
|
def offload_to_cpu(self) -> None:
|
|
if self._stager is not None:
|
|
self._stager.offload()
|
|
else:
|
|
self.remote.to("cpu")
|
|
cache = getattr(self, "_omni_component_cache", None)
|
|
if cache is None:
|
|
torch.accelerator.empty_cache()
|
|
else:
|
|
cache.release_if_needed()
|
|
|
|
@torch.inference_mode()
|
|
def encode_waveform(
|
|
self,
|
|
waveform: torch.Tensor,
|
|
sample_rate: int,
|
|
) -> tuple[torch.Tensor, int]:
|
|
import torchaudio
|
|
|
|
waveform = waveform.float()
|
|
if waveform.ndim == 1:
|
|
waveform = waveform[None]
|
|
if int(sample_rate) != MINIMAX_H3_AUDIO_SAMPLE_RATE:
|
|
waveform = torchaudio.transforms.Resample(
|
|
int(sample_rate),
|
|
MINIMAX_H3_AUDIO_SAMPLE_RATE,
|
|
)(waveform)
|
|
if waveform.shape[0] < MINIMAX_H3_AUDIO_CHANNELS:
|
|
waveform = waveform.repeat(
|
|
MINIMAX_H3_AUDIO_CHANNELS,
|
|
1,
|
|
)
|
|
waveform = waveform[:MINIMAX_H3_AUDIO_CHANNELS]
|
|
device = next(self.model.parameters()).device
|
|
waveform = waveform.to(device)
|
|
|
|
with _AudioVAEDeterminismContext():
|
|
audio = self.model.preprocess(
|
|
waveform.unsqueeze(1),
|
|
MINIMAX_H3_AUDIO_SAMPLE_RATE,
|
|
)
|
|
latent = self.model.encoder(audio)
|
|
if bool(getattr(self.model, "attn_proj", False)):
|
|
latent = self.model.pre_block(latent.transpose(1, 2)).transpose(1, 2)
|
|
latent = self.model.mean_proj(latent).float().cpu()
|
|
|
|
channels = int(self.config_dict["latent_channels"])
|
|
if latent.shape[-1] != channels:
|
|
if latent.shape[1] != channels:
|
|
raise ValueError(f"cannot canonicalize audio latent {tuple(latent.shape)}")
|
|
latent = latent.transpose(1, 2).contiguous()
|
|
mean = torch.tensor(
|
|
self.config_dict["latents_mean"],
|
|
).view(1, 1, channels)
|
|
std = torch.tensor(
|
|
self.config_dict["latents_std"],
|
|
).view(1, 1, channels)
|
|
rows = ((latent - mean) / std).reshape(-1, channels)
|
|
return rows.float(), int(latent.shape[1])
|
|
|
|
@torch.inference_mode()
|
|
def decode_latent(self, latent: torch.Tensor) -> torch.Tensor:
|
|
channels = int(self.config_dict["latent_channels"])
|
|
mean = torch.tensor(
|
|
self.config_dict["latents_mean"],
|
|
device=latent.device,
|
|
dtype=latent.dtype,
|
|
).view(1, channels, 1)
|
|
std = torch.tensor(
|
|
self.config_dict["latents_std"],
|
|
device=latent.device,
|
|
dtype=latent.dtype,
|
|
).view(1, channels, 1)
|
|
waveform = self.remote.decode(latent * std + mean)
|
|
if waveform.ndim != 3 or waveform.shape[1] != 1:
|
|
raise ValueError(f"unexpected decoded audio shape {tuple(waveform.shape)}")
|
|
return waveform.permute(1, 0, 2).contiguous().float()
|
|
|
|
|
|
__all__ = [
|
|
"MiniMaxH3AudioVAE",
|
|
"MiniMaxH3ChunkCallbackPeerError",
|
|
"MiniMaxH3ChunkedDecodeUnsupportedError",
|
|
"MiniMaxH3VideoChunkCallback",
|
|
"MiniMaxH3VideoVAE",
|
|
]
|