Patch for duplicate timesteps in DPM Solver Single Step, until huggingface/diffusers#4231 gets resolved

Author: cmdr2

sdkit/models/model_loader/stable_diffusion/__init__.py

@@ -128,6 +128,7 @@ def load_diffusers_model(context: Context, model_path, config_file_path, convert
     from sdkit.utils import gc, has_amd_gpu
 
     from diffusers.pipelines.stable_diffusion.convert_from_ckpt import download_from_original_stable_diffusion_ckpt
+    from . import diffusers_bugfixes
 
     log.info("loading on diffusers")
 

sdkit/models/model_loader/stable_diffusion/diffusers_bugfixes.py

@@ -0,0 +1,27 @@
+from diffusers.schedulers import DPMSolverSinglestepScheduler
+import numpy as np
+import torch
+
+# temporary patch, while waiting for PR: https://github.com/huggingface/diffusers/pull/4231
+
+old_set_timesteps = DPMSolverSinglestepScheduler.set_timesteps
+
+
+def set_timesteps_remove_duplicates(self, num_inference_steps: int, device=None):
+    old_set_timesteps(self, num_inference_steps, device)
+
+    timesteps = self.timesteps.cpu().detach().numpy().astype(np.int64)
+
+    # when num_inference_steps == num_train_timesteps, we can end up with
+    # duplicates in timesteps.
+    _, unique_indices = np.unique(timesteps, return_index=True)
+    timesteps = timesteps[np.sort(unique_indices)]
+
+    self.timesteps = torch.from_numpy(timesteps).to(device)
+
+    self.num_inference_steps = len(timesteps)
+
+    self.order_list = self.get_order_list(self.num_inference_steps)
+
+
+DPMSolverSinglestepScheduler.set_timesteps = set_timesteps_remove_duplicates

tests/samplers/test_samplers_txt2img.py

@@ -39,7 +39,7 @@ def sampler_test():
             width=512,
             height=512,
             sampler_name=sampler_name,
-            num_inference_steps=25,
+            num_inference_steps=50,
         )[0]
 
         expected_image = Image.open(f"{EXPECTED_DIR}/1.4-txt-{sampler_name}-42-512x512-50-cuda.png")