docstr: Create a wrapper class for the forward SDE (VP type).
  
  ***
  Update: We support discrete-time diffusion models by implementing a picewise linear interpolation for log_alpha_t.
          We recommend to use schedule='discrete' for the discrete-time diffusion models, especially for high-resolution images.
  ***
  
  The forward SDE ensures that the condition distribution q_{t|0}(x_t | x_0) = N ( alpha_t * x_0, sigma_t^2 * I ).
  We further define lambda_t = log(alpha_t) - log(sigma_t), which is the half-logSNR (described in the DPM-Solver paper).
  Therefore, we implement the functions for computing alpha_t, sigma_t and lambda_t. For t in [0, T], we have:
  
      log_alpha_t = self.marginal_log_mean_coeff(t)
      sigma_t = self.marginal_std(t)
      lambda_t = self.marginal_lambda(t)
  
  Moreover, as lambda(t) is an invertible function, we also support its inverse function:
  
      t = self.inverse_lambda(lambda_t)
  
  ===============================================================
  
  We support both discrete-time DPMs (trained on n = 0, 1, ..., N-1) and continuous-time DPMs (trained on t in [t_0, T]).
  
  1. For discrete-time DPMs:
  
      For discrete-time DPMs trained on n = 0, 1, ..., N-1, we convert the discrete steps to continuous time steps by:
          t_i = (i + 1) / N
      e.g. for N = 1000, we have t_0 = 1e-3 and T = t_{N-1} = 1.
      We solve the corresponding diffusion ODE from time T = 1 to time t_0 = 1e-3.
  
      Args:
          betas: A `torch.Tensor`. The beta array for the discrete-time DPM. (See the original DDPM paper for details)
          alphas_cumprod: A `torch.Tensor`. The cumprod alphas for the discrete-time DPM. (See the original DDPM paper for details)
  
      Note that we always have alphas_cumprod = cumprod(betas). Therefore, we only need to set one of `betas` and `alphas_cumprod`.
  
      **Important**:  Please pay special attention for the args for `alphas_cumprod`:
          The `alphas_cumprod` is the \hat{alpha_n} arrays in the notations of DDPM. Specifically, DDPMs assume that
              q_{t_n | 0}(x_{t_n} | x_0) = N ( \sqrt{\hat{alpha_n}} * x_0, (1 - \hat{alpha_n}) * I ).
          Therefore, the notation \hat{alpha_n} is different from the notation alpha_t in DPM-Solver. In fact, we have
              alpha_{t_n} = \sqrt{\hat{alpha_n}},
          and
              log(alpha_{t_n}) = 0.5 * log(\hat{alpha_n}).
  
  
  2. For continuous-time DPMs:
  
      We support two types of VPSDEs: linear (DDPM) and cosine (improved-DDPM). The hyperparameters for the noise
      schedule are the default settings in DDPM and improved-DDPM:
  
      Args:
          beta_min: A `float` number. The smallest beta for the linear schedule.
          beta_max: A `float` number. The largest beta for the linear schedule.
          cosine_s: A `float` number. The hyperparameter in the cosine schedule.
          cosine_beta_max: A `float` number. The hyperparameter in the cosine schedule.
          T: A `float` number. The ending time of the forward process.
  
  ===============================================================
  
  Args:
      schedule: A `str`. The noise schedule of the forward SDE. 'discrete' for discrete-time DPMs,
              'linear' or 'cosine' for continuous-time DPMs.
  Returns:
      A wrapper object of the forward SDE (VP type).
  
  ===============================================================
  
  Example:
  
  # For discrete-time DPMs, given betas (the beta array for n = 0, 1, ..., N - 1):
  >>> ns = NoiseScheduleVP('discrete', betas=betas)
  
  # For discrete-time DPMs, given alphas_cumprod (the \hat{alpha_n} array for n = 0, 1, ..., N - 1):
  >>> ns = NoiseScheduleVP('discrete', alphas_cumprod=alphas_cumprod)
  
  # For continuous-time DPMs (VPSDE), linear schedule:
  >>> ns = NoiseScheduleVP('linear', continuous_beta_0=0.1, continuous_beta_1=20.)
docstr_tok: ["Create","a","wrapper","class","for","the","forward","SDE","(","VP","type",")",".","*","*","*","Update",":","We","support","discrete-time","diffusion","models","by","implementing","a","picewise","linear","interpolation","for","log_alpha_t",".","We","recommend","to","use","schedule='discrete","'","for","the","discrete-time","diffusion","models",",","especially","for","high-resolution","images",".","*","*","*","The","forward","SDE","ensures","that","the","condition","distribution","q_","{","t|0","}","(","x_t","|","x_0",")","=","N","(","alpha_t","*","x_0",",","sigma_t^2","*","I",")",".","We","further","define","lambda_t","=","log","(","alpha_t",")","-","log","(","sigma_t",")",",","which","is","the","half-logSNR","(","described","in","the","DPM-Solver","paper",")",".","Therefore",",","we","implement","the","functions","for","computing","alpha_t",",","sigma_t","and","lambda_t",".","For","t","in","[","0",",","T","]",",","we","have",":","log_alpha_t","=","self.marginal_log_mean_coeff","(","t",")","sigma_t","=","self.marginal_std","(","t",")","lambda_t","=","self.marginal_lambda","(","t",")","Moreover",",","as","lambda","(","t",")","is","an","invertible","function",",","we","also","support","its","inverse","function",":","t","=","self.inverse_lambda","(","lambda_t",")","===============================================================","We","support","both","discrete-time","DPMs","(","trained","on","n","=","0",",","1",",","...",",","N-1",")","and","continuous-time","DPMs","(","trained","on","t","in","[","t_0",",","T","]",")",".","1",".","For","discrete-time","DPMs",":","For","discrete-time","DPMs","trained","on","n","=","0",",","1",",","...",",","N-1",",","we","convert","the","discrete","steps","to","continuous","time","steps","by",":","t_i","=","(","i","+","1",")","\/","N","e.g",".","for","N","=","1000",",","we","have","t_0","=","1e-3","and","T","=","t_","{","N-1","}","=","1",".","We","solve","the","corresponding","diffusion","ODE","from","time","T","=","1","to","time","t_0","=","1e-3",".","Args",":","betas",":","A","`","torch.Tensor","`",".","The","beta","array","for","the","discrete-time","DPM",".","(","See","the","original","DDPM","paper","for","details",")","alphas_cumprod",":","A","`","torch.Tensor","`",".","The","cumprod","alphas","for","the","discrete-time","DPM",".","(","See","the","original","DDPM","paper","for","details",")","Note","that","we","always","have","alphas_cumprod","=","cumprod","(","betas",")",".","Therefore",",","we","only","need","to","set","one","of","`","betas","`","and","`","alphas_cumprod","`",".","*","*","Important","*","*",":","Please","pay","special","attention","for","the","args","for","`","alphas_cumprod","`",":","The","`","alphas_cumprod","`","is","the","\\hat","{","alpha_n","}","arrays","in","the","notations","of","DDPM",".","Specifically",",","DDPMs","assume","that","q_","{","t_n","|","0","}","(","x_","{","t_n","}","|","x_0",")","=","N","(","\\sqrt","{","\\hat","{","alpha_n","}","}","*","x_0",",","(","1","-","\\hat","{","alpha_n","}",")","*","I",")",".","Therefore",",","the","notation","\\hat","{","alpha_n","}","is","different","from","the","notation","alpha_t","in","DPM-Solver",".","In","fact",",","we","have","alpha_","{","t_n","}","=","\\sqrt","{","\\hat","{","alpha_n","}","}",",","and","log","(","alpha_","{","t_n","}",")","=","0.5","*","log","(","\\hat","{","alpha_n","}",")",".","2",".","For","continuous-time","DPMs",":","We","support","two","types","of","VPSDEs",":","linear","(","DDPM",")","and","cosine","(","improved-DDPM",")",".","The","hyperparameters","for","the","noise","schedule","are","the","default","settings","in","DDPM","and","improved-DDPM",":","Args",":","beta_min",":","A","`","float","`","number",".","The","smallest","beta","for","the","linear","schedule",".","beta_max",":","A","`","float","`","number",".","The","largest","beta","for","the","linear","schedule",".","cosine_s",":","A","`","float","`","number",".","The","hyperparameter","in","the","cosine","schedule",".","cosine_beta_max",":","A","`","float","`","number",".","The","hyperparameter","in","the","cosine","schedule",".","T",":","A","`","float","`","number",".","The","ending","time","of","the","forward","process",".","===============================================================","Args",":","schedule",":","A","`","str","`",".","The","noise","schedule","of","the","forward","SDE",".","'discrete","'","for","discrete-time","DPMs",",","'linear","'","or","'cosine","'","for","continuous-time","DPMs",".","Returns",":","A","wrapper","object","of","the","forward","SDE","(","VP","type",")",".","===============================================================","Example",":","#","For","discrete-time","DPMs",",","given","betas","(","the","beta","array","for","n","=","0",",","1",",","...",",","N","-","1",")",":",">",">",">","ns","=","NoiseScheduleVP","(","'discrete","'",",","betas=betas",")","#","For","discrete-time","DPMs",",","given","alphas_cumprod","(","the","\\hat","{","alpha_n","}","array","for","n","=","0",",","1",",","...",",","N","-","1",")",":",">",">",">","ns","=","NoiseScheduleVP","(","'discrete","'",",","alphas_cumprod=alphas_cumprod",")","#","For","continuous-time","DPMs","(","VPSDE",")",",","linear","schedule",":",">",">",">","ns","=","NoiseScheduleVP","(","'linear","'",",","continuous_beta_0=0.1",",","continuous_beta_1=20",".",")"]
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 (VP type).
  
      ***
      Update: We support discrete-time diffusion models by implementing a picewise linear interpolation for log_alpha_t.
              We recommend to use schedule='discrete' for the discrete-time diffusion models, especially for high-resolution images.
      ***
  
      The forward SDE ensures that the condition distribution q_{t|0}(x_t | x_0) = N ( alpha_t * x_0, sigma_t^2 * I ).
      We further define lambda_t = log(alpha_t) - log(sigma_t), which is the half-logSNR (described in the DPM-Solver paper).
      Therefore, we implement the functions for computing alpha_t, sigma_t and lambda_t. For t in [0, T], we have:
  
          log_alpha_t = self.marginal_log_mean_coeff(t)
          sigma_t = self.marginal_std(t)
          lambda_t = self.marginal_lambda(t)
  
      Moreover, as lambda(t) is an invertible function, we also support its inverse function:
  
          t = self.inverse_lambda(lambda_t)
  
      ===============================================================
  
      We support both discrete-time DPMs (trained on n = 0, 1, ..., N-1) and continuous-time DPMs (trained on t in [t_0, T]).
  
      1. For discrete-time DPMs:
  
          For discrete-time DPMs trained on n = 0, 1, ..., N-1, we convert the discrete steps to continuous time steps by:
              t_i = (i + 1) / N
          e.g. for N = 1000, we have t_0 = 1e-3 and T = t_{N-1} = 1.
          We solve the corresponding diffusion ODE from time T = 1 to time t_0 = 1e-3.
  
          Args:
              betas: A `torch.Tensor`. The beta array for the discrete-time DPM. (See the original DDPM paper for details)
              alphas_cumprod: A `torch.Tensor`. The cumprod alphas for the discrete-time DPM. (See the original DDPM paper for details)
  
          Note that we always have alphas_cumprod = cumprod(betas). Therefore, we only need to set one of `betas` and `alphas_cumprod`.
  
          **Important**:  Please pay special attention for the args for `alphas_cumprod`:
              The `alphas_cumprod` is the \hat{alpha_n} arrays in the notations of DDPM. Specifically, DDPMs assume that
                  q_{t_n | 0}(x_{t_n} | x_0) = N ( \sqrt{\hat{alpha_n}} * x_0, (1 - \hat{alpha_n}) * I ).
              Therefore, the notation \hat{alpha_n} is different from the notation alpha_t in DPM-Solver. In fact, we have
                  alpha_{t_n} = \sqrt{\hat{alpha_n}},
              and
                  log(alpha_{t_n}) = 0.5 * log(\hat{alpha_n}).
  
  
      2. For continuous-time DPMs:
  
          We support two types of VPSDEs: linear (DDPM) and cosine (improved-DDPM). The hyperparameters for the noise
          schedule are the default settings in DDPM and improved-DDPM:
  
          Args:
              beta_min: A `float` number. The smallest beta for the linear schedule.
              beta_max: A `float` number. The largest beta for the linear schedule.
              cosine_s: A `float` number. The hyperparameter in the cosine schedule.
              cosine_beta_max: A `float` number. The hyperparameter in the cosine schedule.
              T: A `float` number. The ending time of the forward process.
  
      ===============================================================
  
      Args:
          schedule: A `str`. The noise schedule of the forward SDE. 'discrete' for discrete-time DPMs,
                  'linear' or 'cosine' for continuous-time DPMs.
      Returns:
          A wrapper object of the forward SDE (VP type).
  
      ===============================================================
  
      Example:
  
      # For discrete-time DPMs, given betas (the beta array for n = 0, 1, ..., N - 1):
      >>> ns = NoiseScheduleVP('discrete', betas=betas)
  
      # For discrete-time DPMs, given alphas_cumprod (the \hat{alpha_n} array for n = 0, 1, ..., N - 1):
      >>> ns = NoiseScheduleVP('discrete', alphas_cumprod=alphas_cumprod)
  
      # For continuous-time DPMs (VPSDE), linear schedule:
      >>> ns = NoiseScheduleVP('linear', continuous_beta_0=0.1, continuous_beta_1=20.)
  
      """
  
      if schedule not in ["discrete", "linear", "cosine"]:
          raise ValueError(
              f"Unsupported noise schedule {schedule}. The schedule needs to be 'discrete' or 'linear' or 'cosine'"
          )
  
      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":
              # For the cosine schedule, T = 1 will have numerical issues. So we manually set the ending time T.
              # Note that T = 0.9946 may be not the optimal setting. However, we find it works well.
              self.T = 0.9946
          else:
              self.T = 1.0
  
imports: import torch
  import math
  import tqdm
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)