| 650 | |
| 651 | class DPMScheduler(): |
| 652 | def __init__( |
| 653 | self, |
| 654 | beta_start = 0.00085, |
| 655 | beta_end = 0.012, |
| 656 | num_train_timesteps = 1000, |
| 657 | solver_order = 2, |
| 658 | predict_epsilon = True, |
| 659 | thresholding = False, |
| 660 | dynamic_thresholding_ratio = 0.995, |
| 661 | sample_max_value = 1.0, |
| 662 | algorithm_type = "dpmsolver++", |
| 663 | solver_type = "midpoint", |
| 664 | lower_order_final = True, |
| 665 | device = 'cuda', |
| 666 | steps_offset = 0, |
| 667 | prediction_type = 'epsilon' |
| 668 | ): |
| 669 | # this schedule is very specific to the latent diffusion model. |
| 670 | self.betas = ( |
| 671 | torch.linspace(beta_start**0.5, beta_end**0.5, num_train_timesteps, dtype=torch.float32) ** 2 |
| 672 | ) |
| 673 | |
| 674 | self.device = device |
| 675 | self.alphas = 1.0 - self.betas |
| 676 | self.alphas_cumprod = torch.cumprod(self.alphas, dim=0) |
| 677 | # Currently we only support VP-type noise schedule |
| 678 | self.alpha_t = torch.sqrt(self.alphas_cumprod) |
| 679 | self.sigma_t = torch.sqrt(1 - self.alphas_cumprod) |
| 680 | self.lambda_t = torch.log(self.alpha_t) - torch.log(self.sigma_t) |
| 681 | self.steps_offset = steps_offset |
| 682 | |
| 683 | # standard deviation of the initial noise distribution |
| 684 | self.init_noise_sigma = 1.0 |
| 685 | |
| 686 | self.algorithm_type = algorithm_type |
| 687 | self.predict_epsilon = predict_epsilon |
| 688 | self.thresholding = thresholding |
| 689 | self.dynamic_thresholding_ratio = dynamic_thresholding_ratio |
| 690 | self.sample_max_value = sample_max_value |
| 691 | self.lower_order_final = lower_order_final |
| 692 | self.prediction_type = prediction_type |
| 693 | |
| 694 | # settings for DPM-Solver |
| 695 | if algorithm_type not in ["dpmsolver", "dpmsolver++"]: |
| 696 | raise NotImplementedError(f"{algorithm_type} does is not implemented for {self.__class__}") |
| 697 | if solver_type not in ["midpoint", "heun"]: |
| 698 | raise NotImplementedError(f"{solver_type} does is not implemented for {self.__class__}") |
| 699 | |
| 700 | # setable values |
| 701 | self.num_inference_steps = None |
| 702 | self.solver_order = solver_order |
| 703 | self.num_train_timesteps = num_train_timesteps |
| 704 | self.solver_type = solver_type |
| 705 | |
| 706 | self.first_order_first_coef = [] |
| 707 | self.first_order_second_coef = [] |
| 708 | |
| 709 | self.second_order_first_coef = [] |