From d3802f9fbc049107e91420da3188c1a645c73b61 Mon Sep 17 00:00:00 2001 From: AI-Casanova <54461896+AI-Casanova@users.noreply.github.com> Date: Sun, 16 Jun 2024 18:00:30 -0500 Subject: [PATCH 1/2] SD3 Prompt Parsing, preliminary --- modules/processing_args.py | 6 ++- modules/prompt_parser_diffusers.py | 74 +++++++++++++++++++++++++----- 2 files changed, 68 insertions(+), 12 deletions(-) diff --git a/modules/processing_args.py b/modules/processing_args.py index a8e7c33a3..163c51efd 100644 --- a/modules/processing_args.py +++ b/modules/processing_args.py @@ -106,7 +106,7 @@ def set_pipeline_args(p, model, prompts: list, negative_prompts: list, prompts_2 shared.log.error(f'Sampler timesteps: {e}') else: shared.log.warning(f'Sampler: sampler={model.scheduler.__class__.__name__} timesteps not supported') - if shared.opts.prompt_attention != 'Fixed attention' and ('StableDiffusion' in model.__class__.__name__ or 'StableCascade' in model.__class__.__name__) and 'Onnx' not in model.__class__.__name__ and 'StableDiffusion3' not in model.__class__.__name__: + if shared.opts.prompt_attention != 'Fixed attention' and ('StableDiffusion' in model.__class__.__name__ or 'StableCascade' in model.__class__.__name__) and 'Onnx' not in model.__class__.__name__: try: prompt_parser_diffusers.encode_prompts(model, p, prompts, negative_prompts, steps=steps, clip_skip=clip_skip) parser = shared.opts.prompt_attention @@ -126,6 +126,8 @@ def set_pipeline_args(p, model, prompts: list, negative_prompts: list, prompts_2 args['prompt_embeds_pooled'] = p.positive_pooleds[0].unsqueeze(0) elif 'XL' in model.__class__.__name__ and len(getattr(p, 'positive_pooleds', [])) > 0: args['pooled_prompt_embeds'] = p.positive_pooleds[0] + elif 'StableDiffusion3' in model.__class__.__name__ and len(getattr(p, 'positive_pooleds', [])) > 0: + args['pooled_prompt_embeds'] = p.positive_pooleds[0] else: args['prompt'] = prompts if 'negative_prompt' in possible: @@ -135,6 +137,8 @@ def set_pipeline_args(p, model, prompts: list, negative_prompts: list, prompts_2 args['negative_prompt_embeds_pooled'] = p.negative_pooleds[0].unsqueeze(0) if 'XL' in model.__class__.__name__ and len(getattr(p, 'negative_pooleds', [])) > 0: args['negative_pooled_prompt_embeds'] = p.negative_pooleds[0] + if 'StableDiffusion3' in model.__class__.__name__ and len(getattr(p, 'negative_pooleds', [])) > 0: + args['negative_pooled_prompt_embeds'] = p.negative_pooleds[0] else: if 'PixArtSigmaPipeline' in model.__class__.__name__: # pixart-sigma pipeline throws list-of-list for negative prompt args['negative_prompt'] = negative_prompts[0] diff --git a/modules/prompt_parser_diffusers.py b/modules/prompt_parser_diffusers.py index 3c472b227..0d410af87 100644 --- a/modules/prompt_parser_diffusers.py +++ b/modules/prompt_parser_diffusers.py @@ -217,7 +217,7 @@ def prepare_embedding_providers(pipe, clip_skip) -> list[EmbeddingsProvider]: embeddings_providers = [] if 'StableCascade' in pipe.__class__.__name__: embedding_type = -(clip_skip) - elif 'XL' in pipe.__class__.__name__: + elif 'XL' in pipe.__class__.__name__ or 'SD3' in pipe.__class__.__name__: embedding_type = -(clip_skip + 1) else: embedding_type = clip_skip @@ -257,15 +257,33 @@ def pad_to_same_length(pipe, embeds): embeds[i] = embed return embeds +def split_prompts(prompt, SD3 = False): + if prompt.find("TE2:") != -1: + prompt, prompt2 = prompt.split("TE2:") + else: + prompt2 = prompt + + if prompt.find("TE3:") != -1: + prompt, prompt3 = prompt.split("TE3:") + elif prompt2.find("TE3:") != -1: + prompt2, prompt3 = prompt2.split("TE3:") + else: + prompt3 = prompt + + prompt = prompt.strip() + prompt2 = " " if prompt2.strip() == "" else prompt2.strip() + prompt3 = " " if prompt3.strip() == "" else prompt3.strip() + + if SD3 and prompt3 != " ": + ps, ws = get_prompts_with_weights(prompt3) + prompt3 = ", ".join(ps) + return prompt, prompt2, prompt3 def get_weighted_text_embeddings(pipe, prompt: str = "", neg_prompt: str = "", clip_skip: int = None): device = pipe.device if str(pipe.device) != 'meta' else devices.device - prompt_split = prompt.split("TE2:") - prompt = prompt_split[0] - prompt_2 = prompt_split[-1] - neg_prompt_split = neg_prompt.split("TE2:") - neg_prompt_2 = neg_prompt_split[-1] - neg_prompt = neg_prompt_split[0] + SD3 = hasattr(pipe, 'text_encoder_3') + prompt, prompt_2, prompt_3 = split_prompts(prompt, SD3) + neg_prompt, neg_prompt_2, neg_prompt_3 = split_prompts(prompt, SD3) if prompt != prompt_2: ps = [get_prompts_with_weights(p) for p in [prompt, prompt_2]] @@ -285,8 +303,8 @@ def get_weighted_text_embeddings(pipe, prompt: str = "", neg_prompt: str = "", c embedding_providers = prepare_embedding_providers(pipe, clip_skip) prompt_embeds = [] negative_prompt_embeds = [] - pooled_prompt_embeds = None - negative_pooled_prompt_embeds = None + pooled_prompt_embeds = [] + negative_pooled_prompt_embeds = [] for i in range(len(embedding_providers)): t0 = time.time() text = list(positives[i]) @@ -310,8 +328,25 @@ def get_weighted_text_embeddings(pipe, prompt: str = "", neg_prompt: str = "", c embed, ntokens = embedding_providers[i].get_embeddings_for_weighted_prompt_fragments(text_batch=[negatives[i]], fragment_weights_batch=[negative_weights[i]], device=device, should_return_tokens=True) negative_prompt_embeds.append(embed) debug(f'Prompt: unpadded shape={prompt_embeds[0].shape} TE{i+1} ptokens={torch.count_nonzero(ptokens)} ntokens={torch.count_nonzero(ntokens)} time={(time.time() - t0):.3f}') - - if prompt_embeds[-1].shape[-1] > 768: + if SD3: + t0 = time.time() + for i in range(len(prompt_embeds)): + pooled_prompt_embeds.append(prompt_embeds[i][ + torch.arange(prompt_embeds[i].shape[0], device=device), + (ptokens.to(dtype=torch.int, device=device) == 49407) + .int() + .argmax(dim=-1), + ]) + negative_pooled_prompt_embeds.append(negative_prompt_embeds[i][ + torch.arange(negative_prompt_embeds[i].shape[0], device=device), + (ntokens.to(dtype=torch.int, device=device) == 49407) + .int() + .argmax(dim=-1), + ]) + pooled_prompt_embeds = torch.cat(pooled_prompt_embeds, dim=-1) + negative_pooled_prompt_embeds = torch.cat(negative_pooled_prompt_embeds, dim=-1) + debug(f'Prompt: pooled shape={pooled_prompt_embeds[0].shape} time={(time.time() - t0):.3f}') + elif prompt_embeds[-1].shape[-1] > 768: t0 = time.time() if shared.opts.diffusers_pooled == "weighted": pooled_prompt_embeds = prompt_embeds[-1][ @@ -341,4 +376,21 @@ def get_weighted_text_embeddings(pipe, prompt: str = "", neg_prompt: str = "", c 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]) + if SD3: + t5_prompt_embed = pipe._get_t5_prompt_embeds( + prompt=prompt_3, + num_images_per_prompt=prompt_embeds.shape[0], + device=pipe.device, + ) + prompt_embeds = torch.nn.functional.pad( + prompt_embeds, (0, t5_prompt_embed.shape[-1] - prompt_embeds.shape[-1])) + prompt_embeds = torch.cat([prompt_embeds, t5_prompt_embed], dim=-2) + t5_negative_prompt_embed = pipe._get_t5_prompt_embeds( + prompt=neg_prompt_3, + num_images_per_prompt=prompt_embeds.shape[0], + device=pipe.device, + ) + negative_prompt_embeds = torch.nn.functional.pad( + negative_prompt_embeds, (0, t5_negative_prompt_embed.shape[-1] - negative_prompt_embeds.shape[-1])) + negative_prompt_embeds = torch.cat([negative_prompt_embeds, t5_negative_prompt_embed], dim=-2) return prompt_embeds, pooled_prompt_embeds, negative_prompt_embeds, negative_pooled_prompt_embeds From e3009753e31bcd7c132abd96e2ce6b72c85312eb Mon Sep 17 00:00:00 2001 From: AI-Casanova <54461896+AI-Casanova@users.noreply.github.com> Date: Wed, 19 Jun 2024 23:41:11 -0500 Subject: [PATCH 2/2] Finish SD3 Prompt Parsing, reconfigure Compel Hijack --- modules/prompt_parser_diffusers.py | 54 ++++++++++++++++++------------ modules/sd_models.py | 3 ++ 2 files changed, 35 insertions(+), 22 deletions(-) diff --git a/modules/prompt_parser_diffusers.py b/modules/prompt_parser_diffusers.py index 0d410af87..29cf86441 100644 --- a/modules/prompt_parser_diffusers.py +++ b/modules/prompt_parser_diffusers.py @@ -41,8 +41,26 @@ def compel_hijack(self, token_ids: torch.Tensor, return hidden_state -EmbeddingsProvider._encode_token_ids_to_embeddings = compel_hijack # pylint: disable=protected-access +def sd3_compel_hijack(self, token_ids: torch.Tensor, + attention_mask: typing.Optional[torch.Tensor] = None) -> torch.Tensor: + needs_hidden_states = True + text_encoder_output = self.text_encoder(token_ids, attention_mask, output_hidden_states=needs_hidden_states, return_dict=True) + clip_skip = int(self.returned_embeddings_type) + hidden_state = text_encoder_output.hidden_states[-(clip_skip+1)] + return hidden_state + +def insert_parser_highjack(pipename): + if "StableDiffusion3" in pipename: + EmbeddingsProvider._encode_token_ids_to_embeddings = sd3_compel_hijack # pylint: disable=protected-access + debug("Loading SD3 Parser hijack") + else: + EmbeddingsProvider._encode_token_ids_to_embeddings = compel_hijack # pylint: disable=protected-access + debug("Loading Standard Parser hijack") + + + +insert_parser_highjack("Initialize") # from https://github.com/damian0815/compel/blob/main/src/compel/diffusers_textual_inversion_manager.py class DiffusersTextualInversionManager(BaseTextualInversionManager): @@ -217,7 +235,7 @@ def prepare_embedding_providers(pipe, clip_skip) -> list[EmbeddingsProvider]: embeddings_providers = [] if 'StableCascade' in pipe.__class__.__name__: embedding_type = -(clip_skip) - elif 'XL' in pipe.__class__.__name__ or 'SD3' in pipe.__class__.__name__: + elif 'XL' in pipe.__class__.__name__: embedding_type = -(clip_skip + 1) else: embedding_type = clip_skip @@ -237,7 +255,7 @@ def pad_to_same_length(pipe, embeds): 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 - if shared.opts.diffusers_zeros_prompt_pad: + if shared.opts.diffusers_zeros_prompt_pad or 'StableDiffusion3' in pipe.__class__.__name__: empty_embed = [torch.zeros((1, 77, embeds[0].shape[2]), device=device, dtype=embeds[0].dtype)] else: try: @@ -276,14 +294,15 @@ def split_prompts(prompt, SD3 = False): if SD3 and prompt3 != " ": ps, ws = get_prompts_with_weights(prompt3) - prompt3 = ", ".join(ps) + prompt3 = " ".join(ps) return prompt, prompt2, prompt3 + def get_weighted_text_embeddings(pipe, prompt: str = "", neg_prompt: str = "", clip_skip: int = None): device = pipe.device if str(pipe.device) != 'meta' else devices.device SD3 = hasattr(pipe, 'text_encoder_3') prompt, prompt_2, prompt_3 = split_prompts(prompt, SD3) - neg_prompt, neg_prompt_2, neg_prompt_3 = split_prompts(prompt, SD3) + neg_prompt, neg_prompt_2, neg_prompt_3 = split_prompts(neg_prompt, SD3) if prompt != prompt_2: ps = [get_prompts_with_weights(p) for p in [prompt, prompt_2]] @@ -330,37 +349,28 @@ def get_weighted_text_embeddings(pipe, prompt: str = "", neg_prompt: str = "", c debug(f'Prompt: unpadded shape={prompt_embeds[0].shape} TE{i+1} ptokens={torch.count_nonzero(ptokens)} ntokens={torch.count_nonzero(ntokens)} time={(time.time() - t0):.3f}') if SD3: t0 = time.time() - for i in range(len(prompt_embeds)): - pooled_prompt_embeds.append(prompt_embeds[i][ - torch.arange(prompt_embeds[i].shape[0], device=device), - (ptokens.to(dtype=torch.int, device=device) == 49407) - .int() - .argmax(dim=-1), - ]) - negative_pooled_prompt_embeds.append(negative_prompt_embeds[i][ - torch.arange(negative_prompt_embeds[i].shape[0], device=device), - (ntokens.to(dtype=torch.int, device=device) == 49407) - .int() - .argmax(dim=-1), - ]) + pooled_prompt_embeds.append(embedding_providers[0].get_pooled_embeddings(texts=positives[0] if len(positives[0]) == 1 else [" ".join(positives[0])], device=device)) + pooled_prompt_embeds.append(embedding_providers[1].get_pooled_embeddings(texts=positives[-1] if len(positives[-1]) == 1 else [" ".join(positives[-1])], device=device)) + negative_pooled_prompt_embeds.append(embedding_providers[0].get_pooled_embeddings(texts=negatives[0] if len(negatives[0]) == 1 else [" ".join(negatives[0])], device=device)) + negative_pooled_prompt_embeds.append(embedding_providers[1].get_pooled_embeddings(texts=negatives[-1] if len(negatives[-1]) == 1 else [" ".join(negatives[-1])], device=device)) pooled_prompt_embeds = torch.cat(pooled_prompt_embeds, dim=-1) negative_pooled_prompt_embeds = torch.cat(negative_pooled_prompt_embeds, dim=-1) debug(f'Prompt: pooled shape={pooled_prompt_embeds[0].shape} time={(time.time() - t0):.3f}') elif prompt_embeds[-1].shape[-1] > 768: t0 = time.time() if shared.opts.diffusers_pooled == "weighted": - pooled_prompt_embeds = prompt_embeds[-1][ + pooled_prompt_embeds = embedding_providers[-1].text_encoder.text_projection(prompt_embeds[-1][ torch.arange(prompt_embeds[-1].shape[0], device=device), (ptokens.to(dtype=torch.int, device=device) == 49407) .int() .argmax(dim=-1), - ] - negative_pooled_prompt_embeds = negative_prompt_embeds[-1][ + ]) + negative_pooled_prompt_embeds = embedding_providers[-1].text_encoder.text_projection(negative_prompt_embeds[-1][ torch.arange(negative_prompt_embeds[-1].shape[0], device=device), (ntokens.to(dtype=torch.int, device=device) == 49407) .int() .argmax(dim=-1), - ] + ]) else: try: pooled_prompt_embeds = embedding_providers[-1].get_pooled_embeddings(texts=[prompt_2], device=device) if prompt_embeds[-1].shape[-1] > 768 else None diff --git a/modules/sd_models.py b/modules/sd_models.py index 81214bae5..22926447d 100644 --- a/modules/sd_models.py +++ b/modules/sd_models.py @@ -1158,6 +1158,9 @@ def load_diffuser(checkpoint_info=None, already_loaded_state_dict=None, timer=No sd_model.embedding_db.load_textual_inversion_embeddings(force_reload=True) timer.record("embeddings") + from modules.prompt_parser_diffusers import insert_parser_highjack + insert_parser_highjack(sd_model.__class__.__name__) + set_diffuser_options(sd_model, vae, op) if op == 'model': sd_vae.apply_vae_config(shared.sd_model.sd_checkpoint_info.filename, vae_file, sd_model)