(model, x, sigmas, extra_args=None, callback=None, disable=None, order=4)
| 259 | |
| 260 | @torch.no_grad() |
| 261 | def sample_lms(model, x, sigmas, extra_args=None, callback=None, disable=None, order=4): |
| 262 | extra_args = {} if extra_args is None else extra_args |
| 263 | s_in = x.new_ones([x.shape[0]]) |
| 264 | sigmas_cpu = sigmas.detach().cpu().numpy() |
| 265 | ds = [] |
| 266 | for i in trange(len(sigmas) - 1, disable=disable): |
| 267 | denoised = model(x, sigmas[i] * s_in, **extra_args) |
| 268 | d = to_d(x, sigmas[i], denoised) |
| 269 | ds.append(d) |
| 270 | if len(ds) > order: |
| 271 | ds.pop(0) |
| 272 | if callback is not None: |
| 273 | callback({'x': x, 'i': i, 'sigma': sigmas[i], 'sigma_hat': sigmas[i], 'denoised': denoised}) |
| 274 | cur_order = min(i + 1, order) |
| 275 | coeffs = [linear_multistep_coeff(cur_order, sigmas_cpu, i, j) for j in range(cur_order)] |
| 276 | x = x + sum(coeff * d for coeff, d in zip(coeffs, reversed(ds))) |
| 277 | return x |
| 278 | |
| 279 | |
| 280 | @torch.no_grad() |
nothing calls this directly
no test coverage detected