| import math |
|
|
| import torch |
| import torch.nn as nn |
| import torch.nn.functional as F |
|
|
|
|
| class SinusoidalEmbedding(nn.Module): |
| def __init__(self, embedding_dim, device="cpu"): |
| super().__init__() |
| half_dim = embedding_dim // 2 |
| freqs = math.log(10_000) / (half_dim - 1) |
| freqs = torch.exp(torch.arange(half_dim, device=device) * -freqs) |
| self.register_buffer("freqs", freqs) |
|
|
| def forward(self, x): |
| embeddings = self.freqs.reshape((1,) * len(x.shape) + self.freqs.shape) |
| embeddings = x.unsqueeze(-1) * embeddings |
| return torch.cat((embeddings.cos(), embeddings.sin()), dim=-1) |
|
|
|
|
| class MzArrayEncoder(nn.Module): |
| def __init__(self, lambda_min=1e-3, lambda_max=1e4, emb_dim=128, device="cpu"): |
| super().__init__() |
| half_dim = emb_dim // 2 |
| freqs = math.log(lambda_max / lambda_min) / (half_dim - 1) |
| freqs = torch.exp(torch.arange(half_dim, device=device) * -freqs) |
| freqs = freqs * 2 * torch.pi / lambda_min |
| self.register_buffer("freqs", freqs) |
|
|
| def forward(self, m): |
| embeddings = self.freqs.reshape((1,) * len(m.shape) + self.freqs.shape) |
| embeddings = m.unsqueeze(-1) * embeddings |
| return torch.cat((embeddings.cos(), -embeddings.sin()), dim=-1) |
|
|
|
|
| class AdaNLayers(nn.Module): |
| def __init__(self, emb_dim=128, n=1): |
| super().__init__() |
| self.norm = nn.LayerNorm(emb_dim, elementwise_affine=False) |
| self.mlp = nn.Sequential(nn.SiLU(), nn.Linear(emb_dim, emb_dim * n * 2)) |
| self.n = n |
|
|
| def forward(self, conditioner): |
| x = self.mlp(conditioner) |
| return x.chunk(2 * self.n, dim=-1) |
|
|
|
|
| def modulate(x: torch.Tensor, shift: torch.Tensor, scale: torch.Tensor) -> torch.Tensor: |
| """Modulate normalized features with scale and shift: x * (1 + scale) + shift.""" |
| return x * (1.0 + scale.unsqueeze(1)) + shift.unsqueeze(1) |
|
|
|
|
| class AdaLNZero(nn.Module): |
| """ |
| Adaptive Layer Normalization with Zero-Initialization for Diffusion Transformers (DiT). |
| Predicts (scale, shift, gate) modulation parameters for transformer sublayers. |
| Initializes the projection to zeros so blocks start as identity functions. |
| """ |
| def __init__(self, emb_dim: int = 512, num_params: int = 9): |
| super().__init__() |
| self.mlp = nn.Sequential( |
| nn.SiLU(), |
| nn.Linear(emb_dim, num_params * emb_dim), |
| ) |
| |
| nn.init.zeros_(self.mlp[1].weight) |
| nn.init.zeros_(self.mlp[1].bias) |
| self.num_params = num_params |
|
|
| def forward(self, conditioner: torch.Tensor) -> tuple[torch.Tensor, ...]: |
| return self.mlp(conditioner).chunk(self.num_params, dim=-1) |
|
|
|
|
| class SwiGLU(nn.Module): |
| def __init__(self, d_model: int, d_ff: int, dropout: float = 0.0): |
| super().__init__() |
| self.w_gate_up = nn.Linear(d_model, 2 * d_ff, bias=False) |
| self.w_down = nn.Linear(d_ff, d_model, bias=False) |
| self.dropout = nn.Dropout(dropout) if dropout > 0.0 else nn.Identity() |
|
|
| def forward(self, x: torch.Tensor) -> torch.Tensor: |
| gate, up = self.w_gate_up(x).chunk(2, dim=-1) |
| return self.dropout(self.w_down(F.silu(gate) * up)) |
|
|
|
|
| SwiGLUFFN = SwiGLU |
|
|
