A residual block that can optionally change the number of channels. :param channels: the number of input channels.
| 97 | |
| 98 | |
| 99 | class ResBlock(nn.Module): |
| 100 | """ |
| 101 | A residual block that can optionally change the number of channels. |
| 102 | :param channels: the number of input channels. |
| 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) |
| 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): |