code: def sample_lms(
      model,
      x,
      sigmas,
      extra_args=None,
      callback=None,
      disable=None,
      order=4,
  ):
      extra_args = {} if extra_args is None else extra_args
      s_in = x.new_ones([x.shape[0]])
      sigmas_cpu = sigmas.detach().cpu().numpy()
      ds = []
      for i in trange(len(sigmas) - 1, disable=disable):
          denoised = model(x, sigmas[i] * s_in, **extra_args)
          d = to_d(x, sigmas[i], denoised)
          ds.append(d)
          if len(ds) > order:
              ds.pop(0)
          if callback is not None:
              callback(
                  {
                      "x": x,
                      "i": i,
                      "sigma": sigmas[i],
                      "sigma_hat": sigmas[i],
                      "denoised": denoised,
                  }
              )
          cur_order = min(i + 1, order)
          coeffs = [
              linear_multistep_coeff(cur_order, sigmas_cpu, i, j)
              for j in range(cur_order)
          ]
          x = x + sum(
              coeff * d for coeff, d in zip(coeffs, reversed(ds))
          )
      return x
  
imports: import math
  import torch
  import torchsde
  from scipy import integrate
  from torch import nn
  from torchdiffeq import odeint
  from tqdm.auto import tqdm, trange
label_desc: Use image node_id 16 for calling a global function with example usage: sample_lms(model, x, sigmas, extra_args, callback, disable, order) and returns: x