| 243 | |
| 244 | @register_to_config |
| 245 | def __init__( |
| 246 | self, |
| 247 | in_channels: int = 3, |
| 248 | channels: List[int] = [320, 640, 1280, 1280], |
| 249 | num_res_blocks: int = 2, |
| 250 | downscale_factor: int = 8, |
| 251 | adapter_type: str = "full_adapter", |
| 252 | ): |
| 253 | super().__init__() |
| 254 | |
| 255 | if adapter_type == "full_adapter": |
| 256 | self.adapter = FullAdapter(in_channels, channels, num_res_blocks, downscale_factor) |
| 257 | elif adapter_type == "full_adapter_xl": |
| 258 | self.adapter = FullAdapterXL(in_channels, channels, num_res_blocks, downscale_factor) |
| 259 | elif adapter_type == "light_adapter": |
| 260 | self.adapter = LightAdapter(in_channels, channels, num_res_blocks, downscale_factor) |
| 261 | else: |
| 262 | raise ValueError( |
| 263 | f"Unsupported adapter_type: '{adapter_type}'. Choose either 'full_adapter' or " |
| 264 | "'full_adapter_xl' or 'light_adapter'." |
| 265 | ) |
| 266 | |
| 267 | def forward(self, x: torch.Tensor) -> List[torch.Tensor]: |
| 268 | r""" |