(self, x, t_start, t_end, order=3, rtol=0.05, atol=0.0078, h_init=0.05, pcoeff=0., icoeff=1., dcoeff=0., accept_safety=0.81, eta=0., s_noise=1., noise_sampler=None)
| 425 | return x |
| 426 | |
| 427 | def dpm_solver_adaptive(self, x, t_start, t_end, order=3, rtol=0.05, atol=0.0078, h_init=0.05, pcoeff=0., icoeff=1., dcoeff=0., accept_safety=0.81, eta=0., s_noise=1., noise_sampler=None): |
| 428 | noise_sampler = default_noise_sampler(x) if noise_sampler is None else noise_sampler |
| 429 | if order not in {2, 3}: |
| 430 | raise ValueError('order should be 2 or 3') |
| 431 | forward = t_end > t_start |
| 432 | if not forward and eta: |
| 433 | raise ValueError('eta must be 0 for reverse sampling') |
| 434 | h_init = abs(h_init) * (1 if forward else -1) |
| 435 | atol = torch.tensor(atol) |
| 436 | rtol = torch.tensor(rtol) |
| 437 | s = t_start |
| 438 | x_prev = x |
| 439 | accept = True |
| 440 | pid = PIDStepSizeController(h_init, pcoeff, icoeff, dcoeff, 1.5 if eta else order, accept_safety) |
| 441 | info = {'steps': 0, 'nfe': 0, 'n_accept': 0, 'n_reject': 0} |
| 442 | |
| 443 | while s < t_end - 1e-5 if forward else s > t_end + 1e-5: |
| 444 | eps_cache = {} |
| 445 | t = torch.minimum(t_end, s + pid.h) if forward else torch.maximum(t_end, s + pid.h) |
| 446 | if eta: |
| 447 | sd, su = get_ancestral_step(self.sigma(s), self.sigma(t), eta) |
| 448 | t_ = torch.minimum(t_end, self.t(sd)) |
| 449 | su = (self.sigma(t) ** 2 - self.sigma(t_) ** 2) ** 0.5 |
| 450 | else: |
| 451 | t_, su = t, 0. |
| 452 | |
| 453 | eps, eps_cache = self.eps(eps_cache, 'eps', x, s) |
| 454 | denoised = x - self.sigma(s) * eps |
| 455 | |
| 456 | if order == 2: |
| 457 | x_low, eps_cache = self.dpm_solver_1_step(x, s, t_, eps_cache=eps_cache) |
| 458 | x_high, eps_cache = self.dpm_solver_2_step(x, s, t_, eps_cache=eps_cache) |
| 459 | else: |
| 460 | x_low, eps_cache = self.dpm_solver_2_step(x, s, t_, r1=1 / 3, eps_cache=eps_cache) |
| 461 | x_high, eps_cache = self.dpm_solver_3_step(x, s, t_, eps_cache=eps_cache) |
| 462 | delta = torch.maximum(atol, rtol * torch.maximum(x_low.abs(), x_prev.abs())) |
| 463 | error = torch.linalg.norm((x_low - x_high) / delta) / x.numel() ** 0.5 |
| 464 | accept = pid.propose_step(error) |
| 465 | if accept: |
| 466 | x_prev = x_low |
| 467 | x = x_high + su * s_noise * noise_sampler(self.sigma(s), self.sigma(t)) |
| 468 | s = t |
| 469 | info['n_accept'] += 1 |
| 470 | else: |
| 471 | info['n_reject'] += 1 |
| 472 | info['nfe'] += order |
| 473 | info['steps'] += 1 |
| 474 | |
| 475 | if self.info_callback is not None: |
| 476 | self.info_callback({'x': x, 'i': info['steps'] - 1, 't': s, 't_up': s, 'denoised': denoised, 'error': error, 'h': pid.h, **info}) |
| 477 | |
| 478 | return x, info |
| 479 | |
| 480 | |
| 481 | @torch.no_grad() |
no test coverage detected