From 4bcd2055e34829daf16328eb9d8de1356683d310 Mon Sep 17 00:00:00 2001 From: lzx1413 Date: Tue, 4 Aug 2026 10:03:50 +0000 Subject: [PATCH 01/24] feat(minimax-h3): add model components Port the MiniMax H3 joint audio-video DiT, text encoder, audio VAE, and video VAE components with checkpoint conversion and mixed-precision boundaries. Preserve checkpoint-declared FP32 values across dtype conversion and cover architecture, loading, conversion, and forward contracts with focused model unit tests. Verification: MiniMax H3 CPU suite passed (64 tests, 5 subtests); ruff check and format checks passed. --- telefuser/models/minimax_h3_audio/__init__.py | 6 + .../models/minimax_h3_audio/alias_free.py | 167 +++ .../models/minimax_h3_audio/audio_vae.py | 308 +++++ telefuser/models/minimax_h3_audio/bigvgan.py | 239 ++++ telefuser/models/minimax_h3_audio_vae.py | 147 +++ telefuser/models/minimax_h3_dit.py | 671 ++++++++++ telefuser/models/minimax_h3_encoder.py | 159 +++ telefuser/models/minimax_h3_video/__init__.py | 5 + .../models/minimax_h3_video/attention.py | 157 +++ .../models/minimax_h3_video/base_module.py | 195 +++ telefuser/models/minimax_h3_video/conv.py | 83 ++ telefuser/models/minimax_h3_video/flash.py | 135 ++ telefuser/models/minimax_h3_video/klvae.py | 1152 +++++++++++++++++ telefuser/models/minimax_h3_video/norm.py | 247 ++++ .../models/minimax_h3_video/processor.py | 254 ++++ telefuser/models/minimax_h3_video/vae_cnn.py | 262 ++++ telefuser/models/minimax_h3_video/vae_vit.py | 337 +++++ .../models/minimax_h3_video/vit_utils.py | 100 ++ telefuser/models/minimax_h3_video_vae.py | 152 +++ .../unit/models/test_minimax_h3_audio_vae.py | 101 ++ tests/unit/models/test_minimax_h3_dit.py | 119 ++ tests/unit/models/test_minimax_h3_encoder.py | 25 + .../unit/models/test_minimax_h3_video_vae.py | 61 + 23 files changed, 5082 insertions(+) create mode 100644 telefuser/models/minimax_h3_audio/__init__.py create mode 100644 telefuser/models/minimax_h3_audio/alias_free.py create mode 100644 telefuser/models/minimax_h3_audio/audio_vae.py create mode 100644 telefuser/models/minimax_h3_audio/bigvgan.py create mode 100644 telefuser/models/minimax_h3_audio_vae.py create mode 100644 telefuser/models/minimax_h3_dit.py create mode 100644 telefuser/models/minimax_h3_encoder.py create mode 100644 telefuser/models/minimax_h3_video/__init__.py create mode 100644 telefuser/models/minimax_h3_video/attention.py create mode 100644 telefuser/models/minimax_h3_video/base_module.py create mode 100644 telefuser/models/minimax_h3_video/conv.py create mode 100644 telefuser/models/minimax_h3_video/flash.py create mode 100644 telefuser/models/minimax_h3_video/klvae.py create mode 100644 telefuser/models/minimax_h3_video/norm.py create mode 100644 telefuser/models/minimax_h3_video/processor.py create mode 100644 telefuser/models/minimax_h3_video/vae_cnn.py create mode 100644 telefuser/models/minimax_h3_video/vae_vit.py create mode 100644 telefuser/models/minimax_h3_video/vit_utils.py create mode 100644 telefuser/models/minimax_h3_video_vae.py create mode 100644 tests/unit/models/test_minimax_h3_audio_vae.py create mode 100644 tests/unit/models/test_minimax_h3_dit.py create mode 100644 tests/unit/models/test_minimax_h3_encoder.py create mode 100644 tests/unit/models/test_minimax_h3_video_vae.py diff --git a/telefuser/models/minimax_h3_audio/__init__.py b/telefuser/models/minimax_h3_audio/__init__.py new file mode 100644 index 0000000..9c1905a --- /dev/null +++ b/telefuser/models/minimax_h3_audio/__init__.py @@ -0,0 +1,6 @@ +# SPDX-License-Identifier: Apache-2.0 +"""MiniMax H3 audio VAE implementation.""" + +from .audio_vae import DacAudioVAE + +__all__ = ["DacAudioVAE"] diff --git a/telefuser/models/minimax_h3_audio/alias_free.py b/telefuser/models/minimax_h3_audio/alias_free.py new file mode 100644 index 0000000..e7b82c1 --- /dev/null +++ b/telefuser/models/minimax_h3_audio/alias_free.py @@ -0,0 +1,167 @@ +# SPDX-License-Identifier: Apache-2.0 +# Adapted from https://github.com/junjun3518/alias-free-torch under the Apache License 2.0 + +import math + +import torch +import torch.nn as nn +import torch.nn.functional as F + +if "sinc" in dir(torch): + sinc = torch.sinc +else: + # This code is adopted from adefossez's julius.core.sinc under the MIT License + # https://adefossez.github.io/julius/julius/core.html + def sinc(x: torch.Tensor): + """ + Implementation of sinc, i.e. sin(pi * x) / (pi * x) + __Warning__: Different to julius.sinc, the input is multiplied by `pi`! + """ + return torch.where( + x == 0, + torch.tensor(1.0, device=x.device, dtype=x.dtype), + torch.sin(math.pi * x) / math.pi / x, + ) + + +# This code is adopted from adefossez's julius.lowpass.LowPassFilters under the MIT License +# https://adefossez.github.io/julius/julius/lowpass.html +def kaiser_sinc_filter1d(cutoff, half_width, kernel_size): # return filter [1,1,kernel_size] + even = kernel_size % 2 == 0 + half_size = kernel_size // 2 + + # For kaiser window + delta_f = 4 * half_width + A = 2.285 * (half_size - 1) * math.pi * delta_f + 7.95 + if A > 50.0: + beta = 0.1102 * (A - 8.7) + elif A >= 21.0: + beta = 0.5842 * (A - 21) ** 0.4 + 0.07886 * (A - 21.0) + else: + beta = 0.0 + window = torch.kaiser_window(kernel_size, beta=beta, periodic=False) + + # ratio = 0.5/cutoff -> 2 * cutoff = 1 / ratio + if even: + time = torch.arange(-half_size, half_size) + 0.5 + else: + time = torch.arange(kernel_size) - half_size + if cutoff == 0: + filter_ = torch.zeros_like(time) + else: + filter_ = 2 * cutoff * window * sinc(2 * cutoff * time) + """ + Normalize the filter to avoid leaking the constant component. + """ + filter_ /= filter_.sum() + filter = filter_.view(1, 1, kernel_size) + + return filter + + +class LowPassFilter1d(nn.Module): + def __init__( + self, + cutoff=0.5, + half_width=0.6, + stride: int = 1, + padding: bool = True, + padding_mode: str = "replicate", + kernel_size: int = 12, + ): + """ + kernel_size should be even number for stylegan3 setup, in this implementation, odd number is also possible. + """ + super().__init__() + if cutoff < -0.0: + raise ValueError("Minimum cutoff must be larger than zero.") + if cutoff > 0.5: + raise ValueError("A cutoff above 0.5 does not make sense.") + self.kernel_size = kernel_size + self.even = kernel_size % 2 == 0 + self.pad_left = kernel_size // 2 - int(self.even) + self.pad_right = kernel_size // 2 + self.stride = stride + self.padding = padding + self.padding_mode = padding_mode + filter = kaiser_sinc_filter1d(cutoff, half_width, kernel_size) + self.register_buffer("filter", filter) + + # Input [B, C, T] + def forward(self, x): + _, C, _ = x.shape + + if self.padding: + x = F.pad(x, (self.pad_left, self.pad_right), mode=self.padding_mode) + out = F.conv1d(x, self.filter.expand(C, -1, -1), stride=self.stride, groups=C) + + return out + + +class UpSample1d(nn.Module): + def __init__(self, ratio=2, kernel_size=None): + super().__init__() + self.ratio = ratio + self.kernel_size = int(6 * ratio // 2) * 2 if kernel_size is None else kernel_size + self.stride = ratio + self.pad = self.kernel_size // ratio - 1 + self.pad_left = self.pad * self.stride + (self.kernel_size - self.stride) // 2 + self.pad_right = self.pad * self.stride + (self.kernel_size - self.stride + 1) // 2 + filter = kaiser_sinc_filter1d( + cutoff=0.5 / ratio, + half_width=0.6 / ratio, + kernel_size=self.kernel_size, + ) + self.register_buffer("filter", filter) + + def forward(self, x): + _, C, _ = x.shape + + x = F.pad(x, (self.pad, self.pad), mode="replicate") + x = F.conv_transpose1d(x, self.filter.expand(C, -1, -1), stride=self.stride, groups=C) + x.mul_(self.ratio) + x = x[..., self.pad_left : -self.pad_right] + + return x + + +class DownSample1d(nn.Module): + def __init__(self, ratio=2, kernel_size=None): + super().__init__() + self.ratio = ratio + self.kernel_size = int(6 * ratio // 2) * 2 if kernel_size is None else kernel_size + self.lowpass = LowPassFilter1d( + cutoff=0.5 / ratio, + half_width=0.6 / ratio, + stride=ratio, + kernel_size=self.kernel_size, + ) + + def forward(self, x): + xx = self.lowpass(x) + + return xx + + +class Activation1d(nn.Module): + def __init__( + self, + activation, + up_ratio: int = 2, + down_ratio: int = 2, + up_kernel_size: int = 12, + down_kernel_size: int = 12, + ): + super().__init__() + self.up_ratio = up_ratio + self.down_ratio = down_ratio + self.act = activation + self.upsample = UpSample1d(up_ratio, up_kernel_size) + self.downsample = DownSample1d(down_ratio, down_kernel_size) + + def forward(self, x): + x = self.upsample(x) + x = self.act(x) + x = self.downsample(x) + + return x diff --git a/telefuser/models/minimax_h3_audio/audio_vae.py b/telefuser/models/minimax_h3_audio/audio_vae.py new file mode 100644 index 0000000..4a68213 --- /dev/null +++ b/telefuser/models/minimax_h3_audio/audio_vae.py @@ -0,0 +1,308 @@ +# SPDX-License-Identifier: Apache-2.0 +# DAC-lineage audio VAE: waveform encoder + BigVGAN decoder (inference-only bundle). +import math +from typing import List + +import numpy as np +import torch +import torch.nn.functional as F +from torch import nn +from torch.nn.utils.parametrizations import weight_norm + +from telefuser.ops.attention import attention + +from .bigvgan import AttrDict, BigVGAN + + +class GeGluMlp(nn.Module): + def __init__(self, in_features, hidden_features): + super().__init__() + self.norm = nn.LayerNorm(in_features) + self.act = nn.GELU(approximate="tanh") + self.w0 = nn.Linear(in_features, hidden_features) + self.w1 = nn.Linear(in_features, hidden_features) + self.w2 = nn.Linear(hidden_features, in_features) + + def forward(self, x): + x = self.norm(x) + x = self.act(self.w0(x)).mul_(self.w1(x)) + x = self.w2(x) + return x + + +class CausalAttention(nn.Module): + def __init__(self, in_dim, out_dim, num_heads): + super().__init__() + if in_dim > out_dim: + # assert in_dim // num_heads == out_dim + self.head_dim = in_dim // num_heads + self.qkv = nn.Linear(in_dim, in_dim * 3, bias=False) + self.q_bias = nn.Parameter(torch.zeros(in_dim)) + self.v_bias = nn.Parameter(torch.zeros(in_dim)) + self.register_buffer("zero_k_bias", torch.zeros(in_dim)) + else: + # assert out_dim // num_heads == in_dim + self.head_dim = out_dim // num_heads + self.qkv = nn.Linear(in_dim, out_dim * 3, bias=False) + self.q_bias = nn.Parameter(torch.zeros(out_dim)) + self.v_bias = nn.Parameter(torch.zeros(out_dim)) + self.register_buffer("zero_k_bias", torch.zeros(out_dim)) + + self.in_dim = in_dim + self.out_dim = out_dim + self.num_heads = num_heads + self.scale = self.head_dim**-0.5 + self.proj = nn.Linear(out_dim, out_dim) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + B, N, C = x.shape + qkv = F.linear( + input=x, + weight=self.qkv.weight, + bias=torch.cat((self.q_bias, self.zero_k_bias, self.v_bias)), + ) + q, k, v = qkv.reshape(B, N, 3, self.num_heads, self.head_dim).permute(2, 0, 3, 1, 4).unbind(0) + + x = attention( + q, + k, + v, + input_layout="BNSD", + output_layout="BNSD", + is_causal=True, + scale=self.scale, + ) + + if self.in_dim > self.out_dim: + x = torch.mean(x, dim=1) + if self.in_dim // self.num_heads != self.out_dim: + x = nn.functional.adaptive_avg_pool1d(x, self.out_dim) + else: + x = x.transpose(1, 2).reshape(B, N, -1) + x = self.proj(x) + return x + + +class AttnProjection(nn.Module): + def __init__(self, in_dim, out_dim, num_heads, norm_layer=nn.LayerNorm, mlp_ratio=2): + super().__init__() + assert out_dim % in_dim == 0 or in_dim % out_dim == 0 + self.in_dim = in_dim + self.out_dim = out_dim + self.norm1 = norm_layer(in_dim) + self.attn = CausalAttention(in_dim, out_dim, num_heads) + self.proj = nn.Linear(in_dim, out_dim) + self.norm3 = norm_layer(in_dim) + + self.norm2 = norm_layer(out_dim) + hidden_dim = int(out_dim * mlp_ratio) + self.mlp = GeGluMlp(in_features=out_dim, hidden_features=hidden_dim) + # self.mlp = FeedForward(out_dim, out_dim) + + def forward(self, x): + x = self.proj(self.norm3(x)).add_(self.attn(self.norm1(x))) + return self.mlp(self.norm2(x)).add_(x) + + +def WNConv1d(*args, **kwargs): + return weight_norm(nn.Conv1d(*args, **kwargs)) + + +@torch.jit.script +def snake(x, alpha): + shape = x.shape + x = x.reshape(shape[0], shape[1], -1) + x = x + (alpha + 1e-9).reciprocal() * torch.sin(alpha * x).pow(2) + x = x.reshape(shape) + return x + + +class Snake1d(nn.Module): + def __init__(self, channels): + super().__init__() + self.alpha = nn.Parameter(torch.ones(1, channels, 1)) + + def forward(self, x): + return snake(x, self.alpha) + + +class ResidualUnit(nn.Module): + def __init__(self, dim: int = 16, dilation: int = 1): + super().__init__() + pad = ((7 - 1) * dilation) // 2 + self.block = nn.Sequential( + Snake1d(dim), + WNConv1d(dim, dim, kernel_size=7, dilation=dilation, padding=pad), + Snake1d(dim), + WNConv1d(dim, dim, kernel_size=1), + ) + + def forward(self, x): + y = self.block(x) + pad = (x.shape[-1] - y.shape[-1]) // 2 + if pad > 0: + x = x[..., pad:-pad] + return x + y + + +class EncoderBlock(nn.Module): + def __init__(self, dim: int = 16, stride: int = 1): + super().__init__() + self.block = nn.Sequential( + ResidualUnit(dim // 2, dilation=1), + ResidualUnit(dim // 2, dilation=3), + ResidualUnit(dim // 2, dilation=9), + Snake1d(dim // 2), + WNConv1d( + dim // 2, + dim, + kernel_size=2 * stride, + stride=stride, + padding=math.ceil(stride / 2), + ), + ) + + def forward(self, x): + return self.block(x) + + +class Encoder(nn.Module): + def __init__( + self, + d_model: int = 64, + strides: list = [2, 4, 8, 8], + d_latent: int = 64, + ): + super().__init__() + # Create first convolution + self.block = [WNConv1d(1, d_model, kernel_size=7, padding=3)] + + # Create EncoderBlocks that double channels as they downsample by `stride` + for stride in strides: + d_model *= 2 + self.block += [EncoderBlock(d_model, stride=stride)] + + # Create last convolution + self.block += [ + Snake1d(d_model), + WNConv1d(d_model, d_latent, kernel_size=3, padding=1), + ] + + # Wrap black into nn.Sequential + self.block = nn.Sequential(*self.block) + self.enc_dim = d_model + + def forward(self, x): + return self.block(x) + + +class DacAudioVAE(nn.Module): + def __init__( + self, + encoder_dim: int = 64, + encoder_rates: List[int] = [2, 4, 8, 8], + latent_dim: int = None, + decoder_dim: int = 1536, + decoder_rates: List[int] = [8, 8, 4, 2], + sample_rate: int = 44100, + vae_latent_channels: int = 64, + attn_proj: bool = False, + decoder_type: str = "bigvgan", + ): + super().__init__() + + self.encoder_dim = encoder_dim + self.encoder_rates = encoder_rates + self.decoder_dim = decoder_dim + self.decoder_rates = decoder_rates + self.sample_rate = sample_rate + self.attn_proj = attn_proj + self.decoder_type = decoder_type + + if latent_dim is None: + latent_dim = encoder_dim * (2 ** len(encoder_rates)) + + self.latent_dim = latent_dim + + self.hop_length = np.prod(encoder_rates) + self.encoder = Encoder(encoder_dim, encoder_rates, latent_dim) + + if latent_dim % vae_latent_channels == 0: + self.attn_proj_dim = vae_latent_channels + else: + # smallest power of two >= vae_latent_channels + self.attn_proj_dim = 2 ** int(np.ceil(np.log2(vae_latent_channels))) + + self.mean_proj = nn.Conv1d(self.attn_proj_dim, vae_latent_channels, 1) + self.logs_proj = nn.Conv1d(self.attn_proj_dim, vae_latent_channels, 1) + + self.dec_in_proj = nn.Conv1d(vae_latent_channels, latent_dim, 1) + + if self.decoder_type == "bigvgan": + if sample_rate == 16000: + bigvgan_conf = { + "resblock": "1", + "num_mels": latent_dim, + "upsample_rates": [5, 5, 2, 2, 2, 2], + "upsample_kernel_sizes": [9, 9, 4, 4, 4, 4], + "upsample_initial_channel": decoder_dim, + "resblock_kernel_sizes": [3, 7, 11], + "resblock_dilation_sizes": [[1, 3, 5], [1, 3, 5], [1, 3, 5]], + "use_tanh_at_final": False, + "use_bias_at_final": False, + "activation": "snakebeta", + "snake_logscale": True, + } + elif sample_rate == 32000: + bigvgan_conf = { + "resblock": "1", + "num_mels": latent_dim, + "upsample_rates": [5, 5, 2, 2, 2, 2, 2], + "upsample_kernel_sizes": [9, 9, 4, 4, 4, 4, 4], + "upsample_initial_channel": decoder_dim, + "resblock_kernel_sizes": [3, 7, 11], + "resblock_dilation_sizes": [[1, 3, 5], [1, 3, 5], [1, 3, 5]], + "use_tanh_at_final": False, + "use_bias_at_final": False, + "activation": "snakebeta", + "snake_logscale": True, + } + else: + raise ValueError(f"Invalid sample_rate: {sample_rate}") + + h = AttrDict(**bigvgan_conf) + self.decoder = BigVGAN(h) + else: + raise ValueError(f"Invalid decoder type: {self.decoder_type}") + + if self.attn_proj: + self.pre_block = AttnProjection(latent_dim, self.attn_proj_dim, num_heads=8) + + self.sample_rate = sample_rate + + def preprocess(self, audio_data, sample_rate): + if sample_rate is None: + sample_rate = self.sample_rate + + length = audio_data.shape[-1] + right_pad = math.ceil(length / self.hop_length) * self.hop_length - length + if right_pad: + audio_data = nn.functional.pad(audio_data, (0, right_pad)) + + return audio_data + + def decode(self, z: torch.Tensor): + """Decode given latent codes and return audio data + + Parameters + ---------- + z : Tensor[B x D x T] + Continuous latent representation + + Returns + ------- + Tensor[B x 1 x length] + Decoded audio data. + """ + z = self.dec_in_proj(z) + return self.decoder(z) diff --git a/telefuser/models/minimax_h3_audio/bigvgan.py b/telefuser/models/minimax_h3_audio/bigvgan.py new file mode 100644 index 0000000..fb55863 --- /dev/null +++ b/telefuser/models/minimax_h3_audio/bigvgan.py @@ -0,0 +1,239 @@ +# SPDX-License-Identifier: MIT +# Copyright (c) 2024 NVIDIA CORPORATION. +# Licensed under the MIT license. + +# Adapted from https://github.com/jik876/hifi-gan under the MIT license. + +import torch +import torch.nn as nn +from torch.nn import Conv1d, ConvTranspose1d, Parameter +from torch.nn.utils.parametrizations import weight_norm + +from .alias_free import Activation1d + + +def get_padding(kernel_size, dilation=1): + return int((kernel_size * dilation - dilation) / 2) + + +# Adapted from https://github.com/EdwardDixon/snake under the MIT license. +@torch.jit.script +def snakebeta(x, alpha, beta): + shape = x.shape + x = x.reshape(shape[0], shape[1], -1) + x = x + (beta + 1e-9).reciprocal() * torch.sin(alpha * x).pow(2) + x = x.reshape(shape) + return x + + +class SnakeBeta(nn.Module): + def __init__(self, in_features, alpha=1.0, alpha_trainable=True, alpha_logscale=False): + super(SnakeBeta, self).__init__() + self.in_features = in_features + + self.alpha_logscale = alpha_logscale + if self.alpha_logscale: + self.alpha = Parameter(torch.zeros(in_features) * alpha) + self.beta = Parameter(torch.zeros(in_features) * alpha) + else: + self.alpha = Parameter(torch.ones(in_features) * alpha) + self.beta = Parameter(torch.ones(in_features) * alpha) + + self.alpha.requires_grad = alpha_trainable + self.beta.requires_grad = alpha_trainable + self.no_div_by_zero = 0.000000001 + + def forward(self, x): + alpha = self.alpha.unsqueeze(0).unsqueeze(-1) + beta = self.beta.unsqueeze(0).unsqueeze(-1) + if self.alpha_logscale: + alpha = torch.exp(alpha) + beta = torch.exp(beta) + x = snakebeta(x, alpha, beta) + return x + + +class AttrDict(dict): + def __init__(self, *args, **kwargs): + super(AttrDict, self).__init__(*args, **kwargs) + self.__dict__ = self + + +class AMPBlock1(torch.nn.Module): + """ + AMPBlock applies trainable SnakeBeta periodic activations. + AMPBlock1 follows every dilated convolution with a fixed-dilation Conv1d. + + Args: + h (AttrDict): Hyperparameters. + channels (int): Number of convolution channels. + kernel_size (int): Size of the convolution kernel. Default is 3. + dilation (tuple): Dilation rates. Each layer has two convolutions. + activation (str): Activation function type. Must be 'snakebeta'. + """ + + def __init__( + self, + h: AttrDict, + channels: int, + kernel_size: int = 3, + dilation: tuple = (1, 3, 5), + activation: str = None, + ): + super().__init__() + + self.h = h + + self.convs1 = nn.ModuleList( + [ + weight_norm( + Conv1d( + channels, + channels, + kernel_size, + stride=1, + dilation=d, + padding=get_padding(kernel_size, d), + ) + ) + for d in dilation + ] + ) + + self.convs2 = nn.ModuleList( + [ + weight_norm( + Conv1d( + channels, + channels, + kernel_size, + stride=1, + dilation=1, + padding=get_padding(kernel_size, 1), + ) + ) + for _ in range(len(dilation)) + ] + ) + + self.num_layers = len(self.convs1) + len(self.convs2) # Total number of conv layers + + if activation == "snakebeta": + self.activations = nn.ModuleList( + [ + Activation1d(activation=SnakeBeta(channels, alpha_logscale=h.snake_logscale)) + for _ in range(self.num_layers) + ] + ) + else: + raise NotImplementedError( + "activation incorrectly specified. check the config file and look for 'activation'." + ) + + def forward(self, x): + activation_iter = iter(self.activations) + for c1, c2 in zip(self.convs1, self.convs2): + a1 = next(activation_iter) + a2 = next(activation_iter) + xt = a1(x) + xt = c1(xt) + xt = a2(xt) + xt = c2(xt) + x = xt.add_(x) + + return x + + +class BigVGAN(torch.nn.Module): + """ + BigVGAN is a neural vocoder model that applies anti-aliased periodic activation for residual blocks (resblocks). + + Args: + h (AttrDict): Hyperparameters. + """ + + def __init__(self, h: AttrDict): + super().__init__() + self.h = h + + self.num_kernels = len(h.resblock_kernel_sizes) + self.num_upsamples = len(h.upsample_rates) + + # Pre-conv + self.conv_pre = weight_norm(Conv1d(h.num_mels, h.upsample_initial_channel, 7, 1, padding=3)) + + # Define which AMPBlock to use. BigVGAN uses AMPBlock1 as default + if h.resblock == "1": + resblock_class = AMPBlock1 + else: + raise ValueError(f"Incorrect resblock class specified in hyperparameters. Got {h.resblock}") + + # Transposed conv-based upsamplers. does not apply anti-aliasing + self.ups = nn.ModuleList() + for i, (u, k) in enumerate(zip(h.upsample_rates, h.upsample_kernel_sizes)): + self.ups.append( + nn.ModuleList( + [ + weight_norm( + ConvTranspose1d( + h.upsample_initial_channel // (2**i), + h.upsample_initial_channel // (2 ** (i + 1)), + k, + u, + padding=(k - u) // 2, + ) + ) + ] + ) + ) + + # Residual blocks using anti-aliased multi-periodicity composition modules (AMP) + self.resblocks = nn.ModuleList() + for i in range(len(self.ups)): + ch = h.upsample_initial_channel // (2 ** (i + 1)) + for j, (k, d) in enumerate(zip(h.resblock_kernel_sizes, h.resblock_dilation_sizes)): + self.resblocks.append(resblock_class(h, ch, k, d, activation=h.activation)) + + # Post-conv + if h.activation != "snakebeta": + raise NotImplementedError( + "activation incorrectly specified. check the config file and look for 'activation'." + ) + activation_post = SnakeBeta(ch, alpha_logscale=h.snake_logscale) + + self.activation_post = Activation1d(activation=activation_post) + + # Whether to use bias for the final conv_post. Default to True for backward compatibility + self.use_bias_at_final = h.get("use_bias_at_final", True) + self.conv_post = weight_norm(Conv1d(ch, 1, 7, 1, padding=3, bias=self.use_bias_at_final)) + + # Final tanh activation. Defaults to True for backward compatibility + self.use_tanh_at_final = h.get("use_tanh_at_final", True) + + def forward(self, x): + # Pre-conv + x = self.conv_pre(x) + + for i in range(self.num_upsamples): + # Upsampling + for i_up in range(len(self.ups[i])): + x = self.ups[i][i_up](x) + # AMP blocks + xs = None + for j in range(self.num_kernels): + if xs is None: + xs = self.resblocks[i * self.num_kernels + j](x) + else: + xs += self.resblocks[i * self.num_kernels + j](x) + x = xs.div_(self.num_kernels) + + # Post-conv + x = self.activation_post(x) + x = self.conv_post(x) + # Final tanh activation + if self.use_tanh_at_final: + x.tanh_() + else: + x.clamp_(min=-1.0, max=1.0) # Bound the output to [-1, 1] + + return x diff --git a/telefuser/models/minimax_h3_audio_vae.py b/telefuser/models/minimax_h3_audio_vae.py new file mode 100644 index 0000000..b5692c4 --- /dev/null +++ b/telefuser/models/minimax_h3_audio_vae.py @@ -0,0 +1,147 @@ +# SPDX-License-Identifier: Apache-2.0 +"""MiniMax H3 32 kHz stereo audio VAE.""" + +from __future__ import annotations + +import json +from dataclasses import dataclass +from pathlib import Path +from typing import Any + +import torch + +from telefuser.core.base_model import BaseModel + +from .minimax_h3_audio import DacAudioVAE + + +@dataclass(frozen=True) +class MiniMaxH3AudioVAEConfig: + encoder_dim: int + encoder_rates: tuple[int, ...] + latent_dim: int + decoder_dim: int + decoder_rates: tuple[int, ...] + sample_rate: int + latent_channels: int + output_channels: int + attn_proj: bool + decoder_type: str + latents_mean: tuple[float, ...] + latents_std: tuple[float, ...] + + @classmethod + def from_path(cls, path: str | Path) -> MiniMaxH3AudioVAEConfig: + component_dir = Path(path) + if component_dir.is_file(): + component_dir = component_dir.parent + component = json.loads((component_dir / "config.json").read_text(encoding="utf-8")) + metadata_path = component_dir / component["source_metadata_path"] + metadata = json.loads(metadata_path.read_text(encoding="utf-8"))["metadata"]["kwargs"] + config = cls( + encoder_dim=int(metadata["encoder_dim"]), + encoder_rates=tuple(int(rate) for rate in metadata["encoder_rates"]), + latent_dim=int(metadata["latent_dim"]), + decoder_dim=int(metadata["decoder_dim"]), + decoder_rates=tuple(int(rate) for rate in metadata["decoder_rates"]), + sample_rate=int(metadata["sample_rate"]), + latent_channels=int(metadata["vae_latent_channels"]), + output_channels=int(component["output_channel"]), + attn_proj=bool(metadata["attn_proj"]), + decoder_type=str(metadata["decoder_type"]), + latents_mean=tuple(float(value) for value in component["latents_mean"]), + latents_std=tuple(float(value) for value in component["latents_std"]), + ) + config.validate() + return config + + def validate(self) -> None: + if self.sample_rate != 32_000: + raise ValueError(f"MiniMax H3 audio VAE requires 32000 Hz, got {self.sample_rate}") + if self.latent_channels != 32: + raise ValueError(f"MiniMax H3 audio VAE requires 32 latent channels, got {self.latent_channels}") + if self.output_channels != 2: + raise ValueError(f"MiniMax H3 output requires stereo audio, got {self.output_channels} channels") + if len(self.latents_mean) != self.latent_channels or len(self.latents_std) != self.latent_channels: + raise ValueError("audio VAE latent statistics must contain one value per latent channel") + if any(value <= 0 for value in self.latents_std): + raise ValueError("audio VAE latent standard deviations must be positive") + + +class MiniMaxH3AudioVAE(DacAudioVAE, BaseModel): + """DAC-lineage waveform encoder and BigVGAN decoder used by MiniMax H3.""" + + def __init__(self, config: MiniMaxH3AudioVAEConfig) -> None: + config.validate() + super().__init__( + encoder_dim=config.encoder_dim, + encoder_rates=list(config.encoder_rates), + latent_dim=config.latent_dim, + decoder_dim=config.decoder_dim, + decoder_rates=list(config.decoder_rates), + sample_rate=config.sample_rate, + vae_latent_channels=config.latent_channels, + attn_proj=config.attn_proj, + decoder_type=config.decoder_type, + ) + self.config = config + self.layer_name_list = ["encoder", "decoder"] + + @torch.no_grad() + def encode_mean(self, waveform: torch.Tensor, sample_rate: int | None = None) -> torch.Tensor: + """Encode independent waveform channels to deterministic normalized latents.""" + if waveform.ndim != 3 or waveform.shape[1] != 1: + raise ValueError("waveform must be [channels, 1, samples]") + waveform = self.preprocess(waveform, sample_rate) + parameter = next(self.parameters()) + waveform = waveform.to(device=parameter.device, dtype=parameter.dtype) + hidden = self.encoder(waveform) + if self.attn_proj: + hidden = self.pre_block(hidden.transpose(1, 2)).transpose(1, 2) + latent = self.mean_proj(hidden) + mean = latent.new_tensor(self.config.latents_mean).view(1, -1, 1) + std = latent.new_tensor(self.config.latents_std).view(1, -1, 1) + return latent.sub(mean).div(std) + + @torch.no_grad() + def decode_normalized(self, latent: torch.Tensor) -> torch.Tensor: + """Decode normalized [2, 32, T] latents to [1, 2, samples] stereo.""" + if latent.ndim != 3 or latent.shape[0] != self.config.output_channels: + raise ValueError("audio latent must be [2, 32, T]") + if latent.shape[1] != self.config.latent_channels: + raise ValueError(f"audio latent channel dimension must be {self.config.latent_channels}") + parameter = next(self.parameters()) + latent = latent.to(device=parameter.device, dtype=parameter.dtype) + mean = latent.new_tensor(self.config.latents_mean).view(1, -1, 1) + std = latent.new_tensor(self.config.latents_std).view(1, -1, 1) + waveform = self.decode(latent.mul(std).add(mean)) + return waveform.transpose(0, 1).contiguous() + + @staticmethod + def state_dict_converter(config_path: str | Path) -> MiniMaxH3AudioVAEStateDictConverter: + return MiniMaxH3AudioVAEStateDictConverter(config_path) + + +class MiniMaxH3AudioVAEStateDictConverter: + def __init__(self, config_path: str | Path) -> None: + self.config = MiniMaxH3AudioVAEConfig.from_path(config_path) + + def from_official(self, state_dict: dict[str, torch.Tensor]) -> tuple[dict[str, torch.Tensor], dict[str, Any]]: + converted: dict[str, torch.Tensor] = {} + for name, value in state_dict.items(): + if name.endswith(".weight_g"): + name = name.removesuffix(".weight_g") + ".parametrizations.weight.original0" + elif name.endswith(".weight_v"): + name = name.removesuffix(".weight_v") + ".parametrizations.weight.original1" + converted[name] = value + return converted, {"config": self.config} + + def from_diffusers(self, state_dict: dict[str, torch.Tensor]) -> tuple[dict[str, torch.Tensor], dict[str, Any]]: + return self.from_official(state_dict) + + +__all__ = [ + "MiniMaxH3AudioVAE", + "MiniMaxH3AudioVAEConfig", + "MiniMaxH3AudioVAEStateDictConverter", +] diff --git a/telefuser/models/minimax_h3_dit.py b/telefuser/models/minimax_h3_dit.py new file mode 100644 index 0000000..cfd5bf7 --- /dev/null +++ b/telefuser/models/minimax_h3_dit.py @@ -0,0 +1,671 @@ +# SPDX-License-Identifier: Apache-2.0 +"""MiniMax H3 packed multimodal DiT.""" + +from __future__ import annotations + +import json +import math +import re +from dataclasses import dataclass +from pathlib import Path +from typing import Any + +import torch +import torch.distributed as dist +import torch.nn as nn + +from telefuser.core.base_model import BaseModel +from telefuser.core.config import AttentionConfig +from telefuser.distributed.device_mesh import get_ulysses_group, get_ulysses_world_size +from telefuser.distributed.parallel_shard import sequence_parallel_shard, sequence_parallel_unshard +from telefuser.distributed.ulysses_comm import ulysses_gather_heads, ulysses_scatter_heads +from telefuser.ops.attention import attention + +MINIMAX_H3_ADALN_MODALITY_NUM = 3 +MINIMAX_H3_FP32_PARAM_NAMES = frozenset( + { + "video_patch_proj.weight", + "video_patch_proj.bias", + "audio_patch_proj.weight", + "audio_patch_proj.bias", + "time_embedder.proj_in.weight", + "time_embedder.proj_in.bias", + "time_embedder.proj_out.weight", + "time_embedder.proj_out.bias", + "final_layer.video_out.weight", + "final_layer.video_out.bias", + "final_layer.audio_out.weight", + "final_layer.audio_out.bias", + } +) +MINIMAX_H3_FP32_BUFFER_NAMES = frozenset({"rope.inv_freq"}) + + +@dataclass(frozen=True) +class MiniMaxH3DiTConfig: + hidden_size: int = 5376 + num_layers: int = 50 + token_refiner_num_layers: int = 2 + num_attention_heads: int = 56 + attention_head_dim: int = 128 + ffn_hidden_size: int = 14336 + latents_dim: int = 24 + audio_latents_dim: int = 32 + patch_size: tuple[int, int, int] = (1, 2, 2) + text_dim: int = 5120 + timestep_input_dim: int = 256 + time_embed_hidden_size: int = 5376 + time_embed_dim: int = 2688 + rope_inv_freq_len: int = 16 + norm_eps: float = 1e-5 + qk_norm_eps: float = 1e-5 + final_norm_eps: float = 1e-5 + + @classmethod + def from_json(cls, path: str | Path) -> MiniMaxH3DiTConfig: + payload = json.loads(Path(path).read_text(encoding="utf-8")) + fields = cls.__dataclass_fields__ + values = {key: payload[key] for key in fields if key in payload} + if "patch_size" in values: + values["patch_size"] = tuple(int(value) for value in values["patch_size"]) + return cls(**values) + + def __post_init__(self) -> None: + if self.hidden_size <= 0 or self.num_layers <= 0: + raise ValueError("MiniMax H3 hidden_size and num_layers must be positive") + if self.num_attention_heads <= 0 or self.attention_head_dim <= 0: + raise ValueError("MiniMax H3 attention dimensions must be positive") + if len(self.patch_size) != 3 or any(value <= 0 for value in self.patch_size): + raise ValueError("MiniMax H3 patch_size must contain three positive integers") + if 6 * self.rope_inv_freq_len > self.attention_head_dim: + raise ValueError("MiniMax H3 rotary dimensions must fit inside attention_head_dim") + + @property + def inner_dim(self) -> int: + return self.num_attention_heads * self.attention_head_dim + + @property + def video_patch_dim(self) -> int: + return self.latents_dim * math.prod(self.patch_size) + + @property + def adaln_out_features(self) -> int: + return 18 * self.hidden_size + + @property + def final_adaln_out_features(self) -> int: + return 2 * self.hidden_size + + +def _rms_norm(size: int, eps: float) -> nn.RMSNorm: + return nn.RMSNorm(size, eps=eps, dtype=torch.bfloat16) + + +def _reorder_grouped_qkv_to_qkv( + weight: torch.Tensor, + *, + num_query_groups: int, + heads_per_group: int, + head_dim: int, +) -> torch.Tensor: + per_group = (heads_per_group + 2) * head_dim + if weight.shape[0] != num_query_groups * per_group: + raise ValueError("MiniMax H3 grouped QKV weight has an incompatible output dimension") + rest = weight.shape[1:] + grouped = weight.reshape(num_query_groups, per_group, *rest) + q, k, v = torch.split(grouped, [heads_per_group * head_dim, head_dim, head_dim], dim=1) + return torch.cat( + ( + q.reshape(num_query_groups * heads_per_group * head_dim, *rest), + k.reshape(num_query_groups * head_dim, *rest), + v.reshape(num_query_groups * head_dim, *rest), + ), + dim=0, + ) + + +def _rotate_half(value: torch.Tensor) -> torch.Tensor: + first, second = value.chunk(2, dim=-1) + return torch.cat((-second, first), dim=-1) + + +class MiniMaxH3Rope(nn.Module): + def __init__(self, inv_freq_len: int) -> None: + super().__init__() + inv_freq = 10000.0 ** (-torch.arange(inv_freq_len, dtype=torch.float32) / inv_freq_len) + self.register_buffer("inv_freq", inv_freq, persistent=True) + + def forward(self, position_ids: torch.Tensor) -> torch.Tensor: + if position_ids.ndim != 3 or position_ids.shape[0] != 1 or position_ids.shape[-1] != 3: + raise ValueError("MiniMax H3 position_ids must have shape [1, sequence, 3]") + per_axis = position_ids[0].float().unsqueeze(-1) * self.inv_freq.view(1, 1, -1) + half = torch.cat(tuple(per_axis.unbind(dim=1)), dim=-1) + return torch.cat((half, half), dim=-1) + + +class MiniMaxH3TimeEmbedder(nn.Module): + def __init__(self, config: MiniMaxH3DiTConfig) -> None: + super().__init__() + self.frequency_embedding_size = config.timestep_input_dim + self.proj_in = nn.Linear( + config.timestep_input_dim, + config.time_embed_hidden_size, + dtype=torch.float32, + ) + self.proj_out = nn.Linear( + config.time_embed_hidden_size, + config.time_embed_dim, + dtype=torch.float32, + ) + + def forward(self, timestep: torch.Tensor) -> torch.Tensor: + half = self.frequency_embedding_size // 2 + frequencies = torch.exp( + -math.log(10000.0) * torch.arange(half, dtype=torch.float32, device=timestep.device) / half + ) + args = timestep.float().reshape(-1, 1) * frequencies.reshape(1, -1) + embedding = torch.cat((torch.cos(args), torch.sin(args)), dim=-1) + return self.proj_out(nn.functional.silu(self.proj_in(embedding))) + + +class MiniMaxH3Attention(nn.Module): + def __init__(self, config: MiniMaxH3DiTConfig) -> None: + super().__init__() + self.num_heads = config.num_attention_heads + self.head_dim = config.attention_head_dim + self.inner_dim = config.inner_dim + self.qkv_proj = nn.Linear(config.hidden_size, 3 * self.inner_dim, bias=False, dtype=torch.bfloat16) + self.q_norm = _rms_norm(self.head_dim, config.qk_norm_eps) + self.k_norm = _rms_norm(self.head_dim, config.qk_norm_eps) + self.out_proj = nn.Linear(self.inner_dim, config.hidden_size, bias=False, dtype=torch.bfloat16) + self.ulysses_group: dist.ProcessGroup | None = None + self._communication_events: list[tuple[torch.cuda.Event, torch.cuda.Event]] = [] + + def set_ulysses_group(self, group: dist.ProcessGroup | None) -> None: + self.ulysses_group = group + + def reset_communication_metrics(self) -> None: + self._communication_events.clear() + + def communication_seconds(self) -> float: + return sum(start.elapsed_time(end) for start, end in self._communication_events) / 1000.0 + + def forward( + self, + hidden: torch.Tensor, + *, + sequence_lengths: list[int], + rope_frequencies: torch.Tensor | None, + attention_config: AttentionConfig | None, + ) -> torch.Tensor: + sequence, _ = hidden.shape + qkv = self.qkv_proj(hidden).reshape(sequence, 3, self.num_heads, self.head_dim) + query, key, value = qkv.unbind(dim=1) + query = self.q_norm(query) + key = self.k_norm(key) + if rope_frequencies is not None: + rotary_dim = rope_frequencies.shape[-1] + cosine = rope_frequencies.cos().unsqueeze(1).to(query.dtype) + sine = rope_frequencies.sin().unsqueeze(1).to(query.dtype) + query_rotary, query_pass = query[..., :rotary_dim], query[..., rotary_dim:] + key_rotary, key_pass = key[..., :rotary_dim], key[..., rotary_dim:] + query = torch.cat((query_rotary * cosine + _rotate_half(query_rotary) * sine, query_pass), dim=-1) + key = torch.cat((key_rotary * cosine + _rotate_half(key_rotary) * sine, key_pass), dim=-1) + query = query.unsqueeze(0) + key = key.unsqueeze(0) + value = value.unsqueeze(0) + group = self.ulysses_group + use_ulysses = group is not None and dist.get_world_size(group) > 1 + if use_ulysses: + scatter_start = torch.cuda.Event(enable_timing=True) + scatter_end = torch.cuda.Event(enable_timing=True) + scatter_start.record() + qkv_wait = ulysses_scatter_heads(torch.cat((query, key, value), dim=-1), group) + query, key, value = qkv_wait().chunk(3, dim=-1) + scatter_end.record() + self._communication_events.append((scatter_start, scatter_end)) + output = attention( + query, + key, + value, + attention_config=attention_config, + scale=self.head_dim**-0.5, + sequence_lengths=sequence_lengths, + ) + if use_ulysses: + gather_start = torch.cuda.Event(enable_timing=True) + gather_end = torch.cuda.Event(enable_timing=True) + gather_start.record() + output = ulysses_gather_heads(output, group, num_heads=self.num_heads)() + gather_end.record() + self._communication_events.append((gather_start, gather_end)) + return self.out_proj(output[0].reshape(sequence, self.inner_dim)) + + +class MiniMaxH3MLP(nn.Module): + def __init__(self, config: MiniMaxH3DiTConfig) -> None: + super().__init__() + self.fc1 = nn.Linear(config.hidden_size, 2 * config.ffn_hidden_size, bias=False, dtype=torch.bfloat16) + self.fc2 = nn.Linear(config.ffn_hidden_size, config.hidden_size, bias=False, dtype=torch.bfloat16) + + def forward(self, hidden: torch.Tensor) -> torch.Tensor: + gate, up = self.fc1(hidden).chunk(2, dim=-1) + return self.fc2(nn.functional.silu(gate) * up) + + +class MiniMaxH3AdaLNProjection(nn.Module): + def __init__(self, config: MiniMaxH3DiTConfig, *, expand_ratio: int, modality_count: int) -> None: + super().__init__() + self.expand_ratio = expand_ratio + self.modality_count = modality_count + self.hidden_size = config.hidden_size + self.linear = nn.Linear( + config.time_embed_dim, + expand_ratio * modality_count * config.hidden_size, + dtype=torch.bfloat16, + ) + + def forward(self, embedding: torch.Tensor) -> tuple[torch.Tensor, ...]: + output = self.linear(embedding) + output = output.reshape(-1, self.expand_ratio * self.hidden_size) + return tuple(output.chunk(self.expand_ratio, dim=-1)) + + +def _modulate( + hidden: torch.Tensor, + shift: torch.Tensor, + scale: torch.Tensor, + indices: torch.Tensor, +) -> torch.Tensor: + return hidden * (1 + scale.index_select(0, indices)) + shift.index_select(0, indices) + + +class MiniMaxH3TokenRefinerBlock(nn.Module): + def __init__(self, config: MiniMaxH3DiTConfig) -> None: + super().__init__() + self.norm1 = _rms_norm(config.hidden_size, config.norm_eps) + self.norm2 = _rms_norm(config.hidden_size, config.norm_eps) + self.attn = MiniMaxH3Attention(config) + self.mlp = MiniMaxH3MLP(config) + + def forward( + self, + hidden: torch.Tensor, + *, + sequence_lengths: list[int], + attention_config: AttentionConfig | None, + ) -> torch.Tensor: + hidden = hidden + self.attn( + self.norm1(hidden), + sequence_lengths=sequence_lengths, + rope_frequencies=None, + attention_config=attention_config, + ) + return hidden + self.mlp(self.norm2(hidden)) + + +class MiniMaxH3TokenRefiner(nn.Module): + def __init__(self, config: MiniMaxH3DiTConfig) -> None: + super().__init__() + self.blocks = nn.ModuleList( + [MiniMaxH3TokenRefinerBlock(config) for _ in range(config.token_refiner_num_layers)] + ) + self.final_norm = _rms_norm(config.hidden_size, config.final_norm_eps) + + def forward( + self, + hidden: torch.Tensor, + *, + attention_config: AttentionConfig | None, + ) -> torch.Tensor: + for block in self.blocks: + hidden = block( + hidden, + sequence_lengths=[hidden.shape[0]], + attention_config=attention_config, + ) + return self.final_norm(hidden) + + +class MiniMaxH3DiTBlock(nn.Module): + def __init__(self, config: MiniMaxH3DiTConfig) -> None: + super().__init__() + self.norm1 = _rms_norm(config.hidden_size, config.norm_eps) + self.norm2 = _rms_norm(config.hidden_size, config.norm_eps) + self.attn = MiniMaxH3Attention(config) + self.mlp = MiniMaxH3MLP(config) + self.adaln_proj = MiniMaxH3AdaLNProjection( + config, + expand_ratio=6, + modality_count=MINIMAX_H3_ADALN_MODALITY_NUM, + ) + + def forward( + self, + hidden: torch.Tensor, + *, + adaln_input: torch.Tensor, + combined_indices: torch.Tensor, + sequence_lengths: list[int], + rope_frequencies: torch.Tensor, + attention_config: AttentionConfig | None, + ) -> torch.Tensor: + shift_msa, scale_msa, gate_msa, shift_mlp, scale_mlp, gate_mlp = self.adaln_proj(adaln_input) + residual = hidden + value = _modulate(self.norm1(hidden), shift_msa, scale_msa, combined_indices) + value = self.attn( + value, + sequence_lengths=sequence_lengths, + rope_frequencies=rope_frequencies, + attention_config=attention_config, + ) + hidden = residual + gate_msa.index_select(0, combined_indices) * value + residual = hidden + value = _modulate(self.norm2(hidden), shift_mlp, scale_mlp, combined_indices) + value = self.mlp(value) + return residual + gate_mlp.index_select(0, combined_indices) * value + + +class MiniMaxH3FinalLayer(nn.Module): + def __init__(self, config: MiniMaxH3DiTConfig) -> None: + super().__init__() + self.norm = _rms_norm(config.hidden_size, config.final_norm_eps) + self.adaln_proj = MiniMaxH3AdaLNProjection(config, expand_ratio=2, modality_count=1) + self.video_out = nn.Linear(config.hidden_size, config.video_patch_dim, dtype=torch.float32) + self.audio_out = nn.Linear(config.hidden_size, config.audio_latents_dim, dtype=torch.float32) + + def forward( + self, + hidden: torch.Tensor, + *, + adaln_input: torch.Tensor, + inverse_indices: torch.Tensor, + ) -> tuple[torch.Tensor, torch.Tensor]: + shift, scale = self.adaln_proj(adaln_input) + hidden = _modulate(self.norm(hidden), shift, scale, inverse_indices).float() + return self.video_out(hidden), self.audio_out(hidden) + + +class MiniMaxH3DiT(BaseModel): + """Faithful packed DiT baseline for H3-Base with optional Ulysses SP.""" + + def __init__(self, config: MiniMaxH3DiTConfig | None = None) -> None: + super().__init__() + self.config = config or MiniMaxH3DiTConfig() + config = self.config + self.video_patch_proj = nn.Linear(config.video_patch_dim, config.hidden_size, dtype=torch.float32) + self.audio_patch_proj = nn.Linear(config.audio_latents_dim, config.hidden_size, dtype=torch.float32) + self.condition_proj = nn.Linear(config.text_dim, config.hidden_size, dtype=torch.bfloat16) + self.time_embedder = MiniMaxH3TimeEmbedder(config) + self.rope = MiniMaxH3Rope(config.rope_inv_freq_len) + self.token_refiner = MiniMaxH3TokenRefiner(config) + self.blocks = nn.ModuleList([MiniMaxH3DiTBlock(config) for _ in range(config.num_layers)]) + self.final_layer = MiniMaxH3FinalLayer(config) + self.layer_name_list = ["blocks"] + self.device_mesh: Any | None = None + self.usp_flag = False + + def _preserve_fp32_boundaries(self) -> None: + for name in MINIMAX_H3_FP32_PARAM_NAMES: + parameter = self.get_parameter(name) + if parameter.dtype != torch.float32: + parameter.data = parameter.data.float() + if self.rope.inv_freq.dtype != torch.float32: + self.rope.inv_freq.data = self.rope.inv_freq.data.float() + + def to(self, *args: Any, **kwargs: Any) -> MiniMaxH3DiT: + preserved_parameters = { + name: parameter.detach().clone() + for name in MINIMAX_H3_FP32_PARAM_NAMES + if not (parameter := self.get_parameter(name)).is_meta + } + preserved_buffers = { + name: buffer.detach().clone() + for name in MINIMAX_H3_FP32_BUFFER_NAMES + if not (buffer := self.get_buffer(name)).is_meta + } + result = super().to(*args, **kwargs) + for name, value in preserved_parameters.items(): + parameter = result.get_parameter(name) + parameter.data = value.to(device=parameter.device, dtype=torch.float32) + for name, value in preserved_buffers.items(): + buffer = result.get_buffer(name) + buffer.data = value.to(device=buffer.device, dtype=torch.float32) + result._preserve_fp32_boundaries() + return result + + @staticmethod + def _position_ids(value: Any, name: str) -> torch.Tensor: + position_ids = value.get("position_ids") if isinstance(value, dict) else getattr(value, "position_ids", None) + if position_ids is None: + raise ValueError(f"{name}.position_ids is required") + return position_ids.reshape(-1).long() + + @staticmethod + def _sequence_lengths(packed: Any) -> list[int]: + cu = packed.get("cu_seqlens_q") if isinstance(packed, dict) else packed.cu_seqlens_q + values = [int(value) for value in cu.tolist()] + return [stop - start for start, stop in zip(values[:-1], values[1:], strict=True) if stop > start] + + def forward(self, **kwargs: Any) -> tuple[torch.Tensor, torch.Tensor]: + required = ( + "x", + "audio_x", + "img_position_ids", + "unique_timesteps", + "inverse_indices", + "update_mask", + "prompt_embeds", + "img_pos_info", + "audio_pos_info", + "text_pos_info", + "img_pos_for_infer_output_info", + "packed_seq_params", + ) + missing = [name for name in required if kwargs.get(name) is None] + if missing: + raise ValueError(f"MiniMaxH3DiT.forward missing required inputs: {missing}") + video_state = kwargs["x"] + audio_state = kwargs["audio_x"] + if video_state.ndim != 3 or video_state.shape[0] != 1: + raise ValueError("x must have shape [1, sequence, video_patch_dim]") + sequence = video_state.shape[1] + device = video_state.device + image_positions = self._position_ids(kwargs["img_pos_info"], "img_pos_info").to(device) + audio_positions = self._position_ids(kwargs["audio_pos_info"], "audio_pos_info").to(device) + text_positions = self._position_ids(kwargs["text_pos_info"], "text_pos_info").to(device) + output_positions = self._position_ids( + kwargs["img_pos_for_infer_output_info"], "img_pos_for_infer_output_info" + ).to(device) + + prompt = kwargs["prompt_embeds"].to(device=device, dtype=torch.bfloat16) + live_text = text_positions.numel() + prompt = self.condition_proj(prompt[:live_text]) + prompt = self.token_refiner(prompt, attention_config=self.attention_config) + hidden = torch.zeros(sequence, self.config.hidden_size, device=device, dtype=torch.bfloat16) + hidden.index_copy_(0, text_positions, prompt) + video_rows = video_state[0].index_select(0, image_positions).float() + audio_rows = audio_state[0].index_select(0, audio_positions).float() + hidden.index_copy_(0, image_positions, self.video_patch_proj(video_rows).to(torch.bfloat16)) + hidden.index_copy_(0, audio_positions, self.audio_patch_proj(audio_rows).to(torch.bfloat16)) + + timesteps = kwargs["unique_timesteps"].reshape(-1).to(device) + adaln_input = nn.functional.silu(self.time_embedder(timesteps)).to(torch.bfloat16) + inverse_indices = kwargs["inverse_indices"].reshape(-1).long().to(device) + if inverse_indices.numel() != sequence: + raise ValueError("inverse_indices must cover the full packed sequence") + token_tags = kwargs.get("block_token_tags") + if token_tags is None: + token_tags = kwargs.get("token_tags") + if token_tags is None: + raise ValueError("token_tags or block_token_tags is required") + token_tags = token_tags.reshape(-1).long().to(device).clamp_min(0) + combined_indices = token_tags + inverse_indices * MINIMAX_H3_ADALN_MODALITY_NUM + rope_frequencies = self.rope(kwargs["img_position_ids"].to(device)) + sequence_lengths = self._sequence_lengths(kwargs["packed_seq_params"]) + full_sequence = sequence + if self.usp_flag: + world_size = get_ulysses_world_size(self.device_mesh) + if sequence % world_size: + raise ValueError( + f"MiniMax H3 packed sequence length ({sequence}) must be divisible by Ulysses degree ({world_size})" + ) + inverse_indices = inverse_indices.clone() + sequence_parallel_shard( + self.device_mesh, + [hidden, combined_indices, inverse_indices, rope_frequencies], + [0, 0, 0, 0], + ) + for block in self.blocks: + hidden = block( + hidden, + adaln_input=adaln_input, + combined_indices=combined_indices, + sequence_lengths=sequence_lengths, + rope_frequencies=rope_frequencies, + attention_config=self.attention_config, + ) + video_logits, audio_logits = self.final_layer( + hidden, + adaln_input=adaln_input, + inverse_indices=inverse_indices, + ) + if self.usp_flag: + video_logits, audio_logits = sequence_parallel_unshard( + self.device_mesh, + [video_logits, audio_logits], + [0, 0], + [full_sequence, full_sequence], + ) + video_logits = video_logits.index_select(0, output_positions) + audio_logits = audio_logits.index_select(0, audio_positions) + if not bool(kwargs.get("skip_mask_out_condition", False)): + video_logits = video_logits * kwargs["update_mask"].reshape(-1, 1).to(video_logits) + if kwargs.get("update_audio_mask") is not None: + audio_logits = audio_logits * kwargs["update_audio_mask"].reshape(-1, 1).to(audio_logits) + return video_logits, audio_logits + + def enable_usp(self, device_mesh: Any | None = None) -> None: + self.device_mesh = device_mesh if device_mesh is not None else self.device_mesh + world_size = get_ulysses_world_size(self.device_mesh) + if self.config.num_attention_heads % world_size: + raise ValueError( + f"MiniMax H3 attention heads ({self.config.num_attention_heads}) must be divisible by " + f"Ulysses degree ({world_size})" + ) + group = get_ulysses_group(self.device_mesh) if world_size > 1 else None + self.usp_flag = world_size > 1 + for block in self.blocks: + block.attn.set_ulysses_group(group) + + def reset_communication_metrics(self) -> None: + for block in self.blocks: + block.attn.reset_communication_metrics() + + def communication_seconds(self) -> float: + return sum(block.attn.communication_seconds() for block in self.blocks) + + def get_fsdp_module_names(self) -> list[str]: + return ["blocks"] + + @staticmethod + def state_dict_converter(config_path: str | Path | None = None) -> MiniMaxH3DiTStateDictConverter: + return MiniMaxH3DiTStateDictConverter(config_path=config_path) + + +_BLOCK_INDEX = re.compile(r"^blocks\.(\d+)\.") +_REFINER_INDEX = re.compile(r"^token_refiner\.blocks\.(\d+)\.") + + +class MiniMaxH3DiTStateDictConverter: + def __init__(self, config_path: str | Path | None = None) -> None: + self.config_path = None if config_path is None else Path(config_path) + + def _config(self, state_dict: dict[str, torch.Tensor]) -> MiniMaxH3DiTConfig: + if self.config_path is not None: + return MiniMaxH3DiTConfig.from_json(self.config_path) + q_norm = state_dict["blocks.0.attn.q_norm.weight"] + qkv = state_dict["blocks.0.attn.qkv_proj.weight"] + layers = 1 + max(int(match.group(1)) for key in state_dict if (match := _BLOCK_INDEX.match(key))) + refiners = 1 + max(int(match.group(1)) for key in state_dict if (match := _REFINER_INDEX.match(key))) + video_patch_dim = state_dict["video_patch_proj.weight"].shape[1] + return MiniMaxH3DiTConfig( + hidden_size=state_dict["video_patch_proj.weight"].shape[0], + num_layers=layers, + token_refiner_num_layers=refiners, + num_attention_heads=qkv.shape[0] // (3 * q_norm.numel()), + attention_head_dim=q_norm.numel(), + ffn_hidden_size=state_dict["blocks.0.mlp.fc1.weight"].shape[0] // 2, + latents_dim=video_patch_dim // 4, + audio_latents_dim=state_dict["audio_patch_proj.weight"].shape[1], + text_dim=state_dict["condition_proj.weight"].shape[1], + timestep_input_dim=state_dict["time_embedder.proj_in.weight"].shape[1], + time_embed_hidden_size=state_dict["time_embedder.proj_in.weight"].shape[0], + time_embed_dim=state_dict["time_embedder.proj_out.weight"].shape[0], + rope_inv_freq_len=state_dict["rope.inv_freq"].numel(), + ) + + def from_official(self, state_dict: dict[str, torch.Tensor]) -> tuple[dict[str, torch.Tensor], dict[str, Any]]: + config = self._config(state_dict) + converted = dict(state_dict) + for key, value in state_dict.items(): + if key.endswith(".attn.qkv_proj.weight"): + converted[key] = _reorder_grouped_qkv_to_qkv( + value, + num_query_groups=config.num_attention_heads, + heads_per_group=1, + head_dim=config.attention_head_dim, + ) + return converted, {"config": config} + + def from_diffusers(self, state_dict: dict[str, torch.Tensor]) -> tuple[dict[str, torch.Tensor], dict[str, Any]]: + renamed: dict[str, torch.Tensor] = {} + qkv_parts: dict[str, dict[str, torch.Tensor]] = {} + direct = { + "proj_in.": "video_patch_proj.", + "audio_proj_in.": "audio_patch_proj.", + "context_embedder.": "condition_proj.", + "time_embedder.linear_1.": "time_embedder.proj_in.", + "time_embedder.linear_2.": "time_embedder.proj_out.", + "norm_out.norm.": "final_layer.norm.", + "norm_out.linear.": "final_layer.adaln_proj.linear.", + "proj_out.": "final_layer.video_out.", + "audio_proj_out.": "final_layer.audio_out.", + } + for key, value in state_dict.items(): + target = key + for source, destination in direct.items(): + if target.startswith(source): + target = destination + target[len(source) :] + break + target = target.replace("transformer_blocks.", "blocks.") + target = target.replace("token_refiner.refiner_blocks.", "token_refiner.blocks.") + target = target.replace(".attn.norm_q.", ".attn.q_norm.") + target = target.replace(".attn.norm_k.", ".attn.k_norm.") + target = target.replace(".attn.to_out.0.", ".attn.out_proj.") + target = target.replace(".ff.net.0.proj.", ".mlp.fc1.") + target = target.replace(".ff.net.2.", ".mlp.fc2.") + for part in ("q", "k", "v"): + marker = f".attn.to_{part}." + if marker in target: + prefix, suffix = target.split(marker, 1) + qkv_parts.setdefault(f"{prefix}.attn.qkv_proj.{suffix}", {})[part] = value + break + else: + renamed[target] = value + for target, parts in qkv_parts.items(): + if set(parts) != {"q", "k", "v"}: + raise ValueError(f"incomplete Diffusers QKV weights for {target}") + renamed[target] = torch.cat((parts["q"], parts["k"], parts["v"]), dim=0) + config = self._config(renamed) + return renamed, {"config": config} + + +__all__ = [ + "MINIMAX_H3_FP32_BUFFER_NAMES", + "MINIMAX_H3_FP32_PARAM_NAMES", + "MiniMaxH3DiT", + "MiniMaxH3DiTConfig", + "MiniMaxH3DiTStateDictConverter", + "_reorder_grouped_qkv_to_qkv", +] diff --git a/telefuser/models/minimax_h3_encoder.py b/telefuser/models/minimax_h3_encoder.py new file mode 100644 index 0000000..6cc76fd --- /dev/null +++ b/telefuser/models/minimax_h3_encoder.py @@ -0,0 +1,159 @@ +# SPDX-License-Identifier: Apache-2.0 +"""MiniMax H3 Qwen3-VL layer-50 encoder.""" + +from __future__ import annotations + +import re +from pathlib import Path +from typing import Any + +import torch +import torch.nn as nn +from transformers import Qwen3VLConfig, Qwen3VLModel + +from telefuser.core.base_model import BaseModel + +MINIMAX_H3_QWEN3VL_SELECTED_LM_LAYER = 50 +MINIMAX_H3_QWEN3VL_HIDDEN_DIM = 5120 +_LAYER_WEIGHT_RE = re.compile(r"^model\.language_model\.layers\.(\d+)\.") + + +def _is_unconsumed_checkpoint_weight(name: str) -> bool: + if name == "lm_head.weight" or name.startswith("model.language_model.norm."): + return True + match = _LAYER_WEIGHT_RE.match(name) + return bool(match and int(match.group(1)) >= MINIMAX_H3_QWEN3VL_SELECTED_LM_LAYER) + + +def load_minimax_h3_encoder_config(path: str | Path) -> Qwen3VLConfig: + config_path = Path(path) + source = config_path.parent if config_path.is_file() else config_path + config = Qwen3VLConfig.from_pretrained(source, local_files_only=True) + config.text_config.num_hidden_layers = MINIMAX_H3_QWEN3VL_SELECTED_LM_LAYER + config.text_config.output_hidden_states = False + config.text_config.use_cache = False + return config + + +class MiniMaxH3Encoder(BaseModel): + """Qwen3-VL multimodal backbone ending at unnormalized layer 50.""" + + def __init__(self, config: Qwen3VLConfig) -> None: + super().__init__() + if int(config.text_config.num_hidden_layers) != MINIMAX_H3_QWEN3VL_SELECTED_LM_LAYER: + raise ValueError("MiniMax H3 encoder config must be trimmed to 50 language layers") + if int(config.text_config.hidden_size) != MINIMAX_H3_QWEN3VL_HIDDEN_DIM: + raise ValueError("MiniMax H3 encoder hidden size must be 5120") + self.model = Qwen3VLModel(config) + self.model.language_model.norm = nn.Identity() + self.config = config + self.image_token_id = int(config.image_token_id) + self.video_token_id = int(config.video_token_id) + self.selected_lm_layer = MINIMAX_H3_QWEN3VL_SELECTED_LM_LAYER + self.hidden_dim = MINIMAX_H3_QWEN3VL_HIDDEN_DIM + self.layer_name_list = ["model"] + + @property + def device(self) -> torch.device: + return next(self.parameters()).device + + @torch.no_grad() + def forward( + self, + input_ids: torch.Tensor | None, + position_ids: torch.Tensor | None = None, + attention_mask: torch.Tensor | None = None, + inputs_embeds: torch.Tensor | None = None, + **kwargs: Any, + ) -> torch.Tensor: + outputs = self.model( + input_ids=input_ids, + position_ids=position_ids, + attention_mask=attention_mask, + inputs_embeds=inputs_embeds, + output_attentions=False, + output_hidden_states=False, + return_dict=True, + use_cache=False, + **kwargs, + ) + return outputs.last_hidden_state + + @torch.no_grad() + def encode_ids( + self, + input_ids: torch.Tensor, + *, + pixel_values: torch.Tensor | None = None, + image_grid_thw: torch.Tensor | None = None, + pixel_values_videos: torch.Tensor | None = None, + video_grid_thw: torch.Tensor | None = None, + ) -> torch.Tensor: + if input_ids.ndim != 1: + raise ValueError(f"input_ids must be one-dimensional, got {list(input_ids.shape)}") + if (pixel_values is None) != (image_grid_thw is None): + raise ValueError("pixel_values and image_grid_thw must be provided together") + if (pixel_values_videos is None) != (video_grid_thw is None): + raise ValueError("pixel_values_videos and video_grid_thw must be provided together") + + host_ids = input_ids.to(device="cpu", dtype=torch.long).unsqueeze(0) + host_image_grid = None if image_grid_thw is None else image_grid_thw.to(device="cpu", dtype=torch.long) + host_video_grid = None if video_grid_thw is None else video_grid_thw.to(device="cpu", dtype=torch.long) + position_ids = None + if host_image_grid is not None or host_video_grid is not None: + position_ids, _ = self.model.get_rope_index( + host_ids, + host_image_grid, + host_video_grid, + attention_mask=torch.ones_like(host_ids), + ) + call_kwargs: dict[str, Any] = { + "input_ids": host_ids.to(self.device), + "attention_mask": torch.ones_like(host_ids).to(self.device), + "output_attentions": False, + "output_hidden_states": False, + "return_dict": True, + "use_cache": False, + } + if position_ids is not None: + call_kwargs["position_ids"] = position_ids.to(self.device) + if pixel_values is not None: + call_kwargs["pixel_values"] = pixel_values.to(self.device, torch.bfloat16) + call_kwargs["image_grid_thw"] = host_image_grid + if pixel_values_videos is not None: + call_kwargs["pixel_values_videos"] = pixel_values_videos.to(self.device, torch.bfloat16) + call_kwargs["video_grid_thw"] = host_video_grid + hidden = self.model(**call_kwargs).last_hidden_state[0].to(torch.bfloat16) + expected = (input_ids.numel(), self.hidden_dim) + if tuple(hidden.shape) != expected: + raise ValueError(f"unexpected MiniMax H3 encoder shape {tuple(hidden.shape)}, expected {expected}") + return hidden + + @staticmethod + def state_dict_converter(config_path: str | Path) -> MiniMaxH3EncoderStateDictConverter: + return MiniMaxH3EncoderStateDictConverter(config_path) + + +class MiniMaxH3EncoderStateDictConverter: + def __init__(self, config_path: str | Path) -> None: + self.config = load_minimax_h3_encoder_config(config_path) + + def from_official(self, state_dict: dict[str, torch.Tensor]) -> tuple[dict[str, torch.Tensor], dict[str, Any]]: + converted = { + name: value + for name, value in state_dict.items() + if "rotary_emb.inv_freq" not in name and not _is_unconsumed_checkpoint_weight(name) + } + return converted, {"config": self.config} + + def from_diffusers(self, state_dict: dict[str, torch.Tensor]) -> tuple[dict[str, torch.Tensor], dict[str, Any]]: + return self.from_official(state_dict) + + +__all__ = [ + "MINIMAX_H3_QWEN3VL_HIDDEN_DIM", + "MINIMAX_H3_QWEN3VL_SELECTED_LM_LAYER", + "MiniMaxH3Encoder", + "MiniMaxH3EncoderStateDictConverter", + "load_minimax_h3_encoder_config", +] diff --git a/telefuser/models/minimax_h3_video/__init__.py b/telefuser/models/minimax_h3_video/__init__.py new file mode 100644 index 0000000..28f7bde --- /dev/null +++ b/telefuser/models/minimax_h3_video/__init__.py @@ -0,0 +1,5 @@ +# SPDX-License-Identifier: Apache-2.0 + +from .klvae import AutoencoderKLLegacy + +__all__ = ["AutoencoderKLLegacy"] diff --git a/telefuser/models/minimax_h3_video/attention.py b/telefuser/models/minimax_h3_video/attention.py new file mode 100644 index 0000000..0c7d2d8 --- /dev/null +++ b/telefuser/models/minimax_h3_video/attention.py @@ -0,0 +1,157 @@ +# SPDX-License-Identifier: Apache-2.0 +# Attention module for the MiniMax H3 visual VAE (inference-only bundle). +from typing import Optional + +import torch +import torch.nn as nn +from diffusers.utils import logging + +from .flash import flash_attn +from .vit_utils import apply_rotary_pos_emb_qk + +logger = logging.get_logger(__name__) # pylint: disable=invalid-name + + +def _vit_norm_input(module, hidden_states): + return hidden_states.float() + + +def _apply_qk_norm(module, hidden_states): + if ( + isinstance(module, (nn.LayerNorm, nn.RMSNorm)) + and getattr(module, "weight", None) is None + and getattr(module, "bias", None) is None + and hidden_states.is_cuda + and hidden_states.dtype in (torch.float16, torch.bfloat16) + and not torch.is_grad_enabled() + and not torch.compiler.is_compiling() + ): + # CUDA LayerNorm/RMSNorm accumulates half/bfloat16 inputs in FP32. + # With no affine parameters its half output is bit-identical to the + # released FP32-norm-then-cast recipe, without two full-tensor casts. + with torch.autocast("cuda", enabled=False): + return module(hidden_states) + return module(_vit_norm_input(module, hidden_states)).to(hidden_states.dtype) + + +class Attention(nn.Module): + def __init__( + self, + heads, + dim_head, + embed_dim: Optional[int] = None, + qk_norm_type: Optional[str] = None, + qk_norm_affine: bool = False, + bias: bool = True, + out_bias: Optional[bool] = None, + eps: float = 1e-5, + **kwargs, + ): + super().__init__() + self.dim_head = dim_head + self.heads = heads + self.attn_inner_dim = dim_head * heads + self.embed_dim = embed_dim if embed_dim is not None else self.attn_inner_dim + + out_bias = out_bias if out_bias is not None else bias + + if qk_norm_type is None: + self.norm_q = None + self.norm_k = None + elif qk_norm_type == "layer_norm": + self.norm_q = nn.LayerNorm(dim_head, eps=eps, elementwise_affine=qk_norm_affine) + self.norm_k = nn.LayerNorm(dim_head, eps=eps, elementwise_affine=qk_norm_affine) + elif qk_norm_type == "rms_norm": + self.norm_q = nn.RMSNorm(dim_head, eps=eps, elementwise_affine=qk_norm_affine) + self.norm_k = nn.RMSNorm(dim_head, eps=eps, elementwise_affine=qk_norm_affine) + else: + raise ValueError(f"unknown qk_norm_type: {qk_norm_type}. Should be None,'layer_norm','rms_norm'") + + self.to_qkv = nn.Linear(self.embed_dim, self.attn_inner_dim * 3, bias=bias) + + self.to_out = nn.Linear(self.attn_inner_dim, self.embed_dim, bias=out_bias) + + if len(kwargs) > 0: + logger.warning(f"Unused kwargs: {kwargs}") + + def _perform_attention(self, query, key, value, pack_info): + cu_seqlens = pack_info.get("cu_seqlens", None) + mask_mod = pack_info.get("mask_mod", None) + block_sparse = pack_info.get("block_sparse", None) + valid_seq_len = pack_info.get("valid_seq_len", None) + + if cu_seqlens is not None: + raise NotImplementedError("varlen attention is not supported in this inference-only bundle") + + padded_seq_len = query.shape[1] + if valid_seq_len is not None: + valid_seq_len = int(valid_seq_len) + if not 0 < valid_seq_len <= padded_seq_len: + raise ValueError( + "valid_seq_len must be in (0, padded_seq_len], got " + f"{valid_seq_len} for padded_seq_len={padded_seq_len}" + ) + query = query[:, :valid_seq_len] + key = key[:, :valid_seq_len] + value = value[:, :valid_seq_len] + + if mask_mod is not None: + hidden_states = flash_attn( + query, + key, + value, + mask_mod=mask_mod, + block_sparse=block_sparse, + ) + else: + hidden_states = flash_attn( + query, + key, + value, + ) + + if valid_seq_len is not None and valid_seq_len < padded_seq_len: + hidden_states = torch.cat( + [ + hidden_states, + hidden_states.new_zeros( + hidden_states.shape[0], + padded_seq_len - valid_seq_len, + hidden_states.shape[2], + hidden_states.shape[3], + ), + ], + dim=1, + ) + + return hidden_states + + def perform_attention(self, query, key, value, pack_info={}): + return self._perform_attention(query, key, value, pack_info) + + def forward( + self, + hidden_states: torch.Tensor, + rotary_pos_emb: Optional[torch.Tensor] = None, + pack_info: dict = {}, + ) -> torch.Tensor: + batch_size, seq_len, _ = hidden_states.shape + + qkv = self.to_qkv(hidden_states) + qkv = qkv.view(batch_size, seq_len, -1, 3 * self.dim_head) + query, key, value = torch.chunk(qkv, 3, dim=-1) + + if self.norm_q is not None: + query = _apply_qk_norm(self.norm_q, query) + if self.norm_k is not None: + key = _apply_qk_norm(self.norm_k, key) + + if rotary_pos_emb is not None: + query, key = apply_rotary_pos_emb_qk(query, key, rotary_pos_emb) + + hidden_states = self.perform_attention(query, key, value, pack_info) + + hidden_states = hidden_states.reshape(batch_size, seq_len, -1) + hidden_states = self.to_out(hidden_states) + + return hidden_states diff --git a/telefuser/models/minimax_h3_video/base_module.py b/telefuser/models/minimax_h3_video/base_module.py new file mode 100644 index 0000000..3c58beb --- /dev/null +++ b/telefuser/models/minimax_h3_video/base_module.py @@ -0,0 +1,195 @@ +# SPDX-License-Identifier: Apache-2.0 +# Transformer building blocks for the MiniMax H3 visual VAE ViT decoder. +import math +from typing import Optional + +import torch +import torch.nn as nn +from diffusers.utils import logging +from diffusers.utils.torch_utils import maybe_allow_in_graph + +from telefuser.ops.activations import silu_and_mul + +from .attention import Attention + +logger = logging.get_logger(__name__) # pylint: disable=invalid-name + + +def _vit_norm_input(module, hidden_states): + return hidden_states.float() + + +def _scaled_residual_add(residual, x, scale): + return residual + x * scale + + +class FeedForward(nn.Module): + def __init__( + self, + dim: int, + dim_out: Optional[int] = None, + mult: int = 4, + activation_fn: str = "silu", + bias: bool = True, + use_gated: bool = True, + glu_balanced: bool = False, + ): + super().__init__() + ratio = 2 / 3 if (use_gated and glu_balanced) else 1 + inner_dim = round(dim * mult * ratio) + dim_out = dim_out if dim_out is not None else dim + self.use_gated = use_gated + + if use_gated: + self.w1 = nn.Linear(dim, inner_dim * 2, bias=bias) + else: + self.w1 = nn.Linear(dim, inner_dim, bias=bias) + + if activation_fn == "silu": + self.act_fn = nn.SiLU() + elif activation_fn == "gelu": + self.act_fn = nn.GELU() + elif activation_fn == "gelu-approximate": + self.act_fn = nn.GELU(approximate="tanh") + else: + raise ValueError(f"Unsupported activation function: {activation_fn}") + + self.w2 = nn.Linear(inner_dim, dim_out, bias=bias) + + def _forward_impl(self, hidden_states: torch.Tensor) -> torch.Tensor: + hidden_states = self.w1(hidden_states) + + if self.use_gated: + if isinstance(self.act_fn, nn.SiLU): + hidden_states = silu_and_mul(hidden_states) + else: + gate, hidden_states = hidden_states.chunk(2, dim=-1) + hidden_states = self.act_fn(gate).mul_(hidden_states) + else: + hidden_states = self.act_fn(hidden_states) + + hidden_states = self.w2(hidden_states) + return hidden_states + + def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: + return self._forward_impl(hidden_states) + + +class RotaryEmbeddingND(nn.Module): + def __init__(self, dim, rotary_base=10000, n_dim=3, use_angle=False): + super().__init__() + self.dim = dim + self.n_dim = n_dim + + if dim % (2 * n_dim) != 0: + raise ValueError(f"head_dim {dim} must be divisible by 2 * n_dim {2 * n_dim}") + + if use_angle: + self.angle_scale = 2.0 * math.pi + else: + self.angle_scale = 1.0 + + inv_freq = 1 / rotary_base ** torch.arange(0, 1, 2 * n_dim / dim, dtype=torch.float32) + self.register_buffer("inv_freq", inv_freq, persistent=False) + + def forward(self, img_ids): + B, N, D = img_ids.shape + if D != self.n_dim: + raise ValueError(f"Expected {self.n_dim} dimensions, got {D}") + + with torch.autocast("cuda", enabled=False): + angles = self.angle_scale * img_ids[:, :, :, None] * self.inv_freq.to(img_ids.device)[None, None, None, :] + angles = angles.flatten(2, 3) + angles = angles.tile(2) + angles = angles.unsqueeze(2) + + cos = torch.cos(angles) + sin = torch.sin(angles) + + return cos.to(dtype=img_ids.dtype), sin.to(dtype=img_ids.dtype) + + +@maybe_allow_in_graph +class TransformerBlock(nn.Module): + def __init__( + self, + heads: int, + dim_head: int, + embed_dim: Optional[int] = None, + ffn_glu_balanced: bool = False, + norm_type: str = "layer_norm", + norm_affine: bool = True, + qk_norm_type: str = "rms_norm", + qk_norm_affine: bool = False, + ffn_activation_fn: str = "silu", + ffn_use_gated: bool = True, + use_scale: bool = True, + bias: bool = True, + eps: float = 1e-5, + **kwargs, + ): + super().__init__() + dim = embed_dim if embed_dim is not None else dim_head * heads + self.use_scale = use_scale + + if norm_type == "layer_norm": + norm_class = nn.LayerNorm + elif norm_type == "rms_norm": + norm_class = nn.RMSNorm + else: + raise ValueError(f"unknown norm_type {norm_type}") + + self.norm1 = norm_class( + dim, + elementwise_affine=norm_affine, + eps=eps, + ) + self.attn = Attention( + heads=heads, + dim_head=dim_head, + embed_dim=dim, + qk_norm_type=qk_norm_type, + qk_norm_affine=qk_norm_affine, + bias=bias, + eps=eps, + **kwargs, + ) + if use_scale: + self.scale1 = nn.Parameter(torch.zeros(dim)) + + self.norm2 = norm_class( + dim, + elementwise_affine=norm_affine, + eps=eps, + ) + self.ff = FeedForward( + dim=dim, + activation_fn=ffn_activation_fn, + bias=bias, + use_gated=ffn_use_gated, + glu_balanced=ffn_glu_balanced, + ) + if use_scale: + self.scale2 = nn.Parameter(torch.zeros(dim)) + + def forward( + self, + hidden_states: torch.FloatTensor, + rotary_pos_emb: Optional[torch.FloatTensor] = None, + pack_info: dict = {}, + ): + norm_hidden_states = self.norm1(_vit_norm_input(self.norm1, hidden_states)).to(hidden_states.dtype) + attn_output = self.attn(norm_hidden_states, rotary_pos_emb, pack_info) + if self.use_scale: + hidden_states = _scaled_residual_add(hidden_states, attn_output, self.scale1) + else: + hidden_states = hidden_states + attn_output + + norm_hidden_states = self.norm2(_vit_norm_input(self.norm2, hidden_states)).to(hidden_states.dtype) + ff_output = self.ff(norm_hidden_states) + if self.use_scale: + hidden_states = _scaled_residual_add(hidden_states, ff_output, self.scale2) + else: + hidden_states = hidden_states + ff_output + + return hidden_states diff --git a/telefuser/models/minimax_h3_video/conv.py b/telefuser/models/minimax_h3_video/conv.py new file mode 100644 index 0000000..98a2e48 --- /dev/null +++ b/telefuser/models/minimax_h3_video/conv.py @@ -0,0 +1,83 @@ +# SPDX-License-Identifier: Apache-2.0 +# 3D convolution for the MiniMax H3 visual VAE. +import torch.nn as nn +import torch.nn.functional as F + + +class BaseConv3d(nn.Conv3d): + def __init__( + self, + in_channels, + out_channels, + kernel_size, + stride=1, + padding=0, + bias=True, + padding_mode="zeros", + padding_mode_t=None, + causal=True, + ): + super().__init__( + in_channels, + out_channels, + kernel_size=kernel_size, + stride=stride, + padding=padding, + bias=bias, + padding_mode=padding_mode, + ) + padding_mode = "constant" if padding_mode == "zeros" else padding_mode + padding_mode_t = "constant" if padding_mode_t == "zeros" else padding_mode_t + self.pad_mode = padding_mode + self.pad_mode_t = padding_mode_t or ("constant" if causal else "replicate") + self.causal = causal + + def _apply_temporal_padding(self, x): + B, C, D, H, W = x.shape + if D > 1: + pad_size = ( + 0, + 0, + 0, + 0, + self.padding[0] * 2 if self.causal else self.padding[0], + 0 if self.causal else self.padding[0], + ) + return F.pad(x, pad_size, mode=self.pad_mode_t) + else: + if self.pad_mode_t == "constant": + assert self.causal, "Zeros padding is only supported for causal mode" + return F.pad( + x, + (0, 0, 0, 0, self.kernel_size[0] - 1, 0), + mode="constant", + ) + else: + return x.expand(-1, -1, self.kernel_size[0], -1, -1) + + def _apply_padding(self, x): + if sum(self.padding) == 0: + return x + + x = F.pad( + x, + (self.padding[2], self.padding[2], self.padding[1], self.padding[1], 0, 0), + mode=self.pad_mode, + ) + + x = self._apply_temporal_padding(x) + return x + + def forward(self, x): + if sum(self.padding) == 0: + return super().forward(x) + + x = self._apply_padding(x) + return F.conv3d( + x, + self.weight, + self.bias, + stride=self.stride, + padding=0, + dilation=self.dilation, + ) diff --git a/telefuser/models/minimax_h3_video/flash.py b/telefuser/models/minimax_h3_video/flash.py new file mode 100644 index 0000000..f4b8339 --- /dev/null +++ b/telefuser/models/minimax_h3_video/flash.py @@ -0,0 +1,135 @@ +# SPDX-License-Identifier: Apache-2.0 +# Torch-native attention implemented with PyTorch SDPA instead of FA4/CUTLASS. +import torch + +from telefuser.ops.attention import attention + +_BLOCK_CAUSAL_MASK_MOD_CACHE = {} + + +def _as_bool_mask(mask, *, device): + if not isinstance(mask, torch.Tensor): + mask = torch.as_tensor(mask, device=device) + return mask.to(device=device, dtype=torch.bool) + + +def _ensure_nonempty_rows(mask): + if mask.numel() == 0 or mask.shape[-1] == 0: + return mask + empty = ~mask.any(dim=-1) + mask[..., 0] |= empty + return mask + + +def _sdpa_attention(query, key, value, causal=False, attn_mask=None): + if attn_mask is not None and attn_mask.dim() == 3: + attn_mask = attn_mask.unsqueeze(0) + return attention( + query, + key, + value, + attn_mask=attn_mask, + input_layout="BSND", + output_layout="BSND", + is_causal=causal, + ).nan_to_num(0.0) + + +def _mask_mod_to_dense(mask_mod, batch, heads, q_len, kv_len, device, aux_tensors=None): + q_idx = torch.arange(q_len, device=device).view(q_len, 1) + kv_idx = torch.arange(kv_len, device=device).view(1, kv_len) + dense = torch.empty((batch, heads, q_len, kv_len), dtype=torch.bool, device=device) + for b in range(batch): + b_idx = torch.tensor(b, device=device) + for h in range(heads): + h_idx = torch.tensor(h, device=device) + mask = mask_mod(b_idx, h_idx, q_idx, kv_idx, None, aux_tensors) + dense[b, h] = _as_bool_mask(mask, device=device) + return _ensure_nonempty_rows(dense) + + +######################################################### +# Block causal attention +######################################################### + + +def make_block_causal_mask_mod(num_tokens, block_size, num_special=0, suffix=False): + if num_tokens < 0: + raise ValueError(f"num_tokens must be non-negative, got {num_tokens}") + if block_size <= 0: + raise ValueError(f"block_size must be positive, got {block_size}") + if num_special < 0: + raise ValueError(f"num_special must be non-negative, got {num_special}") + + cache_key = (num_tokens, block_size, num_special, suffix) + if cache_key in _BLOCK_CAUSAL_MASK_MOD_CACHE: + return _BLOCK_CAUSAL_MASK_MOD_CACHE[cache_key] + + if suffix: + + def mask_mod(b, h, q_idx, kv_idx, seqlen_info, aux_tensors): + del b, h, seqlen_info, aux_tensors + q_is_special = q_idx >= num_tokens + kv_is_special = kv_idx >= num_tokens + return q_is_special | kv_is_special | (q_idx // block_size >= kv_idx // block_size) + + else: + + def mask_mod(b, h, q_idx, kv_idx, seqlen_info, aux_tensors): + del b, h, seqlen_info, aux_tensors + q_is_special = q_idx < num_special + kv_is_special = kv_idx < num_special + q_block_idx = (q_idx - num_special) // block_size + kv_block_idx = (kv_idx - num_special) // block_size + return q_is_special | kv_is_special | (q_block_idx >= kv_block_idx) + + mask_mod.block_sparse_cache_key = ( + "block_causal", + num_tokens, + block_size, + num_special, + suffix, + ) + _BLOCK_CAUSAL_MASK_MOD_CACHE[cache_key] = mask_mod + return mask_mod + + +######################################################### +# Public entry point +######################################################### + + +@torch.compiler.disable +def flash_attn( + query: torch.Tensor, + key: torch.Tensor, + value: torch.Tensor, + causal: bool = False, + mask_mod=None, + block_sparse=None, + aux_tensors=None, +) -> torch.Tensor: + use_masked = mask_mod is not None or block_sparse is not None + + if block_sparse is not None and mask_mod is None: + raise ValueError("block_sparse requires mask_mod") + if causal and mask_mod is not None: + raise ValueError("causal must be encoded in mask_mod when using masked attention") + if aux_tensors is not None and not use_masked: + raise ValueError("aux_tensors is only supported with masked attention") + + if use_masked: + batch, q_len, heads, _ = query.shape + kv_len = key.shape[1] + dense_mask = _mask_mod_to_dense( + mask_mod, + batch, + heads, + q_len, + kv_len, + query.device, + aux_tensors=aux_tensors, + ) + return _sdpa_attention(query, key, value, attn_mask=dense_mask) + + return _sdpa_attention(query, key, value, causal=causal) diff --git a/telefuser/models/minimax_h3_video/klvae.py b/telefuser/models/minimax_h3_video/klvae.py new file mode 100644 index 0000000..c72583d --- /dev/null +++ b/telefuser/models/minimax_h3_video/klvae.py @@ -0,0 +1,1152 @@ +# SPDX-License-Identifier: Apache-2.0 +# MiniMax H3 visual VAE: 3D causal CNN encoder + ViT3D decoder (inference-only bundle). +import math +from typing import List, Union + +import numpy as np +import torch +import torch.nn as nn +from PIL import Image +from diffusers.configuration_utils import ConfigMixin, register_to_config +from diffusers.loaders.single_file_model import FromOriginalModelMixin +from diffusers.models import ModelMixin +from diffusers.utils import logging + +from .processor import ( + VAEProcessor, + get_denormalize_transform, + get_normalize_transform, +) +from .vae_cnn import EncoderFCN3D +from .vae_vit import ViT3DDecoder + +logger = logging.get_logger(__name__) # pylint: disable=invalid-name + + +def _resolve_temporal_cat_dtype(): + return None + + +def _resolve_temporal_stream_cat(): + return True + + +def get_tile_parallel_state(): + return 0, 1 + + +class DiagonalGaussianDistribution(object): + def __init__(self, parameters, upcast_fp32=True): + if upcast_fp32: + parameters = parameters.to(dtype=torch.float32) + + self.parameters = parameters + self.mean, self.logvar = torch.chunk(parameters, 2, dim=1) + self.logvar = torch.clamp(self.logvar, -30.0, 20.0) + self.std = self.logvar.mul(0.5).exp_() + + @torch.compiler.disable + def sample(self, generator=None): + noise = torch.randn(self.mean.shape, generator=generator) + noise = noise.to(device=self.parameters.device) + return noise.mul_(self.std).add_(self.mean) + + +class ClsTokenAggregator: + def __init__(self, vae_model): + self.vae = vae_model + self.cls_tokens = [] + + def __enter__(self): + return self + + def __exit__(self, exc_type, exc_val, exc_tb): + if self.cls_tokens and hasattr(self.vae.encoder, "loss_info"): + self.vae.encoder.loss_info["cls_token"] = torch.stack(self.cls_tokens, dim=0).mean(dim=0) + return False + + def collect(self): + if hasattr(self.vae.encoder, "loss_info") and "cls_token" in self.vae.encoder.loss_info: + self.cls_tokens.append(self.vae.encoder.loss_info["cls_token"].clone()) + + def collect_stacked(self, num_tiles, sample_batch_size): + if hasattr(self.vae.encoder, "loss_info") and "cls_token" in self.vae.encoder.loss_info: + cls_token = self.vae.encoder.loss_info["cls_token"] + cls_token = cls_token.unflatten(0, (num_tiles, sample_batch_size)) + self.cls_tokens.extend(token.clone() for token in cls_token) + + +class AutoencoderKL(ModelMixin, ConfigMixin, FromOriginalModelMixin): + r""" + Abstract shared base for the MiniMax H3 visual VAE. + + This class only carries the shared inference machinery (temporal + chunking, tiling, encode/decode entry points). Instantiate the concrete + subclass ``AutoencoderKLLegacy`` via ``from_pretrained`` instead. + """ + + _compilable_modules = ["encoder", "decoder"] + _deprecated_kwargs = [ + "clip_length", + "token_drop", + "isolated_first_frame", + "isolated_last_frame", + "isolated_key_frame", + "encoder_tiling", + "decoder_tiling", + "parallel_tiling", + "stack_tiling", + "tile_size", + "tile_overlap_min", + "decoder_tile_size", + "decoder_tile_overlap_min", + "latent_patch_size", + "crop_mode", + "encoder_parallel", + "decoder_parallel", + "chunk_dim", + ] # legacy config keys accepted by from_pretrained for checkpoint compatibility + + def setup_forward(self, **kwargs): + self.clip_length = kwargs.get("clip_length", 17) + self.token_drop = kwargs.get("token_drop", 0) + self.frame_drop = self.token_drop * self.vae_ratio_t + self.frame_pre_padding = (-self.clip_length) % self.vae_ratio_t + self.tokens_chunk_size = math.ceil(self.clip_length / self.vae_ratio_t) + self.token_overlap = (-self.token_drop) % self.tokens_chunk_size + self.frame_overlap = max(self.token_overlap * self.vae_ratio_t - self.frame_pre_padding, 0) + self.isolated_first_frame = kwargs.get("isolated_first_frame", False) + self.isolated_last_frame = kwargs.get("isolated_last_frame", False) + self.isolated_key_frame = kwargs.get("isolated_key_frame", False) + + self.encoder_tiling = kwargs.get("encoder_tiling", False) + self.decoder_tiling = kwargs.get("decoder_tiling", False) + self.stack_tiling = kwargs.get("stack_tiling", False) + self.tile_size = kwargs.get("tile_size", 256) + self.tile_overlap_min = kwargs.get("tile_overlap_min", 64) + self.decoder_tile_size = kwargs.get("decoder_tile_size", self.tile_size) + self.decoder_tile_overlap_min = kwargs.get("decoder_tile_overlap_min", self.tile_overlap_min) + self._blend_weight_cache = {} + self.latent_patch_size = kwargs.get("latent_patch_size", 1) + self.crop_mode = kwargs.get("crop_mode", "top_left") + self.pixel_norm_type = kwargs.get("pixel_norm_type", "imagenet") + + if kwargs.get("encoder_parallel", False) or kwargs.get("decoder_parallel", False): + raise ValueError( + "MiniMax H3 VAE spatial sharding is unsupported; use complete-tile parallel decode instead" + ) + parallel_tiling = kwargs.get("parallel_tiling", False) + if hasattr(self, "parallel_tiling") and parallel_tiling != self.parallel_tiling: + logger.warning("Do not support changing parallel tiling after initialization") + else: + self.parallel_tiling = parallel_tiling + + processor_kwargs = { + "vae_ratio": self.vae_ratio, + "vae_ratio_t": self.vae_ratio_t, + "clip_length": self.clip_length, + "frame_overlap": self.frame_overlap, + "token_overlap": self.token_overlap, + "tokens_chunk_size": self.tokens_chunk_size, + "isolated_last_frame": self.isolated_last_frame, + "latent_patch_size": self.latent_patch_size, + "crop_mode": self.crop_mode, + "pixel_norm_type": self.pixel_norm_type, + "transform": self.transform, + "transform_rev": self.transform_rev, + "use_3d_conv": self.use_3d_conv, + } + if hasattr(self, "processor"): + for key, value in processor_kwargs.items(): + setattr(self.processor, key, value) + else: + self.processor = VAEProcessor(**processor_kwargs) + + def split_tiles(self, input_len, is_decoder=False): + tile_size = self.decoder_tile_size if is_decoder else self.tile_size + tile_overlap_min = self.decoder_tile_overlap_min if is_decoder else self.tile_overlap_min + + if tile_size >= input_len: + return [0], [input_len], [] + + N = math.ceil(input_len / tile_size) + while True: + overlaps = [tile_overlap_min] * (N - 1) + remaining = tile_size * N - sum(overlaps) - input_len + + if remaining < 0: + N += 1 + else: + break + + remaining_units = remaining // self.vae_ratio + for i in range(remaining_units): + overlaps[i % (N - 1)] += self.vae_ratio + + tile_start_idx = [0] + for i in range(N - 1): + tile_start_idx.append(tile_start_idx[-1] + tile_size - overlaps[i]) + + tile_len = [tile_size] * N + return tile_start_idx, tile_len, overlaps + + def blend(self, a: torch.Tensor, b: torch.Tensor, blend_extent: int, dim: int) -> torch.Tensor: + blend_extent = min(a.shape[dim], b.shape[dim], blend_extent) + + cache_key = (blend_extent, b.device, b.dtype) + weights = self._blend_weight_cache.get(cache_key) + if weights is None: + positions = torch.arange(blend_extent, device=b.device, dtype=b.dtype) + weights = (1 - positions / blend_extent, positions / blend_extent) + self._blend_weight_cache[cache_key] = weights + weight_a, weight_b = weights + + shape = [1] * a.ndim + shape[dim] = blend_extent + weight_a = weight_a.view(shape) + weight_b = weight_b.view(shape) + + slice_a = [slice(None)] * a.ndim + slice_a[dim] = slice(-blend_extent, None) + a_overlap = a[tuple(slice_a)] + + slice_b = [slice(None)] * b.ndim + slice_b[dim] = slice(0, blend_extent) + b_overlap = b[tuple(slice_b)] + + blended = a_overlap * weight_a + blended.add_(b_overlap * weight_b) + + if blend_extent < b.shape[dim]: + slice_b_rest = [slice(None)] * b.ndim + slice_b_rest[dim] = slice(blend_extent, None) + b_rest = b[tuple(slice_b_rest)] + return torch.cat([blended, b_rest], dim=dim) + else: + return blended + + def _assemble_tiles(self, rows, y_overlap, x_overlap): + output_height = sum(row[0].shape[-2] - (y_overlap[i] if i < len(rows) - 1 else 0) for i, row in enumerate(rows)) + output_width = sum( + tile.shape[-1] - (x_overlap[j] if j < len(rows[0]) - 1 else 0) for j, tile in enumerate(rows[0]) + ) + output = rows[0][0].new_empty((*rows[0][0].shape[:-2], output_height, output_width)) + + y_offset = 0 + # released VAE blends vertically before horizontally; preserve that + # order while writing each cropped tile into the final tensor + for i, row in enumerate(rows): + row_height = row[0].shape[-2] - (y_overlap[i] if i < len(rows) - 1 else 0) + x_offset = 0 + for j, tile in enumerate(row): + if i > 0: + tile = self.blend(rows[i - 1][j], tile, y_overlap[i - 1], dim=-2) + if j > 0: + tile = self.blend(row[j - 1], tile, x_overlap[j - 1], dim=-1) + if i < len(rows) - 1: + tile = tile[..., : -y_overlap[i], :] + if j < len(row) - 1: + tile = tile[..., :, : -x_overlap[j]] + + tile_height, tile_width = tile.shape[-2:] + output[ + ..., + y_offset : y_offset + tile_height, + x_offset : x_offset + tile_width, + ].copy_(tile) + x_offset += tile_width + y_offset += row_height + + return output + + def _all_gather_tiled_results(self, tasks, num_tiles): + tile_rank, tile_world_size = get_tile_parallel_state() + + if not tasks: + raise ValueError(f"Found empty tasks on tile rank {tile_rank}") + + expected_tasks = (num_tiles - tile_rank + tile_world_size - 1) // tile_world_size + if len(tasks) != expected_tasks: + raise ValueError( + f"Expected {expected_tasks} tiled tasks on rank {tile_rank}, " + f"got {len(tasks)} for num_tiles={num_tiles}, " + f"world_size={tile_world_size}" + ) + + max_tasks = (num_tiles + tile_world_size - 1) // tile_world_size + if len(tasks) == max_tasks: + stacked = torch.stack(tasks, dim=0) + else: + stacked = tasks[0].new_empty((max_tasks, *tasks[0].shape)) + torch.stack(tasks, dim=0, out=stacked[: len(tasks)]) + stacked[len(tasks) :].zero_() + + # Round-robin tile ownership makes every rank's task count known from + # num_tiles and world size. Pad only the leading task dimension and use a + # single equal-shape all-gather; the previous path paid for a barrier, + # a shape all-gather, and then a padded data all-gather per temporal + # clip. + gathered = [stacked] + + results = [None] * num_tiles + for rank, rank_tensors in enumerate(gathered): + num_rank_tasks = (num_tiles - rank + tile_world_size - 1) // tile_world_size + for k in range(num_rank_tasks): + global_idx = k * tile_world_size + rank + results[global_idx] = rank_tensors[k] + + return results + + def _local_tile_indices(self, num_tiles, tile_rank, tile_world_size): + return list(range(tile_rank, num_tiles, tile_world_size)) + + def _run_tile_tasks(self, tiles, tile_indices, forward_fn, stack_tiling, cls_agg=None): + if stack_tiling and tile_indices: + sample_batch_size = tiles[0].shape[0] + tile_batch = torch.cat([tiles[idx] for idx in tile_indices], dim=0) + output_batch = forward_fn(tile_batch) + output_tiles = output_batch.unflatten(0, (len(tile_indices), sample_batch_size)).unbind(dim=0) + if cls_agg is not None: + cls_agg.collect_stacked(len(tile_indices), sample_batch_size) + return list(output_tiles) + + tasks = [] + for idx in tile_indices: + tasks.append(forward_fn(tiles[idx])) + if cls_agg is not None: + cls_agg.collect() + return tasks + + def tiled_encode(self, x): + if self.parallel_tiling: # Fast online encoding for large videos + tile_rank, tile_world_size = get_tile_parallel_state() + else: + tile_rank, tile_world_size = 0, 1 + + height, width = x.shape[-2], x.shape[-1] + y_idx, y_len, y_overlap = self.split_tiles(height, False) + x_idx, x_len, x_overlap = self.split_tiles(width, False) + + i_max, j_max = len(y_idx), len(x_idx) + num_tiles = i_max * j_max + if tile_world_size > num_tiles: + # Every rank executes this replicated path. When the canvas has + # fewer tiles than ranks, run locally instead of assigning + # empty task lists that cannot participate in the tensor gather. + tile_rank, tile_world_size = 0, 1 + + x_tiles = [] + for i, (i_pos, i_len) in enumerate(zip(y_idx, y_len)): + for j, (j_pos, j_len) in enumerate(zip(x_idx, x_len)): + tile = x[..., i_pos : i_pos + i_len, j_pos : j_pos + j_len] + x_tiles.append(tile) + + with ClsTokenAggregator(self) as agg: + local_tile_indices = self._local_tile_indices(num_tiles, tile_rank, tile_world_size) + stack_tiling = self.stack_tiling and not (self.training and getattr(self.encoder, "mask_enabled", False)) + encoded_tasks = self._run_tile_tasks(x_tiles, local_tile_indices, self.encode, stack_tiling, agg) + + if tile_world_size > 1: + all_encoded = self._all_gather_tiled_results(encoded_tasks, num_tiles) + if agg.cls_tokens: + agg.cls_tokens = self._all_gather_tiled_results(agg.cls_tokens, num_tiles) + else: + all_encoded = encoded_tasks + + rows = [[None for _ in range(j_max)] for _ in range(i_max)] + for idx, encoded in enumerate(all_encoded): + i, j = idx // j_max, idx % j_max + rows[i][j] = encoded.to(x.device) + + latent_y_overlap = [tile_overlap // self.vae_ratio for tile_overlap in y_overlap] + latent_x_overlap = [tile_overlap // self.vae_ratio for tile_overlap in x_overlap] + + z = self._assemble_tiles(rows, latent_y_overlap, latent_x_overlap) + + return z + + def tiled_decode(self, z): + if self.parallel_tiling: # Fast online decoding for large videos + tile_rank, tile_world_size = get_tile_parallel_state() + else: + tile_rank, tile_world_size = 0, 1 + + height, width = ( + z.shape[-2] * self.vae_ratio, + z.shape[-1] * self.vae_ratio, + ) + y_idx, y_len, y_overlap = self.split_tiles(height, True) + x_idx, x_len, x_overlap = self.split_tiles(width, True) + + i_max, j_max = len(y_idx), len(x_idx) + num_tiles = i_max * j_max + if tile_world_size > num_tiles: + # See tiled_encode: small canvases fall back to replicated local + # tiling so no decode rank is assigned an empty task list. + tile_rank, tile_world_size = 0, 1 + + z_tiles = [] + for i, (i_pos, i_len) in enumerate(zip(y_idx, y_len)): + i_pos, i_len = ( + i_pos // self.vae_ratio, + i_len // self.vae_ratio, + ) + for j, (j_pos, j_len) in enumerate(zip(x_idx, x_len)): + j_pos, j_len = (j_pos // self.vae_ratio, j_len // self.vae_ratio) + tile = z[..., i_pos : i_pos + i_len, j_pos : j_pos + j_len] + z_tiles.append(tile) + + local_tile_indices = self._local_tile_indices(num_tiles, tile_rank, tile_world_size) + stack_tiling = self.stack_tiling and not (self.training and getattr(self.decoder, "mask_enabled", False)) + decoded_tasks = self._run_tile_tasks(z_tiles, local_tile_indices, self.decode, stack_tiling) + + if tile_world_size > 1: + all_decoded = self._all_gather_tiled_results(decoded_tasks, num_tiles) + else: + all_decoded = decoded_tasks + + rows = [[None for _ in range(j_max)] for _ in range(i_max)] + for idx, decoded in enumerate(all_decoded): + i, j = idx // j_max, idx % j_max + rows[i][j] = decoded.to(z.device) + + dec = self._assemble_tiles(rows, y_overlap, x_overlap) + return dec + + def _adaptive_encode(self, x): + if self.encoder_tiling: + return self.tiled_encode(x) + else: + return self.encode(x) + + def _adaptive_decode(self, z): + if self.decoder_tiling: + return self.tiled_decode(z) + else: + return self.decode(z) + + def trim_code(self, z, target_codes): + if target_codes < z.shape[2]: + if self.causal_encoder: + z = z[:, :, -target_codes:, :, :] + else: + start_frame = (z.shape[2] - target_codes) // 2 + z = z[:, :, start_frame : start_frame + target_codes, :, :] + return z + + def trim_output(self, dec, target_frames): + if target_frames < dec.shape[2]: + if self.causal_encoder: # This is defined by encoder, not decoder + dec = dec[:, :, -target_frames:, :, :] + else: + start_frame = (dec.shape[2] - target_frames) // 2 + dec = dec[:, :, start_frame : start_frame + target_frames, :, :] + return dec + + def encode_temporal(self, x): + offset_frame = 1 if self.isolated_first_frame and self.frame_pre_padding == 0 else 0 + + frame_num = x.shape[2] + pad_size = (offset_frame - frame_num) % self.clip_length + padded_frame_num = frame_num + pad_size + num_chunks = (padded_frame_num - offset_frame) // self.clip_length + + z_list = [] + for i in range(num_chunks): + start_idx = i * self.clip_length + offset_frame + end_idx = (i + 1) * self.clip_length + offset_frame + clip_x = x[:, :, start_idx : min(end_idx, frame_num), :, :] + if end_idx > frame_num: + pad_frames = x[:, :, -1:].expand(-1, -1, end_idx - frame_num, -1, -1) + clip_x = torch.cat([clip_x, pad_frames], dim=2) + + if self.isolated_key_frame: + key_frame = clip_x[:, :, :1, :, :] + z_key = self._adaptive_encode(key_frame) + + if clip_x.shape[2] > 1: + video_frames = clip_x[:, :, 1:, :, :] + z_video = self._adaptive_encode(video_frames) + z = torch.cat([z_key, z_video], dim=2) + else: + z = z_key + else: + z = self._adaptive_encode(clip_x) + + z_list.append(z) + + z = torch.cat(z_list, dim=2) + if self.token_drop > 0: + z = z[:, :, : -self.token_drop] + + if self.isolated_first_frame: + input_first_frame = x[:, :, :1, :, :] + z_first_frame = self._adaptive_encode(input_first_frame) + + if self.frame_pre_padding == 0: + z = torch.cat([z_first_frame, z], dim=2) + else: + z = torch.cat([z_first_frame, z[:, :, 1:, :, :]], dim=2) + + if self.isolated_last_frame: + last_frame_idx = padded_frame_num - self.frame_drop + offset_frame + if last_frame_idx >= frame_num: + input_last_frame = x[:, :, -1:, :, :] + else: + input_last_frame = x[:, :, last_frame_idx : last_frame_idx + 1, :, :] + z_last_frame = self._adaptive_encode(input_last_frame) + z = torch.cat([z, z_last_frame], dim=2) + + return z + + def _decode_temporal_pad_frames(self, z, pad_tokens): + if pad_tokens <= 0: + return 0 + intra_tail = self.clip_length % self.vae_ratio_t + if intra_tail == 0: + return int(pad_tokens) * int(self.vae_ratio_t) + + z_len_before_pad = z.shape[2] - pad_tokens + return sum( + (intra_tail if (z_len_before_pad + k) % self.tokens_chunk_size == 0 else self.vae_ratio_t) + for k in range(pad_tokens) + ) + + def _decode_temporal_output_frame_plan(self, z, z_head, z_tail, num_chunks, pad_tokens): + chunk_dec = self.tokens_chunk_size * self.vae_ratio_t + split_count = int(self.token_drop > 0) + 1 + total_frames = 0 + final_overlap_frames = 0 + + if z_head is not None: + total_frames += 1 + + for i in range(num_chunks): + t_start_idx = i * self.tokens_chunk_size + t_end_idx = t_start_idx + self.tokens_chunk_size + self.token_overlap + clip_token_len = max(0, min(t_end_idx, z.shape[2]) - min(t_start_idx, z.shape[2])) + if i == 0 and z_head is not None: + clip_token_len += z_head.shape[2] + if i == num_chunks - 1 and z_tail is not None: + clip_token_len += z_tail.shape[2] + + clip_frame_len = clip_token_len * self.vae_ratio_t + if i == 0 and z_head is not None: + clip_frame_len = max(0, clip_frame_len - self.vae_ratio_t) + if i == num_chunks - 1 and z_tail is not None: + clip_frame_len = max(0, clip_frame_len - self.vae_ratio_t) + + for j in range(split_count): + f_start_idx = j * chunk_dec + f_end_idx = min(f_start_idx + chunk_dec, clip_frame_len) + chunk_frames = max(0, f_end_idx - f_start_idx - self.frame_pre_padding) + if j == 0: + total_frames += chunk_frames + else: + final_overlap_frames = chunk_frames + + total_frames += final_overlap_frames + if z_tail is not None: + total_frames += 1 + + pad_frames = self._decode_temporal_pad_frames(z, pad_tokens) + return int(total_frames), int(pad_frames), int(total_frames - pad_frames) + + def _decode_temporal_streaming(self, z, z_head, z_tail, num_chunks, pad_tokens, temporal_cat_dtype): + total_frames, pad_frames, output_frames = self._decode_temporal_output_frame_plan( + z, z_head, z_tail, num_chunks, pad_tokens + ) + if output_frames <= 0: + raise ValueError( + f"decode_temporal streaming planned non-positive output_frames={output_frames} " + f"total_frames={total_frames} pad_frames={pad_frames}" + ) + + chunk_dec = self.tokens_chunk_size * self.vae_ratio_t + split_count = int(self.token_drop > 0) + 1 + dec = None + dec_overlap = None + write_pos = 0 + logical_frames = 0 + dropped_frames = 0 + decoded_count = 0 + + def write_part(part): + nonlocal dec, write_pos, logical_frames, dropped_frames + part_frames = int(part.shape[2]) + if part_frames <= 0: + return + logical_frames += part_frames + if dec is None: + out_shape = list(part.shape) + out_shape[2] = output_frames + dec = torch.empty(out_shape, dtype=part.dtype, device=part.device) + + remaining = int(dec.shape[2]) - write_pos + copy_frames = min(part_frames, max(0, remaining)) + if copy_frames > 0: + dec[:, :, write_pos : write_pos + copy_frames, :, :].copy_(part[:, :, :copy_frames, :, :]) + write_pos += copy_frames + dropped_frames += part_frames - copy_frames + + for i in range(num_chunks): + t_start_idx = i * self.tokens_chunk_size + t_end_idx = t_start_idx + self.tokens_chunk_size + self.token_overlap + clip_z = z[:, :, t_start_idx:t_end_idx, :, :] + + if i == 0 and z_head is not None: + clip_z = torch.cat([z_head, clip_z], dim=2) + + if i == num_chunks - 1 and z_tail is not None: + clip_z = torch.cat([clip_z, z_tail], dim=2) + + clip_dec = self._adaptive_decode(clip_z) + decoded_count += 1 + if temporal_cat_dtype is not None and clip_dec.dtype != temporal_cat_dtype: + clip_dec = clip_dec.to(temporal_cat_dtype) + if clip_dec.device != z.device: + clip_dec = clip_dec.to(z.device) + + dec_tail = None + if i == 0 and z_head is not None: + write_part(clip_dec[:, :, self.vae_ratio_t - 1 : self.vae_ratio_t, :, :]) + clip_dec = clip_dec[:, :, self.vae_ratio_t :, :, :] + + if i == num_chunks - 1 and z_tail is not None: + dec_tail = clip_dec[:, :, -1:, :, :] + clip_dec = clip_dec[:, :, : -self.vae_ratio_t, :, :] + + for j in range(split_count): + f_start_idx = j * chunk_dec + f_end_idx = min(f_start_idx + chunk_dec, clip_dec.shape[2]) + clip_dec_chunk = clip_dec[:, :, f_start_idx:f_end_idx, :, :] + clip_dec_chunk = clip_dec_chunk[:, :, self.frame_pre_padding :, :, :] + + if j == 0: + if dec_overlap is not None: + clip_dec_chunk = self.blend(dec_overlap, clip_dec_chunk, self.frame_overlap, dim=-3) + dec_overlap = None + write_part(clip_dec_chunk) + else: + # Break the view's reference to the full decoded clip so earlier + # temporal chunks can be released before the final output exists. + dec_overlap = clip_dec_chunk.contiguous() + + if i == num_chunks - 1: + if dec_overlap is not None: + write_part(dec_overlap) + dec_overlap = None + if dec_tail is not None: + write_part(dec_tail) + + del clip_dec, clip_z + + if dec is None: + raise RuntimeError("decode_temporal streaming produced no output tensor") + if logical_frames != total_frames or dropped_frames != pad_frames or write_pos != output_frames: + raise RuntimeError( + "decode_temporal streaming frame plan mismatch: " + f"logical_frames={logical_frames} total_frames={total_frames} " + f"dropped_frames={dropped_frames} pad_frames={pad_frames} " + f"write_pos={write_pos} output_frames={output_frames}" + ) + + return dec + + def decode_temporal(self, z): + chunk_dec = self.tokens_chunk_size * self.vae_ratio_t + + isolated_token_num = 0 + if self.isolated_first_frame and self.frame_pre_padding == 0: + isolated_token_num = isolated_token_num + 1 + if self.isolated_last_frame: + isolated_token_num = isolated_token_num + 1 + + pseudo_total_tokens = z.shape[2] - isolated_token_num + self.token_drop + + pad_tokens = 0 + remainder = pseudo_total_tokens % self.tokens_chunk_size + if remainder != 0: + if self.training: + raise ValueError(f"Temporal token length {z.shape[2]} is wrong!") + else: + pad_tokens = self.tokens_chunk_size - remainder + pseudo_total_tokens = pseudo_total_tokens + pad_tokens + + pseudo_num_chunks = pseudo_total_tokens // self.tokens_chunk_size + num_chunks = pseudo_num_chunks - int(self.token_drop > 0) + + z_head = None + if self.isolated_first_frame and self.frame_pre_padding == 0: + z_head = z[:, :, :1, :, :] + z = z[:, :, 1:, :, :] + + z_tail = None + if self.isolated_last_frame: + z_tail = z[:, :, -1:, :, :] + z = z[:, :, :-1, :, :] + + if pad_tokens > 0: + pad_z = z[:, :, -1:, :, :].expand(-1, -1, pad_tokens, -1, -1) + z = torch.cat([z, pad_z], dim=2) + + temporal_cat_dtype = _resolve_temporal_cat_dtype() + if not self.training and _resolve_temporal_stream_cat(): + return self._decode_temporal_streaming(z, z_head, z_tail, num_chunks, pad_tokens, temporal_cat_dtype) + + decoded_tasks = [] + for i in range(num_chunks): + t_start_idx = i * self.tokens_chunk_size + t_end_idx = t_start_idx + self.tokens_chunk_size + self.token_overlap + clip_z = z[:, :, t_start_idx:t_end_idx, :, :] + + if i == 0 and z_head is not None: + clip_z = torch.cat([z_head, clip_z], dim=2) + + if i == num_chunks - 1 and z_tail is not None: + clip_z = torch.cat([clip_z, z_tail], dim=2) + + clip_dec = self._adaptive_decode(clip_z) + if temporal_cat_dtype is not None and clip_dec.dtype != temporal_cat_dtype: + clip_dec = clip_dec.to(temporal_cat_dtype) + + decoded_tasks.append((i, clip_dec)) + + clip_dec_list = [clip_dec.to(z.device) for _, clip_dec in decoded_tasks] + + dec_list = [] + dec_overlap = None + + dec_head = None + if z_head is not None: + dec_head = clip_dec_list[0][:, :, self.vae_ratio_t - 1 : self.vae_ratio_t, :, :] + clip_dec_list[0] = clip_dec_list[0][:, :, self.vae_ratio_t :, :, :] + + dec_tail = None + if z_tail is not None: + dec_tail = clip_dec_list[-1][:, :, -1:, :, :] + clip_dec_list[-1] = clip_dec_list[-1][:, :, : -self.vae_ratio_t, :, :] + + if dec_head is not None: + dec_list.append(dec_head) + + for i in range(num_chunks): + for j in range(int(self.token_drop > 0) + 1): + clip_dec = clip_dec_list[i] + + f_start_idx = j * chunk_dec + f_end_idx = min(f_start_idx + chunk_dec, clip_dec.shape[2]) + clip_dec_chunk = clip_dec[:, :, f_start_idx:f_end_idx, :, :] + clip_dec_chunk = clip_dec_chunk[:, :, self.frame_pre_padding :, :, :] + + if j == 0: + if dec_overlap is not None: + clip_dec_chunk = self.blend(dec_overlap, clip_dec_chunk, self.frame_overlap, dim=-3) + dec_list.append(clip_dec_chunk) + else: + dec_overlap = clip_dec_chunk + + if dec_overlap is not None: + dec_list.append(dec_overlap) + + if dec_tail is not None: + dec_list.append(dec_tail) + + dec = torch.cat(dec_list, dim=2) + + pad_frames = self._decode_temporal_pad_frames(z, pad_tokens) + if pad_frames > 0: + dec = dec[:, :, :-pad_frames, :, :] + + return dec + + def decode_base(self, z, frame_num=None, process_image=False): + if process_image or not self.use_3d_conv: + if not self.use_3d_conv and z.ndim == 5: + z = z.squeeze(2) + + recon = self._adaptive_decode(z) + else: + recon = self.decode_temporal(z) + + if self.use_3d_conv: + if frame_num is not None: + target_frames = frame_num + else: + target_frames = recon.shape[2] + + recon = self.trim_output(recon, target_frames) + if process_image: + recon = recon.squeeze(2) + + return recon + + ######################################################### + # following methods are for inference + ######################################################### + + @torch.no_grad() + def encode_images( + self, + images: Union[List[np.ndarray], List[torch.Tensor]], + transform_input: bool = False, + use_fp16_latent: bool = False, + verbose: bool = False, + ) -> List[torch.Tensor]: + """encode images into latents + + Args: + images (Union[List[np.ndarray], List[torch.Tensor]]): + List of images, single input will be wrapped in a list. + If input is a list of np.ndarray, it should be in shape B * (H, W, 3), dtype uint8. + If input is a list of torch.Tensor, it should be in shape B * (3, H, W), dtype float32. + transform_input (bool, optional): + Whether to transform input using ImageNet std/mean. Defaults to False. + If input is a list of np.ndarray, it will always be set to True. + use_fp16_latent (bool, optional): + Whether to use fp16 latent. Defaults to False. + verbose (bool, optional): + Whether to print debug information. Defaults to False. + + Returns: + List[torch.Tensor]: + List of image latents. + If self.use_3d_conv is True, it should be in shape B * (D, 1, H', W'). + Otherwise, it should be in shape B * (D, H', W'). + """ + + images = self.processor._ensure_list(images) + runtime_owned = False + + if isinstance(images[0], Image.Image): + images = [np.array(image) for image in images] + + if isinstance(images[0], np.ndarray): + device = next(self.parameters()).device + images = self.processor.convert_numpy_to_tensor(images, device) + images = torch.split(images, 1, dim=0) + transform_input = True + runtime_owned = True + + if transform_input: + images = [image.unsqueeze(0) if image.ndim == 3 else image for image in images] + images = [self.processor.transform_tensor(image, runtime_owned=runtime_owned) for image in images] + + prepared = [] + for image_tensor in images: + if image_tensor.ndim == 3: + image_tensor = image_tensor.unsqueeze(0) + _, _, h, w = image_tensor.shape + new_h, new_w = self.processor._align_to_total_patch_size(h, w) + image_tensor = self.processor._crop_to_align(image_tensor, new_h, new_w) + prepared.append(image_tensor) + + if len(prepared) > 1 and len(set(t.shape for t in prepared)) == 1: + stacked = torch.cat(prepared, dim=0) + if verbose: + logger.info(f"batch encode input shape {tuple(stacked.shape)}") + all_latents = self.encode_base(stacked, True) + image_latents = [all_latents[i].contiguous() for i in range(all_latents.shape[0])] + else: + image_latents = [] + for image_tensor in prepared: + if verbose: + logger.info(f"input shape {tuple(image_tensor.shape)}") + image_latent = self.encode_base(image_tensor, True) + image_latents.append(image_latent.squeeze(0).contiguous()) + + if use_fp16_latent: + image_latents = [lat.to(torch.float16) for lat in image_latents] + + if verbose: + for lat in image_latents: + logger.info(f"image latent shape {tuple(lat.shape)}") + + return image_latents + + @torch.no_grad() + def encode_videos( + self, + videos: Union[List[np.ndarray], List[torch.Tensor]], + transform_input: bool = False, + use_fp16_latent: bool = False, + verbose: bool = False, + encode_prefix: bool = False, + ) -> List[torch.Tensor]: + """encode videos into latents + + Args: + videos (Union[List[np.ndarray], List[torch.Tensor]]): + List of videos, single input will be wrapped in a list. + If input is a list of np.ndarray, it should be in shape B * (T, H, W, 3), dtype uint8. + If input is a list of torch.Tensor, it should be in shape B * (3, T, H, W), dtype float32. + transform_input (bool, optional): + Whether to transform input using ImageNet std/mean. Defaults to False. + If input is a list of np.ndarray, it will always be set to True. + use_fp16_latent (bool, optional): + Whether to use fp16 latent. Defaults to False. + verbose (bool, optional): + Whether to print debug information. Defaults to False. + encode_prefix (bool, optional): + Continuation (prefix) mode: prepend normalized + black frames to token alignment, append black frames to chunk + alignment, encode with token_drop disabled, then discard only + the trailing padding tokens. Returns both latents and leading + pad-frame counts. Defaults to False. + + Returns: + List[torch.Tensor]: + List of video latents, shape B * (D, T', H', W'). + With encode_prefix=True, returns + (List[torch.Tensor], List[int]). + """ + + videos = self.processor._ensure_list(videos) + runtime_owned = False + + if isinstance(videos[0], np.ndarray): + device = next(self.parameters()).device + videos = [self.processor.convert_numpy_to_tensor(video, device) for video in videos] + transform_input = True + runtime_owned = True + + if transform_input: + videos = [self.processor.transform_tensor(video, runtime_owned=runtime_owned) for video in videos] + videos = [video.transpose(0, 1) for video in videos] + + if encode_prefix: + if self.isolated_last_frame: + raise ValueError("encode_prefix does not support isolated_last_frame") + + video_latents = [] + prefix_pad_frames = [] + for video in videos: + if video.ndim == 4: + video = video.unsqueeze(0) + _, _, _, h, w = video.shape + new_h, new_w = self.processor._align_to_total_patch_size(h, w) + video = self.processor._crop_to_align(video, new_h, new_w, is_video=True) + + model_alignment = ( + self.token_drop, + self.frame_drop, + self.token_overlap, + self.frame_overlap, + ) + processor_alignment = ( + self.processor.token_overlap, + self.processor.frame_overlap, + ) + self.token_drop = 0 + self.frame_drop = 0 + self.token_overlap = 0 + self.frame_overlap = 0 + self.processor.token_overlap = 0 + self.processor.frame_overlap = 0 + try: + orig_frames = video.shape[2] + leading, trailing, drop_tokens = self.processor.align_video_length_2pass(orig_frames) + _, _, _, cropped_h, cropped_w = video.shape + if leading > 0: + black = self.processor.transform(video.new_zeros(leading, 3, cropped_h, cropped_w)) + black = black.unsqueeze(0).permute(0, 2, 1, 3, 4) + video = torch.cat([black, video], dim=2) + if trailing > 0: + black = self.processor.transform(video.new_zeros(trailing, 3, cropped_h, cropped_w)) + black = black.unsqueeze(0).permute(0, 2, 1, 3, 4) + video = torch.cat([video, black], dim=2) + + if verbose: + logger.info( + f"[encode_prefix] {orig_frames} frames -> " + f"pad leading={leading}, trailing={trailing} -> " + f"{video.shape[2]} frames" + ) + + video_latent = self.encode_base(video, False) + if drop_tokens > 0: + video_latent = video_latent[:, :, :-drop_tokens, :, :] + prefix_pad_frames.append(leading) + finally: + ( + self.token_drop, + self.frame_drop, + self.token_overlap, + self.frame_overlap, + ) = model_alignment + ( + self.processor.token_overlap, + self.processor.frame_overlap, + ) = processor_alignment + + video_latents.append(video_latent.squeeze(0).contiguous()) + + if use_fp16_latent: + video_latents = [lat.to(torch.float16) for lat in video_latents] + if verbose: + for latent in video_latents: + logger.info(f"video latent shape {tuple(latent.shape)}") + return video_latents, prefix_pad_frames + + prepared = [] + for video in videos: + if video.ndim == 4: + video = video.unsqueeze(0) + used_frame_length = self.processor.get_suitable_video_length(video.shape[2], verbose) + _, _, _, h, w = video.shape + new_h, new_w = self.processor._align_to_total_patch_size(h, w) + video = video[:, :, :used_frame_length, :, :] + video = self.processor._crop_to_align(video, new_h, new_w, is_video=True) + prepared.append(video) + + if len(prepared) > 1 and len(set(t.shape for t in prepared)) == 1: + stacked = torch.cat(prepared, dim=0) + if verbose: + logger.info(f"batch encode input shape {tuple(stacked.shape)}") + all_latents = self.encode_base(stacked, False) + video_latents = [all_latents[i].contiguous() for i in range(all_latents.shape[0])] + else: + video_latents = [] + for video in prepared: + if verbose: + logger.info(f"input shape {tuple(video.shape)}") + video_latent = self.encode_base(video, False) + video_latents.append(video_latent.squeeze(0).contiguous()) + + if use_fp16_latent: + video_latents = [lat.to(torch.float16) for lat in video_latents] + + if verbose: + for lat in video_latents: + logger.info(f"video latent shape {tuple(lat.shape)}") + + return video_latents + + +# ============================================================================ +# Legacy CNN VAE +# ============================================================================ + + +class AutoencoderKLLegacy(AutoencoderKL): + r""" + A VAE model (legacy CNN-based) for encoding pixels into latents and decoding latent representations into pixels. + """ + + @register_to_config + def __init__( + self, + in_channels=3, + out_ch=3, + ch=128, + embed_dim=16, + z_channels=16, + use_3d_conv=False, + # cnn vae + zq_ch_encoder=None, + zq_ch_decoder=None, + num_res_blocks=2, + num_res_blocks_decoder=None, + ch_mult=[1, 2, 2, 4, 4, 8], + space_down=[2, 2, 2, 2, 1, 1], + space_up=[1, 2, 2, 2, 2, 1], + time_down=None, + time_up=None, + padding_mode="zeros", + padding_mode_t=None, + use_t_isolated_gn=False, + causal_encoder=True, + causal_decoder=True, + use_vit_decoder=False, + vit_decoder_kwargs=None, + # stats + shift_factor=0.0, + scaling_factor=1.0, + # pixel normalization + pixel_norm_type="imagenet", + # others + **kwargs, + ): + ModelMixin.__init__(self) # NOTE: avoid wrong @register_to_config + + if not use_3d_conv or not use_vit_decoder: + raise NotImplementedError("this release only supports use_3d_conv=True with use_vit_decoder=True") + + self.transform = get_normalize_transform(pixel_norm_type) + self.transform_rev = get_denormalize_transform(pixel_norm_type) + + self.use_3d_conv = use_3d_conv + self.causal_encoder = causal_encoder + self.causal_decoder = causal_decoder + self.slidedec = self.causal_encoder and not self.causal_decoder + + # some registered parameters for simplicity + self.vae_ratio = int(np.cumprod(space_down)[-1]) + self.vae_ratio_t = int(np.cumprod(time_down)[-1]) if time_down else 1 + self.config["vae_ratio"] = self.vae_ratio + self.config["vae_ratio_t"] = self.vae_ratio_t + + # Configure inference-time chunking and tiling. + self.setup_forward(**kwargs) + + # init encoder + encoder_config = { + "double_z": True, + "z_channels": z_channels, + "zq_ch": zq_ch_encoder, + "in_channels": in_channels, + "ch": ch, + "num_res_blocks": num_res_blocks, + "ch_mult": ch_mult, + "space_down": space_down, + "time_down": time_down, + "padding_mode": padding_mode, + "padding_mode_t": padding_mode_t, + "causal": causal_encoder, + "use_t_isolated_gn": use_t_isolated_gn, + } + self.encoder = EncoderFCN3D(**encoder_config) + + # init pointwise quant/post_quant conv + self.quant_conv = nn.Conv3d(z_channels * 2, 2 * embed_dim, 1) + self.post_quant_conv = nn.Conv3d(embed_dim, z_channels, 1) + + self.use_vit_decoder = use_vit_decoder + + # init decoder + vit_kwargs = { + "patch_size": self.vae_ratio, + "in_channels": z_channels, + "out_channels": out_ch, + **(vit_decoder_kwargs or {}), + } + vit_kwargs.setdefault("patch_size_t", self.vae_ratio_t) + vit_kwargs.setdefault("t_causal", causal_decoder) + self.decoder = ViT3DDecoder(**vit_kwargs) + + @torch.no_grad() + def encode(self, x): + return self.quant_conv(self.encoder(x)) + + @torch.no_grad() + def decode(self, z): + z2 = self.post_quant_conv(z) + if self.use_vit_decoder: + return self.decoder(z2) + return self.decoder(z2, z) + + def encode_base(self, input, process_image=False): + if self.use_3d_conv and input.ndim == 4: + input = input.unsqueeze(2) + + if process_image or not self.use_3d_conv: + moments = self._adaptive_encode(input) + else: + moments = self.encode_temporal(input) + + z = DiagonalGaussianDistribution(moments).sample() + + if process_image and self.use_3d_conv: + z = self.trim_code(z, 1) + + return z diff --git a/telefuser/models/minimax_h3_video/norm.py b/telefuser/models/minimax_h3_video/norm.py new file mode 100644 index 0000000..f39a242 --- /dev/null +++ b/telefuser/models/minimax_h3_video/norm.py @@ -0,0 +1,247 @@ +# SPDX-License-Identifier: Apache-2.0 +# Torch-native normalization for the MiniMax H3 visual VAE. +import math + +import torch +import torch.nn as nn +import torch.nn.functional as F + +from .conv import BaseConv3d + + +def _validate_activation(activation): + valid_activations = {"identity", "silu", "relu"} + if activation not in valid_activations: + raise ValueError(f"Unsupported activation: {activation}. Supported: {valid_activations}") + + +def _apply_activation(x, activation): + _validate_activation(activation) + if activation == "identity": + return x + if activation == "silu": + return F.silu(x) + return F.relu(x) + + +def _merge_time_to_batch(x): + batch, channels, depth, height, width = x.shape + return x.permute(0, 2, 1, 3, 4).contiguous().view(batch * depth, channels, 1, height, width) + + +def _split_time_from_batch(x, batch): + batch_depth, channels, _, height, width = x.shape + depth = batch_depth // batch + return x.view(batch, depth, channels, height, width).permute(0, 2, 1, 3, 4).contiguous() + + +def fused_group_norm(x, num_groups, weight, bias, eps=1e-5, activation="silu"): + out = F.group_norm(x, num_groups, weight=weight, bias=bias, eps=eps) + return _apply_activation(out, activation) + + +def fused_spatial_norm( + f, + num_groups, + norm_weight, + norm_bias, + dynamic_scale, + dynamic_bias, + eps=1e-5, + activation="silu", +): + norm_f = F.group_norm( + f, + num_groups, + weight=norm_weight, + bias=norm_bias, + eps=eps, + ) + out = norm_f * dynamic_scale + dynamic_bias + return _apply_activation(out, activation) + + +class DummyAffine(torch.nn.Module): + def __init__(self, num_channels, affine=True): + super().__init__() + if affine: + self.weight = torch.nn.Parameter(torch.ones(num_channels)) + self.bias = torch.nn.Parameter(torch.zeros(num_channels)) + else: + self.register_parameter("weight", None) + self.register_parameter("bias", None) + + def forward(self, input): + if self.weight is None: + return input + shape = [1, -1] + [1] * (input.dim() - 2) + return input * self.weight.view(*shape) + self.bias.view(*shape) + + +class FusedGroupNorm3D(torch.nn.Module): + """Compatibility wrapper implemented with native PyTorch ops.""" + + def __init__( + self, + num_groups, + num_channels, + eps=1e-5, + affine=True, + activation="silu", + cond_channels=None, + use_t_isolated_gn=False, + padding_mode="zeros", + padding_mode_t=None, + causal=True, + ): + super().__init__() + _validate_activation(activation) + self.num_groups = num_groups + self.num_channels = num_channels + self.eps = eps + self.affine = affine + self.activation = activation + self.use_t_isolated_gn = use_t_isolated_gn + + if cond_channels is not None: + self.use_spatial_affine = True + self.norm_layer = DummyAffine(num_channels, affine=affine) + self.conv_y = BaseConv3d( + cond_channels, + num_channels, + kernel_size=1, + padding_mode=padding_mode, + padding_mode_t=padding_mode_t, + causal=causal, + ) + self.conv_b = BaseConv3d( + cond_channels, + num_channels, + kernel_size=1, + padding_mode=padding_mode, + padding_mode_t=padding_mode_t, + causal=causal, + ) + else: + self.use_spatial_affine = False + if self.affine: + self.weight = torch.nn.Parameter(torch.ones(num_channels)) + self.bias = torch.nn.Parameter(torch.zeros(num_channels)) + else: + self.register_parameter("weight", None) + self.register_parameter("bias", None) + + def forward(self, f, cond=None): + need_reshape = self.use_t_isolated_gn and f.dim() == 5 + batch = f.shape[0] if need_reshape else None + f_size = f.shape[-3:] + if need_reshape: + f = _merge_time_to_batch(f) + + if self.use_spatial_affine: + scale = self.conv_y(cond) + bias = self.conv_b(cond) + if math.prod(scale.shape[-3:]) * math.prod(bias.shape[-3:]) > 1: + scale = F.interpolate(scale, size=f_size, mode="nearest") + bias = F.interpolate(bias, size=f_size, mode="nearest") + if need_reshape: + scale = _merge_time_to_batch(scale) + bias = _merge_time_to_batch(bias) + out = fused_spatial_norm( + f, + self.num_groups, + self.norm_layer.weight, + self.norm_layer.bias, + scale, + bias, + self.eps, + self.activation, + ) + else: + if cond is not None: + raise NotImplementedError("Dynamic affine is not defined") + weight = self.weight if self.affine else None + bias = self.bias if self.affine else None + out = fused_group_norm(f, self.num_groups, weight, bias, self.eps, self.activation) + + if need_reshape: + out = _split_time_from_batch(out, batch) + return out + + +class TemporalIsolatedGroupNorm(nn.GroupNorm): + def forward(self, input): + if input.dim() == 5: + batch = input.shape[0] + input = _merge_time_to_batch(input) + output = super().forward(input) + return _split_time_from_batch(output, batch) + return super().forward(input) + + +class SpatialNorm3D(nn.Module): + def __init__( + self, + f_channels, + zq_channels, + padding_mode="zeros", + padding_mode_t=None, + causal=True, + use_t_isolated_gn=False, + ): + super().__init__() + norm_cls = TemporalIsolatedGroupNorm if use_t_isolated_gn else nn.GroupNorm + self.norm_layer = norm_cls(num_groups=32, num_channels=f_channels, eps=1e-6, affine=True) + + self.conv_y = BaseConv3d( + zq_channels, + f_channels, + kernel_size=1, + padding_mode=padding_mode, + padding_mode_t=padding_mode_t, + causal=causal, + ) + self.conv_b = BaseConv3d( + zq_channels, + f_channels, + kernel_size=1, + padding_mode=padding_mode, + padding_mode_t=padding_mode_t, + causal=causal, + ) + + def forward(self, f, zq): + f_size = f.shape[-3:] + norm_f = self.norm_layer(f) + scale = self.conv_y(zq) + bias = self.conv_b(zq) + + if math.prod(scale.shape[-3:]) * math.prod(bias.shape[-3:]) > 1: + scale = F.interpolate(scale, size=f_size, mode="nearest") + bias = F.interpolate(bias, size=f_size, mode="nearest") + + return norm_f * scale + bias + + +def get_spatial_norm_3d( + num_channels, + cond_channels, + *, + padding_mode="zeros", + padding_mode_t=None, + causal=True, + use_t_isolated_gn=False, +): + return SpatialNorm3D( + num_channels, + cond_channels, + padding_mode=padding_mode, + padding_mode_t=padding_mode_t, + causal=causal, + use_t_isolated_gn=use_t_isolated_gn, + ) + + +def get_group_norm_3d(num_channels, use_t_isolated_gn=False): + norm_cls = TemporalIsolatedGroupNorm if use_t_isolated_gn else nn.GroupNorm + return norm_cls(num_groups=32, num_channels=num_channels, eps=1e-6, affine=True) diff --git a/telefuser/models/minimax_h3_video/processor.py b/telefuser/models/minimax_h3_video/processor.py new file mode 100644 index 0000000..467d072 --- /dev/null +++ b/telefuser/models/minimax_h3_video/processor.py @@ -0,0 +1,254 @@ +# SPDX-License-Identifier: Apache-2.0 +# Tensor pre/post-processing for the MiniMax H3 visual VAE. +import math +from typing import Tuple + +import numpy as np +import torch +from diffusers.utils import logging +from einops import rearrange +from torchvision.transforms import Normalize + +NORM_CONFIGS = { + "imagenet": { + "mean": (0.485, 0.456, 0.406), + "std": (0.229, 0.224, 0.225), + }, + "simple": { + "mean": (0.5, 0.5, 0.5), + "std": (0.5, 0.5, 0.5), + }, + "raw": { + "mean": (0.0, 0.0, 0.0), + "std": (1.0, 1.0, 1.0), + }, +} + + +def get_norm_constants( + norm_type: str = "imagenet", +) -> Tuple[Tuple[float, ...], Tuple[float, ...]]: + if norm_type not in NORM_CONFIGS: + raise ValueError(f"Unknown norm_type: {norm_type}. Must be one of {list(NORM_CONFIGS.keys())}") + config = NORM_CONFIGS[norm_type] + return config["mean"], config["std"] + + +def get_normalize_transform(norm_type: str = "imagenet", *, inplace: bool = False) -> Normalize: + mean, std = get_norm_constants(norm_type) + return Normalize(mean, std, inplace=inplace) + + +def get_denormalize_transform(norm_type: str = "imagenet") -> Normalize: + mean, std = get_norm_constants(norm_type) + inv_mean = tuple(-m / s for m, s in zip(mean, std)) + inv_std = tuple(1.0 / s for s in std) + return Normalize(inv_mean, inv_std) + + +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._runtime_owned_transform = ( + get_normalize_transform(pixel_norm_type, inplace=True) if transform is None else None + ) + 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, *, runtime_owned=False): + 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}") + + transform = ( + self._runtime_owned_transform + if runtime_owned and self._runtime_owned_transform is not None + else self.transform + ) + tensor = 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) + tensor = torch.from_numpy(numpy_array) + # Keep decoded uint8 pixels compact across the host-to-device copy. + # Casting the full video on CPU quadruples both the temporary host + # allocation and transfer volume for no loss of information. + if device is not None: + tensor = tensor.to(device) + tensor = tensor.permute(0, 3, 1, 2) + return tensor.to(torch.float32).div_(255.0) diff --git a/telefuser/models/minimax_h3_video/vae_cnn.py b/telefuser/models/minimax_h3_video/vae_cnn.py new file mode 100644 index 0000000..6c40907 --- /dev/null +++ b/telefuser/models/minimax_h3_video/vae_cnn.py @@ -0,0 +1,262 @@ +# SPDX-License-Identifier: Apache-2.0 +# 3D causal CNN encoder for the MiniMax H3 visual VAE (inference-only bundle). +import torch.nn as nn +import torch.nn.functional as F + +from .conv import BaseConv3d +from .norm import get_group_norm_3d, get_spatial_norm_3d + +# ============================================================================ +# 3D CNN Components +# ============================================================================ + + +def norm_silu(x, norm, cond=None): + if cond is None: + return F.silu(norm(x), inplace=True) + else: + return F.silu(norm(x, cond), inplace=True) + + +class Downsample3D(nn.Module): + def __init__( + self, + in_channels, + out_channels, + time_stride=1, + space_stride=2, + padding_mode="zeros", + padding_mode_t=None, + causal=True, + ): + super().__init__() + self.time_stride = time_stride + self.space_stride = space_stride + + assert time_stride in [1, 2] + assert space_stride in [1, 2, 3] + + self.conv = BaseConv3d( + in_channels, + out_channels, + kernel_size=3, + padding=(1, 0, 0), + stride=(time_stride, space_stride, space_stride), + padding_mode=padding_mode, + padding_mode_t=padding_mode_t, + causal=causal, + ) + self.causal = self.conv.causal + self.pad_mode = self.conv.pad_mode + + def forward(self, x): + if self.space_stride == 2: + pad = (0, 1, 0, 1, 0, 0) + x = F.pad(x, pad, mode=self.pad_mode) + return self.conv(x) + + +class ResnetBlock3D(nn.Module): + def __init__( + self, + in_channels, + out_channels=None, + zq_ch=None, + padding_mode="zeros", + padding_mode_t=None, + causal=True, + use_t_isolated_gn=False, + ): + super().__init__() + self.in_channels = in_channels + out_channels = in_channels if out_channels is None else out_channels + self.out_channels = out_channels + + self.use_fused_norm = False + + if zq_ch is None: + self.norm1 = get_group_norm_3d(in_channels, use_t_isolated_gn=use_t_isolated_gn) + self.norm2 = get_group_norm_3d(out_channels, use_t_isolated_gn=use_t_isolated_gn) + else: + self.norm1 = get_spatial_norm_3d( + in_channels, + zq_ch, + padding_mode=padding_mode, + padding_mode_t=padding_mode_t, + causal=causal, + use_t_isolated_gn=use_t_isolated_gn, + ) + self.norm2 = get_spatial_norm_3d( + out_channels, + zq_ch, + padding_mode=padding_mode, + padding_mode_t=padding_mode_t, + causal=causal, + use_t_isolated_gn=use_t_isolated_gn, + ) + + self.conv1 = BaseConv3d( + in_channels, + out_channels, + kernel_size=3, + padding=1, + padding_mode=padding_mode, + padding_mode_t=padding_mode_t, + causal=causal, + ) + + self.conv2 = BaseConv3d( + out_channels, + out_channels, + kernel_size=3, + padding=1, + padding_mode=padding_mode, + padding_mode_t=padding_mode_t, + causal=causal, + ) + + if self.in_channels != self.out_channels: + self.nin_shortcut = BaseConv3d( + in_channels, + out_channels, + kernel_size=1, + padding_mode=padding_mode, + padding_mode_t=padding_mode_t, + causal=causal, + ) + + def forward(self, x, zq=None): + h = x + + if self.use_fused_norm: + h = self.norm1(h, zq) + else: + h = norm_silu(h, self.norm1, zq) + + h = self.conv1(h) + + if self.use_fused_norm: + h = self.norm2(h, zq) + else: + h = norm_silu(h, self.norm2, zq) + + h = self.conv2(h) + + if self.in_channels != self.out_channels: + x = self.nin_shortcut(x) + + return h.add_(x) + + +class EncoderFCN3D(nn.Module): + def __init__( + self, + ch, + ch_mult, + space_down, + time_down, + num_res_blocks, + in_channels, + z_channels, + double_z=False, + zq_ch=None, + padding_mode="zeros", + padding_mode_t=None, + causal=True, + use_t_isolated_gn=False, + ): + super().__init__() + self.ch = ch + self.num_levels = len(ch_mult) + + if isinstance(num_res_blocks, int): + self.num_res_blocks = [num_res_blocks] * self.num_levels + else: + self.num_res_blocks = num_res_blocks + + self.space_down_factors = space_down + self.time_down_factors = time_down + self.in_channels = in_channels + + self.use_fused_norm = False + + block_mid = [ch * ch_mult[i] for i in range(self.num_levels)] + block_in = [block_mid[0]] + block_mid[:-1] + block_out = block_mid + + conv_kwargs = dict( + padding_mode=padding_mode, + padding_mode_t=padding_mode_t, + causal=causal, + ) + + self.conv_in = BaseConv3d(in_channels, block_in[0], kernel_size=3, padding=1, **conv_kwargs) + + self.down = nn.ModuleList() + for i_level in range(self.num_levels): + down = nn.Module() + + down.block = nn.ModuleList() + for i in range(self.num_res_blocks[i_level]): + down.block.append( + ResnetBlock3D( + in_channels=block_in[i_level] if i == 0 else block_mid[i_level], + out_channels=block_mid[i_level], + zq_ch=zq_ch, + use_t_isolated_gn=use_t_isolated_gn, + **conv_kwargs, + ) + ) + + if space_down[i_level] * time_down[i_level] > 1: + down.downsample = Downsample3D( + block_mid[i_level], + block_out[i_level], + time_stride=time_down[i_level], + space_stride=space_down[i_level], + **conv_kwargs, + ) + else: + if block_out[i_level] != block_mid[i_level]: + down.downsample = BaseConv3d( + block_mid[i_level], + block_out[i_level], + kernel_size=1, + **conv_kwargs, + ) + + self.down.append(down) + + if zq_ch is None: + self.norm_out = get_group_norm_3d(block_out[-1], use_t_isolated_gn=use_t_isolated_gn) + else: + self.norm_out = get_spatial_norm_3d( + block_out[-1], + zq_ch, + use_t_isolated_gn=use_t_isolated_gn, + **conv_kwargs, + ) + + self.conv_out = BaseConv3d( + block_out[-1], + 2 * z_channels if double_z else z_channels, + kernel_size=3, + padding=1, + **conv_kwargs, + ) + + def forward(self, x, zq=None): + h = self.conv_in(x) + for i_level in range(self.num_levels): + for i_block in range(self.num_res_blocks[i_level]): + h = self.down[i_level].block[i_block](h, zq) + if hasattr(self.down[i_level], "downsample"): + h = self.down[i_level].downsample(h) + + if self.use_fused_norm: + h = self.norm_out(h, zq) + else: + h = norm_silu(h, self.norm_out, zq) + + h = self.conv_out(h) + return h diff --git a/telefuser/models/minimax_h3_video/vae_vit.py b/telefuser/models/minimax_h3_video/vae_vit.py new file mode 100644 index 0000000..c05bc7e --- /dev/null +++ b/telefuser/models/minimax_h3_video/vae_vit.py @@ -0,0 +1,337 @@ +# SPDX-License-Identifier: Apache-2.0 +# ViT3D decoder for the MiniMax H3 visual VAE (inference-only bundle). +import torch +import torch.distributed as dist +import torch.nn as nn +from diffusers.configuration_utils import ConfigMixin, register_to_config +from diffusers.models.modeling_utils import ModelMixin +from diffusers.utils import logging + +from .base_module import RotaryEmbeddingND, TransformerBlock +from .flash import make_block_causal_mask_mod +from .vit_utils import create_token_ids, prepare_rotary_pos_emb + +logger = logging.get_logger(__name__) + + +def _linear_with_module_dtype(linear, tensor, out_dtype=None): + weight = getattr(linear, "weight", None) + target_dtype = getattr(weight, "dtype", tensor.dtype) + output = linear(tensor.to(target_dtype)) + if out_dtype is not None and output.dtype != out_dtype: + output = output.to(out_dtype) + return output + + +def _pack_tensors_3d(tensors, patch_size, patch_size_t): + batch_size, num_channels_tensors, temporal, height, width = tensors.shape + + tensors = tensors.view( + batch_size, + num_channels_tensors, + temporal // patch_size_t, + patch_size_t, + height // patch_size, + patch_size, + width // patch_size, + patch_size, + ) + tensors = tensors.permute(0, 2, 4, 6, 1, 3, 5, 7) + tensors = tensors.reshape( + batch_size, + (temporal // patch_size_t) * (height // patch_size) * (width // patch_size), + num_channels_tensors * patch_size_t * patch_size * patch_size, + ) + return tensors + + +def _unpack_tensors_3d(tensors, patch_size, patch_size_t, temporal, height, width): + batch_size, num_patches, channels = tensors.shape + num_channels_tensors = channels // (patch_size_t * patch_size * patch_size) + + tensors = tensors.view( + batch_size, + temporal // patch_size_t, + height // patch_size, + width // patch_size, + num_channels_tensors, + patch_size_t, + patch_size, + patch_size, + ) + tensors = tensors.permute(0, 4, 1, 5, 2, 6, 3, 7).contiguous() + tensors = tensors.reshape(batch_size, num_channels_tensors, temporal, height, width) + return tensors + + +class ViTBase(ModelMixin, ConfigMixin): + """Base class for ViT Encoder and Decoder with common functionality.""" + + _no_split_modules = ["TransformerBlock"] + + def _init_weights(self): + def basic_init(m): + if isinstance(m, nn.Linear): + nn.init.xavier_uniform_(m.weight) + if m.bias is not None: + nn.init.constant_(m.bias, 0) + + self.apply(basic_init) + + def init_mask_config(self, dim, is_3d=False): + self._mask_dim = dim + self._mask_is_3d = is_3d + self.register_buffer("mask_token", torch.zeros(1, 1, dim)) + + def set_mask_config(self, mask_config): + self.mask_prob = mask_config.get("mask_prob", 0.0) + self.mask_enabled = self.mask_prob > 0 + self.mask_style = mask_config.get("mask_style", "replace") + if self.mask_enabled and self.mask_style == "drop" and self.mask_prob < 1.0: + logger.warning("mask_style='drop' with mask_prob < 1.0") + if self._mask_is_3d: + self.temporal_scale_range = mask_config.get("temporal_scale_range", (0.3, 0.5)) + self.spatial_scale_range = mask_config.get("spatial_scale_range", (0.1, 0.25)) + self.min_mask_ratio = mask_config.get("min_mask_ratio", 0.75) + self.max_mask_ratio = mask_config.get("max_mask_ratio", 0.95) + else: + self.spatial_scale_range = mask_config.get("spatial_scale_range", (0.15, 0.15)) + self.min_mask_ratio = mask_config.get("min_mask_ratio", 0.5) + self.max_mask_ratio = mask_config.get("max_mask_ratio", 0.75) + self.aspect_ratio_range = mask_config.get("aspect_ratio_range", (0.75, 1.5)) + self.max_retries = mask_config.get("max_retries", 100) + if self.mask_enabled and self.mask_style == "drop" and getattr(self, "t_causal", False): + logger.warning("mask_style='drop' with t_causal may cause issues") + if self.mask_enabled and "mask_token" in self._buffers: + del self._buffers["mask_token"] + self.mask_token = nn.Parameter(torch.randn(1, 1, self._mask_dim) * 0.02) + + def init_suffix_tokens(self, dim, num_register_tokens, has_cls_token=True): + self.num_register_tokens = num_register_tokens + if num_register_tokens > 0: + self.register_tokens = nn.Parameter(torch.randn(1, num_register_tokens, dim) * 0.02) + else: + self.register_tokens = None + if has_cls_token: + self.cls_token = nn.Parameter(torch.randn(1, 1, dim) * 0.02) + + def apply_mask_preprocess(self, hidden_states, img_ids, patch_dims, num_suffix): + if self.training and self.mask_enabled: + raise NotImplementedError("mask modeling is not supported in this inference-only bundle") + return hidden_states, img_ids + + def forward_transformer_blocks(self, hidden_states, rotary_pos_emb, pack_info=None): + if pack_info is None: + pack_info = {} + for block in self.transformer_blocks: + hidden_states = block(hidden_states, rotary_pos_emb, pack_info) + return hidden_states + + def apply_mask_postprocess(self, hidden_states, num_patches): + if self.training and self.mask_enabled and self.mask_style == "drop": + raise NotImplementedError("mask modeling is not supported in this inference-only bundle") + return hidden_states + + +class ViT3DDecoder(ViTBase): + """Vision Transformer Video Decoder using TransformerBlock.""" + + @register_to_config + def __init__( + self, + patch_size: int = 16, + patch_size_t: int = 4, + t_causal: bool = False, + in_channels: int = 16, + out_channels: int = 3, + num_layers: int = 24, + heads: int = 16, + dim_head: int = 64, + norm_type: str = "layer_norm", + norm_affine: bool = True, + qk_norm_type: str = None, + qk_norm_affine: bool = False, + ffn_activation_fn: str = "gelu", + ffn_use_gated: bool = False, + rope_theta: float = 100.0, + rope_dim_ratio: float = 1.0, + bias: bool = True, + eps: float = 1e-5, + num_register_tokens: int = 4, + mask_config: dict = {}, + **kwargs, + ): + super().__init__() + + dim = heads * dim_head + rope_apply_dim = int(dim_head * rope_dim_ratio) + + self.pos_embed = RotaryEmbeddingND(rope_apply_dim, rope_theta, n_dim=3, use_angle=True) + + self.x_embedder = nn.Linear(in_channels, dim) + + self.init_suffix_tokens(dim, num_register_tokens, has_cls_token=False) + + self.t_causal = t_causal + + self.transformer_blocks = nn.ModuleList( + [ + TransformerBlock( + heads=heads, + dim_head=dim_head, + norm_type=norm_type, + norm_affine=norm_affine, + qk_norm_type=qk_norm_type, + qk_norm_affine=qk_norm_affine, + ffn_activation_fn=ffn_activation_fn, + ffn_use_gated=ffn_use_gated, + bias=bias, + eps=eps, + **kwargs, + ) + for _ in range(num_layers) + ] + ) + + self.norm_out = nn.LayerNorm(dim, elementwise_affine=norm_affine, eps=eps) + patch_dim = out_channels * patch_size_t * patch_size * patch_size + self.proj_out = nn.Linear(dim, patch_dim) + + self.init_mask_config(dim, is_3d=True) + self.set_mask_config(mask_config) + + self._rotary_pos_emb_cache = None + self._autocast_linear_dtype = None + + if len(kwargs) > 0 and (not dist.is_initialized() or dist.get_rank() == 0): + logger.warning(f"Unused kwargs: {kwargs}") + + def _apply(self, fn, recurse=True): + result = super()._apply(fn, recurse=recurse) + self._rotary_pos_emb_cache = None + self._autocast_linear_dtype = None + return result + + def prepare_autocast_linear_weights(self, dtype: torch.dtype) -> int: + """Keep decoder-block linear weights in their autocast compute dtype. + + PyTorch autocast does not cache casts for these frozen parameters, so + tiled decode otherwise converts every FP32 weight and bias once per + block invocation. Persisting the rounded values is numerically + equivalent to the per-call autocast conversion. The embedding and + output projections stay FP32 because their calls explicitly disable + autocast. + """ + + if dtype not in (torch.float16, torch.bfloat16): + raise ValueError(f"MiniMax H3 decoder autocast weights require fp16 or bf16, got {dtype}") + if self._autocast_linear_dtype == dtype: + return 0 + + converted = 0 + for block in self.transformer_blocks: + for linear in ( + block.attn.to_qkv, + block.attn.to_out, + block.ff.w1, + block.ff.w2, + ): + if linear.weight.dtype != dtype: + linear.to(dtype=dtype) + converted += 1 + self._autocast_linear_dtype = dtype + return converted + + def forward(self, x: torch.Tensor) -> torch.Tensor: + B, C, latent_T, latent_H, latent_W = x.shape + patch_size = self.config.patch_size + patch_size_t = self.config.patch_size_t + num_suffix = 1 + self.num_register_tokens + + hidden_states = _pack_tensors_3d(x, 1, 1) + latent_size = (latent_T, latent_H, latent_W) + + with torch.autocast("cuda", enabled=False): + hidden_states = _linear_with_module_dtype(self.x_embedder, hidden_states, hidden_states.dtype) + + num_patches = hidden_states.shape[1] + + tokens = [hidden_states] + + if self.register_tokens is not None: + register_tokens = self.register_tokens.expand(B, -1, -1) + tokens.append(register_tokens) + + cls_token = torch.zeros_like(hidden_states[:, 0:1, :]) + tokens.append(cls_token) + hidden_states = torch.cat(tokens, dim=1) + + patch_dims = [latent_T, latent_H, latent_W] + rotary_dtype = ( + torch.get_autocast_dtype("cuda") if x.is_cuda and torch.is_autocast_enabled("cuda") else hidden_states.dtype + ) + cache_enabled = not self.training and not self.mask_enabled and not torch.compiler.is_compiling() + cache_key = ( + B, + latent_T, + latent_H, + latent_W, + num_suffix, + x.device, + x.dtype, + rotary_dtype, + ) + cache_record = self._rotary_pos_emb_cache if cache_enabled else None + cache_hit = cache_record is not None and cache_record[0] == cache_key + if cache_hit: + img_ids = cache_record[1] + else: + img_ids = create_token_ids(latent_size, x.device, x.dtype).expand(B, -1, -1) + suffix_ids = torch.zeros((B, num_suffix, 3), device=x.device, dtype=img_ids.dtype) + img_ids = torch.cat([img_ids, suffix_ids], dim=1) + + hidden_states, img_ids = self.apply_mask_preprocess(hidden_states, img_ids, patch_dims, num_suffix) + cache_img_ids = img_ids + + pack_info = {} + if self.t_causal: + spatial_size = latent_H * latent_W + mask_mod = make_block_causal_mask_mod( + num_tokens=num_patches, + block_size=spatial_size, + suffix=True, + ) + pack_info["mask_mod"] = mask_mod + + if cache_hit: + rotary_pos_emb = cache_record[2] + else: + rotary_pos_emb = prepare_rotary_pos_emb( + self.pos_embed(img_ids), + dtype=rotary_dtype, + ) + if cache_enabled: + self._rotary_pos_emb_cache = ( + cache_key, + cache_img_ids, + rotary_pos_emb, + ) + + for block in self.transformer_blocks: + hidden_states = block(hidden_states, rotary_pos_emb, pack_info) + + hidden_states = self.norm_out(hidden_states) + + hidden_states = self.apply_mask_postprocess(hidden_states, num_patches) + + with torch.autocast("cuda", enabled=False): + output = _linear_with_module_dtype(self.proj_out, hidden_states, hidden_states.dtype) + + output = output[:, :num_patches, :] + + video_t = latent_size[0] * patch_size_t + video_h = latent_size[1] * patch_size + video_w = latent_size[2] * patch_size + output = _unpack_tensors_3d(output, patch_size, patch_size_t, video_t, video_h, video_w) + + return output diff --git a/telefuser/models/minimax_h3_video/vit_utils.py b/telefuser/models/minimax_h3_video/vit_utils.py new file mode 100644 index 0000000..95b98fe --- /dev/null +++ b/telefuser/models/minimax_h3_video/vit_utils.py @@ -0,0 +1,100 @@ +# SPDX-License-Identifier: Apache-2.0 +# ViT runtime helpers for the MiniMax H3 visual VAE. +from collections.abc import Sequence +from typing import Tuple + +import torch + + +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 _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[:2] + + 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:] + scaled = t_rot * cos + scaled.add_(_rotate_half(t_rot) * sin) + t_rot = scaled + t = torch.cat((t_rot, t_pass), dim=-1) + else: + scaled = t * cos + scaled.add_(_rotate_half(t) * sin) + t = scaled + + return t + + +def prepare_rotary_pos_emb( + rotary_pos_emb: Tuple[torch.Tensor, torch.Tensor], + *, + dtype: torch.dtype, +) -> tuple[torch.Tensor, ...]: + """Cast-free eager rotary cache used by the parity path.""" + cos, sin = rotary_pos_emb + del dtype + return cos, sin + + +def apply_rotary_pos_emb(t: torch.Tensor, rotary_pos_emb: Sequence[torch.Tensor]) -> torch.Tensor: + return _apply_rotary_pos_emb_impl(t, rotary_pos_emb) + + +def apply_rotary_pos_emb_qk( + query: torch.Tensor, + key: torch.Tensor, + rotary_pos_emb: Sequence[torch.Tensor], +) -> tuple[torch.Tensor, torch.Tensor]: + """Apply the released NeoX rotary recipe to Q and K.""" + return ( + apply_rotary_pos_emb(query, rotary_pos_emb), + apply_rotary_pos_emb(key, rotary_pos_emb), + ) diff --git a/telefuser/models/minimax_h3_video_vae.py b/telefuser/models/minimax_h3_video_vae.py new file mode 100644 index 0000000..2c91704 --- /dev/null +++ b/telefuser/models/minimax_h3_video_vae.py @@ -0,0 +1,152 @@ +# SPDX-License-Identifier: Apache-2.0 +"""MiniMax H3 f16/t4/d24 visual VAE.""" + +from __future__ import annotations + +import json +from dataclasses import dataclass +from pathlib import Path +from typing import Any + +import torch + +from telefuser.core.base_model import BaseModel + +from .minimax_h3_video import AutoencoderKLLegacy + + +@dataclass(frozen=True) +class MiniMaxH3VideoVAEConfig: + architecture: dict[str, Any] + clip_length: int + token_drop: int + encoder_tiling: bool + decoder_tiling: bool + tile_size: int + tile_overlap_min: int + chunk_dim: int + latent_channels: int + latents_mean: tuple[float, ...] + latents_std: tuple[float, ...] + + @classmethod + def from_path(cls, path: str | Path) -> MiniMaxH3VideoVAEConfig: + component_dir = Path(path) + if component_dir.is_file(): + component_dir = component_dir.parent + component = json.loads((component_dir / "config.json").read_text(encoding="utf-8")) + source_dir = component_dir / component["source_path"] + architecture = json.loads((source_dir / "config.json").read_text(encoding="utf-8")) + config = cls( + architecture=architecture, + clip_length=int(component["vae_clip_length"]), + token_drop=int(component["vae_token_drop"]), + encoder_tiling=bool(component["vae_encoder_tiling"]), + decoder_tiling=bool(component["vae_decoder_tiling"]), + tile_size=int(component["vae_tile_size"]), + tile_overlap_min=int(component["vae_tile_overlap_min"]), + chunk_dim=int(component["vae_chunk_dim"]), + latent_channels=int(component["latent_channels"]), + latents_mean=tuple(float(value) for value in component["latents_mean"]), + latents_std=tuple(float(value) for value in component["latents_std"]), + ) + config.validate() + return config + + def validate(self) -> None: + spatial_ratio = int(self.architecture.get("vae_ratio", 0)) + temporal_ratio = int(self.architecture.get("vae_ratio_t", 0)) + if (spatial_ratio, temporal_ratio, self.latent_channels) != (16, 4, 24): + raise ValueError( + "MiniMax H3 visual VAE requires f16/t4/d24 geometry, got " + f"f{spatial_ratio}/t{temporal_ratio}/d{self.latent_channels}" + ) + if int(self.architecture.get("embed_dim", 0)) != self.latent_channels: + raise ValueError("visual VAE embed_dim must equal latent_channels") + if self.clip_length != 17 or self.token_drop != 3: + raise ValueError("MiniMax H3 visual VAE requires clip_length=17 and token_drop=3") + if len(self.latents_mean) != self.latent_channels or len(self.latents_std) != self.latent_channels: + raise ValueError("visual VAE latent statistics must contain one value per channel") + if any(value <= 0 for value in self.latents_std): + raise ValueError("visual VAE latent standard deviations must be positive") + + def model_kwargs(self) -> dict[str, Any]: + ignored = {"_class_name", "_diffusers_version", "vae_ratio", "vae_ratio_t"} + kwargs = {key: value for key, value in self.architecture.items() if key not in ignored} + kwargs.update( + { + "clip_length": self.clip_length, + "token_drop": self.token_drop, + "encoder_tiling": self.encoder_tiling, + "decoder_tiling": self.decoder_tiling, + "parallel_tiling": False, + "tile_size": self.tile_size, + "tile_overlap_min": self.tile_overlap_min, + "encoder_parallel": False, + "decoder_parallel": False, + "chunk_dim": self.chunk_dim, + } + ) + return kwargs + + +class MiniMaxH3VideoVAE(BaseModel): + """Checkpoint-backed 3D CNN encoder and ViT decoder.""" + + def __init__(self, config: MiniMaxH3VideoVAEConfig) -> None: + super().__init__() + config.validate() + self.model = AutoencoderKLLegacy(**config.model_kwargs()) + self.config = config + self.layer_name_list = ["model"] + + @property + def processor(self) -> Any: + return self.model.processor + + @torch.no_grad() + def encode_images(self, *args: Any, **kwargs: Any) -> list[torch.Tensor]: + return self.model.encode_images(*args, **kwargs) + + @torch.no_grad() + def encode_videos(self, *args: Any, **kwargs: Any) -> list[torch.Tensor]: + return self.model.encode_videos(*args, **kwargs) + + @torch.no_grad() + def decode_base(self, *args: Any, **kwargs: Any) -> torch.Tensor: + return self.model.decode_base(*args, **kwargs) + + def prepare_decoder_autocast_weights(self, dtype: torch.dtype) -> int: + return self.model.decoder.prepare_autocast_linear_weights(dtype) + + @torch.no_grad() + def decode_normalized(self, latent: torch.Tensor) -> torch.Tensor: + if latent.ndim != 5 or latent.shape[1] != self.config.latent_channels: + raise ValueError(f"visual latent must be [B, {self.config.latent_channels}, T, H, W]") + mean = latent.new_tensor(self.config.latents_mean).view(1, -1, 1, 1, 1) + std = latent.new_tensor(self.config.latents_std).view(1, -1, 1, 1, 1) + frames = self.model.decode_base(latent.mul(std).add(mean)) + return self.model.processor.revert_tensor(frames) + + @staticmethod + def state_dict_converter(config_path: str | Path) -> MiniMaxH3VideoVAEStateDictConverter: + return MiniMaxH3VideoVAEStateDictConverter(config_path) + + +class MiniMaxH3VideoVAEStateDictConverter: + def __init__(self, config_path: str | Path) -> None: + self.config = MiniMaxH3VideoVAEConfig.from_path(config_path) + + def from_official(self, state_dict: dict[str, torch.Tensor]) -> tuple[dict[str, torch.Tensor], dict[str, Any]]: + converted = {f"model.{name}": value for name, value in state_dict.items()} + return converted, {"config": self.config} + + def from_diffusers(self, state_dict: dict[str, torch.Tensor]) -> tuple[dict[str, torch.Tensor], dict[str, Any]]: + return self.from_official(state_dict) + + +__all__ = [ + "MiniMaxH3VideoVAE", + "MiniMaxH3VideoVAEConfig", + "MiniMaxH3VideoVAEStateDictConverter", +] diff --git a/tests/unit/models/test_minimax_h3_audio_vae.py b/tests/unit/models/test_minimax_h3_audio_vae.py new file mode 100644 index 0000000..fe9b238 --- /dev/null +++ b/tests/unit/models/test_minimax_h3_audio_vae.py @@ -0,0 +1,101 @@ +import json +from pathlib import Path + +import pytest +import torch + +from telefuser.models.minimax_h3_audio_vae import ( + MiniMaxH3AudioVAE, + MiniMaxH3AudioVAEConfig, + MiniMaxH3AudioVAEStateDictConverter, +) + + +def _config() -> MiniMaxH3AudioVAEConfig: + return MiniMaxH3AudioVAEConfig( + encoder_dim=8, + encoder_rates=(2,), + latent_dim=16, + decoder_dim=32, + decoder_rates=(2,), + sample_rate=32_000, + latent_channels=8, + output_channels=2, + attn_proj=True, + decoder_type="bigvgan", + latents_mean=(0.0,) * 8, + latents_std=(1.0,) * 8, + ) + + +def test_audio_vae_config_reads_released_component(tmp_path: Path) -> None: + (tmp_path / "config.json").write_text( + json.dumps( + { + "source_metadata_path": "metadata.json", + "output_channel": 2, + "latents_mean": [0.0] * 32, + "latents_std": [1.0] * 32, + } + ), + encoding="utf-8", + ) + (tmp_path / "metadata.json").write_text( + '{"metadata":{"kwargs":{"encoder_dim":64,"encoder_rates":[2,4,4,5,5],' + '"latent_dim":2048,"decoder_dim":1024,"decoder_rates":[5,5,2,2,2,2,2],' + '"sample_rate":32000,"vae_latent_channels":32,"attn_proj":true,' + '"decoder_type":"bigvgan"}}}', + encoding="utf-8", + ) + config = MiniMaxH3AudioVAEConfig.from_path(tmp_path) + assert config.sample_rate == 32_000 + assert config.encoder_rates == (2, 4, 4, 5, 5) + assert config.output_channels == 2 + + +def test_audio_vae_config_rejects_non_h3_latent_width() -> None: + config = _config() + with pytest.raises(ValueError, match="32 latent channels"): + config.validate() + + +def test_audio_vae_converter_maps_legacy_weight_norm_keys() -> None: + converter = MiniMaxH3AudioVAEStateDictConverter.__new__(MiniMaxH3AudioVAEStateDictConverter) + converter.config = _config() + state = { + "encoder.block.0.bias": torch.zeros(1), + "encoder.block.0.weight_g": torch.ones(1, 1, 1), + "encoder.block.0.weight_v": torch.ones(1, 1, 3), + } + converted, kwargs = converter.from_official(state) + assert set(converted) == { + "encoder.block.0.bias", + "encoder.block.0.parametrizations.weight.original0", + "encoder.block.0.parametrizations.weight.original1", + } + assert kwargs == {"config": converter.config} + + +def test_decode_normalized_requires_stereo() -> None: + model = MiniMaxH3AudioVAE.__new__(MiniMaxH3AudioVAE) + torch.nn.Module.__init__(model) + model.config = _config() + with pytest.raises(ValueError, match=r"\[2, 32, T\]"): + model.decode_normalized(torch.zeros(1, 8, 4)) + + +def test_decode_normalized_casts_fp32_denoise_output_to_model_dtype() -> None: + model = MiniMaxH3AudioVAE.__new__(MiniMaxH3AudioVAE) + torch.nn.Module.__init__(model) + model.config = _config() + model.probe = torch.nn.Parameter(torch.zeros(1, dtype=torch.bfloat16)) + observed: list[torch.dtype] = [] + + def decode(latent: torch.Tensor) -> torch.Tensor: + observed.append(latent.dtype) + return latent[:, :1] + + model.decode = decode + output = model.decode_normalized(torch.zeros(2, 8, 4, dtype=torch.float32)) + assert observed == [torch.bfloat16] + assert output.shape == (1, 2, 4) diff --git a/tests/unit/models/test_minimax_h3_dit.py b/tests/unit/models/test_minimax_h3_dit.py new file mode 100644 index 0000000..028dd7f --- /dev/null +++ b/tests/unit/models/test_minimax_h3_dit.py @@ -0,0 +1,119 @@ +from unittest.mock import MagicMock, patch + +import pytest +import torch + +from telefuser.models.minimax_h3_dit import ( + MINIMAX_H3_FP32_BUFFER_NAMES, + MINIMAX_H3_FP32_PARAM_NAMES, + MiniMaxH3DiT, + MiniMaxH3DiTConfig, + _reorder_grouped_qkv_to_qkv, +) + + +def _small_config() -> MiniMaxH3DiTConfig: + return MiniMaxH3DiTConfig( + hidden_size=32, + num_layers=2, + token_refiner_num_layers=1, + num_attention_heads=4, + attention_head_dim=8, + ffn_hidden_size=64, + latents_dim=2, + audio_latents_dim=2, + patch_size=(1, 2, 2), + text_dim=16, + timestep_input_dim=8, + time_embed_hidden_size=32, + time_embed_dim=16, + rope_inv_freq_len=1, + ) + + +def test_released_architecture_has_exact_parameter_contract() -> None: + config = MiniMaxH3DiTConfig() + with torch.device("meta"): + model = MiniMaxH3DiT(config) + assert len(model.state_dict()) == 535 + assert model.video_patch_proj.weight.shape == (5376, 96) + assert model.blocks[49].attn.qkv_proj.weight.shape == (3 * 56 * 128, 5376) + assert model.token_refiner.blocks[1].mlp.fc1.weight.shape == (2 * 14336, 5376) + + +def test_mixed_precision_boundaries_match_upstream_contract() -> None: + with torch.device("meta"): + model = MiniMaxH3DiT(_small_config()) + state = model.state_dict() + for name, tensor in state.items(): + if name in MINIMAX_H3_FP32_PARAM_NAMES | MINIMAX_H3_FP32_BUFFER_NAMES: + assert tensor.dtype == torch.float32, name + elif tensor.is_floating_point(): + assert tensor.dtype == torch.bfloat16, name + + +def test_to_preserves_fp32_boundary_values_during_dtype_conversion() -> None: + model = MiniMaxH3DiT(_small_config()) + boundary_names = MINIMAX_H3_FP32_PARAM_NAMES | MINIMAX_H3_FP32_BUFFER_NAMES + with torch.no_grad(): + for index, name in enumerate(sorted(boundary_names), start=1): + tensor = model.state_dict()[name] + tensor.fill_(index * 0.123456789) + expected = {name: model.state_dict()[name].clone() for name in boundary_names} + + model.to(dtype=torch.bfloat16) + + state = model.state_dict() + for name, value in expected.items(): + assert state[name].dtype == torch.float32 + assert torch.equal(state[name], value), name + assert model.blocks[0].attn.qkv_proj.weight.dtype == torch.bfloat16 + + +def test_grouped_qkv_reorder_matches_sglang_vector() -> None: + weight = torch.arange(12, dtype=torch.float32).reshape(12, 1) + actual = _reorder_grouped_qkv_to_qkv( + weight, + num_query_groups=2, + heads_per_group=1, + head_dim=2, + ) + expected = torch.tensor([0, 1, 6, 7, 2, 3, 8, 9, 4, 5, 10, 11], dtype=torch.float32).reshape(12, 1) + torch.testing.assert_close(actual, expected) + + +def test_small_packed_forward_returns_video_and_audio_rows() -> None: + torch.manual_seed(0) + model = MiniMaxH3DiT(_small_config()).eval() + sequence = 8 + video_positions = torch.arange(4, 8) + audio_positions = torch.arange(2, 4) + text_positions = torch.arange(0, 2) + video, audio = model( + x=torch.randn(1, sequence, 8), + audio_x=torch.randn(1, sequence, 2), + img_position_ids=torch.zeros(1, sequence, 3, dtype=torch.float64), + unique_timesteps=torch.tensor([0.5]), + inverse_indices=torch.zeros(sequence, dtype=torch.long), + update_mask=torch.ones(video_positions.numel(), dtype=torch.bool), + token_tags=torch.tensor([1, 1, 2, 2, 0, 0, 0, 0]), + prompt_embeds=torch.randn(2, 16), + img_pos_info={"position_ids": video_positions}, + audio_pos_info={"position_ids": audio_positions}, + text_pos_info={"position_ids": text_positions}, + img_pos_for_infer_output_info={"position_ids": video_positions}, + packed_seq_params={"cu_seqlens_q": torch.tensor([0, sequence], dtype=torch.int32)}, + ) + assert video.shape == (4, 8) + assert audio.shape == (2, 2) + assert torch.isfinite(video).all() + assert torch.isfinite(audio).all() + + +def test_enable_usp_rejects_uneven_head_partition() -> None: + model = MiniMaxH3DiT(_small_config()) + with ( + patch("telefuser.models.minimax_h3_dit.get_ulysses_world_size", return_value=3), + pytest.raises(ValueError, match="must be divisible"), + ): + model.enable_usp(MagicMock()) diff --git a/tests/unit/models/test_minimax_h3_encoder.py b/tests/unit/models/test_minimax_h3_encoder.py new file mode 100644 index 0000000..361cda4 --- /dev/null +++ b/tests/unit/models/test_minimax_h3_encoder.py @@ -0,0 +1,25 @@ +import torch + +from telefuser.models.minimax_h3_encoder import ( + MINIMAX_H3_QWEN3VL_SELECTED_LM_LAYER, + _is_unconsumed_checkpoint_weight, +) + + +def test_encoder_filters_tail_norm_and_lm_head_but_keeps_layer_49() -> None: + assert not _is_unconsumed_checkpoint_weight("model.language_model.layers.49.self_attn.q_proj.weight") + assert _is_unconsumed_checkpoint_weight("model.language_model.layers.50.self_attn.q_proj.weight") + assert _is_unconsumed_checkpoint_weight("model.language_model.layers.63.mlp.down_proj.weight") + assert _is_unconsumed_checkpoint_weight("model.language_model.norm.weight") + assert _is_unconsumed_checkpoint_weight("lm_head.weight") + assert MINIMAX_H3_QWEN3VL_SELECTED_LM_LAYER == 50 + + +def test_encoder_filter_does_not_drop_visual_or_embedding_weights() -> None: + for name in ( + "model.language_model.embed_tokens.weight", + "model.visual.patch_embed.proj.weight", + "model.visual.blocks.26.attn.qkv.weight", + ): + assert not _is_unconsumed_checkpoint_weight(name) + assert torch.bfloat16.is_floating_point diff --git a/tests/unit/models/test_minimax_h3_video_vae.py b/tests/unit/models/test_minimax_h3_video_vae.py new file mode 100644 index 0000000..e80b9d1 --- /dev/null +++ b/tests/unit/models/test_minimax_h3_video_vae.py @@ -0,0 +1,61 @@ +from pathlib import Path + +import pytest +import torch + +from telefuser.models.minimax_h3_video_vae import ( + MiniMaxH3VideoVAEConfig, + MiniMaxH3VideoVAEStateDictConverter, +) + + +def _config() -> MiniMaxH3VideoVAEConfig: + return MiniMaxH3VideoVAEConfig( + architecture={"vae_ratio": 16, "vae_ratio_t": 4, "embed_dim": 24}, + clip_length=17, + token_drop=3, + encoder_tiling=True, + decoder_tiling=True, + tile_size=256, + tile_overlap_min=64, + chunk_dim=-1, + latent_channels=24, + latents_mean=(0.0,) * 24, + latents_std=(1.0,) * 24, + ) + + +def test_video_vae_config_reads_released_component(tmp_path: Path) -> None: + source = tmp_path / "source" + source.mkdir() + (tmp_path / "config.json").write_text( + '{"source_path":"source","vae_clip_length":17,"vae_token_drop":3,' + '"vae_encoder_tiling":1,"vae_decoder_tiling":1,"vae_tile_size":256,' + '"vae_tile_overlap_min":64,"vae_chunk_dim":-1,"latent_channels":24,' + '"latents_mean":[' + ",".join(["0"] * 24) + "]," + '"latents_std":[' + ",".join(["1"] * 24) + "]}", + encoding="utf-8", + ) + (source / "config.json").write_text( + '{"vae_ratio":16,"vae_ratio_t":4,"embed_dim":24}', + encoding="utf-8", + ) + config = MiniMaxH3VideoVAEConfig.from_path(tmp_path) + assert config.latent_channels == 24 + assert config.model_kwargs()["parallel_tiling"] is False + + +def test_video_vae_config_rejects_wrong_geometry() -> None: + config = _config() + config.architecture["vae_ratio_t"] = 8 + with pytest.raises(ValueError, match="f16/t4/d24"): + config.validate() + + +def test_video_vae_converter_prefixes_composed_model() -> None: + converter = MiniMaxH3VideoVAEStateDictConverter.__new__(MiniMaxH3VideoVAEStateDictConverter) + converter.config = _config() + state = {"encoder.conv_in.conv.weight": torch.zeros(1)} + converted, kwargs = converter.from_official(state) + assert set(converted) == {"model.encoder.conv_in.conv.weight"} + assert kwargs == {"config": converter.config} From a66780ec4b9f58ff7f4b24a3bf5a76d2b8f2f92e Mon Sep 17 00:00:00 2001 From: lzx1413 Date: Tue, 4 Aug 2026 10:04:54 +0000 Subject: [PATCH 02/24] feat(minimax-h3): add generation pipeline Implement canonical T2VA, FL2VA, and Ref2VA request planning, material preparation, packed joint denoising, independent audio/video scheduling, VAE processing, and presentation. Add local-checkpoint H100 runners, complete-audio muxing, Ulysses execution support, trajectory capture, and focused CPU plus distributed tests. Verification: MiniMax H3 CPU suite passed (64 tests, 5 subtests); ruff check and format checks passed. --- examples/minimax_h3/common.py | 173 +++++++ examples/minimax_h3/minimax_h3_fl2va_h100.py | 56 +++ examples/minimax_h3/minimax_h3_ref2va_h100.py | 57 +++ telefuser/pipelines/minimax_h3/__init__.py | 29 ++ telefuser/pipelines/minimax_h3/canvas.py | 98 ++++ .../pipelines/minimax_h3/condition_noise.py | 170 +++++++ telefuser/pipelines/minimax_h3/constants.py | 37 ++ telefuser/pipelines/minimax_h3/data.py | 388 ++++++++++++++ telefuser/pipelines/minimax_h3/denoising.py | 337 +++++++++++++ telefuser/pipelines/minimax_h3/material_io.py | 114 +++++ .../pipelines/minimax_h3/packed_sequence.py | 476 ++++++++++++++++++ .../pipelines/minimax_h3/packed_tokens.py | 98 ++++ telefuser/pipelines/minimax_h3/pipeline.py | 267 ++++++++++ .../pipelines/minimax_h3/presentation.py | 274 ++++++++++ .../pipelines/minimax_h3/resolved_plan.py | 426 ++++++++++++++++ telefuser/pipelines/minimax_h3/scheduler.py | 206 ++++++++ .../pipelines/minimax_h3/task_profiles.py | 278 ++++++++++ .../pipelines/minimax_h3/text_encoding.py | 165 ++++++ .../pipelines/minimax_h3/time_request.py | 59 +++ telefuser/pipelines/minimax_h3/vae.py | 319 ++++++++++++ .../test_minimax_h3_distributed.py | 106 ++++ .../unit/pipelines/minimax_h3/test_canvas.py | 34 ++ tests/unit/pipelines/minimax_h3/test_data.py | 159 ++++++ .../minimax_h3/test_packed_sequence.py | 149 ++++++ .../minimax_h3/test_packed_tokens.py | 22 + .../pipelines/minimax_h3/test_parallelism.py | 59 +++ .../pipelines/minimax_h3/test_pipeline.py | 313 ++++++++++++ .../pipelines/minimax_h3/test_scheduler.py | 62 +++ tests/unit/pipelines/minimax_h3/test_vae.py | 83 +++ .../validation/minimax_h3_trajectory_stage.py | 201 ++++++++ 30 files changed, 5215 insertions(+) create mode 100644 examples/minimax_h3/common.py create mode 100644 examples/minimax_h3/minimax_h3_fl2va_h100.py create mode 100644 examples/minimax_h3/minimax_h3_ref2va_h100.py create mode 100644 telefuser/pipelines/minimax_h3/__init__.py create mode 100644 telefuser/pipelines/minimax_h3/canvas.py create mode 100644 telefuser/pipelines/minimax_h3/condition_noise.py create mode 100644 telefuser/pipelines/minimax_h3/constants.py create mode 100644 telefuser/pipelines/minimax_h3/data.py create mode 100644 telefuser/pipelines/minimax_h3/denoising.py create mode 100644 telefuser/pipelines/minimax_h3/material_io.py create mode 100644 telefuser/pipelines/minimax_h3/packed_sequence.py create mode 100644 telefuser/pipelines/minimax_h3/packed_tokens.py create mode 100644 telefuser/pipelines/minimax_h3/pipeline.py create mode 100644 telefuser/pipelines/minimax_h3/presentation.py create mode 100644 telefuser/pipelines/minimax_h3/resolved_plan.py create mode 100644 telefuser/pipelines/minimax_h3/scheduler.py create mode 100644 telefuser/pipelines/minimax_h3/task_profiles.py create mode 100644 telefuser/pipelines/minimax_h3/text_encoding.py create mode 100644 telefuser/pipelines/minimax_h3/time_request.py create mode 100644 telefuser/pipelines/minimax_h3/vae.py create mode 100644 tests/integration/test_minimax_h3_distributed.py create mode 100644 tests/unit/pipelines/minimax_h3/test_canvas.py create mode 100644 tests/unit/pipelines/minimax_h3/test_data.py create mode 100644 tests/unit/pipelines/minimax_h3/test_packed_sequence.py create mode 100644 tests/unit/pipelines/minimax_h3/test_packed_tokens.py create mode 100644 tests/unit/pipelines/minimax_h3/test_parallelism.py create mode 100644 tests/unit/pipelines/minimax_h3/test_pipeline.py create mode 100644 tests/unit/pipelines/minimax_h3/test_scheduler.py create mode 100644 tests/unit/pipelines/minimax_h3/test_vae.py create mode 100644 tools/validation/minimax_h3_trajectory_stage.py diff --git a/examples/minimax_h3/common.py b/examples/minimax_h3/common.py new file mode 100644 index 0000000..3957337 --- /dev/null +++ b/examples/minimax_h3/common.py @@ -0,0 +1,173 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Shared local-checkpoint loader and artifact writer for MiniMax H3 examples.""" + +from __future__ import annotations + +import subprocess +import tempfile +from pathlib import Path + +import torch + +from telefuser.core.config import ModelRuntimeConfig, OffloadConfig, ParallelConfig, WeightOffloadType +from telefuser.core.module_manager import ModuleManager +from telefuser.models.minimax_h3_audio_vae import MiniMaxH3AudioVAE +from telefuser.models.minimax_h3_dit import MiniMaxH3DiT +from telefuser.models.minimax_h3_encoder import MiniMaxH3Encoder +from telefuser.models.minimax_h3_video_vae import MiniMaxH3VideoVAE +from telefuser.pipelines.minimax_h3.pipeline import ( + MiniMaxH3Generation, + MiniMaxH3Pipeline, + MiniMaxH3PipelineConfig, +) +from telefuser.utils.audio import save_wav +from telefuser.utils.video import save_video + + +def _checkpoint_shards(component: Path) -> list[str]: + shards = sorted(str(path) for path in component.glob("model-*.safetensors")) + if not shards: + raise FileNotFoundError(f"no model safetensor shards found in {component}") + return shards + + +def load_minimax_h3_pipeline( + model_root: str | Path, + *, + partition: str, + device: str = "cuda:0", + num_inference_steps: int = 50, + ulysses_degree: int = 1, +) -> MiniMaxH3Pipeline: + if partition not in {"FL2VA", "Ref2VA"}: + raise ValueError("partition must be 'FL2VA' or 'Ref2VA'") + if ulysses_degree not in {1, 2, 4}: + raise ValueError("ulysses_degree must be 1, 2, or 4") + component_root = Path(model_root) / partition + if not component_root.is_dir(): + raise FileNotFoundError(f"MiniMax H3 partition not found: {component_root}") + runtime_device = torch.device(device) + offload = OffloadConfig( + offload_type=WeightOffloadType.MODEL_CPU_OFFLOAD, + pin_cpu_memory=False, + ) + bf16_runtime = ModelRuntimeConfig( + device_type=runtime_device.type, + device_id=runtime_device.index or 0, + torch_dtype=torch.bfloat16, + offload_config=offload, + ) + dit_runtime = ModelRuntimeConfig( + device_type=runtime_device.type, + device_id=runtime_device.index or 0, + torch_dtype=torch.bfloat16, + offload_config=offload, + parallel_config=ParallelConfig( + device_ids=list(range(ulysses_degree)), + sp_ulysses_degree=ulysses_degree, + timeout=1800, + ), + ) + fp32_runtime = ModelRuntimeConfig( + device_type=runtime_device.type, + device_id=runtime_device.index or 0, + torch_dtype=torch.float32, + offload_config=offload, + ) + + manager = ModuleManager(device="cpu", torch_dtype=torch.bfloat16) + transformer_dir = component_root / "transformer" + manager.load_model( + _checkpoint_shards(transformer_dir), + device="cpu", + torch_dtype=torch.bfloat16, + low_cpu_mem_usage=True, + name="minimax_h3_transformer", + model_class=MiniMaxH3DiT, + converter_kwargs={"config_path": transformer_dir / "config.json"}, + ) + encoder_dir = component_root / "text_encoder" + manager.load_model( + _checkpoint_shards(encoder_dir), + device="cpu", + torch_dtype=torch.bfloat16, + low_cpu_mem_usage=True, + name="minimax_h3_text_encoder", + model_class=MiniMaxH3Encoder, + converter_kwargs={"config_path": encoder_dir}, + ) + video_vae_dir = component_root / "video_vae" + manager.load_model( + str(video_vae_dir / "source" / "model.safetensors"), + device="cpu", + torch_dtype=torch.float32, + low_cpu_mem_usage=True, + name="minimax_h3_video_vae", + model_class=MiniMaxH3VideoVAE, + converter_kwargs={"config_path": video_vae_dir}, + ) + audio_vae_dir = component_root / "audio_vae" + manager.load_model( + str(audio_vae_dir / "model.safetensors"), + device="cpu", + torch_dtype=torch.float32, + low_cpu_mem_usage=True, + name="minimax_h3_audio_vae", + model_class=MiniMaxH3AudioVAE, + converter_kwargs={"config_path": audio_vae_dir}, + ) + + pipeline = MiniMaxH3Pipeline(device=device) + pipeline.init( + manager, + MiniMaxH3PipelineConfig( + processor_path=str(component_root / "processor"), + text_encoder_config=bf16_runtime, + dit_config=dit_runtime, + video_vae_config=fp32_runtime, + audio_vae_config=fp32_runtime, + num_inference_steps=num_inference_steps, + ), + ) + return pipeline + + +def save_generation(result: MiniMaxH3Generation, output_path: str | Path) -> None: + output = Path(output_path) + output.parent.mkdir(parents=True, exist_ok=True) + frames = result.video[0].mul(255).clamp(0, 255).to(torch.uint8) + waveform = result.audio[0] + with tempfile.TemporaryDirectory() as directory: + video_path = Path(directory) / "video.mp4" + audio_path = Path(directory) / "audio.wav" + save_wav(waveform, result.audio_sample_rate, str(audio_path)) + save_video( + frames, + str(video_path), + fps=float(result.video_fps), + quality=6, + ) + subprocess.run( + [ + "ffmpeg", + "-y", + "-i", + str(video_path), + "-i", + str(audio_path), + "-map", + "0:v:0", + "-map", + "1:a:0", + "-c:v", + "copy", + "-c:a", + "aac", + str(output), + ], + check=True, + capture_output=True, + ) + + +__all__ = ["load_minimax_h3_pipeline", "save_generation"] diff --git a/examples/minimax_h3/minimax_h3_fl2va_h100.py b/examples/minimax_h3/minimax_h3_fl2va_h100.py new file mode 100644 index 0000000..45a307b --- /dev/null +++ b/examples/minimax_h3/minimax_h3_fl2va_h100.py @@ -0,0 +1,56 @@ +# SPDX-License-Identifier: Apache-2.0 +from __future__ import annotations + +import argparse + +try: + from examples.minimax_h3.common import load_minimax_h3_pipeline, save_generation +except ModuleNotFoundError as exc: + if exc.name != "examples": + raise + from common import load_minimax_h3_pipeline, save_generation + + +def main() -> None: + parser = argparse.ArgumentParser(description="Generate MiniMax H3 T2VA/FL2VA audio-video on H100 GPUs") + parser.add_argument("--model-root", default="/hhb-data/aigc/model_zoo/MiniMaxAI_MiniMax-H3") + parser.add_argument("--image") + parser.add_argument("--last-image") + parser.add_argument("--prompt", required=True) + parser.add_argument("--duration", type=float, default=8.0) + parser.add_argument("--seed", type=int, default=0) + parser.add_argument("--steps", type=int, default=50) + parser.add_argument("--device", default="cuda:0") + parser.add_argument("--ulysses-degree", type=int, choices=(1, 2, 4), default=1) + parser.add_argument("--output", default="minimax_h3_fl2va.mp4") + args = parser.parse_args() + + if args.last_image and not args.image: + parser.error("--last-image requires --image") + conditions = [] + if args.image: + conditions.append({"type": "image", "role": "keyframe", "uri": args.image, "frame_index": 0}) + if args.last_image: + conditions.append({"type": "image", "role": "keyframe", "uri": args.last_image, "frame_index": -1}) + pipeline = load_minimax_h3_pipeline( + args.model_root, + partition="FL2VA", + device=args.device, + num_inference_steps=args.steps, + ulysses_degree=args.ulysses_degree, + ) + try: + result = pipeline( + task="fl2va" if conditions else "t2va", + prompt=args.prompt, + conditions=conditions, + target={"short_edge": 768, "aspect_ratio": "auto", "duration_seconds": args.duration}, + seed=args.seed, + ) + save_generation(result, args.output) + finally: + pipeline.stop() + + +if __name__ == "__main__": + main() diff --git a/examples/minimax_h3/minimax_h3_ref2va_h100.py b/examples/minimax_h3/minimax_h3_ref2va_h100.py new file mode 100644 index 0000000..38a9ac5 --- /dev/null +++ b/examples/minimax_h3/minimax_h3_ref2va_h100.py @@ -0,0 +1,57 @@ +# SPDX-License-Identifier: Apache-2.0 +from __future__ import annotations + +import argparse + +try: + from examples.minimax_h3.common import load_minimax_h3_pipeline, save_generation +except ModuleNotFoundError as exc: + if exc.name != "examples": + raise + from common import load_minimax_h3_pipeline, save_generation + + +def main() -> None: + parser = argparse.ArgumentParser(description="Generate MiniMax H3 Ref2VA audio-video on H100 GPUs") + parser.add_argument("--model-root", default="/hhb-data/aigc/model_zoo/MiniMaxAI_MiniMax-H3") + parser.add_argument("--image", action="append", default=[]) + parser.add_argument("--video", action="append", default=[]) + parser.add_argument("--audio", action="append", default=[]) + parser.add_argument("--prompt", required=True) + parser.add_argument("--duration", type=float, default=5.0) + parser.add_argument("--seed", type=int, default=0) + parser.add_argument("--steps", type=int, default=50) + parser.add_argument("--device", default="cuda:0") + parser.add_argument("--ulysses-degree", type=int, choices=(1, 2, 4), default=1) + parser.add_argument("--output", default="minimax_h3_ref2va.mp4") + args = parser.parse_args() + + conditions = [ + *({"type": "image", "role": "reference", "uri": path} for path in args.image), + *({"type": "video", "role": "reference", "uri": path} for path in args.video), + *({"type": "audio", "role": "reference", "uri": path} for path in args.audio), + ] + if not conditions: + parser.error("at least one --image, --video, or --audio reference is required") + pipeline = load_minimax_h3_pipeline( + args.model_root, + partition="Ref2VA", + device=args.device, + num_inference_steps=args.steps, + ulysses_degree=args.ulysses_degree, + ) + try: + result = pipeline( + task="ref2va", + prompt=args.prompt, + conditions=conditions, + target={"short_edge": 768, "aspect_ratio": "auto", "duration_seconds": args.duration}, + seed=args.seed, + ) + save_generation(result, args.output) + finally: + pipeline.stop() + + +if __name__ == "__main__": + main() diff --git a/telefuser/pipelines/minimax_h3/__init__.py b/telefuser/pipelines/minimax_h3/__init__.py new file mode 100644 index 0000000..91f895d --- /dev/null +++ b/telefuser/pipelines/minimax_h3/__init__.py @@ -0,0 +1,29 @@ +"""MiniMax H3 pipeline contracts and stages.""" + +from telefuser.pipelines.minimax_h3.data import minimax_h3_validate_canonical_request +from telefuser.pipelines.minimax_h3.packed_sequence import ( + minimax_h3_packed_sequence, + minimax_h3_packed_sequence_ref2va_blocks, +) +from telefuser.pipelines.minimax_h3.pipeline import ( + MiniMaxH3Generation, + MiniMaxH3Pipeline, + MiniMaxH3PipelineConfig, +) +from telefuser.pipelines.minimax_h3.resolved_plan import ( + MiniMaxH3ResolvedPlan, + minimax_h3_resolve_plan, +) +from telefuser.pipelines.minimax_h3.scheduler import MiniMaxH3EulerAncestralEta0SchedulerAdapter + +__all__ = [ + "MiniMaxH3EulerAncestralEta0SchedulerAdapter", + "MiniMaxH3Generation", + "MiniMaxH3Pipeline", + "MiniMaxH3PipelineConfig", + "MiniMaxH3ResolvedPlan", + "minimax_h3_packed_sequence", + "minimax_h3_packed_sequence_ref2va_blocks", + "minimax_h3_resolve_plan", + "minimax_h3_validate_canonical_request", +] diff --git a/telefuser/pipelines/minimax_h3/canvas.py b/telefuser/pipelines/minimax_h3/canvas.py new file mode 100644 index 0000000..e1cee40 --- /dev/null +++ b/telefuser/pipelines/minimax_h3/canvas.py @@ -0,0 +1,98 @@ +# SPDX-License-Identifier: Apache-2.0 +"""MiniMax H3 keyframe target-canvas preparation. + +Geometry behavior: +- auto-aspect canvases delegate to the shared adaptive v2 shape resolver; +- cover-crop: aspect-preserving max-scale LANCZOS resize + center crop, + upscaling refused unless explicitly allowed. + +Both the Qwen presentation (pixel_values) and the visual-condition tokenizer consume +the SAME prepared canvas image, so preparation +is cached per request in batch.extra. +""" + +from __future__ import annotations + +from typing import Any + +MINIMAX_H3_CANVAS_MULTIPLE = 32 + + +def minimax_h3_cover_crop_plan( + *, + source_width: int, + source_height: int, + target_width: int, + target_height: int, + allow_upscale: bool, +) -> dict[str, Any]: + """Deterministic aspect-preserving cover-crop transform.""" + if source_width <= 0 or source_height <= 0: + raise ValueError("cover_crop requires positive source width/height") + scale = max(target_width / float(source_width), target_height / float(source_height)) + if scale > 1.0 and not allow_upscale: + raise ValueError( + "target_canvas cover_crop would upscale the source; set " + f"allow_upscale=true (source={source_width}x{source_height}, " + f"target={target_width}x{target_height})" + ) + resized_width = max(target_width, int(round(source_width * scale))) + resized_height = max(target_height, int(round(source_height * scale))) + left = max(0, (resized_width - target_width) // 2) + top = max(0, (resized_height - target_height) // 2) + return { + "scale": scale, + "resized_size": (resized_width, resized_height), + "crop_box": (left, top, left + target_width, top + target_height), + } + + +def minimax_h3_prepare_keyframe_canvas( + image: Any, + *, + target_width: int, + target_height: int, + allow_upscale: bool = False, +) -> Any: + """Prepare a PIL image onto the target canvas. + + Identity (no resample) when the image already IS the canvas. + """ + from PIL import Image + + image = image.convert("RGB") + if image.size == (target_width, target_height): + return image + plan = minimax_h3_cover_crop_plan( + source_width=image.size[0], + source_height=image.size[1], + target_width=target_width, + target_height=target_height, + allow_upscale=allow_upscale, + ) + resized = image.resize(plan["resized_size"], Image.Resampling.LANCZOS) + return resized.crop(plan["crop_box"]) + + +def minimax_h3_stretch_keyframe_canvas( + image: Any, + *, + target_width: int, + target_height: int, +) -> Any: + """Stretch the FL first frame directly onto the resolved target canvas.""" + + from PIL import Image + + image = image.convert("RGB") + if image.size == (target_width, target_height): + return image + return image.resize((target_width, target_height), Image.Resampling.LANCZOS) + + +__all__ = [ + "MINIMAX_H3_CANVAS_MULTIPLE", + "minimax_h3_cover_crop_plan", + "minimax_h3_prepare_keyframe_canvas", + "minimax_h3_stretch_keyframe_canvas", +] diff --git a/telefuser/pipelines/minimax_h3/condition_noise.py b/telefuser/pipelines/minimax_h3/condition_noise.py new file mode 100644 index 0000000..b6d31b2 --- /dev/null +++ b/telefuser/pipelines/minimax_h3/condition_noise.py @@ -0,0 +1,170 @@ +# SPDX-License-Identifier: Apache-2.0 +"""MiniMax H3 visual/audio condition-noise augmentation. + +The request's condition timestep is applied to both the tensor value and the +DiT timestep. Tokenizer artifacts remain clean +and reusable; this module materializes the fixed noised anchors immediately +before the denoise loop. +""" + +from __future__ import annotations + +from collections.abc import Sequence + +import torch + +from telefuser.pipelines.minimax_h3.packed_tokens import ( + minimax_h3_patchify_video_latent, +) + +# Channel-major packed audio rows always carry a stereo layout. +MINIMAX_H3_AUDIO_COND_CHANNELS = 2 + + +def minimax_h3_imgvid_cond_noise_aug_rows( + clean_rows: torch.Tensor, + *, + condition_shapes: Sequence[Sequence[int]], + target_latent_t: int, + imgvid_cond_num_frames: int, + seed: int, + noise_aug: float, +) -> torch.Tensor: + """Apply the imgvid-condition RF noise recipe to packed clean rows. + + ``condition_shapes`` contains ``(latent_t, latent_h, latent_w)`` in packed + visual-condition order. A new CPU generator with the same row seed is + created for every condition. Under the dependent-noise policy, each draw + uses the target temporal length plus the template's imgvid-condition frame + count, then slices the prefix matching the current condition. + """ + + noise_aug = float(noise_aug) + if not 0.0 <= noise_aug <= 1.0: + raise ValueError(f"noise_aug must be in [0, 1], got {noise_aug}") + if noise_aug == 1.0: + return clean_rows + if clean_rows.ndim != 2 or int(clean_rows.shape[1]) != 96: + raise ValueError(f"clean imgvid condition rows must have shape [n, 96], got {list(clean_rows.shape)}") + + target_latent_t = int(target_latent_t) + imgvid_cond_num_frames = int(imgvid_cond_num_frames) + if target_latent_t <= 0: + raise ValueError(f"target_latent_t must be positive, got {target_latent_t}") + if imgvid_cond_num_frames <= 0: + raise ValueError( + f"imgvid_cond_num_frames must be positive when condition rows exist, got {imgvid_cond_num_frames}" + ) + + parsed_shapes: list[tuple[int, int, int]] = [] + expected_rows = 0 + for raw_shape in condition_shapes: + if len(raw_shape) != 3: + raise ValueError( + f"each imgvid condition shape must be (latent_t, latent_h, latent_w), got {list(raw_shape)}" + ) + latent_t, latent_h, latent_w = (int(value) for value in raw_shape) + if latent_t <= 0 or latent_h <= 0 or latent_w <= 0: + raise ValueError(f"imgvid condition shape must be positive, got {list(raw_shape)}") + if latent_h % 2 or latent_w % 2: + raise ValueError( + f"imgvid condition spatial dimensions must be divisible by 2, got {(latent_t, latent_h, latent_w)}" + ) + parsed_shapes.append((latent_t, latent_h, latent_w)) + expected_rows += latent_t * (latent_h // 2) * (latent_w // 2) + if not parsed_shapes: + raise ValueError("condition_shapes must not be empty") + if int(clean_rows.shape[0]) != expected_rows: + raise ValueError( + f"clean imgvid condition rows {int(clean_rows.shape[0])} != shape-derived rows {expected_rows}" + ) + + out: list[torch.Tensor] = [] + row_offset = 0 + timestep = torch.tensor(noise_aug, dtype=torch.float32, device=clean_rows.device) + for latent_t, latent_h, latent_w in parsed_shapes: + full_t = target_latent_t + imgvid_cond_num_frames + if full_t < latent_t: + raise ValueError(f"condition latent_t {latent_t} exceeds the noise draw length {full_t}") + generator = torch.Generator(device="cpu").manual_seed(int(seed)) + noise = torch.randn( + 1, + 24, + full_t, + latent_h, + latent_w, + generator=generator, + dtype=torch.float32, + device="cpu", + )[:, :, :latent_t] + noise_rows = minimax_h3_patchify_video_latent(noise, patch_size=[1, 2, 2]).to( + device=clean_rows.device, dtype=torch.float32 + ) + row_count = int(noise_rows.shape[0]) + clean_part = clean_rows[row_offset : row_offset + row_count].to(torch.float32) + out.append(timestep * clean_part + (1.0 - timestep) * noise_rows) + row_offset += row_count + return (out[0] if len(out) == 1 else torch.cat(out, dim=0)).contiguous() + + +def minimax_h3_audio_cond_noise_aug_rows( + clean_rows: torch.Tensor, + *, + condition_audio_t: Sequence[int], + seed: int, + noise_aug: float, +) -> torch.Tensor: + """Apply the audio-condition RF noise recipe to packed clean rows. + + ``condition_audio_t`` contains the latent T of each audio-bearing + condition in canonical request order. Noise is drawn per condition + element, with a fresh CPU generator seeded with ``seed + 1`` for every + element. Consequently each condition restarts the + same RNG stream; concatenating the rows and drawing once would be + numerically different for ordered multi-reference requests. + + The mix is intentionally evaluated on CPU in fp32 before the packed rows + are transferred to the DiT device. + """ + + noise_aug = float(noise_aug) + if not 0.0 <= noise_aug <= 1.0: + raise ValueError(f"noise_aug must be in [0, 1], got {noise_aug}") + if noise_aug == 1.0: + return clean_rows + if clean_rows.ndim != 2 or int(clean_rows.shape[1]) != 32: + raise ValueError(f"clean audio condition rows must have shape [n, 32], got {list(clean_rows.shape)}") + + audio_channels = MINIMAX_H3_AUDIO_COND_CHANNELS + parsed_audio_t = [int(value) for value in condition_audio_t] + if not parsed_audio_t: + raise ValueError("condition_audio_t must not be empty") + if any(value <= 0 for value in parsed_audio_t): + raise ValueError(f"condition audio latent lengths must be positive, got {parsed_audio_t}") + expected_rows = audio_channels * sum(parsed_audio_t) + if int(clean_rows.shape[0]) != expected_rows: + raise ValueError(f"clean audio condition rows {int(clean_rows.shape[0])} != shape-derived rows {expected_rows}") + + out: list[torch.Tensor] = [] + row_offset = 0 + timestep = torch.tensor(noise_aug, dtype=torch.float32, device="cpu") + for audio_t in parsed_audio_t: + row_count = audio_channels * audio_t + clean_part = clean_rows[row_offset : row_offset + row_count].detach().to(device="cpu", dtype=torch.float32) + generator = torch.Generator(device="cpu").manual_seed(int(seed) + 1) + noise = torch.randn( + clean_part.shape, + generator=generator, + dtype=torch.float32, + device="cpu", + ) + out.append(timestep * clean_part + (1.0 - timestep) * noise) + row_offset += row_count + rows = out[0] if len(out) == 1 else torch.cat(out, dim=0) + return rows.to(device=clean_rows.device, dtype=torch.float32).contiguous() + + +__all__ = [ + "minimax_h3_audio_cond_noise_aug_rows", + "minimax_h3_imgvid_cond_noise_aug_rows", +] diff --git a/telefuser/pipelines/minimax_h3/constants.py b/telefuser/pipelines/minimax_h3/constants.py new file mode 100644 index 0000000..1f59adf --- /dev/null +++ b/telefuser/pipelines/minimax_h3/constants.py @@ -0,0 +1,37 @@ +# SPDX-License-Identifier: Apache-2.0 +from __future__ import annotations + +# Direct-encode text embeddings: {"positive": +# {"hidden_states": Tensor[text_len, 5120] bf16 cpu, "text_len": int}} +MINIMAX_H3_TEXT_EMBEDDINGS_EXTRA_KEY = "minimax_h3_text_embeddings" +# Direct keyframe encode: {"rows": Tensor[n_rows, 96] fp32 cpu, +# "latent_h": int, "latent_w": int, "canvas_height": int, +# "canvas_width": int, "keyframes": [...], +# "semantic_frame_indices": [...], "pixel_frame_indices": [...]} +MINIMAX_H3_KEYFRAME_COND_ROWS_EXTRA_KEY = "minimax_h3_keyframe_cond_rows" +# Direct sigma schedules: {"video": [float], "audio": [float]} +MINIMAX_H3_SIGMAS_EXTRA_KEY = "minimax_h3_sigmas" +# Direct denoise state: {"initial_video_rows", "initial_audio_rows", +# "latent_t", "latent_h", "latent_w", "audio_t"} +MINIMAX_H3_DENOISE_STATE_EXTRA_KEY = "minimax_h3_denoise_state" +# ref2va direct reference encodes. +MINIMAX_H3_REFERENCE_IMAGE_ROWS_EXTRA_KEY = "minimax_h3_reference_image_rows" +MINIMAX_H3_REFERENCE_AUDIO_ROWS_EXTRA_KEY = "minimax_h3_reference_audio_rows" +MINIMAX_H3_REFERENCE_VIDEO_ROWS_EXTRA_KEY = "minimax_h3_reference_video_rows" +MINIMAX_H3_PREPARED_REFERENCE_VIDEO_EXTRA_KEY = "minimax_h3_prepared_reference_video" + +MINIMAX_H3_SUPPORTED_FPS = 24 +MINIMAX_H3_MIN_DURATION_SECONDS = 4.0 +MINIMAX_H3_MAX_DURATION_SECONDS = 15.0 + +# The distilled checkpoint has exactly one positive denoise branch. +MINIMAX_H3_DEFAULT_BRANCHES: tuple = ({"name": "cond_1"},) + +# Audited 4xH200 T2VA profiles. The tuple is +# (warmup steps, residual-difference threshold, max consecutive cached steps). +MINIMAX_H3_QUALITY_PROFILES: dict[str, tuple[int, float, int] | None] = { + "lossless": None, + "high": (4, 0.04, 1), + "medium": (4, 0.12, 3), + "low": (4, 0.24, 3), +} diff --git a/telefuser/pipelines/minimax_h3/data.py b/telefuser/pipelines/minimax_h3/data.py new file mode 100644 index 0000000..dd725ee --- /dev/null +++ b/telefuser/pipelines/minimax_h3/data.py @@ -0,0 +1,388 @@ +# SPDX-License-Identifier: Apache-2.0 +"""MiniMax H3 canonical request validation. + +Entry fail-fast for `minimax_h3.request/v1`: every violation raises ValueError +with the offending field path. Output is a normalized canonical dict (frame +indices validated but semantic -1 preserved, nothing else rewritten — prompt passes through verbatim and +conditions order is semantic, never reordered). +""" + +from __future__ import annotations + +import math +from collections.abc import Mapping, Sequence +from typing import Any + +from telefuser.pipelines.minimax_h3.constants import ( + MINIMAX_H3_MAX_DURATION_SECONDS, + MINIMAX_H3_MIN_DURATION_SECONDS, + MINIMAX_H3_SUPPORTED_FPS, +) +from telefuser.pipelines.minimax_h3.task_profiles import ( + MINIMAX_H3_CONDITION_ROLE_KEYFRAME, + MINIMAX_H3_CONDITION_ROLE_REFERENCE, + MINIMAX_H3_FINITE_ASPECT_RATIOS, + MINIMAX_H3_FL2VA_KEYFRAME_SIGNATURES, + MINIMAX_H3_TASK_FL2VA, + MINIMAX_H3_TASK_REF2VA, + MINIMAX_H3_TASK_T2VA, + MiniMaxH3TaskProfile, + canonical_minimax_h3_task, + minimax_h3_task_profile, +) +from telefuser.pipelines.minimax_h3.time_request import ( + minimax_h3_align_frame_count, +) + +MINIMAX_H3_REQUEST_SCHEMA = "minimax_h3.request/v1" +MINIMAX_H3_MAX_SIGNED_SEED = (1 << 63) - 1 +MINIMAX_H3_REF2VA_MAX_FILES = 12 +MINIMAX_H3_REF2VA_MAX_IMAGES = 9 +MINIMAX_H3_REF2VA_MAX_VIDEOS = 3 +MINIMAX_H3_REF2VA_MAX_AUDIO = 3 +MINIMAX_H3_REF2VA_MIN_CLIP_SECONDS = 2.0 +MINIMAX_H3_REF2VA_MAX_CLIP_SECONDS = 15.0 +MINIMAX_H3_REF2VA_MAX_TOTAL_SECONDS = 15.0 +_ALLOWED_CONDITION_KEYS = frozenset({"type", "uri", "role", "frame_index", "start_time_seconds"}) + + +def _require_str(value: Any, path: str) -> str: + if not isinstance(value, str) or value == "": + raise ValueError(f"{path} must be a non-empty string") + return value + + +def _require_int(value: Any, path: str) -> int: + if isinstance(value, bool) or not isinstance(value, int): + raise ValueError(f"{path} must be an integer") + return value + + +def _optional_positive_finite_float(value: Any, path: str) -> float | None: + if value is None: + return None + if isinstance(value, bool) or not isinstance(value, (int, float)): + raise ValueError(f"{path} must be a number") + normalized = float(value) + if not math.isfinite(normalized) or normalized <= 0.0: + raise ValueError(f"{path} must be a positive finite number") + return normalized + + +def _optional_nonnegative_finite_float(value: Any, path: str) -> float | None: + if value is None: + return None + if isinstance(value, bool) or not isinstance(value, (int, float)): + raise ValueError(f"{path} must be a number") + normalized = float(value) + if not math.isfinite(normalized) or normalized < 0.0: + raise ValueError(f"{path} must be a non-negative finite number") + return normalized + + +def _validate_target(target: Any, *, profile: MiniMaxH3TaskProfile) -> dict[str, Any]: + path = "target" + if not isinstance(target, Mapping): + raise ValueError(f"{path} is required and must be an object") + # The canonical target has a deliberately small projection. Transport + # compatibility keys are ignored; only these three declared values are + # validated and emitted below. + short_edge = _require_int(target.get("short_edge"), f"{path}.short_edge") + if short_edge != 768: + raise ValueError(f"{path}.short_edge must be 768 for minimax_h3, got {short_edge}") + aspect_ratio = _require_str(target.get("aspect_ratio"), f"{path}.aspect_ratio") + if profile.aspect_ratio_forced_auto and aspect_ratio != "auto": + raise ValueError(f'{path}.aspect_ratio must be "auto" for task {profile.task!r}, got {aspect_ratio!r}') + has_duration = target.get("duration_seconds") is not None + if ( + profile.task in {MINIMAX_H3_TASK_T2VA, MINIMAX_H3_TASK_REF2VA} + and aspect_ratio != "auto" + and aspect_ratio not in MINIMAX_H3_FINITE_ASPECT_RATIOS + ): + raise ValueError( + f"{path}.aspect_ratio for task {profile.task!r} must be 'auto' or " + f"one of {list(MINIMAX_H3_FINITE_ASPECT_RATIOS)!r}, got " + f"{aspect_ratio!r}" + ) + if not has_duration: + if not profile.duration_from_audio_reference: + raise ValueError(f"{path}.duration_seconds is required") + # ref2va: duration may derive from a reference audio; the + # audio-condition presence is enforced after conditions validate. + out: dict[str, Any] = { + "short_edge": short_edge, + "aspect_ratio": aspect_ratio, + } + if has_duration: + duration = target["duration_seconds"] + if isinstance(duration, bool) or not isinstance(duration, (int, float)): + raise ValueError(f"{path}.duration_seconds must be a number") + if duration <= 0: + raise ValueError(f"{path}.duration_seconds must be positive") + if not (MINIMAX_H3_MIN_DURATION_SECONDS <= float(duration) <= MINIMAX_H3_MAX_DURATION_SECONDS): + raise ValueError( + f"{path}.duration_seconds must be in " + f"[{MINIMAX_H3_MIN_DURATION_SECONDS:g}, " + f"{MINIMAX_H3_MAX_DURATION_SECONDS:g}], got {duration}" + ) + out["duration_seconds"] = float(duration) + return out + + +def _validate_conditions( + conditions: Any, + *, + profile: MiniMaxH3TaskProfile, + frame_count: int | None, +) -> list[dict[str, Any]]: + path = "conditions" + if conditions is None: + conditions = [] + if not isinstance(conditions, Sequence) or isinstance(conditions, (str, bytes)): + raise ValueError(f"{path} must be a list") + + if not profile.conditions_required: + if len(conditions) > 0: + raise ValueError(f"{path} must be empty for task {profile.task!r} (got {len(conditions)} entries)") + return [] + if len(conditions) == 0: + raise ValueError(f"{path} requires at least one entry for task {profile.task!r}") + if profile.min_condition_count is not None and len(conditions) < profile.min_condition_count: + raise ValueError( + f"{path} requires at least {profile.min_condition_count} entries " + f"for task {profile.task!r}, got {len(conditions)}" + ) + if profile.max_condition_count is not None and len(conditions) > profile.max_condition_count: + raise ValueError( + f"{path} allows at most {profile.max_condition_count} entries " + f"for task {profile.task!r}, got {len(conditions)}" + ) + + aligned_frame_count = minimax_h3_align_frame_count(frame_count) if frame_count is not None else None + normalized: list[dict[str, Any]] = [] + seen_frame_indices: dict[int, int] = {} + for index, cond in enumerate(conditions): + cpath = f"{path}[{index}]" + if not isinstance(cond, Mapping): + raise ValueError(f"{cpath} must be an object") + unknown = set(cond) - _ALLOWED_CONDITION_KEYS + if unknown: + raise ValueError(f"{cpath} has unknown fields: {sorted(unknown)}") + role = _require_str(cond.get("role"), f"{cpath}.role") + if role not in ( + MINIMAX_H3_CONDITION_ROLE_KEYFRAME, + MINIMAX_H3_CONDITION_ROLE_REFERENCE, + ): + raise ValueError(f"{cpath}.role must be keyframe or reference, got {role!r}") + cond_type = _require_str(cond.get("type"), f"{cpath}.type") + try: + rule = profile.rule_for(role=role, condition_type=cond_type) + except ValueError as exc: + raise ValueError(f"{cpath}: {exc}") from exc + uri = _require_str(cond.get("uri"), f"{cpath}.uri") + + entry: dict[str, Any] = {"type": cond_type, "uri": uri, "role": role} + if rule.requires_frame_index: + frame_index = _require_int(cond.get("frame_index"), f"{cpath}.frame_index") + if aligned_frame_count is None: + raise ValueError(f"{cpath}.frame_index requires a resolved target duration") + if frame_index == -1: + resolved = aligned_frame_count - 1 + elif 0 <= frame_index < aligned_frame_count: + resolved = frame_index + else: + raise ValueError( + f"{cpath}.frame_index must be -1 or in " + f"[0, {aligned_frame_count}) after 17n+5 frame alignment, " + f"got {frame_index}" + ) + if resolved in seen_frame_indices: + raise ValueError( + f"{cpath}.frame_index resolves to {resolved}, already " + f"bound by conditions[{seen_frame_indices[resolved]}]" + ) + seen_frame_indices[resolved] = index + # Preserve the request-level semantic index. In particular, -1 is + # the canonical last-frame sentinel; the resolved pixel frame is + # carried separately by MiniMaxH3ResolvedPlan. + entry["frame_index"] = frame_index + elif cond.get("frame_index") is not None: + raise ValueError(f"{cpath}.frame_index is not allowed for role={role!r}") + start_time_seconds = _optional_nonnegative_finite_float( + cond.get("start_time_seconds"), f"{cpath}.start_time_seconds" + ) + if start_time_seconds is not None: + if cond_type not in {"video", "video_audio"}: + raise ValueError(f"{cpath}.start_time_seconds is only allowed for video or video_audio references") + entry["start_time_seconds"] = start_time_seconds + normalized.append(entry) + return normalized + + +def _validate_fl2va_conditions(conditions: Sequence[Mapping[str, Any]]) -> None: + """Enforce the public FL contract after per-entry schema validation.""" + + frame_indices = tuple(condition.get("frame_index") for condition in conditions) + if frame_indices not in MINIMAX_H3_FL2VA_KEYFRAME_SIGNATURES: + raise ValueError( + "conditions for task 'fl2va' must be one or two ordered " + "image/keyframe entries with frame_index [0], [-1], or [0, -1], " + f"got {list(frame_indices)!r}" + ) + + +def _validate_ref2va_conditions(conditions: Sequence[Mapping[str, Any]]) -> None: + if len(conditions) > MINIMAX_H3_REF2VA_MAX_FILES: + raise ValueError( + f"conditions allows at most {MINIMAX_H3_REF2VA_MAX_FILES} files for task 'ref2va', got {len(conditions)}" + ) + image_count = sum(condition["type"] == "image" for condition in conditions) + video_count = sum(condition["type"] in {"video", "video_audio"} for condition in conditions) + audio_count = sum(condition["type"] in {"audio", "video_audio"} for condition in conditions) + if image_count > MINIMAX_H3_REF2VA_MAX_IMAGES: + raise ValueError( + f"conditions allows at most {MINIMAX_H3_REF2VA_MAX_IMAGES} image references, got {image_count}" + ) + if video_count > MINIMAX_H3_REF2VA_MAX_VIDEOS: + raise ValueError( + f"conditions allows at most {MINIMAX_H3_REF2VA_MAX_VIDEOS} video references, got {video_count}" + ) + if audio_count > MINIMAX_H3_REF2VA_MAX_AUDIO: + raise ValueError(f"conditions allows at most {MINIMAX_H3_REF2VA_MAX_AUDIO} audio references, got {audio_count}") + has_visual = any(condition["type"] in {"image", "video", "video_audio"} for condition in conditions) + if audio_count and not has_visual: + raise ValueError("conditions audio references require at least one image or video reference") + + +def minimax_h3_validate_reference_media_facts( + conditions: Sequence[Mapping[str, Any]], + duration_seconds_by_condition: Mapping[int, float], +) -> None: + """Validate Ref2VA clip and aggregate durations after local media probing.""" + + video_total = 0.0 + audio_total = 0.0 + for index, condition in enumerate(conditions): + condition_type = condition.get("type") + if condition_type not in {"video", "video_audio", "audio"}: + continue + if index not in duration_seconds_by_condition: + raise ValueError(f"conditions[{index}] is missing probed duration") + duration = duration_seconds_by_condition[index] + if isinstance(duration, bool) or not isinstance(duration, (int, float)) or not math.isfinite(duration): + raise ValueError(f"conditions[{index}] duration must be a finite number") + duration = float(duration) + if not MINIMAX_H3_REF2VA_MIN_CLIP_SECONDS <= duration <= MINIMAX_H3_REF2VA_MAX_CLIP_SECONDS: + raise ValueError( + f"conditions[{index}] duration must be in " + f"[{MINIMAX_H3_REF2VA_MIN_CLIP_SECONDS:g}, {MINIMAX_H3_REF2VA_MAX_CLIP_SECONDS:g}] seconds, " + f"got {duration:g}" + ) + if condition_type in {"video", "video_audio"}: + video_total += duration + if condition_type in {"audio", "video_audio"}: + audio_total += duration + if video_total > MINIMAX_H3_REF2VA_MAX_TOTAL_SECONDS: + raise ValueError( + f"reference video total duration must not exceed {MINIMAX_H3_REF2VA_MAX_TOTAL_SECONDS:g} seconds, " + f"got {video_total:g}" + ) + if audio_total > MINIMAX_H3_REF2VA_MAX_TOTAL_SECONDS: + raise ValueError( + f"reference audio total duration must not exceed {MINIMAX_H3_REF2VA_MAX_TOTAL_SECONDS:g} seconds, " + f"got {audio_total:g}" + ) + + +def minimax_h3_validate_canonical_request( + *, + task: Any, + prompt: Any, + conditions: Any, + target: Any, + flow_shift: Any = None, + audio_flow_shift: Any = None, + seed: Any = None, + **_extra_kwargs: Any, +) -> dict[str, Any]: + """Validate and normalize a `minimax_h3.request/v1` canonical request. + + Returns the normalized canonical dict; raises ValueError with a field + path on any violation. Conditions order is preserved (it is semantic: + prompt ordinal labels reference it). seed=0 is a legal value. + """ + # Accept transport wrappers and compatibility kwargs at this boundary, but + # never copy them into the canonical request. + del _extra_kwargs + # Normalize the task name before profile lookup so offline callers match + # the adapter behaviour. + task_name = canonical_minimax_h3_task(_require_str(task, "task")) + profile = minimax_h3_task_profile(task_name) + prompt_text = _require_str(prompt, "prompt") + + normalized_target = _validate_target(target, profile=profile) + requested_frame_count = None + if normalized_target.get("duration_seconds") is not None: + requested_frame_count = int(round(float(normalized_target["duration_seconds"]) * MINIMAX_H3_SUPPORTED_FPS)) + normalized_conditions = _validate_conditions( + conditions, + profile=profile, + frame_count=requested_frame_count, + ) + if profile.task == MINIMAX_H3_TASK_FL2VA: + _validate_fl2va_conditions(normalized_conditions) + elif profile.task == MINIMAX_H3_TASK_REF2VA: + _validate_ref2va_conditions(normalized_conditions) + # ref2va accepts ordered material streams containing any mix of + # image/audio/video/video_audio references. Type admission is handled by + # the task profile; temporal ambiguity is validated later when target + # duration is omitted. + if not profile.video_reference_supported: + for index, cond in enumerate(normalized_conditions): + if cond["type"] in ("video", "video_audio"): + raise ValueError( + f"conditions[{index}]: video references are not supported " + f"in v1 for task {profile.task!r} (image/audio only)" + ) + if normalized_target.get("duration_seconds") is None: + # Only reachable for duration_from_audio_reference profiles. + duration_sources = [cond for cond in normalized_conditions if cond["type"] in ("audio", "video", "video_audio")] + if not duration_sources: + raise ValueError( + "target.duration_seconds is required, or exactly one " + "audio reference to derive duration from (including " + f"video/video_audio soundtracks; task {profile.task!r})" + ) + if len(duration_sources) > 1: + raise ValueError("target.duration_seconds is required when multiple audio-bearing references are provided") + + canonical: dict[str, Any] = { + "schema": MINIMAX_H3_REQUEST_SCHEMA, + "task": task_name, + "prompt": prompt_text, + "conditions": normalized_conditions, + "target": normalized_target, + } + normalized_flow_shift = _optional_positive_finite_float(flow_shift, "flow_shift") + normalized_audio_flow_shift = _optional_positive_finite_float(audio_flow_shift, "audio_flow_shift") + if normalized_flow_shift is not None: + canonical["flow_shift"] = normalized_flow_shift + if normalized_audio_flow_shift is not None: + canonical["audio_flow_shift"] = normalized_audio_flow_shift + if seed is not None: + normalized_seed = _require_int(seed, "seed") + if normalized_seed < 0: + raise ValueError(f"seed must be non-negative, got {normalized_seed}") + if normalized_seed > MINIMAX_H3_MAX_SIGNED_SEED: + raise ValueError(f"seed must not exceed the signed int64 maximum, got {normalized_seed}") + canonical["seed"] = normalized_seed + return canonical + + +__all__ = [ + "MINIMAX_H3_REQUEST_SCHEMA", + "MINIMAX_H3_MAX_SIGNED_SEED", + "MINIMAX_H3_SUPPORTED_FPS", + "minimax_h3_validate_canonical_request", + "minimax_h3_validate_reference_media_facts", +] diff --git a/telefuser/pipelines/minimax_h3/denoising.py b/telefuser/pipelines/minimax_h3/denoising.py new file mode 100644 index 0000000..d46800e --- /dev/null +++ b/telefuser/pipelines/minimax_h3/denoising.py @@ -0,0 +1,337 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Packed single-branch MiniMax H3 denoising stage.""" + +from __future__ import annotations + +import time +from dataclasses import dataclass + +import torch + +from telefuser.core.base_stage import BaseStage, with_model_offload +from telefuser.core.config import ModelRuntimeConfig +from telefuser.core.module_manager import ModuleManager +from telefuser.distributed.device_mesh import create_device_mesh_from_config +from telefuser.distributed.fsdp import shard_model +from telefuser.utils.logging import logger + +from .condition_noise import ( + minimax_h3_audio_cond_noise_aug_rows, + minimax_h3_imgvid_cond_noise_aug_rows, +) +from .packed_sequence import ( + minimax_h3_packed_sequence, + minimax_h3_packed_sequence_ref2va_blocks, +) +from .packed_tokens import ( + minimax_h3_patchify_video_latent, + minimax_h3_unpack_audio_tokens, + minimax_h3_unpatchify_video_tokens, +) +from .resolved_plan import MiniMaxH3ResolvedPlan +from .scheduler import MiniMaxH3EulerAncestralEta0SchedulerAdapter +from .text_encoding import MiniMaxH3TextCondition +from .time_request import minimax_h3_time_shift_sigmas +from .vae import MiniMaxH3PreparedCondition + +MINIMAX_H3_IMGVID_COND_TIMESTEP = 0.999 +MINIMAX_H3_AUDIO_REF_COND_TIMESTEP = 1.0 + + +@dataclass(frozen=True) +class MiniMaxH3DenoiseResult: + video_latent: torch.Tensor + audio_latent: torch.Tensor + packed: dict[str, torch.Tensor] + runtime_metrics: dict[str, float | int] + + +class MiniMaxH3DenoisingStage(BaseStage): + def __init__(self, module_manager: ModuleManager, model_runtime_config: ModelRuntimeConfig) -> None: + super().__init__("minimax_h3_denoising", model_runtime_config) + self.transformer = module_manager.fetch_module("minimax_h3_transformer") + if self.transformer is None: + raise ValueError("ModuleManager must contain 'minimax_h3_transformer'") + self.scheduler = MiniMaxH3EulerAncestralEta0SchedulerAdapter() + self.model_names = ["transformer"] + + def parallel_models(self) -> None: + parallel_config = self.model_runtime_config.parallel_config + unsupported = { + "cfg_degree": parallel_config.cfg_degree, + "sp_ring_degree": parallel_config.sp_ring_degree, + "pp_degree": parallel_config.pp_degree, + "tp_degree": parallel_config.tp_degree, + } + invalid = {name: degree for name, degree in unsupported.items() if degree != 1} + if invalid: + raise NotImplementedError(f"MiniMax H3 does not support these parallel degrees yet: {invalid}") + device_mesh = create_device_mesh_from_config(parallel_config) + self.transformer.device_mesh = device_mesh + self.transformer.set_attention_config(self.model_runtime_config.attention_config) + if parallel_config.sp_ulysses_degree > 1: + self.transformer.enable_usp(device_mesh) + if parallel_config.enable_fsdp: + logger.info(f"Enabling FSDP for {self.name}") + fp32_parameters = [ + parameter for parameter in self.transformer.parameters() if parameter.dtype == torch.float32 + ] + self.transformer = shard_model( + module=self.transformer, + device_id=self.device, + wrap_module_names=self.transformer.get_fsdp_module_names(), + param_dtype=self.torch_dtype, + reduce_dtype=self.torch_dtype, + buffer_dtype=torch.float32, + ignored_states=fp32_parameters, + ) + self.onload_models_flag = True + + @staticmethod + def _reference_blocks( + conditions: list[MiniMaxH3PreparedCondition], + ) -> tuple[list[dict[str, object]], torch.Tensor | None, torch.Tensor | None]: + blocks: list[dict[str, object]] = [] + visual: list[torch.Tensor] = [] + audio: list[torch.Tensor] = [] + for condition in conditions: + if condition.kind == "image": + if condition.visual_rows is None: + raise ValueError("reference image is missing visual VAE rows") + blocks.append( + { + "kind": "image", + "latent_h": condition.latent_h, + "latent_w": condition.latent_w, + } + ) + visual.append(condition.visual_rows) + elif condition.kind == "audio": + if condition.audio_rows is None: + raise ValueError("reference audio is missing audio VAE rows") + blocks.append({"kind": "audio", "ref_audio_t": condition.ref_audio_t}) + audio.append(condition.audio_rows) + elif condition.kind in {"video", "video_audio"}: + if condition.visual_rows is None: + raise ValueError("reference video is missing visual VAE rows") + blocks.append( + { + "kind": condition.kind, + "ref_audio_t": condition.ref_audio_t, + "latent_t": condition.latent_t, + "latent_h": condition.latent_h, + "latent_w": condition.latent_w, + } + ) + visual.append(condition.visual_rows) + if condition.audio_rows is not None: + audio.append(condition.audio_rows) + else: + raise ValueError(f"unsupported reference block kind {condition.kind!r}") + visual_rows = None if not visual else torch.cat(visual, dim=0) + audio_rows = None if not audio else torch.cat(audio, dim=0) + return blocks, visual_rows, audio_rows + + @with_model_offload(["transformer"]) + @torch.inference_mode() + def denoise( + self, + *, + plan: MiniMaxH3ResolvedPlan, + text: MiniMaxH3TextCondition, + conditions: list[MiniMaxH3PreparedCondition], + num_inference_steps: int, + ) -> MiniMaxH3DenoiseResult: + shape = plan.shape + latent_t = int(shape["video_latent_t"]) + latent_h = int(shape["height"]) // 16 + latent_w = int(shape["width"]) // 16 + audio_t = int(shape["audio_latent_t"]) + seed = 42 if plan.seed is None else int(plan.seed) + if num_inference_steps < 2: + raise ValueError("num_inference_steps must be at least 2") + + ref_blocks: list[dict[str, object]] | None = None + if plan.task == "ref2va": + ref_blocks, visual_cond, audio_cond = self._reference_blocks(conditions) + packed = minimax_h3_packed_sequence_ref2va_blocks( + text_len=text.text_len, + latent_t=latent_t, + latent_h=latent_h, + latent_w=latent_w, + audio_t=audio_t, + ref_blocks=ref_blocks, + ) + else: + visual_conditions = [condition for condition in conditions if condition.visual_rows is not None] + visual_cond = ( + None + if not visual_conditions + else torch.cat([condition.visual_rows for condition in visual_conditions], dim=0) + ) + audio_cond = None + semantic_indices = tuple( + int(condition.material.frame_index) + for condition in visual_conditions + if condition.material.frame_index is not None + ) + packed = minimax_h3_packed_sequence( + text_len=text.text_len, + latent_t=latent_t, + latent_h=latent_h, + latent_w=latent_w, + audio_t=audio_t, + include_keyframe_cond=bool(visual_conditions), + keyframe_frame_indices=semantic_indices if visual_conditions else None, + frame_count=int(shape["frame_count"]) if visual_conditions else None, + ) + + token_tags = packed["token_tags"].clone() + token_tags[: text.text_len] = text.token_tags + condition_shapes = [ + (condition.latent_t, condition.latent_h, condition.latent_w) + for condition in conditions + if condition.visual_rows is not None + ] + if visual_cond is not None: + visual_cond = minimax_h3_imgvid_cond_noise_aug_rows( + visual_cond, + condition_shapes=condition_shapes, + target_latent_t=latent_t, + imgvid_cond_num_frames=len(condition_shapes), + seed=seed, + noise_aug=MINIMAX_H3_IMGVID_COND_TIMESTEP, + ) + audio_lengths = [condition.ref_audio_t for condition in conditions if condition.audio_rows is not None] + if audio_cond is not None: + audio_cond = minimax_h3_audio_cond_noise_aug_rows( + audio_cond, + condition_audio_t=audio_lengths, + seed=seed, + noise_aug=MINIMAX_H3_AUDIO_REF_COND_TIMESTEP, + ) + + video_generator = torch.Generator(device="cpu").manual_seed(seed) + video_native = torch.randn( + 1, + 24, + latent_t, + latent_h, + latent_w, + generator=video_generator, + dtype=torch.float32, + ) + video_target = minimax_h3_patchify_video_latent(video_native, patch_size=(1, 2, 2)) + audio_generator = torch.Generator(device="cpu").manual_seed(seed) + audio_target = torch.randn(audio_t * 2, 32, generator=audio_generator, dtype=torch.float32) + + device = next(self.transformer.parameters()).device + if device.type == "cuda": + torch.cuda.reset_peak_memory_stats(device) + reset_communication_metrics = getattr(self.transformer, "reset_communication_metrics", None) + if callable(reset_communication_metrics): + reset_communication_metrics() + denoising_started = time.perf_counter() + video_rows = torch.zeros(len(packed["img_pos"]), 96, dtype=torch.float32, device=device) + audio_rows = torch.zeros(len(packed["audio_pos"]), 32, dtype=torch.float32, device=device) + video_update = packed["update_mask"].to(device) + audio_update = packed.get("audio_update_mask", torch.ones(len(packed["audio_pos"]), dtype=torch.bool)).to( + device + ) + video_rows[video_update] = video_target.to(device) + audio_rows[audio_update] = audio_target.to(device) + if visual_cond is not None: + video_rows[~video_update] = visual_cond.to(device) + if audio_cond is not None: + audio_rows[~audio_update] = audio_cond.to(device) + + video_shift = plan.flow_shift or plan.default_flow_shift + audio_shift = plan.audio_flow_shift or plan.default_audio_flow_shift + video_sigmas = minimax_h3_time_shift_sigmas(num_steps=num_inference_steps, shift_scale=video_shift) + audio_sigmas = minimax_h3_time_shift_sigmas(num_steps=num_inference_steps, shift_scale=audio_shift) + img_pos = packed["img_pos"].to(device) + audio_pos = packed["audio_pos"].to(device) + text_pos = packed["text_pos"].to(device) + target_img_pos = img_pos[video_update] + target_audio_row_start = int((~audio_update).sum()) + + for step in range(len(video_sigmas) - 1): + t_video = float(1.0 - video_sigmas[step]) + t_audio = float(1.0 - audio_sigmas[step]) + row_timesteps = torch.full((int(packed["seq_len"]),), t_video, dtype=torch.float32) + row_timesteps[packed["img_pos"][~packed["update_mask"]]] = max(t_video, MINIMAX_H3_IMGVID_COND_TIMESTEP) + row_timesteps[packed["audio_pos"][audio_update.cpu()]] = t_audio + row_timesteps[packed["audio_pos"][~audio_update.cpu()]] = max(t_audio, MINIMAX_H3_AUDIO_REF_COND_TIMESTEP) + unique_timesteps, inverse_indices = torch.unique( + row_timesteps, + sorted=True, + return_inverse=True, + ) + x = torch.zeros(1, int(packed["seq_len"]), 96, dtype=torch.float32, device=device) + audio_x = torch.zeros(1, int(packed["seq_len"]), 32, dtype=torch.float32, device=device) + x[0].index_copy_(0, img_pos, video_rows) + audio_x[0].index_copy_(0, audio_pos, audio_rows) + video_velocity, audio_velocity = self.transformer( + x=x, + audio_x=audio_x, + img_position_ids=packed["img_position_ids"].unsqueeze(0).float().to(device), + unique_timesteps=unique_timesteps.to(device), + inverse_indices=inverse_indices.to(device), + update_mask=video_update, + update_audio_mask=audio_update, + prompt_embeds=text.hidden_states.to(device), + img_pos_info={"position_ids": img_pos}, + audio_pos_info={"position_ids": audio_pos}, + text_pos_info={"position_ids": text_pos}, + img_pos_for_infer_output_info={"position_ids": target_img_pos}, + packed_seq_params={"cu_seqlens_q": packed["cu_seqlens"].to(device)}, + block_token_tags=token_tags.to(device), + skip_mask_out_condition=True, + ) + stepped = self.scheduler.step_denoising( + input_visual_latent=video_rows[video_update], + input_audio_latent=audio_rows[audio_update], + timestep=torch.tensor(t_video, device=device), + video_timestep=torch.tensor(t_video, device=device), + audio_timestep=torch.tensor(t_audio, device=device), + noise_pred_visual=video_velocity, + noise_pred_audio=audio_velocity[target_audio_row_start:], + sigma_curr=video_sigmas[step], + sigma_next=video_sigmas[step + 1], + video_sigma_curr=video_sigmas[step], + video_sigma_next=video_sigmas[step + 1], + audio_sigma_curr=audio_sigmas[step], + audio_sigma_next=audio_sigmas[step + 1], + ) + video_rows[video_update] = stepped["output_visual_latent"] + audio_rows[audio_update] = stepped["output_audio_latent"] + + video_latent = minimax_h3_unpatchify_video_tokens( + video_rows[video_update].cpu(), + latent_shape=(latent_t, latent_h // 2, latent_w // 2, 24), + patch_size=(1, 2, 2), + ) + audio_latent = minimax_h3_unpack_audio_tokens( + audio_rows[audio_update].cpu(), + audio_t=audio_t * 2, + audio_channel=2, + ) + if device.type == "cuda": + torch.cuda.synchronize(device) + peak_allocated = int(torch.cuda.max_memory_allocated(device)) + peak_reserved = int(torch.cuda.max_memory_reserved(device)) + else: + peak_allocated = 0 + peak_reserved = 0 + get_communication_seconds = getattr(self.transformer, "communication_seconds", None) + communication_seconds = float(get_communication_seconds()) if callable(get_communication_seconds) else 0.0 + runtime_metrics: dict[str, float | int] = { + "denoising_seconds": time.perf_counter() - denoising_started, + "peak_allocated_bytes": peak_allocated, + "peak_reserved_bytes": peak_reserved, + "communication_seconds": communication_seconds, + } + return MiniMaxH3DenoiseResult(video_latent, audio_latent, packed, runtime_metrics) + + +__all__ = ["MiniMaxH3DenoiseResult", "MiniMaxH3DenoisingStage"] diff --git a/telefuser/pipelines/minimax_h3/material_io.py b/telefuser/pipelines/minimax_h3/material_io.py new file mode 100644 index 0000000..7d5d9aa --- /dev/null +++ b/telefuser/pipelines/minimax_h3/material_io.py @@ -0,0 +1,114 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Bounded material localization and probing for MiniMax H3 requests.""" + +from __future__ import annotations + +import json +import shutil +import subprocess +import tempfile +import urllib.parse +import urllib.request +from contextlib import contextmanager +from dataclasses import dataclass +from pathlib import Path +from typing import Iterator + +MINIMAX_H3_MAX_MATERIAL_BYTES = 2 * 1024**3 + + +@dataclass(frozen=True) +class MiniMaxH3MaterialFacts: + width: int | None = None + height: int | None = None + duration_seconds: float | None = None + sample_rate: int | None = None + has_audio: bool = False + + +@contextmanager +def minimax_h3_localize_material(uri: str) -> Iterator[Path]: + """Yield a local path for a file path, file URI, or bounded HTTP(S) URI.""" + parsed = urllib.parse.urlparse(uri) + if parsed.scheme in ("", "file"): + path = Path(urllib.request.url2pathname(parsed.path) if parsed.scheme == "file" else uri).expanduser() + if not path.is_file(): + raise FileNotFoundError(f"MiniMax H3 material not found: {path}") + yield path + return + if parsed.scheme not in {"http", "https"}: + raise ValueError(f"unsupported MiniMax H3 material URI scheme {parsed.scheme!r}") + + suffix = Path(parsed.path).suffix + temp_dir = Path(tempfile.mkdtemp(prefix="telefuser-minimax-h3-")) + target = temp_dir / f"material{suffix}" + try: + request = urllib.request.Request(uri, headers={"User-Agent": "TeleFuser/1"}) + with urllib.request.urlopen(request, timeout=60) as response, target.open("wb") as output: + declared = response.headers.get("Content-Length") + if declared is not None and int(declared) > MINIMAX_H3_MAX_MATERIAL_BYTES: + raise ValueError("MiniMax H3 material exceeds the 2 GiB download limit") + copied = 0 + while chunk := response.read(1024 * 1024): + copied += len(chunk) + if copied > MINIMAX_H3_MAX_MATERIAL_BYTES: + raise ValueError("MiniMax H3 material exceeds the 2 GiB download limit") + output.write(chunk) + yield target + finally: + shutil.rmtree(temp_dir, ignore_errors=True) + + +def _probe_av(path: Path) -> dict: + command = [ + "ffprobe", + "-v", + "error", + "-show_entries", + "format=duration:stream=index,codec_type,width,height,sample_rate,duration", + "-of", + "json", + str(path), + ] + result = subprocess.run(command, check=True, capture_output=True, text=True) + return json.loads(result.stdout) + + +def minimax_h3_probe_material(path: Path, condition_type: str) -> MiniMaxH3MaterialFacts: + if condition_type == "image": + from PIL import Image, ImageOps + + with Image.open(path) as source: + image = ImageOps.exif_transpose(source) + return MiniMaxH3MaterialFacts(width=int(image.width), height=int(image.height)) + if condition_type not in {"audio", "video", "video_audio"}: + raise ValueError(f"unsupported MiniMax H3 condition type {condition_type!r}") + payload = _probe_av(path) + streams = payload.get("streams") or [] + video = next((stream for stream in streams if stream.get("codec_type") == "video"), None) + audio = next((stream for stream in streams if stream.get("codec_type") == "audio"), None) + if condition_type in {"video", "video_audio"} and video is None: + raise ValueError(f"{condition_type} material has no video stream: {path}") + if condition_type in {"audio", "video_audio"} and audio is None: + raise ValueError(f"{condition_type} material has no audio stream: {path}") + duration_value = (payload.get("format") or {}).get("duration") + if duration_value is None: + duration_value = (audio or video or {}).get("duration") + if duration_value is None: + raise ValueError(f"media duration is unavailable: {path}") + sample_rate = None if audio is None or audio.get("sample_rate") is None else int(audio["sample_rate"]) + return MiniMaxH3MaterialFacts( + width=None if video is None else int(video["width"]), + height=None if video is None else int(video["height"]), + duration_seconds=float(duration_value), + sample_rate=sample_rate, + has_audio=audio is not None, + ) + + +__all__ = [ + "MINIMAX_H3_MAX_MATERIAL_BYTES", + "MiniMaxH3MaterialFacts", + "minimax_h3_localize_material", + "minimax_h3_probe_material", +] diff --git a/telefuser/pipelines/minimax_h3/packed_sequence.py b/telefuser/pipelines/minimax_h3/packed_sequence.py new file mode 100644 index 0000000..5e14883 --- /dev/null +++ b/telefuser/pipelines/minimax_h3/packed_sequence.py @@ -0,0 +1,476 @@ +# SPDX-License-Identifier: Apache-2.0 +"""MiniMax H3 packed-sequence materialization from the validated workspace +builder, covering fl2va and t2va layouts. + +Layout: [text L | imgvid_cond C | audio A(=t*2ch) | video_target V | pad P]. +Builder rules: +- block-derived position infos, update masks, token tags, and cu_seqlens +- img_position_ids fp64 grid: text rows (row_idx,0,0); video/cond t counter + continues text_len with temporal interp spans (frame_rescale 5/3 x + frame_per_token (1,4,4,4,4)); each spatial sqrt_area axis uses evenly spaced + coordinates excluding the right endpoint, then scales them by INTERP; + audio channel-major blocks pinned to the w-grid extremes. +""" + +from __future__ import annotations + +from collections.abc import Mapping, Sequence +from typing import Any + +import numpy as np +import torch + +from telefuser.pipelines.minimax_h3.task_profiles import ( + MINIMAX_H3_FL2VA_KEYFRAME_SIGNATURES, +) + +MINIMAX_H3_PACKED_SEQUENCE_ALIGNMENT = 64 + +_INTERP = 32 +_T_GROUP = 5 +_FRAME_PER_TOKEN = (1, 4, 4, 4, 4) +_FRAME_RESCALE = 5.0 / 3.0 +_PATCH_H = 2 +_PATCH_W = 2 + + +def _keyframe_cond_frame_indices( + *, + include_keyframe_cond: bool, + keyframe_frame_indices: list[int] | tuple[int, ...] | None, +) -> list[int]: + if not include_keyframe_cond: + if keyframe_frame_indices is not None: + raise ValueError("keyframe_frame_indices must be omitted when keyframe cond is not included") + return [] + if keyframe_frame_indices is None: + raise ValueError("strict fl2va packed layout requires keyframe_frame_indices") + if any(isinstance(value, bool) or not isinstance(value, int) for value in keyframe_frame_indices): + raise ValueError("strict fl2va packed layout requires integer keyframe_frame_indices") + out = list(keyframe_frame_indices) + if tuple(out) not in MINIMAX_H3_FL2VA_KEYFRAME_SIGNATURES: + raise ValueError( + "strict fl2va packed layout requires keyframe_frame_indices in " + f"{MINIMAX_H3_FL2VA_KEYFRAME_SIGNATURES!r}, got {out!r}" + ) + return out + + +def _resolve_keyframe_frame_indices( + frame_indices: Sequence[int], + *, + frame_count: int | None, +) -> list[int]: + if frame_indices and frame_count is None: + raise ValueError("frame_count is required when keyframe_frame_indices are provided") + if frame_count is None: + return [] + if frame_count <= 0: + raise ValueError("frame_count must be positive") + seen: dict[int, int] = {} + resolved: list[int] = [] + for block_index, semantic_index in enumerate(frame_indices): + if semantic_index == -1: + resolved_index = frame_count - 1 + elif 0 <= semantic_index < frame_count: + resolved_index = semantic_index + else: + raise ValueError(f"keyframe frame index {semantic_index} must be -1 or in [0, {frame_count})") + previous = seen.get(resolved_index) + if previous is not None: + raise ValueError( + f"keyframe frame index at block {block_index} resolves to " + f"{resolved_index}, already bound by block {previous}" + ) + seen[resolved_index] = block_index + resolved.append(resolved_index) + return resolved + + +def _temporal_position_span(temporal_length: int) -> float: + """Temporal position span for patch_t=1, in fp64. + + NOTE: intentionally NOT merged with ``_video_t_span``. This variant sums + via numpy (pairwise summation), matching the fl2va anchor computation, + while ``_video_t_span`` sums sequentially, matching the ref2va + t-origin accumulation. The two orders diverge in the last ulp + from n=16 onward, so each path must keep its own summation order. + """ + spans = np.ones(int(temporal_length), dtype=np.float64) * _FRAME_RESCALE + for token_index in range(_T_GROUP): + spans[token_index::_T_GROUP] *= _FRAME_PER_TOKEN[token_index] + return float(spans.sum()) + + +def minimax_h3_packed_sequence( + *, + text_len: int, + latent_t: int, + latent_h: int, + latent_w: int, + audio_t: int, + audio_channel: int = 2, + include_keyframe_cond: bool, + keyframe_frame_indices: list[int] | tuple[int, ...] | None = None, + frame_count: int | None = None, +) -> dict[str, Any]: + """Build the packed-sequence structural fields for one CFG branch. + + The used length is padded up to a multiple of 64. + """ + ph, pw = latent_h // _PATCH_H, latent_w // _PATCH_W + frame_rows = ph * pw + cond_frame_indices = _keyframe_cond_frame_indices( + include_keyframe_cond=include_keyframe_cond, + keyframe_frame_indices=keyframe_frame_indices, + ) + resolved_cond_frame_indices = _resolve_keyframe_frame_indices( + cond_frame_indices, + frame_count=frame_count, + ) + cond_rows = len(cond_frame_indices) * frame_rows + video_rows = latent_t * frame_rows + audio_rows = audio_t * audio_channel + used = text_len + cond_rows + audio_rows + video_rows + seq_len = ( + (used + MINIMAX_H3_PACKED_SEQUENCE_ALIGNMENT - 1) + // MINIMAX_H3_PACKED_SEQUENCE_ALIGNMENT + * MINIMAX_H3_PACKED_SEQUENCE_ALIGNMENT + ) + + text_sl = slice(0, text_len) + cond_sl = slice(text_len, text_len + cond_rows) + audio_sl = slice(cond_sl.stop, cond_sl.stop + audio_rows) + video_sl = slice(audio_sl.stop, audio_sl.stop + video_rows) + target_img_pos = torch.arange(video_sl.start, video_sl.stop) + img_pos = torch.cat([torch.arange(cond_sl.start, cond_sl.stop), target_img_pos]) if cond_rows else target_img_pos + update_mask = torch.zeros(img_pos.shape[0], dtype=torch.bool) + update_mask[cond_rows:] = True + audio_pos = torch.arange(audio_sl.start, audio_sl.stop) + text_pos = torch.arange(0, text_len) + + g = torch.zeros(seq_len, 3, dtype=torch.float64) + g[text_sl, 0] = torch.arange(text_len, dtype=torch.float64) + + t_grid = _video_t_grid(latent_t, float(text_len)) + sqrt_area = np.sqrt(latent_h * latent_w) + h_grid = _axis_from_sqrt_area(latent_h, _PATCH_H, sqrt_area) + w_grid = _axis_from_sqrt_area(latent_w, _PATCH_W, sqrt_area) + hh, ww = torch.meshgrid(h_grid, w_grid, indexing="ij") + frame = torch.stack([hh.reshape(-1), ww.reshape(-1)], dim=-1) + video_g = g[video_sl].view(latent_t, frame_rows, 3) + video_g[:, :, 0] = t_grid[:, None] + video_g[:, :, 1:] = frame[None] + for block_index, pixel_index in enumerate(resolved_cond_frame_indices): + sl = slice( + cond_sl.start + block_index * frame_rows, + cond_sl.start + (block_index + 1) * frame_rows, + ) + if pixel_index == 0: + cond_t = float(text_len) + elif frame_count is not None and pixel_index == frame_count - 1: + cond_t = float(text_len) + _temporal_position_span(latent_t) - _FRAME_RESCALE + else: + raise ValueError( + f"fl2va packed layout only supports first/last keyframe anchors, got resolved frame index {pixel_index}" + ) + g[sl, 0] = cond_t + g[sl, 1:] = frame + audio_t_grid = float(text_len) + torch.arange(audio_t, dtype=torch.float64) + g[audio_sl, 0] = audio_t_grid.repeat(audio_channel) + g[audio_sl.start : audio_sl.start + audio_t, 2] = float(w_grid[0]) + g[audio_sl.start + audio_t : audio_sl.stop, 2] = float(w_grid[-1]) + + token_tags = torch.full((seq_len,), -1, dtype=torch.long) # PADDING + token_tags[text_sl] = 1 # TEXT (fl2va image-segment override happens upstream) + token_tags[audio_sl] = 2 # AUDIO + token_tags[img_pos] = 0 # VIDEO + + cu = torch.tensor([0, used, seq_len], dtype=torch.int32) + return { + "seq_len": seq_len, + "img_pos": img_pos, + "audio_pos": audio_pos, + "text_pos": text_pos, + "update_mask": update_mask, + "img_position_ids": g, + "token_tags": token_tags, + "cu_seqlens": cu, + } + + +def _positive_int( + block: Mapping[str, object], + key: str, + path: str, + *, + allow_zero: bool = False, +) -> int: + value = block.get(key) + if isinstance(value, bool) or not isinstance(value, int): + raise ValueError(f"{path}.{key} must be an integer") + if value < 0 or (value == 0 and not allow_zero): + predicate = "non-negative" if allow_zero else "positive" + raise ValueError(f"{path}.{key} must be {predicate}") + return int(value) + + +def _axis_from_sqrt_area(dim: int, patch: int, sqrt_area: float) -> torch.Tensor: + ratio = dim / sqrt_area + left = (1.0 - ratio) * 1.0 / 2.0 + right = left + ratio * 1.0 + grid = np.linspace(left, right, dim // patch, endpoint=False) * _INTERP + return torch.from_numpy(grid).to(torch.float64) + + +def _video_t_grid(n: int, origin: float) -> torch.Tensor: + spans = torch.tensor( + [_FRAME_RESCALE * _FRAME_PER_TOKEN[k % _T_GROUP] for k in range(n)], + dtype=torch.float64, + ) + return origin + torch.cat([torch.zeros(1, dtype=torch.float64), spans[:-1].cumsum(0)]) + + +def _video_t_span(n: int) -> float: + # Sequential fp64 summation on purpose — see _temporal_position_span for + # why the two span implementations must not be unified. + return sum(_FRAME_RESCALE * _FRAME_PER_TOKEN[k % _T_GROUP] for k in range(n)) + + +def _range_for_slice(sl: slice) -> torch.Tensor: + return torch.arange(sl.start, sl.stop, dtype=torch.long) + + +def _cat_ranges(parts: list[torch.Tensor]) -> torch.Tensor: + if len(parts) == 1: + return parts[0] + if parts: + return torch.cat(parts) + return torch.empty(0, dtype=torch.long) + + +def minimax_h3_packed_sequence_ref2va_blocks( + *, + text_len: int, + latent_t: int, + latent_h: int, + latent_w: int, + audio_t: int, + ref_blocks: Sequence[Mapping[str, object]], + audio_channel: int = 2, + seq_len: int | None = None, +) -> dict[str, Any]: + """General ref2va-family packed layout. + + ``ref_blocks`` are consumed in request/plan order: + - ``{"kind": "image", "latent_h": H, "latent_w": W}`` + - ``{"kind": "audio", "ref_audio_t": T}`` + - ``{"kind": "video"|"video_audio", "ref_audio_t": T, + "latent_t": RT, "latent_h": RH, "latent_w": RW}`` + + Video-bearing blocks pack their audio rows immediately before their video + rows; both share the same temporal origin and advance by the longer of the + audio and video spans. Standalone audio advances the target origin by its + own T, and image blocks advance it by one integer slot. + """ + if not isinstance(ref_blocks, Sequence) or isinstance(ref_blocks, (str, bytes)): + raise ValueError("ref_blocks must be a sequence") + + parsed: list[dict[str, object]] = [] + ref_visual_rows = 0 + ref_audio_rows = 0 + for index, raw in enumerate(ref_blocks): + path = f"ref_blocks[{index}]" + if not isinstance(raw, Mapping): + raise ValueError(f"{path} must be an object") + kind = raw.get("kind", raw.get("type")) + if not isinstance(kind, str) or not kind: + raise ValueError(f"{path}.kind must be a non-empty string") + if kind == "image": + rh = _positive_int(raw, "latent_h", path) + rw = _positive_int(raw, "latent_w", path) + rows = (rh // _PATCH_H) * (rw // _PATCH_W) + item = {"kind": kind, "latent_h": rh, "latent_w": rw, "rows": rows} + ref_visual_rows += rows + elif kind == "audio": + rt = _positive_int(raw, "ref_audio_t", path, allow_zero=True) + rows = rt * audio_channel + item = {"kind": kind, "ref_audio_t": rt, "audio_rows": rows} + ref_audio_rows += rows + elif kind in ("video", "video_audio"): + rt = _positive_int(raw, "ref_audio_t", path, allow_zero=True) + vt = _positive_int(raw, "latent_t", path) + vh = _positive_int(raw, "latent_h", path) + vw = _positive_int(raw, "latent_w", path) + frame_rows = (vh // _PATCH_H) * (vw // _PATCH_W) + audio_rows = rt * audio_channel + video_rows = vt * frame_rows + item = { + "kind": kind, + "ref_audio_t": rt, + "latent_t": vt, + "latent_h": vh, + "latent_w": vw, + "frame_rows": frame_rows, + "audio_rows": audio_rows, + "video_rows": video_rows, + } + ref_audio_rows += audio_rows + ref_visual_rows += video_rows + else: + raise ValueError(f"{path}.kind unsupported for ref2va: {kind!r}") + parsed.append(item) + + ph, pw = latent_h // _PATCH_H, latent_w // _PATCH_W + frame_rows = ph * pw + video_rows = latent_t * frame_rows + audio_rows = audio_t * audio_channel + ref_rows = ref_visual_rows + ref_audio_rows + used = text_len + ref_rows + audio_rows + video_rows + if seq_len is None: + seq_len = ( + (used + MINIMAX_H3_PACKED_SEQUENCE_ALIGNMENT - 1) + // MINIMAX_H3_PACKED_SEQUENCE_ALIGNMENT + * MINIMAX_H3_PACKED_SEQUENCE_ALIGNMENT + ) + if seq_len < used: + raise ValueError(f"seq_len {seq_len} < used rows {used}") + + text_sl = slice(0, text_len) + cursor = text_len + block_slices: list[dict[str, object]] = [] + for item in parsed: + kind = str(item["kind"]) + if kind == "image": + rows = int(item["rows"]) + visual_sl = slice(cursor, cursor + rows) + cursor = visual_sl.stop + block_slices.append({**item, "visual_sl": visual_sl}) + elif kind == "audio": + rows = int(item["audio_rows"]) + audio_sl = slice(cursor, cursor + rows) + cursor = audio_sl.stop + block_slices.append({**item, "audio_sl": audio_sl}) + else: + a_rows = int(item["audio_rows"]) + v_rows = int(item["video_rows"]) + audio_sl = slice(cursor, cursor + a_rows) + visual_sl = slice(audio_sl.stop, audio_sl.stop + v_rows) + cursor = visual_sl.stop + block_slices.append({**item, "audio_sl": audio_sl, "visual_sl": visual_sl}) + + audio_sl = slice(cursor, cursor + audio_rows) + video_sl = slice(audio_sl.stop, audio_sl.stop + video_rows) + ref_img_pos_parts: list[torch.Tensor] = [] + ref_audio_pos_parts: list[torch.Tensor] = [] + g = torch.zeros(seq_len, 3, dtype=torch.float64) + g[text_sl, 0] = torch.arange(text_len, dtype=torch.float64) + + target_area = np.sqrt(latent_h * latent_w) + h_grid = _axis_from_sqrt_area(latent_h, _PATCH_H, target_area) + w_grid = _axis_from_sqrt_area(latent_w, _PATCH_W, target_area) + hh, ww = torch.meshgrid(h_grid, w_grid, indexing="ij") + target_frame = torch.stack([hh.reshape(-1), ww.reshape(-1)], dim=-1) + + t_cursor = float(text_len) + for item in block_slices: + kind = str(item["kind"]) + if kind == "image": + visual_sl = item["visual_sl"] + assert isinstance(visual_sl, slice) + ref_img_pos_parts.append(_range_for_slice(visual_sl)) + rh = int(item["latent_h"]) + rw = int(item["latent_w"]) + area = np.sqrt(rh * rw) + ref_hh, ref_ww = torch.meshgrid( + _axis_from_sqrt_area(rh, _PATCH_H, area), + _axis_from_sqrt_area(rw, _PATCH_W, area), + indexing="ij", + ) + g[visual_sl, 0] = t_cursor + g[visual_sl, 1] = ref_hh.reshape(-1) + g[visual_sl, 2] = ref_ww.reshape(-1) + t_cursor += 1.0 + elif kind == "audio": + audio_ref_sl = item["audio_sl"] + assert isinstance(audio_ref_sl, slice) + ref_t = int(item["ref_audio_t"]) + ref_audio_pos_parts.append(_range_for_slice(audio_ref_sl)) + ref_t_grid = t_cursor + torch.arange(ref_t, dtype=torch.float64) + g[audio_ref_sl, 0] = ref_t_grid.repeat(audio_channel) + if ref_t: + g[audio_ref_sl.start : audio_ref_sl.start + ref_t, 2] = float(w_grid[0]) + g[audio_ref_sl.start + ref_t : audio_ref_sl.stop, 2] = float(w_grid[-1]) + t_cursor += float(ref_t) + else: + audio_ref_sl = item["audio_sl"] + visual_sl = item["visual_sl"] + assert isinstance(audio_ref_sl, slice) + assert isinstance(visual_sl, slice) + ref_t = int(item["ref_audio_t"]) + vt = int(item["latent_t"]) + vh = int(item["latent_h"]) + vw = int(item["latent_w"]) + ref_audio_pos_parts.append(_range_for_slice(audio_ref_sl)) + ref_img_pos_parts.append(_range_for_slice(visual_sl)) + + ref_area = np.sqrt(vh * vw) + rv_h_grid = _axis_from_sqrt_area(vh, _PATCH_H, ref_area) + rv_w_grid = _axis_from_sqrt_area(vw, _PATCH_W, ref_area) + rv_hh, rv_ww = torch.meshgrid(rv_h_grid, rv_w_grid, indexing="ij") + + ref_t_grid = t_cursor + torch.arange(ref_t, dtype=torch.float64) + g[audio_ref_sl, 0] = ref_t_grid.repeat(audio_channel) + if ref_t: + g[audio_ref_sl.start : audio_ref_sl.start + ref_t, 2] = float(rv_w_grid[0]) + g[audio_ref_sl.start + ref_t : audio_ref_sl.stop, 2] = float(rv_w_grid[-1]) + + rv_frame = torch.stack([rv_hh.reshape(-1), rv_ww.reshape(-1)], dim=-1) + rv_g = g[visual_sl].view(vt, int(item["frame_rows"]), 3) + rv_g[:, :, 0] = _video_t_grid(vt, t_cursor)[:, None] + rv_g[:, :, 1:] = rv_frame[None] + t_cursor += max(float(ref_t), _video_t_span(vt)) + + audio_t_grid = t_cursor + torch.arange(audio_t, dtype=torch.float64) + g[audio_sl, 0] = audio_t_grid.repeat(audio_channel) + g[audio_sl.start : audio_sl.start + audio_t, 2] = float(w_grid[0]) + g[audio_sl.start + audio_t : audio_sl.stop, 2] = float(w_grid[-1]) + + video_g = g[video_sl].view(latent_t, frame_rows, 3) + video_g[:, :, 0] = _video_t_grid(latent_t, t_cursor)[:, None] + video_g[:, :, 1:] = target_frame[None] + + target_img_pos = _range_for_slice(video_sl) + target_audio_pos = _range_for_slice(audio_sl) + img_pos = _cat_ranges(ref_img_pos_parts + [target_img_pos]) + audio_pos = _cat_ranges(ref_audio_pos_parts + [target_audio_pos]) + + update_mask = torch.zeros(img_pos.shape[0], dtype=torch.bool) + update_mask[ref_visual_rows:] = True + audio_update_mask = torch.zeros(audio_pos.shape[0], dtype=torch.bool) + audio_update_mask[ref_audio_rows:] = True + text_pos = torch.arange(0, text_len) + + token_tags = torch.full((seq_len,), -1, dtype=torch.long) # PADDING + token_tags[text_sl] = 1 # TEXT + token_tags[audio_pos] = 2 # AUDIO (refs + target) + token_tags[img_pos] = 0 # VIDEO (refs + target) + + cu = torch.tensor([0, used, seq_len], dtype=torch.int32) + return { + "seq_len": seq_len, + "img_pos": img_pos, + "audio_pos": audio_pos, + "text_pos": text_pos, + "update_mask": update_mask, + "audio_update_mask": audio_update_mask, + "img_position_ids": g, + "token_tags": token_tags, + "cu_seqlens": cu, + } + + +__all__ = [ + "minimax_h3_packed_sequence", + "minimax_h3_packed_sequence_ref2va_blocks", +] diff --git a/telefuser/pipelines/minimax_h3/packed_tokens.py b/telefuser/pipelines/minimax_h3/packed_tokens.py new file mode 100644 index 0000000..b4a8292 --- /dev/null +++ b/telefuser/pipelines/minimax_h3/packed_tokens.py @@ -0,0 +1,98 @@ +# SPDX-License-Identifier: Apache-2.0 +from __future__ import annotations + +from collections.abc import Sequence + +import torch + + +def _int_tuple(value: Sequence[int], name: str, length: int) -> tuple[int, ...]: + if len(value) != length: + raise ValueError(f"{name} must have length {length}, got {list(value)!r}") + out = tuple(int(item) for item in value) + if any(item <= 0 for item in out): + raise ValueError(f"{name} values must be positive, got {list(value)!r}") + return out + + +def _rank(tensor: torch.Tensor, name: str, rank: int) -> None: + if tensor.ndim != rank: + raise ValueError(f"{name} must be rank {rank}, got shape={list(tensor.shape)}") + + +def minimax_h3_patchify_video_latent( + latent: torch.Tensor, + *, + patch_size: Sequence[int], +) -> torch.Tensor: + """Pack SGLang video latent [B,C,T,H,W] into DiT token rows.""" + + _rank(latent, "video latent", 5) + pt, ph, pw = _int_tuple(patch_size, "patch_size", 3) + batch, channel, full_t, full_h, full_w = (int(dim) for dim in latent.shape) + if full_t % pt or full_h % ph or full_w % pw: + raise ValueError( + "video latent spatial/time dims must be divisible by patch_size: " + f"shape={list(latent.shape)}, patch_size={[pt, ph, pw]}" + ) + t, h, w = full_t // pt, full_h // ph, full_w // pw + packed = latent.reshape(batch, channel, t, pt, h, ph, w, pw) + packed = torch.einsum("nctrhpwq->nthwcrpq", packed) + return packed.reshape(batch * t * h * w, channel * pt * ph * pw).contiguous() + + +def minimax_h3_unpatchify_video_tokens( + rows: torch.Tensor, + *, + latent_shape: Sequence[int], + patch_size: Sequence[int], +) -> torch.Tensor: + """Unpack DiT video token rows into SGLang latent [B,C,T,H,W].""" + + _rank(rows, "video token rows", 2) + t, h, w, channel = _int_tuple(latent_shape, "latent_shape", 4) + pt, ph, pw = _int_tuple(patch_size, "patch_size", 3) + expected_dim = pt * ph * pw * channel + if int(rows.shape[-1]) != expected_dim: + raise ValueError( + f"video token dim {int(rows.shape[-1])} != patch volume * channel " + f"{expected_dim} for latent_shape={list(latent_shape)}, " + f"patch_size={[pt, ph, pw]}" + ) + rows_per_sample = t * h * w + if int(rows.shape[0]) % rows_per_sample: + raise ValueError( + f"video rows {int(rows.shape[0])} must be divisible by t*h*w " + f"{rows_per_sample} for latent_shape={list(latent_shape)}" + ) + packed = rows.reshape(-1, t, h, w, channel, pt, ph, pw) + latent = torch.einsum("nthwcrpq->nctrhpwq", packed) + return latent.reshape(-1, channel, t * pt, h * ph, w * pw).contiguous() + + +def minimax_h3_unpack_audio_tokens( + rows: torch.Tensor, + *, + audio_t: int, + audio_channel: int, +) -> torch.Tensor: + """Unpack DiT audio token rows into SGLang audio VAE latent [C,latent_dim,T].""" + + _rank(rows, "audio token rows", 2) + audio_t = int(audio_t) + audio_channel = int(audio_channel) + if audio_t <= 0 or audio_channel <= 0: + raise ValueError(f"audio_t and audio_channel must be positive, got {audio_t=} {audio_channel=}") + if int(rows.shape[0]) != audio_t: + raise ValueError(f"audio rows {int(rows.shape[0])} != audio_t {audio_t}") + if audio_t % audio_channel: + raise ValueError(f"audio_t must be divisible by audio_channel, got {audio_t=} {audio_channel=}") + native = rows.reshape(audio_channel, audio_t // audio_channel, int(rows.shape[-1])) + return native.permute(0, 2, 1).contiguous() + + +__all__ = [ + "minimax_h3_patchify_video_latent", + "minimax_h3_unpack_audio_tokens", + "minimax_h3_unpatchify_video_tokens", +] diff --git a/telefuser/pipelines/minimax_h3/pipeline.py b/telefuser/pipelines/minimax_h3/pipeline.py new file mode 100644 index 0000000..083f9ab --- /dev/null +++ b/telefuser/pipelines/minimax_h3/pipeline.py @@ -0,0 +1,267 @@ +# SPDX-License-Identifier: Apache-2.0 +"""MiniMax H3 local T2VA, FL2VA, and Ref2VA pipeline.""" + +from __future__ import annotations + +from contextlib import ExitStack +from dataclasses import dataclass, field, replace +from pathlib import Path +from typing import Any + +import torch +import torch.distributed as dist +from transformers import AutoProcessor + +from telefuser.core.base_pipeline import BasePipeline +from telefuser.core.config import ModelRuntimeConfig +from telefuser.core.module_manager import ModuleManager +from telefuser.worker import ParallelWorker + +from .data import ( + minimax_h3_validate_canonical_request, + minimax_h3_validate_reference_media_facts, +) +from .denoising import MiniMaxH3DenoisingStage +from .material_io import ( + MiniMaxH3MaterialFacts, + minimax_h3_localize_material, + minimax_h3_probe_material, +) +from .resolved_plan import ( + MiniMaxH3ResolvedPlan, + minimax_h3_resolve_plan, + minimax_h3_resolve_spatial_shape, +) +from .text_encoding import MiniMaxH3TextEncodingStage +from .vae import MiniMaxH3PreparedCondition, MiniMaxH3VAEStage + + +def _fp32_runtime_config() -> ModelRuntimeConfig: + return ModelRuntimeConfig(torch_dtype=torch.float32) + + +@dataclass +class MiniMaxH3PipelineConfig: + processor_path: str + text_encoder_config: ModelRuntimeConfig = field(default_factory=ModelRuntimeConfig) + dit_config: ModelRuntimeConfig = field(default_factory=ModelRuntimeConfig) + video_vae_config: ModelRuntimeConfig = field(default_factory=_fp32_runtime_config) + audio_vae_config: ModelRuntimeConfig = field(default_factory=_fp32_runtime_config) + num_inference_steps: int = 50 + + def __post_init__(self) -> None: + if not self.processor_path: + raise ValueError("processor_path is required") + if self.num_inference_steps < 2: + raise ValueError("num_inference_steps must be at least 2") + + +@dataclass(frozen=True) +class MiniMaxH3Generation: + video: torch.Tensor + audio: torch.Tensor + video_fps: int + audio_sample_rate: int + plan: MiniMaxH3ResolvedPlan + packed_sequence_length: int + runtime_metrics: dict[str, float | int] + + +class MiniMaxH3Pipeline(BasePipeline): + def __init__(self, device: str | torch.device = "cuda", torch_dtype: torch.dtype = torch.bfloat16) -> None: + super().__init__(device=device, torch_dtype=torch_dtype) + self.config: MiniMaxH3PipelineConfig | None = None + self.text_stage: MiniMaxH3TextEncodingStage | None = None + self.vae_stage: MiniMaxH3VAEStage | None = None + self.denoising_stage: MiniMaxH3DenoisingStage | ParallelWorker | None = None + + def init(self, module_manager: ModuleManager, config: MiniMaxH3PipelineConfig) -> None: + self.config = config + processor = AutoProcessor.from_pretrained( + config.processor_path, + local_files_only=True, + trust_remote_code=False, + ) + self.text_stage = MiniMaxH3TextEncodingStage( + module_manager, + config.text_encoder_config, + processor=processor, + ) + self.vae_stage = MiniMaxH3VAEStage( + module_manager, + config.video_vae_config, + config.audio_vae_config, + ) + denoising_stage = MiniMaxH3DenoisingStage(module_manager, config.dit_config) + if config.dit_config.parallel_config.world_size > 1 and not dist.is_initialized(): + self.denoising_stage = ParallelWorker(denoising_stage) + else: + self.denoising_stage = denoising_stage + if dist.is_initialized(): + denoising_stage.parallel_models() + self._model_info = module_manager.get_model_info() + + def _get_stages(self) -> list[object]: + return [stage for stage in (self.text_stage, self.vae_stage, self.denoising_stage) if stage is not None] + + def stop(self) -> None: + for stage in self._get_stages(): + close = getattr(stage, "close", None) + if callable(close): + close() + + @staticmethod + def _resolve_stage_result(value: Any) -> Any: + return value() if callable(value) else value + + @staticmethod + def _resolve_deferred_plan( + canonical: dict[str, Any], + facts: dict[int, MiniMaxH3MaterialFacts], + ) -> tuple[dict[str, Any], MiniMaxH3ResolvedPlan]: + if canonical["target"].get("duration_seconds") is None: + duration_sources = [ + (index, condition) + for index, condition in enumerate(canonical["conditions"]) + if condition["type"] in {"audio", "video", "video_audio"} + ] + index, condition = duration_sources[0] + duration = facts[index].duration_seconds + if duration is None: + raise ValueError(f"conditions[{index}] has no probed duration") + effective = duration - float(condition.get("start_time_seconds", 0.0)) + canonical = minimax_h3_validate_canonical_request( + task=canonical["task"], + prompt=canonical["prompt"], + conditions=canonical["conditions"], + target={**canonical["target"], "duration_seconds": effective}, + flow_shift=canonical.get("flow_shift"), + audio_flow_shift=canonical.get("audio_flow_shift"), + seed=canonical.get("seed"), + ) + plan = minimax_h3_resolve_plan(canonical) + if plan.shape.get("geometry") == "deferred": + first = plan.materials[0] + item_facts = facts[int(first.condition_index)] + if item_facts.width is None or item_facts.height is None: + raise ValueError("deferred FL2VA geometry requires image width and height") + shape = dict(plan.shape) + shape.update( + minimax_h3_resolve_spatial_shape( + width=item_facts.width, + height=item_facts.height, + ) + ) + plan = replace(plan, shape=shape) + if plan.shape.get("geometry") != "resolved_v2": + raise ValueError("MiniMax H3 target geometry must resolve before model execution") + return canonical, plan + + @staticmethod + def _condition_labels( + prepared: list[MiniMaxH3PreparedCondition], + ) -> list[tuple[str, int]]: + counters = {"image": 0, "audio": 0, "video": 0} + labels: list[tuple[str, int]] = [] + for condition in prepared: + if condition.kind == "image": + counters["image"] += 1 + labels.append(("image", counters["image"])) + elif condition.kind == "audio": + counters["audio"] += 1 + labels.append(("audio", counters["audio"])) + elif condition.kind in {"video", "video_audio"}: + if condition.has_audio: + counters["audio"] += 1 + labels.append(("audio", counters["audio"])) + counters["video"] += 1 + labels.append(("video", counters["video"])) + return labels + + @torch.inference_mode() + def __call__( + self, + *, + task: str, + prompt: str, + conditions: list[dict[str, Any]] | None, + target: dict[str, Any], + seed: int | None = None, + flow_shift: float | None = None, + audio_flow_shift: float | None = None, + num_inference_steps: int | None = None, + ) -> MiniMaxH3Generation: + if self.config is None or self.text_stage is None or self.vae_stage is None or self.denoising_stage is None: + raise RuntimeError("MiniMaxH3Pipeline.init must be called before generation") + canonical = minimax_h3_validate_canonical_request( + task=task, + prompt=prompt, + conditions=[] if conditions is None else conditions, + target=target, + seed=seed, + flow_shift=flow_shift, + audio_flow_shift=audio_flow_shift, + ) + with ExitStack() as stack: + paths: dict[int, Path] = {} + facts: dict[int, MiniMaxH3MaterialFacts] = {} + for index, condition in enumerate(canonical["conditions"]): + path = stack.enter_context(minimax_h3_localize_material(condition["uri"])) + paths[index] = path + facts[index] = minimax_h3_probe_material(path, condition["type"]) + duration = facts[index].duration_seconds + start = float(condition.get("start_time_seconds", 0.0)) + if duration is not None and start >= duration: + raise ValueError(f"conditions[{index}].start_time_seconds must be less than media duration") + if task == "ref2va": + duration_facts = { + index: float(item.duration_seconds) + for index, item in facts.items() + if item.duration_seconds is not None + } + minimax_h3_validate_reference_media_facts(canonical["conditions"], duration_facts) + canonical, plan = self._resolve_deferred_plan(canonical, facts) + prepared = self.vae_stage.prepare_media(plan, paths, facts) + images = [item.image for item in prepared if item.image is not None] + videos = [item.video_frames for item in prepared if item.video_frames is not None] + text = self.text_stage.encode( + task=plan.task, + prompt=plan.prompt, + images=images, + videos=videos, + condition_labels=self._condition_labels(prepared), + ) + if any(item.image is not None or item.video_frames is not None for item in prepared): + prepared = self.vae_stage.encode_visual(prepared) + duration_seconds = float(plan.shape["frame_count"]) / float(plan.shape["fps"]) + if any(item.has_audio for item in prepared): + prepared = self.vae_stage.encode_audio(prepared, paths, facts, duration_seconds) + steps = self.config.num_inference_steps if num_inference_steps is None else num_inference_steps + if isinstance(steps, bool) or not isinstance(steps, int) or steps < 2: + raise ValueError("num_inference_steps must be an integer of at least 2") + denoised = self._resolve_stage_result( + self.denoising_stage.denoise( + plan=plan, + text=text, + conditions=prepared, + num_inference_steps=steps, + ) + ) + video = self.vae_stage.decode_video(denoised.video_latent) + audio = self.vae_stage.decode_audio(denoised.audio_latent) + return MiniMaxH3Generation( + video=video[:, : int(plan.shape["frame_count"])], + audio=audio, + video_fps=24, + audio_sample_rate=32_000, + plan=plan, + packed_sequence_length=int(denoised.packed["seq_len"]), + runtime_metrics=denoised.runtime_metrics, + ) + + +__all__ = [ + "MiniMaxH3Generation", + "MiniMaxH3Pipeline", + "MiniMaxH3PipelineConfig", +] diff --git a/telefuser/pipelines/minimax_h3/presentation.py b/telefuser/pipelines/minimax_h3/presentation.py new file mode 100644 index 0000000..9dcd383 --- /dev/null +++ b/telefuser/pipelines/minimax_h3/presentation.py @@ -0,0 +1,274 @@ +# SPDX-License-Identifier: Apache-2.0 +"""MiniMax H3 Qwen presentation building. + +Builds the positive presentation token stream: +- fl2va: ': ' label + vision block (<|vision_start|> + + N*<|image_pad|> + <|vision_end|>) + prompt text. +- t2va: prompt text only (no vision block). +Prompt text passes through verbatim (no stripping or rewriting). + +All presentation variants are emitted through the shared ``_Presentation`` +accumulator so ids and AdaLN token tags cannot drift apart. +""" + +from __future__ import annotations + +from collections.abc import Sequence +from typing import Any + +import torch + +VISION_START = "<|vision_start|>" +VISION_END = "<|vision_end|>" +IMAGE_PAD = "<|image_pad|>" +VIDEO_PAD = "<|video_pad|>" + +_TEXT_TAG = 1 +_VIDEO_TAG = 0 + + +def _text_ids(tokenizer: Any, text: str) -> list[int]: + return list(tokenizer(text, add_special_tokens=False)["input_ids"]) + + +def _vision_block_ids(tokenizer: Any, pad_token: str, count: int) -> list[int]: + return ( + [tokenizer.convert_tokens_to_ids(VISION_START)] + + [tokenizer.convert_tokens_to_ids(pad_token)] * int(count) + + [tokenizer.convert_tokens_to_ids(VISION_END)] + ) + + +class _Presentation: + """Accumulates aligned (ids, token_tags) presentation segments.""" + + def __init__(self) -> None: + self.ids: list[int] = [] + self.tags: list[int] = [] + + def text(self, token_ids: list[int]) -> None: + self.ids += token_ids + self.tags += [_TEXT_TAG] * len(token_ids) + + def vision(self, token_ids: list[int]) -> None: + self.ids += token_ids + self.tags += [_VIDEO_TAG] * len(token_ids) + + def build(self) -> tuple[torch.Tensor, torch.Tensor]: + return ( + torch.tensor(self.ids, dtype=torch.long), + torch.tensor(self.tags, dtype=torch.long), + ) + + +def _timestamped_video_blocks( + presentation: _Presentation, + tokenizer: Any, + *, + counts: Sequence[int], + timestamps: Sequence[float], + context: str, +) -> None: + """Emit per-temporal-block ``<{t:.1f} seconds>`` text + VIDEO vision.""" + + counts = [int(value) for value in counts] + timestamps = [float(value) for value in timestamps] + if not counts or len(counts) != len(timestamps): + raise ValueError(f"{context}video block token counts and timestamps must align") + for count, timestamp in zip(counts, timestamps): + if count <= 0: + raise ValueError(f"{context}video block token count must be positive") + presentation.text(_text_ids(tokenizer, f"<{timestamp:.1f} seconds>")) + presentation.vision(_vision_block_ids(tokenizer, VIDEO_PAD, count)) + + +def minimax_h3_text_only_ids(tokenizer: Any, prompt: str) -> torch.Tensor: + """t2va presentation: verbatim prompt, no special tokens.""" + if not prompt: + raise ValueError("prompt must be non-empty") + return torch.tensor(_text_ids(tokenizer, prompt), dtype=torch.long) + + +def minimax_h3_multi_image_presentation( + tokenizer: Any, + *, + prompt: str, + image_token_counts: list[int], +) -> tuple[torch.Tensor, torch.Tensor]: + if not image_token_counts: + raise ValueError("image_token_counts must be non-empty") + presentation = _Presentation() + for index, count in enumerate(image_token_counts, start=1): + if int(count) <= 0: + raise ValueError("image_token_count must be positive") + presentation.text(_text_ids(tokenizer, f": ")) + presentation.vision(_vision_block_ids(tokenizer, IMAGE_PAD, count)) + presentation.text(_text_ids(tokenizer, prompt)) + return presentation.build() + + +def minimax_h3_ref2va_presentation( + tokenizer: Any, + *, + prompt: str, + condition_labels: list[tuple[str, int]], + image_token_count: int | list[int] | None, +) -> tuple[torch.Tensor, torch.Tensor]: + """ref2va positive presentation: + + per condition in request order — image i: ``: `` label followed + by the vision block; audio j: ``