(self, x, y)
| 122 | ) |
| 123 | |
| 124 | def forward(self, x, y): |
| 125 | shift_mlp, scale_mlp, gate_mlp = self.adaLN_modulation(y).chunk(3, dim=-1) |
| 126 | h = modulate(self.in_ln(x), shift_mlp, scale_mlp) |
| 127 | h = self.mlp(h) |
| 128 | return x + gate_mlp * h |
| 129 | |
| 130 | |
| 131 | class FinalLayer(nn.Module): |
nothing calls this directly
no test coverage detected