From 716a69237cefb385f71105dbbf50e92d664e0f42 Mon Sep 17 00:00:00 2001
From: space-nuko <24979496+space-nuko@users.noreply.github.com>
Date: Sat, 11 Feb 2023 06:18:34 -0800
Subject: support SD2.X models

---
 modules/models/diffusion/uni_pc/sampler.py | 5 ++++-
 1 file changed, 4 insertions(+), 1 deletion(-)

(limited to 'modules/models/diffusion/uni_pc')

diff --git a/modules/models/diffusion/uni_pc/sampler.py b/modules/models/diffusion/uni_pc/sampler.py
index 0bef6eed..708a9b2b 100644
--- a/modules/models/diffusion/uni_pc/sampler.py
+++ b/modules/models/diffusion/uni_pc/sampler.py
@@ -80,10 +80,13 @@ class UniPCSampler(object):
 
         ns = NoiseScheduleVP('discrete', alphas_cumprod=self.alphas_cumprod)
 
+        # SD 1.X is "noise", SD 2.X is "v"
+        model_type = "v" if self.model.parameterization == "v" else "noise"
+
         model_fn = model_wrapper(
             lambda x, t, c: self.model.apply_model(x, t, c),
             ns,
-            model_type="noise",
+            model_type=model_type,
             guidance_type="classifier-free",
             #condition=conditioning,
             #unconditional_condition=unconditional_conditioning,
-- 
cgit v1.2.1