(self, z, temb=None)
| 601 | padding=1) |
| 602 | |
| 603 | def forward(self, z, temb=None): |
| 604 | #assert z.shape[1:] == self.z_shape[1:] |
| 605 | outputs = {} |
| 606 | self.last_z_shape = z.shape |
| 607 | |
| 608 | # timestep embedding |
| 609 | #temb = None |
| 610 | |
| 611 | # z to block_in |
| 612 | h = self.conv_in(z) |
| 613 | if 'conv_in' in self.output_key: |
| 614 | outputs['conv_in'] = h.clone() |
| 615 | |
| 616 | # middle |
| 617 | h = self.mid.block_1(h, temb) |
| 618 | h = self.mid.attn_1(h) |
| 619 | h = self.mid.block_2(h, temb) |
| 620 | |
| 621 | if 'mid' in self.output_key: |
| 622 | outputs['mid'] = h.clone() |
| 623 | |
| 624 | # upsampling |
| 625 | for i_level in reversed(range(self.num_resolutions)): |
| 626 | for i_block in range(self.num_res_blocks+1): |
| 627 | h = self.up[i_level].block[i_block](h, temb) |
| 628 | if len(self.up[i_level].attn) > 0: |
| 629 | h = self.up[i_level].attn[i_block](h) |
| 630 | if i_level != 0: |
| 631 | h = self.up[i_level].upsample(h) |
| 632 | |
| 633 | block_idx = self.num_resolutions - i_level - 1 |
| 634 | if f'up_block{block_idx}' in self.output_key: |
| 635 | outputs[f'up_block{block_idx}'] = h.clone() |
| 636 | |
| 637 | # end |
| 638 | if self.give_pre_end: |
| 639 | return h |
| 640 | |
| 641 | h = self.norm_out(h) |
| 642 | if 'norm_out' in self.output_key: |
| 643 | outputs['norm_out'] = h.clone() |
| 644 | |
| 645 | h = nonlinearity(h) |
| 646 | if 'nonlinear' in self.output_key: |
| 647 | outputs['nonlinear'] = h.clone() |
| 648 | |
| 649 | h = self.conv_out(h) |
| 650 | outputs['final'] = h |
| 651 | return outputs |
| 652 | |
| 653 | |
| 654 | class DecoderNoAttn(nn.Module): |
nothing calls this directly
no test coverage detected