(
self,
in_channels: int,
prev_output_channel: int,
out_channels: int,
temb_channels: int,
resolution_idx: int = None,
dropout: float = 0.0,
num_layers: int = 1,
resnet_eps: float = 1e-6,
resnet_time_scale_shift: str = "default",
resnet_act_fn: str = "swish",
resnet_groups: int = 32,
resnet_pre_norm: bool = True,
attention_head_dim: int = 1,
output_scale_factor: float = 1.0,
upsample_type: str = "conv",
)
| 2184 | |
| 2185 | class AttnUpBlock2D(nn.Module): |
| 2186 | def __init__( |
| 2187 | self, |
| 2188 | in_channels: int, |
| 2189 | prev_output_channel: int, |
| 2190 | out_channels: int, |
| 2191 | temb_channels: int, |
| 2192 | resolution_idx: int = None, |
| 2193 | dropout: float = 0.0, |
| 2194 | num_layers: int = 1, |
| 2195 | resnet_eps: float = 1e-6, |
| 2196 | resnet_time_scale_shift: str = "default", |
| 2197 | resnet_act_fn: str = "swish", |
| 2198 | resnet_groups: int = 32, |
| 2199 | resnet_pre_norm: bool = True, |
| 2200 | attention_head_dim: int = 1, |
| 2201 | output_scale_factor: float = 1.0, |
| 2202 | upsample_type: str = "conv", |
| 2203 | ): |
| 2204 | super().__init__() |
| 2205 | resnets = [] |
| 2206 | attentions = [] |
| 2207 | |
| 2208 | self.upsample_type = upsample_type |
| 2209 | |
| 2210 | if attention_head_dim is None: |
| 2211 | logger.warning( |
| 2212 | f"It is not recommend to pass `attention_head_dim=None`. Defaulting `attention_head_dim` to `in_channels`: {out_channels}." |
| 2213 | ) |
| 2214 | attention_head_dim = out_channels |
| 2215 | |
| 2216 | for i in range(num_layers): |
| 2217 | res_skip_channels = in_channels if (i == num_layers - 1) else out_channels |
| 2218 | resnet_in_channels = prev_output_channel if i == 0 else out_channels |
| 2219 | |
| 2220 | resnets.append( |
| 2221 | ResnetBlock2D( |
| 2222 | in_channels=resnet_in_channels + res_skip_channels, |
| 2223 | out_channels=out_channels, |
| 2224 | temb_channels=temb_channels, |
| 2225 | eps=resnet_eps, |
| 2226 | groups=resnet_groups, |
| 2227 | dropout=dropout, |
| 2228 | time_embedding_norm=resnet_time_scale_shift, |
| 2229 | non_linearity=resnet_act_fn, |
| 2230 | output_scale_factor=output_scale_factor, |
| 2231 | pre_norm=resnet_pre_norm, |
| 2232 | ) |
| 2233 | ) |
| 2234 | attentions.append( |
| 2235 | Attention( |
| 2236 | out_channels, |
| 2237 | heads=out_channels // attention_head_dim, |
| 2238 | dim_head=attention_head_dim, |
| 2239 | rescale_output_factor=output_scale_factor, |
| 2240 | eps=resnet_eps, |
| 2241 | norm_num_groups=resnet_groups, |
| 2242 | residual_connection=True, |
| 2243 | bias=True, |
nothing calls this directly
no test coverage detected