Upload sd_models.py

#3
by korp123 - opened
Files changed (1) hide show
  1. 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