(self, x0, t, denoise_fn, noise_fn=torch.randn, constrain_fn=lambda x, t:x)
| 228 | return img_t |
| 229 | |
| 230 | def reconstruct(self, x0, t, denoise_fn, noise_fn=torch.randn, constrain_fn=lambda x, t:x): |
| 231 | |
| 232 | assert t >= 1 |
| 233 | |
| 234 | t_vec = torch.empty(x0.shape[0], dtype=torch.int64, device=x0.device).fill_(t-1) |
| 235 | encoding = self.q_sample(x0, t_vec) |
| 236 | |
| 237 | img_t = encoding |
| 238 | |
| 239 | for k in reversed(range(0,t)): |
| 240 | img_t = constrain_fn(img_t, k) |
| 241 | t_ = torch.empty(x0.shape[0], dtype=torch.int64, device=x0.device).fill_(k) |
| 242 | img_t = self.p_sample(denoise_fn=denoise_fn, data=img_t, t=t_, noise_fn=noise_fn, |
| 243 | clip_denoised=False, return_pred_xstart=False, use_var=True).detach() |
| 244 | |
| 245 | |
| 246 | return img_t |
| 247 | |
| 248 | |
| 249 | class PVCNN2(PVCNN2Base): |
nothing calls this directly
no test coverage detected