Add dropdown for scheduler type

1 files changed, 16 insertions, 10 deletions

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"),