diff --git a/modules/processing_diffusers.py b/modules/processing_diffusers.py index 01b676e40..9dcfdc954 100644 --- a/modules/processing_diffusers.py +++ b/modules/processing_diffusers.py @@ -7,8 +7,7 @@ import modules.sd_models as sd_models import modules.images as images from modules.lora_diffusers import lora_state, unload_diffusers_lora from modules.processing import StableDiffusionProcessing -from compel import Compel, ReturnedEmbeddingsType - +import modules.prompt_parser_diffusers as prompt_parser_diffusers try: import diffusers @@ -24,28 +23,6 @@ def encode_prompt(encoder, prompt): shared.log.debug(f'Diffuser encoder: {encoder.__class__.__name__} dict={getattr(cfg, "vocab_size", None)} layers={getattr(cfg, "num_hidden_layers", None)} tokens={getattr(cfg, "max_position_embeddings", None)}') embeds = prompt return embeds - -def compel_encode_prompt(pipeline, prompt, negative_prompt): - compel = Compel( - truncate_long_prompts=True, - tokenizer=[ - pipeline.tokenizer, - pipeline.tokenizer_2 - ], - text_encoder=[ - pipeline.text_encoder, - pipeline.text_encoder_2 - ], - returned_embeddings_type=ReturnedEmbeddingsType.PENULTIMATE_HIDDEN_STATES_NON_NORMALIZED, - requires_pooled=[ - False, - True - ] - ) - prompt_embed, pooled = compel(prompt) - negative_embed, negative_pooled = compel(negative_prompt) - [prompt_embed, negative_embed] = compel.pad_conditioning_tensors_to_same_length([prompt_embed, negative_embed]) - return prompt_embed, pooled, negative_embed, negative_pooled def process_diffusers(p: StableDiffusionProcessing, seeds, prompts, negative_prompts): @@ -74,7 +51,7 @@ def process_diffusers(p: StableDiffusionProcessing, seeds, prompts, negative_pro return latents - def set_pipeline_args(model, prompt, negative_prompt, **kwargs): + def set_pipeline_args(model, prompt, negative_prompt, prompt_2=None, negative_prompt_2=None, refiner=False, **kwargs): args = {} pipeline = model signature = inspect.signature(type(pipeline).__call__) @@ -85,18 +62,20 @@ def process_diffusers(p: StableDiffusionProcessing, seeds, prompts, negative_pro pooled = None negative_embed = None negative_pooled = None - if shared.opts.data['prompt_attention'] == 'Compel parser' and (shared.opts.diffusers_pipeline == shared.pipelines[1] or shared.opts.diffusers_pipeline == shared.pipelines[7]): #Gated for SDXL only - prompt_embed, pooled, negative_embed, negative_pooled = compel_encode_prompt(model, prompt, negative_prompt) + if shared.opts.data['prompt_attention'] == 'Compel parser': + prompt_embed, pooled, negative_embed, negative_pooled = prompt_parser_diffusers.compel_encode_prompt(model, prompt, negative_prompt, prompt_2, negative_prompt_2, refiner) if 'prompt' in possible: if hasattr(model, 'text_encoder') and 'prompt_embeds' in possible and prompt_embed is not None: args['prompt_embeds'] = prompt_embed args['pooled_prompt_embeds'] = pooled + args['prompt_2'] = None #Cannot pass prompts when passing embeds else: args['prompt'] = prompt if 'negative_prompt' in possible: if hasattr(model, 'text_encoder') and 'negative_prompt_embeds' in possible and negative_embed is not None: args['negative_prompt_embeds'] = negative_embed args['negative_pooled_prompt_embeds'] = negative_pooled + args['negative_prompt_2'] = None else: args['negative_prompt'] = negative_prompt if 'num_inference_steps' in possible: @@ -118,10 +97,6 @@ def process_diffusers(p: StableDiffusionProcessing, seeds, prompts, negative_pro args[arg] = kwargs[arg] else: pass - if prompt_embed is not None: #Cannot pass prompts when passing embeds - del args['prompt_2'] - del args['negative_prompt_2'] - # shared.log.debug(f'Diffuser not supported: pipeline={pipeline.__class__.__name__} task={sd_models.get_diffusers_task(model)} arg={arg}') # shared.log.debug(f'Diffuser pipeline: {pipeline.__class__.__name__} possible={possible}') clean = args.copy() @@ -182,6 +157,7 @@ def process_diffusers(p: StableDiffusionProcessing, seeds, prompts, negative_pro denoising_start=0 if refiner_enabled and p.refiner_start > 0 and p.refiner_start < 1 else None, denoising_end=p.refiner_start if refiner_enabled and p.refiner_start > 0 and p.refiner_start < 1 else None, output_type='latent' if hasattr(shared.sd_model, 'vae') else 'np', + refiner=False, **task_specific_kwargs ) output = shared.sd_model(**pipe_args) # pylint: disable=not-callable @@ -235,6 +211,7 @@ def process_diffusers(p: StableDiffusionProcessing, seeds, prompts, negative_pro denoising_end=1 if p.refiner_start > 0 and p.refiner_start < 1 else None, image=output.images[i], output_type='latent' if hasattr(shared.sd_refiner, 'vae') else 'np', + refiner=True ) refiner_output = shared.sd_refiner(**pipe_args) # pylint: disable=not-callable if not shared.state.interrupted and not shared.state.skipped: @@ -244,6 +221,7 @@ def process_diffusers(p: StableDiffusionProcessing, seeds, prompts, negative_pro if shared.opts.diffusers_move_refiner and not shared.sd_refiner.has_accelerate: shared.log.debug('Diffusers: Moving refiner model to CPU') shared.sd_refiner.to(devices.cpu) + devices.torch_gc() else: results = output.images @@ -252,4 +230,4 @@ def process_diffusers(p: StableDiffusionProcessing, seeds, prompts, negative_pro - return results + return results \ No newline at end of file diff --git a/modules/prompt_parser_diffusers.py b/modules/prompt_parser_diffusers.py new file mode 100644 index 000000000..10c1bc3db --- /dev/null +++ b/modules/prompt_parser_diffusers.py @@ -0,0 +1,36 @@ +import torch +import modules.shared as shared +from compel import Compel, ReturnedEmbeddingsType + +def compel_encode_prompt(pipeline, prompt, negative_prompt, prompt_2=None, negative_prompt_2=None, refiner=False): + if "XL" not in pipeline.__class__.__name__: + print(f"Compel parser is not configured for: {pipeline.__class__.__name__}") + return None, None, None, None + compel_te1 = Compel( + tokenizer=pipeline.tokenizer, + text_encoder=pipeline.text_encoder, + returned_embeddings_type=ReturnedEmbeddingsType.PENULTIMATE_HIDDEN_STATES_NON_NORMALIZED, + requires_pooled=False, + ) + + compel_te2 = Compel( + tokenizer=pipeline.tokenizer_2, + text_encoder=pipeline.text_encoder_2, + returned_embeddings_type=ReturnedEmbeddingsType.PENULTIMATE_HIDDEN_STATES_NON_NORMALIZED, + requires_pooled=True, + ) + if not refiner: + positive_te1 = compel_te1(prompt) + positive_te2, 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) + if refiner: + positive, pooled = compel_te2(prompt) + negative, negative_pooled = compel_te2(negative_prompt) + + + [prompt_embed, negative_embed] = compel_te2.pad_conditioning_tensors_to_same_length([positive, negative]) + return prompt_embed, pooled, negative_embed, negative_pooled \ No newline at end of file