| 12 | self.NUM_DDIM_STEPS = NUM_DDIM_STEPS |
| 13 | |
| 14 | def next_step(self, model_output: Union[torch.FloatTensor, np.ndarray], timestep: int, |
| 15 | sample: Union[torch.FloatTensor, np.ndarray]): |
| 16 | timestep, next_timestep = min( |
| 17 | timestep - self.scheduler.config.num_train_timesteps // self.scheduler.num_inference_steps, 999), timestep |
| 18 | alpha_prod_t = self.scheduler.alphas_cumprod[timestep] if timestep >= 0 else self.scheduler.final_alpha_cumprod |
| 19 | alpha_prod_t_next = self.scheduler.alphas_cumprod[next_timestep] |
| 20 | beta_prod_t = 1 - alpha_prod_t |
| 21 | next_original_sample = (sample - beta_prod_t ** 0.5 * model_output) / alpha_prod_t ** 0.5 |
| 22 | next_sample_direction = (1 - alpha_prod_t_next) ** 0.5 * model_output |
| 23 | next_sample = alpha_prod_t_next ** 0.5 * next_original_sample + next_sample_direction |
| 24 | return next_sample |
| 25 | |
| 26 | def get_noise_pred_single(self, latents, t, context, ref_images_pil=None, pose_cond_fea=None): |
| 27 | noise_pred = self.model( |