From ccf9749df5add45052e909cd062ab66ff359c6a5 Mon Sep 17 00:00:00 2001
From: Bryce Drennan <github20210803@accounts.brycedrennan.com>
Date: Sun, 18 Dec 2022 00:00:38 -0800
Subject: [PATCH] fix: performance improvement. disable ema (#139)

A configuration `use_ema: False` was became necessary in the newer Stable Diffusion code but was missing from the 1.5 config.
---
 .gitignore                                    |  4 +++-
 imaginairy/cmds.py                            |  6 ++++++
 .../configs/stable-diffusion-v1-inpaint.yaml  |  1 +
 imaginairy/configs/stable-diffusion-v1.yaml   |  7 ++++---
 imaginairy/model_manager.py                   |  3 ++-
 imaginairy/modules/attention.py               | 19 +++++++++++--------
 6 files changed, 27 insertions(+), 13 deletions(-)

diff --git a/.gitignore b/.gitignore
index 3cfb13a..43b7559 100644
--- a/.gitignore
+++ b/.gitignore
@@ -21,4 +21,6 @@ gfpgan/**
 tests/vastai_cli.py
 /tests/test_output_local_cuda/
 /testing_support/
-.unison*
\ No newline at end of file
+.unison*
+*.kgrind
+*.pyprof
\ No newline at end of file
diff --git a/imaginairy/cmds.py b/imaginairy/cmds.py
index 62740b2..5ff7e29 100644
--- a/imaginairy/cmds.py
+++ b/imaginairy/cmds.py
@@ -298,3 +298,9 @@ aimg.add_command(imagine_cmd, name="imagine")
 
 if __name__ == "__main__":
     imagine_cmd()  # noqa
+    # from cProfile import Profile
+    # from pyprof2calltree import convert, visualize
+    # profiler = Profile()
+    # profiler.runctx("imagine_cmd.main(standalone_mode=False)", locals(), globals())
+    # convert(profiler.getstats(), 'imagine.kgrind')
+    # visualize(profiler.getstats())
diff --git a/imaginairy/configs/stable-diffusion-v1-inpaint.yaml b/imaginairy/configs/stable-diffusion-v1-inpaint.yaml
index b652ab9..9ff2419 100644
--- a/imaginairy/configs/stable-diffusion-v1-inpaint.yaml
+++ b/imaginairy/configs/stable-diffusion-v1-inpaint.yaml
@@ -16,6 +16,7 @@ model:
     monitor: val/loss_simple_ema
     scale_factor: 0.18215
     finetune_keys: null
+    use_ema: False
 
     scheduler_config: # 10000 warm-up steps
       target: ldm.lr_scheduler.LambdaLinearScheduler
diff --git a/imaginairy/configs/stable-diffusion-v1.yaml b/imaginairy/configs/stable-diffusion-v1.yaml
index 83af01c..61d6021 100644
--- a/imaginairy/configs/stable-diffusion-v1.yaml
+++ b/imaginairy/configs/stable-diffusion-v1.yaml
@@ -1,5 +1,5 @@
 model:
-  base_learning_rate: 1.0e-04
+  base_learning_rate: 1.0e-4
   target: imaginairy.modules.diffusion.ddpm.LatentDiffusion
   params:
     linear_start: 0.00085
@@ -11,10 +11,11 @@ model:
     cond_stage_key: "txt"
     image_size: 64
     channels: 4
-    cond_stage_trainable: false   # Note: different from the one we trained before
+    cond_stage_trainable: false
     conditioning_key: crossattn
     monitor: val/loss_simple_ema
     scale_factor: 0.18215
+    use_ema: False # we set this to false because this is an inference only config
 
     scheduler_config: # 10000 warm-up steps
       target: imaginairy.lr_scheduler.LambdaLinearScheduler
@@ -28,6 +29,7 @@ model:
     unet_config:
       target: imaginairy.modules.diffusion.openaimodel.UNetModel
       params:
+        use_checkpoint: True
         image_size: 32 # unused
         in_channels: 4
         out_channels: 4
@@ -39,7 +41,6 @@ model:
         use_spatial_transformer: True
         transformer_depth: 1
         context_dim: 768
-        use_checkpoint: True
         legacy: False
 
     first_stage_config:
diff --git a/imaginairy/model_manager.py b/imaginairy/model_manager.py
index cbcb357..3b3dd1c 100644
--- a/imaginairy/model_manager.py
+++ b/imaginairy/model_manager.py
@@ -80,7 +80,8 @@ def load_model_from_config(config, weights_location):
     except FileNotFoundError as e:
         if e.errno == 2:
             logger.error(
-                f'Error: "{ckpt_path}" not a valid path to model weights.\nPreconfigured models you can use: {MODEL_SHORT_NAMES}.')
+                f'Error: "{ckpt_path}" not a valid path to model weights.\nPreconfigured models you can use: {MODEL_SHORT_NAMES}.'
+            )
             sys.exit(1)
         raise e
     except RuntimeError as e:
diff --git a/imaginairy/modules/attention.py b/imaginairy/modules/attention.py
index 9d79e5b..cf7f3af 100644
--- a/imaginairy/modules/attention.py
+++ b/imaginairy/modules/attention.py
@@ -449,16 +449,19 @@ class SpatialTransformer(nn.Module):
         b, c, h, w = x.shape  # noqa
         x_in = x
         x = self.norm(x)
-        if not self.use_linear:
-            x = self.proj_in(x)
-        x = rearrange(x, "b c h w -> b (h w) c").contiguous()
         if self.use_linear:
+            x = rearrange(x, "b c h w -> b (h w) c").contiguous()
             x = self.proj_in(x)
-        for i, block in enumerate(self.transformer_blocks):
-            x = block(x, context=context[i])
-        if self.use_linear:
+            for i, block in enumerate(self.transformer_blocks):
+                x = block(x, context=context[i])
             x = self.proj_out(x)
-        x = rearrange(x, "b (h w) c -> b c h w", h=h, w=w).contiguous()
-        if not self.use_linear:
+            x = rearrange(x, "b (h w) c -> b c h w", h=h, w=w).contiguous()
+        else:
+            x = self.proj_in(x)
+            x = rearrange(x, "b c h w -> b (h w) c")
+            for i, block in enumerate(self.transformer_blocks):
+                x = block(x, context=context[i])
+            x = rearrange(x, "b (h w) c -> b c h w", h=h, w=w)
             x = self.proj_out(x)
+
         return x + x_in