From f1975b0213f5be400889ec04b3891d1cb571fe20 Mon Sep 17 00:00:00 2001
From: AUTOMATIC1111 <16777216c@gmail.com>
Date: Sun, 6 Aug 2023 17:01:07 +0300
Subject: initial refiner support

---
 modules/sd_samplers_common.py | 19 ++++++++++++++++++-
 1 file changed, 18 insertions(+), 1 deletion(-)

(limited to 'modules/sd_samplers_common.py')

diff --git a/modules/sd_samplers_common.py b/modules/sd_samplers_common.py
index 39586b40..3f3e83e3 100644
--- a/modules/sd_samplers_common.py
+++ b/modules/sd_samplers_common.py
@@ -2,7 +2,7 @@ from collections import namedtuple
 import numpy as np
 import torch
 from PIL import Image
-from modules import devices, images, sd_vae_approx, sd_samplers, sd_vae_taesd, shared
+from modules import devices, images, sd_vae_approx, sd_samplers, sd_vae_taesd, shared, sd_models
 from modules.shared import opts, state
 
 SamplerData = namedtuple('SamplerData', ['name', 'constructor', 'aliases', 'options'])
@@ -127,3 +127,20 @@ def replace_torchsde_browinan():
 
 
 replace_torchsde_browinan()
+
+
+def apply_refiner(sampler):
+    completed_ratio = sampler.step / sampler.steps
+    if completed_ratio > shared.opts.sd_refiner_switch_at and shared.sd_model.sd_checkpoint_info.title != shared.opts.sd_refiner_checkpoint:
+        refiner_checkpoint_info = sd_models.get_closet_checkpoint_match(shared.opts.sd_refiner_checkpoint)
+        if refiner_checkpoint_info is None:
+            raise Exception(f'Could not find checkpoint with name {shared.opts.sd_refiner_checkpoint}')
+
+        with sd_models.SkipWritingToConfig():
+            sd_models.reload_model_weights(info=refiner_checkpoint_info)
+
+        devices.torch_gc()
+
+        sampler.update_inner_model()
+
+        sampler.p.setup_conds()
-- 
cgit v1.2.1