(
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,
)
| 232 | """ |
| 233 | |
| 234 | def __init__( |
| 235 | self, |
| 236 | channels, |
| 237 | emb_channels, |
| 238 | dropout, |
| 239 | out_channels=None, |
| 240 | use_conv=False, |
| 241 | use_scale_shift_norm=False, |
| 242 | dims=2, |
| 243 | use_checkpoint=False, |
| 244 | up=False, |
| 245 | down=False, |
| 246 | ): |
| 247 | super().__init__() |
| 248 | self.channels = channels |
| 249 | self.emb_channels = emb_channels |
| 250 | self.dropout = dropout |
| 251 | self.out_channels = out_channels or channels |
| 252 | self.use_conv = use_conv |
| 253 | self.use_checkpoint = use_checkpoint |
| 254 | self.use_scale_shift_norm = use_scale_shift_norm |
| 255 | |
| 256 | self.in_layers = nn.Sequential( |
| 257 | normalization(channels), |
| 258 | nn.SiLU(), |
| 259 | conv_nd(dims, channels, self.out_channels, 3, padding=1), |
| 260 | ) |
| 261 | |
| 262 | self.updown = up or down |
| 263 | |
| 264 | if up: |
| 265 | self.h_upd = Upsample(channels, False, dims) |
| 266 | self.x_upd = Upsample(channels, False, dims) |
| 267 | elif down: |
| 268 | self.h_upd = Downsample(channels, False, dims) |
| 269 | self.x_upd = Downsample(channels, False, dims) |
| 270 | else: |
| 271 | self.h_upd = self.x_upd = nn.Identity() |
| 272 | |
| 273 | self.emb_layers = nn.Sequential( |
| 274 | nn.SiLU(), |
| 275 | linear( |
| 276 | emb_channels, |
| 277 | 2 * self.out_channels if use_scale_shift_norm else self.out_channels, |
| 278 | ), |
| 279 | ) |
| 280 | |
| 281 | if self.out_channels % 32 == 0: |
| 282 | self.out_layers = nn.Sequential( |
| 283 | normalization(self.out_channels), |
| 284 | nn.SiLU(), |
| 285 | nn.Dropout(p=dropout), |
| 286 | zero_module( |
| 287 | conv_nd(dims, self.out_channels, self.out_channels, 3, padding=1) |
| 288 | ), |
| 289 | ) |
| 290 | else: |
| 291 | self.out_layers = nn.Sequential( |
nothing calls this directly
no test coverage detected