From 889274f39b25523cc4bdabcae5bb3c6ea2e1acfa Mon Sep 17 00:00:00 2001 From: Vladimir Mandic Date: Tue, 5 Sep 2023 08:49:37 -0400 Subject: [PATCH] post merge cleanup --- modules/prompt_parser_diffusers.py | 91 ++++++++---------------------- 1 file changed, 24 insertions(+), 67 deletions(-) diff --git a/modules/prompt_parser_diffusers.py b/modules/prompt_parser_diffusers.py index 665b2ae31..758777da4 100644 --- a/modules/prompt_parser_diffusers.py +++ b/modules/prompt_parser_diffusers.py @@ -5,7 +5,7 @@ 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 @@ -34,82 +34,39 @@ 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. + def maybe_convert_prompt(self, prompt: typing.Union[str, typing.List[str]], tokenizer = "PreTrainedTokenizer"): + prompts = [prompt] if not isinstance(prompt, typing.List) else prompt + prompts = [self._maybe_convert_prompt(p, tokenizer) for p in prompts] + if not isinstance(prompt, typing.List): + return prompts[0] + return prompts - 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. + def _maybe_convert_prompt(self, prompt: str, tokenizer = "PreTrainedTokenizer"): + 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 - 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]: + def expand_textual_inversion_token_ids_if_necessary(self, token_ids: typing.List[int]) -> typing.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, @@ -191,7 +148,7 @@ def compel_encode_prompt( textual_inversion_manager=textual_inversion_manager ) - if not is_refiner and shared.sd_model_type == "sdxl": + if 'XL' in pipeline.__class__.__name__ and not is_refiner: 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) @@ -205,7 +162,7 @@ def compel_encode_prompt( [prompt_embed, negative_embed] = compel_te2.pad_conditioning_tensors_to_same_length([positive, negative]) return prompt_embed, positive_pooled, negative_embed, negative_pooled - if is_refiner and shared.sd_refiner_type == "sdxl": + elif 'XL' in pipeline.__class__.__name__ and is_refiner: 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)