| """
|
| Modified: https://github.com/datnguyen-tien204
|
| Flux Project: https://github.com/black-forest-labs/flux/blob/main/src/flux/modules/layers.py
|
| """
|
| import math
|
| import torch
|
| from torch import nn, Tensor
|
| from einops import rearrange
|
| from flux.utils.math import attention, rope
|
|
|
|
|
| class EmbedND(nn.Module):
|
| def __init__(self, dim: int, theta: int, axes_dim: list[int]):
|
| super().__init__()
|
| self.dim = dim
|
| self.theta = theta
|
| self.axes_dim = axes_dim
|
|
|
| def forward(self, ids: Tensor) -> Tensor:
|
| n_axes = ids.shape[-1]
|
| emb = torch.cat(
|
| [rope(ids[..., i], self.axes_dim[i], self.theta) for i in range(n_axes)],
|
| dim=-3,
|
| )
|
| return emb.unsqueeze(1)
|
|
|
|
|
| def timestep_embedding(t: Tensor, dim, max_period=10000, time_factor: float = 1000.0):
|
| t = time_factor * t
|
| half = dim // 2
|
| freqs = torch.exp(-math.log(max_period) * torch.arange(start=0, end=half, dtype=torch.float32) / half).to(t.device)
|
| args = t[:, None].float() * freqs[None]
|
| embedding = torch.cat([torch.cos(args), torch.sin(args)], dim=-1)
|
| if dim % 2:
|
| embedding = torch.cat([embedding, torch.zeros_like(embedding[:, :1])], dim=-1)
|
| if torch.is_floating_point(t):
|
| embedding = embedding.to(t)
|
| return embedding
|
|
|
|
|
| class MLPEmbedder(nn.Module):
|
| def __init__(self, in_dim: int, hidden_dim: int):
|
| super().__init__()
|
| self.in_layer = nn.Linear(in_dim, hidden_dim, bias=True)
|
| self.silu = nn.SiLU()
|
| self.out_layer = nn.Linear(hidden_dim, hidden_dim, bias=True)
|
|
|
| def forward(self, x: Tensor) -> Tensor:
|
| return self.out_layer(self.silu(self.in_layer(x)))
|
|
|
|
|
| class RMSNorm(torch.nn.Module):
|
| def __init__(self, dim: int):
|
| super().__init__()
|
| self.scale = nn.Parameter(torch.ones(dim))
|
|
|
| def forward(self, x: Tensor):
|
| x_dtype = x.dtype
|
| x = x.float()
|
| rrms = torch.rsqrt(torch.mean(x ** 2, dim=-1, keepdim=True) + 1e-6)
|
|
|
|
|
| if x.shape[-1] != self.scale.shape[0]:
|
| print(f"!!! SHAPE MISMATCH ERROR !!!")
|
| print(f"Input x shape: {x.shape}")
|
| print(f"Scale shape: {self.scale.shape}")
|
|
|
|
|
| return (x * rrms).to(dtype=x_dtype) * self.scale
|
|
|
|
|
| class QKNorm(torch.nn.Module):
|
| def __init__(self, dim: int):
|
| super().__init__()
|
| self.query_norm = RMSNorm(dim)
|
| self.key_norm = RMSNorm(dim)
|
|
|
| def forward(self, q: Tensor, k: Tensor, v: Tensor) -> tuple[Tensor, Tensor]:
|
| q = self.query_norm(q)
|
| k = self.key_norm(k)
|
| return q.to(v), k.to(v)
|
|
|
|
|
| class SelfAttention(nn.Module):
|
| def __init__(self, dim: int, num_heads: int = 8, qkv_bias: bool = False):
|
| super().__init__()
|
| self.num_heads = num_heads
|
| head_dim = dim // num_heads
|
| self.qkv = nn.Linear(dim, dim * 3, bias=qkv_bias)
|
| self.norm = QKNorm(head_dim)
|
| self.proj = nn.Linear(dim, dim)
|
|
|
| def forward(self, x: Tensor, pe: Tensor) -> Tensor:
|
| qkv = self.qkv(x)
|
| q, k, v = rearrange(qkv, "B L (K H D) -> K B H L D", K=3, H=self.num_heads)
|
| q, k = self.norm(q, k, v)
|
| x = attention(q, k, v, pe=pe)
|
| x = self.proj(x)
|
| return x |
