Merge pull request #6304 from vladmandic/add-cross-attention-info

add cross-attention info

1 files changed, 11 insertions, 1 deletions

diff --git a/modules/sd_hijack.py b/modules/sd_hijack.py
index edcbaf52..fa2cd4bb 100644
--- a/modules/sd_hijack.py
+++ b/modules/sd_hijack.py
@@ -35,26 +35,35 @@ def apply_optimizations():
 

     ldm.modules.diffusionmodules.model.nonlinearity = silu

     ldm.modules.diffusionmodules.openaimodel.th = sd_hijack_unet.th

+    

+    optimization_method = None

 

     if cmd_opts.force_enable_xformers or (cmd_opts.xformers and shared.xformers_available and torch.version.cuda and (6, 0) <= torch.cuda.get_device_capability(shared.device) <= (9, 0)):

         print("Applying xformers cross attention optimization.")

         ldm.modules.attention.CrossAttention.forward = sd_hijack_optimizations.xformers_attention_forward

         ldm.modules.diffusionmodules.model.AttnBlock.forward = sd_hijack_optimizations.xformers_attnblock_forward

+        optimization_method = 'xformers'

     elif cmd_opts.opt_split_attention_v1:

         print("Applying v1 cross attention optimization.")

         ldm.modules.attention.CrossAttention.forward = sd_hijack_optimizations.split_cross_attention_forward_v1

+        optimization_method = 'V1'

     elif not cmd_opts.disable_opt_split_attention and (cmd_opts.opt_split_attention_invokeai or not torch.cuda.is_available()):

         if not invokeAI_mps_available and shared.device.type == 'mps':

             print("The InvokeAI cross attention optimization for MPS requires the psutil package which is not installed.")

             print("Applying v1 cross attention optimization.")

             ldm.modules.attention.CrossAttention.forward = sd_hijack_optimizations.split_cross_attention_forward_v1

+            optimization_method = 'V1'

         else:

             print("Applying cross attention optimization (InvokeAI).")

             ldm.modules.attention.CrossAttention.forward = sd_hijack_optimizations.split_cross_attention_forward_invokeAI

+            optimization_method = 'InvokeAI'

     elif not cmd_opts.disable_opt_split_attention and (cmd_opts.opt_split_attention or torch.cuda.is_available()):

         print("Applying cross attention optimization (Doggettx).")

         ldm.modules.attention.CrossAttention.forward = sd_hijack_optimizations.split_cross_attention_forward

         ldm.modules.diffusionmodules.model.AttnBlock.forward = sd_hijack_optimizations.cross_attention_attnblock_forward

+        optimization_method = 'Doggettx'

+

+    return optimization_method

 

 

 def undo_optimizations():

@@ -75,6 +84,7 @@ class StableDiffusionModelHijack:
     layers = None

     circular_enabled = False

     clip = None

+    optimization_method = None

 

     embedding_db = modules.textual_inversion.textual_inversion.EmbeddingDatabase(cmd_opts.embeddings_dir)

 

@@ -94,7 +104,7 @@ class StableDiffusionModelHijack:
             m.cond_stage_model.model.token_embedding = EmbeddingsWithFixes(m.cond_stage_model.model.token_embedding, self)

             m.cond_stage_model = sd_hijack_open_clip.FrozenOpenCLIPEmbedderWithCustomWords(m.cond_stage_model, self)

 

-        apply_optimizations()

+        self.optimization_method = apply_optimizations()

 

         self.clip = m.cond_stage_model