| 103 | """ |
| 104 | |
| 105 | def __init__( |
| 106 | self, |
| 107 | channels |
| 108 | ): |
| 109 | super().__init__() |
| 110 | self.channels = channels |
| 111 | |
| 112 | self.in_ln = nn.LayerNorm(channels, eps=1e-6) |
| 113 | self.mlp = nn.Sequential( |
| 114 | nn.Linear(channels, channels, bias=True), |
| 115 | nn.SiLU(), |
| 116 | nn.Linear(channels, channels, bias=True), |
| 117 | ) |
| 118 | |
| 119 | self.adaLN_modulation = nn.Sequential( |
| 120 | nn.SiLU(), |
| 121 | nn.Linear(channels, 3 * channels, bias=True) |
| 122 | ) |
| 123 | |
| 124 | def forward(self, x, y): |
| 125 | shift_mlp, scale_mlp, gate_mlp = self.adaLN_modulation(y).chunk(3, dim=-1) |