(self, z)
| 533 | padding=1) |
| 534 | |
| 535 | def forward(self, z): |
| 536 | #assert z.shape[1:] == self.z_shape[1:] |
| 537 | self.last_z_shape = z.shape |
| 538 | |
| 539 | # timestep embedding |
| 540 | temb = None |
| 541 | |
| 542 | # z to block_in |
| 543 | h = self.conv_in(z) |
| 544 | |
| 545 | # middle |
| 546 | h = self.mid.block_1(h, temb) |
| 547 | h = self.mid.attn_1(h) |
| 548 | h = self.mid.block_2(h, temb) |
| 549 | |
| 550 | # upsampling |
| 551 | for i_level in reversed(range(self.num_resolutions)): |
| 552 | for i_block in range(self.num_res_blocks+1): |
| 553 | h = self.up[i_level].block[i_block](h, temb) |
| 554 | if len(self.up[i_level].attn) > 0: |
| 555 | h = self.up[i_level].attn[i_block](h) |
| 556 | if i_level != 0: |
| 557 | h = self.up[i_level].upsample(h) |
| 558 | |
| 559 | # end |
| 560 | if self.give_pre_end: |
| 561 | return h |
| 562 | |
| 563 | h = self.norm_out(h) |
| 564 | h = nonlinearity(h) |
| 565 | h = self.conv_out(h) |
| 566 | if self.tanh_out: |
| 567 | h = torch.tanh(h) |
| 568 | return h |
| 569 | |
| 570 | |
| 571 | class SimpleDecoder(nn.Module): |
nothing calls this directly
no test coverage detected