| 1016 | |
| 1017 | |
| 1018 | class AttnDownBlock2D(nn.Module): |
| 1019 | def __init__( |
| 1020 | self, |
| 1021 | in_channels: int, |
| 1022 | out_channels: int, |
| 1023 | temb_channels: int, |
| 1024 | dropout: float = 0.0, |
| 1025 | num_layers: int = 1, |
| 1026 | resnet_eps: float = 1e-6, |
| 1027 | resnet_time_scale_shift: str = "default", |
| 1028 | resnet_act_fn: str = "swish", |
| 1029 | resnet_groups: int = 32, |
| 1030 | resnet_pre_norm: bool = True, |
| 1031 | attention_head_dim: int = 1, |
| 1032 | output_scale_factor: float = 1.0, |
| 1033 | downsample_padding: int = 1, |
| 1034 | downsample_type: str = "conv", |
| 1035 | ): |
| 1036 | super().__init__() |
| 1037 | resnets = [] |
| 1038 | attentions = [] |
| 1039 | self.downsample_type = downsample_type |
| 1040 | |
| 1041 | if attention_head_dim is None: |
| 1042 | logger.warning( |
| 1043 | f"It is not recommend to pass `attention_head_dim=None`. Defaulting `attention_head_dim` to `in_channels`: {out_channels}." |
| 1044 | ) |
| 1045 | attention_head_dim = out_channels |
| 1046 | |
| 1047 | for i in range(num_layers): |
| 1048 | in_channels = in_channels if i == 0 else out_channels |
| 1049 | resnets.append( |
| 1050 | ResnetBlock2D( |
| 1051 | in_channels=in_channels, |
| 1052 | out_channels=out_channels, |
| 1053 | temb_channels=temb_channels, |
| 1054 | eps=resnet_eps, |
| 1055 | groups=resnet_groups, |
| 1056 | dropout=dropout, |
| 1057 | time_embedding_norm=resnet_time_scale_shift, |
| 1058 | non_linearity=resnet_act_fn, |
| 1059 | output_scale_factor=output_scale_factor, |
| 1060 | pre_norm=resnet_pre_norm, |
| 1061 | ) |
| 1062 | ) |
| 1063 | attentions.append( |
| 1064 | Attention( |
| 1065 | out_channels, |
| 1066 | heads=out_channels // attention_head_dim, |
| 1067 | dim_head=attention_head_dim, |
| 1068 | rescale_output_factor=output_scale_factor, |
| 1069 | eps=resnet_eps, |
| 1070 | norm_num_groups=resnet_groups, |
| 1071 | residual_connection=True, |
| 1072 | bias=True, |
| 1073 | upcast_softmax=True, |
| 1074 | _from_deprecated_attn_block=True, |
| 1075 | ) |
no outgoing calls
no test coverage detected
searching dependent graphs…