| 1370 | |
| 1371 | |
| 1372 | class DownEncoderBlock2D(nn.Module): |
| 1373 | def __init__( |
| 1374 | self, |
| 1375 | in_channels: int, |
| 1376 | out_channels: int, |
| 1377 | dropout: float = 0.0, |
| 1378 | num_layers: int = 1, |
| 1379 | resnet_eps: float = 1e-6, |
| 1380 | resnet_time_scale_shift: str = "default", |
| 1381 | resnet_act_fn: str = "swish", |
| 1382 | resnet_groups: int = 32, |
| 1383 | resnet_pre_norm: bool = True, |
| 1384 | output_scale_factor: float = 1.0, |
| 1385 | add_downsample: bool = True, |
| 1386 | downsample_padding: int = 1, |
| 1387 | ): |
| 1388 | super().__init__() |
| 1389 | resnets = [] |
| 1390 | |
| 1391 | for i in range(num_layers): |
| 1392 | in_channels = in_channels if i == 0 else out_channels |
| 1393 | if resnet_time_scale_shift == "spatial": |
| 1394 | resnets.append( |
| 1395 | ResnetBlockCondNorm2D( |
| 1396 | in_channels=in_channels, |
| 1397 | out_channels=out_channels, |
| 1398 | temb_channels=None, |
| 1399 | eps=resnet_eps, |
| 1400 | groups=resnet_groups, |
| 1401 | dropout=dropout, |
| 1402 | time_embedding_norm="spatial", |
| 1403 | non_linearity=resnet_act_fn, |
| 1404 | output_scale_factor=output_scale_factor, |
| 1405 | ) |
| 1406 | ) |
| 1407 | else: |
| 1408 | resnets.append( |
| 1409 | ResnetBlock2D( |
| 1410 | in_channels=in_channels, |
| 1411 | out_channels=out_channels, |
| 1412 | temb_channels=None, |
| 1413 | eps=resnet_eps, |
| 1414 | groups=resnet_groups, |
| 1415 | dropout=dropout, |
| 1416 | time_embedding_norm=resnet_time_scale_shift, |
| 1417 | non_linearity=resnet_act_fn, |
| 1418 | output_scale_factor=output_scale_factor, |
| 1419 | pre_norm=resnet_pre_norm, |
| 1420 | ) |
| 1421 | ) |
| 1422 | |
| 1423 | self.resnets = nn.ModuleList(resnets) |
| 1424 | |
| 1425 | if add_downsample: |
| 1426 | self.downsamplers = nn.ModuleList( |
| 1427 | [ |
| 1428 | Downsample2D( |
| 1429 | out_channels, use_conv=True, out_channels=out_channels, padding=downsample_padding, name="op" |
no outgoing calls
no test coverage detected
searching dependent graphs…