mirror of
https://github.com/vladmandic/automatic
synced 2026-09-15 02:58:44 +02:00
Major Refactor prompt_parser_diffusers.py
Prompt_2 has been mostly overridden by `TE2` keyword in prompt
This commit is contained in:
AI-Casanova
GitHub
@@ -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
|
||||
|
||||
|
||||
Reference in New Issue
Block a user