fred-dev/comfy_ui_ali

1import torch2from comfy.ldm.modules.diffusionmodules.util import make_beta_schedule3import math4 5def rescale_zero_terminal_snr_sigmas(sigmas):6    alphas_cumprod = 1 / ((sigmas * sigmas) + 1)7    alphas_bar_sqrt = alphas_cumprod.sqrt()8 9    # Store old values.10    alphas_bar_sqrt_0 = alphas_bar_sqrt[0].clone()11    alphas_bar_sqrt_T = alphas_bar_sqrt[-1].clone()12 13    # Shift so the last timestep is zero.14    alphas_bar_sqrt -= (alphas_bar_sqrt_T)15 16    # Scale so the first timestep is back to the old value.17    alphas_bar_sqrt *= alphas_bar_sqrt_0 / (alphas_bar_sqrt_0 - alphas_bar_sqrt_T)18 19    # Convert alphas_bar_sqrt to betas20    alphas_bar = alphas_bar_sqrt**2  # Revert sqrt21    alphas_bar[-1] = 4.8973451890853435e-0822    return ((1 - alphas_bar) / alphas_bar) ** 0.523 24class EPS:25    def calculate_input(self, sigma, noise):26        sigma = sigma.view(sigma.shape[:1] + (1,) * (noise.ndim - 1))27        return noise / (sigma ** 2 + self.sigma_data ** 2) ** 0.528 29    def calculate_denoised(self, sigma, model_output, model_input):30        sigma = sigma.view(sigma.shape[:1] + (1,) * (model_output.ndim - 1))31        return model_input - model_output * sigma32 33    def noise_scaling(self, sigma, noise, latent_image, max_denoise=False):34        sigma = sigma.view(sigma.shape[:1] + (1,) * (noise.ndim - 1))35        if max_denoise:36            noise = noise * torch.sqrt(1.0 + sigma ** 2.0)37        else:38            noise = noise * sigma39 40        noise += latent_image41        return noise42 43    def inverse_noise_scaling(self, sigma, latent):44        return latent45 46class V_PREDICTION(EPS):47    def calculate_denoised(self, sigma, model_output, model_input):48        sigma = sigma.view(sigma.shape[:1] + (1,) * (model_output.ndim - 1))49        return model_input * self.sigma_data ** 2 / (sigma ** 2 + self.sigma_data ** 2) - model_output * sigma * self.sigma_data / (sigma ** 2 + self.sigma_data ** 2) ** 0.550 51class EDM(V_PREDICTION):52    def calculate_denoised(self, sigma, model_output, model_input):53        sigma = sigma.view(sigma.shape[:1] + (1,) * (model_output.ndim - 1))54        return model_input * self.sigma_data ** 2 / (sigma ** 2 + self.sigma_data ** 2) + model_output * sigma * self.sigma_data / (sigma ** 2 + self.sigma_data ** 2) ** 0.555 56class CONST:57    def calculate_input(self, sigma, noise):58        return noise59 60    def calculate_denoised(self, sigma, model_output, model_input):61        sigma = sigma.view(sigma.shape[:1] + (1,) * (model_output.ndim - 1))62        return model_input - model_output * sigma63 64    def noise_scaling(self, sigma, noise, latent_image, max_denoise=False):65        sigma = sigma.view(sigma.shape[:1] + (1,) * (noise.ndim - 1))66        return sigma * noise + (1.0 - sigma) * latent_image67 68    def inverse_noise_scaling(self, sigma, latent):69        sigma = sigma.view(sigma.shape[:1] + (1,) * (latent.ndim - 1))70        return latent / (1.0 - sigma)71 72class ModelSamplingDiscrete(torch.nn.Module):73    def __init__(self, model_config=None, zsnr=None):74        super().__init__()75 76        if model_config is not None:77            sampling_settings = model_config.sampling_settings78        else:79            sampling_settings = {}80 81        beta_schedule = sampling_settings.get("beta_schedule", "linear")82        linear_start = sampling_settings.get("linear_start", 0.00085)83        linear_end = sampling_settings.get("linear_end", 0.012)84        timesteps = sampling_settings.get("timesteps", 1000)85 86        if zsnr is None:87            zsnr = sampling_settings.get("zsnr", False)88 89        self._register_schedule(given_betas=None, beta_schedule=beta_schedule, timesteps=timesteps, linear_start=linear_start, linear_end=linear_end, cosine_s=8e-3, zsnr=zsnr)90        self.sigma_data = 1.091 92    def _register_schedule(self, given_betas=None, beta_schedule="linear", timesteps=1000,93                          linear_start=1e-4, linear_end=2e-2, cosine_s=8e-3, zsnr=False):94        if given_betas is not None:95            betas = given_betas96        else:97            betas = make_beta_schedule(beta_schedule, timesteps, linear_start=linear_start, linear_end=linear_end, cosine_s=cosine_s)98        alphas = 1. - betas99        alphas_cumprod = torch.cumprod(alphas, dim=0)100 101        timesteps, = betas.shape102        self.num_timesteps = int(timesteps)103        self.linear_start = linear_start104        self.linear_end = linear_end105 106        # self.register_buffer('betas', torch.tensor(betas, dtype=torch.float32))107        # self.register_buffer('alphas_cumprod', torch.tensor(alphas_cumprod, dtype=torch.float32))108        # self.register_buffer('alphas_cumprod_prev', torch.tensor(alphas_cumprod_prev, dtype=torch.float32))109 110        sigmas = ((1 - alphas_cumprod) / alphas_cumprod) ** 0.5111        if zsnr:112            sigmas = rescale_zero_terminal_snr_sigmas(sigmas)113 114        self.set_sigmas(sigmas)115 116    def set_sigmas(self, sigmas):117        self.register_buffer('sigmas', sigmas.float())118        self.register_buffer('log_sigmas', sigmas.log().float())119 120    @property121    def sigma_min(self):122        return self.sigmas[0]123 124    @property125    def sigma_max(self):126        return self.sigmas[-1]127 128    def timestep(self, sigma):129        log_sigma = sigma.log()130        dists = log_sigma.to(self.log_sigmas.device) - self.log_sigmas[:, None]131        return dists.abs().argmin(dim=0).view(sigma.shape).to(sigma.device)132 133    def sigma(self, timestep):134        t = torch.clamp(timestep.float().to(self.log_sigmas.device), min=0, max=(len(self.sigmas) - 1))135        low_idx = t.floor().long()136        high_idx = t.ceil().long()137        w = t.frac()138        log_sigma = (1 - w) * self.log_sigmas[low_idx] + w * self.log_sigmas[high_idx]139        return log_sigma.exp().to(timestep.device)140 141    def percent_to_sigma(self, percent):142        if percent <= 0.0:143            return 999999999.9144        if percent >= 1.0:145            return 0.0146        percent = 1.0 - percent147        return self.sigma(torch.tensor(percent * 999.0)).item()148 149class ModelSamplingDiscreteEDM(ModelSamplingDiscrete):150    def timestep(self, sigma):151        return 0.25 * sigma.log()152 153    def sigma(self, timestep):154        return (timestep / 0.25).exp()155 156class ModelSamplingContinuousEDM(torch.nn.Module):157    def __init__(self, model_config=None):158        super().__init__()159        if model_config is not None:160            sampling_settings = model_config.sampling_settings161        else:162            sampling_settings = {}163 164        sigma_min = sampling_settings.get("sigma_min", 0.002)165        sigma_max = sampling_settings.get("sigma_max", 120.0)166        sigma_data = sampling_settings.get("sigma_data", 1.0)167        self.set_parameters(sigma_min, sigma_max, sigma_data)168 169    def set_parameters(self, sigma_min, sigma_max, sigma_data):170        self.sigma_data = sigma_data171        sigmas = torch.linspace(math.log(sigma_min), math.log(sigma_max), 1000).exp()172 173        self.register_buffer('sigmas', sigmas) #for compatibility with some schedulers174        self.register_buffer('log_sigmas', sigmas.log())175 176    @property177    def sigma_min(self):178        return self.sigmas[0]179 180    @property181    def sigma_max(self):182        return self.sigmas[-1]183 184    def timestep(self, sigma):185        return 0.25 * sigma.log()186 187    def sigma(self, timestep):188        return (timestep / 0.25).exp()189 190    def percent_to_sigma(self, percent):191        if percent <= 0.0:192            return 999999999.9193        if percent >= 1.0:194            return 0.0195        percent = 1.0 - percent196 197        log_sigma_min = math.log(self.sigma_min)198        return math.exp((math.log(self.sigma_max) - log_sigma_min) * percent + log_sigma_min)199 200 201class ModelSamplingContinuousV(ModelSamplingContinuousEDM):202    def timestep(self, sigma):203        return sigma.atan() / math.pi * 2204 205    def sigma(self, timestep):206        return (timestep * math.pi / 2).tan()207 208 209def time_snr_shift(alpha, t):210    if alpha == 1.0:211        return t212    return alpha * t / (1 + (alpha - 1) * t)213 214class ModelSamplingDiscreteFlow(torch.nn.Module):215    def __init__(self, model_config=None):216        super().__init__()217        if model_config is not None:218            sampling_settings = model_config.sampling_settings219        else:220            sampling_settings = {}221 222        self.set_parameters(shift=sampling_settings.get("shift", 1.0), multiplier=sampling_settings.get("multiplier", 1000))223 224    def set_parameters(self, shift=1.0, timesteps=1000, multiplier=1000):225        self.shift = shift226        self.multiplier = multiplier227        ts = self.sigma((torch.arange(1, timesteps + 1, 1) / timesteps) * multiplier)228        self.register_buffer('sigmas', ts)229 230    @property231    def sigma_min(self):232        return self.sigmas[0]233 234    @property235    def sigma_max(self):236        return self.sigmas[-1]237 238    def timestep(self, sigma):239        return sigma * self.multiplier240 241    def sigma(self, timestep):242        return time_snr_shift(self.shift, timestep / self.multiplier)243 244    def percent_to_sigma(self, percent):245        if percent <= 0.0:246            return 1.0247        if percent >= 1.0:248            return 0.0249        return time_snr_shift(self.shift, 1.0 - percent)250 251class StableCascadeSampling(ModelSamplingDiscrete):252    def __init__(self, model_config=None):253        super().__init__()254 255        if model_config is not None:256            sampling_settings = model_config.sampling_settings257        else:258            sampling_settings = {}259 260        self.set_parameters(sampling_settings.get("shift", 1.0))261 262    def set_parameters(self, shift=1.0, cosine_s=8e-3):263        self.shift = shift264        self.cosine_s = torch.tensor(cosine_s)265        self._init_alpha_cumprod = torch.cos(self.cosine_s / (1 + self.cosine_s) * torch.pi * 0.5) ** 2266 267        #This part is just for compatibility with some schedulers in the codebase268        self.num_timesteps = 10000269        sigmas = torch.empty((self.num_timesteps), dtype=torch.float32)270        for x in range(self.num_timesteps):271            t = (x + 1) / self.num_timesteps272            sigmas[x] = self.sigma(t)273 274        self.set_sigmas(sigmas)275 276    def sigma(self, timestep):277        alpha_cumprod = (torch.cos((timestep + self.cosine_s) / (1 + self.cosine_s) * torch.pi * 0.5) ** 2 / self._init_alpha_cumprod)278 279        if self.shift != 1.0:280            var = alpha_cumprod281            logSNR = (var/(1-var)).log()282            logSNR += 2 * torch.log(1.0 / torch.tensor(self.shift))283            alpha_cumprod = logSNR.sigmoid()284 285        alpha_cumprod = alpha_cumprod.clamp(0.0001, 0.9999)286        return ((1 - alpha_cumprod) / alpha_cumprod) ** 0.5287 288    def timestep(self, sigma):289        var = 1 / ((sigma * sigma) + 1)290        var = var.clamp(0, 1.0)291        s, min_var = self.cosine_s.to(var.device), self._init_alpha_cumprod.to(var.device)292        t = (((var * min_var) ** 0.5).acos() / (torch.pi * 0.5)) * (1 + s) - s293        return t294 295    def percent_to_sigma(self, percent):296        if percent <= 0.0:297            return 999999999.9298        if percent >= 1.0:299            return 0.0300 301        percent = 1.0 - percent302        return self.sigma(torch.tensor(percent))303 304 305def flux_time_shift(mu: float, sigma: float, t):306    return math.exp(mu) / (math.exp(mu) + (1 / t - 1) ** sigma)307 308class ModelSamplingFlux(torch.nn.Module):309    def __init__(self, model_config=None):310        super().__init__()311        if model_config is not None:312            sampling_settings = model_config.sampling_settings313        else:314            sampling_settings = {}315 316        self.set_parameters(shift=sampling_settings.get("shift", 1.15))317 318    def set_parameters(self, shift=1.15, timesteps=10000):319        self.shift = shift320        ts = self.sigma((torch.arange(1, timesteps + 1, 1) / timesteps))321        self.register_buffer('sigmas', ts)322 323    @property324    def sigma_min(self):325        return self.sigmas[0]326 327    @property328    def sigma_max(self):329        return self.sigmas[-1]330 331    def timestep(self, sigma):332        return sigma333 334    def sigma(self, timestep):335        return flux_time_shift(self.shift, 1.0, timestep)336 337    def percent_to_sigma(self, percent):338        if percent <= 0.0:339            return 1.0340        if percent >= 1.0:341            return 0.0342        return flux_time_shift(self.shift, 1.0, 1.0 - percent)343