| 284 | v = torch.randint_like(x, 2) * 2 - 1 |
| 285 | fevals = 0 |
| 286 | def ode_fn(sigma, x): |
| 287 | nonlocal fevals |
| 288 | with torch.enable_grad(): |
| 289 | x = x[0].detach().requires_grad_() |
| 290 | denoised = model(x, sigma * s_in, **extra_args) |
| 291 | d = to_d(x, sigma, denoised) |
| 292 | fevals += 1 |
| 293 | grad = torch.autograd.grad((d * v).sum(), x)[0] |
| 294 | d_ll = (v * grad).flatten(1).sum(1) |
| 295 | return d.detach(), d_ll |
| 296 | x_min = x, x.new_zeros([x.shape[0]]) |
| 297 | t = x.new_tensor([sigma_min, sigma_max]) |
| 298 | sol = odeint(ode_fn, x_min, t, atol=atol, rtol=rtol, method='dopri5') |