(self, x, t, data_batch, aux_info={}, num_samp=1, class_free_guide_w=0.0, apply_guidance=True, guide_clean=False, eval_final_guide_loss=False)
| 842 | |
| 843 | @torch.no_grad() |
| 844 | def p_sample(self, x, t, data_batch, aux_info={}, num_samp=1, class_free_guide_w=0.0, apply_guidance=True, guide_clean=False, eval_final_guide_loss=False): |
| 845 | b, *_, device = *x.shape, x.device |
| 846 | with_func = torch.no_grad |
| 847 | if self.current_perturbation_guidance.current_guidance is not None and apply_guidance and guide_clean == "video_diff": |
| 848 | # will need to take grad wrt noisy |
| 849 | x = x.detach() |
| 850 | x.requires_grad_() |
| 851 | with_func = torch.enable_grad |
| 852 | |
| 853 | with with_func(): |
| 854 | # get prior mean and variance for next step |
| 855 | model_mean, _, model_log_variance, q_posterior_in = self.p_mean_variance(x=x, t=t, aux_info=aux_info, class_free_guide_w=class_free_guide_w) |
| 856 | |
| 857 | sigma = (0.5 * model_log_variance).exp() |
| 858 | |
| 859 | # no noise when t == 0 |
| 860 | # i.e. use the mean of the distribution predicted at the final step rather than sampling. |
| 861 | nonzero_mask = (1 - (t == 0).float()).reshape(b, *((1,) * (len(x.shape) - 1))) |
| 862 | |
| 863 | if self.current_perturbation_guidance.current_guidance is not None and apply_guidance and guide_clean: |
| 864 | # want to guide the predicted clean traj from model, not the noisy one |
| 865 | x_initial = q_posterior_in[0] |
| 866 | return_grad_of = x |
| 867 | else: |
| 868 | x_initial = model_mean.clone().detach() |
| 869 | return_grad_of = x_initial |
| 870 | x_initial.requires_grad_() |
| 871 | |
| 872 | guide_losses = dict() |
| 873 | x_guidance = None |
| 874 | # consider intermediate and final guidance (for ablation) separately |
| 875 | if apply_guidance and self.guidance_optimization_params is not None: |
| 876 | if t[0] == 0: |
| 877 | apply_guidance = self.apply_guidance_output |
| 878 | if apply_guidance: |
| 879 | opt_params = self.final_step_opt_params |
| 880 | else: |
| 881 | apply_guidance = self.apply_guidance_intermediate |
| 882 | if apply_guidance: |
| 883 | assert self.guidance_optimization_params['grad_steps'] > 0 |
| 884 | perturb_th = self.guidance_optimization_params['perturb_th'] |
| 885 | apply_guidance_output = self.apply_guidance_output |
| 886 | lr = self.guidance_optimization_params['lr'] |
| 887 | |
| 888 | if perturb_th is not None: |
| 889 | # gradually decrease clip bounds from 1 to perturb_th |
| 890 | sig_scale = (torch.sigmoid(10 * t[0] / self.n_timesteps) - 1/2) * 2 |
| 891 | perturb_th = sig_scale * (4-perturb_th) + perturb_th |
| 892 | # print(t[0].item(), 'perturb_th', perturb_th) |
| 893 | if not apply_guidance_output: |
| 894 | perturb_th = perturb_th * nonzero_mask |
| 895 | else: |
| 896 | if t[0] == 0 and not apply_guidance_output: |
| 897 | perturb_th = nonzero_mask * sigma |
| 898 | else: |
| 899 | perturb_th = sigma |
| 900 | |
| 901 | if lr is None: |
no test coverage detected