| 188 | |
| 189 | |
| 190 | class AdaLayerNormSingle(torch.nn.Module): |
| 191 | def __init__(self, dim): |
| 192 | super().__init__() |
| 193 | self.silu = torch.nn.SiLU() |
| 194 | self.linear = torch.nn.Linear(dim, 3 * dim, bias=True) |
| 195 | self.norm = torch.nn.LayerNorm(dim, elementwise_affine=False, eps=1e-6) |
| 196 | |
| 197 | |
| 198 | def forward(self, x, emb): |
| 199 | emb = self.linear(self.silu(emb)) |
| 200 | shift_msa, scale_msa, gate_msa = emb.chunk(3, dim=1) |
| 201 | x = self.norm(x) * (1 + scale_msa[:, None]) + shift_msa[:, None] |
| 202 | return x, gate_msa |
| 203 | |
| 204 | |
| 205 |