(
self,
in_channels: int,
out_channels: int,
temb_channels: int,
resolution_idx: int,
dropout: float = 0.0,
num_layers: int = 5,
resnet_eps: float = 1e-5,
resnet_act_fn: str = "gelu",
resnet_group_size: int | None = 32,
add_upsample: bool = True,
)
| 3247 | |
| 3248 | class KUpBlock2D(nn.Module): |
| 3249 | def __init__( |
| 3250 | self, |
| 3251 | in_channels: int, |
| 3252 | out_channels: int, |
| 3253 | temb_channels: int, |
| 3254 | resolution_idx: int, |
| 3255 | dropout: float = 0.0, |
| 3256 | num_layers: int = 5, |
| 3257 | resnet_eps: float = 1e-5, |
| 3258 | resnet_act_fn: str = "gelu", |
| 3259 | resnet_group_size: int | None = 32, |
| 3260 | add_upsample: bool = True, |
| 3261 | ): |
| 3262 | super().__init__() |
| 3263 | resnets = [] |
| 3264 | k_in_channels = 2 * out_channels |
| 3265 | k_out_channels = in_channels |
| 3266 | num_layers = num_layers - 1 |
| 3267 | |
| 3268 | for i in range(num_layers): |
| 3269 | in_channels = k_in_channels if i == 0 else out_channels |
| 3270 | groups = in_channels // resnet_group_size |
| 3271 | groups_out = out_channels // resnet_group_size |
| 3272 | |
| 3273 | resnets.append( |
| 3274 | ResnetBlockCondNorm2D( |
| 3275 | in_channels=in_channels, |
| 3276 | out_channels=k_out_channels if (i == num_layers - 1) else out_channels, |
| 3277 | temb_channels=temb_channels, |
| 3278 | eps=resnet_eps, |
| 3279 | groups=groups, |
| 3280 | groups_out=groups_out, |
| 3281 | dropout=dropout, |
| 3282 | non_linearity=resnet_act_fn, |
| 3283 | time_embedding_norm="ada_group", |
| 3284 | conv_shortcut_bias=False, |
| 3285 | ) |
| 3286 | ) |
| 3287 | |
| 3288 | self.resnets = nn.ModuleList(resnets) |
| 3289 | |
| 3290 | if add_upsample: |
| 3291 | self.upsamplers = nn.ModuleList([KUpsample2D()]) |
| 3292 | else: |
| 3293 | self.upsamplers = None |
| 3294 | |
| 3295 | self.gradient_checkpointing = False |
| 3296 | self.resolution_idx = resolution_idx |
| 3297 | |
| 3298 | def forward( |
| 3299 | self, |
nothing calls this directly
no test coverage detected