Class ZoeScheduler

trainer/scheduler.py:10–55

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
        https://arxiv.org/abs/2303.09556."""
        snr = self.get_snr(timesteps)
        min_snr_gamma = torch.min(snr, torch.ones_like(snr) * gamma)
        if self.config.prediction_type == 'epsilon':
            return min_snr_gamma / (snr + 1e-8)
        elif self.config.prediction_type == 'sample':
            return min_snr_gamma
        elif self.config.prediction_type == 'v_prediction':
            return min_snr_gamma / (snr + 1)

    def get_min_snr_weight_modified(self,
                                    timesteps,
                                    gamma: float = 5.0,
                                    min_weight: float = 0.01):
        """modified loss weighting, so that the weighting at zero snr is not
        zero the original min snr weighting in paper
        https://arxiv.org/abs/2303.09556."""
        snr = self.get_snr(timesteps)
        min_snr_gamma = torch.min(snr, torch.ones_like(snr) * gamma)
        if self.config.prediction_type == 'epsilon':
            return torch.max(min_snr_gamma / (snr + 1e-8), min_weight)
        elif self.config.prediction_type == 'sample':
            return torch.max(min_snr_gamma, min_weight)
        elif self.config.prediction_type == 'v_prediction':
            return torch.max(min_snr_gamma / (snr + 1), min_weight)

    def get_min_snr_weight_cogvideo(self, timesteps):
        """modified loss weighting, so that the weighting at zero snr is not
        zero the original min snr weighting in paper
        https://arxiv.org/abs/2303.09556."""
        snr = self.get_snr(timesteps)
        if self.config.prediction_type == 'epsilon':
            return (snr + 1) / (snr + 1e-8)
        elif self.config.prediction_type == 'sample':
            return snr + 1
        elif self.config.prediction_type == 'v_prediction':
            return torch.ones_like(snr)


class CogVideoXDDIMScheduler(CogVideoXDDIMSchedulerBase, ZoeScheduler):