| 1449 | |
| 1450 | |
| 1451 | class AttnDownEncoderBlock2D(nn.Module): |
| 1452 | def __init__( |
| 1453 | self, |
| 1454 | in_channels: int, |
| 1455 | out_channels: int, |
| 1456 | dropout: float = 0.0, |
| 1457 | num_layers: int = 1, |
| 1458 | resnet_eps: float = 1e-6, |
| 1459 | resnet_time_scale_shift: str = "default", |
| 1460 | resnet_act_fn: str = "swish", |
| 1461 | resnet_groups: int = 32, |
| 1462 | resnet_pre_norm: bool = True, |
| 1463 | attention_head_dim: int = 1, |
| 1464 | output_scale_factor: float = 1.0, |
| 1465 | add_downsample: bool = True, |
| 1466 | downsample_padding: int = 1, |
| 1467 | ): |
| 1468 | super().__init__() |
| 1469 | resnets = [] |
| 1470 | attentions = [] |
| 1471 | |
| 1472 | if attention_head_dim is None: |
| 1473 | logger.warning( |
| 1474 | f"It is not recommend to pass `attention_head_dim=None`. Defaulting `attention_head_dim` to `in_channels`: {out_channels}." |
| 1475 | ) |
| 1476 | attention_head_dim = out_channels |
| 1477 | |
| 1478 | for i in range(num_layers): |
| 1479 | in_channels = in_channels if i == 0 else out_channels |
| 1480 | if resnet_time_scale_shift == "spatial": |
| 1481 | resnets.append( |
| 1482 | ResnetBlockCondNorm2D( |
| 1483 | in_channels=in_channels, |
| 1484 | out_channels=out_channels, |
| 1485 | temb_channels=None, |
| 1486 | eps=resnet_eps, |
| 1487 | groups=resnet_groups, |
| 1488 | dropout=dropout, |
| 1489 | time_embedding_norm="spatial", |
| 1490 | non_linearity=resnet_act_fn, |
| 1491 | output_scale_factor=output_scale_factor, |
| 1492 | ) |
| 1493 | ) |
| 1494 | else: |
| 1495 | resnets.append( |
| 1496 | ResnetBlock2D( |
| 1497 | in_channels=in_channels, |
| 1498 | out_channels=out_channels, |
| 1499 | temb_channels=None, |
| 1500 | eps=resnet_eps, |
| 1501 | groups=resnet_groups, |
| 1502 | dropout=dropout, |
| 1503 | time_embedding_norm=resnet_time_scale_shift, |
| 1504 | non_linearity=resnet_act_fn, |
| 1505 | output_scale_factor=output_scale_factor, |
| 1506 | pre_norm=resnet_pre_norm, |
| 1507 | ) |
| 1508 | ) |
no outgoing calls
no test coverage detected
searching dependent graphs…