File size: 8,369 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 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211 212 213 214 215 216 217 218 219 220 221 222 223 224 225 226 227 228 229 230 231 232 233 234 235 | # SPDX-License-Identifier: Apache-2.0
# Tensor pre/post-processing for the MiniMax H3 visual VAE.
import math
import numpy as np
import torch
from diffusers.utils import logging
from einops import rearrange
from .normalize import get_normalize_transform, get_denormalize_transform
logger = logging.get_logger(__name__) # pylint: disable=invalid-name
class VAEProcessor:
def __init__(
self,
*,
vae_ratio,
vae_ratio_t,
clip_length,
frame_overlap,
token_overlap,
tokens_chunk_size,
isolated_last_frame,
latent_patch_size,
crop_mode,
pixel_norm_type="imagenet",
transform=None,
transform_rev=None,
use_3d_conv=False,
):
self.vae_ratio = vae_ratio
self.vae_ratio_t = vae_ratio_t
self.clip_length = clip_length
self.frame_overlap = frame_overlap
self.token_overlap = token_overlap
self.tokens_chunk_size = tokens_chunk_size
self.isolated_last_frame = isolated_last_frame
self.latent_patch_size = latent_patch_size
self.crop_mode = crop_mode
self.transform = transform or get_normalize_transform(pixel_norm_type)
self.transform_rev = transform_rev or get_denormalize_transform(pixel_norm_type)
self.use_3d_conv = use_3d_conv
def _ensure_list(self, data):
return data if isinstance(data, list) else [data]
def _align_to_total_patch_size(self, h, w):
total_patch_size = self.latent_patch_size * self.vae_ratio
new_h = (h // total_patch_size) * total_patch_size
new_w = (w // total_patch_size) * total_patch_size
return new_h, new_w
def _crop_to_align(self, tensor, new_h, new_w, is_video=False):
if is_video:
_, _, _, h, w = tensor.shape
else:
_, _, h, w = tensor.shape
if self.crop_mode == "center":
top = (h - new_h) // 2
left = (w - new_w) // 2
else:
top = 0
left = 0
if is_video:
return tensor[:, :, :, top : top + new_h, left : left + new_w]
else:
return tensor[:, :, top : top + new_h, left : left + new_w]
def _align_target_token(self, T, mode):
intra_tail = self.clip_length % self.vae_ratio_t
min_frames = intra_tail or self.vae_ratio_t
full_chunks = T // self.clip_length
remainder = T % self.clip_length
if remainder == 0:
return max(T, min_frames)
if mode == "pad":
aligned_r = (
math.ceil((remainder - intra_tail) / self.vae_ratio_t) * self.vae_ratio_t
+ intra_tail
)
if aligned_r > self.clip_length:
return (full_chunks + 1) * self.clip_length + intra_tail
return full_chunks * self.clip_length + aligned_r
else: # trim
k = (remainder - intra_tail) // self.vae_ratio_t
if k >= 0:
target = full_chunks * self.clip_length + k * self.vae_ratio_t + intra_tail
return max(target, min_frames)
elif full_chunks > 0:
return full_chunks * self.clip_length
else:
return min_frames
def _align_target(self, T, mode, granularity):
if granularity == "chunk":
step = self.clip_length
tail = self.frame_overlap
if self.isolated_last_frame:
tail += 1
k = math.ceil((T - tail) / step) if mode == "pad" else (T - tail) // step
return max(k, 1) * step + tail
isolated_extra = 1 if self.isolated_last_frame else 0
return self._align_target_token(T - isolated_extra, mode) + isolated_extra
def align_video_length(self, video_length, mode="pad", granularity="chunk"):
target = self._align_target(video_length, mode, granularity)
delta = target - video_length
if delta > 0 and mode == "trim":
raise ValueError(
f"Cannot trim {video_length} frames to valid length {target}: "
f"not enough frames (granularity={granularity})"
)
return delta
def align_video_length_2pass(self, video_length):
"""Return the leading/trailing frame pads and trailing latent drop.
This is the continuation-prefix (2-pass) alignment. The caller temporarily disables the model's normal token
drop and keeps these mirrored processor fields at zero.
"""
if self.isolated_last_frame:
raise ValueError(
"align_video_length_2pass does not support isolated_last_frame"
)
if self.token_overlap != 0 or self.frame_overlap != 0:
raise ValueError(
"align_video_length_2pass requires token_drop=0 alignment"
)
leading = self.align_video_length(
video_length, mode="pad", granularity="token"
)
token_aligned = video_length + leading
trailing = self.align_video_length(
token_aligned, mode="pad", granularity="chunk"
)
if trailing > 0:
intra_tail = self.clip_length % self.vae_ratio_t
full_chunks = token_aligned // self.clip_length
remainder = token_aligned % self.clip_length
real_tokens = full_chunks * self.tokens_chunk_size
if remainder > 0:
real_tokens += (
(remainder - intra_tail) // self.vae_ratio_t + 1
)
drop_tokens = (
self.get_latent_length(token_aligned + trailing) - real_tokens
)
else:
drop_tokens = 0
return leading, trailing, drop_tokens
def get_suitable_video_length(self, video_length, verbose=False):
used_frame_length = video_length + self.align_video_length(
video_length, mode="trim", granularity="chunk"
)
if verbose:
logger.info(
f"Pick first {used_frame_length} frames from {video_length}-frame video"
)
return used_frame_length
def get_latent_length(self, video_length):
tail_frame = self.frame_overlap
tail_token = self.token_overlap
if self.isolated_last_frame:
tail_frame += 1
tail_token += 1
video_length = self.get_suitable_video_length(video_length)
latent_length = (
int((video_length - tail_frame) // self.clip_length)
* self.tokens_chunk_size
+ tail_token
)
return latent_length
def transform_tensor(self, tensor):
B, T = None, None
if tensor.ndim == 5:
if tensor.shape[2] == 3:
tensor = tensor.transpose(1, 2)
B, _, T, _, _ = tensor.shape
tensor = rearrange(tensor, "b c t h w -> (b t) c h w")
elif tensor.ndim == 4:
if tensor.shape[0] == 3:
tensor = tensor.transpose(0, 1)
elif tensor.ndim == 3:
tensor = tensor.unsqueeze(0)
else:
raise ValueError(f"Unsupported tensor shape: {tensor.shape}")
tensor = self.transform(tensor)
if B is not None and T is not None:
tensor = rearrange(tensor, "(b t) c h w -> b c t h w", b=B, t=T)
return tensor.contiguous()
def revert_tensor(self, tensor):
B, T = None, None
if self.use_3d_conv:
tensor = tensor.unsqueeze(2) if tensor.ndim == 4 else tensor
B, _, T, _, _ = tensor.shape
tensor = rearrange(tensor, "b c t h w -> (b t) c h w")
tensor_rev = self.transform_rev(tensor).clamp(0, 1)
if B is not None:
tensor_rev = rearrange(tensor_rev, "(b t) c h w -> b c t h w", b=B, t=T)
return tensor_rev.contiguous()
@staticmethod
def convert_numpy_to_tensor(numpy_array, device=None):
if isinstance(numpy_array, list):
numpy_array = np.stack(numpy_array, axis=0)
numpy_array = numpy_array.astype(np.float32)
tensor = torch.from_numpy(numpy_array)
tensor = tensor.permute(0, 3, 1, 2)
tensor = tensor / 255.0
if device is not None:
tensor = tensor.to(device)
return tensor
|