From e34b19bb80d949bb29227654fea9b1c65d5a13e4 Mon Sep 17 00:00:00 2001 From: AI-Casanova <54461896+AI-Casanova@users.noreply.github.com> Date: Sun, 3 Sep 2023 21:49:43 -0500 Subject: [PATCH] Add Textual Inversion loading for SDXL --- modules/prompt_parser_diffusers.py | 88 ++++++++++++++++++- .../textual_inversion/textual_inversion.py | 23 ++++- 2 files changed, 107 insertions(+), 4 deletions(-) diff --git a/modules/prompt_parser_diffusers.py b/modules/prompt_parser_diffusers.py index 4bfda29ad..665b2ae31 100644 --- a/modules/prompt_parser_diffusers.py +++ b/modules/prompt_parser_diffusers.py @@ -2,8 +2,10 @@ import os import typing import torch from compel import Compel, ReturnedEmbeddingsType +from compel.embeddings_provider import BaseTextualInversionManager import modules.shared as shared import modules.prompt_parser as prompt_parser +from typing import Callable, Dict, List, Optional, Union debug_output = os.environ.get('SD_PROMPT_DEBUG', None) debug = shared.log.info if debug_output is not None else lambda *args, **kwargs: None @@ -32,6 +34,83 @@ CLIP_SKIP_MAPPING = { } + +#from https://github.com/damian0815/compel/blob/main/src/compel/diffusers_textual_inversion_manager.py +class DiffusersTextualInversionManager(BaseTextualInversionManager): + """ + A textual inversion manager for use with diffusers. + """ + def __init__(self, pipe): + self.pipe = pipe + + #from https://github.com/huggingface/diffusers/blob/705c592ea98ba4e288d837b9cba2767623c78603/src/diffusers/loaders.py#L599 + def maybe_convert_prompt(self, prompt: Union[str, List[str]], tokenizer: "PreTrainedTokenizer"): + r""" + Processes prompts that include a special token corresponding to a multi-vector textual inversion embedding to + be replaced with multiple special tokens each corresponding to one of the vectors. If the prompt has no textual + inversion token or if the textual inversion token is a single vector, the input prompt is returned. + + Parameters: + prompt (`str` or list of `str`): + The prompt or prompts to guide the image generation. + tokenizer (`PreTrainedTokenizer`): + The tokenizer responsible for encoding the prompt into input tokens. + + Returns: + `str` or list of `str`: The converted prompt + """ + if not isinstance(prompt, List): + prompts = [prompt] + else: + prompts = prompt + + prompts = [self._maybe_convert_prompt(p, tokenizer) for p in prompts] + + if not isinstance(prompt, List): + return prompts[0] + + return prompts + + def _maybe_convert_prompt(self, prompt: str, tokenizer: "PreTrainedTokenizer"): + r""" + Maybe convert a prompt into a "multi vector"-compatible prompt. If the prompt includes a token that corresponds + to a multi-vector textual inversion embedding, this function will process the prompt so that the special token + is replaced with multiple special tokens each corresponding to one of the vectors. If the prompt has no textual + inversion token or a textual inversion token that is a single vector, the input prompt is simply returned. + + Parameters: + prompt (`str`): + The prompt to guide the image generation. + tokenizer (`PreTrainedTokenizer`): + The tokenizer responsible for encoding the prompt into input tokens. + + Returns: + `str`: The converted prompt + """ + tokens = tokenizer.tokenize(prompt) + unique_tokens = set(tokens) + for token in unique_tokens: + if token in tokenizer.added_tokens_encoder: + replacement = token + i = 1 + while f"{token}_{i}" in tokenizer.added_tokens_encoder: + replacement += f" {token}_{i}" + i += 1 + + prompt = prompt.replace(token, replacement) + + return prompt + #end of Diffusers code + + def expand_textual_inversion_token_ids_if_necessary(self, token_ids: List[int]) -> List[int]: + if len(token_ids) == 0: + return token_ids + + prompt = self.pipe.tokenizer.decode(token_ids) + prompt = self.maybe_convert_prompt(prompt, self.pipe.tokenizer) + return self.pipe.tokenizer.encode(prompt, add_special_tokens=False) + #end of Compel code + def compel_encode_prompts( pipeline, prompts: list, @@ -100,17 +179,20 @@ def compel_encode_prompt( prompt_2 = convert_to_compel(prompt_2) negative_prompt_2 = convert_to_compel(negative_prompt_2) + textual_inversion_manager = DiffusersTextualInversionManager(pipeline) + compel_te1 = Compel( tokenizer=pipeline.tokenizer, text_encoder=pipeline.text_encoder, returned_embeddings_type=embedding_type, requires_pooled=False, # truncate_long_prompts=False, - device=shared.device + device=shared.device, + textual_inversion_manager=textual_inversion_manager ) if not is_refiner and shared.sd_model_type == "sdxl": - compel_te2 = Compel(tokenizer=pipeline.tokenizer_2, text_encoder=pipeline.text_encoder_2, returned_embeddings_type=embedding_type, requires_pooled=True, device=shared.device) + compel_te2 = Compel(tokenizer=pipeline.tokenizer_2, text_encoder=pipeline.text_encoder_2, returned_embeddings_type=embedding_type, requires_pooled=True, device=shared.device, textual_inversion_manager=textual_inversion_manager) positive_te1 = compel_te1(prompt) positive_te2, positive_pooled = compel_te2(prompt_2) positive = torch.cat((positive_te1, positive_te2), dim=-1) @@ -124,7 +206,7 @@ def compel_encode_prompt( return prompt_embed, positive_pooled, negative_embed, negative_pooled if is_refiner and shared.sd_refiner_type == "sdxl": - compel_te2 = Compel(tokenizer=pipeline.tokenizer_2, text_encoder=pipeline.text_encoder_2, returned_embeddings_type=embedding_type, requires_pooled=True, device=shared.device) + compel_te2 = Compel(tokenizer=pipeline.tokenizer_2, text_encoder=pipeline.text_encoder_2, returned_embeddings_type=embedding_type, requires_pooled=True, device=shared.device, textual_inversion_manager=textual_inversion_manager) positive, positive_pooled = compel_te2(prompt) negative, negative_pooled = compel_te2(negative_prompt) diff --git a/modules/textual_inversion/textual_inversion.py b/modules/textual_inversion/textual_inversion.py index b01fe41e0..803b3163a 100644 --- a/modules/textual_inversion/textual_inversion.py +++ b/modules/textual_inversion/textual_inversion.py @@ -140,7 +140,28 @@ class EmbeddingDatabase: name = os.path.basename(fn) embedding = Embedding(vec=None, name=name) try: - pipe.load_textual_inversion(path, cache_dir=shared.opts.diffusers_dir, local_files_only=True) + if hasattr(pipe,"load_textual_inversion"): + pipe.load_textual_inversion(path, cache_dir=shared.opts.diffusers_dir, local_files_only=True) + elif "safetensors" in path: + embeddings_dict = {} + from safetensors.torch import safe_open + + with safe_open(path, framework="pt") as f: + for k in f.keys(): + embeddings_dict[k] = f.get_tensor(k) + for i in range(len(embeddings_dict["clip_l"])): + if i == 0: + token = name + else: + token = f"{name}_{i}" + pipe.tokenizer.add_tokens(token) + token_id = pipe.tokenizer.convert_tokens_to_ids(token) + pipe.text_encoder.resize_token_embeddings(len(pipe.tokenizer)) + pipe.text_encoder_2.resize_token_embeddings(len(pipe.tokenizer)) + pipe.text_encoder.get_input_embeddings().weight.data[token_id] = embeddings_dict["clip_l"][i] + pipe.text_encoder_2.get_input_embeddings().weight.data[token_id] = embeddings_dict["clip_g"][i] + else: + raise NotImplementedError self.word_embeddings[name] = embedding except Exception: self.skipped_embeddings[name] = embedding