| 2473 | |
| 2474 | class UpBlock2D(nn.Module): |
| 2475 | def __init__( |
| 2476 | self, |
| 2477 | in_channels: int, |
| 2478 | prev_output_channel: int, |
| 2479 | out_channels: int, |
| 2480 | temb_channels: int, |
| 2481 | resolution_idx: int | None = None, |
| 2482 | dropout: float = 0.0, |
| 2483 | num_layers: int = 1, |
| 2484 | resnet_eps: float = 1e-6, |
| 2485 | resnet_time_scale_shift: str = "default", |
| 2486 | resnet_act_fn: str = "swish", |
| 2487 | resnet_groups: int = 32, |
| 2488 | resnet_pre_norm: bool = True, |
| 2489 | output_scale_factor: float = 1.0, |
| 2490 | add_upsample: bool = True, |
| 2491 | ): |
| 2492 | super().__init__() |
| 2493 | resnets = [] |
| 2494 | |
| 2495 | for i in range(num_layers): |
| 2496 | res_skip_channels = in_channels if (i == num_layers - 1) else out_channels |
| 2497 | resnet_in_channels = prev_output_channel if i == 0 else out_channels |
| 2498 | |
| 2499 | resnets.append( |
| 2500 | ResnetBlock2D( |
| 2501 | in_channels=resnet_in_channels + res_skip_channels, |
| 2502 | out_channels=out_channels, |
| 2503 | temb_channels=temb_channels, |
| 2504 | eps=resnet_eps, |
| 2505 | groups=resnet_groups, |
| 2506 | dropout=dropout, |
| 2507 | time_embedding_norm=resnet_time_scale_shift, |
| 2508 | non_linearity=resnet_act_fn, |
| 2509 | output_scale_factor=output_scale_factor, |
| 2510 | pre_norm=resnet_pre_norm, |
| 2511 | ) |
| 2512 | ) |
| 2513 | |
| 2514 | self.resnets = nn.ModuleList(resnets) |
| 2515 | |
| 2516 | if add_upsample: |
| 2517 | self.upsamplers = nn.ModuleList([Upsample2D(out_channels, use_conv=True, out_channels=out_channels)]) |
| 2518 | else: |
| 2519 | self.upsamplers = None |
| 2520 | |
| 2521 | self.gradient_checkpointing = False |
| 2522 | self.resolution_idx = resolution_idx |
| 2523 | |
| 2524 | def forward( |
| 2525 | self, |