(
self,
channels,
emb_channels,
dropout,
out_channels=None,
use_conv=False,
use_scale_shift_norm=False,
dims=2,
use_checkpoint=False,
up=False,
down=False,
)
| 178 | """ |
| 179 | |
| 180 | def __init__( |
| 181 | self, |
| 182 | channels, |
| 183 | emb_channels, |
| 184 | dropout, |
| 185 | out_channels=None, |
| 186 | use_conv=False, |
| 187 | use_scale_shift_norm=False, |
| 188 | dims=2, |
| 189 | use_checkpoint=False, |
| 190 | up=False, |
| 191 | down=False, |
| 192 | ): |
| 193 | super().__init__() |
| 194 | self.channels = channels |
| 195 | self.emb_channels = emb_channels |
| 196 | self.dropout = dropout |
| 197 | self.out_channels = out_channels or channels |
| 198 | self.use_conv = use_conv |
| 199 | self.use_checkpoint = use_checkpoint |
| 200 | self.use_scale_shift_norm = use_scale_shift_norm |
| 201 | |
| 202 | self.in_layers = nn.Sequential( |
| 203 | normalization(channels), |
| 204 | nn.SiLU(), |
| 205 | conv_nd(dims, channels, self.out_channels, 3, padding=1), |
| 206 | ) |
| 207 | |
| 208 | self.updown = up or down |
| 209 | |
| 210 | if up: |
| 211 | self.h_upd = Upsample(channels, False, dims) |
| 212 | self.x_upd = Upsample(channels, False, dims) |
| 213 | elif down: |
| 214 | self.h_upd = Downsample(channels, False, dims) |
| 215 | self.x_upd = Downsample(channels, False, dims) |
| 216 | else: |
| 217 | self.h_upd = self.x_upd = nn.Identity() |
| 218 | |
| 219 | self.emb_layers = nn.Sequential( |
| 220 | nn.SiLU(), |
| 221 | linear( |
| 222 | emb_channels, |
| 223 | 2 * self.out_channels if use_scale_shift_norm else self.out_channels, |
| 224 | ), |
| 225 | ) |
| 226 | self.out_layers = nn.Sequential( |
| 227 | normalization(self.out_channels), |
| 228 | nn.SiLU(), |
| 229 | nn.Dropout(p=dropout), |
| 230 | zero_module( |
| 231 | conv_nd(dims, self.out_channels, self.out_channels, 3, padding=1) |
| 232 | ), |
| 233 | ) |
| 234 | |
| 235 | if self.out_channels == channels: |
| 236 | self.skip_connection = nn.Identity() |
| 237 | elif use_conv: |
nothing calls this directly
no test coverage detected