Merge branch 'dev' into startup-profile

1 files changed, 28 insertions, 9 deletions

diff --git a/modules/sd_hijack.py b/modules/sd_hijack.py
index 08d31080..3b6f95ce 100644
--- a/modules/sd_hijack.py
+++ b/modules/sd_hijack.py
@@ -3,7 +3,7 @@ from torch.nn.functional import silu
 from types import MethodType

 

 import modules.textual_inversion.textual_inversion

-from modules import devices, sd_hijack_optimizations, shared, script_callbacks, errors

+from modules import devices, sd_hijack_optimizations, shared, script_callbacks, errors, sd_unet

 from modules.hypernetworks import hypernetwork

 from modules.shared import cmd_opts

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

@@ -43,11 +43,16 @@ def list_optimizers():
     optimizers.extend(new_optimizers)

 

 

-def apply_optimizations():

+def apply_optimizations(option=None):

     global current_optimizer

 

     undo_optimizations()

 

+    if len(optimizers) == 0:

+        # a script can access the model very early, and optimizations would not be filled by then

+        current_optimizer = None

+        return ''

+

     ldm.modules.diffusionmodules.model.nonlinearity = silu

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

 

@@ -55,7 +60,7 @@ def apply_optimizations():
         current_optimizer.undo()

         current_optimizer = None

 

-    selection = shared.opts.cross_attention_optimization

+    selection = option or shared.opts.cross_attention_optimization

     if selection == "Automatic" and len(optimizers) > 0:

         matching_optimizer = next(iter([x for x in optimizers if x.cmd_opt and getattr(shared.cmd_opts, x.cmd_opt, False)]), optimizers[0])

     else:

@@ -63,15 +68,19 @@ def apply_optimizations():
 

     if selection == "None":

         matching_optimizer = None

+    elif selection == "Automatic" and shared.cmd_opts.disable_opt_split_attention:

+        matching_optimizer = None

     elif matching_optimizer is None:

         matching_optimizer = optimizers[0]

 

     if matching_optimizer is not None:

-        print(f"Applying optimization: {matching_optimizer.name}")

+        print(f"Applying attention optimization: {matching_optimizer.name}... ", end='')

         matching_optimizer.apply()

+        print("done.")

         current_optimizer = matching_optimizer

         return current_optimizer.name

     else:

+        print("Disabling attention optimization")

         return ''

 

 

@@ -149,6 +158,13 @@ class StableDiffusionModelHijack:
     def __init__(self):

         self.embedding_db.add_embedding_dir(cmd_opts.embeddings_dir)

 

+    def apply_optimizations(self, option=None):

+        try:

+            self.optimization_method = apply_optimizations(option)

+        except Exception as e:

+            errors.display(e, "applying cross attention optimization")

+            undo_optimizations()

+

     def hijack(self, m):

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

             model_embeddings = m.cond_stage_model.roberta.embeddings

@@ -168,11 +184,7 @@ class StableDiffusionModelHijack:
         if m.cond_stage_key == "edit":

             sd_hijack_unet.hijack_ddpm_edit()

 

-        try:

-            self.optimization_method = apply_optimizations()

-        except Exception as e:

-            errors.display(e, "applying cross attention optimization")

-            undo_optimizations()

+        self.apply_optimizations()

 

         self.clip = m.cond_stage_model

 

@@ -185,6 +197,11 @@ class StableDiffusionModelHijack:
 

         self.layers = flatten(m)

 

+        if not hasattr(ldm.modules.diffusionmodules.openaimodel, 'copy_of_UNetModel_forward_for_webui'):

+            ldm.modules.diffusionmodules.openaimodel.copy_of_UNetModel_forward_for_webui = ldm.modules.diffusionmodules.openaimodel.UNetModel.forward

+

+        ldm.modules.diffusionmodules.openaimodel.UNetModel.forward = sd_unet.UNetModel_forward

+

     def undo_hijack(self, m):

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

             m.cond_stage_model = m.cond_stage_model.wrapped

@@ -206,6 +223,8 @@ class StableDiffusionModelHijack:
         self.layers = None

         self.clip = None

 

+        ldm.modules.diffusionmodules.openaimodel.UNetModel.forward = ldm.modules.diffusionmodules.openaimodel.copy_of_UNetModel_forward_for_webui

+

     def apply_circular(self, enable):

         if self.circular_enabled == enable:

             return