(self, z, skipfeatures, temb=None)
| 858 | padding=1) |
| 859 | |
| 860 | def forward(self, z, skipfeatures, temb=None): |
| 861 | #assert z.shape[1:] == self.z_shape[1:] |
| 862 | outputs = {} |
| 863 | self.last_z_shape = z.shape |
| 864 | skiptype = self.skipconnect_type |
| 865 | |
| 866 | h = self.conv_in(z) # (512, 32, 32) |
| 867 | if 'conv_in' in self.output_key: |
| 868 | outputs['conv_in'] = h.clone() |
| 869 | |
| 870 | # middle, (512, 32, 32) |
| 871 | h = self.mid.block_1(h, temb) |
| 872 | h = self.mid.attn_1(h) |
| 873 | h = self.mid.block_2(h, temb) |
| 874 | |
| 875 | if 'mid' in self.output_key: |
| 876 | outputs['mid'] = h.clone() |
| 877 | |
| 878 | # upsampling |
| 879 | # 2 torch.Size([40, 512, 32, 32]) |
| 880 | # 1 torch.Size([40, 256, 64, 64]) |
| 881 | # 0 torch.Size([40, 128, 128, 128]) |
| 882 | # ---------------------------------- |
| 883 | # conv: torch.Size([40, 128, 32, 32]) |
| 884 | # layer1: torch.Size([40, 512, 32, 32]) |
| 885 | # layer2: torch.Size([40, 1024, 32, 32]) |
| 886 | skipfeat = None |
| 887 | for i, i_level in enumerate(reversed(range(self.num_resolutions))): |
| 888 | if i != 0: |
| 889 | skipfeat = skipfeatures[self.skipconnect_layer[i-1]] if self.skipconnect_layer[i-1] is not None else None |
| 890 | for i_block in range(self.num_res_blocks+1): |
| 891 | h = self.up[i_level].block[i_block](h, temb, skipfeat, skiptype=skiptype) |
| 892 | if len(self.up[i_level].attn) > 0: |
| 893 | h = self.up[i_level].attn[i_block](h) |
| 894 | if i_level != 0: |
| 895 | h = self.up[i_level].upsample(h) |
| 896 | |
| 897 | block_idx = self.num_resolutions - i_level - 1 |
| 898 | if f'up_block{block_idx}' in self.output_key: |
| 899 | outputs[f'up_block{block_idx}'] = h.clone() |
| 900 | |
| 901 | # end |
| 902 | if self.give_pre_end: |
| 903 | return h |
| 904 | |
| 905 | h = self.norm_out(h) |
| 906 | if 'norm_out' in self.output_key: |
| 907 | outputs['norm_out'] = h.clone() |
| 908 | |
| 909 | h = nonlinearity(h) |
| 910 | if 'nonlinear' in self.output_key: |
| 911 | outputs['nonlinear'] = h.clone() |
| 912 | |
| 913 | h = self.conv_out(h) |
| 914 | outputs['final'] = h |
| 915 | return outputs |
| 916 | |
| 917 |
nothing calls this directly
no test coverage detected