OPENGIT/stable-diffusion-webui
1
0
Fork 0
mirror of https://github.com/AUTOMATIC1111/stable-diffusion-webui.git synced 2024-12-29 19:05:05 +08:00
Code Issues Packages Projects Releases Wiki Activity

Merge pull request #16569 from AUTOMATIC1111/feat/ztsnr-auto

Browse Source 
Automatically enable ztSNR based on existence of key in `state_dict`
This commit is contained in:
w-e-w 2024-10-20 11:40:02 +09:00 committed by GitHub
commit 5865da28d1
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
1 changed files with 5 additions and 1 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

6
modules/sd_models.py
View File

@ -423,6 +423,10 @@ def load_model_weights(model, checkpoint_info: CheckpointInfo, state_dict, timer
set_model_type(model, state_dict)
set_model_fields(model)
if 'ztsnr' in state_dict:
model.ztsnr = True
else:
model.ztsnr = False
if model.is_sdxl:
sd_models_xl.extend_sdxl(model)
@ -661,7 +665,7 @@ def apply_alpha_schedule_override(sd_model, p=None):
p.extra_generation_params['Downcast alphas_cumprod'] = opts.use_downcasted_alpha_bar
sd_model.alphas_cumprod = sd_model.alphas_cumprod.half().to(shared.device)
if opts.sd_noise_schedule == "Zero Terminal SNR":
if opts.sd_noise_schedule == "Zero Terminal SNR" or (hasattr(sd_model, 'ztsnr') and sd_model.ztsnr):
if p is not None:
p.extra_generation_params['Noise Schedule'] = opts.sd_noise_schedule
sd_model.alphas_cumprod = rescale_zero_terminal_snr_abar(sd_model.alphas_cumprod).to(shared.device)