From 669314dff10d18eb64e94bd15df74cb1882be173 Mon Sep 17 00:00:00 2001
From: lvmin <lyuminzhang@outlook.com>
Date: Thu, 10 Aug 2023 10:42:36 -0700
Subject: [PATCH] i

---
 modules/core.py              |  6 +++---
 modules/default_pipeline.py  |  2 +-
 modules/samplers_advanced.py | 13 ++++++++++++-
 3 files changed, 16 insertions(+), 5 deletions(-)

diff --git a/modules/core.py b/modules/core.py
index f4db8ed0..0a9c0548 100644
--- a/modules/core.py
+++ b/modules/core.py
@@ -178,7 +178,6 @@ def ksampler_with_refiner(model, positive, negative, refiner, refiner_positive,
         noise_mask = prepare_mask(noise_mask, noise.shape, device)
 
     comfy.model_management.load_model_gpu(model)
-    real_model = model.model
 
     noise = noise.to(device)
     latent_image = latent_image.to(device)
@@ -188,8 +187,9 @@ def ksampler_with_refiner(model, positive, negative, refiner, refiner_positive,
 
     models = load_additional_models(positive, negative, model.model_dtype())
 
-    sampler = KSamplerWithRefiner(real_model, steps=steps, device=device, sampler=sampler_name, scheduler=scheduler,
-                       denoise=denoise, model_options=model.model_options)
+    sampler = KSamplerWithRefiner(model=model.model, refiner_model=refiner.model, steps=steps, device=device,
+                                  sampler=sampler_name, scheduler=scheduler,
+                                  denoise=denoise, model_options=model.model_options)
 
     samples = sampler.sample(noise, positive_copy, negative_copy, cfg=cfg, latent_image=latent_image,
                              start_step=start_step, last_step=last_step, force_full_denoise=force_full_denoise,
diff --git a/modules/default_pipeline.py b/modules/default_pipeline.py
index 3aba6651..bdb6bdf6 100644
--- a/modules/default_pipeline.py
+++ b/modules/default_pipeline.py
@@ -27,7 +27,7 @@ def process(positive_prompt, negative_prompt, width=1024, height=1024, batch_siz
         model=xl_base.unet,
         positive=positive_conditions,
         negative=negative_conditions,
-        refiner=xl_refiner,
+        refiner=xl_refiner.unet,
         refiner_positive=positive_conditions_refiner,
         refiner_negative=negative_conditions_refiner,
         refiner_switch_step=20,
diff --git a/modules/samplers_advanced.py b/modules/samplers_advanced.py
index 2318cde0..c77d672e 100644
--- a/modules/samplers_advanced.py
+++ b/modules/samplers_advanced.py
@@ -7,15 +7,26 @@ class KSamplerWithRefiner:
                 "lms", "dpm_fast", "dpm_adaptive", "dpmpp_2s_ancestral", "dpmpp_sde", "dpmpp_sde_gpu",
                 "dpmpp_2m", "dpmpp_2m_sde", "dpmpp_2m_sde_gpu", "ddim", "uni_pc", "uni_pc_bh2"]
 
-    def __init__(self, model, steps, device, sampler=None, scheduler=None, denoise=None, model_options={}):
+    def __init__(self, model, refiner_model, steps, device, sampler=None, scheduler=None, denoise=None, model_options={}):
         self.model = model
+        self.refiner_model = refiner_model
+
         self.model_denoise = CFGNoisePredictor(self.model)
+        self.refiner_model_denoise = CFGNoisePredictor(self.refiner_model)
+
         if self.model.model_type == model_base.ModelType.V_PREDICTION:
             self.model_wrap = CompVisVDenoiser(self.model_denoise, quantize=True)
         else:
             self.model_wrap = k_diffusion_external.CompVisDenoiser(self.model_denoise, quantize=True)
 
+        if self.refiner_model.model_type == model_base.ModelType.V_PREDICTION:
+            self.refiner_model_wrap = CompVisVDenoiser(self.refiner_model_denoise, quantize=True)
+        else:
+            self.refiner_model_wrap = k_diffusion_external.CompVisDenoiser(self.refiner_model_denoise, quantize=True)
+
         self.model_k = KSamplerX0Inpaint(self.model_wrap)
+        self.refiner_model_k = KSamplerX0Inpaint(self.refiner_model_wrap)
+
         self.device = device
         if scheduler not in self.SCHEDULERS:
             scheduler = self.SCHEDULERS[0]