(
self,
*,
ch,
out_ch,
ch_mult=(1, 2, 4, 8),
num_res_blocks,
attn_resolutions,
dropout=0.0,
resamp_with_conv=True,
in_channels,
resolution,
z_channels,
double_z=True,
use_linear_attn=False,
**ignore_kwargs,
)
| 305 | |
| 306 | class Encoder(nn.Module): |
| 307 | def __init__( |
| 308 | self, |
| 309 | *, |
| 310 | ch, |
| 311 | out_ch, |
| 312 | ch_mult=(1, 2, 4, 8), |
| 313 | num_res_blocks, |
| 314 | attn_resolutions, |
| 315 | dropout=0.0, |
| 316 | resamp_with_conv=True, |
| 317 | in_channels, |
| 318 | resolution, |
| 319 | z_channels, |
| 320 | double_z=True, |
| 321 | use_linear_attn=False, |
| 322 | **ignore_kwargs, |
| 323 | ): |
| 324 | super().__init__() |
| 325 | ### setup attention type |
| 326 | if Config.attn_mode == AttnMode.SDP: |
| 327 | attn_type = "sdp" |
| 328 | elif Config.attn_mode == AttnMode.XFORMERS: |
| 329 | attn_type = "xformers" |
| 330 | else: |
| 331 | attn_type = "vanilla" |
| 332 | if use_linear_attn: |
| 333 | attn_type = "linear" |
| 334 | self.ch = ch |
| 335 | self.temb_ch = 0 |
| 336 | self.num_resolutions = len(ch_mult) |
| 337 | self.num_res_blocks = num_res_blocks |
| 338 | self.resolution = resolution |
| 339 | self.in_channels = in_channels |
| 340 | |
| 341 | # downsampling |
| 342 | self.conv_in = torch.nn.Conv2d( |
| 343 | in_channels, self.ch, kernel_size=3, stride=1, padding=1 |
| 344 | ) |
| 345 | |
| 346 | curr_res = resolution |
| 347 | in_ch_mult = (1,) + tuple(ch_mult) |
| 348 | self.in_ch_mult = in_ch_mult |
| 349 | self.down = nn.ModuleList() |
| 350 | for i_level in range(self.num_resolutions): |
| 351 | block = nn.ModuleList() |
| 352 | attn = nn.ModuleList() |
| 353 | block_in = ch * in_ch_mult[i_level] |
| 354 | block_out = ch * ch_mult[i_level] |
| 355 | for i_block in range(self.num_res_blocks): |
| 356 | block.append( |
| 357 | ResnetBlock( |
| 358 | in_channels=block_in, |
| 359 | out_channels=block_out, |
| 360 | temb_channels=self.temb_ch, |
| 361 | dropout=dropout, |
| 362 | ) |
| 363 | ) |
| 364 | block_in = block_out |
no test coverage detected