DPM-Solver. See https://arxiv.org/abs/2206.00927.
| 331 | |
| 332 | |
| 333 | class DPMSolver(nn.Module): |
| 334 | """DPM-Solver. See https://arxiv.org/abs/2206.00927.""" |
| 335 | |
| 336 | def __init__(self, model, extra_args=None, eps_callback=None, info_callback=None): |
| 337 | super().__init__() |
| 338 | self.model = model |
| 339 | self.extra_args = {} if extra_args is None else extra_args |
| 340 | self.eps_callback = eps_callback |
| 341 | self.info_callback = info_callback |
| 342 | |
| 343 | def t(self, sigma): |
| 344 | return -sigma.log() |
| 345 | |
| 346 | def sigma(self, t): |
| 347 | return t.neg().exp() |
| 348 | |
| 349 | def eps(self, eps_cache, key, x, t, *args, **kwargs): |
| 350 | if key in eps_cache: |
| 351 | return eps_cache[key], eps_cache |
| 352 | sigma = self.sigma(t) * x.new_ones([x.shape[0]]) |
| 353 | eps = (x - self.model(x, sigma, *args, **self.extra_args, **kwargs)) / self.sigma(t) |
| 354 | if self.eps_callback is not None: |
| 355 | self.eps_callback() |
| 356 | return eps, {key: eps, **eps_cache} |
| 357 | |
| 358 | def dpm_solver_1_step(self, x, t, t_next, eps_cache=None): |
| 359 | eps_cache = {} if eps_cache is None else eps_cache |
| 360 | h = t_next - t |
| 361 | eps, eps_cache = self.eps(eps_cache, 'eps', x, t) |
| 362 | x_1 = x - self.sigma(t_next) * h.expm1() * eps |
| 363 | return x_1, eps_cache |
| 364 | |
| 365 | def dpm_solver_2_step(self, x, t, t_next, r1=1 / 2, eps_cache=None): |
| 366 | eps_cache = {} if eps_cache is None else eps_cache |
| 367 | h = t_next - t |
| 368 | eps, eps_cache = self.eps(eps_cache, 'eps', x, t) |
| 369 | s1 = t + r1 * h |
| 370 | u1 = x - self.sigma(s1) * (r1 * h).expm1() * eps |
| 371 | eps_r1, eps_cache = self.eps(eps_cache, 'eps_r1', u1, s1) |
| 372 | x_2 = x - self.sigma(t_next) * h.expm1() * eps - self.sigma(t_next) / (2 * r1) * h.expm1() * (eps_r1 - eps) |
| 373 | return x_2, eps_cache |
| 374 | |
| 375 | def dpm_solver_3_step(self, x, t, t_next, r1=1 / 3, r2=2 / 3, eps_cache=None): |
| 376 | eps_cache = {} if eps_cache is None else eps_cache |
| 377 | h = t_next - t |
| 378 | eps, eps_cache = self.eps(eps_cache, 'eps', x, t) |
| 379 | s1 = t + r1 * h |
| 380 | s2 = t + r2 * h |
| 381 | u1 = x - self.sigma(s1) * (r1 * h).expm1() * eps |
| 382 | eps_r1, eps_cache = self.eps(eps_cache, 'eps_r1', u1, s1) |
| 383 | u2 = x - self.sigma(s2) * (r2 * h).expm1() * eps - self.sigma(s2) * (r2 / r1) * ((r2 * h).expm1() / (r2 * h) - 1) * (eps_r1 - eps) |
| 384 | eps_r2, eps_cache = self.eps(eps_cache, 'eps_r2', u2, s2) |
| 385 | x_3 = x - self.sigma(t_next) * h.expm1() * eps - self.sigma(t_next) / r2 * (h.expm1() / h - 1) * (eps_r2 - eps) |
| 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: |
no outgoing calls
no test coverage detected