code: def restart_sampler(
      model,
      x,
      sigmas,
      extra_args=None,
      callback=None,
      disable=None,
      s_noise=1.0,
      restart_list=None,
  ):
      """Implements restart sampling in Restart Sampling for Improving Generative Processes (2023)
      Restart_list format: {min_sigma: [ restart_steps, restart_times, max_sigma]}
      If restart_list is None: will choose restart_list automatically, otherwise will use the given restart_list
      """
      extra_args = {} if extra_args is None else extra_args
      s_in = x.new_ones([x.shape[0]])
      step_id = 0
      from k_diffusion.sampling import to_d, get_sigmas_karras
  
      def heun_step(x, old_sigma, new_sigma, second_order=True):
          nonlocal step_id
          denoised = model(x, old_sigma * s_in, **extra_args)
          d = to_d(x, old_sigma, denoised)
          if callback is not None:
              callback(
                  {
                      "x": x,
                      "i": step_id,
                      "sigma": new_sigma,
                      "sigma_hat": old_sigma,
                      "denoised": denoised,
                  }
              )
          dt = new_sigma - old_sigma
          if new_sigma == 0 or not second_order:
              # Euler method
              x = x + d * dt
          else:
              # Heun's method
              x_2 = x + d * dt
              denoised_2 = model(x_2, new_sigma * s_in, **extra_args)
              d_2 = to_d(x_2, new_sigma, denoised_2)
              d_prime = (d + d_2) / 2
              x = x + d_prime * dt
          step_id += 1
          return x
  
      steps = sigmas.shape[0] - 1
      if restart_list is None:
          if steps >= 20:
              restart_steps = 9
              restart_times = 1
              if steps >= 36:
                  restart_steps = steps // 4
                  restart_times = 2
              sigmas = get_sigmas_karras(
                  steps - restart_steps * restart_times,
                  sigmas[-2].item(),
                  sigmas[0].item(),
                  device=sigmas.device,
              )
              restart_list = {
                  0.1: [restart_steps + 1, restart_times, 2]
              }
          else:
              restart_list = {}
  
      restart_list = {
          int(torch.argmin(abs(sigmas - key), dim=0)): value
          for key, value in restart_list.items()
      }
  
      step_list = []
      for i in range(len(sigmas) - 1):
          step_list.append((sigmas[i], sigmas[i + 1]))
          if i + 1 in restart_list:
              restart_steps, restart_times, restart_max = restart_list[
                  i + 1
              ]
              min_idx = i + 1
              max_idx = int(
                  torch.argmin(abs(sigmas - restart_max), dim=0)
              )
              if max_idx < min_idx:
                  sigma_restart = get_sigmas_karras(
                      restart_steps,
                      sigmas[min_idx].item(),
                      sigmas[max_idx].item(),
                      device=sigmas.device,
                  )[:-1]
                  while restart_times > 0:
                      restart_times -= 1
                      step_list.extend(
                          zip(sigma_restart[:-1], sigma_restart[1:])
                      )
  
      last_sigma = None
      for old_sigma, new_sigma in tqdm.tqdm(step_list, disable=disable):
          if last_sigma is None:
              last_sigma = old_sigma
          elif last_sigma < old_sigma:
              x = (
                  x
                  + k_diffusion.sampling.torch.randn_like(x)
                  * s_noise
                  * (old_sigma**2 - last_sigma**2) ** 0.5
              )
          x = heun_step(x, old_sigma, new_sigma)
          last_sigma = new_sigma
  
      return x
  
imports: import torch
  import tqdm
  import k_diffusion.sampling
label_desc: Use image node_id 1 for calling a global function with example usage: restart_sampler(model, x, sigmas, extra_args, callback, disable, s_noise, restart_list) and returns: x, x