class_docstr: ## DDPM Sampler
  
  This extends the [`DiffusionSampler` base class](index.html).
  
  DDPM samples images by repeatedly removing noise by sampling step by step from
  $p_ heta(x_{t-1} | x_t)$,
  
  egin{align}
  
  p_  heta(x_{t-1} | x_t) &= \mathcal{N}ig(x_{t-1}; \mu_     heta(x_t, t),   ildeeta_t \mathbf{I} ig) \
  
  \mu_t(x_t, t) &= rac{\sqrt{arlpha_{t-1}}eta_t}{1 - arlpha_t}x_0
                       + rac{\sqrt{lpha_t}(1 - arlpha_{t-1})}{1-arlpha_t}x_t \
  
      ildeeta_t &= rac{1 - arlpha_{t-1}}{1 - arlpha_t} eta_t \
  
  x_0 &= rac{1}{\sqrt{arlpha_t}} x_t -  \Big(\sqrt{rac{1}{arlpha_t} - 1}\Big)\epsilon_  heta \
  
  \end{align}
class_docstr_tok: ["#","#","DDPM","Sampler","This","extends","the","[","`","DiffusionSampler","`","base","class","]","(","index.html",")",".","DDPM","samples","images","by","repeatedly","removing","noise","by","sampling","step","by","step","from","$","p_","heta","(","x_","{","t-1","}","|","x_t",")","$",",","\begin","{","align","}","p_","heta","(","x_","{","t-1","}","|","x_t",")","&","=","\\mathcal","{","N","}","\big","(","x_","{","t-1","}",";","\\mu_","heta","(","x_t",",","t",")",",","ilde\beta_t","\\mathbf","{","I","}","\big",")","\\","\\mu_t","(","x_t",",","t",")","&","=","rac","{","\\sqrt","{","\bar\u0007lpha_","{","t-1","}","}","\beta_t","}","{","1","-","\bar\u0007lpha_t","}","x_0","+","rac","{","\\sqrt","{","\u0007lpha_t","}","(","1","-","\bar\u0007lpha_","{","t-1","}",")","}","{","1-\bar\u0007lpha_t","}","x_t","\\","ilde\beta_t","&","=","rac","{","1","-","\bar\u0007lpha_","{","t-1","}","}","{","1","-","\bar\u0007lpha_t","}","\beta_t","\\","x_0","&","=","rac","{","1","}","{","\\sqrt","{","\bar\u0007lpha_t","}","}","x_t","-","\\Big","(","\\sqrt","{","rac","{","1","}","{","\bar\u0007lpha_t","}","-","1","}","\\Big",")","\\epsilon_","heta","\\","\\end","{","align","}"]
docstr: ### Sample $x_{t-1}$ from $p_   heta(x_{t-1} | x_t)$
  
  :param x: is $x_t$ of shape `[batch_size, channels, height, width]`
  :param c: is the conditional embeddings $c$ of shape `[batch_size, emb_size]`
  :param t: is $t$ of shape `[batch_size]`
  :param step: is the step $t$ as an integer
  :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_","{","t-1","}","$","from","$","p_","heta","(","x_","{","t-1","}","|","x_t",")","$",":","param","x",":","is","$","x_t","$","of","shape","`","[","batch_size",",","channels",",","height",",","width","]","`",":","param","c",":","is","the","conditional","embeddings","$","c","$","of","shape","`","[","batch_size",",","emb_size","]","`",":","param","t",":","is","$","t","$","of","shape","`","[","batch_size","]","`",":","param","step",":","is","the","step","$","t","$","as","an","integer",":","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,
      repeat_noise: bool = False,
      temperature: float = 1.0,
      uncond_scale: float = 1.0,
      uncond_cond: Optional[torch.Tensor] = None,
  ):
      """
      ### Sample $x_{t-1}$ from $p_\theta(x_{t-1} | x_t)$
  
      :param x: is $x_t$ of shape `[batch_size, channels, height, width]`
      :param c: is the conditional embeddings $c$ of shape `[batch_size, emb_size]`
      :param t: is $t$ of shape `[batch_size]`
      :param step: is the step $t$ as an integer
      :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$
      e_t = self.get_eps(
          x, t, c, uncond_scale=uncond_scale, uncond_cond=uncond_cond
      )
  
      # Get batch size
      bs = x.shape[0]
  
      # $\frac{1}{\sqrt{\bar\alpha_t}}$
      sqrt_recip_alpha_bar = x.new_full(
          (bs, 1, 1, 1), self.sqrt_recip_alpha_bar[step]
      )
      # $\sqrt{\frac{1}{\bar\alpha_t} - 1}$
      sqrt_recip_m1_alpha_bar = x.new_full(
          (bs, 1, 1, 1), self.sqrt_recip_m1_alpha_bar[step]
      )
  
      # Calculate $x_0$ with current $\epsilon_\theta$
      #
      # $$x_0 = \frac{1}{\sqrt{\bar\alpha_t}} x_t -  \Big(\sqrt{\frac{1}{\bar\alpha_t} - 1}\Big)\epsilon_\theta$$
      x0 = sqrt_recip_alpha_bar * x - sqrt_recip_m1_alpha_bar * e_t
  
      # $\frac{\sqrt{\bar\alpha_{t-1}}\beta_t}{1 - \bar\alpha_t}$
      mean_x0_coef = x.new_full((bs, 1, 1, 1), self.mean_x0_coef[step])
      # $\frac{\sqrt{\alpha_t}(1 - \bar\alpha_{t-1})}{1-\bar\alpha_t}$
      mean_xt_coef = x.new_full((bs, 1, 1, 1), self.mean_xt_coef[step])
  
      # Calculate $\mu_t(x_t, t)$
      #
      # $$\mu_t(x_t, t) = \frac{\sqrt{\bar\alpha_{t-1}}\beta_t}{1 - \bar\alpha_t}x_0
      #    + \frac{\sqrt{\alpha_t}(1 - \bar\alpha_{t-1})}{1-\bar\alpha_t}x_t$$
      mean = mean_x0_coef * x0 + mean_xt_coef * x
      # $\log \tilde\beta_t$
      log_var = x.new_full((bs, 1, 1, 1), self.log_var[step])
  
      # Do not add noise when $t = 1$ (final step sampling process).
      # Note that `step` is `0` when $t = 1$)
      if step == 0:
          noise = 0
      # If same noise is used for all samples in the batch
      elif repeat_noise:
          noise = torch.randn((1, *x.shape[1:]))
      # Different noise for each sample
      else:
          noise = torch.randn(x.shape)
  
      # Multiply noise by the temperature
      noise = noise * temperature
  
      # Sample from,
      #
      # $$p_\theta(x_{t-1} | x_t) = \mathcal{N}\big(x_{t-1}; \mu_\theta(x_t, t), \tilde\beta_t \mathbf{I} \big)$$
      x_prev = mean + (0.5 * log_var).exp() * noise
  
      #
      return x_prev, x0, e_t
  
file_path: annotated_deep_learning_paper_implementations/labml_nn/diffusion/stable_diffusion/sampler/ddpm.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 DDPMSampler obj's underlying member method code with example usage: obj.p_sample(x, c, t, step, repeat_noise, temperature, uncond_scale, uncond_cond) and returns: x_prev,  x0,  e_t