# SPDX-License-Identifier: Apache-2.0 # Torch-native normalization for the MiniMax H3 visual VAE. import math import os import torch import torch.distributed as dist import torch.nn as nn import torch.nn.functional as F from .conv import SpatialParallelConv3d from .parallel import all_reduce, get_parallel_state 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 = SpatialParallelConv3d( cond_channels, num_channels, kernel_size=1, padding_mode=padding_mode, padding_mode_t=padding_mode_t, causal=causal, ) self.conv_b = SpatialParallelConv3d( 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 SpatialParallelGroupNorm(nn.GroupNorm): def __init__( self, *args, **kwargs, ): super().__init__(*args, **kwargs) self.spatial_parallel = False def _compute_stats(self, input): batch, channels = input.shape[0], input.shape[1] spatial_dims = input.shape[2:] spatial_size = math.prod(spatial_dims) groups = self.num_groups x = input.reshape(batch, groups, channels // groups, -1).to(torch.float32) local_sum = x.sum(dim=(2, 3)) local_square_sum = (x * x).sum(dim=(2, 3)) local_n = (channels // groups) * spatial_size local_n_tensor = torch.full_like(local_sum, float(local_n)) stats = torch.stack([local_sum, local_square_sum, local_n_tensor], dim=0) local_process_group = get_parallel_state()["local_process_group"] stats = all_reduce(stats, dist.ReduceOp.SUM, local_process_group) total_sum = stats[0] total_square_sum = stats[1] total_n = stats[2] mean = total_sum / total_n var = (total_square_sum / total_n) - mean**2 return mean, var def forward(self, input): if not self.spatial_parallel: return nn.GroupNorm.forward(self, input) batch, channels = input.shape[0], input.shape[1] orig_shape = input.shape mean, var = self._compute_stats(input) x = input.reshape(batch, self.num_groups, channels // self.num_groups, -1) mean = mean.unsqueeze(-1).unsqueeze(-1) var = var.unsqueeze(-1).unsqueeze(-1) x = (x - mean) / torch.sqrt(var + self.eps) x = x.reshape(orig_shape) if self.affine: shape = [1, -1] + [1] * (len(orig_shape) - 2) x *= self.weight.view(*shape) x += self.bias.view(*shape) return x class TemporalIsolatedSpatialParallelGroupNorm(SpatialParallelGroupNorm): 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 = ( TemporalIsolatedSpatialParallelGroupNorm if use_t_isolated_gn else SpatialParallelGroupNorm ) self.norm_layer = norm_cls( num_groups=32, num_channels=f_channels, eps=1e-6, affine=True ) self.conv_y = SpatialParallelConv3d( zq_channels, f_channels, kernel_size=1, padding_mode=padding_mode, padding_mode_t=padding_mode_t, causal=causal, ) self.conv_b = SpatialParallelConv3d( 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, ): if os.environ.get("MINIMAX_H3_USE_FUSED_NORM", "false").lower() == "true": return FusedGroupNorm3D( num_groups=32, num_channels=num_channels, eps=1e-6, affine=True, cond_channels=cond_channels, use_t_isolated_gn=use_t_isolated_gn, padding_mode=padding_mode, padding_mode_t=padding_mode_t, causal=causal, ) 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): if os.environ.get("MINIMAX_H3_USE_FUSED_NORM", "false").lower() == "true": return FusedGroupNorm3D( num_groups=32, num_channels=num_channels, eps=1e-6, affine=True, use_t_isolated_gn=use_t_isolated_gn, ) norm_cls = ( TemporalIsolatedSpatialParallelGroupNorm if use_t_isolated_gn else SpatialParallelGroupNorm ) return norm_cls(num_groups=32, num_channels=num_channels, eps=1e-6, affine=True)