From 7882f76da45de7279c7db0dd17b6aca82b7ddf46 Mon Sep 17 00:00:00 2001
From: Kohaku-Blueleaf <59680068+KohakuBlueleaf@users.noreply.github.com>
Date: Mon, 22 May 2023 23:26:28 +0800
Subject: Replace karras by k_diffusion, fix gen info

---
 modules/sd_samplers_kdiffusion.py | 3 ++-
 1 file changed, 2 insertions(+), 1 deletion(-)

(limited to 'modules/sd_samplers_kdiffusion.py')

diff --git a/modules/sd_samplers_kdiffusion.py b/modules/sd_samplers_kdiffusion.py
index 441c040e..4d8f57a7 100644
--- a/modules/sd_samplers_kdiffusion.py
+++ b/modules/sd_samplers_kdiffusion.py
@@ -44,6 +44,7 @@ sampler_extra_params = {
     'sample_dpm_2': ['s_churn', 's_tmin', 's_tmax', 's_noise'],
 }
 
+k_diffusion_samplers_map = {x.name: x for x in samplers_data_k_diffusion}
 k_diffusion_scheduler = {
     'karras': k_diffusion.sampling.get_sigmas_karras,
     'exponential': k_diffusion.sampling.get_sigmas_exponential,
@@ -310,7 +311,7 @@ class KDiffusionSampler:
 
         if p.sampler_noise_scheduler_override:
             sigmas = p.sampler_noise_scheduler_override(steps)
-        elif p.enable_karras:
+        elif p.enable_custom_k_sched:
             print(p.k_sched_type, p.sigma_min, p.sigma_max, p.rho)
             sigmas_func = k_diffusion_scheduler[p.k_sched_type]
             sigmas_kwargs = {
-- 
cgit v1.2.1