(
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,
norm_type: str = "layer_norm",
transformer_layers_per_block: int = 1,
num_attention_heads: Optional[int] = None,
resnet_groups: Optional[int] = None,
cross_attention_dim: Optional[int] = None,
downsample_padding: Optional[int] = 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",
attention_pre_only: bool = False,
resnet_skip_time_act: bool = False,
resnet_out_scale_factor: float = 1.0,
cross_attention_norm: str | None = None,
attention_head_dim: Optional[int] = None,
use_attention_ffn: bool = True,
downsample_type: str | None = None,
dropout: float = 0.0,
)
| 1594 | |
| 1595 | |
| 1596 | def get_down_block( |
| 1597 | down_block_type: str, |
| 1598 | num_layers: int, |
| 1599 | in_channels: int, |
| 1600 | out_channels: int, |
| 1601 | temb_channels: int, |
| 1602 | add_downsample: bool, |
| 1603 | resnet_eps: float, |
| 1604 | resnet_act_fn: str, |
| 1605 | norm_type: str = "layer_norm", |
| 1606 | transformer_layers_per_block: int = 1, |
| 1607 | num_attention_heads: Optional[int] = None, |
| 1608 | resnet_groups: Optional[int] = None, |
| 1609 | cross_attention_dim: Optional[int] = None, |
| 1610 | downsample_padding: Optional[int] = None, |
| 1611 | dual_cross_attention: bool = False, |
| 1612 | use_linear_projection: bool = False, |
| 1613 | only_cross_attention: bool = False, |
| 1614 | upcast_attention: bool = False, |
| 1615 | resnet_time_scale_shift: str = "default", |
| 1616 | attention_type: str = "default", |
| 1617 | attention_pre_only: bool = False, |
| 1618 | resnet_skip_time_act: bool = False, |
| 1619 | resnet_out_scale_factor: float = 1.0, |
| 1620 | cross_attention_norm: str | None = None, |
| 1621 | attention_head_dim: Optional[int] = None, |
| 1622 | use_attention_ffn: bool = True, |
| 1623 | downsample_type: str | None = None, |
| 1624 | dropout: float = 0.0, |
| 1625 | ): |
| 1626 | # If attn head dim is not defined, we default it to the number of heads |
| 1627 | if attention_head_dim is None: |
| 1628 | logger.warning( |
| 1629 | f"It is recommended to provide `attention_head_dim` when calling `get_down_block`. Defaulting `attention_head_dim` to {num_attention_heads}." |
| 1630 | ) |
| 1631 | attention_head_dim = num_attention_heads |
| 1632 | |
| 1633 | down_block_type = down_block_type[7:] if down_block_type.startswith("UNetRes") else down_block_type |
| 1634 | if down_block_type == "DownBlock2D": |
| 1635 | return DownBlock2D( |
| 1636 | num_layers=num_layers, |
| 1637 | in_channels=in_channels, |
| 1638 | out_channels=out_channels, |
| 1639 | temb_channels=temb_channels, |
| 1640 | dropout=dropout, |
| 1641 | add_downsample=add_downsample, |
| 1642 | resnet_eps=resnet_eps, |
| 1643 | resnet_act_fn=resnet_act_fn, |
| 1644 | resnet_groups=resnet_groups, |
| 1645 | downsample_padding=downsample_padding, |
| 1646 | resnet_time_scale_shift=resnet_time_scale_shift, |
| 1647 | ) |
| 1648 | elif down_block_type == "CrossAttnDownBlock2D": |
| 1649 | if cross_attention_dim is None: |
| 1650 | raise ValueError("cross_attention_dim must be specified for CrossAttnDownBlock2D") |
| 1651 | return CrossAttnDownBlock2D( |
| 1652 | num_layers=num_layers, |
| 1653 | transformer_layers_per_block=transformer_layers_per_block, |
no test coverage detected
searching dependent graphs…