| 1145 | |
| 1146 | |
| 1147 | class CrossAttnDownBlock2D(nn.Module): |
| 1148 | def __init__( |
| 1149 | self, |
| 1150 | in_channels: int, |
| 1151 | out_channels: int, |
| 1152 | temb_channels: int, |
| 1153 | dropout: float = 0.0, |
| 1154 | num_layers: int = 1, |
| 1155 | transformer_layers_per_block: int | tuple[int] = 1, |
| 1156 | resnet_eps: float = 1e-6, |
| 1157 | resnet_time_scale_shift: str = "default", |
| 1158 | resnet_act_fn: str = "swish", |
| 1159 | resnet_groups: int = 32, |
| 1160 | resnet_pre_norm: bool = True, |
| 1161 | num_attention_heads: int = 1, |
| 1162 | cross_attention_dim: int = 1280, |
| 1163 | output_scale_factor: float = 1.0, |
| 1164 | downsample_padding: int = 1, |
| 1165 | add_downsample: bool = True, |
| 1166 | dual_cross_attention: bool = False, |
| 1167 | use_linear_projection: bool = False, |
| 1168 | only_cross_attention: bool = False, |
| 1169 | upcast_attention: bool = False, |
| 1170 | attention_type: str = "default", |
| 1171 | ): |
| 1172 | super().__init__() |
| 1173 | resnets = [] |
| 1174 | attentions = [] |
| 1175 | |
| 1176 | self.has_cross_attention = True |
| 1177 | self.num_attention_heads = num_attention_heads |
| 1178 | if isinstance(transformer_layers_per_block, int): |
| 1179 | transformer_layers_per_block = [transformer_layers_per_block] * num_layers |
| 1180 | |
| 1181 | for i in range(num_layers): |
| 1182 | in_channels = in_channels if i == 0 else out_channels |
| 1183 | resnets.append( |
| 1184 | ResnetBlock2D( |
| 1185 | in_channels=in_channels, |
| 1186 | out_channels=out_channels, |
| 1187 | temb_channels=temb_channels, |
| 1188 | eps=resnet_eps, |
| 1189 | groups=resnet_groups, |
| 1190 | dropout=dropout, |
| 1191 | time_embedding_norm=resnet_time_scale_shift, |
| 1192 | non_linearity=resnet_act_fn, |
| 1193 | output_scale_factor=output_scale_factor, |
| 1194 | pre_norm=resnet_pre_norm, |
| 1195 | ) |
| 1196 | ) |
| 1197 | if not dual_cross_attention: |
| 1198 | attentions.append( |
| 1199 | Transformer2DModel( |
| 1200 | num_attention_heads, |
| 1201 | out_channels // num_attention_heads, |
| 1202 | in_channels=out_channels, |
| 1203 | num_layers=transformer_layers_per_block[i], |
| 1204 | cross_attention_dim=cross_attention_dim, |
no outgoing calls
no test coverage detected
searching dependent graphs…