add k_dpmpp_2_a and k_dpmpp_2 solvers options (#1389)

* add k_dpmpp_2_a and k_dpmpp_2 solvers options * update frontend Co-authored-by: Victor <[email protected]> Co-authored-by: Lincoln Stein <[email protected]>
invoke-ai · Nov 30, 2022 · 2443e5d · 2443e5d
1 parent a9aa4e4
commit 2443e5d
Show file tree

Hide file tree

Showing 12 changed files with 21 additions and 1,127 deletions.
diff --git a/backend/modules/parameters.py b/backend/modules/parameters.py
@@ -5,6 +5,8 @@
     "ddim",
     "k_dpm_2_a",
     "k_dpm_2",
+    "k_dpmpp_2_a",
+    "k_dpmpp_2",
     "k_euler_a",
     "k_euler",
     "k_heun",

diff --git a/frontend/dist/assets/index.2b7cd976.js b/frontend/dist/assets/index.2b7cd976.js
diff --git a/frontend/dist/assets/index.678a45f6.js b/frontend/dist/assets/index.678a45f6.js
diff --git a/frontend/dist/index.html b/frontend/dist/index.html
@@ -6,7 +6,7 @@
   <meta name="viewport" content="width=device-width, initial-scale=1.0" />
   <title>InvokeAI - A Stable Diffusion Toolkit</title>
   <link rel="shortcut icon" type="icon" href="./assets/favicon.0d253ced.ico" />
-  <script type="module" crossorigin src="./assets/index.2b7cd976.js"></script>
+  <script type="module" crossorigin src="./assets/index.f1b8da92.js"></script>
   <link rel="stylesheet" href="./assets/index.f999e69e.css">
 </head>
 

diff --git a/frontend/src/app/constants.ts b/frontend/src/app/constants.ts
@@ -9,6 +9,8 @@ export const SAMPLERS: Array<string> = [
   'k_lms',
   'k_dpm_2',
   'k_dpm_2_a',
+  'k_dpmpp_2',
+  'k_dpmpp_2_a',
   'k_euler',
   'k_euler_a',
   'k_heun',

diff --git a/frontend/src/app/invokeai.d.ts b/frontend/src/app/invokeai.d.ts
@@ -46,6 +46,8 @@ export declare type CommonGeneratedImageMetadata = {
     | 'ddim'
     | 'k_dpm_2_a'
     | 'k_dpm_2'
+    | 'k_dpmpp_2_a'
+    | 'k_dpmpp_2'
     | 'k_euler_a'
     | 'k_euler'
     | 'k_heun'

diff --git a/ldm/generate.py b/ldm/generate.py
@@ -120,7 +120,7 @@ def new_func(*args, **kw):
           safety_checker:bool = activate safety checker [False]
 
           # this value is sticky and maintained between generation calls
-          sampler_name:str  = ['ddim', 'k_dpm_2_a', 'k_dpm_2', 'k_euler_a', 'k_euler', 'k_heun', 'k_lms', 'plms']  // k_lms
+          sampler_name:str  = ['ddim', 'k_dpm_2_a', 'k_dpm_2', 'k_dpmpp_2', 'k_dpmpp_2_a', 'k_euler_a', 'k_euler', 'k_heun', 'k_lms', 'plms']  // k_lms
 
           # these are deprecated - use conf and model instead
           weights     = path to model weights ('models/ldm/stable-diffusion-v1/model.ckpt')
@@ -974,6 +974,10 @@ def _set_sampler(self):
             self.sampler = KSampler(self.model, 'dpm_2_ancestral', device=self.device)
         elif self.sampler_name == 'k_dpm_2':
             self.sampler = KSampler(self.model, 'dpm_2', device=self.device)
+        elif self.sampler_name == 'k_dpmpp_2_a':
+            self.sampler = KSampler(self.model, 'dpmpp_2s_ancestral', device=self.device)
+        elif self.sampler_name == 'k_dpmpp_2':
+            self.sampler = KSampler(self.model, 'dpmpp_2m', device=self.device)
         elif self.sampler_name == 'k_euler_a':
             self.sampler = KSampler(self.model, 'euler_ancestral', device=self.device)
         elif self.sampler_name == 'k_euler':

diff --git a/ldm/invoke/args.py b/ldm/invoke/args.py
@@ -101,6 +101,8 @@
     'ddim',
     'k_dpm_2_a',
     'k_dpm_2',
+    'k_dpmpp_2_a',
+    'k_dpmpp_2',
     'k_euler_a',
     'k_euler',
     'k_heun',

diff --git a/ldm/models/diffusion/ddpm.py b/ldm/models/diffusion/ddpm.py
@@ -1490,7 +1490,7 @@ def p_losses(self, x_start, cond, t, noise=None):
         )
         loss_dict.update({f'{prefix}/loss_simple': loss_simple.mean()})
 
-        logvar_t = self.logvar[t].to(self.device)
+        logvar_t = self.logvar[t.item()].to(self.device)
         loss = loss_simple / torch.exp(logvar_t) + logvar_t
         # loss = loss_simple / torch.exp(self.logvar) + self.logvar
         if self.learn_logvar:

diff --git a/scripts/legacy_api.py b/scripts/legacy_api.py
@@ -385,6 +385,8 @@ def write_log_message(results, log_path):
     'ddim',
     'k_dpm_2_a',
     'k_dpm_2',
+    'k_dpmpp_2_a',
+    'k_dpmpp_2',
     'k_euler_a',
     'k_euler',
     'k_heun',

diff --git a/static/dream_web/index.html b/static/dream_web/index.html
@@ -50,6 +50,8 @@ <h1>Stable Diffusion Dream Server</h1>
           <option value="k_lms" selected>KLMS</option>
           <option value="k_dpm_2">KDPM_2</option>
           <option value="k_dpm_2_a">KDPM_2A</option>
+          <option value="k_dpmpp_2">KDPMPP_2</option>
+          <option value="k_dpmpp_2_a">KDPMPP_2A</option>
           <option value="k_euler">KEULER</option>
           <option value="k_euler_a">KEULER_A</option>
           <option value="k_heun">KHEUN</option>

diff --git a/static/legacy_web/index.html b/static/legacy_web/index.html
@@ -39,6 +39,8 @@ <h1>Stable Diffusion Dream Server</h1>
             <option value="k_lms" selected>KLMS</option>
             <option value="k_dpm_2">KDPM_2</option>
             <option value="k_dpm_2_a">KDPM_2A</option>
+            <option value="k_dpmpp_2">KDPMPP_2</option>
+            <option value="k_dpmpp_2_a">KDPMPP_2A</option>
             <option value="k_euler">KEULER</option>
 	    <option value="k_euler_a">KEULER_A</option>
             <option value="k_heun">KHEUN</option>