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: ### Sample $x_{ au_{i-1}}$ given $\epsilon_     heta(x_{        au_i})$
docstr_tok: ["#","#","#","Sample","$","x_","{","au_","{","i-1","}","}","$","given","$","\\epsilon_","heta","(","x_","{","au_i","}",")","$"]
code: def get_x_prev_and_pred_x0(
      self,
      e_t: torch.Tensor,
      index: int,
      x: torch.Tensor,
      *,
      temperature: float,
      repeat_noise: bool,
  ):
      """
      ### Sample $x_{\tau_{i-1}}$ given $\epsilon_\theta(x_{\tau_i})$
      """
  
      # $\alpha_{\tau_i}$
      alpha = self.ddim_alpha[index]
      # $\alpha_{\tau_{i-1}}$
      alpha_prev = self.ddim_alpha_prev[index]
      # $\sigma_{\tau_i}$
      sigma = self.ddim_sigma[index]
      # $\sqrt{1 - \alpha_{\tau_i}}$
      sqrt_one_minus_alpha = self.ddim_sqrt_one_minus_alpha[index]
  
      # Current prediction for $x_0$,
      # $$\frac{x_{\tau_i} - \sqrt{1 - \alpha_{\tau_i}}\epsilon_\theta(x_{\tau_i})}{\sqrt{\alpha_{\tau_i}}}$$
      pred_x0 = (x - sqrt_one_minus_alpha * e_t) / (alpha**0.5)
      # Direction pointing to $x_t$
      # $$\sqrt{1 - \alpha_{\tau_{i- 1}} - \sigma_{\tau_i}^2} \cdot \epsilon_\theta(x_{\tau_i})$$
      dir_xt = (1.0 - alpha_prev - sigma**2).sqrt() * e_t
  
      # No noise is added, when $\eta = 0$
      if sigma == 0.0:
          noise = 0.0
      # If same noise is used for all samples in the batch
      elif repeat_noise:
          noise = torch.randn((1, *x.shape[1:]), device=x.device)
          # Different noise for each sample
      else:
          noise = torch.randn(x.shape, device=x.device)
  
      # Multiply noise by the temperature
      noise = noise * temperature
  
      #  \begin{align}
      #     x_{\tau_{i-1}} &= \sqrt{\alpha_{\tau_{i-1}}}\Bigg(
      #             \frac{x_{\tau_i} - \sqrt{1 - \alpha_{\tau_i}}\epsilon_\theta(x_{\tau_i})}{\sqrt{\alpha_{\tau_i}}}
      #             \Bigg) \\
      #             &+ \sqrt{1 - \alpha_{\tau_{i- 1}} - \sigma_{\tau_i}^2} \cdot \epsilon_\theta(x_{\tau_i}) \\
      #             &+ \sigma_{\tau_i} \epsilon_{\tau_i}
      #  \end{align}
      x_prev = (alpha_prev**0.5) * pred_x0 + dir_xt + sigma * noise
  
      #
      return x_prev, pred_x0
  
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 4 for calling the DDIMSampler obj's underlying member method code with example usage: obj.get_x_prev_and_pred_x0(e_t, index, x) and returns: x_prev,  pred_x0