Spaces:
Runtime error
Runtime error
File size: 17,520 Bytes
c19ca42 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211 212 213 214 215 216 217 218 219 220 221 222 223 224 225 226 227 228 229 230 231 232 233 234 235 236 237 238 239 240 241 242 243 244 245 246 247 248 249 250 251 252 253 254 255 256 257 258 259 260 261 262 263 264 265 266 267 268 269 270 271 272 273 274 275 276 277 278 279 280 281 282 283 284 285 286 287 288 289 290 291 292 293 294 295 296 297 298 299 300 301 302 303 304 305 306 307 308 309 310 311 312 313 314 315 316 317 318 319 320 321 322 323 324 325 326 327 328 329 330 331 332 333 334 335 336 337 338 339 340 341 342 343 344 345 346 347 348 349 350 351 352 353 354 355 356 357 358 359 360 361 362 363 364 365 366 367 368 369 370 371 372 373 374 375 376 377 378 379 380 381 |
from typing import List, Union
import os
import time
from collections import namedtuple
import torch
import safetensors.torch
from PIL import Image
from modules import shared, devices, sd_models, errors
from modules.textual_inversion.image_embedding import embedding_from_b64, extract_image_data_embed
from modules.files_cache import directory_files, directory_mtime, extension_filter
debug = shared.log.trace if os.environ.get('SD_TI_DEBUG', None) is not None else lambda *args, **kwargs: None
debug('Trace: TEXTUAL INVERSION')
TokenToAdd = namedtuple("TokenToAdd", ["clip_l", "clip_g"])
TextualInversionTemplate = namedtuple("TextualInversionTemplate", ["name", "path"])
textual_inversion_templates = {}
def list_textual_inversion_templates():
textual_inversion_templates.clear()
for root, _dirs, fns in os.walk(shared.opts.embeddings_templates_dir):
for fn in fns:
path = os.path.join(root, fn)
textual_inversion_templates[fn] = TextualInversionTemplate(fn, path)
return textual_inversion_templates
def list_embeddings(*dirs):
is_ext = extension_filter(['.SAFETENSORS', '.PT' ] + ( ['.PNG', '.WEBP', '.JXL', '.AVIF', '.BIN' ] if shared.backend != shared.Backend.DIFFUSERS else [] ))
is_not_preview = lambda fp: not next(iter(os.path.splitext(fp))).upper().endswith('.PREVIEW') # pylint: disable=unnecessary-lambda-assignment
return list(filter(lambda fp: is_ext(fp) and is_not_preview(fp) and os.stat(fp).st_size > 0, directory_files(*dirs)))
class Embedding:
def __init__(self, vec, name, filename=None, step=None):
self.vec = vec
self.name = name
self.tag = name
self.step = step
self.filename = filename
self.basename = os.path.relpath(filename, shared.opts.embeddings_dir) if filename is not None else None
self.shape = None
self.vectors = 0
self.cached_checksum = None
self.sd_checkpoint = None
self.sd_checkpoint_name = None
self.optimizer_state_dict = None
def save(self, filename):
embedding_data = {
"string_to_token": {"*": 265},
"string_to_param": {"*": self.vec},
"name": self.name,
"step": self.step,
"sd_checkpoint": self.sd_checkpoint,
"sd_checkpoint_name": self.sd_checkpoint_name,
}
torch.save(embedding_data, filename)
if shared.opts.save_optimizer_state and self.optimizer_state_dict is not None:
optimizer_saved_dict = {
'hash': self.checksum(),
'optimizer_state_dict': self.optimizer_state_dict,
}
torch.save(optimizer_saved_dict, f"{filename}.optim")
def checksum(self):
if self.cached_checksum is not None:
return self.cached_checksum
def const_hash(a):
r = 0
for v in a:
r = (r * 281 ^ int(v) * 997) & 0xFFFFFFFF
return r
self.cached_checksum = f'{const_hash(self.vec.reshape(-1) * 100) & 0xffff:04x}'
return self.cached_checksum
class DirWithTextualInversionEmbeddings:
def __init__(self, path):
self.path = path
self.mtime = None
def has_changed(self):
if not os.path.isdir(self.path):
return False
return directory_mtime(self.path) != self.mtime
def update(self):
if not os.path.isdir(self.path):
return
self.mtime = directory_mtime(self.path)
def convert_embedding(tensor, text_encoder, text_encoder_2):
with torch.no_grad():
vectors = []
clip_l_embeds = text_encoder.get_input_embeddings().weight.data.clone().to(device=devices.device)
tensor = tensor.to(device=devices.device)
for vec in tensor:
values, indices = torch.max(torch.nan_to_num(torch.cosine_similarity(vec.unsqueeze(0), clip_l_embeds)), 0)
if values < 0.707: # Arbitrary similarity to cutoff, here 45 degrees
indices *= 0 # Use SDXL padding vector 0
vectors.append(indices)
vectors = torch.stack(vectors)
output = text_encoder_2.get_input_embeddings().weight.data[vectors]
return output
class EmbeddingDatabase:
def __init__(self):
self.ids_lookup = {}
self.word_embeddings = {}
self.skipped_embeddings = {}
self.expected_shape = -1
self.embedding_dirs = {}
self.previously_displayed_embeddings = ()
self.embeddings_used = []
def add_embedding_dir(self, path):
self.embedding_dirs[path] = DirWithTextualInversionEmbeddings(path)
def clear_embedding_dirs(self):
self.embedding_dirs.clear()
def register_embedding(self, embedding, model):
self.word_embeddings[embedding.name] = embedding
if hasattr(model, 'cond_stage_model'):
ids = model.cond_stage_model.tokenize([embedding.name])[0]
elif hasattr(model, 'tokenizer'):
ids = model.tokenizer.convert_tokens_to_ids(embedding.name)
if type(ids) != list:
ids = [ids]
first_id = ids[0]
if first_id not in self.ids_lookup:
self.ids_lookup[first_id] = []
self.ids_lookup[first_id] = sorted(self.ids_lookup[first_id] + [(ids, embedding)], key=lambda x: len(x[0]), reverse=True)
return embedding
def get_expected_shape(self):
if shared.backend == shared.Backend.DIFFUSERS:
return 0
if shared.sd_model is None:
shared.log.error('Model not loaded')
return 0
vec = shared.sd_model.cond_stage_model.encode_embedding_init_text(",", 1)
return vec.shape[1]
def load_diffusers_embedding(self, filename: Union[str, List[str]]):
_loaded_pre = len(self.word_embeddings)
embeddings_to_load = []
loaded_embeddings = {}
skipped_embeddings = []
if shared.sd_model is None:
return 0
tokenizer = getattr(shared.sd_model, 'tokenizer', None)
tokenizer_2 = getattr(shared.sd_model, 'tokenizer_2', None)
clip_l = getattr(shared.sd_model, 'text_encoder', None)
clip_g = getattr(shared.sd_model, 'text_encoder_2', None)
if clip_g and tokenizer_2:
model_type = 'SDXL'
elif clip_l and tokenizer:
model_type = 'SD'
else:
return 0
filenames = list(filename)
exts = [".SAFETENSORS", '.BIN', '.PT', '.PNG', '.WEBP', '.JXL', '.AVIF']
for filename in filenames:
# debug(f'Embedding check: {filename}')
fullname = filename
filename = os.path.basename(fullname)
fn, ext = os.path.splitext(filename)
name = os.path.basename(fn)
embedding = Embedding(vec=None, name=name, filename=fullname)
tokenizer_vocab = tokenizer.get_vocab()
try:
if ext.upper() not in exts:
raise ValueError(f'extension `{ext}` is invalid, expected one of: {exts}')
if name in tokenizer.get_vocab() or f"{name}_1" in tokenizer.get_vocab():
loaded_embeddings[name] = embedding
debug(f'Embedding already loaded: {name}')
embeddings_to_load.append(embedding)
except Exception as e:
skipped_embeddings.append(embedding)
debug(f'Embedding skipped: "{name}" {e}')
continue
embeddings_to_load = sorted(embeddings_to_load, key=lambda e: exts.index(os.path.splitext(e.filename)[1].upper()))
tokens_to_add = {}
for embedding in embeddings_to_load:
try:
if embedding.name in tokens_to_add or embedding.name in loaded_embeddings:
raise ValueError('duplicate token')
embeddings_dict = {}
_, ext = os.path.splitext(embedding.filename)
if ext.upper() in ['.SAFETENSORS']:
with safetensors.torch.safe_open(embedding.filename, framework="pt") as f: # type: ignore
for k in f.keys():
embeddings_dict[k] = f.get_tensor(k)
else: # fallback for sd1.5 pt embeddings
embeddings_dict["clip_l"] = self.load_from_file(embedding.filename, embedding.filename)
if 'clip_l' not in embeddings_dict:
raise ValueError('Invalid Embedding, dict missing required key `clip_l`')
if 'clip_g' not in embeddings_dict and model_type == "SDXL" and shared.opts.diffusers_convert_embed:
embeddings_dict["clip_g"] = convert_embedding(embeddings_dict["clip_l"], clip_l, clip_g)
if 'clip_g' in embeddings_dict:
embedding_type = 'SDXL'
else:
embedding_type = 'SD'
if embedding_type != model_type:
raise ValueError(f'Unable to load {embedding_type} Embedding "{embedding.name}" into {model_type} Model')
_tokens_to_add = {}
for i in range(len(embeddings_dict["clip_l"])):
if len(clip_l.get_input_embeddings().weight.data[0]) == len(embeddings_dict["clip_l"][i]):
token = embedding.name if i == 0 else f"{embedding.name}_{i}"
if token in tokenizer_vocab:
raise RuntimeError(f'Multi-Vector Embedding would add pre-existing Token in Vocabulary: {token}')
if token in tokens_to_add:
raise RuntimeError(f'Multi-Vector Embedding would add duplicate Token to Add: {token}')
_tokens_to_add[token] = TokenToAdd(
embeddings_dict["clip_l"][i],
embeddings_dict["clip_g"][i] if 'clip_g' in embeddings_dict else None
)
if not _tokens_to_add:
raise ValueError('no valid tokens to add')
tokens_to_add.update(_tokens_to_add)
loaded_embeddings[embedding.name] = embedding
except Exception as e:
debug(f"Embedding loading: {embedding.filename} {e}")
continue
if len(tokens_to_add) > 0:
tokenizer.add_tokens(list(tokens_to_add.keys()))
clip_l.resize_token_embeddings(len(tokenizer))
if model_type == 'SDXL':
tokenizer_2.add_tokens(list(tokens_to_add.keys())) # type: ignore
clip_g.resize_token_embeddings(len(tokenizer_2)) # type: ignore
unk_token_id = tokenizer.convert_tokens_to_ids(tokenizer.unk_token)
for token, data in tokens_to_add.items():
token_id = tokenizer.convert_tokens_to_ids(token)
if token_id > unk_token_id:
clip_l.get_input_embeddings().weight.data[token_id] = data.clip_l
if model_type == 'SDXL':
clip_g.get_input_embeddings().weight.data[token_id] = data.clip_g # type: ignore
for embedding in loaded_embeddings.values():
if not embedding:
continue
self.register_embedding(embedding, shared.sd_model)
if embedding in embeddings_to_load:
embeddings_to_load.remove(embedding)
skipped_embeddings.extend(embeddings_to_load)
for embedding in skipped_embeddings:
if loaded_embeddings.get(embedding.name, None) == embedding:
continue
self.skipped_embeddings[embedding.name] = embedding
try:
if model_type == 'SD':
debug(f"Embeddings loaded: text-encoder={shared.sd_model.text_encoder.get_input_embeddings().weight.data.shape[0]}")
if model_type == 'SDXL':
debug(f"Embeddings loaded: text-encoder-1={shared.sd_model.text_encoder.get_input_embeddings().weight.data.shape[0]} text-encoder-2={shared.sd_model.text_encoder_2.get_input_embeddings().weight.data.shape[0]}")
except Exception:
pass
return len(self.word_embeddings) - _loaded_pre
def load_from_file(self, path, filename):
name, ext = os.path.splitext(filename)
ext = ext.upper()
if ext in ['.PNG', '.WEBP', '.JXL', '.AVIF']:
if '.preview' in filename.lower():
return
embed_image = Image.open(path)
if hasattr(embed_image, 'text') and 'sd-ti-embedding' in embed_image.text:
data = embedding_from_b64(embed_image.text['sd-ti-embedding'])
else:
data = extract_image_data_embed(embed_image)
if not data: # if data is None, means this is not an embeding, just a preview image
return
elif ext in ['.BIN', '.PT']:
data = torch.load(path, map_location="cpu")
elif ext in ['.SAFETENSORS']:
data = safetensors.torch.load_file(path, device="cpu")
else:
return
# textual inversion embeddings
if 'string_to_param' in data:
param_dict = data['string_to_param']
param_dict = getattr(param_dict, '_parameters', param_dict) # fix for torch 1.12.1 loading saved file from torch 1.11
assert len(param_dict) == 1, 'embedding file has multiple terms in it'
emb = next(iter(param_dict.items()))[1]
# diffuser concepts
elif type(data) == dict and type(next(iter(data.values()))) == torch.Tensor:
if len(data.keys()) != 1:
self.skipped_embeddings[name] = Embedding(None, name=name, filename=path)
return
emb = next(iter(data.values()))
if len(emb.shape) == 1:
emb = emb.unsqueeze(0)
else:
raise RuntimeError(f"Couldn't identify {filename} as textual inversion embedding")
if shared.backend == shared.Backend.DIFFUSERS:
return emb
vec = emb.detach().to(devices.device, dtype=torch.float32)
# name = data.get('name', name)
embedding = Embedding(vec=vec, name=name, filename=path)
embedding.tag = data.get('name', None)
embedding.step = data.get('step', None)
embedding.sd_checkpoint = data.get('sd_checkpoint', None)
embedding.sd_checkpoint_name = data.get('sd_checkpoint_name', None)
embedding.vectors = vec.shape[0]
embedding.shape = vec.shape[-1]
if self.expected_shape == -1 or self.expected_shape == embedding.shape:
self.register_embedding(embedding, shared.sd_model)
else:
self.skipped_embeddings[name] = embedding
def load_from_dir(self, embdir):
if sd_models.model_data.sd_model is None:
shared.log.info('Skipping embeddings load: model not loaded')
return
if not os.path.isdir(embdir.path):
return
file_paths = list_embeddings(embdir.path)
if shared.backend == shared.Backend.DIFFUSERS:
self.load_diffusers_embedding(file_paths)
else:
for file_path in file_paths:
try:
fn = os.path.basename(file_path)
self.load_from_file(file_path, fn)
except Exception as e:
errors.display(e, f'Load embeding={fn}')
continue
def load_textual_inversion_embeddings(self, force_reload=False):
if shared.sd_model is None:
return
t0 = time.time()
if not force_reload:
need_reload = False
for embdir in self.embedding_dirs.values():
if embdir.has_changed():
need_reload = True
break
if not need_reload:
return
self.ids_lookup.clear()
self.word_embeddings.clear()
self.skipped_embeddings.clear()
self.embeddings_used.clear()
self.expected_shape = self.get_expected_shape()
for embdir in self.embedding_dirs.values():
self.load_from_dir(embdir)
embdir.update()
# re-sort word_embeddings because load_from_dir may not load in alphabetic order.
# using a temporary copy so we don't reinitialize self.word_embeddings in case other objects have a reference to it.
sorted_word_embeddings = {e.name: e for e in sorted(self.word_embeddings.values(), key=lambda e: e.name.lower())}
self.word_embeddings.clear()
self.word_embeddings.update(sorted_word_embeddings)
displayed_embeddings = (tuple(self.word_embeddings.keys()), tuple(self.skipped_embeddings.keys()))
if self.previously_displayed_embeddings != displayed_embeddings:
self.previously_displayed_embeddings = displayed_embeddings
t1 = time.time()
shared.log.info(f"Load embeddings: loaded={len(self.word_embeddings)} skipped={len(self.skipped_embeddings)} time={t1-t0:.2f}")
def find_embedding_at_position(self, tokens, offset):
token = tokens[offset]
possible_matches = self.ids_lookup.get(token, None)
if possible_matches is None:
return None, None
for ids, embedding in possible_matches:
if tokens[offset:offset + len(ids)] == ids:
return embedding, len(ids)
return None, None
|