Add prompt_parser_diffusers.py

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