FL2VA/video_vae/func.py
5.7 KB · 164 lines · python Raw
1 # SPDX-License-Identifier: Apache-2.0
2 # Token-id and rotary-embedding helpers for the MiniMax H3 visual VAE.
3 import os
4 import torch
5 from typing import Tuple
6
7 from diffusers.utils import logging
8
9 logger = logging.get_logger(__name__) # pylint: disable=invalid-name
10
11
12 def 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_type
19 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 continue
33
34 if _id_type == "length_normalized":
35 coords = torch.arange(0.5, _dim_size, dtype=dtype, device=device)
36 coords = coords / _dim_size
37 coords = 2.0 * coords - 1.0
38 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
50 def _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
55 def _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 None
59 return value not in ("0", "false", "no", "off", "disabled")
60
61
62 def _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"] = backend
68 if mode and mode.lower() not in ("default", "none"):
69 kwargs["mode"] = mode
70 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"] = dynamic
74 return kwargs
75
76
77 def _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
82 def _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_emb
86
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 t
104
105 _COMPILED_APPLY_ROTARY_POS_EMB = None
106 _APPLY_ROTARY_POS_EMB_COMPILE_DISABLED = False
107
108
109 def _get_apply_rotary_pos_emb_impl():
110 global _COMPILED_APPLY_ROTARY_POS_EMB, _APPLY_ROTARY_POS_EMB_COMPILE_DISABLED
111 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_impl
115 if _COMPILED_APPLY_ROTARY_POS_EMB is not None:
116 return _COMPILED_APPLY_ROTARY_POS_EMB
117 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 = True
123 return _apply_rotary_pos_emb_impl
124
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, **kwargs
129 )
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 raise
134 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 = True
139 _COMPILED_APPLY_ROTARY_POS_EMB = None
140 return _apply_rotary_pos_emb_impl
141 return _COMPILED_APPLY_ROTARY_POS_EMB
142
143
144 def 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_DISABLED
148 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_EMB
154 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 = True
161 _COMPILED_APPLY_ROTARY_POS_EMB = None
162 return _apply_rotary_pos_emb_impl(t, rotary_pos_emb)
163 raise
164