(self, z, skipfeatures, temb=None)
| 1017 | self.zero_test = ignorekwargs.get("zero_test", False) |
| 1018 | |
| 1019 | def forward(self, z, skipfeatures, temb=None): |
| 1020 | #assert z.shape[1:] == self.z_shape[1:] |
| 1021 | outputs = {} |
| 1022 | self.last_z_shape = z.shape |
| 1023 | |
| 1024 | h = self.conv_in(z) # (512, 32, 32) |
| 1025 | if 'conv_in' in self.output_key: |
| 1026 | outputs['conv_in'] = h.clone() |
| 1027 | |
| 1028 | # middle, (512, 32, 32) |
| 1029 | h = self.mid.block_1(h, temb) |
| 1030 | h = self.mid.attn_1(h) |
| 1031 | h = self.mid.block_2(h, temb) |
| 1032 | |
| 1033 | if 'mid' in self.output_key: |
| 1034 | outputs['mid'] = h.clone() |
| 1035 | |
| 1036 | # upsampling |
| 1037 | # 2 torch.Size([40, 512, 32, 32]) + torch.Size([40, 512, 32, 32]) |
| 1038 | # 1 torch.Size([40, 256, 64, 64]) + torch.Size([40, 128, 64, 64]) |
| 1039 | # 0 torch.Size([40, 128, 128, 128]) |
| 1040 | # ---------------------------------- |
| 1041 | # conv: torch.Size([40, 128, 64, 64]) |
| 1042 | # layer1: torch.Size([40, 512, 32, 32]) |
| 1043 | # layer2: torch.Size([40, 1024, 32, 32]) |
| 1044 | skipfeat = None |
| 1045 | for i, i_level in enumerate(reversed(range(self.num_resolutions))): |
| 1046 | if i != 0: |
| 1047 | skipfeat = skipfeatures[self.skipconnect_layer[i-1]] if self.skipconnect_layer[i-1] is not None else None |
| 1048 | if str(i_level) in self.skip_proj: |
| 1049 | skipfeat = self.skip_proj[str(i_level)](skipfeat) |
| 1050 | if self.use_skipconnect_proj: |
| 1051 | skipfeat = self.up[i_level].proj(skipfeat) |
| 1052 | if self.zero_test: |
| 1053 | skipfeat = torch.zeros_like(skipfeat) |
| 1054 | if self.skipconnect_type == 'sum': |
| 1055 | h = h + skipfeat |
| 1056 | elif self.skipconnect_type == 'concat': |
| 1057 | h = torch.cat([h, skipfeat], dim=1) |
| 1058 | else: |
| 1059 | raise NotImplementedError |
| 1060 | for i_block in range(self.num_res_blocks+1): |
| 1061 | h = self.up[i_level].block[i_block](h, temb) |
| 1062 | if len(self.up[i_level].attn) > 0: |
| 1063 | h = self.up[i_level].attn[i_block](h) |
| 1064 | if i_level != 0: |
| 1065 | h = self.up[i_level].upsample(h) |
| 1066 | |
| 1067 | block_idx = self.num_resolutions - i_level - 1 |
| 1068 | if f'up_block{block_idx}' in self.output_key: |
| 1069 | outputs[f'up_block{block_idx}'] = h.clone() |
| 1070 | |
| 1071 | # end |
| 1072 | if self.give_pre_end: |
| 1073 | return h |
| 1074 | |
| 1075 | h = self.norm_out(h) |
| 1076 | if 'norm_out' in self.output_key: |
nothing calls this directly
no test coverage detected