class_docstr: Utilities for training and sampling diffusion models.
  
  Ported directly from here, and then adapted over time to further experimentation.
  https://github.com/hojonathanho/diffusion/blob/1e0dceb3b3495bbe19116a5e1b3596cd0706c543/diffusion_tf/diffusion_utils_2.py#L42
  
  :param betas: a 1-D numpy array of betas for each diffusion timestep,
                starting at T and going to 1.
  :param model_mean_type: a ModelMeanType determining what the model outputs.
  :param model_var_type: a ModelVarType determining how variance is output.
  :param loss_type: a LossType determining the loss function to use.
  :param rescale_timesteps: if True, pass floating point timesteps into the
                            model so that they are always scaled like in the
                            original paper (0 to 1000).
docstr: Apply the model to get p(x_{t-1} | x_t), as well as a prediction of
  the initial x, x_0.
  
  :param model: the model, which takes a signal and a batch of timesteps
                as input.
  :param x: the [N x C x ...] tensor at time t.
  :param t: a 1-D Tensor of timesteps.
  :param clip_denoised: if True, clip the denoised signal into [-1, 1].
  :param denoised_fn: if not None, a function which applies to the
      x_start prediction before it is used to sample. Applies before
      clip_denoised.
  :param model_kwargs: if not None, a dict of extra keyword arguments to
      pass to the model. This can be used for conditioning.
  :return: a dict with the following keys:
           - 'mean': the model mean output.
           - 'variance': the model variance output.
           - 'log_variance': the log of 'variance'.
           - 'pred_xstart': the prediction for x_0.
docstr_tok: ["Apply","the","model","to","get","p","(","x_","{","t-1","}","|","x_t",")",",","as","well","as","a","prediction","of","the","initial","x",",","x_0",".",":","param","model",":","the","model",",","which","takes","a","signal","and","a","batch","of","timesteps","as","input",".",":","param","x",":","the","[","N","x","C","x","...","]","tensor","at","time","t.",":","param","t",":","a","1-D","Tensor","of","timesteps",".",":","param","clip_denoised",":","if","True",",","clip","the","denoised","signal","into","[","-1",",","1","]",".",":","param","denoised_fn",":","if","not","None",",","a","function","which","applies","to","the","x_start","prediction","before","it","is","used","to","sample",".","Applies","before","clip_denoised",".",":","param","model_kwargs",":","if","not","None",",","a","dict","of","extra","keyword","arguments","to","pass","to","the","model",".","This","can","be","used","for","conditioning",".",":","return",":","a","dict","with","the","following","keys",":","-","'mean","'",":","the","model","mean","output",".","-","'variance","'",":","the","model","variance","output",".","-","'log_variance","'",":","the","log","of","'variance","'",".","-","'pred_xstart","'",":","the","prediction","for","x_0","."]
code: def p_mean_variance(
      self,
      model,
      x,
      t,
      clip_denoised=True,
      denoised_fn=None,
      model_kwargs=None,
  ):
      """
      Apply the model to get p(x_{t-1} | x_t), as well as a prediction of
      the initial x, x_0.
  
      :param model: the model, which takes a signal and a batch of timesteps
                    as input.
      :param x: the [N x C x ...] tensor at time t.
      :param t: a 1-D Tensor of timesteps.
      :param clip_denoised: if True, clip the denoised signal into [-1, 1].
      :param denoised_fn: if not None, a function which applies to the
          x_start prediction before it is used to sample. Applies before
          clip_denoised.
      :param model_kwargs: if not None, a dict of extra keyword arguments to
          pass to the model. This can be used for conditioning.
      :return: a dict with the following keys:
               - 'mean': the model mean output.
               - 'variance': the model variance output.
               - 'log_variance': the log of 'variance'.
               - 'pred_xstart': the prediction for x_0.
      """
      if model_kwargs is None:
          model_kwargs = {}
  
      B, C = x.shape[:2]
      assert t.shape == (B,)
      model_output = model(x, self._scale_timesteps(t), **model_kwargs)
  
      if self.model_var_type in [
          ModelVarType.LEARNED,
          ModelVarType.LEARNED_RANGE,
      ]:
          assert model_output.shape == (B, C * 2, *x.shape[2:])
          model_output, model_var_values = th.split(
              model_output, C, dim=1
          )
          if self.model_var_type == ModelVarType.LEARNED:
              model_log_variance = model_var_values
              model_variance = th.exp(model_log_variance)
          else:
              min_log = _extract_into_tensor(
                  self.posterior_log_variance_clipped, t, x.shape
              )
              max_log = _extract_into_tensor(
                  np.log(self.betas), t, x.shape
              )
              # The model_var_values is [-1, 1] for [min_var, max_var].
              frac = (model_var_values + 1) / 2
              model_log_variance = frac * max_log + (1 - frac) * min_log
              model_variance = th.exp(model_log_variance)
      else:
          model_variance, model_log_variance = {
              # for fixedlarge, we set the initial (log-)variance like so
              # to get a better decoder log likelihood.
              ModelVarType.FIXED_LARGE: (
                  np.append(self.posterior_variance[1], self.betas[1:]),
                  np.log(
                      np.append(
                          self.posterior_variance[1], self.betas[1:]
                      )
                  ),
              ),
              ModelVarType.FIXED_SMALL: (
                  self.posterior_variance,
                  self.posterior_log_variance_clipped,
              ),
          }[self.model_var_type]
          model_variance = _extract_into_tensor(
              model_variance, t, x.shape
          )
          model_log_variance = _extract_into_tensor(
              model_log_variance, t, x.shape
          )
  
      def process_xstart(x):
          if denoised_fn is not None:
              x = denoised_fn(x)
          if clip_denoised:
              return x.clamp(-1, 1)
          return x
  
      if self.model_mean_type == ModelMeanType.PREVIOUS_X:
          pred_xstart = process_xstart(
              self._predict_xstart_from_xprev(
                  x_t=x, t=t, xprev=model_output
              )
          )
          model_mean = model_output
      elif self.model_mean_type in [
          ModelMeanType.START_X,
          ModelMeanType.EPSILON,
      ]:
          if self.model_mean_type == ModelMeanType.START_X:
              pred_xstart = process_xstart(model_output)
          else:
              pred_xstart = process_xstart(
                  self._predict_xstart_from_eps(
                      x_t=x, t=t, eps=model_output
                  )
              )
          model_mean, _, _ = self.q_posterior_mean_variance(
              x_start=pred_xstart, x_t=x, t=t
          )
      else:
          raise NotImplementedError(self.model_mean_type)
  
      assert (
          model_mean.shape
          == model_log_variance.shape
          == pred_xstart.shape
          == x.shape
      )
      return {
          "mean": model_mean,
          "variance": model_variance,
          "log_variance": model_log_variance,
          "pred_xstart": pred_xstart,
      }
  
imports: import enum
  import math
  import numpy
  import torch
  from .nn import mean_flat
  from .losses import normal_kl, discretized_gaussian_log_likelihood
label_desc: Use image node_id 5 for calling the GaussianDiffusion obj's underlying member method code with example usage: obj.p_mean_variance(model, x, t, clip_denoised, denoised_fn, model_kwargs) and returns: dict, x, x