FL2VA/video_vae/norm.py
10.4 KB · 358 lines · python Raw
1 # SPDX-License-Identifier: Apache-2.0
2 # Torch-native normalization for the MiniMax H3 visual VAE.
3 import math
4 import os
5
6 import torch
7 import torch.distributed as dist
8 import torch.nn as nn
9 import torch.nn.functional as F
10
11 from .conv import SpatialParallelConv3d
12 from .parallel import all_reduce, get_parallel_state
13
14
15 def _validate_activation(activation):
16 valid_activations = {"identity", "silu", "relu"}
17 if activation not in valid_activations:
18 raise ValueError(
19 f"Unsupported activation: {activation}. Supported: {valid_activations}"
20 )
21
22
23 def _apply_activation(x, activation):
24 _validate_activation(activation)
25 if activation == "identity":
26 return x
27 if activation == "silu":
28 return F.silu(x)
29 return F.relu(x)
30
31
32 def _merge_time_to_batch(x):
33 batch, channels, depth, height, width = x.shape
34 return (
35 x.permute(0, 2, 1, 3, 4)
36 .contiguous()
37 .view(batch * depth, channels, 1, height, width)
38 )
39
40
41 def _split_time_from_batch(x, batch):
42 batch_depth, channels, _, height, width = x.shape
43 depth = batch_depth // batch
44 return (
45 x.view(batch, depth, channels, height, width)
46 .permute(0, 2, 1, 3, 4)
47 .contiguous()
48 )
49
50
51 def fused_group_norm(x, num_groups, weight, bias, eps=1e-5, activation="silu"):
52 out = F.group_norm(x, num_groups, weight=weight, bias=bias, eps=eps)
53 return _apply_activation(out, activation)
54
55
56 def fused_spatial_norm(
57 f,
58 num_groups,
59 norm_weight,
60 norm_bias,
61 dynamic_scale,
62 dynamic_bias,
63 eps=1e-5,
64 activation="silu",
65 ):
66 norm_f = F.group_norm(
67 f,
68 num_groups,
69 weight=norm_weight,
70 bias=norm_bias,
71 eps=eps,
72 )
73 out = norm_f * dynamic_scale + dynamic_bias
74 return _apply_activation(out, activation)
75
76
77 class DummyAffine(torch.nn.Module):
78 def __init__(self, num_channels, affine=True):
79 super().__init__()
80 if affine:
81 self.weight = torch.nn.Parameter(torch.ones(num_channels))
82 self.bias = torch.nn.Parameter(torch.zeros(num_channels))
83 else:
84 self.register_parameter("weight", None)
85 self.register_parameter("bias", None)
86
87 def forward(self, input):
88 if self.weight is None:
89 return input
90 shape = [1, -1] + [1] * (input.dim() - 2)
91 return input * self.weight.view(*shape) + self.bias.view(*shape)
92
93
94 class FusedGroupNorm3D(torch.nn.Module):
95 """Compatibility wrapper implemented with native PyTorch ops."""
96
97 def __init__(
98 self,
99 num_groups,
100 num_channels,
101 eps=1e-5,
102 affine=True,
103 activation="silu",
104 cond_channels=None,
105 use_t_isolated_gn=False,
106 padding_mode="zeros",
107 padding_mode_t=None,
108 causal=True,
109 ):
110 super().__init__()
111 _validate_activation(activation)
112 self.num_groups = num_groups
113 self.num_channels = num_channels
114 self.eps = eps
115 self.affine = affine
116 self.activation = activation
117 self.use_t_isolated_gn = use_t_isolated_gn
118
119 if cond_channels is not None:
120 self.use_spatial_affine = True
121 self.norm_layer = DummyAffine(num_channels, affine=affine)
122 self.conv_y = SpatialParallelConv3d(
123 cond_channels,
124 num_channels,
125 kernel_size=1,
126 padding_mode=padding_mode,
127 padding_mode_t=padding_mode_t,
128 causal=causal,
129 )
130 self.conv_b = SpatialParallelConv3d(
131 cond_channels,
132 num_channels,
133 kernel_size=1,
134 padding_mode=padding_mode,
135 padding_mode_t=padding_mode_t,
136 causal=causal,
137 )
138 else:
139 self.use_spatial_affine = False
140 if self.affine:
141 self.weight = torch.nn.Parameter(torch.ones(num_channels))
142 self.bias = torch.nn.Parameter(torch.zeros(num_channels))
143 else:
144 self.register_parameter("weight", None)
145 self.register_parameter("bias", None)
146
147 def forward(self, f, cond=None):
148 need_reshape = self.use_t_isolated_gn and f.dim() == 5
149 batch = f.shape[0] if need_reshape else None
150 f_size = f.shape[-3:]
151 if need_reshape:
152 f = _merge_time_to_batch(f)
153
154 if self.use_spatial_affine:
155 scale = self.conv_y(cond)
156 bias = self.conv_b(cond)
157 if math.prod(scale.shape[-3:]) * math.prod(bias.shape[-3:]) > 1:
158 scale = F.interpolate(scale, size=f_size, mode="nearest")
159 bias = F.interpolate(bias, size=f_size, mode="nearest")
160 if need_reshape:
161 scale = _merge_time_to_batch(scale)
162 bias = _merge_time_to_batch(bias)
163 out = fused_spatial_norm(
164 f,
165 self.num_groups,
166 self.norm_layer.weight,
167 self.norm_layer.bias,
168 scale,
169 bias,
170 self.eps,
171 self.activation,
172 )
173 else:
174 if cond is not None:
175 raise NotImplementedError("Dynamic affine is not defined")
176 weight = self.weight if self.affine else None
177 bias = self.bias if self.affine else None
178 out = fused_group_norm(
179 f, self.num_groups, weight, bias, self.eps, self.activation
180 )
181
182 if need_reshape:
183 out = _split_time_from_batch(out, batch)
184 return out
185
186
187 class SpatialParallelGroupNorm(nn.GroupNorm):
188 def __init__(
189 self,
190 *args,
191 **kwargs,
192 ):
193 super().__init__(*args, **kwargs)
194 self.spatial_parallel = False
195
196 def _compute_stats(self, input):
197 batch, channels = input.shape[0], input.shape[1]
198 spatial_dims = input.shape[2:]
199 spatial_size = math.prod(spatial_dims)
200
201 groups = self.num_groups
202 x = input.reshape(batch, groups, channels // groups, -1).to(torch.float32)
203
204 local_sum = x.sum(dim=(2, 3))
205 local_square_sum = (x * x).sum(dim=(2, 3))
206 local_n = (channels // groups) * spatial_size
207 local_n_tensor = torch.full_like(local_sum, float(local_n))
208
209 stats = torch.stack([local_sum, local_square_sum, local_n_tensor], dim=0)
210
211 local_process_group = get_parallel_state()["local_process_group"]
212 stats = all_reduce(stats, dist.ReduceOp.SUM, local_process_group)
213
214 total_sum = stats[0]
215 total_square_sum = stats[1]
216 total_n = stats[2]
217
218 mean = total_sum / total_n
219 var = (total_square_sum / total_n) - mean**2
220 return mean, var
221
222 def forward(self, input):
223 if not self.spatial_parallel:
224 return nn.GroupNorm.forward(self, input)
225
226 batch, channels = input.shape[0], input.shape[1]
227 orig_shape = input.shape
228
229 mean, var = self._compute_stats(input)
230 x = input.reshape(batch, self.num_groups, channels // self.num_groups, -1)
231
232 mean = mean.unsqueeze(-1).unsqueeze(-1)
233 var = var.unsqueeze(-1).unsqueeze(-1)
234 x = (x - mean) / torch.sqrt(var + self.eps)
235 x = x.reshape(orig_shape)
236
237 if self.affine:
238 shape = [1, -1] + [1] * (len(orig_shape) - 2)
239 x *= self.weight.view(*shape)
240 x += self.bias.view(*shape)
241
242 return x
243
244
245 class TemporalIsolatedSpatialParallelGroupNorm(SpatialParallelGroupNorm):
246 def forward(self, input):
247 if input.dim() == 5:
248 batch = input.shape[0]
249 input = _merge_time_to_batch(input)
250 output = super().forward(input)
251 return _split_time_from_batch(output, batch)
252 return super().forward(input)
253
254
255
256
257
258
259
260
261 class SpatialNorm3D(nn.Module):
262 def __init__(
263 self,
264 f_channels,
265 zq_channels,
266 padding_mode="zeros",
267 padding_mode_t=None,
268 causal=True,
269 use_t_isolated_gn=False,
270 ):
271 super().__init__()
272 norm_cls = (
273 TemporalIsolatedSpatialParallelGroupNorm
274 if use_t_isolated_gn
275 else SpatialParallelGroupNorm
276 )
277 self.norm_layer = norm_cls(
278 num_groups=32, num_channels=f_channels, eps=1e-6, affine=True
279 )
280
281 self.conv_y = SpatialParallelConv3d(
282 zq_channels,
283 f_channels,
284 kernel_size=1,
285 padding_mode=padding_mode,
286 padding_mode_t=padding_mode_t,
287 causal=causal,
288 )
289 self.conv_b = SpatialParallelConv3d(
290 zq_channels,
291 f_channels,
292 kernel_size=1,
293 padding_mode=padding_mode,
294 padding_mode_t=padding_mode_t,
295 causal=causal,
296 )
297
298 def forward(self, f, zq):
299 f_size = f.shape[-3:]
300 norm_f = self.norm_layer(f)
301 scale = self.conv_y(zq)
302 bias = self.conv_b(zq)
303
304 if math.prod(scale.shape[-3:]) * math.prod(bias.shape[-3:]) > 1:
305 scale = F.interpolate(scale, size=f_size, mode="nearest")
306 bias = F.interpolate(bias, size=f_size, mode="nearest")
307
308 return norm_f * scale + bias
309
310
311 def get_spatial_norm_3d(
312 num_channels,
313 cond_channels,
314 *,
315 padding_mode="zeros",
316 padding_mode_t=None,
317 causal=True,
318 use_t_isolated_gn=False,
319 ):
320 if os.environ.get("MINIMAX_H3_USE_FUSED_NORM", "false").lower() == "true":
321 return FusedGroupNorm3D(
322 num_groups=32,
323 num_channels=num_channels,
324 eps=1e-6,
325 affine=True,
326 cond_channels=cond_channels,
327 use_t_isolated_gn=use_t_isolated_gn,
328 padding_mode=padding_mode,
329 padding_mode_t=padding_mode_t,
330 causal=causal,
331 )
332 return SpatialNorm3D(
333 num_channels,
334 cond_channels,
335 padding_mode=padding_mode,
336 padding_mode_t=padding_mode_t,
337 causal=causal,
338 use_t_isolated_gn=use_t_isolated_gn,
339 )
340
341
342 def get_group_norm_3d(num_channels, use_t_isolated_gn=False):
343 if os.environ.get("MINIMAX_H3_USE_FUSED_NORM", "false").lower() == "true":
344 return FusedGroupNorm3D(
345 num_groups=32,
346 num_channels=num_channels,
347 eps=1e-6,
348 affine=True,
349 use_t_isolated_gn=use_t_isolated_gn,
350 )
351
352 norm_cls = (
353 TemporalIsolatedSpatialParallelGroupNorm
354 if use_t_isolated_gn
355 else SpatialParallelGroupNorm
356 )
357 return norm_cls(num_groups=32, num_channels=num_channels, eps=1e-6, affine=True)
358