r""" See [`T2IAdapter`] for more information.
| 467 | |
| 468 | |
| 469 | class LightAdapter(nn.Module): |
| 470 | r""" |
| 471 | See [`T2IAdapter`] for more information. |
| 472 | """ |
| 473 | |
| 474 | def __init__( |
| 475 | self, |
| 476 | in_channels: int = 3, |
| 477 | channels: List[int] = [320, 640, 1280], |
| 478 | num_res_blocks: int = 4, |
| 479 | downscale_factor: int = 8, |
| 480 | ): |
| 481 | super().__init__() |
| 482 | |
| 483 | in_channels = in_channels * downscale_factor**2 |
| 484 | |
| 485 | self.unshuffle = nn.PixelUnshuffle(downscale_factor) |
| 486 | |
| 487 | self.body = nn.ModuleList( |
| 488 | [ |
| 489 | LightAdapterBlock(in_channels, channels[0], num_res_blocks), |
| 490 | *[ |
| 491 | LightAdapterBlock(channels[i], channels[i + 1], num_res_blocks, down=True) |
| 492 | for i in range(len(channels) - 1) |
| 493 | ], |
| 494 | LightAdapterBlock(channels[-1], channels[-1], num_res_blocks, down=True), |
| 495 | ] |
| 496 | ) |
| 497 | |
| 498 | self.total_downscale_factor = downscale_factor * (2 ** len(channels)) |
| 499 | |
| 500 | def forward(self, x: torch.Tensor) -> List[torch.Tensor]: |
| 501 | r""" |
| 502 | This method takes the input tensor x and performs downscaling and appends it in list of feature tensors. Each |
| 503 | feature tensor corresponds to a different level of processing within the LightAdapter. |
| 504 | """ |
| 505 | x = self.unshuffle(x) |
| 506 | |
| 507 | features = [] |
| 508 | |
| 509 | for block in self.body: |
| 510 | x = block(x) |
| 511 | features.append(x) |
| 512 | |
| 513 | return features |
| 514 | |
| 515 | |
| 516 | class LightAdapterBlock(nn.Module): |