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 <victorca25@users.noreply.github.com>
Co-authored-by: Lincoln Stein <lincoln.stein@gmail.com>

backend/modules/parameters.py

@@ -5,6 +5,8 @@ SAMPLER_CHOICES = [
     "ddim",
     "k_dpm_2_a",
     "k_dpm_2",
+    "k_dpmpp_2_a",
+    "k_dpmpp_2",
     "k_euler_a",
     "k_euler",
     "k_heun",

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>
 

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',

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'

ldm/generate.py

@@ -120,7 +120,7 @@ gr = Generate(
           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 @@ class Generate:
             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':

ldm/invoke/args.py

@@ -101,6 +101,8 @@ SAMPLER_CHOICES = [
     'ddim',
     'k_dpm_2_a',
     'k_dpm_2',
+    'k_dpmpp_2_a',
+    'k_dpmpp_2',
     'k_euler_a',
     'k_euler',
     'k_heun',

ldm/models/diffusion/ddpm.py

@@ -1490,7 +1490,7 @@ class LatentDiffusion(DDPM):
         )
         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:

scripts/legacy_api.py

@@ -385,6 +385,8 @@ SAMPLER_CHOICES = [
     'ddim',
     'k_dpm_2_a',
     'k_dpm_2',
+    'k_dpmpp_2_a',
+    'k_dpmpp_2',
     'k_euler_a',
     'k_euler',
     'k_heun',

static/dream_web/index.html

@@ -50,6 +50,8 @@
           <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>

static/legacy_web/index.html

@@ -39,6 +39,8 @@
             <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>