| import math |
|
|
| from einops import rearrange |
| import torch |
| from torch import nn, Tensor |
| import torch.nn.functional as F |
|
|
| from diffusers.configuration_utils import ConfigMixin, register_to_config |
| from diffusers.models.modeling_utils import ModelMixin |
| from src.models.mmdit.vae.layers import RMSNorm, swish, AttnBlock, DecoderOutput, EncoderOutput, DiagonalGaussianDistribution |
|
|
| CACHE_T = 1 |
|
|
|
|
| def _cached_conv(conv, x, feat_cache, feat_idx): |
| if feat_cache is None: |
| return conv(x) |
| idx = feat_idx[0] |
| tail = x[:, :, -CACHE_T:].clone() |
| x = conv(x, cache_x=feat_cache[idx]) |
| feat_cache[idx] = tail |
| feat_idx[0] += 1 |
| return x |
|
|
|
|
| class ChunkCausalConv3d(nn.Conv3d): |
| def __init__( |
| self, |
| in_channels: int, |
| out_channels: int, |
| kernel_size: int | tuple[int, int, int], |
| stride: int | tuple[int, int, int] = 1, |
| padding: int | tuple[int, int, int] = 0 |
| ): |
| super().__init__( |
| in_channels, |
| out_channels, |
| kernel_size, |
| stride=stride, |
| padding=padding, |
| ) |
| |
| assert self.padding[0] == 1, "Causal padding only supports padding of 1 in temporal dimension." |
| self._padding = (self.padding[2], self.padding[2], |
| self.padding[1], self.padding[1], 0, 0) |
| self.padding = (0, 0, 0) |
|
|
| def forward(self, x: Tensor, cache_x: Tensor | None = None) -> Tensor: |
| if cache_x is not None: |
| padding_front = cache_x.to(x.device) |
| else: |
| assert x.shape[ |
| 2] == 1, f"Input temporal dimension is expected to be 1 when cache_x is None, got {x.shape[2]}." |
| padding_front = x |
| x = torch.cat([padding_front, x, x[:, :, -1:, :, :]], dim=2) |
| x = F.pad(x, self._padding) |
| return super().forward(x) |
|
|
|
|
| class ResidualBlock(nn.Module): |
| def __init__(self, in_channels: int, out_channels: int): |
| super().__init__() |
|
|
| self.norm1 = RMSNorm(in_channels, channel_first=True, images=False) |
| self.conv1 = ChunkCausalConv3d( |
| in_channels, out_channels, kernel_size=3, stride=1, padding=1) |
| self.norm2 = RMSNorm(out_channels, channel_first=True, images=False) |
| self.conv2 = ChunkCausalConv3d( |
| out_channels, out_channels, kernel_size=3, stride=1, padding=1) |
| self.nin_shortcut = nn.Identity() if in_channels == out_channels else nn.Conv3d(in_channels, out_channels, 1) |
|
|
| def forward(self, x, feat_cache: Tensor | None = None, feat_idx: int | None = None): |
| shortcut = self.nin_shortcut(x) |
|
|
| x = self.norm1(x) |
| x = swish(x) |
| x = _cached_conv(self.conv1, x, feat_cache, feat_idx) |
|
|
| x = self.norm2(x) |
| x = swish(x) |
| x = _cached_conv(self.conv2, x, feat_cache, feat_idx) |
|
|
| return x + shortcut |
|
|
|
|
| class DownsampleBlock(nn.Module): |
| def __init__(self, in_channels: int, out_channels: int, temporal_downsample: bool): |
| super().__init__() |
| factor = 2 * 2 * 2 if temporal_downsample else 1 * 2 * 2 |
| self.conv = ChunkCausalConv3d( |
| in_channels, out_channels // factor, kernel_size=3, stride=1, padding=1) |
|
|
| self.temporal_downsample = temporal_downsample |
| self.group_size = factor * in_channels // out_channels |
|
|
| def forward(self, x: Tensor, feat_cache: Tensor | None = None, feat_idx: int | None = None, first_chunk: bool = False) -> Tensor: |
| r1 = 2 if self.temporal_downsample else 1 |
|
|
| if self.temporal_downsample and first_chunk: |
| shortcut = torch.cat([x[:, :, :1, :, :], x], dim=2) |
| else: |
| shortcut = x |
| shortcut = rearrange( |
| shortcut, "b c (f r1) (h r2) (w r3) -> b (r1 r2 r3 c) f h w", r1=r1, r2=2, r3=2) |
|
|
| x = _cached_conv(self.conv, x, feat_cache, feat_idx) |
|
|
| if self.temporal_downsample and first_chunk: |
| x = torch.cat([x[:, :, :1, :, :], x], dim=2) |
| x = rearrange( |
| x, "b c (f r1) (h r2) (w r3) -> b (r1 r2 r3 c) f h w", r1=r1, r2=2, r3=2) |
|
|
| B, _, T, H, W = shortcut.shape |
| shortcut = shortcut.view( |
| B, x.shape[1], self.group_size, T, H, W).mean(dim=2) |
| return x + shortcut |
|
|
|
|
| class UpsampleBlock(nn.Module): |
| def __init__(self, in_channels: int, out_channels: int, temporal_upsample: bool): |
| super().__init__() |
| factor = 2 * 2 * 2 if temporal_upsample else 1 * 2 * 2 |
| self.conv = ChunkCausalConv3d( |
| in_channels, out_channels * factor, kernel_size=3, stride=1, padding=1) |
|
|
| self.temporal_upsample = temporal_upsample |
| self.repeats = factor * out_channels // in_channels |
|
|
| def forward(self, x: Tensor, feat_cache: Tensor | None = None, feat_idx: int | None = None, first_chunk: bool = False) -> Tensor: |
| r1 = 2 if self.temporal_upsample else 1 |
|
|
| shortcut = x.repeat_interleave(repeats=self.repeats, dim=1) |
| shortcut = rearrange( |
| shortcut, "b (r1 r2 r3 c) f h w -> b c (f r1) (h r2) (w r3)", r1=r1, r2=2, r3=2) |
|
|
| x = _cached_conv(self.conv, x, feat_cache, feat_idx) |
|
|
| x = rearrange( |
| x, "b (r1 r2 r3 c) f h w -> b c (f r1) (h r2) (w r3)", r1=r1, r2=2, r3=2) |
|
|
| x += shortcut |
| if self.temporal_upsample and first_chunk: |
| x = x[:, :, 1:, :, :] |
|
|
| return x |
|
|
|
|
| class Encoder(nn.Module): |
| def __init__( |
| self, |
| in_channels: int, |
| z_channels: int, |
| num_res_blocks: int, |
| block_in_channels: tuple[int, ...], |
| temporal_downsample: tuple[bool, ...], |
| ): |
| super().__init__() |
|
|
| self.conv_in = ChunkCausalConv3d( |
| in_channels, block_in_channels[0], kernel_size=3, stride=1, padding=1) |
|
|
| self.down_blocks = nn.ModuleList([]) |
| for i_level, block_in in enumerate(block_in_channels): |
| for _ in range(num_res_blocks): |
| self.down_blocks.append(ResidualBlock( |
| in_channels=block_in, out_channels=block_in)) |
|
|
| if i_level != len(block_in_channels) - 1: |
| block_out = block_in_channels[i_level + 1] |
| self.down_blocks.append(DownsampleBlock( |
| block_in, block_out, temporal_downsample[i_level])) |
|
|
| self.mid_blocks = nn.ModuleList([ |
| ResidualBlock(in_channels=block_in, |
| out_channels=block_in), |
| AttnBlock(block_in), |
| ResidualBlock(in_channels=block_in, |
| out_channels=block_in), |
| ]) |
|
|
| self.norm_out = RMSNorm(block_in, channel_first=True, images=False) |
| self.conv_out = ChunkCausalConv3d( |
| block_in, 2 * z_channels, kernel_size=3, stride=1, padding=1) |
|
|
| def forward(self, x: Tensor, feat_cache: Tensor | None = None, feat_idx: int | None = None, first_chunk: bool = False) -> Tensor: |
| x = _cached_conv(self.conv_in, x, feat_cache, feat_idx) |
|
|
| for block in self.down_blocks: |
| if isinstance(block, DownsampleBlock): |
| x = block(x, feat_cache=feat_cache, |
| feat_idx=feat_idx, first_chunk=first_chunk) |
| else: |
| x = block(x, feat_cache=feat_cache, feat_idx=feat_idx) |
| for block in self.mid_blocks: |
| if isinstance(block, ResidualBlock): |
| x = block(x, feat_cache=feat_cache, feat_idx=feat_idx) |
| else: |
| x = block(x) |
|
|
| x = self.norm_out(x) |
| x = swish(x) |
| x = _cached_conv(self.conv_out, x, feat_cache, feat_idx) |
| return x |
|
|
|
|
| class Decoder(nn.Module): |
| def __init__( |
| self, |
| z_channels: int, |
| out_channels: int, |
| num_res_blocks: int, |
| block_in_channels: tuple[int, ...], |
| temporal_upsample: tuple[bool, ...], |
| ): |
| super().__init__() |
|
|
| block_in = block_in_channels[0] |
| self.conv_in = ChunkCausalConv3d( |
| z_channels, block_in, kernel_size=3, stride=1, padding=1) |
|
|
| self.mid_blocks = nn.ModuleList([ |
| ResidualBlock(in_channels=block_in, |
| out_channels=block_in), |
| AttnBlock(block_in), |
| ResidualBlock(in_channels=block_in, |
| out_channels=block_in), |
| ]) |
|
|
| self.up_blocks = nn.ModuleList([]) |
| for i_level, block_in in enumerate(block_in_channels): |
| for _ in range(num_res_blocks + 1): |
| self.up_blocks.append(ResidualBlock( |
| in_channels=block_in, out_channels=block_in)) |
|
|
| if i_level != len(block_in_channels) - 1: |
| block_out = block_in_channels[i_level + 1] |
| self.up_blocks.append(UpsampleBlock( |
| block_in, block_out, temporal_upsample[i_level])) |
|
|
| self.norm_out = RMSNorm(block_in, channel_first=True, images=False) |
| self.conv_out = ChunkCausalConv3d( |
| block_in, out_channels, kernel_size=3, stride=1, padding=1) |
|
|
| def forward(self, x: Tensor, feat_cache: Tensor | None = None, feat_idx: int | None = None, first_chunk: bool = False) -> Tensor: |
| x = _cached_conv(self.conv_in, x, feat_cache, feat_idx) |
|
|
| for block in self.mid_blocks: |
| if isinstance(block, ResidualBlock): |
| x = block(x, feat_cache=feat_cache, feat_idx=feat_idx) |
| else: |
| x = block(x) |
|
|
| for block in self.up_blocks: |
| if isinstance(block, UpsampleBlock): |
| x = block(x, feat_cache=feat_cache, |
| feat_idx=feat_idx, first_chunk=first_chunk) |
| else: |
| x = block(x, feat_cache=feat_cache, feat_idx=feat_idx) |
|
|
| x = self.norm_out(x) |
| x = swish(x) |
| x = _cached_conv(self.conv_out, x, feat_cache, feat_idx) |
| return x |
|
|
|
|
| class Stem(nn.Module): |
| def __init__(self, channels: int, stride: int = 3, group: int = 2): |
| super().__init__() |
| self.stride = stride |
| self.group = group |
| self.proj = nn.Conv2d( |
| channels * stride * stride, |
| channels * group * group, |
| kernel_size=1, |
| bias=False, |
| ) |
|
|
| def forward(self, inputs: Tensor) -> Tensor: |
| batch_size, channels, num_frames, height, width = inputs.shape |
| if height % self.stride != 0 or width % self.stride != 0: |
| return inputs |
| output_height = height * self.group // self.stride |
| output_width = width * self.group // self.stride |
| features = inputs.permute(0, 2, 1, 3, 4).reshape( |
| batch_size * num_frames, channels, height, width |
| ) |
| features = F.pixel_unshuffle(features, self.stride) |
| features = self.proj(features) |
| features = F.pixel_shuffle(features, self.group) |
| return features.reshape( |
| batch_size, num_frames, channels, output_height, output_width |
| ).permute(0, 2, 1, 3, 4) |
|
|
|
|
| class HeadResBlock(nn.Module): |
| def __init__(self, channels: int): |
| super().__init__() |
| self.dw = nn.Conv2d( |
| channels, channels, kernel_size=3, padding=1, groups=channels |
| ) |
| self.pw = nn.Conv2d(channels, channels, kernel_size=1) |
| self.act = nn.ReLU(inplace=False) |
|
|
| def forward(self, inputs: Tensor) -> Tensor: |
| return inputs + self.pw(self.act(self.dw(inputs))) |
|
|
|
|
| class Head(nn.Module): |
| def __init__( |
| self, |
| channels: int, |
| scale: float = 1.5, |
| hidden: int = 32, |
| num_blocks: int = 4, |
| mid_channels: int = 12, |
| ): |
| super().__init__() |
| self.scale = float(scale) |
| self.conv_in = nn.Conv2d(channels, hidden, kernel_size=3, padding=1) |
| self.act = nn.ReLU(inplace=False) |
| self.blocks = nn.Sequential(*[HeadResBlock(hidden) for _ in range(num_blocks)]) |
| self.reduce = nn.Conv2d(hidden, mid_channels, kernel_size=3, padding=1) |
| self.conv_out = nn.Conv2d(mid_channels, channels, kernel_size=3, padding=1) |
|
|
| def forward(self, inputs: Tensor) -> Tensor: |
| batch_size, channels, num_frames, height, width = inputs.shape |
| output_height = round(height * self.scale) |
| output_width = round(width * self.scale) |
| frames = inputs.permute(0, 2, 1, 3, 4).reshape( |
| batch_size * num_frames, channels, height, width |
| ) |
| features = self.act(self.conv_in(frames)) |
| features = self.blocks(features) |
| features = self.reduce(features) |
| features = F.interpolate( |
| features, size=(output_height, output_width), mode="bilinear", align_corners=False |
| ) |
| residual = self.conv_out(features) |
| base = F.interpolate( |
| frames, size=(output_height, output_width), mode="bilinear", align_corners=False |
| ) |
| output = (base + residual).clamp(-1.0, 1.0) |
| return output.reshape( |
| batch_size, num_frames, channels, output_height, output_width |
| ).permute(0, 2, 1, 3, 4) |
|
|
|
|
| class XVAEChunkCausal(ModelMixin, ConfigMixin): |
|
|
| @register_to_config |
| def __init__( |
| self, |
| in_channels: int, |
| out_channels: int, |
| patch_size: int, |
| latent_channels: int, |
| layers_per_block: int, |
| block_in_channels: tuple[int, ...], |
| temporal_downsample: tuple[bool, ...], |
| chunk_size: int, |
| latents_mean: tuple[float, ...] = None, |
| latents_std: tuple[float, ...] = None, |
| enable_slicing: bool = False, |
| ): |
| super().__init__() |
| self.patch_size = patch_size |
| self.ffactor_temporal = 2 ** sum(temporal_downsample[:len(block_in_channels) - 1]) |
| self.chunk_size = chunk_size |
|
|
| self.stem = Stem(in_channels) |
| encoder_spatial_factor = patch_size * 2 ** (len(block_in_channels) - 1) |
| if encoder_spatial_factor % self.stem.group: |
| raise ValueError("Encoder spatial compression must be divisible by the stem output group.") |
| self.ffactor_spatial = encoder_spatial_factor * self.stem.stride // self.stem.group |
| if chunk_size <= 0 or chunk_size % self.ffactor_temporal: |
| raise ValueError(f"VAE chunk_size must be a positive multiple of {self.ffactor_temporal}.") |
|
|
| self.encoder = Encoder( |
| in_channels=in_channels * (patch_size**2), |
| z_channels=latent_channels, |
| num_res_blocks=layers_per_block, |
| block_in_channels=block_in_channels, |
| temporal_downsample=temporal_downsample, |
| ) |
| self.decoder = Decoder( |
| z_channels=latent_channels, |
| out_channels=out_channels * (patch_size**2), |
| num_res_blocks=layers_per_block, |
| block_in_channels=tuple(reversed(block_in_channels)), |
| temporal_upsample=temporal_downsample, |
| ) |
|
|
| self.head = Head(out_channels, scale=self.stem.stride / self.stem.group) |
|
|
| self.use_slicing = enable_slicing |
| self._enc_conv_num = sum(isinstance(m, ChunkCausalConv3d) for m in self.encoder.modules()) |
| self._dec_conv_num = sum(isinstance(m, ChunkCausalConv3d) for m in self.decoder.modules()) |
|
|
| def enable_slicing(self): |
| self.use_slicing = True |
|
|
| def disable_slicing(self): |
| self.use_slicing = False |
|
|
| @staticmethod |
| def patchify(x, patch_size: int) -> Tensor: |
| if patch_size == 1: |
| return x |
| x = rearrange(x, "b c t (h r1) (w r2) -> b (c r1 r2) t h w", |
| r1=patch_size, r2=patch_size) |
| return x |
|
|
| @staticmethod |
| def unpatchify(x, patch_size: int) -> Tensor: |
| if patch_size == 1: |
| return x |
| x = rearrange(x, "b (r1 r2 c) t h w -> b c t (h r1) (w r2)", |
| r1=patch_size, r2=patch_size) |
| return x |
|
|
| def clear_cache(self): |
| self._enc_feat_map = [None] * self._enc_conv_num |
| self._dec_feat_map = [None] * self._dec_conv_num |
|
|
| def _validate_input(self, inputs: Tensor): |
| if inputs.ndim != 5: |
| raise ValueError(f"VAE input must have shape (batch, channels, frames, height, width), got {inputs.shape}.") |
| num_frames, height, width = inputs.shape[2:] |
| if num_frames <= 0 or (num_frames - 1) % self.ffactor_temporal: |
| raise ValueError(f"Temporal dimension must be {self.ffactor_temporal}n+1, got {num_frames}.") |
| if height <= 0 or width <= 0 or height % self.ffactor_spatial or width % self.ffactor_spatial: |
| raise ValueError(f"Spatial dimensions must be positive multiples of {self.ffactor_spatial}, got {(height, width)}.") |
|
|
| def _encode(self, x: Tensor): |
| self._validate_input(x) |
| x = self.stem(x) |
| x = self.patchify(x, self.patch_size) |
|
|
| out = [] |
| self.clear_cache() |
| iter_ = 1 + math.ceil((x.shape[2] - 1) / self.chunk_size) |
| for i in range(iter_): |
| if i == 0: |
| h = self.encoder( |
| x[:, :, :1, :, :], |
| feat_cache=self._enc_feat_map, |
| feat_idx=[0], |
| first_chunk=True |
| ) |
| else: |
| h = self.encoder( |
| x[:, :, 1 + (i - 1) * self.chunk_size: 1 + |
| i * self.chunk_size, :, :], |
| feat_cache=self._enc_feat_map, |
| feat_idx=[0], |
| first_chunk=False |
| ) |
| out.append(h) |
| out = torch.cat(out, dim=2) |
| self.clear_cache() |
| return out |
|
|
| def encode(self, x: Tensor, return_dict: bool = True): |
| if self.use_slicing and x.shape[0] > 1: |
| encoded_slices = [self._encode(x_slice) for x_slice in x.split(1)] |
| h = torch.cat(encoded_slices) |
| else: |
| h = self._encode(x) |
|
|
| posterior = DiagonalGaussianDistribution(h) |
| if not return_dict: |
| return (posterior,) |
|
|
| return EncoderOutput(latent_dist=posterior) |
|
|
| def _decode(self, z: Tensor): |
| latent_chunk_size = self.chunk_size // self.ffactor_temporal |
|
|
| self.clear_cache() |
| decoded = [] |
| iter_ = 1 + math.ceil((z.shape[2] - 1) / latent_chunk_size) |
| for i in range(iter_): |
| if i == 0: |
| h = self.decoder( |
| z[:, :, :1, :, :], |
| feat_cache=self._dec_feat_map, |
| feat_idx=[0], |
| first_chunk=True |
| ) |
| else: |
| h = self.decoder( |
| z[:, :, 1 + (i - 1) * latent_chunk_size: 1 + |
| i * latent_chunk_size, :, :], |
| feat_cache=self._dec_feat_map, |
| feat_idx=[0], |
| first_chunk=False |
| ) |
| decoded.append(h) |
| decoded = torch.cat(decoded, dim=2) |
| self.clear_cache() |
|
|
| decoded = self.unpatchify(decoded, self.patch_size) |
| decoded = self.head(decoded) |
| return decoded |
|
|
| def decode(self, z: Tensor, return_dict: bool = True): |
| if self.use_slicing and z.shape[0] > 1: |
| decoded_slices = [self._decode(z_slice) for z_slice in z.split(1)] |
| decoded = torch.cat(decoded_slices) |
| else: |
| decoded = self._decode(z) |
|
|
| if not return_dict: |
| return (decoded,) |
|
|
| return DecoderOutput(sample=decoded) |
|
|
| def forward( |
| self, |
| sample: torch.Tensor, |
| sample_posterior: bool = False, |
| return_dict: bool = True, |
| ): |
| posterior = self.encode(sample).latent_dist |
| z = posterior.sample() if sample_posterior else posterior.mode() |
| dec = self.decode(z).sample |
| return DecoderOutput(sample=dec, posterior=posterior) if return_dict else (dec, posterior) |
|
|
