(
down_block_type: str,
num_layers: int,
in_channels: int,
out_channels: int,
temb_channels: int,
add_downsample: bool,
)
| 624 | |
| 625 | |
| 626 | def get_down_block( |
| 627 | down_block_type: str, |
| 628 | num_layers: int, |
| 629 | in_channels: int, |
| 630 | out_channels: int, |
| 631 | temb_channels: int, |
| 632 | add_downsample: bool, |
| 633 | ) -> DownBlockType: |
| 634 | if down_block_type == "DownResnetBlock1D": |
| 635 | return DownResnetBlock1D( |
| 636 | in_channels=in_channels, |
| 637 | num_layers=num_layers, |
| 638 | out_channels=out_channels, |
| 639 | temb_channels=temb_channels, |
| 640 | add_downsample=add_downsample, |
| 641 | ) |
| 642 | elif down_block_type == "DownBlock1D": |
| 643 | return DownBlock1D(out_channels=out_channels, in_channels=in_channels) |
| 644 | elif down_block_type == "AttnDownBlock1D": |
| 645 | return AttnDownBlock1D(out_channels=out_channels, in_channels=in_channels) |
| 646 | elif down_block_type == "DownBlock1DNoSkip": |
| 647 | return DownBlock1DNoSkip(out_channels=out_channels, in_channels=in_channels) |
| 648 | raise ValueError(f"{down_block_type} does not exist.") |
| 649 | |
| 650 | |
| 651 | def get_up_block( |
no test coverage detected