(
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,
give_pre_end=False,
tanh_out=False,
use_linear_attn=False,
**ignorekwargs,
)
| 428 | |
| 429 | class Decoder(nn.Module): |
| 430 | def __init__( |
| 431 | self, |
| 432 | *, |
| 433 | ch, |
| 434 | out_ch, |
| 435 | ch_mult=(1, 2, 4, 8), |
| 436 | num_res_blocks, |
| 437 | attn_resolutions, |
| 438 | dropout=0.0, |
| 439 | resamp_with_conv=True, |
| 440 | in_channels, |
| 441 | resolution, |
| 442 | z_channels, |
| 443 | give_pre_end=False, |
| 444 | tanh_out=False, |
| 445 | use_linear_attn=False, |
| 446 | **ignorekwargs, |
| 447 | ): |
| 448 | super().__init__() |
| 449 | ### setup attention type |
| 450 | if Config.attn_mode == AttnMode.SDP: |
| 451 | attn_type = "sdp" |
| 452 | elif Config.attn_mode == AttnMode.XFORMERS: |
| 453 | attn_type = "xformers" |
| 454 | else: |
| 455 | attn_type = "vanilla" |
| 456 | if use_linear_attn: |
| 457 | attn_type = "linear" |
| 458 | self.ch = ch |
| 459 | self.temb_ch = 0 |
| 460 | self.num_resolutions = len(ch_mult) |
| 461 | self.num_res_blocks = num_res_blocks |
| 462 | self.resolution = resolution |
| 463 | self.in_channels = in_channels |
| 464 | self.give_pre_end = give_pre_end |
| 465 | self.tanh_out = tanh_out |
| 466 | |
| 467 | # compute in_ch_mult, block_in and curr_res at lowest res |
| 468 | in_ch_mult = (1,) + tuple(ch_mult) |
| 469 | block_in = ch * ch_mult[self.num_resolutions - 1] |
| 470 | curr_res = resolution // 2 ** (self.num_resolutions - 1) |
| 471 | self.z_shape = (1, z_channels, curr_res, curr_res) |
| 472 | |
| 473 | # z to block_in |
| 474 | self.conv_in = torch.nn.Conv2d( |
| 475 | z_channels, block_in, kernel_size=3, stride=1, padding=1 |
| 476 | ) |
| 477 | |
| 478 | # middle |
| 479 | self.mid = nn.Module() |
| 480 | self.mid.block_1 = ResnetBlock( |
| 481 | in_channels=block_in, |
| 482 | out_channels=block_in, |
| 483 | temb_channels=self.temb_ch, |
| 484 | dropout=dropout, |
| 485 | ) |
| 486 | self.mid.attn_1 = make_attn(block_in, attn_type=attn_type) |
| 487 | self.mid.block_2 = ResnetBlock( |
nothing calls this directly
no test coverage detected