(self, x, t_start, t_end, nfe, eta=0., s_noise=1., noise_sampler=None)
| 386 | return x_3, eps_cache |
| 387 | |
| 388 | def dpm_solver_fast(self, x, t_start, t_end, nfe, eta=0., s_noise=1., noise_sampler=None): |
| 389 | noise_sampler = default_noise_sampler(x) if noise_sampler is None else noise_sampler |
| 390 | if not t_end > t_start and eta: |
| 391 | raise ValueError('eta must be 0 for reverse sampling') |
| 392 | |
| 393 | m = math.floor(nfe / 3) + 1 |
| 394 | ts = torch.linspace(t_start, t_end, m + 1, device=x.device) |
| 395 | |
| 396 | if nfe % 3 == 0: |
| 397 | orders = [3] * (m - 2) + [2, 1] |
| 398 | else: |
| 399 | orders = [3] * (m - 1) + [nfe % 3] |
| 400 | |
| 401 | for i in range(len(orders)): |
| 402 | eps_cache = {} |
| 403 | t, t_next = ts[i], ts[i + 1] |
| 404 | if eta: |
| 405 | sd, su = get_ancestral_step(self.sigma(t), self.sigma(t_next), eta) |
| 406 | t_next_ = torch.minimum(t_end, self.t(sd)) |
| 407 | su = (self.sigma(t_next) ** 2 - self.sigma(t_next_) ** 2) ** 0.5 |
| 408 | else: |
| 409 | t_next_, su = t_next, 0. |
| 410 | |
| 411 | eps, eps_cache = self.eps(eps_cache, 'eps', x, t) |
| 412 | denoised = x - self.sigma(t) * eps |
| 413 | if self.info_callback is not None: |
| 414 | self.info_callback({'x': x, 'i': i, 't': ts[i], 't_up': t, 'denoised': denoised}) |
| 415 | |
| 416 | if orders[i] == 1: |
| 417 | x, eps_cache = self.dpm_solver_1_step(x, t, t_next_, eps_cache=eps_cache) |
| 418 | elif orders[i] == 2: |
| 419 | x, eps_cache = self.dpm_solver_2_step(x, t, t_next_, eps_cache=eps_cache) |
| 420 | else: |
| 421 | x, eps_cache = self.dpm_solver_3_step(x, t, t_next_, eps_cache=eps_cache) |
| 422 | |
| 423 | x = x + su * s_noise * noise_sampler(self.sigma(t), self.sigma(t_next)) |
| 424 | |
| 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 |
no test coverage detected