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