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).
class_docstr_tok: ["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: Compute training losses for a single timestep.
  
  :param model: the model to evaluate loss on.
  :param x_start: the [N x C x ...] tensor of inputs.
  :param t: a batch of timestep indices.
  :param model_kwargs: if not None, a dict of extra keyword arguments to
      pass to the model. This can be used for conditioning.
  :param noise: if specified, the specific Gaussian noise to try to remove.
  :return: a dict with the key "loss" containing a tensor of shape [N].
           Some mean or variance settings may also have other keys.
docstr_tok: ["Compute","training","losses","for","a","single","timestep",".",":","param","model",":","the","model","to","evaluate","loss","on",".",":","param","x_start",":","the","[","N","x","C","x","...","]","tensor","of","inputs",".",":","param","t",":","a","batch","of","timestep","indices",".",":","param","model_kwargs",":","if","not","None",",","a","dict","of","extra","keyword","arguments","to","pass","to","the","model",".","This","can","be","used","for","conditioning",".",":","param","noise",":","if","specified",",","the","specific","Gaussian","noise","to","try","to","remove",".",":","return",":","a","dict","with","the","key","``","loss","''","containing","a","tensor","of","shape","[","N","]",".","Some","mean","or","variance","settings","may","also","have","other","keys","."]
code: def autoregressive_training_losses(
      self,
      model,
      x_start,
      t,
      model_output_keys,
      gd_out_key,
      model_kwargs=None,
      noise=None,
  ):
      """
      Compute training losses for a single timestep.
  
      :param model: the model to evaluate loss on.
      :param x_start: the [N x C x ...] tensor of inputs.
      :param t: a batch of timestep indices.
      :param model_kwargs: if not None, a dict of extra keyword arguments to
          pass to the model. This can be used for conditioning.
      :param noise: if specified, the specific Gaussian noise to try to remove.
      :return: a dict with the key "loss" containing a tensor of shape [N].
               Some mean or variance settings may also have other keys.
      """
      if model_kwargs is None:
          model_kwargs = {}
      if noise is None:
          noise = th.randn_like(x_start)
      x_t = self.q_sample(x_start, t, noise=noise)
      terms = {}
      if (
          self.loss_type == LossType.KL
          or self.loss_type == LossType.RESCALED_KL
      ):
          assert False  # not currently supported for this type of diffusion.
      elif (
          self.loss_type == LossType.MSE
          or self.loss_type == LossType.RESCALED_MSE
      ):
          model_outputs = model(
              x_t, x_start, self._scale_timesteps(t), **model_kwargs
          )
          terms.update(
              {k: o for k, o in zip(model_output_keys, model_outputs)}
          )
          model_output = terms[gd_out_key]
          if self.model_var_type in [
              ModelVarType.LEARNED,
              ModelVarType.LEARNED_RANGE,
          ]:
              B, C = x_t.shape[:2]
              assert model_output.shape == (B, C, 2, *x_t.shape[2:])
              model_output, model_var_values = (
                  model_output[:, :, 0],
                  model_output[:, :, 1],
              )
              # Learn the variance using the variational bound, but don't let
              # it affect our mean prediction.
              frozen_out = th.cat(
                  [model_output.detach(), model_var_values], dim=1
              )
              terms["vb"] = self._vb_terms_bpd(
                  model=lambda *args, r=frozen_out: r,
                  x_start=x_start,
                  x_t=x_t,
                  t=t,
                  clip_denoised=False,
              )["output"]
              if self.loss_type == LossType.RESCALED_MSE:
                  # Divide by 1000 for equivalence with initial implementation.
                  # Without a factor of 1/1000, the VB term hurts the MSE term.
                  terms["vb"] *= self.num_timesteps / 1000.0
  
          if self.model_mean_type == ModelMeanType.PREVIOUS_X:
              target = self.q_posterior_mean_variance(
                  x_start=x_start, x_t=x_t, t=t
              )[0]
              x_start_pred = torch.zeros(x_start)  # Not supported.
          elif self.model_mean_type == ModelMeanType.START_X:
              target = x_start
              x_start_pred = model_output
          elif self.model_mean_type == ModelMeanType.EPSILON:
              target = noise
              x_start_pred = self._predict_xstart_from_eps(
                  x_t, t, model_output
              )
          else:
              raise NotImplementedError(self.model_mean_type)
          assert model_output.shape == target.shape == x_start.shape
          terms["mse"] = mean_flat((target - model_output) ** 2)
          terms["x_start_predicted"] = x_start_pred
          if "vb" in terms:
              terms["loss"] = terms["mse"] + terms["vb"]
          else:
              terms["loss"] = terms["mse"]
      else:
          raise NotImplementedError(self.loss_type)
  
      return terms
  
imports: import enum
  import math
  import numpy
  import torch
  import torch
  from tqdm import tqdm
  from TTS.tts.layers.tortoise.dpm_solver import DPM_Solver, NoiseScheduleVP, model_wrapper
label_desc: Use image node_id 23 for calling the GaussianDiffusion obj's underlying member method code with example usage: obj.autoregressive_training_losses(model, x_start, t, model_output_keys, gd_out_key, model_kwargs, noise) and returns: terms