MiniMaxAI/MiniMax-H3

1# SPDX-License-Identifier: Apache-2.02# Token-id and rotary-embedding helpers for the MiniMax H3 visual VAE.3import os4import torch5from typing import Tuple6 7from diffusers.utils import logging8 9logger = logging.get_logger(__name__)  # pylint: disable=invalid-name10 11 12def create_token_ids(patch_dims, device, dtype, id_type="length_normalized", flatten=True):13    coords_list = []14 15    if isinstance(id_type, str):16        id_type_list = [id_type] * len(patch_dims)17    elif isinstance(id_type, list):18        id_type_list = id_type19        if len(id_type_list) != len(patch_dims):20            raise ValueError("id_type list must match patch_dims")21    else:22        raise ValueError("id_type must be a string or a list")23 24    if "area_normalized" in id_type_list or id_type == "area_normalized":25        raise NotImplementedError(26            "area_normalized id_type is not supported in this inference-only bundle"27        )28 29    for _dim_size, _id_type in zip(patch_dims, id_type_list):30        if isinstance(_dim_size, torch.Tensor):31            coords_list.append(_dim_size.to(device=device, dtype=dtype))32            continue33 34        if _id_type == "length_normalized":35            coords = torch.arange(0.5, _dim_size, dtype=dtype, device=device)36            coords = coords / _dim_size37            coords = 2.0 * coords - 1.038        else:39            coords = torch.arange(_dim_size, dtype=dtype, device=device)40 41        coords_list.append(coords)42 43    coords = torch.stack(torch.meshgrid(*coords_list, indexing="ij"), dim=-1)44    if flatten:45        coords = coords.flatten(0, len(patch_dims) - 1)46 47    return coords.unsqueeze(0)48 49 50def _env_flag(name, default="0"):51    value = os.environ.get(name, default)52    return str(value).strip().lower() in ("1", "true", "yes", "on")53 54 55def _env_optional_bool(name, default=""):56    value = str(os.environ.get(name, default)).strip().lower()57    if value in ("", "default", "auto", "none", "unset"):58        return None59    return value not in ("0", "false", "no", "off", "disabled")60 61 62def _vit_torch_compile_kwargs(prefix):63    kwargs = {}64    backend = os.environ.get(f"{prefix}_BACKEND", "inductor").strip()65    mode = os.environ.get(f"{prefix}_MODE", "reduce-overhead").strip()66    if backend and backend.lower() not in ("default", "none"):67        kwargs["backend"] = backend68    if mode and mode.lower() not in ("default", "none"):69        kwargs["mode"] = mode70    kwargs["fullgraph"] = _env_flag(f"{prefix}_FULLGRAPH", "0")71    dynamic = _env_optional_bool(f"{prefix}_DYNAMIC")72    if dynamic is not None:73        kwargs["dynamic"] = dynamic74    return kwargs75 76 77def _rotate_half(x: torch.Tensor) -> torch.Tensor:78    x1, x2 = torch.chunk(x, 2, dim=-1)79    return torch.cat((-x2, x1), dim=-1)80 81 82def _apply_rotary_pos_emb_impl(83    t: torch.Tensor, rotary_pos_emb: Tuple[torch.Tensor, torch.Tensor]84) -> torch.Tensor:85    cos, sin = rotary_pos_emb86 87    if cos.dim() != 4:88        raise ValueError(f"cos must be [B, N, 1, D], got {cos.shape}")89 90    cos = cos.to(t.dtype)91    sin = sin.to(t.dtype)92 93    rot_dim = cos.shape[-1]94    t_dim = t.shape[-1]95 96    if rot_dim < t_dim:97        t_rot, t_pass = t[..., :rot_dim], t[..., rot_dim:]98        t_rot = (t_rot * cos) + (_rotate_half(t_rot) * sin)99        t = torch.cat((t_rot, t_pass), dim=-1)100    else:101        t = (t * cos) + (_rotate_half(t) * sin)102 103    return t104 105_COMPILED_APPLY_ROTARY_POS_EMB = None106_APPLY_ROTARY_POS_EMB_COMPILE_DISABLED = False107 108 109def _get_apply_rotary_pos_emb_impl():110    global _COMPILED_APPLY_ROTARY_POS_EMB, _APPLY_ROTARY_POS_EMB_COMPILE_DISABLED111    if _APPLY_ROTARY_POS_EMB_COMPILE_DISABLED or not _env_flag(112        "MINIMAX_H3_VAE_DECODER_VIT_ROPE_TORCH_COMPILE", "0"113    ):114        return _apply_rotary_pos_emb_impl115    if _COMPILED_APPLY_ROTARY_POS_EMB is not None:116        return _COMPILED_APPLY_ROTARY_POS_EMB117    if not hasattr(torch, "compile"):118        message = "torch.compile is unavailable; falling back to eager ViT rotary embedding"119        if _env_flag("MINIMAX_H3_VAE_DECODER_VIT_ROPE_TORCH_COMPILE_FATAL", "0"):120            raise RuntimeError(message)121        logger.warning(f"[ViTRope] {message}")122        _APPLY_ROTARY_POS_EMB_COMPILE_DISABLED = True123        return _apply_rotary_pos_emb_impl124 125    kwargs = _vit_torch_compile_kwargs("MINIMAX_H3_VAE_DECODER_VIT_ROPE_TORCH_COMPILE")126    try:127        _COMPILED_APPLY_ROTARY_POS_EMB = torch.compile(128            _apply_rotary_pos_emb_impl, **kwargs129        )130        logger.info(f"[ViTRope] torch.compile enabled kwargs={kwargs}")131    except Exception as exc:132        if _env_flag("MINIMAX_H3_VAE_DECODER_VIT_ROPE_TORCH_COMPILE_FATAL", "0"):133            raise134        logger.warning(135            f"[ViTRope] torch.compile setup failed: {type(exc).__name__}: {exc}; "136            "falling back to eager"137        )138        _APPLY_ROTARY_POS_EMB_COMPILE_DISABLED = True139        _COMPILED_APPLY_ROTARY_POS_EMB = None140        return _apply_rotary_pos_emb_impl141    return _COMPILED_APPLY_ROTARY_POS_EMB142 143 144def apply_rotary_pos_emb(145    t: torch.Tensor, rotary_pos_emb: Tuple[torch.Tensor, torch.Tensor]146) -> torch.Tensor:147    global _COMPILED_APPLY_ROTARY_POS_EMB, _APPLY_ROTARY_POS_EMB_COMPILE_DISABLED148    fn = _get_apply_rotary_pos_emb_impl()149    try:150        return fn(t, rotary_pos_emb)151    except Exception as exc:152        if (153            fn is _COMPILED_APPLY_ROTARY_POS_EMB154            and not _env_flag("MINIMAX_H3_VAE_DECODER_VIT_ROPE_TORCH_COMPILE_FATAL", "0")155        ):156            logger.warning(157                f"[ViTRope] compiled call failed: {type(exc).__name__}: {exc}; "158                "disabling compile and retrying eager"159            )160            _APPLY_ROTARY_POS_EMB_COMPILE_DISABLED = True161            _COMPILED_APPLY_ROTARY_POS_EMB = None162            return _apply_rotary_pos_emb_impl(t, rotary_pos_emb)163        raise164