diffusion_models.gaussian_diffusion.beta_schedulers
===================================================

.. py:module:: diffusion_models.gaussian_diffusion.beta_schedulers




Module Contents
---------------

.. py:class:: BaseBetaScheduler(steps, enforce_zero_terminal_snr = False, initialize = True)

   Initializes a beta scheduler.

   BaseBetaScheduler is an abstract base class for different beta scheduler
   implementations. It defines the interface that all beta schedulers should
   adhere to.

   :param steps: The number of steps for the beta.
   :param enforce_zero_terminal_snr: Whether to enforce zero terminal SNR inline
                                     with `"Common Diffusion Noise Schedules and Sample Steps are Flawed"
                                     <https://arxiv.org/abs/2305.08891>`_.

                                     Defaults to ``False``.
   :param initialize: Whether to initialize the beta scheduler. If this is
                      set to ``False``, you will need to manually set ``self.betas`` and
                      ``self.alpha_bars``. Otherwise, they are initialized using your
                      ``sample_betas`` and ``compute_alpha_bar`` methods.

   .. warning::

      Do not instantiate this class directly. Instead, build your own Beta
      scheduler by inheriting from BaseBetaScheduler.
      (see :class:`~.LinearBetaScheduler`)


   .. py:attribute:: steps
      :type:  int

      The number of steps for the beta scheduler.


   .. py:attribute:: betas
      :value: None


      The :math:`\beta` computed according to :meth:`~.BaseBetaScheduler.sample_betas`.


   .. py:attribute:: alpha_bars
      :value: None


      The :math:`\bar{\alpha}` computed according to :meth:`~.BaseBetaScheduler.compute_alpha_bar`.


   .. py:method:: enforce_zero_terminal_snr()

      Enforce terminal SNR by adjusting :math:`\beta` and :math:`\bar{\alpha}`.

      This method enforces zero terminal SNR according to
      `"Common Diffusion Noise Schedules and Sample Steps are Flawed"
      <https://arxiv.org/abs/2305.08891>`_.



   .. py:method:: sample_betas()
      :abstractmethod:


      Compute :math:`\beta` for noise scheduling.

      :returns: A torch tensor of the :math:`\beta` values.



   .. py:method:: compute_alpha_bar()
      :abstractmethod:


      Compute :math:`\bar{\alpha}` for noise scheduling.

      :returns: A torch tensor of the :math:`\bar{\alpha}` values.



   .. py:method:: to(device)

      Moves the beta scheduler to the given device.

      :param device: The device to which the method should move the object.
                     Default is "cpu".



   .. py:method:: from_tensors(steps, betas, alpha_bars)
      :classmethod:


      Instantiate a beta scheduler from tensors.

      Instantiate a beta scheduler from tensors. This is particularly useful for
      loading checkpoints.

      :param steps: The number of steps for the beta scheduler.
      :param betas: The pre-computed beta values for the noise scheduler.
      :param alpha_bars: The pre-computed alpha bar values for the noise scheduler.

      Returns:




.. py:class:: LinearBetaScheduler(beta_start = 0.0001, beta_end = 0.02, steps = 1000, enforce_zero_terminal_snr = True)



   A Linear Beta scheduler.

   A simple linear beta scheduler with betas linearly spaced between
   ``beta_start`` and ``beta_end``.

   :param beta_start: The starting value of the betas.
   :param beta_end: The end value of the betas.
   :param steps: The number of steps for the beta scheduler. This is also the number
                 of betas.
   :param enforce_zero_terminal_snr: Whether to enforce zero terminal SNR.


   .. py:attribute:: beta_start
      :type:  int

      The starting value of the betas.


   .. py:attribute:: beta_end
      :type:  int

      The end value of the betas.


   .. py:method:: sample_betas()

      Return linearly spaced betas between ``self.beta_start`` and ``self.beta_end``.



   .. py:method:: compute_alpha_bar()

      Return :math:`\bar{\alpha}` computed from the beta values.



.. py:class:: CosineBetaScheduler(offset = 0.008, steps = 1000, max_beta = 0.999)



   A Cosine Beta scheduler.

   A Cosine Beta Scheduler based on the following formulas:

   .. math::
     :nowrap:

     \begin{equation}
     \left\{ \begin{aligned}
       \bar{\alpha}_t &= \frac{f(t)}{f(0)} \\
       \beta_t &= 1 - \frac{\bar{\alpha}_t}{\bar{\alpha}_t -1}
     \end{aligned} \right.
     \end{equation}

   where

     .. math::

       f(t) = \cos(\frac{t/T + s}{1 + s} * \frac{\pi}{2})^2

   where

     .. math::
       :nowrap:

       \begin{equation}
       \left\{ \begin{aligned}

         s & \text{ is the offset} \\
         T & \text{ is the number of steps}

       \end{aligned} \right.
       \end{equation}

   :param offset: The offset :math:`s` defined above.
   :param steps: The number of steps for the beta scheduler.
   :param max_beta: The maximum beta values. Higher values will be clipped.


   .. py:attribute:: offset
      :type:  float

      The offset :math:`s` defined above.


   .. py:attribute:: max_beta
      :type:  Optional[float]

      The maximum beta values. Higher values will be clipped.


   .. py:attribute:: steps
      :type:  int

      The number of steps for the beta scheduler.


   .. py:method:: f(t)

      A helper function to compute the :math:`\bar{\alpha}_t`.

      :param t: The timestep to compute.

      :returns:

                .. math::

                  f(t) = \cos(\frac{t/T + s}{1 + s} * \frac{\pi}{2})^2



   .. py:method:: sample_betas()

      Compute :math:`\beta` for noise scheduling.

      :returns: A torch tensor of the :math:`\beta` values.



   .. py:method:: compute_alpha_bar()

      Compute :math:`\bar{\alpha}` for noise scheduling.

      :returns: A torch tensor of the :math:`\bar{\alpha}` values.