| import torch |
| import torch.nn as nn |
| import torch.nn.functional as F |
|
|
| class PatchEmbed(nn.Module): |
| def __init__(self, patch_size, in_channels, embed_dim): |
| super().__init__() |
| self.proj = nn.Conv2d(in_channels, embed_dim, kernel_size=patch_size, stride=patch_size) |
|
|
| def forward(self, x): |
| B, C, latent_h, latent_w = x.shape |
| x = self.proj(x) |
| x = x.flatten(2).transpose(1, 2) |
| return x, latent_h, latent_w |
|
|
| class SinusoidalPositionEmbeddings(nn.Module): |
| def __init__(self, dim): |
| super().__init__() |
| self.dim = dim |
|
|
| def forward(self, time): |
| device = time.device |
| half_dim = self.dim // 2 |
|
|
| embeddings = torch.log(torch.tensor(10000.0)) / (half_dim - 1) |
| embeddings = torch.exp(torch.arange(half_dim, device=device) * -embeddings) |
|
|
| embeddings = time[:, None] * embeddings[None, :] |
|
|
| embeddings = torch.cat( |
| (embeddings.sin(), embeddings.cos()), dim=-1 |
| ) |
| return embeddings |
|
|
| class LabelEmbedding(nn.Module): |
| def __init__(self, num_classes, dim): |
| super().__init__() |
|
|
| self.embedding = nn.Embedding(num_classes + 1, dim) |
| self.null_classes = num_classes |
|
|
| def forward(self, labels): |
| return self.embedding(labels) |
|
|
| class DiTBlock(nn.Module): |
| def __init__(self, dim, heads, mlp_dim): |
| super().__init__() |
|
|
| self.norm1 = nn.LayerNorm(dim) |
| self.attn = nn.MultiheadAttention(dim, heads, batch_first=True) |
| self.norm2 = nn.LayerNorm(dim) |
| self.mlp = nn.Sequential( |
| nn.Linear(dim, mlp_dim), |
| nn.GELU(), |
| nn.Linear(mlp_dim, dim) |
| ) |
|
|
| def forward(self, x): |
| h = self.norm1(x) |
| attn_out, _ = self.attn(h, h, h) |
| x = x + attn_out |
| x = x + self.mlp(self.norm2(x)) |
| return x |
|
|
| class Unpatchify(nn.Module): |
| def __init__(self, patch_size, dim, out_channels): |
| super().__init__() |
| self.patch_size = patch_size |
| self.out_channels = out_channels |
| self.proj = nn.Linear(dim, out_channels * patch_size * patch_size) |
|
|
| def forward(self, x, latent_h, latent_w): |
| B, N, D = x.shape |
| x = self.proj(x) |
| H = latent_h // self.patch_size |
| W = latent_w // self.patch_size |
| x = x.view(B, H, W, self.out_channels, self.patch_size, self.patch_size) |
| x = x.permute(0, 3, 1, 4, 2, 5) |
| x = x.reshape(B, self.out_channels, H*self.patch_size, W*self.patch_size) |
| return x |
|
|
| class TinyDit(nn.Module): |
| def __init__(self, latent_channels, dim, depth, patch_size, num_classes): |
| super().__init__() |
| self.patch_embed = PatchEmbed(patch_size=patch_size, in_channels=latent_channels, embed_dim=dim) |
| self.label_embed = LabelEmbedding(num_classes=num_classes, dim=dim) |
| self.time_embed = nn.Sequential( |
| SinusoidalPositionEmbeddings(dim), |
| nn.Linear(dim, dim), |
| nn.GELU(), |
| nn.Linear(dim, dim) |
| ) |
| self.pos_embed = nn.Parameter(torch.randn(1, latent_channels, dim)) |
| self.blocks = nn.ModuleList([DiTBlock(dim=dim, heads=8, mlp_dim=dim*4) for _ in range(depth)]) |
| self.unpatchify = Unpatchify(patch_size=patch_size, dim=dim, out_channels=latent_channels) |
|
|
| def forward(self, x, t, labels=None): |
| x, H, W = self.patch_embed(x) |
| x = x + self.pos_embed |
| t = self.time_embed(t) |
| if labels is not None: |
| t = t + self.label_embed(labels) |
| x = x + t[:, None, :] |
| for block in self.blocks: |
| x = block(x) |
| x = self.unpatchify(x, H, W) |
| return x |
