(
self,
in_channels: int = 3,
channels: List[int] = [320, 640, 1280, 1280],
num_res_blocks: int = 2,
downscale_factor: int = 16,
)
| 344 | """ |
| 345 | |
| 346 | def __init__( |
| 347 | self, |
| 348 | in_channels: int = 3, |
| 349 | channels: List[int] = [320, 640, 1280, 1280], |
| 350 | num_res_blocks: int = 2, |
| 351 | downscale_factor: int = 16, |
| 352 | ): |
| 353 | super().__init__() |
| 354 | |
| 355 | in_channels = in_channels * downscale_factor**2 |
| 356 | |
| 357 | self.unshuffle = nn.PixelUnshuffle(downscale_factor) |
| 358 | self.conv_in = nn.Conv2d(in_channels, channels[0], kernel_size=3, padding=1) |
| 359 | |
| 360 | self.body = [] |
| 361 | # blocks to extract XL features with dimensions of [320, 64, 64], [640, 64, 64], [1280, 32, 32], [1280, 32, 32] |
| 362 | for i in range(len(channels)): |
| 363 | if i == 1: |
| 364 | self.body.append(AdapterBlock(channels[i - 1], channels[i], num_res_blocks)) |
| 365 | elif i == 2: |
| 366 | self.body.append(AdapterBlock(channels[i - 1], channels[i], num_res_blocks, down=True)) |
| 367 | else: |
| 368 | self.body.append(AdapterBlock(channels[i], channels[i], num_res_blocks)) |
| 369 | |
| 370 | self.body = nn.ModuleList(self.body) |
| 371 | # XL has only one downsampling AdapterBlock. |
| 372 | self.total_downscale_factor = downscale_factor * 2 |
| 373 | |
| 374 | def forward(self, x: torch.Tensor) -> List[torch.Tensor]: |
| 375 | r""" |
nothing calls this directly
no test coverage detected