Add UniPC / Euler Karras / DPMPP_2 Karras / DEIS / DDPM Schedulers

invokeai/app/invocations/generate.py

@@ -52,7 +52,7 @@ class TextToImageInvocation(BaseInvocation, SDImageInvocation):
     width:       int = Field(default=512, multiple_of=8, gt=0, description="The width of the resulting image", )
     height:      int = Field(default=512, multiple_of=8, gt=0, description="The height of the resulting image", )
     cfg_scale: float = Field(default=7.5, ge=1, description="The Classifier-Free Guidance, higher values may result in a result closer to the prompt", )
-    scheduler: SAMPLER_NAME_VALUES = Field(default="k_lms", description="The scheduler to use" )
+    scheduler: SAMPLER_NAME_VALUES = Field(default="lms", description="The scheduler to use" )
     model:       str = Field(default="", description="The model to use (currently ignored)")
     # fmt: on
 

invokeai/app/invocations/latent.py

@@ -17,6 +17,7 @@
 from ...backend.image_util.seamless import configure_model_padding
 from ...backend.prompting.conditioning import get_uc_and_c_and_ec
 from ...backend.stable_diffusion.diffusers_pipeline import ConditioningData, StableDiffusionGeneratorPipeline, image_resized_to_grid_as_tensor
+from ...backend.stable_diffusion.schedulers import SCHEDULER_MAP
 from .baseinvocation import BaseInvocation, BaseInvocationOutput, InvocationContext, InvocationConfig
 import numpy as np
 from ..services.image_storage import ImageType
@@ -52,29 +53,20 @@ class NoiseOutput(BaseInvocationOutput):
     #fmt: on
 
 
-# TODO: this seems like a hack
-scheduler_map = dict(
-    ddim=diffusers.DDIMScheduler,
-    dpmpp_2=diffusers.DPMSolverMultistepScheduler,
-    k_dpm_2=diffusers.KDPM2DiscreteScheduler,
-    k_dpm_2_a=diffusers.KDPM2AncestralDiscreteScheduler,
-    k_dpmpp_2=diffusers.DPMSolverMultistepScheduler,
-    k_euler=diffusers.EulerDiscreteScheduler,
-    k_euler_a=diffusers.EulerAncestralDiscreteScheduler,
-    k_heun=diffusers.HeunDiscreteScheduler,
-    k_lms=diffusers.LMSDiscreteScheduler,
-    plms=diffusers.PNDMScheduler,
-)
-
-
 SAMPLER_NAME_VALUES = Literal[
-    tuple(list(scheduler_map.keys()))
+    tuple(list(SCHEDULER_MAP.keys()))
 ]
 
 
 def get_scheduler(scheduler_name:str, model: StableDiffusionGeneratorPipeline)->Scheduler:
-    scheduler_class = scheduler_map.get(scheduler_name,'ddim')
-    scheduler = scheduler_class.from_config(model.scheduler.config)
+    scheduler_class, scheduler_extra_config = SCHEDULER_MAP.get(scheduler_name, SCHEDULER_MAP['ddim'])
+
+    scheduler_config = model.scheduler.config
+    if "_backup" in scheduler_config:
+        scheduler_config = scheduler_config["_backup"]
+    scheduler_config = {**scheduler_config, **scheduler_extra_config, "_backup": scheduler_config}
+    scheduler = scheduler_class.from_config(scheduler_config)
+
     # hack copied over from generate.py
     if not hasattr(scheduler, 'uses_inpainting_model'):
         scheduler.uses_inpainting_model = lambda: False
@@ -148,7 +140,7 @@ class TextToLatentsInvocation(BaseInvocation):
     noise: Optional[LatentsField] = Field(description="The noise to use")
     steps:       int = Field(default=10, gt=0, description="The number of steps to use to generate the image")
     cfg_scale: float = Field(default=7.5, gt=0, description="The Classifier-Free Guidance, higher values may result in a result closer to the prompt", )
-    scheduler: SAMPLER_NAME_VALUES = Field(default="k_lms", description="The scheduler to use" )
+    scheduler: SAMPLER_NAME_VALUES = Field(default="lms", description="The scheduler to use" )
     model:       str = Field(default="", description="The model to use (currently ignored)")
     seamless:   bool = Field(default=False, description="Whether or not to generate an image that can tile without seams", )
     seamless_axes: str = Field(default="", description="The axes to tile the image on, 'x' and/or 'y'")
@@ -216,7 +208,7 @@ def get_conditioning_data(self, context: InvocationContext, model: StableDiffusi
                 h_symmetry_time_pct=None,#h_symmetry_time_pct,
                 v_symmetry_time_pct=None#v_symmetry_time_pct,
             ),
-        ).add_scheduler_args_if_applicable(model.scheduler, eta=None)#ddim_eta)
+        ).add_scheduler_args_if_applicable(model.scheduler, eta=0.0)#ddim_eta)
         return conditioning_data
 
 
@@ -293,11 +285,7 @@ def step_callback(state: PipelineIntermediateState):
             latent, device=model.device, dtype=latent.dtype
         )
 
-        timesteps, _ = model.get_img2img_timesteps(
-            self.steps,
-            self.strength,
-            device=model.device,
-        )
+        timesteps, _ = model.get_img2img_timesteps(self.steps, self.strength)
 
         result_latents, result_attention_map_saver = model.latents_from_embeddings(
             latents=initial_latents,

invokeai/backend/args.py

@@ -108,17 +108,20 @@
 
 SAMPLER_CHOICES = [
     "ddim",
-    "k_dpm_2_a",
-    "k_dpm_2",
-    "k_dpmpp_2_a",
-    "k_dpmpp_2",
-    "k_euler_a",
-    "k_euler",
-    "k_heun",
-    "k_lms",
-    "plms",
-    # diffusers:
+    "ddpm",
+    "deis",
+    "lms",
     "pndm",
+    "heun",
+    "euler",
+    "euler_k",
+    "euler_a",
+    "kdpm_2",
+    "kdpm_2_a",
+    "dpmpp_2s",
+    "dpmpp_2m",
+    "dpmpp_2m_k",
+    "unipc",
 ]
 
 PRECISION_CHOICES = [
@@ -631,7 +634,7 @@ def _create_arg_parser(self):
             choices=SAMPLER_CHOICES,
             metavar="SAMPLER_NAME",
             help=f'Set the default sampler. Supported samplers: {", ".join(SAMPLER_CHOICES)}',
-            default="k_lms",
+            default="lms",
         )
         render_group.add_argument(
             "--log_tokenization",

invokeai/backend/generate.py

@@ -37,6 +37,7 @@
 from .prompting import get_uc_and_c_and_ec
 from .prompting.conditioning import log_tokenization
 from .stable_diffusion import HuggingFaceConceptsLibrary
+from .stable_diffusion.schedulers import SCHEDULER_MAP
 from .util import choose_precision, choose_torch_device
 
 def fix_func(orig):
@@ -141,7 +142,7 @@ def __init__(
         model=None,
         conf="configs/models.yaml",
         embedding_path=None,
-        sampler_name="k_lms",
+        sampler_name="lms",
         ddim_eta=0.0,  # deterministic
         full_precision=False,
         precision="auto",
@@ -1047,29 +1048,12 @@ def is_legacy_model(self, model_name) -> bool:
     def _set_scheduler(self):
         default = self.model.scheduler
 
-        # See https://github.com/huggingface/diffusers/issues/277#issuecomment-1371428672
-        scheduler_map = dict(
-            ddim=diffusers.DDIMScheduler,
-            dpmpp_2=diffusers.DPMSolverMultistepScheduler,
-            k_dpm_2=diffusers.KDPM2DiscreteScheduler,
-            k_dpm_2_a=diffusers.KDPM2AncestralDiscreteScheduler,
-            # DPMSolverMultistepScheduler is technically not `k_` anything, as it is neither
-            # the k-diffusers implementation nor included in EDM (Karras 2022), but we can
-            # provide an alias for compatibility.
-            k_dpmpp_2=diffusers.DPMSolverMultistepScheduler,
-            k_euler=diffusers.EulerDiscreteScheduler,
-            k_euler_a=diffusers.EulerAncestralDiscreteScheduler,
-            k_heun=diffusers.HeunDiscreteScheduler,
-            k_lms=diffusers.LMSDiscreteScheduler,
-            plms=diffusers.PNDMScheduler,
-        )
-
-        if self.sampler_name in scheduler_map:
-            sampler_class = scheduler_map[self.sampler_name]
+        if self.sampler_name in SCHEDULER_MAP:
+            sampler_class, sampler_extra_config = SCHEDULER_MAP[self.sampler_name]
             msg = (
                 f"Setting Sampler to {self.sampler_name} ({sampler_class.__name__})"
             )
-            self.sampler = sampler_class.from_config(self.model.scheduler.config)
+            self.sampler = sampler_class.from_config({**self.model.scheduler.config, **sampler_extra_config})
         else:
             msg = (
                 f" Unsupported Sampler: {self.sampler_name} "+

invokeai/backend/generator/base.py

@@ -31,6 +31,7 @@
 from ..safety_checker import SafetyChecker
 from ..prompting.conditioning import get_uc_and_c_and_ec
 from ..stable_diffusion.diffusers_pipeline import StableDiffusionGeneratorPipeline
+from ..stable_diffusion.schedulers import SCHEDULER_MAP
 
 downsampling = 8
 
@@ -71,19 +72,6 @@ class InvokeAIGeneratorOutput:
 # we are interposing a wrapper around the original Generator classes so that
 # old code that calls Generate will continue to work.
 class InvokeAIGenerator(metaclass=ABCMeta):
-    scheduler_map = dict(
-        ddim=diffusers.DDIMScheduler,
-        dpmpp_2=diffusers.DPMSolverMultistepScheduler,
-        k_dpm_2=diffusers.KDPM2DiscreteScheduler,
-        k_dpm_2_a=diffusers.KDPM2AncestralDiscreteScheduler,
-        k_dpmpp_2=diffusers.DPMSolverMultistepScheduler,
-        k_euler=diffusers.EulerDiscreteScheduler,
-        k_euler_a=diffusers.EulerAncestralDiscreteScheduler,
-        k_heun=diffusers.HeunDiscreteScheduler,
-        k_lms=diffusers.LMSDiscreteScheduler,
-        plms=diffusers.PNDMScheduler,
-    )
-
     def __init__(self,
                  model_info: dict,
                  params: InvokeAIGeneratorBasicParams=InvokeAIGeneratorBasicParams(),
@@ -175,14 +163,20 @@ def schedulers(self)->List[str]:
         '''
         Return list of all the schedulers that we currently handle.
         '''
-        return list(self.scheduler_map.keys())
+        return list(SCHEDULER_MAP.keys())
 
     def load_generator(self, model: StableDiffusionGeneratorPipeline, generator_class: Type[Generator]):
         return generator_class(model, self.params.precision)
 
     def get_scheduler(self, scheduler_name:str, model: StableDiffusionGeneratorPipeline)->Scheduler:
-        scheduler_class = self.scheduler_map.get(scheduler_name,'ddim')
-        scheduler = scheduler_class.from_config(model.scheduler.config)
+        scheduler_class, scheduler_extra_config = SCHEDULER_MAP.get(scheduler_name, SCHEDULER_MAP['ddim'])
+
+        scheduler_config = model.scheduler.config
+        if "_backup" in scheduler_config:
+            scheduler_config = scheduler_config["_backup"]
+        scheduler_config = {**scheduler_config, **scheduler_extra_config, "_backup": scheduler_config}
+        scheduler = scheduler_class.from_config(scheduler_config)
+
         # hack copied over from generate.py
         if not hasattr(scheduler, 'uses_inpainting_model'):
             scheduler.uses_inpainting_model = lambda: False

invokeai/backend/model_management/convert_ckpt_to_diffusers.py

@@ -47,6 +47,7 @@
     LDMTextToImagePipeline,
     LMSDiscreteScheduler,
     PNDMScheduler,
+    UniPCMultistepScheduler,
     StableDiffusionPipeline,
     UNet2DConditionModel,
 )
@@ -1209,6 +1210,8 @@ def load_pipeline_from_original_stable_diffusion_ckpt(
             scheduler = EulerAncestralDiscreteScheduler.from_config(scheduler.config)
         elif scheduler_type == "dpm":
             scheduler = DPMSolverMultistepScheduler.from_config(scheduler.config)
+        elif scheduler_type == 'unipc':
+            scheduler = UniPCMultistepScheduler.from_config(scheduler.config)
         elif scheduler_type == "ddim":
             scheduler = scheduler
         else:

invokeai/backend/stable_diffusion/diffusers_pipeline.py

@@ -509,10 +509,13 @@ def latents_from_embeddings(
         run_id=None,
         callback: Callable[[PipelineIntermediateState], None] = None,
     ) -> tuple[torch.Tensor, Optional[AttentionMapSaver]]:
+        if self.scheduler.config.get("cpu_only", False):
+            scheduler_device = torch.device('cpu')
+        else:
+            scheduler_device = self._model_group.device_for(self.unet)
+
         if timesteps is None:
-            self.scheduler.set_timesteps(
-                num_inference_steps, device=self._model_group.device_for(self.unet)
-            )
+            self.scheduler.set_timesteps(num_inference_steps, device=scheduler_device)
             timesteps = self.scheduler.timesteps
         infer_latents_from_embeddings = GeneratorToCallbackinator(
             self.generate_latents_from_embeddings, PipelineIntermediateState
@@ -725,12 +728,8 @@ def img2img_from_latents_and_embeddings(
         noise: torch.Tensor,
         run_id=None,
         callback=None,
-    ) -> InvokeAIStableDiffusionPipelineOutput:
-        timesteps, _ = self.get_img2img_timesteps(
-            num_inference_steps,
-            strength,
-            device=self._model_group.device_for(self.unet),
-        )
+    ) -> InvokeAIStableDiffusionPipelineOutput:        
+        timesteps, _ = self.get_img2img_timesteps(num_inference_steps, strength)
         result_latents, result_attention_maps = self.latents_from_embeddings(
             latents=initial_latents if strength < 1.0 else torch.zeros_like(
                 initial_latents, device=initial_latents.device, dtype=initial_latents.dtype
@@ -756,13 +755,19 @@ def img2img_from_latents_and_embeddings(
             return self.check_for_safety(output, dtype=conditioning_data.dtype)
 
     def get_img2img_timesteps(
-        self, num_inference_steps: int, strength: float, device
+        self, num_inference_steps: int, strength: float, device=None
     ) -> (torch.Tensor, int):
         img2img_pipeline = StableDiffusionImg2ImgPipeline(**self.components)
         assert img2img_pipeline.scheduler is self.scheduler
-        img2img_pipeline.scheduler.set_timesteps(num_inference_steps, device=device)
+
+        if self.scheduler.config.get("cpu_only", False):
+            scheduler_device = torch.device('cpu')
+        else:
+            scheduler_device = self._model_group.device_for(self.unet)
+
+        img2img_pipeline.scheduler.set_timesteps(num_inference_steps, device=scheduler_device)
         timesteps, adjusted_steps = img2img_pipeline.get_timesteps(
-            num_inference_steps, strength, device=device
+            num_inference_steps, strength, device=scheduler_device
         )
         # Workaround for low strength resulting in zero timesteps.
         # TODO: submit upstream fix for zero-step img2img
@@ -796,9 +801,7 @@ def inpaint_from_embeddings(
         if init_image.dim() == 3:
             init_image = init_image.unsqueeze(0)
 
-        timesteps, _ = self.get_img2img_timesteps(
-            num_inference_steps, strength, device=device
-        )
+        timesteps, _ = self.get_img2img_timesteps(num_inference_steps, strength)
 
         # 6. Prepare latent variables
         # can't quite use upstream StableDiffusionImg2ImgPipeline.prepare_latents

invokeai/backend/stable_diffusion/schedulers/__init__.py

@@ -0,0 +1 @@
+from .schedulers import SCHEDULER_MAP

invokeai/backend/stable_diffusion/schedulers/schedulers.py

@@ -0,0 +1,22 @@
+from diffusers import DDIMScheduler, DPMSolverMultistepScheduler, KDPM2DiscreteScheduler, \
+    KDPM2AncestralDiscreteScheduler, EulerDiscreteScheduler, EulerAncestralDiscreteScheduler, \
+    HeunDiscreteScheduler, LMSDiscreteScheduler, PNDMScheduler, UniPCMultistepScheduler, \
+    DPMSolverSinglestepScheduler, DEISMultistepScheduler, DDPMScheduler
+
+SCHEDULER_MAP = dict(
+    ddim=(DDIMScheduler, dict()),
+    ddpm=(DDPMScheduler, dict()),
+    deis=(DEISMultistepScheduler, dict()),
+    lms=(LMSDiscreteScheduler, dict()),
+    pndm=(PNDMScheduler, dict()),
+    heun=(HeunDiscreteScheduler, dict()),
+    euler=(EulerDiscreteScheduler, dict(use_karras_sigmas=False)),
+    euler_k=(EulerDiscreteScheduler, dict(use_karras_sigmas=True)),
+    euler_a=(EulerAncestralDiscreteScheduler, dict()),
+    kdpm_2=(KDPM2DiscreteScheduler, dict()),
+    kdpm_2_a=(KDPM2AncestralDiscreteScheduler, dict()),
+    dpmpp_2s=(DPMSolverSinglestepScheduler, dict()),
+    dpmpp_2m=(DPMSolverMultistepScheduler, dict(use_karras_sigmas=False)),
+    dpmpp_2m_k=(DPMSolverMultistepScheduler, dict(use_karras_sigmas=True)),
+    unipc=(UniPCMultistepScheduler, dict(cpu_only=True))
+)

invokeai/backend/web/modules/parameters.py

@@ -4,17 +4,20 @@
 
 SAMPLER_CHOICES = [
     "ddim",
-    "k_dpm_2_a",
-    "k_dpm_2",
-    "k_dpmpp_2_a",
-    "k_dpmpp_2",
-    "k_euler_a",
-    "k_euler",
-    "k_heun",
-    "k_lms",
-    "plms",
-    # diffusers:
+    "ddpm",
+    "deis",
+    "lms",
     "pndm",
+    "heun",
+    "euler",
+    "euler_k",
+    "euler_a",
+    "kdpm_2",
+    "kdpm_2_a",
+    "dpmpp_2s",
+    "dpmpp_2m",
+    "dpmpp_2m_k",
+    "unipc",
 ]