From 7ba9bc2fdbfae8115294962510492faafeb48573 Mon Sep 17 00:00:00 2001
From: "Alex \"mcmonkey\" Goodwin" <git_commits@alexgoodwin.dev>
Date: Sun, 18 Dec 2022 19:16:42 -0800
Subject: fix dpm2 in img2img as well

---
 modules/sd_samplers.py | 3 +++
 1 file changed, 3 insertions(+)

(limited to 'modules/sd_samplers.py')

diff --git a/modules/sd_samplers.py b/modules/sd_samplers.py
index ae3d8bfa..1a1b8919 100644
--- a/modules/sd_samplers.py
+++ b/modules/sd_samplers.py
@@ -454,6 +454,9 @@ class KDiffusionSampler:
         else:
             sigmas = self.model_wrap.get_sigmas(steps)
 
+        if self.funcname in ['sample_dpm_2_ancestral', 'sample_dpm_2']:
+            sigmas = torch.cat([sigmas[:-2], sigmas[-1:]])
+
         sigma_sched = sigmas[steps - t_enc - 1:]
         xi = x + noise * sigma_sched[0]
         
-- 
cgit v1.2.1