| """ |
| Reference code |
| [FLUX] https://github.com/black-forest-labs/flux/blob/main/src/flux/modules/autoencoder.py |
| [WANX] https://github.com/Wan-Video/Wan2.1/blob/main/wan/modules/vae.py |
| """ |
|
|
| from dataclasses import dataclass |
|
|
| from einops import rearrange |
|
|
| import torch |
|
|
| from torch import Tensor, nn |
|
|
| import torch.nn.functional as F |
|
|
| from diffusers.utils.torch_utils import randn_tensor |
|
|
| class DiagonalGaussianDistribution: |
| def __init__(self, parameters: torch.Tensor): |
| self.mean, logvar = parameters.chunk(2, dim=1) |
| self.std = (0.5 * logvar.clamp(-30.0, 20.0)).exp() |
|
|
| def sample(self, generator: torch.Generator | None = None) -> torch.Tensor: |
| sample = randn_tensor( |
| self.mean.shape, |
| generator=generator, |
| device=self.mean.device, |
| dtype=self.mean.dtype, |
| ) |
| return self.mean + self.std * sample |
|
|
| 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): |
| """Frame-level self-attention""" |
|
|
| def __init__(self, in_channels: int): |
| super().__init__() |
|
|
| 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, _, 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) |
|
|
