DPM-Solver++(2M) SDE.
(model, x, sigmas, cfg_val, cfg_interval, extra_args=None, callback=None, disable=None, eta=1., s_noise=1., noise_sampler=None, solver_type='midpoint')
| 653 | |
| 654 | @torch.no_grad() |
| 655 | def sample_dpmpp_2m_sde_cfg(model, x, sigmas, cfg_val, cfg_interval, extra_args=None, callback=None, disable=None, eta=1., s_noise=1., noise_sampler=None, solver_type='midpoint'): |
| 656 | """DPM-Solver++(2M) SDE.""" |
| 657 | |
| 658 | if solver_type not in {'heun', 'midpoint'}: |
| 659 | raise ValueError('solver_type must be \'heun\' or \'midpoint\'') |
| 660 | |
| 661 | sigma_min, sigma_max = sigmas[sigmas > 0].min(), sigmas.max() |
| 662 | noise_sampler = BrownianTreeNoiseSampler(x, sigma_min, sigma_max) if noise_sampler is None else noise_sampler |
| 663 | extra_args = {} if extra_args is None else extra_args |
| 664 | s_in = x.new_ones([x.shape[0]]) |
| 665 | |
| 666 | # print("sigma min ", sigma_min) |
| 667 | # print("sigma max ", sigma_max) |
| 668 | # exit(1) |
| 669 | |
| 670 | old_denoised = None |
| 671 | h_last = None |
| 672 | |
| 673 | for i in trange(len(sigmas) - 1, disable=disable): |
| 674 | |
| 675 | #switch the cfg on or off depending on where in the schedule we are |
| 676 | #from https://arxiv.org/pdf/2404.07724 |
| 677 | cfg_interval_min=cfg_interval[0] |
| 678 | cfg_interval_max=cfg_interval[1] |
| 679 | if sigmas[i]>cfg_interval_min and sigmas[i]<cfg_interval_max: |
| 680 | cfg_val_cur=cfg_val |
| 681 | else: |
| 682 | cfg_val_cur=1.0 |
| 683 | # print("sigmas[i]",sigmas[i]) |
| 684 | # print("cfg_interval", cfg_interval) |
| 685 | print("cfg_val_cur",cfg_val_cur, " for sigma ", sigmas[i]) |
| 686 | |
| 687 | |
| 688 | # denoised = model(x, sigmas[i] * s_in, **extra_args) |
| 689 | if cfg_val_cur==1.0: |
| 690 | #no need to run the unconditional model, just run the conditional one |
| 691 | denoised=model(x, sigmas[i] * s_in, **extra_args) |
| 692 | else: |
| 693 | denoised_uncond = model(x, sigmas[i] * s_in) |
| 694 | denoised_cond = model(x, sigmas[i] * s_in, **extra_args) |
| 695 | denoised = denoised_uncond + cfg_val_cur*(denoised_cond - denoised_uncond) |
| 696 | if callback is not None: |
| 697 | callback({'x': x, 'i': i, 'sigma': sigmas[i], 'sigma_hat': sigmas[i], 'denoised': denoised}) |
| 698 | if sigmas[i + 1] == 0: |
| 699 | # Denoising step |
| 700 | x = denoised |
| 701 | else: |
| 702 | # DPM-Solver++(2M) SDE |
| 703 | t, s = -sigmas[i].log(), -sigmas[i + 1].log() |
| 704 | h = s - t |
| 705 | eta_h = eta * h |
| 706 | |
| 707 | x = sigmas[i + 1] / sigmas[i] * (-eta_h).exp() * x + (-h - eta_h).expm1().neg() * denoised |
| 708 | |
| 709 | if old_denoised is not None: |
| 710 | r = h_last / h |
| 711 | if solver_type == 'heun': |
| 712 | x = x + ((-h - eta_h).expm1().neg() / (-h - eta_h) + 1) * (1 / r) * (denoised - old_denoised) |
nothing calls this directly
no test coverage detected