Fix Cascade empty prompt encode

This commit is contained in:
Disty0
2024-08-09 00:32:19 +03:00
parent e1c40384ad
commit 6431296ec1
+11 -4
View File
@@ -253,6 +253,8 @@ def prepare_embedding_providers(pipe, clip_skip) -> list[EmbeddingsProvider]:
if getattr(pipe, "prior_pipe", None) is not None and getattr(pipe.prior_pipe, "tokenizer", None) is not None and getattr(pipe.prior_pipe, "text_encoder", None) is not None:
provider = EmbeddingsProvider(padding_attention_mask_value=0, tokenizer=pipe.prior_pipe.tokenizer, text_encoder=pipe.prior_pipe.text_encoder, truncate=False, returned_embeddings_type=embedding_type, device=device)
embeddings_providers.append(provider)
no_mask_provider = EmbeddingsProvider(padding_attention_mask_value=1, tokenizer=pipe.prior_pipe.tokenizer, text_encoder=pipe.prior_pipe.text_encoder, truncate=False, returned_embeddings_type=embedding_type, device=device)
embeddings_providers.append(no_mask_provider)
elif getattr(pipe, "tokenizer", None) is not None and getattr(pipe, "text_encoder", None) is not None:
provider = EmbeddingsProvider(tokenizer=pipe.tokenizer, text_encoder=pipe.text_encoder, truncate=False, returned_embeddings_type=embedding_type, device=device)
embeddings_providers.append(provider)
@@ -262,7 +264,7 @@ def prepare_embedding_providers(pipe, clip_skip) -> list[EmbeddingsProvider]:
return embeddings_providers
def pad_to_same_length(pipe, embeds):
def pad_to_same_length(pipe, embeds, empty_embedding_providers=None):
if not hasattr(pipe, 'encode_prompt') and 'StableCascade' not in pipe.__class__.__name__:
return embeds
device = pipe.device if str(pipe.device) != 'meta' else devices.device
@@ -271,8 +273,8 @@ def pad_to_same_length(pipe, embeds):
else:
try:
if 'StableCascade' in pipe.__class__.__name__:
empty_embed = pipe.prior_pipe.encode_prompt(device, 1, 1, False, prompt="")
empty_embed = [torch.nn.functional.normalize(empty_embed[0])]
empty_embed = empty_embedding_providers[0].get_embeddings_for_weighted_prompt_fragments(text_batch=[[""]], fragment_weights_batch=[[1]], should_return_tokens=False, device=device)
empty_embed = [empty_embed]
else:
empty_embed = pipe.encode_prompt("")
except TypeError: # SD1.5
@@ -331,6 +333,11 @@ def get_weighted_text_embeddings(pipe, prompt: str = "", neg_prompt: str = "", c
negative_weights.pop(0)
embedding_providers = prepare_embedding_providers(pipe, clip_skip)
empty_embedding_providers = None
if 'StableCascade' in pipe.__class__.__name__:
empty_embedding_providers = [embedding_providers[1]]
embedding_providers = [embedding_providers[0]]
prompt_embeds = []
negative_prompt_embeds = []
pooled_prompt_embeds = []
@@ -400,7 +407,7 @@ def get_weighted_text_embeddings(pipe, prompt: str = "", neg_prompt: str = "", c
negative_pooled_prompt_embeds = None
debug(f'Prompt: positive={prompt_embeds.shape if prompt_embeds is not None else None} pooled={pooled_prompt_embeds.shape if pooled_prompt_embeds is not None else None} negative={negative_prompt_embeds.shape if negative_prompt_embeds is not None else None} pooled={negative_pooled_prompt_embeds.shape if negative_pooled_prompt_embeds is not None else None}')
if prompt_embeds.shape[1] != negative_prompt_embeds.shape[1]:
[prompt_embeds, negative_prompt_embeds] = pad_to_same_length(pipe, [prompt_embeds, negative_prompt_embeds])
[prompt_embeds, negative_prompt_embeds] = pad_to_same_length(pipe, [prompt_embeds, negative_prompt_embeds], empty_embedding_providers=empty_embedding_providers)
if SD3:
device = pipe.device if str(pipe.device) != 'meta' else devices.device
t5_prompt_embed = pipe._get_t5_prompt_embeds( # pylint: disable=protected-access