text: import torch
  from tqdm import trange
  
  from . import utils
  
  
  @torch.no_grad()
  def sample(model, x, steps, eta, extra_args):
      """Draws samples from a model given starting noise."""
      ts = x.new_ones([x.shape[0]])
  
      # Create the noise schedule
      alphas, sigmas = utils.t_to_alpha_sigma(steps)
  
      # The sampling loop
      for i in trange(len(steps)):
  
          # Get the model output (v, the predicted velocity)
          with torch.cuda.amp.autocast():
              v = model(x, ts * steps[i], **extra_args).float()
  
          # Predict the noise and the denoised image
          pred = x * alphas[i] - v * sigmas[i]
          eps = x * sigmas[i] + v * alphas[i]
  
          # If we are not on the last timestep, compute the noisy image for the
          # next timestep.
          if i < len(steps) - 1:
              # If eta > 0, adjust the scaling factor for the predicted noise
              # downward according to the amount of additional noise to add
              ddim_sigma = eta * (sigmas[i + 1]**2 / sigmas[i]**2).sqrt() * \
                  (1 - alphas[i]**2 / alphas[i + 1]**2).sqrt()
              adjusted_sigma = (sigmas[i + 1]**2 - ddim_sigma**2).sqrt()
  
              # Recombine the predicted noise and predicted denoised image in the
              # correct proportions for the next step
              x = pred * alphas[i + 1] + eps * adjusted_sigma
  
              # Add the correct amount of fresh noise
              if eta:
                  x += torch.randn_like(x) * ddim_sigma
  
      # If we are on the last timestep, output the denoised image
      return pred
  
  
  @torch.no_grad()
  def cond_sample(model, x, steps, eta, extra_args, cond_fn):
      """Draws guided samples from a model given starting noise."""
      ts = x.new_ones([x.shape[0]])
  
      # Create the noise schedule
      alphas, sigmas = utils.t_to_alpha_sigma(steps)
  
      # The sampling loop
      for i in trange(len(steps)):
  
          # Get the model output
          with torch.enable_grad():
              x = x.detach().requires_grad_()
              with torch.cuda.amp.autocast():
                  v = model(x, ts * steps[i], **extra_args)
  
              if steps[i] < 1:
                  pred = x * alphas[i] - v * sigmas[i]
                  cond_grad = cond_fn(x, ts * steps[i], pred, **extra_args).detach()
                  v = v.detach() - cond_grad * (sigmas[i] / alphas[i])
              else:
                  v = v.detach()
  
          # Predict the noise and the denoised image
          pred = x * alphas[i] - v * sigmas[i]
          eps = x * sigmas[i] + v * alphas[i]
  
          # If we are not on the last timestep, compute the noisy image for the
          # next timestep.
          if i < len(steps) - 1:
              # If eta > 0, adjust the scaling factor for the predicted noise
              # downward according to the amount of additional noise to add
              ddim_sigma = eta * (sigmas[i + 1]**2 / sigmas[i]**2).sqrt() * \
                  (1 - alphas[i]**2 / alphas[i + 1]**2).sqrt()
              adjusted_sigma = (sigmas[i + 1]**2 - ddim_sigma**2).sqrt()
  
              # Recombine the predicted noise and predicted denoised image in the
              # correct proportions for the next step
              x = pred * alphas[i + 1] + eps * adjusted_sigma
  
              # Add the correct amount of fresh noise
              if eta:
                  x += torch.randn_like(x) * ddim_sigma
  
      # If we are on the last timestep, output the denoised image
      return pred
  