xinyun/ComfyUI
1
0
Fork 0
mirror of https://git.datalinker.icu/comfyanonymous/ComfyUI synced 2026-09-13 10:57:06 +08:00
Code Issues Packages Projects Releases Wiki Activity

Automatically detect and set zero terminal SNR

Browse Source
This commit is contained in:
catboxanon 2024-11-10 14:14:04 -05:00
parent bdeb1c171c
commit 62d13fb2f3
5 changed files with 30 additions and 19 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

6
comfy/sd.py
View File

@ -682,6 +682,12 @@ def load_state_dict_guess_config(sd, output_vae=True, output_clip=True, output_c
if output_model:
model_patcher = comfy.model_patcher.ModelPatcher(model, load_device=load_device, offload_device=model_management.unet_offload_device())
if model_config.ztsnr:
class ModelSamplingAdvanced(comfy.model_sampling.ModelSamplingDiscrete, comfy.model_sampling.V_PREDICTION):
pass
model_sampling = ModelSamplingAdvanced(model.model_config)
model_sampling.set_sigmas(comfy.utils.rescale_zero_terminal_snr_sigmas(model_sampling.sigmas))
model_patcher.add_object_patch("model_sampling", model_sampling)
if inital_load_device != torch.device("cpu"):
logging.info("loaded straight to GPU")
model_management.load_models_gpu([model_patcher], force_full_load=True)

2
comfy/supported_models.py
View File

@ -197,6 +197,8 @@ class SDXL(supported_models_base.BASE):
self.sampling_settings["sigma_min"] = float(state_dict["edm_vpred.sigma_min"].item())
return model_base.ModelType.V_PREDICTION_EDM
elif "v_pred" in state_dict:
if "ztsnr" in state_dict:
self.ztsnr = True
return model_base.ModelType.V_PREDICTION
else:
return model_base.ModelType.EPS

1
comfy/supported_models_base.py
View File

@ -40,6 +40,7 @@ class BASE:
clip_vision_prefix = None
noise_aug_config = None
sampling_settings = {}
ztsnr = False
latent_format = latent_formats.LatentFormat
vae_key_prefix = ["first_stage_model."]
text_encoder_key_prefix = ["cond_stage_model."]

19
comfy/utils.py
View File

@ -869,3 +869,22 @@ def reshape_mask(input_mask, output_shape):
mask = mask.repeat((1, output_shape[1]) + (1,) * dims)[:,:output_shape[1]]
mask = comfy.utils.repeat_to_batch_size(mask, output_shape[0])
return mask
def rescale_zero_terminal_snr_sigmas(sigmas):
alphas_cumprod = 1 / ((sigmas * sigmas) + 1)
alphas_bar_sqrt = alphas_cumprod.sqrt()
# Store old values.
alphas_bar_sqrt_0 = alphas_bar_sqrt[0].clone()
alphas_bar_sqrt_T = alphas_bar_sqrt[-1].clone()
# Shift so the last timestep is zero.
alphas_bar_sqrt -= (alphas_bar_sqrt_T)
# Scale so the first timestep is back to the old value.
alphas_bar_sqrt *= alphas_bar_sqrt_0 / (alphas_bar_sqrt_0 - alphas_bar_sqrt_T)
# Convert alphas_bar_sqrt to betas
alphas_bar = alphas_bar_sqrt**2 # Revert sqrt
alphas_bar[-1] = 4.8973451890853435e-08
return ((1 - alphas_bar) / alphas_bar) ** 0.5

21
comfy_extras/nodes_model_advanced.py
View File

@ -5,6 +5,8 @@ import comfy.latent_formats
import nodes
import torch
from comfy.utils import rescale_zero_terminal_snr_sigmas
class LCM(comfy.model_sampling.EPS):
def calculate_denoised(self, sigma, model_output, model_input):
timestep = self.timestep(sigma).view(sigma.shape[:1] + (1,) * (model_output.ndim - 1))
@ -51,25 +53,6 @@ class ModelSamplingDiscreteDistilled(comfy.model_sampling.ModelSamplingDiscrete)
return log_sigma.exp().to(timestep.device)
def rescale_zero_terminal_snr_sigmas(sigmas):
alphas_cumprod = 1 / ((sigmas * sigmas) + 1)
alphas_bar_sqrt = alphas_cumprod.sqrt()
# Store old values.
alphas_bar_sqrt_0 = alphas_bar_sqrt[0].clone()
alphas_bar_sqrt_T = alphas_bar_sqrt[-1].clone()
# Shift so the last timestep is zero.
alphas_bar_sqrt -= (alphas_bar_sqrt_T)
# Scale so the first timestep is back to the old value.
alphas_bar_sqrt *= alphas_bar_sqrt_0 / (alphas_bar_sqrt_0 - alphas_bar_sqrt_T)
# Convert alphas_bar_sqrt to betas
alphas_bar = alphas_bar_sqrt**2 # Revert sqrt
alphas_bar[-1] = 4.8973451890853435e-08
return ((1 - alphas_bar) / alphas_bar) ** 0.5
class ModelSamplingDiscrete:
@classmethod
def INPUT_TYPES(s):