code: def make_beta_schedule(
      schedule,
      n_timestep,
      linear_start=1e-4,
      linear_end=2e-2,
      cosine_s=8e-3,
  ):
      if schedule == "linear":
          betas = (
              torch.linspace(
                  linear_start**0.5,
                  linear_end**0.5,
                  n_timestep,
                  dtype=torch.float64,
              )
              ** 2
          )
  
      elif schedule == "cosine":
          timesteps = (
              torch.arange(n_timestep + 1, dtype=torch.float64)
              / n_timestep
              + cosine_s
          )
          alphas = timesteps / (1 + cosine_s) * np.pi / 2
          alphas = torch.cos(alphas).pow(2)
          alphas = alphas / alphas[0]
          betas = 1 - alphas[1:] / alphas[:-1]
          betas = np.clip(betas, a_min=0, a_max=0.999)
  
      elif schedule == "sqrt_linear":
          betas = torch.linspace(
              linear_start, linear_end, n_timestep, dtype=torch.float64
          )
      elif schedule == "sqrt":
          betas = (
              torch.linspace(
                  linear_start,
                  linear_end,
                  n_timestep,
                  dtype=torch.float64,
              )
              ** 0.5
          )
      else:
          raise ValueError(f"schedule '{schedule}' unknown.")
      return betas.numpy()
  
imports: from collections import OrderedDict
  from copy import copy
  import numpy
  import tensorrt
  import torch
  from polygraphy import cuda
  from polygraphy.backend.common import bytes_from_path
  from polygraphy.backend.trt import engine_from_bytes
  from polygraphy.backend.trt import util
label_desc: Use image node_id 2 for calling a global function with example usage: make_beta_schedule(schedule, n_timestep, linear_start, linear_end, cosine_s) and returns: betas