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}}$
  
  :param x: is $x_{       au_i}$ of shape `[batch_size, channels, height, width]`
  :param c: is the conditional embeddings $c$ of shape `[batch_size, emb_size]`
  :param t: is $  au_i$ of shape `[batch_size]`
  :param step: is the step $      au_i$ as an integer
  :param index: is index $i$ in the list $[       au_1,   au_2, \dots,    au_S]$
  :param repeat_noise: specified whether the noise should be same for all samples in the batch
  :param temperature: is the noise temperature (random noise gets multiplied by this)
  :param uncond_scale: is the unconditional guidance scale $s$. This is used for
      $\epsilon_  heta(x_t, c) = s\epsilon_       ext{cond}(x_t, c) + (s - 1)\epsilon_    ext{cond}(x_t, c_u)$
  :param uncond_cond: is the conditional embedding for empty prompt $c_u$
docstr_tok: ["#","#","#","Sample","$","x_","{","au_","{","i-1","}","}","$",":","param","x",":","is","$","x_","{","au_i","}","$","of","shape","`","[","batch_size",",","channels",",","height",",","width","]","`",":","param","c",":","is","the","conditional","embeddings","$","c","$","of","shape","`","[","batch_size",",","emb_size","]","`",":","param","t",":","is","$","au_i","$","of","shape","`","[","batch_size","]","`",":","param","step",":","is","the","step","$","au_i","$","as","an","integer",":","param","index",":","is","index","$","i","$","in","the","list","$","[","au_1",",","au_2",",","\\dots",",","au_S","]","$",":","param","repeat_noise",":","specified","whether","the","noise","should","be","same","for","all","samples","in","the","batch",":","param","temperature",":","is","the","noise","temperature","(","random","noise","gets","multiplied","by","this",")",":","param","uncond_scale",":","is","the","unconditional","guidance","scale","$","s","$",".","This","is","used","for","$","\\epsilon_","heta","(","x_t",",","c",")","=","s\\epsilon_","ext","{","cond","}","(","x_t",",","c",")","+","(","s","-","1",")","\\epsilon_","ext","{","cond","}","(","x_t",",","c_u",")","$",":","param","uncond_cond",":","is","the","conditional","embedding","for","empty","prompt","$","c_u","$"]
code: def p_sample(
      self,
      x: torch.Tensor,
      c: torch.Tensor,
      t: torch.Tensor,
      step: int,
      index: int,
      *,
      repeat_noise: bool = False,
      temperature: float = 1.0,
      uncond_scale: float = 1.0,
      uncond_cond: Optional[torch.Tensor] = None,
  ):
      """
      ### Sample $x_{\tau_{i-1}}$
  
      :param x: is $x_{\tau_i}$ of shape `[batch_size, channels, height, width]`
      :param c: is the conditional embeddings $c$ of shape `[batch_size, emb_size]`
      :param t: is $\tau_i$ of shape `[batch_size]`
      :param step: is the step $\tau_i$ as an integer
      :param index: is index $i$ in the list $[\tau_1, \tau_2, \dots, \tau_S]$
      :param repeat_noise: specified whether the noise should be same for all samples in the batch
      :param temperature: is the noise temperature (random noise gets multiplied by this)
      :param uncond_scale: is the unconditional guidance scale $s$. This is used for
          $\epsilon_\theta(x_t, c) = s\epsilon_\text{cond}(x_t, c) + (s - 1)\epsilon_\text{cond}(x_t, c_u)$
      :param uncond_cond: is the conditional embedding for empty prompt $c_u$
      """
  
      # Get $\epsilon_\theta(x_{\tau_i})$
      e_t = self.get_eps(
          x, t, c, uncond_scale=uncond_scale, uncond_cond=uncond_cond
      )
  
      # Calculate $x_{\tau_{i - 1}}$ and predicted $x_0$
      x_prev, pred_x0 = self.get_x_prev_and_pred_x0(
          e_t,
          index,
          x,
          temperature=temperature,
          repeat_noise=repeat_noise,
      )
  
      #
      return x_prev, pred_x0, e_t
  
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 3 for calling the DDIMSampler obj's underlying member method code with example usage: obj.p_sample(x, c, t, step, index) and returns: x_prev,  pred_x0,  e_t