(self, z)
| 626 | padding=1) |
| 627 | |
| 628 | def forward(self, z): |
| 629 | #assert z.shape[1:] == self.z_shape[1:] |
| 630 | self.last_z_shape = z.shape |
| 631 | |
| 632 | # timestep embedding |
| 633 | temb = None |
| 634 | |
| 635 | # z to block_in |
| 636 | h = self.conv_in(z) |
| 637 | |
| 638 | # middle |
| 639 | h = self.mid.block_1(h, temb) |
| 640 | h = self.mid.attn_1(h) |
| 641 | h = self.mid.block_2(h, temb) |
| 642 | |
| 643 | # upsampling |
| 644 | for i_level in reversed(range(self.num_resolutions)): |
| 645 | for i_block in range(self.num_res_blocks+1): |
| 646 | h = self.up[i_level].block[i_block](h, temb) |
| 647 | if len(self.up[i_level].attn) > 0: |
| 648 | h = self.up[i_level].attn[i_block](h) |
| 649 | if i_level != 0: |
| 650 | h = self.up[i_level].upsample(h) |
| 651 | |
| 652 | # end |
| 653 | if self.give_pre_end: |
| 654 | return h |
| 655 | |
| 656 | h = self.norm_out(h) |
| 657 | h = nonlinearity(h) |
| 658 | h = self.conv_out(h) |
| 659 | if self.tanh_out: |
| 660 | h = torch.tanh(h) |
| 661 | return h |
| 662 | |
| 663 | |
| 664 | class SimpleDecoder(nn.Module): |
nothing calls this directly
no test coverage detected