(self, z)
| 640 | padding=1) |
| 641 | |
| 642 | def forward(self, z): |
| 643 | #assert z.shape[1:] == self.z_shape[1:] |
| 644 | self.last_z_shape = z.shape |
| 645 | |
| 646 | # timestep embedding |
| 647 | temb = None |
| 648 | |
| 649 | # z to block_in |
| 650 | h = self.conv_in(z) |
| 651 | |
| 652 | # middle |
| 653 | h = self.mid.block_1(h, temb) |
| 654 | h = self.mid.attn_1(h) |
| 655 | h = self.mid.block_2(h, temb) |
| 656 | |
| 657 | # upsampling |
| 658 | for i_level in reversed(range(self.num_resolutions)): |
| 659 | for i_block in range(self.num_res_blocks+1): |
| 660 | h = self.up[i_level].block[i_block](h, temb) |
| 661 | if len(self.up[i_level].attn) > 0: |
| 662 | h = self.up[i_level].attn[i_block](h) |
| 663 | if i_level != 0: |
| 664 | h = self.up[i_level].upsample(h) |
| 665 | |
| 666 | # end |
| 667 | if self.give_pre_end: |
| 668 | return h |
| 669 | |
| 670 | h = self.norm_out(h) |
| 671 | h = nonlinearity(h) |
| 672 | h = self.conv_out(h) |
| 673 | if self.tanh_out: |
| 674 | h = torch.tanh(h) |
| 675 | return h |
| 676 | |
| 677 | class Decoder_Mix(nn.Module): |
| 678 | def __init__(self, *, ch, out_ch, ch_mult=(1,2,4,8), num_res_blocks, |
nothing calls this directly
no test coverage detected