vllm-omni/vllm_omni/model_executor/models/minimax_h3/encoder.py

# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
from __future__ import annotations

from collections.abc import Iterable, Mapping, Sequence
from dataclasses import dataclass
from pathlib import Path
from typing import Any

import torch
from vllm.config import VllmConfig
from vllm.distributed.parallel_state import get_tp_group

from vllm_omni.diffusion.distributed.utils import get_local_device
from vllm_omni.diffusion.models.minimax_h3.vae import (
    MiniMaxH3AudioVAE,
    MiniMaxH3VideoVAE,
)
from vllm_omni.model_executor.models.minimax_h3.conditioning import (
    MINIMAX_H3_ENCODER_LAYOUT_KEY,
    MINIMAX_H3_ENCODER_REQUEST_KEY,
    MiniMaxH3EncoderMediaConditioning,
    MiniMaxH3EncoderMediaInput,
)
from vllm_omni.model_executor.models.minimax_h3.encoder_processing import encode_media
from vllm_omni.model_executor.models.minimax_h3.text_encoder import (
    MiniMaxH3TextEncoderBackbone,
)
from vllm_omni.model_executor.models.output_templates import OmniOutput

_MEDIA_CONDITIONING_CACHE_KEY = "_minimax_h3_encoder_media_conditioning"
_MEDIA_CACHE_KEY = "_minimax_h3_encoder_media"
_COMPONENT_CONFIG_KEY = "minimax_h3_encoder_components"
_COMPONENT_ROLES = frozenset({"text_encoder", "video_vae", "audio_vae"})


def _partition_root(model_path: str) -> Path:
    path = Path(model_path)
    if path.name in {"text_encoder", "video_vae", "audio_vae"}:
        return path.parent
    return path


@dataclass(frozen=True)
class MiniMaxH3EncoderComponentConfig:
    text_parallel_mode: str
    video_parallel_mode: str
    audio_parallel_mode: str

    def __post_init__(self) -> None:
        if self.text_parallel_mode != "tp":
            raise ValueError("MiniMax H3 text encoder currently supports parallel_mode='tp' only")
        if self.video_parallel_mode not in {"leader", "patch"}:
            raise ValueError("MiniMax H3 video VAE parallel_mode must be 'leader' or 'patch'")
        if self.audio_parallel_mode != "leader":
            raise ValueError("MiniMax H3 audio VAE currently supports parallel_mode='leader' only")

    @staticmethod
    def _parallel_mode(raw: Mapping[str, Any], role: str) -> str:
        value = raw[role]
        if not isinstance(value, Mapping):
            raise TypeError(f"MiniMax H3 encoder role {role!r} must be a mapping")
        if set(value) != {"parallel_mode"}:
            raise ValueError(f"MiniMax H3 encoder role {role!r} requires only parallel_mode")
        return str(value["parallel_mode"])

    @classmethod
    def from_vllm_config(cls, vllm_config: VllmConfig) -> MiniMaxH3EncoderComponentConfig:
        hf_config = vllm_config.model_config.hf_config
        raw = getattr(hf_config, _COMPONENT_CONFIG_KEY, None)
        if not isinstance(raw, Mapping):
            raise ValueError(f"MiniMax H3 Encoder Stage requires explicit {_COMPONENT_CONFIG_KEY}")
        roles = set(raw)
        if roles != _COMPONENT_ROLES:
            missing = sorted(_COMPONENT_ROLES - roles)
            unknown = sorted(roles - _COMPONENT_ROLES)
            raise ValueError(
                f"{_COMPONENT_CONFIG_KEY} requires exactly {sorted(_COMPONENT_ROLES)}; "
                f"missing={missing}, unknown={unknown}"
            )
        return cls(
            text_parallel_mode=cls._parallel_mode(raw, "text_encoder"),
            video_parallel_mode=cls._parallel_mode(raw, "video_vae"),
            audio_parallel_mode=cls._parallel_mode(raw, "audio_vae"),
        )


class MiniMaxH3Encoder(MiniMaxH3TextEncoderBackbone):
    requires_request_sample_eligibility = True
    enable_update_additional_information = True
    omni_payload_at_request_end = True

    def __init__(self, *, vllm_config: VllmConfig, prefix: str = "model") -> None:
        super().__init__(vllm_config=vllm_config, prefix=prefix)
        try:
            tp_group = get_tp_group()
            tp_rank = int(tp_group.rank_in_group)
        except (AssertionError, RuntimeError):
            tp_group = None
            tp_rank = 0
        self._component_leader = tp_rank == 0
        self.component_config = MiniMaxH3EncoderComponentConfig.from_vllm_config(vllm_config)
        root = _partition_root(vllm_config.model_config.model)
        missing = [name for name in ("text_encoder", "video_vae", "audio_vae") if not (root / name).is_dir()]
        if missing:
            raise RuntimeError(f"MiniMax H3 Encoder Stage requires all three encoders; missing {missing}")
        device = get_local_device()
        self.video_vae: MiniMaxH3VideoVAE | None = None
        self.audio_vae: MiniMaxH3AudioVAE | None = None
        if self._component_leader or self.component_config.video_parallel_mode == "patch":
            self.video_vae = MiniMaxH3VideoVAE(
                str(root / "video_vae"),
                device=device,
                encode_only=True,
                trust_remote_code=vllm_config.model_config.trust_remote_code,
            )
            if self.component_config.video_parallel_mode == "patch":
                if tp_group is None:
                    raise RuntimeError("MiniMax H3 video VAE patch parallelism requires an initialized TP group")
                self.video_vae.set_parallel_size(
                    tp_group.world_size,
                    process_group=tp_group.device_group,
                )
        if self._component_leader:
            self.audio_vae = MiniMaxH3AudioVAE(
                str(root / "audio_vae"),
                device=device,
                encode_only=True,
                trust_remote_code=vllm_config.model_config.trust_remote_code,
            )

    @staticmethod
    def _media_input(info: dict[str, Any]) -> MiniMaxH3EncoderMediaInput:
        cached = info.get(_MEDIA_CACHE_KEY)
        if isinstance(cached, MiniMaxH3EncoderMediaInput):
            return cached
        meta = info.get("meta")
        metadata = meta.get(MINIMAX_H3_ENCODER_REQUEST_KEY) if isinstance(meta, Mapping) else None
        if not isinstance(metadata, Mapping):
            raise RuntimeError("MiniMax H3 unified encoder requires request metadata")

        hidden_states = info.get("hidden_states")
        raw_media = hidden_states.get("layers") if isinstance(hidden_states, Mapping) else None
        if raw_media is None:
            media_tensors: list[torch.Tensor] = []
        elif isinstance(raw_media, Mapping):
            try:
                ordered_items = sorted(raw_media.items(), key=lambda item: int(item[0]))
            except (TypeError, ValueError) as exc:
                raise RuntimeError("MiniMax H3 encoder media keys must be integer indices") from exc
            if [int(key) for key, _ in ordered_items] != list(range(len(ordered_items))):
                raise RuntimeError("MiniMax H3 encoder media indices must be contiguous from zero")
            media_tensors = [value for _, value in ordered_items]
        else:
            raise RuntimeError("MiniMax H3 encoder media must be a namespaced tensor mapping")
        if any(not isinstance(value, torch.Tensor) for value in media_tensors):
            raise RuntimeError("MiniMax H3 encoder media must contain only tensors")
        media = MiniMaxH3EncoderMediaInput.from_mm_tensors(media_tensors, metadata)
        if media.to_metadata() != dict(metadata):
            raise RuntimeError("MiniMax H3 encoder media does not match request metadata")
        info[_MEDIA_CACHE_KEY] = media
        return media

    def _encode_media(self, media: MiniMaxH3EncoderMediaInput) -> MiniMaxH3EncoderMediaConditioning | None:
        if not self._component_leader and self.component_config.video_parallel_mode == "leader":
            return None
        return encode_media(
            media,
            video_vae=self.video_vae,
            audio_vae=self.audio_vae,
            emit_conditioning=self._component_leader,
        )

    @staticmethod
    def _has_encoder_input(info: object) -> bool:
        if not isinstance(info, Mapping):
            return False
        meta = info.get("meta")
        return isinstance(meta, Mapping) and isinstance(meta.get(MINIMAX_H3_ENCODER_REQUEST_KEY), Mapping)

    def make_omni_output(
        self,
        model_outputs: torch.Tensor | OmniOutput,
        **kwargs: Any,
    ) -> OmniOutput:
        output = super().make_omni_output(model_outputs, **kwargs)
        infos = kwargs.get("model_intermediate_buffer") or []
        if not isinstance(infos, Sequence) or not infos:
            return output
        sample_eligible = kwargs.get("request_sample_eligible")
        if sample_eligible is None:
            sample_eligible = [True] * len(infos)
        if len(sample_eligible) != len(infos):
            raise RuntimeError(
                f"MiniMax H3 unified encoder received {len(sample_eligible)} sampling flags for {len(infos)} requests"
            )
        per_request: list[tuple[torch.Tensor, torch.Tensor, torch.Tensor] | None] = []
        release_workspace = False
        for eligible, info in zip(sample_eligible, infos, strict=True):
            if not eligible:
                per_request.append(None)
                continue
            if not isinstance(info, dict) or not self._has_encoder_input(info):
                raise RuntimeError("MiniMax H3 unified encoder request metadata is missing from the runner")
            release_workspace = True
            if _MEDIA_CONDITIONING_CACHE_KEY not in info:
                info[_MEDIA_CONDITIONING_CACHE_KEY] = self._encode_media(self._media_input(info))
            if not self._component_leader:
                continue
            cached = info.get(_MEDIA_CONDITIONING_CACHE_KEY)
            if not isinstance(cached, MiniMaxH3EncoderMediaConditioning):
                raise RuntimeError("MiniMax H3 encoder components did not produce conditioning")
            per_request.append(cached.to_omni_components())

        if release_workspace:
            torch.accelerator.empty_cache()
        if not self._component_leader:
            return output

        if not any(item is not None for item in per_request):
            return output
        empty = torch.empty((0,), dtype=torch.float32)
        multimodal_outputs = dict(output.multimodal_outputs or {})
        dotted_key = f"kv_metadata.{MINIMAX_H3_ENCODER_LAYOUT_KEY}"
        multimodal_outputs[dotted_key] = [
            item[2] if item is not None else torch.empty((0,), dtype=torch.int64) for item in per_request
        ]

        embed = dict(multimodal_outputs.get("embed") or {})
        embed["embedding"] = [item[0] if item is not None else empty for item in per_request]
        embed["speech_feat"] = [item[1] if item is not None else empty for item in per_request]
        multimodal_outputs["embed"] = embed
        return OmniOutput(
            text_hidden_states=output.text_hidden_states,
            multimodal_outputs=multimodal_outputs,
            intermediate_tensors=output.intermediate_tensors,
            next_token_id=output.next_token_id,
        )

    def load_weights(
        self,
        weights: Iterable[tuple[str, torch.Tensor]],
    ) -> set[str]:
        loaded = super().load_weights(weights)
        for component_name in ("video_vae", "audio_vae"):
            component = getattr(self, component_name)
            if component is not None:
                loaded.update(f"{component_name}.{name}" for name, _ in component.named_parameters())
        return loaded


__all__ = ["MiniMaxH3Encoder"]