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: ### Painting Loop
  
  :param x: is $x_{S'}$ of shape `[batch_size, channels, height, width]`
  :param cond: is the conditional embeddings $c$
  :param t_start: is the sampling step to start from, $S'$
  :param orig: is the original image in latent page which we are in paining.
      If this is not provided, it'll be an image to image transformation.
  :param mask: is the mask to keep the original image.
  :param orig_noise: is fixed noise to be added to the original image.
  :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: ["#","#","#","Painting","Loop",":","param","x",":","is","$","x_","{","S","'","}","$","of","shape","`","[","batch_size",",","channels",",","height",",","width","]","`",":","param","cond",":","is","the","conditional","embeddings","$","c","$",":","param","t_start",":","is","the","sampling","step","to","start","from",",","$","S","'","$",":","param","orig",":","is","the","original","image","in","latent","page","which","we","are","in","paining",".","If","this","is","not","provided",",","it","'ll","be","an","image","to","image","transformation",".",":","param","mask",":","is","the","mask","to","keep","the","original","image",".",":","param","orig_noise",":","is","fixed","noise","to","be","added","to","the","original","image",".",":","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 paint(
      self,
      x: torch.Tensor,
      cond: torch.Tensor,
      t_start: int,
      *,
      orig: Optional[torch.Tensor] = None,
      mask: Optional[torch.Tensor] = None,
      orig_noise: Optional[torch.Tensor] = None,
      uncond_scale: float = 1.0,
      uncond_cond: Optional[torch.Tensor] = None,
  ):
      """
      ### Painting Loop
  
      :param x: is $x_{S'}$ of shape `[batch_size, channels, height, width]`
      :param cond: is the conditional embeddings $c$
      :param t_start: is the sampling step to start from, $S'$
      :param orig: is the original image in latent page which we are in paining.
          If this is not provided, it'll be an image to image transformation.
      :param mask: is the mask to keep the original image.
      :param orig_noise: is fixed noise to be added to the original image.
      :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  batch size
      bs = x.shape[0]
  
      # Time steps to sample at $\tau_{S`}, \tau_{S' - 1}, \dots, \tau_1$
      time_steps = np.flip(self.time_steps[:t_start])
  
      for i, step in monit.enum("Paint", time_steps):
          # Index $i$ in the list $[\tau_1, \tau_2, \dots, \tau_S]$
          index = len(time_steps) - i - 1
          # Time step $\tau_i$
          ts = x.new_full((bs,), step, dtype=torch.long)
  
          # Sample $x_{\tau_{i-1}}$
          x, _, _ = self.p_sample(
              x,
              cond,
              ts,
              step,
              index=index,
              uncond_scale=uncond_scale,
              uncond_cond=uncond_cond,
          )
  
          # Replace the masked area with original image
          if orig is not None:
              # Get the $q_{\sigma,\tau}(x_{\tau_i}|x_0)$ for original image in latent space
              orig_t = self.q_sample(orig, index, noise=orig_noise)
              # Replace the masked area
              x = orig_t * mask + x * (1 - mask)
  
      #
      return x
  
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 6 for calling the DDIMSampler obj's underlying member method code with example usage: obj.paint(x, cond, t_start) and returns: x