anodyine/vllm-omni
1
0 Fork 0
Code Issues 4 Pull Requests 2 Actions Packages Projects Releases Wiki Activity
Files
e3be42e052e88d052d21cf7a24ee84c70f020fd6
vllm-omni/vllm_omni/diffusion/models/minimax_h3/pipeline_minimax_h3.py
T

2551 lines
109 KiB
Python
Raw Blame History

# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
"""vLLM-Omni pipeline for MiniMax H3 FL2VA and Ref2VA partitions."""
from __future__ import annotations
import json
import math
import os
from collections.abc import Iterable, Mapping, Sequence
from contextlib import contextmanager
from dataclasses import fields, replace
from itertools import groupby
from pathlib import Path
from typing import TYPE_CHECKING, Any, ClassVar
import numpy as np
import torch
import torch.distributed as dist
import torch.nn as nn
from transformers import Qwen2TokenizerFast, Qwen3VLProcessor
from vllm.logger import init_logger
from vllm.model_executor.layers.quantization.base_config import QuantizationConfig
from vllm_omni.diffusion import envs
from vllm_omni.diffusion.cache.cachedit import (
CacheDiTBackend,
RequestScopedCacheDiTRuntime,
)
from vllm_omni.diffusion.data import DiffusionOutput, OmniDiffusionConfig
from vllm_omni.diffusion.distributed.parallel_state import get_world_group, init_world_group
from vllm_omni.diffusion.distributed.utils import get_local_device
from vllm_omni.diffusion.forward_context import DenoiseProgressMixin
from vllm_omni.diffusion.model_loader.diffusers_loader import (
DiffusersPipelineLoader,
)
from vllm_omni.diffusion.models.interface import (
SupportAudioInput,
SupportAudioOutput,
SupportImageInput,
SupportsComponentDiscovery,
)
from vllm_omni.diffusion.models.progress_bar import ProgressBarMixin
from vllm_omni.diffusion.offloader import (
BoundedAllocatorCache,
OffloadPlan,
apply_sequential_offload,
remove_sequential_offload,
sequential_offload_component,
)
from vllm_omni.diffusion.offloader.config import (
DIT_COMPONENT,
TEXT_ENCODER_COMPONENT,
OffloadStrategy,
resolve_offload,
should_offload_component,
)
from vllm_omni.diffusion.offloader.module_collector import ModuleDiscovery
from vllm_omni.diffusion.profiler.diffusion_pipeline_profiler import (
DiffusionPipelineProfilerMixin,
)
from vllm_omni.diffusion.sched.sigma_schedule import DMD2SigmaSchedule
from vllm_omni.diffusion.utils.media_utils import normalize_preencode_batch_frames, normalize_video_codec_options
from vllm_omni.diffusion.worker.request_batch import DiffusionRequestBatch
from vllm_omni.errors import OmniClientError, client_error_from_metadata
from vllm_omni.model_executor.model_loader.weight_utils import (
download_weights_from_hf_specific,
)
from vllm_omni.model_executor.models.minimax_h3.checkpoint import (
is_minimax_h3_modular,
resolve_minimax_h3_partition,
)
from vllm_omni.model_executor.models.minimax_h3.conditioning import (
MiniMaxH3EncoderConditioning,
MiniMaxH3EncoderMediaConditioning,
MiniMaxH3EncoderMediaInput,
MiniMaxH3TextConditioning,
)
from vllm_omni.model_executor.models.minimax_h3.encoder_processing import (
PreparedEncoderInputs,
encode_media,
prepare_encoder_inputs,
)
from vllm_omni.model_executor.models.minimax_h3.preprocessing import build_minimax_h3_presentation
from vllm_omni.model_executor.models.minimax_h3.reference_video import (
MINIMAX_H3_PREPARED_REFERENCE_VIDEOS_KEY,
deserialize_prepared_reference_videos,
)
from vllm_omni.platforms import current_omni_platform
from vllm_omni.quantization import (
resolve_component_quant_config as _resolve_component_quant_config,
)
from vllm_omni.quantization.component_config import (
resolve_encoder_quant_config as _resolve_encoder_quant_config,
)
from .batched_packing import minimax_h3_batched_forward_kwargs
from .condition_noise import (
minimax_h3_audio_cond_noise_aug_rows,
minimax_h3_imgvid_cond_noise_aug_rows,
)
from .denoise_loop import (
MiniMaxH3DenoiseBranch,
minimax_h3_denoise_loop,
minimax_h3_prepare_denoise_rows,
minimax_h3_publish_denoise_progress,
)
from .encoder import MiniMaxH3Qwen3VLEncoder
from .fasth3 import FastH3WeightFusion, resolve_fasth3_fusion
from .fasth3_checkpoint import FastH3CheckpointSpec
from .lora import TurboSpec, load_minimax_h3_turbo_lora
from .minimax_h3_transformer import (
MiniMaxH3Attention,
MiniMaxH3DiTModel,
_attention_isolates_packed_requests,
)
from .npu.lora import (
MINIMAX_H3_NATIVE_INFERENCE_STEPS,
load_minimax_h3_native_lora,
)
from .packed_sequence import (
MINIMAX_H3_MAX_PAD_SEQ_LEN,
MINIMAX_H3_SEQ_ALIGN,
minimax_h3_packed_sequence,
minimax_h3_packed_sequence_ref2va_blocks,
)
from .packed_tokens import (
minimax_h3_patchify_video_latent,
minimax_h3_unpack_audio_tokens,
minimax_h3_unpatchify_video_tokens,
)
from .quality_policy import MINIMAX_H3_GENERIC_CACHE_KEY, MiniMaxH3QualityPolicy
from .scheduling_minimax_h3_euler_ancestral import (
minimax_h3_euler_eta0_step,
minimax_h3_rf_v_to_x0,
)
from .time_request import (
MINIMAX_H3_SHAPE_PLANNER,
minimax_h3_time_shift_sigmas,
)
from .vae import MiniMaxH3AudioVAE, MiniMaxH3VideoVAE
if TYPE_CHECKING:
from vllm_omni.diffusion.worker.input_batch import InputBatch
from vllm_omni.diffusion.worker.utils import StepRequestState
logger = init_logger(__name__)
if TYPE_CHECKING:
from vllm.lora.lora_model import LoRAModel
from vllm.lora.peft_helper import PEFTHelper
from vllm_omni.lora.request import LoRARequest
MINIMAX_H3_FPS = 24
MINIMAX_H3_AUDIO_SAMPLE_RATE = 32000
MINIMAX_H3_IMGVID_COND_TIMESTEP = 0.999
MINIMAX_H3_AUDIO_REF_COND_TIMESTEP = 1.0
MINIMAX_H3_DOWNLOAD_PATTERNS = [
"FL2VA/**",
"Ref2VA/model_index.json",
"Ref2VA/transformer/**",
]
MINIMAX_H3_TASK_DOWNLOAD_PATTERNS = {
"fl2va": ["FL2VA/**"],
"ref2va": ["Ref2VA/**"],
}
MINIMAX_H3_DIFFUSION_DOWNLOAD_PATTERNS = {
"fl2va": [
"FL2VA/model_index.json",
"FL2VA/transformer/**",
"FL2VA/video_vae/**",
"FL2VA/audio_vae/**",
],
"ref2va": [
"Ref2VA/model_index.json",
"Ref2VA/transformer/**",
"Ref2VA/video_vae/**",
"Ref2VA/audio_vae/**",
],
"combined": [
"FL2VA/model_index.json",
"FL2VA/transformer/**",
"FL2VA/video_vae/**",
"FL2VA/audio_vae/**",
"Ref2VA/model_index.json",
"Ref2VA/transformer/**",
],
}
def _resolve_minimax_h3_text_encoder_quant_config(
quant_config: QuantizationConfig | None,
) -> QuantizationConfig | None:
resolved = _resolve_component_quant_config(quant_config, "text_encoder")
return _resolve_encoder_quant_config(resolved)
def _minimax_h3_partition_for_task(
task_type: str | None,
model: str | None = None,
) -> str:
return resolve_minimax_h3_partition(model or "", task_type, auto_partition="combined")
def _resolve_minimax_h3_model_root(
model: str,
revision: str | None,
partition: str,
*,
load_text_encoder: bool,
) -> Path:
path = Path(model)
if path.is_dir():
if path.name in {"FL2VA", "Ref2VA"}:
return path.parent
return path
if is_minimax_h3_modular(model, revision):
allow_patterns = ["modular_model_index.json", "fastvideo_inference.json", "provenance.json", "transformer/**"]
if load_text_encoder:
allow_patterns += ["text_encoder/**", "tokenizer/**", "processor/**"]
elif load_text_encoder:
allow_patterns = (
MINIMAX_H3_DOWNLOAD_PATTERNS if partition == "combined" else MINIMAX_H3_TASK_DOWNLOAD_PATTERNS[partition]
)
else:
allow_patterns = MINIMAX_H3_DIFFUSION_DOWNLOAD_PATTERNS[partition]
return Path(
download_weights_from_hf_specific(
model_name_or_path=model,
cache_dir=None,
allow_patterns=allow_patterns,
revision=revision,
require_all=True,
)
)
_MINIMAX_H3_DENOISE_INPUT_KEYS = (
"task",
"text_embeddings",
"text_tags",
"seed",
"latent_t",
"latent_h",
"latent_w",
"audio_t",
"num_frames",
"num_steps",
"video_shift",
"audio_shift",
"base_schedule",
"visual_condition",
"visual_condition_shape",
"audio_condition",
"ref_audio_t",
"ref_blocks",
"visual_condition_shapes",
"audio_condition_lengths",
"keyframe_frame_indices",
"pad_seq_len",
)
# ``StepRequestState.extra`` keys owned by the step-execution path.
_STEP_BRANCH = "minimax_h3_branch"
_STEP_AUDIO_ROWS = "minimax_h3_audio_rows"
_STEP_AUDIO_NOISE_PRED = "minimax_h3_audio_noise_pred"
_STEP_SIGMAS_VIDEO = "minimax_h3_sigmas_video"
_STEP_SIGMAS_AUDIO = "minimax_h3_sigmas_audio"
_STEP_COND_ANCHOR = "minimax_h3_cond_anchor"
_STEP_AUDIO_ANCHOR = "minimax_h3_audio_anchor"
_STEP_SHAPE = "minimax_h3_shape"
_STEP_TRANSFORMER = "minimax_h3_transformer"
def _minimax_h3_step_schedule(state: StepRequestState) -> dict[str, float]:
"""Return the sigma/timestep values this request needs for its current step.
Mirrors the per-iteration arithmetic of ``minimax_h3_denoise_loop`` so step
mode and request mode advance identically.
"""
step = int(state.step_index)
sigmas_video = state.extra[_STEP_SIGMAS_VIDEO]
sigmas_audio = state.extra[_STEP_SIGMAS_AUDIO]
sigma_video = float(sigmas_video[step])
sigma_audio = float(sigmas_audio[step])
t_video = 1.0 - sigma_video
t_audio = 1.0 - sigma_audio
return {
"sigma_video": sigma_video,
"sigma_video_next": float(sigmas_video[step + 1]),
"sigma_audio": sigma_audio,
"sigma_audio_next": float(sigmas_audio[step + 1]),
"t_video": t_video,
"t_audio": t_audio,
"imgvid_cond_timestep": max(t_video, MINIMAX_H3_IMGVID_COND_TIMESTEP),
"audio_ref_cond_timestep": max(t_audio, MINIMAX_H3_AUDIO_REF_COND_TIMESTEP),
}
def _read_base_schedule(release: Mapping[str, Any]) -> DMD2SigmaSchedule | None:
"""Read a partition's distilled schedule. An absent key means legacy uniform."""
return DMD2SigmaSchedule.from_metadata(release)
def resolve_minimax_h3_diffusion_model_path(
model: str,
revision: str | None,
task_type: str | None,
) -> str:
"""Resolve a repository root or Hub ID to its startup partition."""
partition = (
"combined"
if str(task_type or "").lower() == "combined"
else resolve_minimax_h3_partition(model, task_type, auto_partition="fl2va")
)
model_root = _resolve_minimax_h3_model_root(
model,
revision,
partition,
load_text_encoder=False,
)
if is_minimax_h3_modular(str(model_root), revision):
return str(model_root)
if partition == "combined":
return str(model_root)
subdir = "Ref2VA" if partition == "ref2va" else "FL2VA"
return str(model_root / subdir)
def _minimax_h3_post_process(output, output_type: str = "np"):
"""Convert the joint video/audio output without capturing worker state.
The callable crosses the multiprocessing result queue, so it must remain a
module-level function that the standard pickle module can resolve.
``_prepare_minimax_h3_video_output`` already quantises the video to uint8
frames on the accelerator, so there is nothing left to scale or transpose
here.
"""
if not isinstance(output, tuple) or len(output) != 2:
return output
video, audio = output
if isinstance(video, (bytes, bytearray, memoryview)):
video = [video]
if isinstance(video, list) and all(isinstance(item, (bytes, bytearray, memoryview)) for item in video):
encoded_videos = [bytes(item) for item in video]
else:
encoded_videos = None
if encoded_videos is not None:
return {
"video": encoded_videos,
"audio": [None] * len(encoded_videos),
"audio_sample_rate": MINIMAX_H3_AUDIO_SAMPLE_RATE,
"fps": MINIMAX_H3_FPS,
}
if video.dtype != torch.uint8 or video.ndim != 5 or video.shape[-1] not in (3, 4):
# Float or channel-first frames would reach the muxer as a black or
# banded video rather than as an error.
raise ValueError(
f"MiniMax-H3 post-processing expects (B, T, H, W, C) uint8, got {tuple(video.shape)} {video.dtype}"
)
if output_type == "latent":
return output
if output_type == "np":
video = video.detach().cpu().numpy()
audio = audio.detach().float().cpu().numpy()
video = [sample for sample in video]
return {
"video": video,
"audio": audio,
"audio_sample_rate": MINIMAX_H3_AUDIO_SAMPLE_RATE,
"fps": MINIMAX_H3_FPS,
}
def _prepare_minimax_h3_video_output(video: torch.Tensor) -> torch.Tensor:
"""Quantize decoded frames in place before worker-to-engine transfer."""
video = video.detach().float()
video.clamp_(0, 1).mul_(255).round_()
return video.permute(0, 2, 3, 4, 1).to(
dtype=torch.uint8,
memory_format=torch.contiguous_format,
)
def _register_dlo_component_cache(cache: BoundedAllocatorCache, *components: Any) -> None:
for component in components:
if component is not None:
component.set_omni_component_cache(cache)
def get_minimax_h3_post_process_func(
od_config: OmniDiffusionConfig,
):
del od_config
return _minimax_h3_post_process
def _resolve_minimax_h3_num_outputs(value: Any) -> int:
if value is None:
return 1
if isinstance(value, bool) or not isinstance(value, (int, np.integer)):
raise OmniClientError("MiniMax H3 num_outputs_per_prompt must be an integer in [1, 10]")
value = int(value)
if not 1 <= value <= 10:
raise OmniClientError(f"MiniMax H3 num_outputs_per_prompt must be in [1, 10], got {value}")
return value
def _resolve_pad_seq_len(value: object) -> int | None:
"""Validate the optional packed-length pin from ``extra_args``.
The packer itself only needs a value that covers the used rows. A request
field needs two more guards: an unaligned value would silently open yet
another compiled shape, and an unbounded one would size every structural
tensor the packer allocates.
"""
if value is None:
return None
if isinstance(value, bool) or not isinstance(value, (int, np.integer)):
raise OmniClientError("MiniMax H3 pad_seq_len must be an integer")
pinned = int(value)
if pinned <= 0:
raise OmniClientError(f"MiniMax H3 pad_seq_len must be positive, got {pinned}")
if pinned % MINIMAX_H3_SEQ_ALIGN:
raise OmniClientError(f"MiniMax H3 pad_seq_len must be a multiple of {MINIMAX_H3_SEQ_ALIGN}, got {pinned}")
if pinned > MINIMAX_H3_MAX_PAD_SEQ_LEN:
raise OmniClientError(f"MiniMax H3 pad_seq_len must be at most {MINIMAX_H3_MAX_PAD_SEQ_LEN}, got {pinned}")
return pinned
def _minimax_h3_output_seeds(seed: int, num_outputs: int) -> list[int]:
return [int(seed) + output_index for output_index in range(int(num_outputs))]
def _dit_rank_world() -> tuple[Any, int, int]:
if not dist.is_initialized():
return None, 0, 1
group = get_world_group().device_group
return group, dist.get_rank(group), dist.get_world_size(group)
def _broadcast_rank0_exception(exc: Exception | None) -> None:
"""Synchronize a rank-0-only exception across every DiT rank.
H3 reference-video preparation runs only on rank 0; the other DiT ranks
return ``None`` without touching disk. When rank 0 raises inside that
path it exits :meth:`prepare_encode` before reaching the downstream
``dist.broadcast`` calls, and non-zero ranks then hang on those
collectives forever. Every rank calls this helper right after the
rank-0-only work, before any subsequent collective, so all ranks either
raise the same error together or all continue.
"""
group, rank, world_size = _dit_rank_world()
if world_size == 1:
if exc is not None:
raise exc
return
if rank == 0:
if exc is None:
payload: list[Any] = [None]
else:
payload = [
{
"type": type(exc).__name__,
"message": str(exc),
"status_code": getattr(exc, "status_code", None),
"error_type": getattr(exc, "error_type", None),
}
]
else:
payload = [None]
dist.broadcast_object_list(payload, src=0, group=group)
info = payload[0]
if info is None:
return
if rank == 0:
assert exc is not None
raise exc
# Rebuild a matching client-facing error on non-zero ranks so the runner's
# per-request try/except records the same 4xx status as rank 0. The exact
# subclass need not survive the wire; the message and status suffice.
status_code = info.get("status_code")
error_type = info.get("error_type")
message = f"[rank 0] {info['type']}: {info['message']}"
if status_code is not None:
raise client_error_from_metadata(
message,
status_code=int(status_code),
error_type=error_type,
)
raise RuntimeError(message)
def _broadcast_tensor(
tensor: torch.Tensor | None,
*,
dtype: torch.dtype,
device: torch.device,
) -> torch.Tensor:
group, rank, world_size = _dit_rank_world()
if world_size == 1:
if tensor is None:
raise ValueError("source tensor is required for single-rank execution")
return tensor.to(device=device, dtype=dtype)
shape = torch.zeros(5, dtype=torch.long, device=device)
if rank == 0:
if tensor is None:
raise ValueError("rank 0 must provide a tensor to broadcast")
shape[0] = tensor.ndim
shape[1 : tensor.ndim + 1] = torch.tensor(
tensor.shape,
device=device,
)
dist.broadcast(shape, src=0, group=group)
ndim = int(shape[0].item())
tensor_shape = tuple(int(v) for v in shape[1 : ndim + 1].tolist())
if rank == 0:
output = tensor.to(device=device, dtype=dtype).contiguous()
else:
output = torch.empty(tensor_shape, device=device, dtype=dtype)
dist.broadcast(output, src=0, group=group)
return output
class _SingleRankEncoderGroup:
"""Lightweight encoder group for ``text_encoder_tp_size == 1``.
Avoids creating a distributed ``GroupCoordinator`` with a single-member
rank set, which would assert on every other DiT rank that is not part of
the group. The pipeline and encoder only use the attributes below, and
all ``world_size == 1`` code paths short-circuit before any collective.
"""
world_size: int = 1
ranks: list[int] = [0]
def __init__(self, rank: int) -> None:
self.rank_in_group = 0 if rank == 0 else -1
self.device_group = None
class MiniMaxH3Pipeline(
nn.Module,
DenoiseProgressMixin,
ProgressBarMixin,
DiffusionPipelineProfilerMixin,
SupportImageInput,
SupportAudioInput,
SupportAudioOutput,
SupportsComponentDiscovery,
):
"""CFG-distilled joint video/audio generation for MiniMax H3."""
supports_step_execution: ClassVar[bool] = True
_dit_modules: ClassVar[list[str]] = ["transformer", "transformers_ref"]
_encoder_modules: ClassVar[list[str]] = ["text_encoder"]
_vae_modules: ClassVar[list[str]] = ["video_vae", "audio_vae"]
_offload_plan: ClassVar[OffloadPlan] = OffloadPlan(
offload_submodules={"token_refiner": "blocks"},
resident_dit_paths=frozenset({"transformer"}),
encoder_component_types={"text_encoder": TEXT_ENCODER_COMPONENT},
encoder_block_attrs={"text_encoder": ("vision.blocks", "text_model.layers")},
on_demand_component_paths=frozenset({"text_encoder", "video_vae", "audio_vae"}),
)
_PROFILER_TARGETS: ClassVar[list[str]] = [
"encode_prompt",
"_encode_local_media",
"diffuse",
"decode",
"video_vae.decode_latent",
"audio_vae.decode_latent",
"prepare_encode",
"denoise_step",
"post_decode",
]
dummy_run_num_frames: ClassVar[int] = 0
# Only distilled releases pin a schedule, so the default keeps the legacy
# uniform path available to partially constructed pipelines.
_base_schedule_by_partition: ClassVar[Mapping[str, DMD2SigmaSchedule | None]] = {}
# Set from --lora-path during construction; absent means no FastH3 adapter.
_fasth3: FastH3WeightFusion | None = None
_fasth3_checkpoint: FastH3CheckpointSpec | None = None
def _load_diffusion_lora_adapter(
self,
*,
lora_request: LoRARequest,
lora_path: str | Path,
dtype: torch.dtype,
) -> tuple[LoRAModel, PEFTHelper] | None:
# A cache eviction may be followed by a different adapter reusing the
# same client-supplied ID. Every real load replaces the classification.
self._turbo_lora_specs.pop(lora_request.lora_int_id, None)
self._native_lora_adapter_ids.discard(lora_request.lora_int_id)
self._lora_sigma_schedules.pop(lora_request.lora_int_id, None)
od_config = getattr(self, "od_config", None)
offload_modes = []
if od_config is not None:
resolved_offload = resolve_offload(od_config)
if resolved_offload.offloads(DIT_COMPONENT):
if resolved_offload.strategy is OffloadStrategy.MODEL_LEVEL:
offload_modes.append("model-level CPU offload")
elif resolved_offload.strategy is OffloadStrategy.LAYER_WISE:
offload_modes.append("layerwise offload")
loaded = load_minimax_h3_turbo_lora(
partition=self.partition,
lora_request=lora_request,
lora_path=lora_path,
dtype=dtype,
unsupported_offload_mode=" or ".join(offload_modes) or None,
)
if loaded is not None:
lora_model, peft_helper, turbo_spec = loaded
self._turbo_lora_specs[lora_request.lora_int_id] = turbo_spec
return lora_model, peft_helper
# Selection is by the artifact's safetensors ``key_format``, not by the
# running platform: the native loader is checkpoint-format parsing with
# no ``torch_npu`` dependency, so it needs no ``current_omni_platform``
# dispatch and binds the same adapter on NPU, CUDA and CPU.
native_loaded = load_minimax_h3_native_lora(
partition=self.partition,
lora_request=lora_request,
lora_path=lora_path,
dtype=dtype,
unsupported_offload_mode=" or ".join(offload_modes) or None,
)
if native_loaded is not None:
lora_model, peft_helper, sigma_schedule = native_loaded
self._native_lora_adapter_ids.add(lora_request.lora_int_id)
self._lora_sigma_schedules[lora_request.lora_int_id] = sigma_schedule
return lora_model, peft_helper
return None
def _validate_diffusion_lora_binding(
self,
*,
lora_model: LoRAModel,
bound_lora_names: frozenset[str],
) -> None:
if lora_model.id in self._turbo_lora_specs:
missing = sorted(set(lora_model.loras) - bound_lora_names)
if missing:
raise ValueError(
"MiniMax-H3 Turbo LoRA binding is incomplete: "
f"bound={len(bound_lora_names)}/{len(lora_model.loras)}, missing={missing[:5]}"
)
return
if lora_model.id not in self._native_lora_adapter_ids:
return
missing = sorted(set(lora_model.loras) - bound_lora_names)
if missing:
raise ValueError(
"MiniMax-H3 native LoRA binding is incomplete: "
f"bound={len(bound_lora_names)}/{len(lora_model.loras)}, missing={missing[:5]}"
)
def _active_turbo_spec(self, sampling: Any) -> TurboSpec | None:
"""Return the spec of the Turbo adapter this request actually applies.
A recognized adapter at scale 0 contributes nothing, so it neither
constrains the task nor imposes its sampler contract.
"""
lora_request = sampling.lora_request
if lora_request is None or math.isclose(0.0, float(sampling.lora_scale)):
return None
return self._turbo_lora_specs.get(lora_request.lora_int_id)
def _has_active_native_lora(self, sampling: Any) -> bool:
lora_request = sampling.lora_request
return (
lora_request is not None
and not math.isclose(0.0, float(sampling.lora_scale))
and lora_request.lora_int_id in self._native_lora_adapter_ids
)
def _validate_native_sampling(self, sampling: Any, *, task: str) -> None:
if task != "t2va":
raise OmniClientError("MiniMax-H3 native LoRA supports T2VA requests only")
# Derive the expected count from the adapter's own schedule so the
# message can never disagree with the schedule the denoise loop runs.
schedule = self._lora_sigma_schedules.get(sampling.lora_request.lora_int_id)
expected_steps = MINIMAX_H3_NATIVE_INFERENCE_STEPS if schedule is None else schedule.num_inference_steps
# Only request mode can take the count from the adapter schedule: step
# mode admits the request in ``StepScheduler``, which reads
# ``num_inference_steps`` off it before any pipeline hook runs. Reject
# omission there rather than advertise a contract that would either fail
# admission or disagree with the denoise loop.
od_config = getattr(self, "od_config", None)
omission_allowed = not getattr(od_config, "step_execution", False)
or_omitted = " or omitted" if omission_allowed else ""
sigma_steps = sampling.num_inference_steps
if sigma_steps is None:
if omission_allowed:
return
raise OmniClientError(
f"MiniMax-H3 native LoRA requires an explicit num_inference_steps={expected_steps} "
"under step execution, because the step scheduler derives the total step count from "
"the request before the adapter schedule is known"
)
if int(sigma_steps) == expected_steps + 1:
raise OmniClientError(
"MiniMax-H3 native LoRA uses the distilled interval-count contract; "
f"num_inference_steps must be {expected_steps}{or_omitted}, not {expected_steps + 1}"
)
if int(sigma_steps) != expected_steps:
raise OmniClientError(
f"MiniMax-H3 native LoRA requires num_inference_steps={expected_steps} "
f"(one denoiser evaluation per sigma interval){or_omitted}"
)
def _sigma_schedule_for_request(self, sampling: Any, task: str) -> DMD2SigmaSchedule | None:
lora_request = sampling.lora_request
if (
lora_request is not None
and not math.isclose(0.0, float(sampling.lora_scale))
and lora_request.lora_int_id in self._lora_sigma_schedules
):
adapter_schedule = self._lora_sigma_schedules[lora_request.lora_int_id]
checkpoint_schedule = self._base_schedule_for_task(task)
if checkpoint_schedule is not None:
raise OmniClientError(
"MiniMax-H3 native LoRA cannot be activated on a checkpoint that already pins base_schedule"
)
return adapter_schedule
return self._base_schedule_for_task(task)
def _validate_turbo_sampling(self, sampling: Any, spec: TurboSpec) -> None:
"""Hold a request to the contract of the artifact that is loaded.
Sigma-point count and both flow shifts vary across the Turbo family, so
each is checked against the adapter's own spec rather than a single
published configuration.
"""
extra = sampling.extra_args or {}
sigma_points = sampling.num_inference_steps
if sigma_points != spec.sigma_points:
raise OmniClientError(
f"{spec.filename} is a {spec.denoise_steps}-step artifact and requires "
f"num_inference_steps={spec.sigma_points} "
f"({spec.sigma_points} sigma points produce {spec.denoise_steps} denoiser evaluations)"
)
try:
video_shift = float(extra.get("flow_shift", self.default_video_shift))
except (TypeError, ValueError) as exc:
raise OmniClientError(f"{spec.filename} requires flow_shift={spec.video_shift:g}") from exc
if not math.isclose(video_shift, spec.video_shift):
raise OmniClientError(f"{spec.filename} requires flow_shift={spec.video_shift:g}")
try:
audio_shift = float(extra.get("audio_flow_shift", self.default_audio_shift))
except (TypeError, ValueError) as exc:
raise OmniClientError(f"{spec.filename} requires audio_flow_shift={spec.audio_shift:g}") from exc
if not math.isclose(audio_shift, spec.audio_shift):
raise OmniClientError(f"{spec.filename} requires audio_flow_shift={spec.audio_shift:g}")
def adopt_cache_dit_backend(self, backend: CacheDiTBackend) -> None:
"""Adopt runner-installed generic Cache-DiT for request transitions."""
self._cache_dit_runtime.adopt(
backend,
installation_key=MINIMAX_H3_GENERIC_CACHE_KEY,
)
def is_cache_dit_enabled(self) -> bool:
"""Return the request-scoped Cache-DiT installation state."""
return self._cache_dit_runtime.is_enabled
def __init__(
self,
*,
od_config: OmniDiffusionConfig,
prefix: str = "",
) -> None:
del prefix
super().__init__()
self.od_config = od_config
self.parallel_config = od_config.parallel_config
if int(self.parallel_config.cfg_parallel_size) != 1:
raise ValueError("MiniMax-H3 is CFG-distilled and has no negative branch; cfg_parallel_size must be 1")
self.device = get_local_device()
self.load_text_encoder = od_config.model_loaded.get("text_encoder", True)
self.load_vae_encoder = od_config.model_loaded.get("vae_encoder", True)
if self.load_vae_encoder is False and self.load_text_encoder is True:
raise ValueError(
"MiniMax H3 does not support local text encoding with external media conditioning; "
"set text_encoder=false or enable vae_encoder"
)
self._encoder_modules = ["text_encoder"] if self.load_text_encoder else []
self._PROFILER_TARGETS = list(type(self)._PROFILER_TARGETS)
encoder_block_attrs = dict(self._offload_plan.encoder_block_attrs)
if not self.load_text_encoder:
encoder_block_attrs.pop("text_encoder", None)
on_demand_component_paths = set(self._offload_plan.on_demand_component_paths)
if not self.load_text_encoder:
on_demand_component_paths.discard("text_encoder")
self._offload_plan = replace(
self._offload_plan,
encoder_block_attrs=encoder_block_attrs,
on_demand_component_paths=frozenset(on_demand_component_paths),
)
if not self.load_text_encoder:
self._PROFILER_TARGETS.remove("encode_prompt")
if not self.load_vae_encoder:
self._PROFILER_TARGETS.remove("_encode_local_media")
modular = is_minimax_h3_modular(str(od_config.model), od_config.revision)
self.partition = _minimax_h3_partition_for_task(
getattr(od_config, "task_type", None),
str(od_config.model),
)
if modular and str(od_config.task_type or "auto").lower() == "auto":
self.partition = "fl2va"
self._turbo_lora_specs: dict[int, TurboSpec] = {}
self._native_lora_adapter_ids: set[int] = set()
self._lora_sigma_schedules: dict[int, DMD2SigmaSchedule] = {}
model_root = _resolve_minimax_h3_model_root(
str(od_config.model),
od_config.revision,
self.partition,
load_text_encoder=self.load_text_encoder,
)
if modular:
model_path = model_root
self._fasth3_checkpoint = FastH3CheckpointSpec.from_metadata(
json.loads((model_root / "fastvideo_inference.json").read_text(encoding="utf-8"))
)
self._fasth3_checkpoint.check_serving_contract(partition=self.partition, od_config=od_config)
release = self._fasth3_checkpoint.release_metadata()
vae_model_path = self._fasth3_checkpoint.resolve_native_vaes(model_root)
else:
model_path = model_root / ("Ref2VA" if self.partition == "ref2va" else "FL2VA")
model_index = json.loads((model_path / "model_index.json").read_text(encoding="utf-8"))
release = model_index.get("_minimax_h3") or {}
vae_model_path = model_path
partition = str(release.get("partition", "")).lower()
expected_partition = "ref2va" if self.partition == "ref2va" else "fl2va"
if partition != expected_partition:
raise ValueError(f"invalid MiniMax-H3 {expected_partition} partition at {model_path}")
supported_tasks = {str(task).lower() for task in release.get("tasks", [])}
if not supported_tasks:
supported_tasks = {"ref2va"} if partition == "ref2va" else {"t2va", "fl2va"}
ref2va_model_path = None
if self.partition == "combined":
ref2va_model_path = model_root / "Ref2VA"
ref2va_index_path = ref2va_model_path / "model_index.json"
if not ref2va_index_path.is_file():
raise ValueError(f"Ref2VA partition not found at {ref2va_model_path}")
ref2va_index = json.loads(ref2va_index_path.read_text(encoding="utf-8"))
ref2va_release = ref2va_index.get("_minimax_h3") or {}
if str(ref2va_release.get("partition", "")).lower() != "ref2va":
raise ValueError(f"invalid MiniMax-H3 ref2va partition at {ref2va_model_path}")
supported_tasks.update(str(task).lower() for task in ref2va_release.get("tasks", ["ref2va"]))
self.supported_tasks = frozenset(supported_tasks)
shifts = release.get("sigma_shift_scales") or {}
self.default_video_shift = float(shifts.get("video", 12.0))
self.default_audio_shift = float(shifts.get("audio", 3.0))
# Distilled releases pin their own few-step rectified-flow positions; the
# uniform schedule derived from num_inference_steps does not match what
# such a checkpoint was trained on. Each partition carries its own
# contract, so a distilled FL2VA must not drag Ref2VA onto its schedule.
self._base_schedule_by_partition = {expected_partition: _read_base_schedule(release)}
if ref2va_model_path is not None:
self._base_schedule_by_partition["ref2va"] = _read_base_schedule(ref2va_release)
self.weights_sources = [
DiffusersPipelineLoader.ComponentSource(
model_or_path=str(model_path),
subfolder="transformer",
revision=od_config.revision,
prefix="transformer.",
fall_back_to_pt=False,
)
]
self._dit_modules = ["transformer"]
if ref2va_model_path is not None:
self.weights_sources.append(
DiffusersPipelineLoader.ComponentSource(
model_or_path=str(ref2va_model_path),
subfolder="transformer",
revision=od_config.revision,
prefix="transformers_ref.",
fall_back_to_pt=False,
)
)
self._dit_modules.append("transformers_ref")
transformer_quant_config = _resolve_component_quant_config(
od_config.quantization_config,
"transformer",
)
self.transformer = MiniMaxH3DiTModel(
od_config,
quant_config=transformer_quant_config,
diffusers_weights=modular,
)
if self._fasth3_checkpoint is not None:
self.transformer.enable_vsa_gates(sparsity=self._fasth3_checkpoint.vsa_sparsity)
logger.info(
"FastH3 V2 full checkpoint: 8 transformer forwards, video/audio shifts 10/3, VSA sparsity=0.8 tile=64"
)
if ref2va_model_path is not None:
self.transformers_ref = MiniMaxH3DiTModel(
od_config,
quant_config=transformer_quant_config,
diffusers_weights=modular,
)
self._fasth3 = resolve_fasth3_fusion(od_config, self.transformer)
if self._fasth3 is not None and self._fasth3.requires_vsa:
# The artifact assigns a compression gate per DiT block, so those
# modules have to exist before load_weights streams them in. Only
# the ``transformer.`` stream is fused, and ``check_task`` admits
# T2VA only, so the Ref2VA DiT would carry 50 gates that nothing
# ever fills or reads.
self.transformer.enable_vsa_gates()
if self._fasth3 is not None:
self._fasth3.check_serving_contract(
partition=self.partition,
od_config=od_config,
video_shift=self.default_video_shift,
audio_shift=self.default_audio_shift,
)
if self.load_text_encoder:
self.tokenizer = Qwen2TokenizerFast.from_pretrained(
str(model_path),
subfolder="tokenizer",
local_files_only=os.path.isdir(model_path),
)
self.processor = Qwen3VLProcessor.from_pretrained(
str(model_path),
subfolder="processor",
local_files_only=os.path.isdir(model_path),
)
else:
self.tokenizer = None
self.processor = None
_, rank, dit_world = _dit_rank_world()
self._dit_rank = rank
if self.load_text_encoder:
text_encoder_tp_size = int(getattr(self.parallel_config, "text_encoder_tp_size", 1))
if text_encoder_tp_size < 1:
raise ValueError(f"text_encoder_tp_size must be >= 1, got {text_encoder_tp_size}")
if text_encoder_tp_size > dit_world:
raise ValueError(
f"text_encoder_tp_size must not exceed the DiT group size ({dit_world}), got {text_encoder_tp_size}"
)
# The Qwen3-VL text model uses 64 attention heads / 8 KV heads.
if 64 % text_encoder_tp_size or 8 % text_encoder_tp_size:
raise ValueError(
"text_encoder_tp_size must divide both Qwen3-VL "
f"num_attention_heads (64) and num_key_value_heads (8), "
f"got {text_encoder_tp_size}"
)
self.text_encoder_tp_size = text_encoder_tp_size
self.text_encoder_group = self._build_text_encoder_group(text_encoder_tp_size)
self.text_encoder = MiniMaxH3Qwen3VLEncoder(
os.path.join(model_path, "text_encoder"),
device=self.device,
load_model=rank < text_encoder_tp_size,
encoder_group=self.text_encoder_group,
quant_config=_resolve_minimax_h3_text_encoder_quant_config(od_config.quantization_config),
)
if rank < text_encoder_tp_size:
self.weights_sources.append(
DiffusersPipelineLoader.ComponentSource(
model_or_path=str(model_path),
subfolder="text_encoder",
revision=od_config.revision,
prefix="text_encoder.",
fall_back_to_pt=False,
)
)
else:
self.text_encoder_tp_size = 0
self.text_encoder_group = None
self.text_encoder = None
self._encoder_modules = []
legacy_manual_components = getattr(od_config, "diffusion_offload_config", None) is None and bool(
od_config.enable_layerwise_offload or getattr(od_config, "enable_distributed_layerwise_offload", False)
)
# Preserve the legacy MiniMax-H3 low-residency path. The compact API
# deliberately limits explicit component selection to dit/text_encoder,
# so VAEs stay resident for new configurations.
component_load_device = torch.device("cpu") if legacy_manual_components else self.device
self.video_vae = MiniMaxH3VideoVAE(
os.path.join(vae_model_path, "video_vae"),
device=self.device,
load_device=component_load_device,
decode_only=not self.load_vae_encoder,
trust_remote_code=od_config.trust_remote_code,
)
self.audio_vae = MiniMaxH3AudioVAE(
os.path.join(vae_model_path, "audio_vae"),
device=self.device,
load_device=component_load_device,
decode_only=not self.load_vae_encoder,
trust_remote_code=od_config.trust_remote_code,
)
# Registry-side VAE patch-parallel discovery uses ``pipeline.vae``.
self.vae = self.video_vae
self._dlo_component_cache = None
offloads_text_encoder = should_offload_component(od_config, TEXT_ENCODER_COMPONENT)
needs_component_cache = legacy_manual_components or offloads_text_encoder
if getattr(od_config, "enable_distributed_layerwise_offload", False) and needs_component_cache:
self._dlo_component_cache = BoundedAllocatorCache(self.device)
if legacy_manual_components:
_register_dlo_component_cache(
self._dlo_component_cache,
self.text_encoder,
self.video_vae,
self.audio_vae,
)
elif offloads_text_encoder:
_register_dlo_component_cache(self._dlo_component_cache, self.text_encoder)
self._quality_policy = MiniMaxH3QualityPolicy(od_config)
self._cache_dit_runtime = RequestScopedCacheDiTRuntime(self)
self.setup_diffusion_pipeline_profiler(
enable_diffusion_pipeline_profiler=(od_config.enable_diffusion_pipeline_profiler)
)
def load_weights(
self,
weights: Iterable[tuple[str, torch.Tensor]],
) -> set[str]:
def source_prefix(item: tuple[str, torch.Tensor]) -> str:
name, _ = item
prefix = name.partition(".")[0] + "."
if prefix in {"transformer.", "transformers_ref.", "text_encoder."}:
return prefix
raise ValueError(f"unexpected MiniMax-H3 weight {name!r}")
loaded_with_prefix: set[str] = set()
loaded_prefixes: set[str] = set()
transformer_loaded: set[str] = set()
for prefix, grouped_weights in groupby(weights, key=source_prefix):
if prefix in loaded_prefixes:
raise ValueError(f"MiniMax-H3 weight source {prefix!r} is not contiguous")
loaded_prefixes.add(prefix)
component = getattr(self, prefix.removesuffix("."))
if component is None:
raise ValueError(f"MiniMax-H3 component {prefix.removesuffix('.')!r} is disabled in this deployment")
stream = ((name[len(prefix) :], tensor) for name, tensor in grouped_weights)
if prefix == "transformer." and self._fasth3 is not None:
# Fuse before the model shards anything, which is also the only
# point where the checkpoint's fused QKV/MLP layouts are intact.
stream = self._fasth3.apply(stream)
loaded = component.load_weights(stream)
if prefix == "transformer.":
transformer_loaded = set(loaded)
if prefix != "text_encoder.":
component.post_load_weights()
loaded_with_prefix.update(prefix + name for name in loaded)
# Both VAEs load eagerly in ``__init__`` rather than through
# ``weights_sources``. The text encoder uses the shared component
# loader so online quantization and offload processing follow the same
# path as the DiT.
for component_name in ("video_vae", "audio_vae"):
component = getattr(self, component_name)
if component is None:
continue
loaded_with_prefix.update(f"{component_name}.{name}" for name, _ in component.named_parameters())
if self._fasth3 is not None:
# load_weights only warns on a parameter the model does not have, so
# close the adapter against what the DiT actually consumed.
self._fasth3.validate_fully_applied(transformer_loaded)
if self._fasth3_checkpoint is not None:
required_gates = {
f"blocks.{i}.attn.to_gate_compress.weight" for i in range(self.transformer.arch.num_layers)
}
if missing_gates := required_gates - transformer_loaded:
raise ValueError(f"FastH3 V2 checkpoint is missing compression gates: {sorted(missing_gates)}")
return loaded_with_prefix
@property
def lora_is_fused(self) -> bool:
"""True when --lora-path was consumed as a load-time weight fusion."""
return self._fasth3 is not None
def _transformer_for_task(self, task: str) -> MiniMaxH3DiTModel:
if task == "ref2va" and hasattr(self, "transformers_ref"):
return self.transformers_ref
return self.transformer
def _resolve_sigma_positions(self, task: str, sampling: Any) -> tuple[tuple[float, ...] | None, int]:
"""Pick the rectified-flow positions this request denoises on.
Returns them explicitly, or ``None`` to leave the uniform ladder to be
derived from the step count, together with the count the rest of the
request speaks in.
"""
if self._fasth3 is not None:
# A fused student carries its own positions; the checkpoint
# underneath it is the many-step teacher, whose schedule does not
# apply. Its five points bound four transformer forwards, and
# forwards is the unit ``check_request``, the pinned-schedule branch
# below and Cache-DiT all speak in.
positions = self._fasth3.base_schedule
return positions, len(positions) - 1
sigma_schedule = self._sigma_schedule_for_request(sampling, task)
if sigma_schedule is None:
return None, int(sampling.num_inference_steps or 50)
# The schedule lists sigma boundaries; the denoise loop runs one step
# per interval, and that count is what requests and Cache-DiT speak in.
num_steps = sigma_schedule.num_inference_steps
requested_steps = sampling.num_inference_steps
if requested_steps is not None and int(requested_steps) != num_steps:
raise OmniClientError(
"this MiniMax H3 checkpoint pins a distilled sigma schedule; num_inference_steps "
f"must be {num_steps} or omitted, got {int(requested_steps)}"
)
return sigma_schedule.base_schedule, num_steps
def _base_schedule_for_task(self, task: str) -> DMD2SigmaSchedule | None:
"""Return the distilled schedule of the partition that serves ``task``."""
partition = "ref2va" if task == "ref2va" else "fl2va"
return self._base_schedule_by_partition.get(partition)
def _resolve_task(
self,
requested: str | None,
multi_modal_data: dict[str, Any] | None = None,
*,
turbo_spec: TurboSpec | None = None,
has_native_lora: bool = False,
) -> str:
multi_modal_data = multi_modal_data or {}
if requested is None:
# A Ref2VA-only startup has no FL2VA transformer; preserve its
# historical implicit default even for image-only references.
if self.partition == "ref2va":
requested = "ref2va"
elif multi_modal_data.get("video") is not None or multi_modal_data.get("audio") is not None:
requested = "ref2va"
elif multi_modal_data.get("image") is not None:
requested = "fl2va"
else:
requested = "t2va"
task = str(requested).lower()
if task not in self.supported_tasks:
raise OmniClientError(
f"checkpoint partition {self.partition!r} supports {sorted(self.supported_tasks)}, got task={task!r}"
)
if turbo_spec is not None and task not in turbo_spec.supported_tasks:
raise OmniClientError(
f"{turbo_spec.filename} is a {turbo_spec.task_family} Turbo artifact and serves "
f"{sorted(turbo_spec.supported_tasks)}, got task={task!r}"
)
if has_native_lora and task != "t2va":
raise OmniClientError("MiniMax-H3 native LoRA supports T2VA requests only")
if self._fasth3 is not None:
self._fasth3.check_task(task)
return task
def _build_text_encoder_group(self, text_encoder_tp_size: int) -> Any:
"""Create the encoder tensor-parallel process group.
The encoder group covers the first ``text_encoder_tp_size`` DiT ranks
(the DiT group is always global ranks ``[0, dit_world)``). Every rank
participates in ``new_group`` so the collective completes; ranks
outside the group never run encoder collectives. For a single-rank
encoder we return a lightweight placeholder so non-encoder ranks do
not need to join a ``GroupCoordinator`` that would assert on ranks
outside the group.
"""
if text_encoder_tp_size == 1:
return _SingleRankEncoderGroup(rank=self._dit_rank)
ranks = list(range(text_encoder_tp_size))
return init_world_group(
ranks=ranks,
local_rank=envs.LOCAL_RANK,
backend=current_omni_platform.dist_backend,
)
def _encoder_group_broadcast_tensor(
self,
tensor: torch.Tensor | None,
*,
dtype: torch.dtype,
device: torch.device,
) -> torch.Tensor:
"""Broadcast a tensor from encoder rank 0 over the encoder TP group."""
group = self.text_encoder_group
if group.world_size == 1:
if tensor is None:
raise ValueError("source tensor is required for single-rank execution")
return tensor.to(device=device, dtype=dtype)
shape = torch.zeros(8, dtype=torch.long, device=device)
if group.rank_in_group == 0:
if tensor is None:
raise ValueError("encoder rank 0 must provide a tensor to broadcast")
shape[0] = tensor.ndim
shape[1 : tensor.ndim + 1] = torch.tensor(tensor.shape, device=device)
torch.distributed.broadcast(shape, src=group.ranks[0], group=group.device_group)
ndim = int(shape[0].item())
tensor_shape = tuple(int(value) for value in shape[1 : ndim + 1].tolist())
if group.rank_in_group == 0:
output = tensor.to(device=device, dtype=dtype).contiguous()
else:
output = torch.empty(tensor_shape, device=device, dtype=dtype)
torch.distributed.broadcast(output, src=group.ranks[0], group=group.device_group)
return output
def _distribute_encode_inputs(
self,
ids: torch.Tensor | None,
vision_kwargs: dict[str, torch.Tensor],
) -> torch.Tensor:
"""Fan out encode inputs from encoder rank 0 to the encoder TP ranks.
Mutates ``vision_kwargs`` in place so every encoder rank ends up with
the same vision tensors, and returns the broadcast ``input_ids``.
"""
keys = ("pixel_values", "image_grid_thw", "pixel_values_videos", "video_grid_thw")
key_dtypes = {
"pixel_values": torch.bfloat16,
"pixel_values_videos": torch.bfloat16,
"image_grid_thw": torch.long,
"video_grid_thw": torch.long,
}
group = self.text_encoder_group
device = self.device
if group.world_size == 1:
if ids is None:
raise ValueError("encoder rank 0 must produce input ids")
return ids.to(device=device, dtype=torch.long)
mask = torch.zeros(len(keys), dtype=torch.long, device=device)
if group.rank_in_group == 0:
for index, key in enumerate(keys):
mask[index] = 1 if key in vision_kwargs else 0
torch.distributed.broadcast(mask, src=group.ranks[0], group=group.device_group)
if group.rank_in_group == 0:
ids = self._encoder_group_broadcast_tensor(ids, dtype=torch.long, device=device)
else:
ids = self._encoder_group_broadcast_tensor(None, dtype=torch.long, device=device)
for index, key in enumerate(keys):
if mask[index].item() == 0:
continue
source = vision_kwargs.get(key) if group.rank_in_group == 0 else None
vision_kwargs[key] = self._encoder_group_broadcast_tensor(
source,
dtype=key_dtypes[key],
device=device,
)
return ids
def _encode_text_hidden(
self,
input_ids: torch.Tensor,
vision_kwargs: dict[str, torch.Tensor],
) -> torch.Tensor:
if getattr(self, "_model_cpu_offload_modules", None):
# Invoke nn.Module.__call__ so the generic model-level offloader
# swaps the resident DiT and encoder.
return self.text_encoder(input_ids, **vision_kwargs)
if self._uses_manual_component_offload(self.text_encoder):
with self._component_on_device(self.text_encoder):
return self.text_encoder.encode_ids(input_ids, **vision_kwargs)
# Keep Qwen resident when it is not selected for layerwise offload.
self.text_encoder.load_to_device()
return self.text_encoder.encode_ids(input_ids, **vision_kwargs)
def _uses_manual_component_offload(self, component: nn.Module) -> bool:
od_config = getattr(self, "od_config", None)
if od_config is None:
return False
if getattr(od_config, "diffusion_offload_config", None) is None:
return bool(
getattr(od_config, "enable_layerwise_offload", False)
or getattr(od_config, "enable_distributed_layerwise_offload", False)
)
return component is getattr(self, "text_encoder", None) and should_offload_component(
od_config, TEXT_ENCODER_COMPONENT
)
def enable_omni_model_cpu_offload(
self,
*,
device: torch.device,
pin_memory: bool,
use_hsdp: bool,
offload_components: frozenset[str] | None = None,
) -> None:
if getattr(self, "_model_cpu_offload_modules", None):
return
components = ModuleDiscovery.discover(self)
dits = components.dits
stages = [*components.encoders, *components.vaes]
modules = [*dits, *stages]
selection_options: dict[str, Any] = {}
if offload_components is not None:
if DIT_COMPONENT in offload_components and not dits:
raise ValueError("MiniMax-H3 has no loaded DiT for selected module offload")
if TEXT_ENCODER_COMPONENT in offload_components and not components.encoders:
raise ValueError("MiniMax-H3 has no loaded text encoder for selected module offload")
selection_options = {
"offload_dit_modules": dits if DIT_COMPONENT in offload_components else (),
"offload_encoder_modules": (
components.encoders if TEXT_ENCODER_COMPONENT in offload_components else ()
),
}
apply_sequential_offload(
dit_modules=dits,
encoder_modules=stages,
device=device,
pin_memory=pin_memory,
use_hsdp=use_hsdp,
offload_initial_dits=offload_components is None or DIT_COMPONENT in offload_components,
**selection_options,
)
self._model_cpu_offload_modules = modules
logger.info(
"MiniMax-H3 model-level CPU offload enabled for selected components: %s",
sorted(offload_components) if offload_components is not None else "legacy full topology",
)
def disable_omni_model_cpu_offload(self) -> None:
modules = getattr(self, "_model_cpu_offload_modules", None)
if not modules:
return
remove_sequential_offload(modules)
self._model_cpu_offload_modules = []
@contextmanager
def _component_on_device(self, component: nn.Module):
if getattr(self, "_model_cpu_offload_modules", None):
with sequential_offload_component(component):
yield
return
staged = self._uses_manual_component_offload(component)
try:
if staged:
component.load_to_device()
yield
except BaseException:
if staged:
try:
component.offload_to_cpu()
except BaseException:
logger.exception("Failed to release %s after component failure", component.__class__.__name__)
cache = getattr(self, "_dlo_component_cache", None)
if cache is not None:
try:
cache.release_if_needed(force=True)
except BaseException:
logger.exception("Failed to release retained allocator cache after component failure")
raise
else:
if staged:
try:
component.offload_to_cpu()
except BaseException:
cache = getattr(self, "_dlo_component_cache", None)
if cache is not None:
try:
cache.release_if_needed(force=True)
except BaseException:
logger.exception("Failed to release retained allocator cache after offload failure")
raise
@staticmethod
def _is_output_owner_rank() -> bool:
"""Whether this rank's output is returned by the diffusion executor."""
return not dist.is_available() or not dist.is_initialized() or dist.get_rank() == 0
def _offload_model_cpu_stage_output(self, tensor: torch.Tensor) -> torch.Tensor:
"""Release a decoded output's storage before a later seed reloads the DiT.
Model-level CPU offload makes the VAE hook evict the DiT before decode.
For a multi-output request, however, the next seed reloads the DiT while
the preceding decoded output would otherwise still occupy the device.
The reply rank performs the D2H copy early; ranks whose outputs are not
consumed retain only a metadata placeholder for the final concatenation.
Call this on the tensor that is actually returned to the engine. Audio is
released inside ``decode``; video is released by the callers of ``decode``
only after ``_prepare_minimax_h3_video_output`` has quantized it, so the
early D2H copy carries ``uint8`` frames rather than the decoded floats.
"""
if not getattr(self, "_model_cpu_offload_modules", None):
return tensor
if self._is_output_owner_rank():
return tensor.cpu()
return torch.empty(tensor.shape, dtype=tensor.dtype, device="meta")
def _initial_noise(
self,
*,
seed: int,
latent_t: int,
latent_h: int,
latent_w: int,
audio_t: int,
) -> tuple[torch.Tensor, torch.Tensor]:
video_generator = torch.Generator(device="cpu").manual_seed(seed)
video = torch.randn(
1,
24,
latent_t,
latent_h,
latent_w,
generator=video_generator,
dtype=torch.float32,
)
video_rows = minimax_h3_patchify_video_latent(
video,
patch_size=(1, 2, 2),
)
audio_generator = torch.Generator(device="cpu").manual_seed(seed)
audio_rows = torch.randn(
audio_t * 2,
32,
generator=audio_generator,
dtype=torch.float32,
)
return video_rows, audio_rows
@contextmanager
def _resident_dit_layers_on_device(self, *, enabled: bool = True):
controller = getattr(self, "_dlo_residency_controller", None)
if controller is not None and enabled:
controller.load_resident_layers()
try:
yield
finally:
if controller is not None and enabled:
controller.offload_resident_layers()
def _build_denoise_inputs(
self,
*,
task: str,
text_embeddings: torch.Tensor,
text_tags: torch.Tensor,
seed: int,
latent_t: int,
latent_h: int,
latent_w: int,
audio_t: int,
num_frames: int,
num_steps: int,
video_shift: float,
audio_shift: float,
base_schedule: Sequence[float] | None,
visual_condition: torch.Tensor | None,
visual_condition_shape: tuple[int, int, int] | None,
audio_condition: torch.Tensor | None,
ref_audio_t: int | None,
ref_blocks: list[dict[str, Any]] | None = None,
visual_condition_shapes: list[tuple[int, int, int]] | None = None,
audio_condition_lengths: list[int] | None = None,
keyframe_frame_indices: list[int] | None = None,
pad_seq_len: int | None = None,
) -> dict[str, Any]:
"""Build the packed layout, initial rows, anchors, and sigma schedules.
Shared by request-mode :meth:`diffuse` and step-mode
:meth:`prepare_encode` so both paths start from identical state.
"""
initial_video, initial_audio = self._initial_noise(
seed=seed,
latent_t=latent_t,
latent_h=latent_h,
latent_w=latent_w,
audio_t=audio_t,
)
if task == "ref2va":
if ref_blocks is None:
if visual_condition_shape is None or ref_audio_t is None:
raise ValueError("ref2va condition metadata is missing")
_, ref_h, ref_w = visual_condition_shape
ref_blocks = [
{"kind": "image", "latent_h": ref_h, "latent_w": ref_w},
{"kind": "audio", "ref_audio_t": ref_audio_t},
]
packed = minimax_h3_packed_sequence_ref2va_blocks(
text_len=int(text_embeddings.shape[0]),
latent_t=latent_t,
latent_h=latent_h,
latent_w=latent_w,
audio_t=audio_t,
ref_blocks=ref_blocks,
seq_len=pad_seq_len,
)
else:
packed = minimax_h3_packed_sequence(
text_len=int(text_embeddings.shape[0]),
latent_t=latent_t,
latent_h=latent_h,
latent_w=latent_w,
audio_t=audio_t,
include_keyframe_cond=task == "fl2va",
keyframe_frame_indices=keyframe_frame_indices if task == "fl2va" else None,
frame_count=num_frames if task == "fl2va" else None,
seq_len=pad_seq_len,
)
# Report the effective shape at info level exactly when a request pins
# it, so a deployment can confirm the pin landed without raising the
# log level for every request.
log = logger.info if pad_seq_len is not None else logger.debug
log(
"MiniMax H3 packed sequence: task=%s pad_seq_len=%s used=%d seq_len=%d",
task,
pad_seq_len,
int(packed["cu_seqlens"][1]),
int(packed["seq_len"]),
)
tags = packed["token_tags"].clone()
tags[packed["text_pos"]] = text_tags.cpu()
branch = MiniMaxH3DenoiseBranch(
packed=packed,
text_embeddings=text_embeddings,
token_tags=tags,
device=self.device,
)
visual_anchor = visual_condition
if visual_anchor is not None:
condition_shapes = visual_condition_shapes
if condition_shapes is None and visual_condition_shape is not None:
condition_shapes = [visual_condition_shape]
if not condition_shapes:
raise ValueError("visual condition shape is missing")
visual_anchor = minimax_h3_imgvid_cond_noise_aug_rows(
visual_anchor,
condition_shapes=condition_shapes,
target_latent_t=latent_t,
imgvid_cond_num_frames=len(condition_shapes),
seed=seed,
noise_aug=MINIMAX_H3_IMGVID_COND_TIMESTEP,
)
full_video = torch.zeros(
branch.img_pos.shape[0],
96,
dtype=torch.float32,
)
full_video[branch.update_mask] = initial_video
initial_video = full_video
audio_anchor = audio_condition
if audio_anchor is not None:
condition_audio_t = audio_condition_lengths
if condition_audio_t is None and ref_audio_t is not None:
condition_audio_t = [ref_audio_t]
if not condition_audio_t:
raise ValueError("reference audio length is missing")
audio_anchor = minimax_h3_audio_cond_noise_aug_rows(
audio_anchor,
condition_audio_t=condition_audio_t,
seed=seed,
noise_aug=MINIMAX_H3_AUDIO_REF_COND_TIMESTEP,
)
full_audio = torch.zeros(
branch.audio_pos.shape[0],
32,
dtype=torch.float32,
)
full_audio[branch.audio_update_mask] = initial_audio
initial_audio = full_audio
video_sigmas = minimax_h3_time_shift_sigmas(
num_steps=num_steps,
shift_scale=video_shift,
base_schedule=base_schedule,
)
audio_sigmas = minimax_h3_time_shift_sigmas(
num_steps=num_steps,
shift_scale=audio_shift,
base_schedule=base_schedule,
)
return {
"branch": branch,
# The request-mode loop moves these onto the device itself; step mode
# keeps them resident across steps, so normalize once for both.
"video_rows": initial_video.to(device=self.device, dtype=torch.float32),
"audio_rows": initial_audio.to(device=self.device, dtype=torch.float32),
"cond_anchor": (
None if visual_anchor is None else visual_anchor.to(device=self.device, dtype=torch.float32)
),
"audio_anchor": (
None if audio_anchor is None else audio_anchor.to(device=self.device, dtype=torch.float32)
),
"sigmas_video": video_sigmas,
"sigmas_audio": audio_sigmas,
}
def _unpack_denoised_rows(
self,
branch: MiniMaxH3DenoiseBranch,
video_rows: torch.Tensor,
audio_rows: torch.Tensor,
*,
latent_t: int,
latent_h: int,
latent_w: int,
audio_t: int,
) -> tuple[torch.Tensor, torch.Tensor]:
"""Select the target rows and unpack them back into VAE latents."""
target_video = video_rows[branch.update_mask_dev]
video_latent = minimax_h3_unpatchify_video_tokens(
target_video,
latent_shape=(
latent_t,
latent_h // 2,
latent_w // 2,
24,
),
patch_size=(1, 2, 2),
)
target_audio = audio_rows[branch.audio_update_mask_dev]
audio_latent = minimax_h3_unpack_audio_tokens(
target_audio,
audio_t=audio_t * 2,
audio_channel=2,
)
return video_latent, audio_latent
def diffuse(
self,
*,
task: str,
text_embeddings: torch.Tensor,
text_tags: torch.Tensor,
seed: int,
latent_t: int,
latent_h: int,
latent_w: int,
audio_t: int,
num_frames: int,
num_steps: int,
video_shift: float,
audio_shift: float,
base_schedule: Sequence[float] | None,
visual_condition: torch.Tensor | None,
visual_condition_shape: tuple[int, int, int] | None,
audio_condition: torch.Tensor | None,
ref_audio_t: int | None,
ref_blocks: list[dict[str, Any]] | None = None,
visual_condition_shapes: list[tuple[int, int, int]] | None = None,
audio_condition_lengths: list[int] | None = None,
keyframe_frame_indices: list[int] | None = None,
pad_seq_len: int | None = None,
) -> tuple[torch.Tensor, torch.Tensor]:
inputs = self._build_denoise_inputs(
task=task,
text_embeddings=text_embeddings,
text_tags=text_tags,
seed=seed,
latent_t=latent_t,
latent_h=latent_h,
latent_w=latent_w,
audio_t=audio_t,
num_frames=num_frames,
num_steps=num_steps,
video_shift=video_shift,
audio_shift=audio_shift,
base_schedule=base_schedule,
visual_condition=visual_condition,
visual_condition_shape=visual_condition_shape,
audio_condition=audio_condition,
ref_audio_t=ref_audio_t,
ref_blocks=ref_blocks,
visual_condition_shapes=visual_condition_shapes,
audio_condition_lengths=audio_condition_lengths,
keyframe_frame_indices=keyframe_frame_indices,
pad_seq_len=pad_seq_len,
)
branch = inputs["branch"]
transformer = self._transformer_for_task(task)
with self._resident_dit_layers_on_device(enabled=transformer is self.transformer):
with self.progress_bar(total=len(inputs["sigmas_video"]) - 1) as progress:
video_rows, audio_rows = minimax_h3_denoise_loop(
model=transformer,
positive=branch,
initial_video_rows=inputs["video_rows"],
initial_audio_rows=inputs["audio_rows"],
keyframe_cond_rows=inputs["cond_anchor"],
audio_ref_rows=inputs["audio_anchor"],
sigmas_video=inputs["sigmas_video"],
sigmas_audio=inputs["sigmas_audio"],
device=self.device,
imgvid_cond_noise_aug_for_inference=(MINIMAX_H3_IMGVID_COND_TIMESTEP),
audio_cond_noise_aug_for_inference=(MINIMAX_H3_AUDIO_REF_COND_TIMESTEP),
on_step=lambda step, video, audio: progress.update(),
)
return self._unpack_denoised_rows(
branch,
video_rows,
audio_rows,
latent_t=latent_t,
latent_h=latent_h,
latent_w=latent_w,
audio_t=audio_t,
)
def decode_to_mp4(
self,
video_latent: torch.Tensor,
audio_latent: torch.Tensor,
*,
height: int,
width: int,
max_pending: int = 2,
batch_frames: int = 17,
video_codec_options: dict[str, str] | None = None,
) -> bytes:
"""Decode and encode one output on the worker without full-video materialization.
Audio is decoded first so the incremental mux session can attach its audio
stream before temporal video chunks arrive. The callback receives committed
float32 ``BCTHW`` frames, performs the requested-size crop and uint8
conversion in the worker, then applies bounded backpressure to the encoder.
"""
from vllm_omni.diffusion.utils.media_utils import ChunkedMP4Encoder
if batch_frames <= 0:
raise ValueError("batch_frames must be positive")
with self._component_on_device(self.audio_vae):
audio = self.audio_vae.decode_latent(audio_latent)
audio_np = audio.detach().float().cpu().numpy()
if audio_np.ndim == 3 and audio_np.shape[0] == 1:
audio_np = audio_np[0]
encoder = ChunkedMP4Encoder(
width=width,
height=height,
fps=MINIMAX_H3_FPS,
audio_waveform=audio_np,
audio_sample_rate=MINIMAX_H3_AUDIO_SAMPLE_RATE,
max_pending=max_pending,
video_codec_options=video_codec_options,
)
pending_chunks: list[torch.Tensor] = []
pending_frames = 0
def flush_pending() -> None:
nonlocal pending_frames
if not pending_chunks:
return
batched = torch.cat(pending_chunks, dim=1)
encoder.push(batched[0].cpu().numpy())
pending_chunks.clear()
pending_frames = 0
def on_chunk(frames: torch.Tensor) -> None:
nonlocal pending_frames
prepared = _prepare_minimax_h3_video_output(frames[..., :height, :width])
if prepared.shape[0] != 1:
raise ValueError("MiniMax H3 chunked MP4 encoding currently expects one output per decoder")
pending_chunks.append(prepared)
pending_frames += int(prepared.shape[1])
if pending_frames >= batch_frames:
flush_pending()
try:
with self._component_on_device(self.video_vae):
with current_omni_platform.create_autocast_context(
device_type=self.device.type,
dtype=torch.float16,
enabled=True,
):
self.video_vae.decode_with_chunks(video_latent, on_chunk=on_chunk)
flush_pending()
return encoder.finish()
except BaseException:
encoder.abort()
raise
def decode(
self,
video_latent: torch.Tensor,
audio_latent: torch.Tensor,
*,
height: int,
width: int,
) -> tuple[torch.Tensor, torch.Tensor]:
with self._component_on_device(self.video_vae):
with current_omni_platform.create_autocast_context(
device_type=self.device.type,
dtype=torch.float16,
enabled=True,
):
video = self.video_vae.decode_latent(video_latent)
video = video[..., :height, :width].contiguous()
with self._component_on_device(self.audio_vae):
audio = self.audio_vae.decode_latent(audio_latent)
audio = self._offload_model_cpu_stage_output(audio)
return video, audio
@staticmethod
def _extract_prompt(raw_prompt: Any) -> tuple[str, dict[str, Any]]:
"""Split a request prompt into its text and multimodal parts."""
if isinstance(raw_prompt, str):
prompt = raw_prompt
multi_modal_data: dict[str, Any] = {}
else:
prompt = str(raw_prompt.get("prompt") or "")
multi_modal_data = raw_prompt.get("multi_modal_data") or {}
if not prompt:
raise OmniClientError("MiniMax H3 requires a non-empty prompt")
return prompt, multi_modal_data
@staticmethod
def _extract_text_conditioning(raw_prompt: Any) -> MiniMaxH3TextConditioning | None:
if isinstance(raw_prompt, str):
return None
additional_information = raw_prompt.get("additional_information") or {}
encoder_output = additional_information.get("encoder_output")
if encoder_output is None:
return None
if not isinstance(encoder_output, Mapping):
raise OmniClientError("MiniMax H3 encoder output must be a mapping")
try:
if "hidden_states" in encoder_output and "token_tags" in encoder_output:
return MiniMaxH3TextConditioning.from_payload(encoder_output)
conditioning = MiniMaxH3EncoderConditioning.from_omni_payload(encoder_output)
return MiniMaxH3TextConditioning(conditioning.hidden_states, conditioning.token_tags)
except ValueError as exc:
raise OmniClientError(str(exc)) from exc
@staticmethod
def _extract_prepared_reference_videos(raw_prompt: Any) -> list[dict[str, Any]] | None:
if isinstance(raw_prompt, str):
return None
additional_information = raw_prompt.get("additional_information") or {}
meta = additional_information.get("meta") or {}
descriptor = meta.get(MINIMAX_H3_PREPARED_REFERENCE_VIDEOS_KEY)
if descriptor is None:
return None
if not isinstance(descriptor, str):
raise OmniClientError("MiniMax H3 prepared-reference-video descriptor must be a string")
try:
_, videos = deserialize_prepared_reference_videos(descriptor)
except ValueError as exc:
raise OmniClientError(str(exc)) from exc
return videos
def encode_prompt(self, prepared: PreparedEncoderInputs | None) -> tuple[torch.Tensor, torch.Tensor]:
_, rank, _ = _dit_rank_world()
ids = tags = hidden = None
vision_kwargs: dict[str, torch.Tensor] = {}
error = None
if rank == 0:
try:
if prepared is None:
raise ValueError("rank 0 must prepare MiniMax H3 text inputs")
if prepared.images:
vision = self.processor.image_processor(images=prepared.images, return_tensors="pt")
vision_kwargs.update(
pixel_values=vision["pixel_values"],
image_grid_thw=vision["image_grid_thw"],
)
if prepared.qwen_videos:
vision = self.processor.video_processor(
videos=[frames for frames, _ in prepared.qwen_videos],
do_sample_frames=False,
return_tensors="pt",
)
vision_kwargs.update(
pixel_values_videos=vision["pixel_values_videos"],
video_grid_thw=vision["video_grid_thw"],
)
ids, tags = build_minimax_h3_presentation(
self.tokenizer,
prompt=prepared.prompt,
task=prepared.media.task,
condition_labels=prepared.condition_labels,
image_grid_thw=vision_kwargs.get("image_grid_thw"),
video_grid_thw=vision_kwargs.get("video_grid_thw"),
video_timestamps=prepared.video_timestamps,
merge_size=int(self.processor.image_processor.merge_size),
)
except Exception as exc:
error = exc
_broadcast_rank0_exception(error)
if rank < self.text_encoder_tp_size:
ids = self._distribute_encode_inputs(ids, vision_kwargs)
hidden = self._encode_text_hidden(ids, vision_kwargs)
return (
_broadcast_tensor(hidden, dtype=torch.bfloat16, device=self.device),
_broadcast_tensor(tags, dtype=torch.long, device=self.device),
)
def _distribute_media_inputs(self, media: MiniMaxH3EncoderMediaInput | None) -> MiniMaxH3EncoderMediaInput:
group, rank, world_size = _dit_rank_world()
if world_size == 1:
assert media is not None
return media
tensors = media.to_mm_tensors() if media is not None else []
header = [(media.to_metadata(), [value.dtype for value in tensors])] if media is not None else [None]
dist.broadcast_object_list(header, src=0, group=group)
metadata, dtypes = header[0]
received = []
for index, dtype in enumerate(dtypes):
value = _broadcast_tensor(tensors[index] if rank == 0 else None, dtype=dtype, device=self.device)
if rank != 0:
received.append(value.cpu())
del value
if rank == 0:
assert media is not None
return media
return MiniMaxH3EncoderMediaInput.from_mm_tensors(received, metadata)
def _broadcast_media_conditioning(
self, conditioning: MiniMaxH3EncoderMediaConditioning | None
) -> MiniMaxH3EncoderMediaConditioning:
group, rank, world_size = _dit_rank_world()
if world_size == 1:
assert conditioning is not None
return conditioning
tensor_names = ("visual_condition", "audio_condition")
header = [None]
if rank == 0:
assert conditioning is not None
metadata = {
item.name: getattr(conditioning, item.name)
for item in fields(conditioning)
if item.name not in tensor_names
}
dtypes = {
name: value.dtype if (value := getattr(conditioning, name)) is not None else None
for name in tensor_names
}
header[0] = (metadata, dtypes)
dist.broadcast_object_list(header, src=0, group=group)
metadata, dtypes = header[0]
tensors = {}
for name, dtype in dtypes.items():
source = getattr(conditioning, name) if rank == 0 else None
tensors[name] = _broadcast_tensor(source, dtype=dtype, device=self.device) if dtype is not None else None
return MiniMaxH3EncoderMediaConditioning(**metadata, **tensors)
def _encode_local_media(self, media: MiniMaxH3EncoderMediaInput | None) -> MiniMaxH3EncoderMediaConditioning:
group, rank, world_size = _dit_rank_world()
distributed_video = self.video_vae.is_distributed_enabled()
if distributed_video:
media = self._distribute_media_inputs(media)
conditioning = None
error = None
try:
if rank == 0 or distributed_video:
assert media is not None
conditioning = encode_media(
media,
video_vae=self.video_vae,
audio_vae=self.audio_vae,
emit_conditioning=rank == 0,
component_scope=self._component_on_device,
)
except ValueError as exc:
error = OmniClientError(str(exc))
except Exception as exc:
error = exc
if world_size > 1:
# All participants finish the codec phase before any latent broadcast.
errors = [None] * world_size
info = (str(error), isinstance(error, OmniClientError)) if error is not None else None
dist.all_gather_object(errors, info, group=group)
for info in errors:
if info is not None:
if error is not None:
raise error
message, is_client_error = info
raise OmniClientError(message) if is_client_error else RuntimeError(message)
elif error is not None:
raise error
return self._broadcast_media_conditioning(conditioning)
def _prepare_local_conditioning(
self,
raw_prompt: Any,
sampling: Any,
*,
require_external_text: bool = False,
) -> MiniMaxH3EncoderConditioning:
group, rank, world_size = _dit_rank_world()
prepared = text_conditioning = None
error = None
if rank == 0:
try:
_, multi_modal_data = self._extract_prompt(raw_prompt)
turbo_spec = self._active_turbo_spec(sampling)
has_native_lora = self._has_active_native_lora(sampling)
task = self._resolve_task(
(sampling.extra_args or {}).get("task"),
multi_modal_data,
turbo_spec=turbo_spec,
has_native_lora=has_native_lora,
)
if turbo_spec is not None:
self._validate_turbo_sampling(sampling, turbo_spec)
if has_native_lora:
self._validate_native_sampling(sampling, task=task)
text_conditioning = self._extract_text_conditioning(raw_prompt)
if require_external_text and text_conditioning is None:
raise OmniClientError(
"MiniMax H3 diffusion stage requires text encoder conditioning when text_encoder is not loaded"
)
prepared = prepare_encoder_inputs(
raw_prompt,
sampling,
task=task,
prepared_reference_videos=self._extract_prepared_reference_videos(raw_prompt),
)
except Exception as exc:
error = exc
_broadcast_rank0_exception(error)
reuse_text = [text_conditioning is not None]
if world_size > 1:
dist.broadcast_object_list(reuse_text, src=0, group=group)
if reuse_text[0]:
hidden = _broadcast_tensor(
text_conditioning.hidden_states if rank == 0 else None,
dtype=torch.bfloat16,
device=self.device,
)
tags = _broadcast_tensor(
text_conditioning.token_tags if rank == 0 else None, dtype=torch.long, device=self.device
)
else:
hidden, tags = self.encode_prompt(prepared)
media = self._encode_local_media(prepared.media if prepared is not None else None)
return MiniMaxH3EncoderConditioning.from_components(MiniMaxH3TextConditioning(hidden, tags), media)
def _prepare_request_inputs(self, raw_prompt: Any, sampling: Any) -> dict[str, Any]:
if self.load_text_encoder:
conditioning = self._prepare_local_conditioning(raw_prompt, sampling)
elif self.load_vae_encoder:
conditioning = self._prepare_local_conditioning(
raw_prompt,
sampling,
require_external_text=True,
)
else:
conditioning = self._extract_encoder_conditioning(raw_prompt)
return self._prepare_encoder_conditioning_inputs(conditioning, sampling)
@staticmethod
def _extract_encoder_conditioning(prompt: Any) -> MiniMaxH3EncoderConditioning:
if isinstance(prompt, list):
prompt = prompt[0] if prompt else None
additional_information = prompt.get("additional_information") if isinstance(prompt, Mapping) else None
payload = additional_information.get("encoder_output") if isinstance(additional_information, Mapping) else None
if not isinstance(payload, Mapping) or not payload:
raise OmniClientError("MiniMax H3 diffusion stage requires encoder conditioning from the encoder stage")
try:
return MiniMaxH3EncoderConditioning.from_omni_payload(payload)
except ValueError as exc:
raise OmniClientError(str(exc)) from exc
def _prepare_encoder_conditioning_inputs(
self,
conditioning: MiniMaxH3EncoderConditioning,
sampling: Any,
) -> dict[str, Any]:
extra = sampling.extra_args or {}
requested_task = extra.get("task")
if requested_task is not None and str(requested_task).lower() != conditioning.task:
raise OmniClientError(
f"MiniMax H3 encoder task {conditioning.task!r} does not match diffusion request {requested_task!r}"
)
turbo_spec = self._active_turbo_spec(sampling)
has_native_lora = self._has_active_native_lora(sampling)
task = self._resolve_task(
conditioning.task,
turbo_spec=turbo_spec,
has_native_lora=has_native_lora,
)
if turbo_spec is not None:
self._validate_turbo_sampling(sampling, turbo_spec)
if has_native_lora:
self._validate_native_sampling(sampling, task=task)
if self._fasth3_checkpoint is not None:
self._fasth3_checkpoint.check_request(
sampling,
step_execution=bool(getattr(self.od_config, "step_execution", False)),
)
if self._fasth3 is not None:
self._fasth3.check_request(
sampling,
video_shift=self.default_video_shift,
audio_shift=self.default_audio_shift,
)
if conditioning.height % 32 or conditioning.width % 32:
raise OmniClientError(
f"MiniMax H3 encoder canvas must be divisible by 32, got {conditioning.width}x{conditioning.height}"
)
expected_latent_t = MINIMAX_H3_SHAPE_PLANNER.video_latent_t(conditioning.num_frames)
expected_audio_t = MINIMAX_H3_SHAPE_PLANNER.audio_latent_t(conditioning.num_frames / MINIMAX_H3_FPS)
if (conditioning.latent_t, conditioning.audio_t) != (expected_latent_t, expected_audio_t):
raise OmniClientError(
"MiniMax H3 encoder latent shape does not match its output frame count: "
f"got ({conditioning.latent_t}, {conditioning.audio_t}), expected "
f"({expected_latent_t}, {expected_audio_t})"
)
visual_shapes = list(conditioning.visual_condition_shapes) or None
audio_lengths = list(conditioning.audio_condition_lengths) or None
base_schedule, num_steps = self._resolve_sigma_positions(task, sampling)
quality_plan = self._quality_policy.resolve(
quality=sampling.quality,
num_inference_steps=num_steps,
extra_args=extra,
)
self._cache_dit_runtime.prepare(quality_plan.cache_dit)
return {
"task": task,
"height": conditioning.height,
"width": conditioning.width,
"num_frames": conditioning.num_frames,
"latent_t": conditioning.latent_t,
"latent_h": conditioning.height // 16,
"latent_w": conditioning.width // 16,
"audio_t": conditioning.audio_t,
"text_embeddings": conditioning.hidden_states.to(device=self.device, dtype=torch.bfloat16),
"text_tags": conditioning.token_tags.to(device=self.device, dtype=torch.long),
"visual_condition": (
conditioning.visual_condition.to(device=self.device)
if conditioning.visual_condition is not None
else None
),
"visual_condition_shape": visual_shapes[0] if visual_shapes and len(visual_shapes) == 1 else None,
"audio_condition": (
conditioning.audio_condition.to(device=self.device)
if conditioning.audio_condition is not None
else None
),
"ref_audio_t": audio_lengths[0] if audio_lengths and len(audio_lengths) == 1 else None,
"ref_blocks": list(conditioning.ref_blocks) or None,
"visual_condition_shapes": visual_shapes,
"audio_condition_lengths": audio_lengths,
"keyframe_frame_indices": list(conditioning.keyframe_frame_indices) or None,
"pad_seq_len": _resolve_pad_seq_len(extra.get("pad_seq_len")),
"seed": int(sampling.seed if sampling.seed is not None else 42),
"num_steps": num_steps,
"video_shift": float(extra.get("flow_shift", self.default_video_shift)),
"audio_shift": float(extra.get("audio_flow_shift", self.default_audio_shift)),
"base_schedule": base_schedule,
"num_outputs": _resolve_minimax_h3_num_outputs(sampling.num_outputs_per_prompt),
"preencode_mp4": bool(extra.get("preencode_mp4", False)),
"preencode_batch_frames": (
normalize_preencode_batch_frames(extra.get("preencode_batch_frames", 17))
if extra.get("preencode_mp4", False)
else 17
),
"video_codec_options": normalize_video_codec_options(
extra.get("video_codec_options", {"preset": "ultrafast", "threads": "0"})
if extra.get("preencode_mp4", False)
else None
),
}
@staticmethod
def _denoise_kwargs(context: dict[str, Any]) -> dict[str, Any]:
"""Select the denoise-input arguments from a prepared request context."""
return {key: context[key] for key in _MINIMAX_H3_DENOISE_INPUT_KEYS}
@torch.no_grad()
def forward(self, request: DiffusionRequestBatch) -> DiffusionOutput:
if len(request.prompts) != 1:
raise OmniClientError("MiniMax H3 supports one request at a time")
context = self._prepare_request_inputs(
request.prompts[0],
request.sampling_params,
)
denoise_kwargs = self._denoise_kwargs(context)
num_outputs = context["num_outputs"]
videos = []
audios = []
for output_seed in _minimax_h3_output_seeds(context["seed"], num_outputs):
video_latent, audio_latent = self.diffuse(**{**denoise_kwargs, "seed": output_seed})
if context["preencode_mp4"]:
videos.append(
self.decode_to_mp4(
video_latent,
audio_latent,
height=context["height"],
width=context["width"],
video_codec_options=context["video_codec_options"],
batch_frames=context["preencode_batch_frames"],
)
)
audios.append(None)
else:
video, audio = self.decode(
video_latent,
audio_latent,
height=context["height"],
width=context["width"],
)
# Rebind rather than append the expression: the local would
# otherwise keep the decoded frames' device storage alive for the
# rest of the iteration, and the next seed's diffuse() -- the DiT
# reload this release exists to make room for -- runs before the
# loop rebinds it. post_decode() rebinds for the same reason.
video = self._offload_model_cpu_stage_output(_prepare_minimax_h3_video_output(video))
videos.append(video)
audios.append(audio)
if videos and isinstance(videos[0], bytes):
video = videos[0] if len(videos) == 1 else videos
audio = None
else:
video = videos[0] if len(videos) == 1 else torch.cat(videos, dim=0)
audio = audios[0] if len(audios) == 1 else torch.cat(audios, dim=0)
return DiffusionOutput(
output=(video, audio),
post_process_func=get_minimax_h3_post_process_func(self.od_config),
stage_durations=(self.stage_durations if hasattr(self, "_stage_durations") else {}),
)
# ------------------------------------------------------------------
# Step-wise execution (continuous batching)
# ------------------------------------------------------------------
@staticmethod
def _packed_batch_supported(transformer: MiniMaxH3DiTModel) -> bool:
"""Whether every attention in this DiT honors multi-document cu_seqlens.
A packed batch is only isolated if *all* of them do: the token refiner
runs under its own attention role and can resolve to a different backend
from the DiT blocks. Ring sequence parallelism dispatches through
``RingParallelAttention``, whose kernels ignore the packed
``cu_seqlens`` metadata regardless of the configured backend; packing
multiple requests under ring would let attention cross document
boundaries, so any layer running ring disqualifies the batch.
The gate probes a per-backend capability rather than a fixed backend
name: FLASH_ATTN, for example, only isolates arbitrary N-document
packed cu_seqlens on CUDA/ROCm/MUSA. Its NPU path only accepts a
``[real, pad]`` two-document layout and its XPU path ignores
cu_seqlens outright — either would silently attend across request
boundaries.
"""
attentions = [module for module in transformer.modules() if isinstance(module, MiniMaxH3Attention)]
if not attentions:
return False
return all(_attention_isolates_packed_requests(module.attention) for module in attentions)
def prepare_encode(self, state: StepRequestState, **kwargs: Any) -> StepRequestState:
"""Run every request-level stage once and seed the per-request step state."""
del kwargs
# Two request-mode features have no place in the shared step contract:
# a request state carries exactly one latent tensor, and distributed
# layerwise offload streams the DiT around one whole denoise loop rather
# than around a single scheduler-driven step.
num_outputs = _resolve_minimax_h3_num_outputs(state.sampling.num_outputs_per_prompt)
if num_outputs != 1:
raise OmniClientError(
f"MiniMax H3 step execution produces one output per request, got num_outputs_per_prompt={num_outputs}"
)
if getattr(self, "_dlo_residency_controller", None) is not None:
raise ValueError(
"MiniMax H3 step execution is not compatible with distributed layerwise offload; "
"the resident-layer window spans a whole denoise loop, so per-step streaming would "
"reload the DiT every step. Drop --step-execution or --enable-distributed-layerwise-offload."
)
# Request-scoped Cache-DiT (quality=high) mutates hook state on the
# shared transformer rather than on ``StepRequestState``. In step mode
# two requests can interleave denoise steps, or be co-batched into a
# single forward, and the second one would then re-enter the DiT with
# cache buffers shaped for the first. Reject the profile here rather
# than let it corrupt outputs at runtime; startup-configured Cache-DiT
# is already blocked in ``DiffusionModelRunner.execute_stepwise``.
if getattr(state.sampling, "quality", None) == "high":
raise OmniClientError(
"MiniMax H3 step execution does not support the high-quality Cache-DiT profile "
"(quality=high); its hooks live on the shared transformer, so interleaved or "
"co-batched requests would reuse incompatible cache state. Drop --step-execution "
"or omit quality=high."
)
context = self._prepare_request_inputs(
state.prompt,
state.sampling,
)
inputs = self._build_denoise_inputs(**self._denoise_kwargs(context))
sigmas_video = inputs["sigmas_video"]
sigmas_audio = inputs["sigmas_audio"]
if len(sigmas_video) < 2:
raise OmniClientError(
f"MiniMax H3 step execution needs at least one denoise step, got num_inference_steps="
f"{len(sigmas_video) - 1}"
)
branch = inputs["branch"]
video_rows, audio_rows, cond_anchor, audio_anchor = minimax_h3_prepare_denoise_rows(
positive=branch,
initial_video_rows=inputs["video_rows"],
initial_audio_rows=inputs["audio_rows"],
keyframe_cond_rows=inputs["cond_anchor"],
audio_ref_rows=inputs["audio_anchor"],
device=self.device,
)
# The denoise loop consumes sigma pairs, so the schedule carries one more
# point than there are steps. ``timesteps`` holds the video branch because
# the shared contract gives a request exactly one timestep sequence; the
# audio schedule rides along in ``extra``.
state.timesteps = torch.tensor(
[1.0 - sigma for sigma in sigmas_video[:-1]],
dtype=torch.float32,
device=self.device,
)
state.step_index = 0
# Video rows are the batched tensor the runner slices per request; audio
# rows have a different width, so they stay request-private.
state.latents = video_rows
state.do_true_cfg = False # H3 checkpoints are CFG-distilled.
state.extra.update(
{
_STEP_BRANCH: branch,
_STEP_TRANSFORMER: self._transformer_for_task(context["task"]),
_STEP_AUDIO_ROWS: audio_rows,
_STEP_COND_ANCHOR: cond_anchor,
_STEP_AUDIO_ANCHOR: audio_anchor,
_STEP_SIGMAS_VIDEO: sigmas_video,
_STEP_SIGMAS_AUDIO: sigmas_audio,
_STEP_SHAPE: {
"height": context["height"],
"width": context["width"],
"latent_t": context["latent_t"],
"latent_h": context["latent_h"],
"latent_w": context["latent_w"],
"audio_t": context["audio_t"],
"preencode_mp4": context.get("preencode_mp4", False),
"preencode_batch_frames": context.get("preencode_batch_frames", 17),
"video_codec_options": context.get("video_codec_options"),
},
}
)
return state
def denoise_step(
self,
input_batch: InputBatch,
*,
states: Sequence[StepRequestState] | None = None,
**kwargs: Any,
) -> torch.Tensor | None:
"""Run one denoise forward covering every request in the batch.
Requests are concatenated into a single packed sequence that keeps one
attention document each, so the whole batch costs one DiT forward.
Backends that ignore ``cu_seqlens`` cannot express that isolation, so
they fall back to one forward per request.
"""
del kwargs
batch_states = list(states if states is not None else input_batch.states)
branches = [state.extra[_STEP_BRANCH] for state in batch_states]
video_rows = [state.latents for state in batch_states]
audio_rows = [state.extra[_STEP_AUDIO_ROWS] for state in batch_states]
schedules = [_minimax_h3_step_schedule(state) for state in batch_states]
transformers = [state.extra[_STEP_TRANSFORMER] for state in batch_states]
mixed_transformers = len({id(transformer) for transformer in transformers}) > 1
# Both execution modes must publish denoise progress for step-gated
# attention features. Requests can differ in both step index and sigma
# schedule, so a batch that is not at one single point has nothing to
# publish and those gates stay dense -- which is their safe default.
progress = {
(state.step_index, schedule["sigma_video"], len(state.extra[_STEP_SIGMAS_VIDEO]) - 1)
for state, schedule in zip(batch_states, schedules)
}
minimax_h3_publish_denoise_progress(*(progress.pop() if len(progress) == 1 else (None, None, None)))
if len(batch_states) > 1 and (mixed_transformers or not self._packed_batch_supported(transformers[0])):
if mixed_transformers:
logger.warning_once(
"MiniMax H3 step batch contains requests for different task-specific DiTs; "
"running %d requests one forward at a time.",
len(batch_states),
)
elif any(
getattr(getattr(module, "attention", None), "use_ring", False)
for module in transformers[0].modules()
if isinstance(module, MiniMaxH3Attention)
):
logger.warning_once(
"MiniMax H3 step batching is disabled when ring attention is active: "
"the ring kernels ignore packed cu_seqlens and would attend across request "
"boundaries. Running %d requests one forward at a time.",
len(batch_states),
)
else:
logger.warning_once(
"MiniMax H3 step batching needs every attention on a backend that isolates "
"packed multi-document cu_seqlens (see AttentionBackend."
"supports_multi_doc_packed_varlen); running %d requests one forward at a time.",
len(batch_states),
)
video_parts: list[torch.Tensor] = []
audio_parts: list[torch.Tensor] = []
for index, branch in enumerate(branches):
forward_kwargs = branch.forward_kwargs(
video_rows=video_rows[index],
audio_rows=audio_rows[index],
t_video=schedules[index]["t_video"],
t_audio=schedules[index]["t_audio"],
imgvid_cond_timestep=schedules[index]["imgvid_cond_timestep"],
audio_ref_cond_timestep=schedules[index]["audio_ref_cond_timestep"],
)
request_video, request_audio = transformers[index](**forward_kwargs)
video_parts.append(request_video)
audio_parts.append(request_audio)
video_velocity = torch.cat(video_parts)
audio_velocity = torch.cat(audio_parts)
else:
forward_kwargs = minimax_h3_batched_forward_kwargs(
branches=branches,
video_rows=video_rows,
audio_rows=audio_rows,
t_video=[schedule["t_video"] for schedule in schedules],
t_audio=[schedule["t_audio"] for schedule in schedules],
imgvid_cond_timesteps=[schedule["imgvid_cond_timestep"] for schedule in schedules],
audio_ref_cond_timesteps=[schedule["audio_ref_cond_timestep"] for schedule in schedules],
)
logger.debug(
"MiniMax H3 denoise step: %d request(s) packed into %d rows",
len(batch_states),
int(forward_kwargs["x"].shape[1]),
)
video_velocity, audio_velocity = transformers[0](**forward_kwargs)
# The shared contract carries one velocity tensor per step, and audio rows
# are a different width than video rows, so hand the audio branch to
# step_scheduler() through request-private state.
audio_parts_by_request = torch.split(audio_velocity, [int(branch.audio_pos.shape[0]) for branch in branches])
for state, request_audio in zip(batch_states, audio_parts_by_request, strict=True):
state.extra[_STEP_AUDIO_NOISE_PRED] = request_audio
return video_velocity
def step_scheduler(self, state: StepRequestState, noise_pred: torch.Tensor, **kwargs: Any) -> None:
"""Apply one Euler-eta0 update to this request's video and audio rows."""
del kwargs
# denoise_step() stages the audio half of this step's velocity; popping
# it keeps a second step_scheduler() call from reusing a stale one.
audio_noise_pred = state.extra.pop(_STEP_AUDIO_NOISE_PRED)
branch = state.extra[_STEP_BRANCH]
schedule = _minimax_h3_step_schedule(state)
update = branch.update_mask_dev
audio_update = branch.audio_update_mask_dev
video_rows = state.latents
audio_rows = state.extra[_STEP_AUDIO_ROWS]
cond_anchor = state.extra[_STEP_COND_ANCHOR]
audio_anchor = state.extra[_STEP_AUDIO_ANCHOR]
device = video_rows.device
x0_video = minimax_h3_rf_v_to_x0(
video_rows[update],
noise_pred.float()[update],
torch.tensor(schedule["t_video"], dtype=torch.float32, device=device),
)
new_video = minimax_h3_euler_eta0_step(
video_rows[update],
x0_video,
sigma_curr=schedule["sigma_video"],
sigma_next=schedule["sigma_video_next"],
)
video_rows = video_rows.clone()
video_rows[update] = new_video
if cond_anchor is not None:
video_rows[~update] = cond_anchor # per-step imgvid cond reset
x0_audio = minimax_h3_rf_v_to_x0(
audio_rows[audio_update],
audio_noise_pred.float()[audio_update],
torch.tensor(schedule["t_audio"], dtype=torch.float32, device=device),
)
new_audio = minimax_h3_euler_eta0_step(
audio_rows[audio_update],
x0_audio,
sigma_curr=schedule["sigma_audio"],
sigma_next=schedule["sigma_audio_next"],
)
audio_rows = audio_rows.clone()
audio_rows[audio_update] = new_audio
if audio_anchor is not None:
audio_rows[~audio_update] = audio_anchor # per-step audio ref reset
state.latents = video_rows
state.extra[_STEP_AUDIO_ROWS] = audio_rows
state.step_index += 1
def post_decode(self, state: StepRequestState, **kwargs: Any) -> DiffusionOutput:
"""Unpack the denoised rows and run the joint video/audio VAE decode."""
del kwargs
shape = state.extra[_STEP_SHAPE]
video_latent, audio_latent = self._unpack_denoised_rows(
state.extra[_STEP_BRANCH],
state.latents,
state.extra[_STEP_AUDIO_ROWS],
latent_t=shape["latent_t"],
latent_h=shape["latent_h"],
latent_w=shape["latent_w"],
audio_t=shape["audio_t"],
)
if shape.get("preencode_mp4", False):
video = self.decode_to_mp4(
video_latent,
audio_latent,
height=shape["height"],
width=shape["width"],
video_codec_options=shape.get("video_codec_options"),
batch_frames=shape.get("preencode_batch_frames", 17),
)
audio = None
else:
video, audio = self.decode(
video_latent,
audio_latent,
height=shape["height"],
width=shape["width"],
)
video = self._offload_model_cpu_stage_output(_prepare_minimax_h3_video_output(video))
return DiffusionOutput(
output=(video, audio),
post_process_func=get_minimax_h3_post_process_func(self.od_config),
stage_durations=(self.stage_durations if hasattr(self, "_stage_durations") else {}),
)
__all__ = [
"MiniMaxH3Pipeline",
"get_minimax_h3_post_process_func",
]
Reference in new issue View Git Blame Copy Permalink