Spaces:
Runtime error
Runtime error
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 | |