code: def register_schedule(
      self,
      given_betas=None,
      beta_schedule="linear",
      timesteps=1000,
      linear_start=1e-4,
      linear_end=2e-2,
      cosine_s=8e-3,
  ):
      if exists(given_betas):
          betas = given_betas
      else:
          betas = make_beta_schedule(
              beta_schedule,
              timesteps,
              linear_start=linear_start,
              linear_end=linear_end,
              cosine_s=cosine_s,
          )
      alphas = 1.0 - betas
      alphas_cumprod = np.cumprod(alphas, axis=0)
      alphas_cumprod_prev = np.append(1.0, alphas_cumprod[:-1])
  
      (timesteps,) = betas.shape
      self.num_timesteps = int(timesteps)
      self.linear_start = linear_start
      self.linear_end = linear_end
      assert (
          alphas_cumprod.shape[0] == self.num_timesteps
      ), "alphas have to be defined for each timestep"
  
      to_torch = partial(torch.tensor, dtype=torch.float32)
  
      self.register_buffer("betas", to_torch(betas))
      self.register_buffer("alphas_cumprod", to_torch(alphas_cumprod))
      self.register_buffer(
          "alphas_cumprod_prev", to_torch(alphas_cumprod_prev)
      )
  
      # calculations for diffusion q(x_t | x_{t-1}) and others
      self.register_buffer(
          "sqrt_alphas_cumprod", to_torch(np.sqrt(alphas_cumprod))
      )
      self.register_buffer(
          "sqrt_one_minus_alphas_cumprod",
          to_torch(np.sqrt(1.0 - alphas_cumprod)),
      )
      self.register_buffer(
          "log_one_minus_alphas_cumprod",
          to_torch(np.log(1.0 - alphas_cumprod)),
      )
      self.register_buffer(
          "sqrt_recip_alphas_cumprod",
          to_torch(np.sqrt(1.0 / alphas_cumprod)),
      )
      self.register_buffer(
          "sqrt_recipm1_alphas_cumprod",
          to_torch(np.sqrt(1.0 / alphas_cumprod - 1)),
      )
  
      # calculations for posterior q(x_{t-1} | x_t, x_0)
      posterior_variance = (1 - self.v_posterior) * betas * (
          1.0 - alphas_cumprod_prev
      ) / (1.0 - alphas_cumprod) + self.v_posterior * betas
      # above: equal to 1. / (1. / (1. - alpha_cumprod_tm1) + alpha_t / beta_t)
      self.register_buffer(
          "posterior_variance", to_torch(posterior_variance)
      )
      # below: log calculation clipped because the posterior variance is 0 at the beginning of the diffusion chain
      self.register_buffer(
          "posterior_log_variance_clipped",
          to_torch(np.log(np.maximum(posterior_variance, 1e-20))),
      )
      self.register_buffer(
          "posterior_mean_coef1",
          to_torch(
              betas
              * np.sqrt(alphas_cumprod_prev)
              / (1.0 - alphas_cumprod)
          ),
      )
      self.register_buffer(
          "posterior_mean_coef2",
          to_torch(
              (1.0 - alphas_cumprod_prev)
              * np.sqrt(alphas)
              / (1.0 - alphas_cumprod)
          ),
      )
  
      if self.parameterization == "eps":
          lvlb_weights = self.betas**2 / (
              2
              * self.posterior_variance
              * to_torch(alphas)
              * (1 - self.alphas_cumprod)
          )
      elif self.parameterization == "x0":
          lvlb_weights = (
              0.5
              * np.sqrt(torch.Tensor(alphas_cumprod))
              / (2.0 * 1 - torch.Tensor(alphas_cumprod))
          )
      else:
          raise NotImplementedError("mu not supported")
      # TODO how to choose this term
      lvlb_weights[0] = lvlb_weights[1]
      self.register_buffer(
          "lvlb_weights", lvlb_weights, persistent=False
      )
      assert not torch.isnan(self.lvlb_weights).all()
  
file_path: NeMo/nemo/collections/multimodal/models/text_to_image/dreambooth/util.py
imports: from functools import partial
  import numpy
  import torch
  import torch.nn
  from nemo.collections.multimodal.modules.stable_diffusion.diffusionmodules.util import extract_into_tensor, make_beta_schedule
  from nemo.collections.multimodal.parts.stable_diffusion.utils import default, exists
  from nemo.core.classes.common import Serialization
label_desc: Use image node_id 2 for calling the sd_noise_scheduler obj's underlying member method code with example usage: obj.register_schedule(given_betas, beta_schedule, timesteps, linear_start, linear_end, cosine_s) without return types