Method get_snr

trainer/scheduler.py:12–15

(self, timesteps)

class ZoeScheduler:

    def get_snr(self, timesteps):
        alphas_cumprod = self.alphas_cumprod[timesteps]
        sigma_power = 1 - alphas_cumprod
        return alphas_cumprod / sigma_power

    def get_min_snr_weight(self, timesteps, gamma: float = 5.0):
        """the original min snr weighting in paper