r""" See [`T2IAdapter`] for more information.
| 339 | |
| 340 | |
| 341 | class FullAdapterXL(nn.Module): |
| 342 | r""" |
| 343 | See [`T2IAdapter`] for more information. |
| 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""" |
| 376 | This method takes the tensor x as input and processes it through FullAdapterXL model. It consists of operations |
| 377 | including unshuffling pixels, applying convolution layer and appending each block into list of feature tensors. |
| 378 | """ |
| 379 | x = self.unshuffle(x) |
| 380 | x = self.conv_in(x) |
| 381 | |
| 382 | features = [] |
| 383 | |
| 384 | for block in self.body: |
| 385 | x = block(x) |
| 386 | features.append(x) |
| 387 | |
| 388 | return features |
| 389 | |
| 390 | |
| 391 | class AdapterBlock(nn.Module): |