From 08d094bb9588347ee1d462b1b595ca4cb515f104 Mon Sep 17 00:00:00 2001 From: AI-Casanova <54461896+AI-Casanova@users.noreply.github.com> Date: Sun, 1 Oct 2023 15:52:42 -0500 Subject: [PATCH] Major Refactor prompt_parser_diffusers.py Prompt_2 has been mostly overridden by `TE2` keyword in prompt --- modules/prompt_parser_diffusers.py | 211 ++++++++++++++--------------- 1 file changed, 98 insertions(+), 113 deletions(-) diff --git a/modules/prompt_parser_diffusers.py b/modules/prompt_parser_diffusers.py index d28d32528..c9218b66a 100644 --- a/modules/prompt_parser_diffusers.py +++ b/modules/prompt_parser_diffusers.py @@ -1,31 +1,13 @@ import os import typing import torch -from compel import Compel, ReturnedEmbeddingsType -from compel.embeddings_provider import BaseTextualInversionManager -from modules import shared, devices, prompt_parser - +from compel import ReturnedEmbeddingsType +from compel.embeddings_provider import BaseTextualInversionManager, EmbeddingsProvider +from modules import shared, prompt_parser debug_output = os.environ.get('SD_PROMPT_DEBUG', None) debug = shared.log.info if debug_output is not None else lambda *args, **kwargs: None - -def convert_to_compel(prompt: str): - if prompt is None: - return None - all_schedules = prompt_parser.get_learned_conditioning_prompt_schedules([prompt], 100)[0] - output_list = prompt_parser.parse_prompt_attention(all_schedules[0][1]) - converted_prompt = [] - for subprompt, weight in output_list: - if subprompt != " ": - if weight == 1: - converted_prompt.append(subprompt) - else: - converted_prompt.append(f"({subprompt}){weight}") - converted_prompt = " ".join(converted_prompt) - return converted_prompt - - CLIP_SKIP_MAPPING = { None: ReturnedEmbeddingsType.LAST_HIDDEN_STATES_NORMALIZED, 1: ReturnedEmbeddingsType.LAST_HIDDEN_STATES_NORMALIZED, @@ -33,22 +15,23 @@ CLIP_SKIP_MAPPING = { } -#from https://github.com/damian0815/compel/blob/main/src/compel/diffusers_textual_inversion_manager.py +# from https://github.com/damian0815/compel/blob/main/src/compel/diffusers_textual_inversion_manager.py class DiffusersTextualInversionManager(BaseTextualInversionManager): - def __init__(self, pipe): + def __init__(self, pipe, tokenizer): self.pipe = pipe + self.tokenizer = tokenizer if hasattr(self.pipe, 'embedding_db'): self.pipe.embedding_db.embeddings_used.clear() - #from https://github.com/huggingface/diffusers/blob/705c592ea98ba4e288d837b9cba2767623c78603/src/diffusers/loaders.py#L599 - def maybe_convert_prompt(self, prompt: typing.Union[str, typing.List[str]], tokenizer = "PreTrainedTokenizer"): + # from https://github.com/huggingface/diffusers/blob/705c592ea98ba4e288d837b9cba2767623c78603/src/diffusers/loaders.py#L599 + 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 - def _maybe_convert_prompt(self, prompt: str, tokenizer = "PreTrainedTokenizer"): + def _maybe_convert_prompt(self, prompt: str, tokenizer="PreTrainedTokenizer"): tokens = tokenizer.tokenize(prompt) unique_tokens = set(tokens) for token in unique_tokens: @@ -70,33 +53,30 @@ class DiffusersTextualInversionManager(BaseTextualInversionManager): return token_ids prompt = self.pipe.tokenizer.decode(token_ids) prompt = self.maybe_convert_prompt(prompt, self.pipe.tokenizer) + print(prompt) return self.pipe.tokenizer.encode(prompt, add_special_tokens=False) -def compel_encode_prompts( - pipeline, - prompts: list, - negative_prompts: list, - prompts_2: typing.Optional[list] = None, - negative_prompts_2: typing.Optional[list] = None, - is_refiner: bool = None, - clip_skip: typing.Optional[int] = None, +def encode_prompts( + pipeline, + prompts: list, + negative_prompts: list, + clip_skip: typing.Optional[int] = None, ): - prompt_embeds = [] - positive_pooleds = [] - negative_embeds = [] - negative_pooleds = [] - for i in range(len(prompts)): - prompt_embed, positive_pooled, negative_embed, negative_pooled = compel_encode_prompt(pipeline, - prompts[i], - negative_prompts[i], - prompts_2[i] if prompts_2 is not None else None, - negative_prompts_2[i] if negative_prompts_2 is not None else None, - is_refiner, clip_skip) - prompt_embeds.append(prompt_embed) - positive_pooleds.append(positive_pooled) - negative_embeds.append(negative_embed) - negative_pooleds.append(negative_pooled) + if 'StableDiffusion' not in pipeline.__class__.__name__: + shared.log.warning(f"Prompt parser not supported: {pipeline.__class__.__name__}") + return None, None, None, None + else: + prompt_embeds = [] + positive_pooleds = [] + negative_embeds = [] + negative_pooleds = [] + for i in range(len(prompts)): + prompt_embed, positive_pooled, negative_embed, negative_pooled = get_weighted_text_embeddings_sdxl(pipeline,prompts[i], negative_prompts[i], clip_skip) + prompt_embeds.append(prompt_embed) + positive_pooleds.append(positive_pooled) + negative_embeds.append(negative_embed) + negative_pooleds.append(negative_pooled) if prompt_embeds is not None: prompt_embeds = torch.cat(prompt_embeds, dim=0) @@ -109,76 +89,81 @@ def compel_encode_prompts( return prompt_embeds, positive_pooleds, negative_embeds, negative_pooleds -def compel_encode_prompt( - pipeline, - prompt: str, - negative_prompt: str, - prompt_2: typing.Optional[str] = None, - negative_prompt_2: typing.Optional[str] = None, - is_refiner: bool = None, - clip_skip: typing.Optional[int] = None, -): - if 'StableDiffusion' not in pipeline.__class__.__name__: - shared.log.warning(f"Prompt parser: Compel not supported: {pipeline.__class__.__name__}") - return (None, None, None, None) +def get_prompts_with_weights(prompt: str): + prompt = DiffusersTextualInversionManager(shared.sd_model, + shared.sd_model.tokenizer or shared.sd_model.tokenizer_2).maybe_convert_prompt( + prompt, shared.sd_model.tokenizer or shared.sd_model.tokenizer_2) + texts_and_weights = prompt_parser.parse_prompt_attention(prompt) + texts = [t for t, w in texts_and_weights] + text_weights = [w for t, w in texts_and_weights] + return texts, text_weights - if 'XL' in pipeline.__class__.__name__ and not is_refiner: + +def prepare_embedding_providers(pipe, clip_skip): + embeddings_providers = [] + if 'XL' in pipe.__class__.__name__: embedding_type = ReturnedEmbeddingsType.PENULTIMATE_HIDDEN_STATES_NON_NORMALIZED - if clip_skip is not None and clip_skip > 1: - shared.log.warning(f"Prompt parser SDXL unsupported: clip_skip={clip_skip}") - elif 'XL' in pipeline.__class__.__name__ and is_refiner: - embedding_type = ReturnedEmbeddingsType.PENULTIMATE_HIDDEN_STATES_NON_NORMALIZED - if clip_skip is not None and clip_skip > 1: - shared.log.warning(f"Prompt parser SDXL unsupported: clip_skip={clip_skip}") else: - embedding_type = CLIP_SKIP_MAPPING.get(clip_skip, ReturnedEmbeddingsType.PENULTIMATE_HIDDEN_STATES_NORMALIZED) - if clip_skip not in CLIP_SKIP_MAPPING: - shared.log.warning(f"Prompt parser unsupported: clip_skip={clip_skip} expected={set(CLIP_SKIP_MAPPING.keys())}") + if clip_skip > 2: + shared.log.warning(f"Prompt parser unsupported: clip_skip={clip_skip}") + clip_skip = 2 + embedding_type = CLIP_SKIP_MAPPING[clip_skip] + if hasattr(pipe, "tokenizer") and hasattr(pipe, "text_encoder"): + embeddings_providers.append( + EmbeddingsProvider(tokenizer=pipe.tokenizer, text_encoder=pipe.text_encoder, + truncate=False, + returned_embeddings_type=embedding_type)) + if hasattr(pipe, "tokenizer_2") and hasattr(pipe, "text_encoder_2"): + embeddings_providers.append( + EmbeddingsProvider(tokenizer=pipe.tokenizer_2, text_encoder=pipe.text_encoder_2, + truncate=False, + returned_embeddings_type=embedding_type)) + return embeddings_providers - if shared.opts.prompt_attention != "Compel parser": - prompt = convert_to_compel(prompt) - negative_prompt = convert_to_compel(negative_prompt) - prompt_2 = convert_to_compel(prompt_2) - negative_prompt_2 = convert_to_compel(negative_prompt_2) - textual_inversion_manager_te1 = 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=devices.device, - textual_inversion_manager=textual_inversion_manager_te1 - ) +def get_weighted_text_embeddings_sdxl( + pipe, + prompt: str = "", + neg_prompt: str = "", + clip_skip: int = None +): + prompt_2 = prompt.split("TE2:")[-1] + neg_prompt_2 = neg_prompt.split("TE2:")[-1] + prompt = prompt.split("TE2:")[0] + neg_prompt = neg_prompt.split("TE2:")[0] - if 'XL' in pipeline.__class__.__name__ and not is_refiner: - # TODO textual_inversion_manager=textual_inversion_manager_te2 - DiffusersTextualInversionManager needs to use tokenizer_2 - compel_te2 = Compel(tokenizer=pipeline.tokenizer_2, text_encoder=pipeline.text_encoder_2, returned_embeddings_type=embedding_type, requires_pooled=True, device=devices.device) - positive_te1 = compel_te1(prompt) - positive_te2, positive_pooled = compel_te2(prompt_2) - positive = torch.cat((positive_te1, positive_te2), dim=-1) - negative_te1 = compel_te1(negative_prompt) - negative_te2, negative_pooled = compel_te2(negative_prompt_2) - negative = torch.cat((negative_te1, negative_te2), dim=-1) + ps = [get_prompts_with_weights(p) for p in [prompt, prompt_2]] + positives = [t for t, w in ps] + positive_weights = [w for t, w in ps] + ns = [get_prompts_with_weights(p) for p in [neg_prompt, neg_prompt_2]] + negatives = [t for t, w in ns] + negative_weights = [w for t, w in ns] - parsed = compel_te1.parse_prompt_string(prompt) - debug(f"Prompt parser Compel: {parsed}") - [prompt_embed, negative_embed] = compel_te2.pad_conditioning_tensors_to_same_length([positive, negative]) - return prompt_embed, positive_pooled, negative_embed, negative_pooled + if hasattr(pipe, "tokenizer_2") and not hasattr(pipe, "tokenizer"): + positives.pop(0) + positive_weights.pop(0) + negatives.pop(0) + negative_weights.pop(0) - 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=devices.device) - positive, positive_pooled = compel_te2(prompt) - negative, negative_pooled = compel_te2(negative_prompt) + embedding_providers = prepare_embedding_providers(pipe, clip_skip) + prompt_embeds = [] + negative_prompt_embeds = [] + for i in range(len(embedding_providers)): + prompt_embeds.append( + embedding_providers[i].get_embeddings_for_weighted_prompt_fragments(text_batch=[positives[i]], + fragment_weights_batch=[ + positive_weights[i]], + device=pipe.device)) + negative_prompt_embeds.append( + embedding_providers[i].get_embeddings_for_weighted_prompt_fragments(text_batch=[negatives[i]], + fragment_weights_batch=[ + negative_weights[i]], + device=pipe.device)) + prompt_embeds = torch.cat(prompt_embeds, dim=-1) if len(prompt_embeds) > 1 else prompt_embeds[0] + negative_prompt_embeds = torch.cat(negative_prompt_embeds, dim=-1) if len(negative_prompt_embeds) > 1 else negative_prompt_embeds[0] - parsed = compel_te1.parse_prompt_string(prompt) - debug(f"Prompt parser Compel: {parsed}") - [prompt_embed, negative_embed] = compel_te2.pad_conditioning_tensors_to_same_length([positive, negative]) - return prompt_embed, positive_pooled, negative_embed, negative_pooled - - # neither base+sdxl nor refiner+sdxl - positive = compel_te1(prompt) - negative = compel_te1(negative_prompt) - [prompt_embed, negative_embed] = compel_te1.pad_conditioning_tensors_to_same_length([positive, negative]) - return prompt_embed, None, negative_embed, None + pooled_prompt_embeds = embedding_providers[-1].get_pooled_embeddings(texts=[prompt_2], device=pipe.device) if prompt_embeds.shape[-1] > 768 else None + negative_pooled_prompt_embeds = embedding_providers[-1].get_pooled_embeddings(texts=[neg_prompt_2], + device=pipe.device) if negative_prompt_embeds.shape[-1] > 768 else None + return prompt_embeds, pooled_prompt_embeds, negative_prompt_embeds, negative_pooled_prompt_embeds +