| 305 | |
| 306 | |
| 307 | class UVitBlock(nn.Module): |
| 308 | def __init__( |
| 309 | self, |
| 310 | channels, |
| 311 | num_res_blocks: int, |
| 312 | hidden_size, |
| 313 | hidden_dropout, |
| 314 | ln_elementwise_affine, |
| 315 | layer_norm_eps, |
| 316 | use_bias, |
| 317 | block_num_heads, |
| 318 | attention_dropout, |
| 319 | downsample: bool, |
| 320 | upsample: bool, |
| 321 | ): |
| 322 | super().__init__() |
| 323 | |
| 324 | if downsample: |
| 325 | self.downsample = Downsample2D( |
| 326 | channels, |
| 327 | use_conv=True, |
| 328 | padding=0, |
| 329 | name="Conv2d_0", |
| 330 | kernel_size=2, |
| 331 | norm_type="rms_norm", |
| 332 | eps=layer_norm_eps, |
| 333 | elementwise_affine=ln_elementwise_affine, |
| 334 | bias=use_bias, |
| 335 | ) |
| 336 | else: |
| 337 | self.downsample = None |
| 338 | |
| 339 | self.res_blocks = nn.ModuleList( |
| 340 | [ |
| 341 | ConvNextBlock( |
| 342 | channels, |
| 343 | layer_norm_eps, |
| 344 | ln_elementwise_affine, |
| 345 | use_bias, |
| 346 | hidden_dropout, |
| 347 | hidden_size, |
| 348 | ) |
| 349 | for i in range(num_res_blocks) |
| 350 | ] |
| 351 | ) |
| 352 | |
| 353 | self.attention_blocks = nn.ModuleList( |
| 354 | [ |
| 355 | SkipFFTransformerBlock( |
| 356 | channels, |
| 357 | block_num_heads, |
| 358 | channels // block_num_heads, |
| 359 | hidden_size, |
| 360 | use_bias, |
| 361 | attention_dropout, |
| 362 | channels, |
| 363 | attention_bias=use_bias, |
| 364 | attention_out_bias=use_bias, |