FL2VA/video_vae/conv.py
4.4 KB · 160 lines · python Raw
1 # SPDX-License-Identifier: Apache-2.0
2 # Spatial-parallel 3D convolution for the MiniMax H3 visual VAE.
3 import torch
4 import torch.nn as nn
5 import torch.nn.functional as F
6
7 from .parallel import get_parallel_state, exchange_borders
8
9
10
11
12 class BaseConv3d(nn.Conv3d):
13 def __init__(
14 self,
15 in_channels,
16 out_channels,
17 kernel_size,
18 stride=1,
19 padding=0,
20 bias=True,
21 padding_mode="zeros",
22 padding_mode_t=None,
23 causal=True,
24 ):
25 super().__init__(
26 in_channels,
27 out_channels,
28 kernel_size=kernel_size,
29 stride=stride,
30 padding=padding,
31 bias=bias,
32 padding_mode=padding_mode,
33 )
34 padding_mode = "constant" if padding_mode == "zeros" else padding_mode
35 padding_mode_t = "constant" if padding_mode_t == "zeros" else padding_mode_t
36 self.pad_mode = padding_mode
37 self.pad_mode_t = padding_mode_t or ("constant" if causal else "replicate")
38 self.causal = causal
39
40 def _apply_temporal_padding(self, x):
41 B, C, D, H, W = x.shape
42 if D > 1:
43 pad_size = (
44 0,
45 0,
46 0,
47 0,
48 self.padding[0] * 2 if self.causal else self.padding[0],
49 0 if self.causal else self.padding[0],
50 )
51 return F.pad(x, pad_size, mode=self.pad_mode_t)
52 else:
53 if self.pad_mode_t == "constant":
54 assert self.causal, "Zeros padding is only supported for causal mode"
55 zeros = torch.zeros_like(x[:, :, :1, :, :]).expand(
56 -1, -1, self.kernel_size[0] - 1, -1, -1
57 )
58 return torch.cat([zeros, x], dim=2)
59 else:
60 return x.expand(-1, -1, self.kernel_size[0], -1, -1)
61
62 def _apply_padding(self, x):
63 if sum(self.padding) == 0:
64 return x
65
66 x = F.pad(
67 x,
68 (self.padding[2], self.padding[2], self.padding[1], self.padding[1], 0, 0),
69 mode=self.pad_mode,
70 )
71
72 x = self._apply_temporal_padding(x)
73 return x
74
75 def forward(self, x):
76 if sum(self.padding) == 0:
77 return super().forward(x)
78
79 x = self._apply_padding(x)
80 return F.conv3d(
81 x,
82 self.weight,
83 self.bias,
84 stride=self.stride,
85 padding=0,
86 dilation=self.dilation,
87 )
88
89
90 class SpatialParallelConv3d(BaseConv3d):
91 def __init__(
92 self,
93 in_channels,
94 out_channels,
95 kernel_size,
96 stride=1,
97 padding=0,
98 bias=True,
99 padding_mode="zeros",
100 padding_mode_t=None,
101 causal=True,
102 ):
103 super().__init__(
104 in_channels,
105 out_channels,
106 kernel_size=kernel_size,
107 stride=stride,
108 padding=padding,
109 bias=bias,
110 padding_mode=padding_mode,
111 padding_mode_t=padding_mode_t,
112 causal=causal,
113 )
114 self.spatial_parallel = False
115 self.chunk_dim = -1
116
117 def _exchange_borders(self, x, sp_rank, sp_size):
118 if self.chunk_dim == -1:
119 pad = self.padding[2]
120 elif self.chunk_dim == -2:
121 pad = self.padding[1]
122 else:
123 raise ValueError(f"Invalid chunk dimension: {self.chunk_dim}")
124
125 if pad == 0:
126 return x
127
128 local_process_group = get_parallel_state()["sp_process_group"]
129 return exchange_borders(
130 x,
131 pad,
132 self.pad_mode,
133 sp_rank,
134 sp_size,
135 local_process_group,
136 dim=self.chunk_dim,
137 )
138
139 def _apply_padding(self, x):
140 if not self.spatial_parallel:
141 return super()._apply_padding(x)
142
143 state = get_parallel_state()
144
145 x = self._exchange_borders(x, state["sp_rank"], state["sp_size"])
146
147 if self.chunk_dim == -1:
148 x = F.pad(
149 x, (0, 0, self.padding[1], self.padding[1], 0, 0), mode=self.pad_mode
150 )
151 elif self.chunk_dim == -2:
152 x = F.pad(
153 x, (self.padding[2], self.padding[2], 0, 0, 0, 0), mode=self.pad_mode
154 )
155 else:
156 raise ValueError(f"Invalid chunk dimension: {self.chunk_dim}")
157
158 x = self._apply_temporal_padding(x)
159 return x
160