(
mid_block_type: str,
num_layers: int,
in_channels: int,
mid_channels: int,
out_channels: int,
embed_dim: int,
add_downsample: bool,
)
| 669 | |
| 670 | |
| 671 | def get_mid_block( |
| 672 | mid_block_type: str, |
| 673 | num_layers: int, |
| 674 | in_channels: int, |
| 675 | mid_channels: int, |
| 676 | out_channels: int, |
| 677 | embed_dim: int, |
| 678 | add_downsample: bool, |
| 679 | ) -> MidBlockType: |
| 680 | if mid_block_type == "MidResTemporalBlock1D": |
| 681 | return MidResTemporalBlock1D( |
| 682 | num_layers=num_layers, |
| 683 | in_channels=in_channels, |
| 684 | out_channels=out_channels, |
| 685 | embed_dim=embed_dim, |
| 686 | add_downsample=add_downsample, |
| 687 | ) |
| 688 | elif mid_block_type == "ValueFunctionMidBlock1D": |
| 689 | return ValueFunctionMidBlock1D(in_channels=in_channels, out_channels=out_channels, embed_dim=embed_dim) |
| 690 | elif mid_block_type == "UNetMidBlock1D": |
| 691 | return UNetMidBlock1D(in_channels=in_channels, mid_channels=mid_channels, out_channels=out_channels) |
| 692 | raise ValueError(f"{mid_block_type} does not exist.") |
| 693 | |
| 694 | |
| 695 | def get_out_block( |
no test coverage detected