Skip to content

Commit d9d9414

Browse files
authored
patch DDPM.register_betas so that users can put given_betas in model yaml
1 parent 5ef669d commit d9d9414

File tree

1 file changed

+14
-1
lines changed

1 file changed

+14
-1
lines changed

modules/sd_models.py

+14-1
Original file line numberDiff line numberDiff line change
@@ -7,7 +7,7 @@
77
import torch
88
import re
99
import safetensors.torch
10-
from omegaconf import OmegaConf
10+
from omegaconf import OmegaConf, ListConfig
1111
from os import mkdir
1212
from urllib import request
1313
import ldm.modules.midas as midas
@@ -17,6 +17,7 @@
1717
from modules import paths, shared, modelloader, devices, script_callbacks, sd_vae, sd_disable_initialization, errors, hashes, sd_models_config, sd_unet, sd_models_xl, cache, extra_networks, processing, lowvram, sd_hijack
1818
from modules.timer import Timer
1919
import tomesd
20+
import numpy as np
2021

2122
model_dir = "Stable-diffusion"
2223
model_path = os.path.abspath(os.path.join(paths.models_path, model_dir))
@@ -132,6 +133,7 @@ def setup_model():
132133
os.makedirs(model_path, exist_ok=True)
133134

134135
enable_midas_autodownload()
136+
patch_given_betas()
135137

136138

137139
def checkpoint_tiles(use_short=False):
@@ -453,6 +455,17 @@ def load_model_wrapper(model_type):
453455
midas.api.load_model = load_model_wrapper
454456

455457

458+
def patch_given_betas():
459+
original_register_schedule = ldm.models.diffusion.ddpm.DDPM.register_schedule
460+
def patched_register_schedule(*args, **kwargs):
461+
if args[1] is not None and isinstance(args[1], ListConfig):
462+
modified_args = list(args) # Convert args tuple to a list
463+
modified_args[1] = np.array(args[1]) # Modify the desired element
464+
args = tuple(modified_args) # Convert the list back to a tuple
465+
original_register_schedule(*args, **kwargs)
466+
ldm.models.diffusion.ddpm.DDPM.register_schedule = patched_register_schedule
467+
468+
456469
def repair_config(sd_config):
457470

458471
if not hasattr(sd_config.model.params, "use_ema"):

0 commit comments

Comments
 (0)