(
down_block_type: str,
num_layers: int,
in_channels: int,
out_channels: int,
temb_channels: int,
add_downsample: bool,
resnet_eps: float,
resnet_act_fn: str,
transformer_layers_per_block: int = 1,
num_attention_heads: int | None = None,
resnet_groups: int | None = None,
cross_attention_dim: int | None = None,
downsample_padding: int | None = None,
dual_cross_attention: bool = False,
use_linear_projection: bool = False,
only_cross_attention: bool = False,
upcast_attention: bool = False,
resnet_time_scale_shift: str = "default",
attention_type: str = "default",
resnet_skip_time_act: bool = False,
resnet_out_scale_factor: float = 1.0,
cross_attention_norm: str | None = None,
attention_head_dim: int | None = None,
downsample_type: str | None = None,
dropout: float = 0.0,
)
| 41 | |
| 42 | |
| 43 | def get_down_block( |
| 44 | down_block_type: str, |
| 45 | num_layers: int, |
| 46 | in_channels: int, |
| 47 | out_channels: int, |
| 48 | temb_channels: int, |
| 49 | add_downsample: bool, |
| 50 | resnet_eps: float, |
| 51 | resnet_act_fn: str, |
| 52 | transformer_layers_per_block: int = 1, |
| 53 | num_attention_heads: int | None = None, |
| 54 | resnet_groups: int | None = None, |
| 55 | cross_attention_dim: int | None = None, |
| 56 | downsample_padding: int | None = None, |
| 57 | dual_cross_attention: bool = False, |
| 58 | use_linear_projection: bool = False, |
| 59 | only_cross_attention: bool = False, |
| 60 | upcast_attention: bool = False, |
| 61 | resnet_time_scale_shift: str = "default", |
| 62 | attention_type: str = "default", |
| 63 | resnet_skip_time_act: bool = False, |
| 64 | resnet_out_scale_factor: float = 1.0, |
| 65 | cross_attention_norm: str | None = None, |
| 66 | attention_head_dim: int | None = None, |
| 67 | downsample_type: str | None = None, |
| 68 | dropout: float = 0.0, |
| 69 | ): |
| 70 | # If attn head dim is not defined, we default it to the number of heads |
| 71 | if attention_head_dim is None: |
| 72 | logger.warning( |
| 73 | f"It is recommended to provide `attention_head_dim` when calling `get_down_block`. Defaulting `attention_head_dim` to {num_attention_heads}." |
| 74 | ) |
| 75 | attention_head_dim = num_attention_heads |
| 76 | |
| 77 | down_block_type = down_block_type[7:] if down_block_type.startswith("UNetRes") else down_block_type |
| 78 | if down_block_type == "DownBlock2D": |
| 79 | return DownBlock2D( |
| 80 | num_layers=num_layers, |
| 81 | in_channels=in_channels, |
| 82 | out_channels=out_channels, |
| 83 | temb_channels=temb_channels, |
| 84 | dropout=dropout, |
| 85 | add_downsample=add_downsample, |
| 86 | resnet_eps=resnet_eps, |
| 87 | resnet_act_fn=resnet_act_fn, |
| 88 | resnet_groups=resnet_groups, |
| 89 | downsample_padding=downsample_padding, |
| 90 | resnet_time_scale_shift=resnet_time_scale_shift, |
| 91 | ) |
| 92 | elif down_block_type == "ResnetDownsampleBlock2D": |
| 93 | return ResnetDownsampleBlock2D( |
| 94 | num_layers=num_layers, |
| 95 | in_channels=in_channels, |
| 96 | out_channels=out_channels, |
| 97 | temb_channels=temb_channels, |
| 98 | dropout=dropout, |
| 99 | add_downsample=add_downsample, |
| 100 | resnet_eps=resnet_eps, |
no test coverage detected
searching dependent graphs…