| 344 | |
| 345 | |
| 346 | class DownBlock2D(nn.Module): |
| 347 | |
| 348 | def __init__( |
| 349 | self, |
| 350 | in_channels, |
| 351 | out_channels, |
| 352 | temb_channels, |
| 353 | num_layers=1, |
| 354 | resnet_eps=1e-6, |
| 355 | resnet_time_scale_shift="default", |
| 356 | resnet_act_fn="swish", |
| 357 | resnet_groups=32, |
| 358 | add_downsample=True, |
| 359 | ): |
| 360 | super().__init__() |
| 361 | resnets = [] |
| 362 | |
| 363 | for i in range(num_layers): |
| 364 | in_channels = in_channels if i == 0 else out_channels |
| 365 | resnets.append( |
| 366 | ResnetBlock2D( |
| 367 | in_channels=in_channels, |
| 368 | out_channels=out_channels, |
| 369 | temb_channels=temb_channels, |
| 370 | eps=resnet_eps, |
| 371 | groups=resnet_groups, |
| 372 | time_embedding_norm=resnet_time_scale_shift, |
| 373 | )) |
| 374 | |
| 375 | self.resnets = nn.ModuleList(resnets) |
| 376 | |
| 377 | if add_downsample: |
| 378 | self.downsamplers = nn.ModuleList([Downsample2D(out_channels)]) |
| 379 | else: |
| 380 | self.downsamplers = None |
| 381 | |
| 382 | def forward(self, hidden_states, temb=None): |
| 383 | output_states = () |
| 384 | |
| 385 | for resnet in self.resnets: |
| 386 | hidden_states = resnet(hidden_states, temb) |
| 387 | output_states += (hidden_states, ) |
| 388 | |
| 389 | if self.downsamplers is not None: |
| 390 | for downsampler in self.downsamplers: |
| 391 | hidden_states = downsampler(hidden_states) |
| 392 | |
| 393 | output_states = output_states + (hidden_states,) |
| 394 | |
| 395 | |
| 396 | return hidden_states, output_states |
| 397 | |
| 398 | |
| 399 | class ResnetBlock2D(nn.Module): |