From e6269cba7fd84a76b2bd0012cb954f947a79b6a5 Mon Sep 17 00:00:00 2001
From: Kohaku-Blueleaf <59680068+KohakuBlueleaf@users.noreply.github.com>
Date: Mon, 22 May 2023 23:02:05 +0800
Subject: Add dropdown for scheduler type

---
 modules/ui.py | 26 ++++++++++++++++----------
 1 file changed, 16 insertions(+), 10 deletions(-)

(limited to 'modules/ui.py')

diff --git a/modules/ui.py b/modules/ui.py
index a65f8d85..28d4f1d1 100644
--- a/modules/ui.py
+++ b/modules/ui.py
@@ -514,9 +514,10 @@ def create_ui():
                     elif category == "karras_scheduler":
                         with FormGroup(visible=False, elem_id="txt2img_karras_scheduler") as t2i_k_sched_options:
                             with FormRow(elem_id="txt2img_karras_scheduler_row1", variant="compact"):
-                                t2i_k_sched_sigma_max = gr.Slider(minimum=0.0, maximum=0.5, step=0.05, label='sigma min', value=0.1, elem_id="txt2img_sigma_min")
-                                t2i_k_sched_sigma_min = gr.Slider(minimum=5.0, maximum=50.0, step=0.1, label='sigma max', value=10.0, elem_id="txt2img_sigma_max")
-                                t2i_k_sched_rho = gr.Slider(minimum=3.0, maximum=10.0, step=0.5, label='rho', value=7.0, elem_id="txt2img_rho")
+                                t2i_k_sched_type = gr.Dropdown(label="Type", elem_id="t2i_k_sched_type", choices=['karras', 'exponential', 'polyexponential'], value='karras')
+                                t2i_k_sched_sigma_min = gr.Slider(minimum=0.0, maximum=0.5, step=0.05, label='sigma min', value=0.1, elem_id="txt2img_sigma_min")
+                                t2i_k_sched_sigma_max = gr.Slider(minimum=5.0, maximum=50.0, step=0.1, label='sigma max', value=10.0, elem_id="txt2img_sigma_max")
+                                t2i_k_sched_rho = gr.Slider(minimum=0.5, maximum=10.0, step=0.1, label='rho', value=7.0, elem_id="txt2img_rho")
 
                     elif category == "batch":
                         if not opts.dimensions_and_batch_together:
@@ -587,8 +588,9 @@ def create_ui():
                     hr_negative_prompt,
                     override_settings,
                     t2i_enable_k_sched,
-                    t2i_k_sched_sigma_max,
+                    t2i_k_sched_type,
                     t2i_k_sched_sigma_min,
+                    t2i_k_sched_sigma_max,
                     t2i_k_sched_rho
 
                 ] + custom_inputs,
@@ -674,7 +676,8 @@ def create_ui():
                 (hr_prompt, "Hires prompt"),
                 (hr_negative_prompt, "Hires negative prompt"),
                 (hr_prompts_container, lambda d: gr.update(visible=True) if d.get("Hires prompt", "") != "" or d.get("Hires negative prompt", "") != "" else gr.update()),
-                (t2i_enable_k_sched, "Enable CustomKarras Schedule"),
+                (t2i_enable_k_sched, "Enable Custom Karras Schedule"),
+                (t2i_k_sched_type, "Karras Scheduler Type"),
                 (t2i_k_sched_sigma_max, "Karras Scheduler sigma_max"),
                 (t2i_k_sched_sigma_min, "Karras Scheduler sigma_min"),
                 (t2i_k_sched_rho, "Karras Scheduler rho"),
@@ -874,9 +877,10 @@ def create_ui():
                     elif category == "karras_scheduler":
                         with FormGroup(visible=False, elem_id="img2img_karras_scheduler") as i2i_k_sched_options:
                             with FormRow(elem_id="img2img_karras_scheduler_row1", variant="compact"):
-                                i2i_k_sched_sigma_max = gr.Slider(minimum=0.0, maximum=0.5, step=0.05, label='sigma min', value=0.1, elem_id="txt2img_sigma_min")
-                                i2i_k_sched_sigma_min = gr.Slider(minimum=5.0, maximum=50.0, step=0.1, label='sigma max', value=10.0, elem_id="txt2img_sigma_max")
-                                i2i_k_sched_rho = gr.Slider(minimum=3.0, maximum=10.0, step=0.5, label='rho', value=7.0, elem_id="txt2img_rho")
+                                i2i_k_sched_type = gr.Dropdown(label="Type", elem_id="t2i_k_sched_type", choices=['karras', 'exponential', 'polyexponential'], value='karras')
+                                i2i_k_sched_sigma_min = gr.Slider(minimum=0.0, maximum=0.5, step=0.05, label='sigma min', value=0.1, elem_id="txt2img_sigma_min")
+                                i2i_k_sched_sigma_max = gr.Slider(minimum=5.0, maximum=50.0, step=0.1, label='sigma max', value=10.0, elem_id="txt2img_sigma_max")
+                                i2i_k_sched_rho = gr.Slider(minimum=0.5, maximum=10.0, step=0.1, label='rho', value=7.0, elem_id="txt2img_rho")
 
                     elif category == "batch":
                         if not opts.dimensions_and_batch_together:
@@ -981,8 +985,9 @@ def create_ui():
                     img2img_batch_inpaint_mask_dir,
                     override_settings,
                     i2i_enable_k_sched,
-                    i2i_k_sched_sigma_max,
+                    i2i_k_sched_type,
                     i2i_k_sched_sigma_min,
+                    i2i_k_sched_sigma_max,
                     i2i_k_sched_rho
                 ] + custom_inputs,
                 outputs=[
@@ -1085,7 +1090,8 @@ def create_ui():
                 (steps, "Steps"),
                 (sampler_index, "Sampler"),
                 (restore_faces, "Face restoration"),
-                (i2i_enable_k_sched, "Enable Karras Schedule"),
+                (i2i_enable_k_sched, "Enable Custom Karras Schedule"),
+                (i2i_k_sched_type, "Karras Scheduler Type"),
                 (i2i_k_sched_sigma_max, "Karras Scheduler sigma_max"),
                 (i2i_k_sched_sigma_min, "Karras Scheduler sigma_min"),
                 (i2i_k_sched_rho, "Karras Scheduler rho"),
-- 
cgit v1.2.1