MCPcopy Create free account
hub / github.com/huggingface/diffusers / AttnDownBlock2D

Class AttnDownBlock2D

src/diffusers/models/unets/unet_2d_blocks.py:1018–1144  ·  view source on GitHub ↗

Source from the content-addressed store, hash-verified

1016
1017
1018class AttnDownBlock2D(nn.Module):
1019 def __init__(
1020 self,
1021 in_channels: int,
1022 out_channels: int,
1023 temb_channels: int,
1024 dropout: float = 0.0,
1025 num_layers: int = 1,
1026 resnet_eps: float = 1e-6,
1027 resnet_time_scale_shift: str = "default",
1028 resnet_act_fn: str = "swish",
1029 resnet_groups: int = 32,
1030 resnet_pre_norm: bool = True,
1031 attention_head_dim: int = 1,
1032 output_scale_factor: float = 1.0,
1033 downsample_padding: int = 1,
1034 downsample_type: str = "conv",
1035 ):
1036 super().__init__()
1037 resnets = []
1038 attentions = []
1039 self.downsample_type = downsample_type
1040
1041 if attention_head_dim is None:
1042 logger.warning(
1043 f"It is not recommend to pass `attention_head_dim=None`. Defaulting `attention_head_dim` to `in_channels`: {out_channels}."
1044 )
1045 attention_head_dim = out_channels
1046
1047 for i in range(num_layers):
1048 in_channels = in_channels if i == 0 else out_channels
1049 resnets.append(
1050 ResnetBlock2D(
1051 in_channels=in_channels,
1052 out_channels=out_channels,
1053 temb_channels=temb_channels,
1054 eps=resnet_eps,
1055 groups=resnet_groups,
1056 dropout=dropout,
1057 time_embedding_norm=resnet_time_scale_shift,
1058 non_linearity=resnet_act_fn,
1059 output_scale_factor=output_scale_factor,
1060 pre_norm=resnet_pre_norm,
1061 )
1062 )
1063 attentions.append(
1064 Attention(
1065 out_channels,
1066 heads=out_channels // attention_head_dim,
1067 dim_head=attention_head_dim,
1068 rescale_output_factor=output_scale_factor,
1069 eps=resnet_eps,
1070 norm_num_groups=resnet_groups,
1071 residual_connection=True,
1072 bias=True,
1073 upcast_softmax=True,
1074 _from_deprecated_attn_block=True,
1075 )

Callers 1

get_down_blockFunction · 0.85

Calls

no outgoing calls

Tested by

no test coverage detected

Used in the wild real call sites across dependent graphs

searching dependent graphs…