(self, z)
| 539 | padding=1) |
| 540 | |
| 541 | def forward(self, z): |
| 542 | #assert z.shape[1:] == self.z_shape[1:] |
| 543 | self.last_z_shape = z.shape |
| 544 | |
| 545 | # print(f'decoder-input={z.shape}') |
| 546 | # timestep embedding |
| 547 | temb = None |
| 548 | |
| 549 | # z to block_in |
| 550 | h = self.conv_in(z) |
| 551 | # print(f'decoder-conv in feat={h.shape}') |
| 552 | |
| 553 | # middle |
| 554 | h = self.mid.block_1(h, temb) |
| 555 | h = self.mid.attn_1(h) |
| 556 | h = self.mid.block_2(h, temb) |
| 557 | # print(f'decoder-mid feat={h.shape}') |
| 558 | |
| 559 | # upsampling |
| 560 | for i_level in reversed(range(self.num_resolutions)): |
| 561 | for i_block in range(self.num_res_blocks+1): |
| 562 | h = self.up[i_level].block[i_block](h, temb) |
| 563 | if len(self.up[i_level].attn) > 0: |
| 564 | h = self.up[i_level].attn[i_block](h) |
| 565 | # print(f'decoder-up feat={h.shape}') |
| 566 | if i_level != 0: |
| 567 | h = self.up[i_level].upsample(h) |
| 568 | # print(f'decoder-upsample feat={h.shape}') |
| 569 | |
| 570 | # end |
| 571 | if self.give_pre_end: |
| 572 | return h |
| 573 | |
| 574 | h = self.norm_out(h) |
| 575 | h = nonlinearity(h) |
| 576 | h = self.conv_out(h) |
| 577 | # print(f'decoder-conv_out feat={h.shape}') |
| 578 | if self.tanh_out: |
| 579 | h = torch.tanh(h) |
| 580 | return h |
| 581 | |
| 582 | |
| 583 | class SimpleDecoder(nn.Module): |
nothing calls this directly
no test coverage detected