docstr: Create a wrapper class for the forward SDE.
docstr_tok: ["Create","a","wrapper","class","for","the","forward","SDE","."]
code: def __init__(
      self,
      schedule="discrete",
      betas=None,
      alphas_cumprod=None,
      continuous_beta_0=0.1,
      continuous_beta_1=20.0,
  ):
      """Create a wrapper class for the forward SDE."""
  
      if schedule not in ["discrete", "linear", "cosine"]:
          raise ValueError(
              "Unsupported noise schedule {}. The schedule needs to be 'discrete' or 'linear' or 'cosine'".format(
                  schedule
              )
          )
  
      self.schedule = schedule
      if schedule == "discrete":
          if betas is not None:
              log_alphas = 0.5 * torch.log(1 - betas).cumsum(dim=0)
          else:
              assert alphas_cumprod is not None
              log_alphas = 0.5 * torch.log(alphas_cumprod)
          self.total_N = len(log_alphas)
          self.T = 1.0
          self.t_array = torch.linspace(0.0, 1.0, self.total_N + 1)[
              1:
          ].reshape((1, -1))
          self.log_alpha_array = log_alphas.reshape(
              (
                  1,
                  -1,
              )
          )
      else:
          self.total_N = 1000
          self.beta_0 = continuous_beta_0
          self.beta_1 = continuous_beta_1
          self.cosine_s = 0.008
          self.cosine_beta_max = 999.0
          self.cosine_t_max = (
              math.atan(
                  self.cosine_beta_max * (1.0 + self.cosine_s) / math.pi
              )
              * 2.0
              * (1.0 + self.cosine_s)
              / math.pi
              - self.cosine_s
          )
          self.cosine_log_alpha_0 = math.log(
              math.cos(
                  self.cosine_s / (1.0 + self.cosine_s) * math.pi / 2.0
              )
          )
          self.schedule = schedule
          if schedule == "cosine":
              self.T = 0.9946
          else:
              self.T = 1.0
  
imports: import math
  import torch
  from nemo.collections.multimodal.modules.stable_diffusion.diffusionmodules.util import expand_dims, interpolate_fn
label_desc: Use image node_id 1 to create a new NoiseScheduleVP object  with example: obj = NoiseScheduleVP(schedule, betas, alphas_cumprod, continuous_beta_0, continuous_beta_1)