From 3327244da277a806e05c596f94248e62770e2103 Mon Sep 17 00:00:00 2001
From: cmdr2 <secondary.cmdr2@gmail.com>
Date: Wed, 9 Oct 2024 11:43:45 +0530
Subject: [PATCH] Scheduler selection in the UI; Store Scheduler and Distilled
 Guidance in metadata

---
 ui/easydiffusion/backends/sdkit_common.py |  1 +
 ui/easydiffusion/backends/webui/impl.py   |  3 ++-
 ui/easydiffusion/types.py                 |  1 +
 ui/easydiffusion/utils/save_utils.py      |  2 ++
 ui/index.html                             | 20 ++++++++++++++++++++
 ui/media/js/auto-save.js                  |  2 ++
 ui/media/js/dnd.js                        | 11 +++++++++++
 ui/media/js/main.js                       |  9 +++++++++
 8 files changed, 48 insertions(+), 1 deletion(-)

diff --git a/ui/easydiffusion/backends/sdkit_common.py b/ui/easydiffusion/backends/sdkit_common.py
index bda1414c..1feae994 100644
--- a/ui/easydiffusion/backends/sdkit_common.py
+++ b/ui/easydiffusion/backends/sdkit_common.py
@@ -82,6 +82,7 @@ def generate_images(
     callback=None,
     controlnet_filter=None,
     distilled_guidance_scale: float = 3.5,
+    scheduler_name: str = "simple",
     output_type="pil",
     **req,
 ):
diff --git a/ui/easydiffusion/backends/webui/impl.py b/ui/easydiffusion/backends/webui/impl.py
index 6c8c66bd..2853e8fb 100644
--- a/ui/easydiffusion/backends/webui/impl.py
+++ b/ui/easydiffusion/backends/webui/impl.py
@@ -160,6 +160,7 @@ def generate_images(
     preserve_init_image_color_profile=False,
     strict_mask_border=False,
     sampler_name: str = "euler_a",
+    scheduler_name: str = "simple",
     hypernetwork_strength: float = 0,
     tiling=None,
     lora_alpha: Union[float, List[float]] = 0,
@@ -178,7 +179,7 @@ def generate_images(
         "prompt": prompt,
         "negative_prompt": negative_prompt,
         "sampler_name": sampler_name,
-        "scheduler": "simple",
+        "scheduler": scheduler_name,
         "steps": num_inference_steps,
         "seed": seed,
         "cfg_scale": guidance_scale,
diff --git a/ui/easydiffusion/types.py b/ui/easydiffusion/types.py
index c26e4fbc..bc0ccabf 100644
--- a/ui/easydiffusion/types.py
+++ b/ui/easydiffusion/types.py
@@ -26,6 +26,7 @@ class GenerateImageRequest(BaseModel):
     strict_mask_border: bool = False
 
     sampler_name: str = None  # "ddim", "plms", "heun", "euler", "euler_a", "dpm2", "dpm2_a", "lms"
+    scheduler_name: str = None
     hypernetwork_strength: float = 0
     lora_alpha: Union[float, List[float]] = 0
     tiling: str = None  # None, "x", "y", "xy"
diff --git a/ui/easydiffusion/utils/save_utils.py b/ui/easydiffusion/utils/save_utils.py
index 216ec899..29c84f22 100644
--- a/ui/easydiffusion/utils/save_utils.py
+++ b/ui/easydiffusion/utils/save_utils.py
@@ -34,10 +34,12 @@ TASK_TEXT_MAPPING = {
     "control_alpha": "ControlNet Strength",
     "use_vae_model": "VAE model",
     "sampler_name": "Sampler",
+    "scheduler_name": "Scheduler",
     "width": "Width",
     "height": "Height",
     "num_inference_steps": "Steps",
     "guidance_scale": "Guidance Scale",
+    "distilled_guidance_scale": "Distilled Guidance",
     "prompt_strength": "Prompt Strength",
     "use_lora_model": "LoRA model",
     "lora_alpha": "LoRA Strength",
diff --git a/ui/index.html b/ui/index.html
index 1284d785..72abe26c 100644
--- a/ui/index.html
+++ b/ui/index.html
@@ -337,6 +337,26 @@
                         <a href="https://github.com/easydiffusion/easydiffusion/wiki/How-to-Use#samplers" target="_blank"><i class="fa-solid fa-circle-question help-btn"><span class="simple-tooltip top-left">Click to learn more about samplers</span></i></a>
                     </td></tr>
                     <tr class="pl-5 warning-label displayNone" id="fluxSamplerWarning"><td></td><td>Please avoid 'Euler Ancestral' with Flux!</td></tr>
+                    <tr id="schedulerSelection" class="pl-5 gated-feature" data-feature-keys="backend_webui"><td><label for="scheduler_name">Scheduler:</label></td><td>
+                        <select id="scheduler_name" name="scheduler_name">
+                            <option value="automatic">Automatic</option>
+                            <option value="uniform">Uniform</option>
+                            <option value="karras">Karras</option>
+                            <option value="exponential">Exponential</option>
+                            <option value="polyexponential">Polyexponential</option>
+                            <option value="sgm_uniform">SGM Uniform</option>
+                            <option value="kl_optimal">KL Optimal</option>
+                            <option value="align_your_steps">Align Your Steps</option>
+                            <option value="simple" selected>Simple</option>
+                            <option value="normal">Normal</option>
+                            <option value="ddim">DDIM</option>
+                            <option value="beta">Beta</option>
+                            <option value="turbo">Turbo</option>
+                            <option value="align_your_steps_GITS">Align Your Steps GITS</option>
+                            <option value="align_your_steps_11">Align Your Steps 11</option>
+                            <option value="align_your_steps_32">Align Your Steps 32</option>
+                        </select>
+                    </td></tr>
                     <tr class="pl-5"><td><label>Image Size: </label></td><td id="image-size-options">
                         <select id="width" name="width" value="512">
                             <option value="128">128</option>
diff --git a/ui/media/js/auto-save.js b/ui/media/js/auto-save.js
index 4ed28b89..3c064352 100644
--- a/ui/media/js/auto-save.js
+++ b/ui/media/js/auto-save.js
@@ -16,10 +16,12 @@ const SETTINGS_IDS_LIST = [
     "clip_skip",
     "vae_model",
     "sampler_name",
+    "scheduler_name",
     "width",
     "height",
     "num_inference_steps",
     "guidance_scale",
+    "distilled_guidance_scale",
     "prompt_strength",
     "tiling",
     "output_format",
diff --git a/ui/media/js/dnd.js b/ui/media/js/dnd.js
index 5cb517fe..9c9c246e 100644
--- a/ui/media/js/dnd.js
+++ b/ui/media/js/dnd.js
@@ -131,6 +131,15 @@ const TASK_MAPPING = {
         readUI: () => parseFloat(guidanceScaleField.value),
         parse: (val) => parseFloat(val),
     },
+    distilled_guidance_scale: {
+        name: "Distilled Guidance",
+        setUI: (distilled_guidance_scale) => {
+            distilledGuidanceScaleField.value = distilled_guidance_scale
+            updateDistilledGuidanceScaleSlider()
+        },
+        readUI: () => parseFloat(distilledGuidanceScaleField.value),
+        parse: (val) => parseFloat(val),
+    },
     prompt_strength: {
         name: "Prompt Strength",
         setUI: (prompt_strength) => {
@@ -590,11 +599,13 @@ const TASK_TEXT_MAPPING = {
     seed: "Seed",
     num_inference_steps: "Steps",
     guidance_scale: "Guidance Scale",
+    distilled_guidance_scale: "Distilled Guidance",
     prompt_strength: "Prompt Strength",
     use_face_correction: "Use Face Correction",
     use_upscale: "Use Upscaling",
     upscale_amount: "Upscale By",
     sampler_name: "Sampler",
+    scheduler_name: "Scheduler",
     negative_prompt: "Negative Prompt",
     use_stable_diffusion_model: "Stable Diffusion model",
     use_hypernetwork_model: "Hypernetwork model",
diff --git a/ui/media/js/main.js b/ui/media/js/main.js
index 3f372911..d1e6be56 100644
--- a/ui/media/js/main.js
+++ b/ui/media/js/main.js
@@ -12,6 +12,10 @@ const taskConfigSetup = {
         seed: { value: ({ seed }) => seed, label: "Seed" },
         dimensions: { value: ({ reqBody }) => `${reqBody?.width}x${reqBody?.height}`, label: "Dimensions" },
         sampler_name: "Sampler",
+        scheduler_name: {
+            label: "Scheduler",
+            visible: ({ reqBody }) => reqBody?.scheduler_name,
+        },
         num_inference_steps: "Inference Steps",
         guidance_scale: "Guidance Scale",
         distilled_guidance_scale: {
@@ -119,6 +123,8 @@ let promptStrengthSlider = document.querySelector("#prompt_strength_slider")
 let promptStrengthField = document.querySelector("#prompt_strength")
 let samplerField = document.querySelector("#sampler_name")
 let samplerSelectionContainer = document.querySelector("#samplerSelection")
+let schedulerField = document.querySelector("#scheduler_name")
+let schedulerSelectionContainer = document.querySelector("#schedulerSelection")
 let useFaceCorrectionField = document.querySelector("#use_face_correction")
 let gfpganModelField = new ModelDropdown(document.querySelector("#gfpgan_model"), ["gfpgan", "codeformer"], "", false)
 let useUpscalingField = document.querySelector("#use_upscale")
@@ -1431,6 +1437,9 @@ function getCurrentUserRequest() {
     if (stableDiffusionModelField.value.toLowerCase().includes("flux")) {
         newTask.reqBody.distilled_guidance_scale = parseFloat(distilledGuidanceScaleField.value)
     }
+    if (schedulerSelectionContainer.style.display !== "none") {
+        newTask.reqBody.scheduler_name = schedulerField.value
+    }
 
     return newTask
 }