class_docstr: ## DDIM Sampler
  
  This extends the [`DiffusionSampler` base class](index.html).
  
  DDIM samples images by repeatedly removing noise by sampling step by step using,
  
  egin{align}
  x_{ au_{i-1}} &= \sqrt{lpha_{      au_{i-1}}}\Bigg(
          rac{x_{    au_i} - \sqrt{1 - lpha_{       au_i}}\epsilon_ heta(x_{        au_i})}{\sqrt{lpha_{   au_i}}}
          \Bigg) \
          &+ \sqrt{1 - lpha_{        au_{i- 1}} - \sigma_{   au_i}^2} \cdot \epsilon_        heta(x_{        au_i}) \
          &+ \sigma_{ au_i} \epsilon_{        au_i}
  \end{align}
  
  where $\epsilon_{   au_i}$ is random noise,
  $   au$ is a subsequence of $[1,2,\dots,T]$ of length $S$,
  and
  $$\sigma_{  au_i} =
  \eta \sqrt{rac{1 - lpha_{ au_{i-1}}}{1 - lpha_{  au_i}}}
  \sqrt{1 - rac{lpha_{      au_i}}{lpha_{  au_{i-1}}}}$$
  
  Note that, $lpha_t$ in DDIM paper refers to ${\color{lightgreen}arlpha_t}$ from [DDPM](ddpm.html).
docstr: :param model: is the model to predict noise $\epsilon_  ext{cond}(x_t, c)$
  :param n_steps: is the number of DDIM sampling steps, $S$
  :param ddim_discretize: specifies how to extract $      au$ from $[1,2,\dots,T]$.
      It can be either `uniform` or `quad`.
  :param ddim_eta: is $\eta$ used to calculate $\sigma_{  au_i}$. $\eta = 0$ makes the
      sampling process deterministic.
docstr_tok: [":","param","model",":","is","the","model","to","predict","noise","$","\\epsilon_","ext","{","cond","}","(","x_t",",","c",")","$",":","param","n_steps",":","is","the","number","of","DDIM","sampling","steps",",","$","S","$",":","param","ddim_discretize",":","specifies","how","to","extract","$","au","$","from","$","[","1,2",",","\\dots",",","T","]","$",".","It","can","be","either","`","uniform","`","or","`","quad","`",".",":","param","ddim_eta",":","is","$","\\eta","$","used","to","calculate","$","\\sigma_","{","au_i","}","$",".","$","\\eta","=","0","$","makes","the","sampling","process","deterministic","."]
code: def __init__(
      self,
      model: LatentDiffusion,
      n_steps: int,
      ddim_discretize: str = "uniform",
      ddim_eta: float = 0.0,
  ):
      """
      :param model: is the model to predict noise $\epsilon_\text{cond}(x_t, c)$
      :param n_steps: is the number of DDIM sampling steps, $S$
      :param ddim_discretize: specifies how to extract $\tau$ from $[1,2,\dots,T]$.
          It can be either `uniform` or `quad`.
      :param ddim_eta: is $\eta$ used to calculate $\sigma_{\tau_i}$. $\eta = 0$ makes the
          sampling process deterministic.
      """
      super().__init__(model)
      # Number of steps, $T$
      self.n_steps = model.n_steps
  
      # Calculate $\tau$ to be uniformly distributed across $[1,2,\dots,T]$
      if ddim_discretize == "uniform":
          c = self.n_steps // n_steps
          self.time_steps = (
              np.asarray(list(range(0, self.n_steps, c))) + 1
          )
      # Calculate $\tau$ to be quadratically distributed across $[1,2,\dots,T]$
      elif ddim_discretize == "quad":
          self.time_steps = (
              (np.linspace(0, np.sqrt(self.n_steps * 0.8), n_steps))
              ** 2
          ).astype(int) + 1
      else:
          raise NotImplementedError(ddim_discretize)
  
      with torch.no_grad():
          # Get ${\color{lightgreen}\bar\alpha_t}$
          alpha_bar = self.model.alpha_bar
  
          # $\alpha_{\tau_i}$
          self.ddim_alpha = (
              alpha_bar[self.time_steps].clone().to(torch.float32)
          )
          # $\sqrt{\alpha_{\tau_i}}$
          self.ddim_alpha_sqrt = torch.sqrt(self.ddim_alpha)
          # $\alpha_{\tau_{i-1}}$
          self.ddim_alpha_prev = torch.cat(
              [alpha_bar[0:1], alpha_bar[self.time_steps[:-1]]]
          )
  
          # $$\sigma_{\tau_i} =
          # \eta \sqrt{\frac{1 - \alpha_{\tau_{i-1}}}{1 - \alpha_{\tau_i}}}
          # \sqrt{1 - \frac{\alpha_{\tau_i}}{\alpha_{\tau_{i-1}}}}$$
          self.ddim_sigma = (
              ddim_eta
              * (
                  (1 - self.ddim_alpha_prev)
                  / (1 - self.ddim_alpha)
                  * (1 - self.ddim_alpha / self.ddim_alpha_prev)
              )
              ** 0.5
          )
  
          # $\sqrt{1 - \alpha_{\tau_i}}$
          self.ddim_sqrt_one_minus_alpha = (
              1.0 - self.ddim_alpha
          ) ** 0.5
  
file_path: annotated_deep_learning_paper_implementations/labml_nn/diffusion/stable_diffusion/sampler/ddim.py
imports: from typing import Optional, List
  import numpy
  import torch
  from labml import monit
  from labml_nn.diffusion.stable_diffusion.latent_diffusion import LatentDiffusion
  from labml_nn.diffusion.stable_diffusion.sampler import DiffusionSampler
label_desc: Use image node_id 1 to create a new DDIMSampler object from inherited base classes: DiffusionSampler with example: obj = DDIMSampler(model, n_steps, ddim_discretize, ddim_eta)