(
up_block_type: str, num_layers: int, in_channels: int, out_channels: int, temb_channels: int, add_upsample: bool
)
| 649 | |
| 650 | |
| 651 | def get_up_block( |
| 652 | up_block_type: str, num_layers: int, in_channels: int, out_channels: int, temb_channels: int, add_upsample: bool |
| 653 | ) -> UpBlockType: |
| 654 | if up_block_type == "UpResnetBlock1D": |
| 655 | return UpResnetBlock1D( |
| 656 | in_channels=in_channels, |
| 657 | num_layers=num_layers, |
| 658 | out_channels=out_channels, |
| 659 | temb_channels=temb_channels, |
| 660 | add_upsample=add_upsample, |
| 661 | ) |
| 662 | elif up_block_type == "UpBlock1D": |
| 663 | return UpBlock1D(in_channels=in_channels, out_channels=out_channels) |
| 664 | elif up_block_type == "AttnUpBlock1D": |
| 665 | return AttnUpBlock1D(in_channels=in_channels, out_channels=out_channels) |
| 666 | elif up_block_type == "UpBlock1DNoSkip": |
| 667 | return UpBlock1DNoSkip(in_channels=in_channels, out_channels=out_channels) |
| 668 | raise ValueError(f"{up_block_type} does not exist.") |
| 669 | |
| 670 | |
| 671 | def get_mid_block( |
no test coverage detected