Upload sd_models.py
Browse files- sd_models.py +311 -0
sd_models.py
ADDED
@@ -0,0 +1,311 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
1 |
+
import collections
|
2 |
+
import os.path
|
3 |
+
import sys
|
4 |
+
import gc
|
5 |
+
from collections import namedtuple
|
6 |
+
import torch
|
7 |
+
import re
|
8 |
+
import safetensors.torch
|
9 |
+
from omegaconf import OmegaConf
|
10 |
+
|
11 |
+
from ldm.util import instantiate_from_config
|
12 |
+
|
13 |
+
from modules import shared, modelloader, devices, script_callbacks, sd_vae
|
14 |
+
from modules.paths import models_path
|
15 |
+
from modules.sd_hijack_inpainting import do_inpainting_hijack, should_hijack_inpainting
|
16 |
+
|
17 |
+
model_dir = "Stable-diffusion"
|
18 |
+
model_path = os.path.abspath(os.path.join(models_path, model_dir))
|
19 |
+
|
20 |
+
CheckpointInfo = namedtuple("CheckpointInfo", ['filename', 'title', 'hash', 'model_name', 'config'])
|
21 |
+
checkpoints_list = {}
|
22 |
+
checkpoints_loaded = collections.OrderedDict()
|
23 |
+
|
24 |
+
try:
|
25 |
+
# this silences the annoying "Some weights of the model checkpoint were not used when initializing..." message at start.
|
26 |
+
|
27 |
+
from transformers import logging, CLIPModel
|
28 |
+
|
29 |
+
logging.set_verbosity_error()
|
30 |
+
except Exception:
|
31 |
+
pass
|
32 |
+
|
33 |
+
|
34 |
+
def setup_model():
|
35 |
+
if not os.path.exists(model_path):
|
36 |
+
os.makedirs(model_path)
|
37 |
+
|
38 |
+
list_models()
|
39 |
+
|
40 |
+
|
41 |
+
def checkpoint_tiles():
|
42 |
+
convert = lambda name: int(name) if name.isdigit() else name.lower()
|
43 |
+
alphanumeric_key = lambda key: [convert(c) for c in re.split('([0-9]+)', key)]
|
44 |
+
return sorted([x.title for x in checkpoints_list.values()], key = alphanumeric_key)
|
45 |
+
|
46 |
+
|
47 |
+
def list_models():
|
48 |
+
checkpoints_list.clear()
|
49 |
+
model_list = modelloader.load_models(model_path=model_path, command_path=shared.cmd_opts.ckpt_dir, ext_filter=[".ckpt", ".safetensors"])
|
50 |
+
|
51 |
+
def modeltitle(path, shorthash):
|
52 |
+
abspath = os.path.abspath(path)
|
53 |
+
|
54 |
+
if shared.cmd_opts.ckpt_dir is not None and abspath.startswith(shared.cmd_opts.ckpt_dir):
|
55 |
+
name = abspath.replace(shared.cmd_opts.ckpt_dir, '')
|
56 |
+
elif abspath.startswith(model_path):
|
57 |
+
name = abspath.replace(model_path, '')
|
58 |
+
else:
|
59 |
+
name = os.path.basename(path)
|
60 |
+
|
61 |
+
if name.startswith("\\") or name.startswith("/"):
|
62 |
+
name = name[1:]
|
63 |
+
|
64 |
+
shortname = os.path.splitext(name.replace("/", "_").replace("\\", "_"))[0]
|
65 |
+
|
66 |
+
return f'{name} [{shorthash}]', shortname
|
67 |
+
|
68 |
+
cmd_ckpt = shared.cmd_opts.ckpt
|
69 |
+
if os.path.exists(cmd_ckpt):
|
70 |
+
h = model_hash(cmd_ckpt)
|
71 |
+
title, short_model_name = modeltitle(cmd_ckpt, h)
|
72 |
+
checkpoints_list[title] = CheckpointInfo(cmd_ckpt, title, h, short_model_name, shared.cmd_opts.config)
|
73 |
+
shared.opts.data['sd_model_checkpoint'] = title
|
74 |
+
elif cmd_ckpt is not None and cmd_ckpt != shared.default_sd_model_file:
|
75 |
+
print(f"Checkpoint in --ckpt argument not found (Possible it was moved to {model_path}: {cmd_ckpt}", file=sys.stderr)
|
76 |
+
for filename in model_list:
|
77 |
+
h = model_hash(filename)
|
78 |
+
title, short_model_name = modeltitle(filename, h)
|
79 |
+
|
80 |
+
basename, _ = os.path.splitext(filename)
|
81 |
+
config = basename + ".yaml"
|
82 |
+
if not os.path.exists(config):
|
83 |
+
config = shared.cmd_opts.config
|
84 |
+
|
85 |
+
checkpoints_list[title] = CheckpointInfo(filename, title, h, short_model_name, config)
|
86 |
+
|
87 |
+
|
88 |
+
def get_closet_checkpoint_match(searchString):
|
89 |
+
applicable = sorted([info for info in checkpoints_list.values() if searchString in info.title], key = lambda x:len(x.title))
|
90 |
+
if len(applicable) > 0:
|
91 |
+
return applicable[0]
|
92 |
+
return None
|
93 |
+
|
94 |
+
|
95 |
+
def model_hash(filename):
|
96 |
+
try:
|
97 |
+
with open(filename, "rb") as file:
|
98 |
+
import hashlib
|
99 |
+
m = hashlib.sha256()
|
100 |
+
|
101 |
+
file.seek(0x100000)
|
102 |
+
m.update(file.read(0x10000))
|
103 |
+
return m.hexdigest()[0:8]
|
104 |
+
except FileNotFoundError:
|
105 |
+
return 'NOFILE'
|
106 |
+
|
107 |
+
|
108 |
+
def select_checkpoint():
|
109 |
+
model_checkpoint = shared.opts.sd_model_checkpoint
|
110 |
+
checkpoint_info = checkpoints_list.get(model_checkpoint, None)
|
111 |
+
if checkpoint_info is not None:
|
112 |
+
return checkpoint_info
|
113 |
+
|
114 |
+
if len(checkpoints_list) == 0:
|
115 |
+
print(f"No checkpoints found. When searching for checkpoints, looked at:", file=sys.stderr)
|
116 |
+
if shared.cmd_opts.ckpt is not None:
|
117 |
+
print(f" - file {os.path.abspath(shared.cmd_opts.ckpt)}", file=sys.stderr)
|
118 |
+
print(f" - directory {model_path}", file=sys.stderr)
|
119 |
+
if shared.cmd_opts.ckpt_dir is not None:
|
120 |
+
print(f" - directory {os.path.abspath(shared.cmd_opts.ckpt_dir)}", file=sys.stderr)
|
121 |
+
print(f"Can't run without a checkpoint. Find and place a .ckpt file into any of those locations. The program will exit.", file=sys.stderr)
|
122 |
+
exit(1)
|
123 |
+
|
124 |
+
checkpoint_info = next(iter(checkpoints_list.values()))
|
125 |
+
if model_checkpoint is not None:
|
126 |
+
print(f"Checkpoint {model_checkpoint} not found; loading fallback {checkpoint_info.title}", file=sys.stderr)
|
127 |
+
|
128 |
+
return checkpoint_info
|
129 |
+
|
130 |
+
|
131 |
+
chckpoint_dict_replacements = {
|
132 |
+
'cond_stage_model.transformer.embeddings.': 'cond_stage_model.transformer.text_model.embeddings.',
|
133 |
+
'cond_stage_model.transformer.encoder.': 'cond_stage_model.transformer.text_model.encoder.',
|
134 |
+
'cond_stage_model.transformer.final_layer_norm.': 'cond_stage_model.transformer.text_model.final_layer_norm.',
|
135 |
+
}
|
136 |
+
|
137 |
+
|
138 |
+
def transform_checkpoint_dict_key(k):
|
139 |
+
for text, replacement in chckpoint_dict_replacements.items():
|
140 |
+
if k.startswith(text):
|
141 |
+
k = replacement + k[len(text):]
|
142 |
+
|
143 |
+
return k
|
144 |
+
|
145 |
+
|
146 |
+
def get_state_dict_from_checkpoint(pl_sd):
|
147 |
+
pl_sd = pl_sd.pop("state_dict", pl_sd)
|
148 |
+
pl_sd.pop("state_dict", None)
|
149 |
+
|
150 |
+
sd = {}
|
151 |
+
for k, v in pl_sd.items():
|
152 |
+
new_key = transform_checkpoint_dict_key(k)
|
153 |
+
|
154 |
+
if new_key is not None:
|
155 |
+
sd[new_key] = v
|
156 |
+
|
157 |
+
pl_sd.clear()
|
158 |
+
pl_sd.update(sd)
|
159 |
+
|
160 |
+
return pl_sd
|
161 |
+
|
162 |
+
|
163 |
+
def read_state_dict(checkpoint_file, print_global_state=False, map_location=None):
|
164 |
+
_, extension = os.path.splitext(checkpoint_file)
|
165 |
+
if extension.lower() == ".safetensors":
|
166 |
+
pl_sd = safetensors.torch.load_file(checkpoint_file, device=map_location or shared.weight_load_location)
|
167 |
+
else:
|
168 |
+
pl_sd = torch.load(checkpoint_file, map_location=map_location or shared.weight_load_location)
|
169 |
+
|
170 |
+
if print_global_state and "global_step" in pl_sd:
|
171 |
+
print(f"Global Step: {pl_sd['global_step']}")
|
172 |
+
|
173 |
+
sd = get_state_dict_from_checkpoint(pl_sd)
|
174 |
+
return sd
|
175 |
+
|
176 |
+
|
177 |
+
def load_model_weights(model, checkpoint_info, vae_file="auto"):
|
178 |
+
checkpoint_file = checkpoint_info.filename
|
179 |
+
sd_model_hash = checkpoint_info.hash
|
180 |
+
|
181 |
+
cache_enabled = shared.opts.sd_checkpoint_cache > 0
|
182 |
+
|
183 |
+
if cache_enabled and checkpoint_info in checkpoints_loaded:
|
184 |
+
# use checkpoint cache
|
185 |
+
print(f"Loading weights [{sd_model_hash}] from cache")
|
186 |
+
model.load_state_dict(checkpoints_loaded[checkpoint_info])
|
187 |
+
else:
|
188 |
+
# load from file
|
189 |
+
print(f"Loading weights [{sd_model_hash}] from {checkpoint_file}")
|
190 |
+
|
191 |
+
sd = read_state_dict(checkpoint_file)
|
192 |
+
model.load_state_dict(sd, strict=False)
|
193 |
+
del sd
|
194 |
+
|
195 |
+
if cache_enabled:
|
196 |
+
# cache newly loaded model
|
197 |
+
checkpoints_loaded[checkpoint_info] = model.state_dict().copy()
|
198 |
+
|
199 |
+
if shared.cmd_opts.opt_channelslast:
|
200 |
+
model.to(memory_format=torch.channels_last)
|
201 |
+
|
202 |
+
if not shared.cmd_opts.no_half:
|
203 |
+
vae = model.first_stage_model
|
204 |
+
|
205 |
+
# with --no-half-vae, remove VAE from model when doing half() to prevent its weights from being converted to float16
|
206 |
+
if shared.cmd_opts.no_half_vae:
|
207 |
+
model.first_stage_model = None
|
208 |
+
|
209 |
+
model.half()
|
210 |
+
model.first_stage_model = vae
|
211 |
+
|
212 |
+
devices.dtype = torch.float32 if shared.cmd_opts.no_half else torch.float16
|
213 |
+
devices.dtype_vae = torch.float32 if shared.cmd_opts.no_half or shared.cmd_opts.no_half_vae else torch.float16
|
214 |
+
|
215 |
+
model.first_stage_model.to(devices.dtype_vae)
|
216 |
+
|
217 |
+
# clean up cache if limit is reached
|
218 |
+
if cache_enabled:
|
219 |
+
while len(checkpoints_loaded) > shared.opts.sd_checkpoint_cache + 1: # we need to count the current model
|
220 |
+
checkpoints_loaded.popitem(last=False) # LRU
|
221 |
+
|
222 |
+
model.sd_model_hash = sd_model_hash
|
223 |
+
model.sd_model_checkpoint = checkpoint_file
|
224 |
+
model.sd_checkpoint_info = checkpoint_info
|
225 |
+
|
226 |
+
vae_file = sd_vae.resolve_vae(checkpoint_file, vae_file=vae_file)
|
227 |
+
sd_vae.load_vae(model, vae_file)
|
228 |
+
|
229 |
+
|
230 |
+
def load_model(checkpoint_info=None):
|
231 |
+
from modules import lowvram, sd_hijack
|
232 |
+
checkpoint_info = checkpoint_info or select_checkpoint()
|
233 |
+
|
234 |
+
if checkpoint_info.config != shared.cmd_opts.config:
|
235 |
+
print(f"Loading config from: {checkpoint_info.config}")
|
236 |
+
|
237 |
+
if shared.sd_model:
|
238 |
+
sd_hijack.model_hijack.undo_hijack(shared.sd_model)
|
239 |
+
shared.sd_model = None
|
240 |
+
gc.collect()
|
241 |
+
devices.torch_gc()
|
242 |
+
|
243 |
+
sd_config = OmegaConf.load(checkpoint_info.config)
|
244 |
+
|
245 |
+
if should_hijack_inpainting(checkpoint_info):
|
246 |
+
# Hardcoded config for now...
|
247 |
+
sd_config.model.target = "ldm.models.diffusion.ddpm.LatentInpaintDiffusion"
|
248 |
+
sd_config.model.params.use_ema = False
|
249 |
+
sd_config.model.params.conditioning_key = "hybrid"
|
250 |
+
sd_config.model.params.unet_config.params.in_channels = 9
|
251 |
+
|
252 |
+
# Create a "fake" config with a different name so that we know to unload it when switching models.
|
253 |
+
checkpoint_info = checkpoint_info._replace(config=checkpoint_info.config.replace(".yaml", "-inpainting.yaml"))
|
254 |
+
|
255 |
+
do_inpainting_hijack()
|
256 |
+
|
257 |
+
if shared.cmd_opts.no_half:
|
258 |
+
sd_config.model.params.unet_config.params.use_fp16 = False
|
259 |
+
|
260 |
+
sd_model = instantiate_from_config(sd_config.model)
|
261 |
+
load_model_weights(sd_model, checkpoint_info)
|
262 |
+
|
263 |
+
if shared.cmd_opts.lowvram or shared.cmd_opts.medvram:
|
264 |
+
lowvram.setup_for_low_vram(sd_model, shared.cmd_opts.medvram)
|
265 |
+
else:
|
266 |
+
sd_model.to(shared.device)
|
267 |
+
|
268 |
+
sd_hijack.model_hijack.hijack(sd_model)
|
269 |
+
|
270 |
+
sd_model.eval()
|
271 |
+
shared.sd_model = sd_model
|
272 |
+
|
273 |
+
script_callbacks.model_loaded_callback(sd_model)
|
274 |
+
|
275 |
+
print(f"Model loaded.")
|
276 |
+
return sd_model
|
277 |
+
|
278 |
+
|
279 |
+
def reload_model_weights(sd_model=None, info=None):
|
280 |
+
from modules import lowvram, devices, sd_hijack
|
281 |
+
checkpoint_info = info or select_checkpoint()
|
282 |
+
|
283 |
+
if not sd_model:
|
284 |
+
sd_model = shared.sd_model
|
285 |
+
|
286 |
+
if sd_model.sd_model_checkpoint == checkpoint_info.filename:
|
287 |
+
return
|
288 |
+
|
289 |
+
if sd_model.sd_checkpoint_info.config != checkpoint_info.config or should_hijack_inpainting(checkpoint_info) != should_hijack_inpainting(sd_model.sd_checkpoint_info):
|
290 |
+
del sd_model
|
291 |
+
checkpoints_loaded.clear()
|
292 |
+
load_model(checkpoint_info)
|
293 |
+
return shared.sd_model
|
294 |
+
|
295 |
+
if shared.cmd_opts.lowvram or shared.cmd_opts.medvram:
|
296 |
+
lowvram.send_everything_to_cpu()
|
297 |
+
else:
|
298 |
+
sd_model.to(devices.cpu)
|
299 |
+
|
300 |
+
sd_hijack.model_hijack.undo_hijack(sd_model)
|
301 |
+
|
302 |
+
load_model_weights(sd_model, checkpoint_info)
|
303 |
+
|
304 |
+
sd_hijack.model_hijack.hijack(sd_model)
|
305 |
+
script_callbacks.model_loaded_callback(sd_model)
|
306 |
+
|
307 |
+
if not shared.cmd_opts.lowvram and not shared.cmd_opts.medvram:
|
308 |
+
sd_model.to(devices.device)
|
309 |
+
|
310 |
+
print(f"Weights loaded.")
|
311 |
+
return sd_model
|