(
self,
in_channels: int,
out_channels: int,
temb_channels: int,
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,
output_scale_factor: float = 1.0,
add_downsample: bool = True,
downsample_padding: int = 1,
)
| 1293 | |
| 1294 | class DownBlock2D(nn.Module): |
| 1295 | def __init__( |
| 1296 | self, |
| 1297 | in_channels: int, |
| 1298 | out_channels: int, |
| 1299 | temb_channels: int, |
| 1300 | dropout: float = 0.0, |
| 1301 | num_layers: int = 1, |
| 1302 | resnet_eps: float = 1e-6, |
| 1303 | resnet_time_scale_shift: str = "default", |
| 1304 | resnet_act_fn: str = "swish", |
| 1305 | resnet_groups: int = 32, |
| 1306 | resnet_pre_norm: bool = True, |
| 1307 | output_scale_factor: float = 1.0, |
| 1308 | add_downsample: bool = True, |
| 1309 | downsample_padding: int = 1, |
| 1310 | ): |
| 1311 | super().__init__() |
| 1312 | resnets = [] |
| 1313 | |
| 1314 | for i in range(num_layers): |
| 1315 | in_channels = in_channels if i == 0 else out_channels |
| 1316 | resnets.append( |
| 1317 | ResnetBlock2D( |
| 1318 | in_channels=in_channels, |
| 1319 | out_channels=out_channels, |
| 1320 | temb_channels=temb_channels, |
| 1321 | eps=resnet_eps, |
| 1322 | groups=resnet_groups, |
| 1323 | dropout=dropout, |
| 1324 | time_embedding_norm=resnet_time_scale_shift, |
| 1325 | non_linearity=resnet_act_fn, |
| 1326 | output_scale_factor=output_scale_factor, |
| 1327 | pre_norm=resnet_pre_norm, |
| 1328 | ) |
| 1329 | ) |
| 1330 | |
| 1331 | self.resnets = nn.ModuleList(resnets) |
| 1332 | |
| 1333 | if add_downsample: |
| 1334 | self.downsamplers = nn.ModuleList( |
| 1335 | [ |
| 1336 | Downsample2D( |
| 1337 | out_channels, use_conv=True, out_channels=out_channels, padding=downsample_padding, name="op" |
| 1338 | ) |
| 1339 | ] |
| 1340 | ) |
| 1341 | else: |
| 1342 | self.downsamplers = None |
| 1343 | |
| 1344 | self.gradient_checkpointing = False |
| 1345 | |
| 1346 | def forward( |
| 1347 | self, hidden_states: torch.Tensor, temb: torch.Tensor | None = None, *args, **kwargs |
nothing calls this directly
no test coverage detected