import math from dataclasses import dataclass from einops import rearrange import numpy as np import torch from torch import nn, Tensor import torch.nn.functional as F from diffusers.utils.torch_utils import randn_tensor from diffusers.configuration_utils import ConfigMixin, register_to_config from diffusers.models.modeling_utils import ModelMixin CACHE_T = 1 class DiagonalGaussianDistribution: def __init__(self, parameters: torch.Tensor, deterministic: bool = False): if parameters.ndim == 3: dim = 2 elif parameters.ndim == 5 or parameters.ndim == 4: dim = 1 else: raise NotImplementedError self.parameters = parameters self.mean, self.logvar = torch.chunk(parameters, 2, dim=dim) self.logvar = torch.clamp(self.logvar, -30.0, 20.0) self.deterministic = deterministic self.std = torch.exp(0.5 * self.logvar) self.var = torch.exp(self.logvar) if self.deterministic: self.var = self.std = torch.zeros_like( self.mean, device=self.parameters.device, dtype=self.parameters.dtype ) def sample(self, generator: torch.Generator | None = None) -> torch.Tensor: sample = randn_tensor( self.mean.shape, generator=generator, device=self.parameters.device, dtype=self.parameters.dtype, ) x = self.mean + self.std * sample return x def kl(self, other: "DiagonalGaussianDistribution" = None) -> torch.Tensor: if self.deterministic: return torch.Tensor([0.0]) else: reduce_dim = list(range(1, self.mean.ndim)) if other is None: return 0.5 * torch.sum( torch.pow(self.mean, 2) + self.var - 1.0 - self.logvar, dim=reduce_dim, ) else: return 0.5 * torch.sum( torch.pow(self.mean - other.mean, 2) / other.var + self.var / other.var - 1.0 - self.logvar + other.logvar, dim=reduce_dim, ) def nll(self, sample: torch.Tensor, dims: tuple[int, ...] = [1, 2, 3]) -> torch.Tensor: if self.deterministic: return torch.Tensor([0.0]) logtwopi = np.log(2.0 * np.pi) return 0.5 * torch.sum( logtwopi + self.logvar + torch.pow(sample - self.mean, 2) / self.var, dim=dims, ) def mode(self) -> torch.Tensor: return self.mean @dataclass class EncoderOutput: latent_dist: DiagonalGaussianDistribution = None @dataclass class DecoderOutput: sample: torch.FloatTensor = None posterior: DiagonalGaussianDistribution | None = None def swish(x: Tensor) -> Tensor: return x * torch.sigmoid(x) class RMSNorm(nn.Module): def __init__(self, dim: int, channel_first: bool = True, images: bool = False, bias: bool = False): super().__init__() broadcastable_dims = (1, 1, 1) if not images else (1, 1) shape = (dim, *broadcastable_dims) if channel_first else (dim, ) self.channel_first = channel_first self.scale = dim ** 0.5 self.gamma = nn.Parameter(torch.ones(shape)) self.bias = nn.Parameter(torch.zeros(shape)) if bias else 0.0 def forward(self, x: Tensor) -> Tensor: return F.normalize(x, dim=1 if self.channel_first else -1) * self.scale * self.gamma + self.bias class AttnBlock(nn.Module): def __init__(self, in_channels: int): super().__init__() self.in_channels = in_channels self.norm = RMSNorm(in_channels, channel_first=True, images=False) self.q = nn.Conv3d(in_channels, in_channels, kernel_size=1) self.k = nn.Conv3d(in_channels, in_channels, kernel_size=1) self.v = nn.Conv3d(in_channels, in_channels, kernel_size=1) self.proj_out = nn.Conv3d(in_channels, in_channels, kernel_size=1) def attention(self, x: Tensor) -> Tensor: b, c, t, h, w = x.shape x = self.norm(x) q = self.q(x) k = self.k(x) v = self.v(x) q = rearrange(q, "b c t h w -> (b t) 1 (h w) c") k = rearrange(k, "b c t h w -> (b t) 1 (h w) c") v = rearrange(v, "b c t h w -> (b t) 1 (h w) c") x = F.scaled_dot_product_attention(q, k, v) x = rearrange(x, "(b t) 1 (h w) c -> b c t h w", b=b, t=t, h=h, w=w) x = self.proj_out(x) return x def forward(self, x: Tensor) -> Tensor: return x + self.attention(x) class ChunkCausalConv3d(nn.Conv3d): def __init__( self, chunk_size: int, 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, ) self.chunk_size = chunk_size 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: padding = list(self._padding) 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, padding) return super().forward(x) class ResidualBlock(nn.Module): def __init__(self, chunk_size: int, in_channels: int, out_channels: int): super().__init__() self.in_channels = in_channels self.out_channels = out_channels self.norm1 = RMSNorm(in_channels, channel_first=True, images=False) self.conv1 = ChunkCausalConv3d( chunk_size, in_channels, out_channels, kernel_size=3, stride=1, padding=1) self.norm2 = RMSNorm(out_channels, channel_first=True, images=False) self.conv2 = ChunkCausalConv3d( chunk_size, out_channels, out_channels, kernel_size=3, stride=1, padding=1) if self.in_channels != self.out_channels: self.nin_shortcut = ChunkCausalConv3d( chunk_size, in_channels, out_channels, kernel_size=1, stride=1, padding=0) def forward(self, x, feat_cache: Tensor | None = None, feat_idx: int | None = None): shortcut = x x = self.norm1(x) x = swish(x) if feat_cache is not None: idx = feat_idx[0] cache_x = x[:, :, -CACHE_T:, :, :].clone() x = self.conv1(x, cache_x=feat_cache[idx]) feat_cache[idx] = cache_x feat_idx[0] += 1 else: x = self.conv1(x) x = self.norm2(x) x = swish(x) if feat_cache is not None: idx = feat_idx[0] cache_x = x[:, :, -CACHE_T:, :, :].clone() x = self.conv2(x, cache_x=feat_cache[idx]) feat_cache[idx] = cache_x feat_idx[0] += 1 else: x = self.conv2(x) if self.in_channels != self.out_channels: shortcut = self.nin_shortcut(shortcut) return x + shortcut class DownsampleBlock(nn.Module): def __init__(self, chunk_size: int, 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( chunk_size, 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) if feat_cache is not None: idx = feat_idx[0] cache_x = x[:, :, -CACHE_T:, :, :].clone() x = self.conv(x, cache_x=feat_cache[idx]) feat_cache[idx] = cache_x feat_idx[0] += 1 else: x = self.conv(x) if self.temporal_downsample and first_chunk: x = torch.cat([x[:, :, :1, :, :], x], dim=2) else: x = x 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, C, 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, chunk_size: int, 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( chunk_size, 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) if feat_cache is not None: idx = feat_idx[0] cache_x = x[:, :, -CACHE_T:, :, :].clone() x = self.conv(x, cache_x=feat_cache[idx]) feat_cache[idx] = cache_x feat_idx[0] += 1 else: x = self.conv(x) 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, ...], chunk_size: int, ): super().__init__() self.z_channels = z_channels self.block_in_channels = block_in_channels self.num_res_blocks = num_res_blocks cur_chunk_size = chunk_size self.conv_in = ChunkCausalConv3d( cur_chunk_size, 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(self.num_res_blocks): self.down_blocks.append(ResidualBlock( cur_chunk_size, 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( cur_chunk_size, block_in, block_out, temporal_downsample[i_level])) if temporal_downsample[i_level]: cur_chunk_size //= 2 self.mid_blocks = nn.ModuleList([ ResidualBlock(cur_chunk_size, in_channels=block_in, out_channels=block_in), AttnBlock(block_in), ResidualBlock(cur_chunk_size, in_channels=block_in, out_channels=block_in), ]) self.norm_out = RMSNorm(block_in, channel_first=True, images=False) self.conv_out = ChunkCausalConv3d( cur_chunk_size, 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: if feat_cache is not None: idx = feat_idx[0] cache_x = x[:, :, -CACHE_T:, :, :].clone() x = self.conv_in(x, cache_x=feat_cache[idx]) feat_cache[idx] = cache_x feat_idx[0] += 1 else: x = self.conv_in(x) 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) if feat_cache is not None: idx = feat_idx[0] cache_x = x[:, :, -CACHE_T:, :, :].clone() x = self.conv_out(x, cache_x=feat_cache[idx]) feat_cache[idx] = cache_x feat_idx[0] += 1 else: x = self.conv_out(x) 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, ...], chunk_size: int, ): super().__init__() self.z_channels = z_channels self.block_in_channels = block_in_channels self.num_res_blocks = num_res_blocks cur_chunk_size = chunk_size // (2 ** sum(temporal_upsample)) block_in = block_in_channels[0] self.conv_in = ChunkCausalConv3d( cur_chunk_size, z_channels, block_in, kernel_size=3, stride=1, padding=1) self.mid_blocks = nn.ModuleList([ ResidualBlock(cur_chunk_size, in_channels=block_in, out_channels=block_in), AttnBlock(block_in), ResidualBlock(cur_chunk_size, 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(self.num_res_blocks + 1): self.up_blocks.append(ResidualBlock( cur_chunk_size, 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( cur_chunk_size, block_in, block_out, temporal_upsample[i_level])) if temporal_upsample[i_level]: cur_chunk_size *= 2 self.norm_out = RMSNorm(block_in, channel_first=True, images=False) self.conv_out = ChunkCausalConv3d( cur_chunk_size, 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: if feat_cache is not None: idx = feat_idx[0] cache_x = x[:, :, -CACHE_T:, :, :].clone() x = self.conv_in(x, cache_x=feat_cache[idx]) feat_cache[idx] = cache_x feat_idx[0] += 1 else: x = self.conv_in(x) 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) if feat_cache is not None: idx = feat_idx[0] cache_x = x[:, :, -CACHE_T:, :, :].clone() x = self.conv_out(x, cache_x=feat_cache[idx]) feat_cache[idx] = cache_x feat_idx[0] += 1 else: x = self.conv_out(x) 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, x: Tensor) -> Tensor: b, c, t, h, w = x.shape if h % self.stride != 0 or w % self.stride != 0: return x oh, ow = h * self.group // self.stride, w * self.group // self.stride z = x.permute(0, 2, 1, 3, 4).reshape(b * t, c, h, w) z = F.pixel_unshuffle(z, self.stride) z = self.proj(z) z = F.pixel_shuffle(z, self.group) return z.reshape(b, t, c, oh, ow).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, x: Tensor) -> Tensor: return x + self.pw(self.act(self.dw(x))) 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, x: Tensor) -> Tensor: b, c, t, h, w = x.shape oh, ow = round(h * self.scale), round(w * self.scale) z = x.permute(0, 2, 1, 3, 4).reshape(b * t, c, h, w) f = self.act(self.conv_in(z)) f = self.blocks(f) f = self.reduce(f) f = F.interpolate(f, size=(oh, ow), mode="bilinear", align_corners=False) residual = self.conv_out(f) base = F.interpolate(z, size=(oh, ow), mode="bilinear", align_corners=False) out = (base + residual).clamp(-1.0, 1.0) return out.reshape(b, t, c, oh, ow).permute(0, 2, 1, 3, 4) class XVAEChunkCausal(ModelMixin, ConfigMixin): """For more technical details on high-resolution causal VAE decoding, see: https://github.com/xin1u/UltraFlash """ @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.in_channels = in_channels self.out_channels = out_channels self.patch_size = patch_size self.latent_channels = latent_channels self.ffactor_spatial = patch_size * 2 ** (len(temporal_downsample) - 1) self.ffactor_temporal = 2 ** sum(temporal_downsample) self.chunk_size = chunk_size self.latents_mean = latents_mean self.latents_std = latents_std self.stem = Stem(in_channels) 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, chunk_size=chunk_size, ) 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, chunk_size=chunk_size, ) self.head = Head(out_channels) self.use_slicing = enable_slicing 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): if not hasattr(self, "_enc_conv_num") or not hasattr(self, "_dec_conv_num"): 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()) self._enc_conv_idx = [0] self._dec_conv_idx = [0] self._enc_feat_map = [None] * self._enc_conv_num self._dec_feat_map = [None] * self._dec_conv_num def _encode(self, x: Tensor): x = self.stem(x) b, c, t, h, w = x.shape assert t % self.ffactor_temporal == 1, f"Temporal dimension must be {self.ffactor_temporal}n + 1, got {x.shape}." assert h % self.ffactor_spatial == 0 and w % self.ffactor_spatial == 0, f"Spatial dimensions must be divisible by {self.ffactor_spatial}, got {x.shape}." 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_): self._enc_conv_idx = [0] if i == 0: h = self.encoder( x[:, :, :1, :, :], feat_cache=self._enc_feat_map, feat_idx=self._enc_conv_idx, 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=self._enc_conv_idx, 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): assert len( x.shape) == 5, f"Input tensor must be 5D (b, c, t, h, w), got {x.shape}." 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_): self._dec_conv_idx = [0] if i == 0: h = self.decoder( z[:, :, :1, :, :], feat_cache=self._dec_feat_map, feat_idx=self._dec_conv_idx, 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=self._dec_conv_idx, 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)