Merge branch 'master' of github.com:AUTOMATIC1111/stable-diffusion-webui into gamepad

1 files changed, 43 insertions, 23 deletions

diff --git a/modules/sd_hijack.py b/modules/sd_hijack.py
index 690a9ec2..f9652d21 100644
--- a/modules/sd_hijack.py
+++ b/modules/sd_hijack.py
@@ -5,9 +5,7 @@ import modules.textual_inversion.textual_inversion
 from modules import devices, sd_hijack_optimizations, shared, sd_hijack_checkpoint

 from modules.hypernetworks import hypernetwork

 from modules.shared import cmd_opts

-from modules import sd_hijack_clip, sd_hijack_open_clip, sd_hijack_unet

-

-from modules.sd_hijack_optimizations import invokeAI_mps_available

+from modules import sd_hijack_clip, sd_hijack_open_clip, sd_hijack_unet, sd_hijack_xlmr, xlmr

 

 import ldm.modules.attention

 import ldm.modules.diffusionmodules.model

@@ -35,26 +33,34 @@ 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_sub_quad_attention:

+        print("Applying sub-quadratic cross attention optimization.")

+        ldm.modules.attention.CrossAttention.forward = sd_hijack_optimizations.sub_quad_attention_forward

+        ldm.modules.diffusionmodules.model.AttnBlock.forward = sd_hijack_optimizations.sub_quad_attnblock_forward

+        optimization_method = 'sub-quadratic'

     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

-    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

-        else:

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

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

+        optimization_method = 'V1'

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

+        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():

@@ -64,9 +70,11 @@ def undo_optimizations():
 

 

 def fix_checkpoint():

-    ldm.modules.attention.BasicTransformerBlock.forward = sd_hijack_checkpoint.BasicTransformerBlock_forward

-    ldm.modules.diffusionmodules.openaimodel.ResBlock.forward = sd_hijack_checkpoint.ResBlock_forward

-    ldm.modules.diffusionmodules.openaimodel.AttentionBlock.forward = sd_hijack_checkpoint.AttentionBlock_forward

+    """checkpoints are now added and removed in embedding/hypernet code, since torch doesn't want

+    checkpoints to be added when not training (there's a warning)"""

+

+    pass

+

 

 class StableDiffusionModelHijack:

     fixes = None

@@ -74,22 +82,31 @@ class StableDiffusionModelHijack:
     layers = None

     circular_enabled = False

     clip = None

+    optimization_method = None

+

+    embedding_db = modules.textual_inversion.textual_inversion.EmbeddingDatabase()

 

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

+    def __init__(self):

+        self.embedding_db.add_embedding_dir(cmd_opts.embeddings_dir)

 

     def hijack(self, m):

-        if type(m.cond_stage_model) == ldm.modules.encoders.modules.FrozenCLIPEmbedder:

+        if type(m.cond_stage_model) == xlmr.BertSeriesModelWithTransformation:

+            model_embeddings = m.cond_stage_model.roberta.embeddings

+            model_embeddings.token_embedding = EmbeddingsWithFixes(model_embeddings.word_embeddings, self)

+            m.cond_stage_model = sd_hijack_xlmr.FrozenXLMREmbedderWithCustomWords(m.cond_stage_model, self)

+

+        elif type(m.cond_stage_model) == ldm.modules.encoders.modules.FrozenCLIPEmbedder:

             model_embeddings = m.cond_stage_model.transformer.text_model.embeddings

             model_embeddings.token_embedding = EmbeddingsWithFixes(model_embeddings.token_embedding, self)

             m.cond_stage_model = sd_hijack_clip.FrozenCLIPEmbedderWithCustomWords(m.cond_stage_model, self)

+

         elif type(m.cond_stage_model) == ldm.modules.encoders.modules.FrozenOpenCLIPEmbedder:

             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)

 

-        self.clip = m.cond_stage_model

+        self.optimization_method = apply_optimizations()

 

-        apply_optimizations()

-        fix_checkpoint()

+        self.clip = m.cond_stage_model

 

         def flatten(el):

             flattened = [flatten(children) for children in el.children()]

@@ -101,7 +118,10 @@ class StableDiffusionModelHijack:
         self.layers = flatten(m)

 

     def undo_hijack(self, m):

-        if type(m.cond_stage_model) == sd_hijack_clip.FrozenCLIPEmbedderWithCustomWords:

+        if type(m.cond_stage_model) == xlmr.BertSeriesModelWithTransformation:

+            m.cond_stage_model = m.cond_stage_model.wrapped 

+

+        elif type(m.cond_stage_model) == sd_hijack_clip.FrozenCLIPEmbedderWithCustomWords:

             m.cond_stage_model = m.cond_stage_model.wrapped

 

             model_embeddings = m.cond_stage_model.transformer.text_model.embeddings

@@ -127,10 +147,10 @@ class StableDiffusionModelHijack:
     def clear_comments(self):

         self.comments = []

 

-    def tokenize(self, text):

-        _, remade_batch_tokens, _, _, _, token_count = self.clip.process_text([text])

-        return remade_batch_tokens[0], token_count, sd_hijack_clip.get_target_prompt_token_count(token_count)

+    def get_prompt_lengths(self, text):

+        _, token_count = self.clip.process_texts([text])

 

+        return token_count, self.clip.get_target_prompt_token_count(token_count)

 

 

 class EmbeddingsWithFixes(torch.nn.Module):