From dde51c0cefaf36857a54441f68d313ad13ed3e67 Mon Sep 17 00:00:00 2001
From: patriceac <48073125+patriceac@users.noreply.github.com>
Date: Sat, 29 Apr 2023 19:25:10 -0700
Subject: [PATCH 1/6] Toast notifications for ED

Adding support for toast notifications for use in Core and user plugins.
---
 ui/media/css/main.css | 51 ++++++++++++++++++++++++++++++++++++++
 ui/media/js/utils.js  | 57 +++++++++++++++++++++++++++++++++++++++++++
 2 files changed, 108 insertions(+)

diff --git a/ui/media/css/main.css b/ui/media/css/main.css
index 08dea664..ba513237 100644
--- a/ui/media/css/main.css
+++ b/ui/media/css/main.css
@@ -1302,3 +1302,54 @@ body.wait-pause {
 .displayNone {
     display:none !important;
 }
+
+/* TOAST NOTIFICATIONS */
+.toast-notification {
+    position: fixed;
+    bottom: 10px;
+    right: -300px;
+    width: 300px;
+    background-color: #333;
+    color: #fff;
+    padding: 10px 20px;
+    border-radius: 5px;
+    box-shadow: 0 0 10px rgba(0, 0, 0, 0.5);
+    z-index: 9999;
+    animation: slideInRight 0.5s ease forwards;
+    transition: bottom 0.5s ease; // Add a transition to smoothly reposition the toasts
+}
+
+.toast-notification-error {
+    color: red;
+}
+
+@keyframes slideInRight {
+    from {
+        right: -300px;
+    }
+    to {
+        right: 10px;
+    }
+}
+
+.toast-notification.hide {
+    animation: slideOutRight 0.5s ease forwards;
+}
+
+@keyframes slideOutRight {
+    from {
+        right: 10px;
+    }
+    to {
+        right: -300px;
+    }
+}
+
+@keyframes slideDown {
+    from {
+        bottom: 10px;
+    }
+    to {
+        bottom: 0;
+    }
+}
diff --git a/ui/media/js/utils.js b/ui/media/js/utils.js
index 6558f7d2..d1578d8e 100644
--- a/ui/media/js/utils.js
+++ b/ui/media/js/utils.js
@@ -840,3 +840,60 @@ function createTab(request) {
         }
     })
 }
+
+/* TOAST NOTIFICATIONS */
+function showToast(message, duration = 5000, error = false) {
+    const toast = document.createElement("div");
+    toast.classList.add("toast-notification");
+    if (error === true) {
+        toast.classList.add("toast-notification-error");
+    }
+    toast.innerHTML = message;
+    document.body.appendChild(toast);
+
+    // Set the position of the toast on the screen
+    const toastCount = document.querySelectorAll(".toast-notification").length;
+    const toastHeight = toast.offsetHeight;
+    const previousToastsHeight = Array.from(document.querySelectorAll(".toast-notification"))
+        .slice(0, -1) // exclude current toast
+        .reduce((totalHeight, toast) => totalHeight + toast.offsetHeight + 10, 0); // add 10 pixels for spacing
+    toast.style.bottom = `${10 + previousToastsHeight}px`;
+    toast.style.right = "10px";
+
+    // Delay the removal of the toast until animation has completed
+    const removeToast = () => {
+        toast.classList.add("hide");
+        const removeTimeoutId = setTimeout(() => {
+            toast.remove();
+            // Adjust the position of remaining toasts
+            const remainingToasts = document.querySelectorAll(".toast-notification");
+            const removedToastBottom = toast.getBoundingClientRect().bottom;
+        
+            remainingToasts.forEach((toast) => {
+                if (toast.getBoundingClientRect().bottom < removedToastBottom) {
+                    toast.classList.add("slide-down");
+                }
+            });
+        
+            // Wait for the slide-down animation to complete
+            setTimeout(() => {
+                // Remove the slide-down class after the animation has completed
+                const slidingToasts = document.querySelectorAll(".slide-down");
+                slidingToasts.forEach((toast) => {
+                    toast.classList.remove("slide-down");
+                });
+        
+                // Adjust the position of remaining toasts again, in case there are multiple toasts being removed at once
+                const remainingToastsDown = document.querySelectorAll(".toast-notification");
+                let heightSoFar = 0;
+                remainingToastsDown.forEach((toast) => {
+                    toast.style.bottom = `${10 + heightSoFar}px`;
+                    heightSoFar += toast.offsetHeight + 10; // add 10 pixels for spacing
+                });
+            }, 0); // The duration of the slide-down animation (in milliseconds)
+        }, 500);
+    };
+
+    // Remove the toast after specified duration
+    setTimeout(removeToast, duration);
+}

From 654749de403ed2cc7de60d273dc6b07962580125 Mon Sep 17 00:00:00 2001
From: patriceac <48073125+patriceac@users.noreply.github.com>
Date: Sat, 29 Apr 2023 19:26:11 -0700
Subject: [PATCH 2/6] Revert "Toast notifications for ED"

This reverts commit dde51c0cefaf36857a54441f68d313ad13ed3e67.
---
 ui/media/css/main.css | 51 --------------------------------------
 ui/media/js/utils.js  | 57 -------------------------------------------
 2 files changed, 108 deletions(-)

diff --git a/ui/media/css/main.css b/ui/media/css/main.css
index ba513237..08dea664 100644
--- a/ui/media/css/main.css
+++ b/ui/media/css/main.css
@@ -1302,54 +1302,3 @@ body.wait-pause {
 .displayNone {
     display:none !important;
 }
-
-/* TOAST NOTIFICATIONS */
-.toast-notification {
-    position: fixed;
-    bottom: 10px;
-    right: -300px;
-    width: 300px;
-    background-color: #333;
-    color: #fff;
-    padding: 10px 20px;
-    border-radius: 5px;
-    box-shadow: 0 0 10px rgba(0, 0, 0, 0.5);
-    z-index: 9999;
-    animation: slideInRight 0.5s ease forwards;
-    transition: bottom 0.5s ease; // Add a transition to smoothly reposition the toasts
-}
-
-.toast-notification-error {
-    color: red;
-}
-
-@keyframes slideInRight {
-    from {
-        right: -300px;
-    }
-    to {
-        right: 10px;
-    }
-}
-
-.toast-notification.hide {
-    animation: slideOutRight 0.5s ease forwards;
-}
-
-@keyframes slideOutRight {
-    from {
-        right: 10px;
-    }
-    to {
-        right: -300px;
-    }
-}
-
-@keyframes slideDown {
-    from {
-        bottom: 10px;
-    }
-    to {
-        bottom: 0;
-    }
-}
diff --git a/ui/media/js/utils.js b/ui/media/js/utils.js
index d1578d8e..6558f7d2 100644
--- a/ui/media/js/utils.js
+++ b/ui/media/js/utils.js
@@ -840,60 +840,3 @@ function createTab(request) {
         }
     })
 }
-
-/* TOAST NOTIFICATIONS */
-function showToast(message, duration = 5000, error = false) {
-    const toast = document.createElement("div");
-    toast.classList.add("toast-notification");
-    if (error === true) {
-        toast.classList.add("toast-notification-error");
-    }
-    toast.innerHTML = message;
-    document.body.appendChild(toast);
-
-    // Set the position of the toast on the screen
-    const toastCount = document.querySelectorAll(".toast-notification").length;
-    const toastHeight = toast.offsetHeight;
-    const previousToastsHeight = Array.from(document.querySelectorAll(".toast-notification"))
-        .slice(0, -1) // exclude current toast
-        .reduce((totalHeight, toast) => totalHeight + toast.offsetHeight + 10, 0); // add 10 pixels for spacing
-    toast.style.bottom = `${10 + previousToastsHeight}px`;
-    toast.style.right = "10px";
-
-    // Delay the removal of the toast until animation has completed
-    const removeToast = () => {
-        toast.classList.add("hide");
-        const removeTimeoutId = setTimeout(() => {
-            toast.remove();
-            // Adjust the position of remaining toasts
-            const remainingToasts = document.querySelectorAll(".toast-notification");
-            const removedToastBottom = toast.getBoundingClientRect().bottom;
-        
-            remainingToasts.forEach((toast) => {
-                if (toast.getBoundingClientRect().bottom < removedToastBottom) {
-                    toast.classList.add("slide-down");
-                }
-            });
-        
-            // Wait for the slide-down animation to complete
-            setTimeout(() => {
-                // Remove the slide-down class after the animation has completed
-                const slidingToasts = document.querySelectorAll(".slide-down");
-                slidingToasts.forEach((toast) => {
-                    toast.classList.remove("slide-down");
-                });
-        
-                // Adjust the position of remaining toasts again, in case there are multiple toasts being removed at once
-                const remainingToastsDown = document.querySelectorAll(".toast-notification");
-                let heightSoFar = 0;
-                remainingToastsDown.forEach((toast) => {
-                    toast.style.bottom = `${10 + heightSoFar}px`;
-                    heightSoFar += toast.offsetHeight + 10; // add 10 pixels for spacing
-                });
-            }, 0); // The duration of the slide-down animation (in milliseconds)
-        }, 500);
-    };
-
-    // Remove the toast after specified duration
-    setTimeout(removeToast, duration);
-}

From a25364732b3f74dbfd5083cc4d7a0f84ccc5d9d5 Mon Sep 17 00:00:00 2001
From: patriceac <48073125+patriceac@users.noreply.github.com>
Date: Wed, 17 May 2023 02:04:20 -0700
Subject: [PATCH 3/6] Support for CodeFormer

Depends on https://github.com/easydiffusion/sdkit/pull/34.
---
 ui/easydiffusion/model_manager.py | 3 ++-
 ui/easydiffusion/renderer.py      | 4 +++-
 2 files changed, 5 insertions(+), 2 deletions(-)

diff --git a/ui/easydiffusion/model_manager.py b/ui/easydiffusion/model_manager.py
index 7bf56575..2bb4852d 100644
--- a/ui/easydiffusion/model_manager.py
+++ b/ui/easydiffusion/model_manager.py
@@ -107,11 +107,12 @@ def resolve_model_to_use(model_name: str = None, model_type: str = None):
 
 
 def reload_models_if_necessary(context: Context, task_data: TaskData):
+    face_correction_model = "codeformer" if "codeformer" in task_data.use_face_correction.lower() else "gfpgan"
     model_paths_in_req = {
         "stable-diffusion": task_data.use_stable_diffusion_model,
         "vae": task_data.use_vae_model,
         "hypernetwork": task_data.use_hypernetwork_model,
-        "gfpgan": task_data.use_face_correction,
+        face_correction_model: task_data.use_face_correction,
         "realesrgan": task_data.use_upscale,
         "nsfw_checker": True if task_data.block_nsfw else None,
         "lora": task_data.use_lora_model,
diff --git a/ui/easydiffusion/renderer.py b/ui/easydiffusion/renderer.py
index f685d038..e1176c8b 100644
--- a/ui/easydiffusion/renderer.py
+++ b/ui/easydiffusion/renderer.py
@@ -158,7 +158,9 @@ def filter_images(task_data: TaskData, images: list, user_stopped):
     filters_to_apply = []
     if task_data.block_nsfw:
         filters_to_apply.append("nsfw_checker")
-    if task_data.use_face_correction and "gfpgan" in task_data.use_face_correction.lower():
+    if task_data.use_face_correction and "codeformer" in task_data.use_face_correction.lower():
+        filters_to_apply.append("codeformer")
+    elif task_data.use_face_correction and "gfpgan" in task_data.use_face_correction.lower():
         filters_to_apply.append("gfpgan")
     if task_data.use_upscale and "realesrgan" in task_data.use_upscale.lower():
         filters_to_apply.append("realesrgan")

From a5a1d3358918626e2860db1effcabea3787fee2b Mon Sep 17 00:00:00 2001
From: patriceac <48073125+patriceac@users.noreply.github.com>
Date: Sun, 21 May 2023 18:32:48 -0700
Subject: [PATCH 4/6] Fix face restoration model selection

---
 ui/easydiffusion/model_manager.py | 10 ++++++++--
 1 file changed, 8 insertions(+), 2 deletions(-)

diff --git a/ui/easydiffusion/model_manager.py b/ui/easydiffusion/model_manager.py
index 2bb4852d..a0b2489a 100644
--- a/ui/easydiffusion/model_manager.py
+++ b/ui/easydiffusion/model_manager.py
@@ -107,12 +107,18 @@ def resolve_model_to_use(model_name: str = None, model_type: str = None):
 
 
 def reload_models_if_necessary(context: Context, task_data: TaskData):
-    face_correction_model = "codeformer" if "codeformer" in task_data.use_face_correction.lower() else "gfpgan"
+    if hasattr(task_data, 'use_face_correction') and task_data.use_face_correction:
+        face_correction_model = "codeformer" if "codeformer" in task_data.use_face_correction.lower() else "gfpgan"
+        face_correction_value = task_data.use_face_correction
+    else:
+        face_correction_model = "gfpgan"
+        face_correction_value = None
+
     model_paths_in_req = {
         "stable-diffusion": task_data.use_stable_diffusion_model,
         "vae": task_data.use_vae_model,
         "hypernetwork": task_data.use_hypernetwork_model,
-        face_correction_model: task_data.use_face_correction,
+        face_correction_model: face_correction_value,
         "realesrgan": task_data.use_upscale,
         "nsfw_checker": True if task_data.block_nsfw else None,
         "lora": task_data.use_lora_model,

From 0860e35d17e15c3cf879b657fd5788b06192c6e4 Mon Sep 17 00:00:00 2001
From: cmdr2 <secondary.cmdr2@gmail.com>
Date: Thu, 1 Jun 2023 16:50:01 +0530
Subject: [PATCH 5/6] sdkit 1.0.101 - CodeFormer as an option to improve faces

---
 scripts/check_models.py           |  4 ++++
 scripts/check_modules.py          |  2 +-
 ui/easydiffusion/model_manager.py | 20 +++++++++++++-------
 ui/easydiffusion/renderer.py      |  3 +++
 ui/easydiffusion/types.py         |  3 ++-
 ui/index.html                     |  9 +++++++--
 ui/media/css/main.css             |  4 ++--
 ui/media/js/main.js               | 27 ++++++++++++++++++++++++---
 ui/media/js/searchable-models.js  | 13 ++++++++++++-
 9 files changed, 68 insertions(+), 17 deletions(-)

diff --git a/scripts/check_models.py b/scripts/check_models.py
index 4b8d68c6..a2186727 100644
--- a/scripts/check_models.py
+++ b/scripts/check_models.py
@@ -24,6 +24,9 @@ models_to_check = {
     "vae": [
         {"file_name": "vae-ft-mse-840000-ema-pruned.ckpt", "model_id": "vae-ft-mse-840000-ema-pruned"},
     ],
+    "codeformer": [
+        {"file_name": "codeformer.pth", "model_id": "codeformer-0.1.0"},
+    ],
 }
 MODEL_EXTENSIONS = {  # copied from easydiffusion/model_manager.py
     "stable-diffusion": [".ckpt", ".safetensors"],
@@ -32,6 +35,7 @@ MODEL_EXTENSIONS = {  # copied from easydiffusion/model_manager.py
     "gfpgan": [".pth"],
     "realesrgan": [".pth"],
     "lora": [".ckpt", ".safetensors"],
+    "codeformer": [".pth"],
 }
 
 
diff --git a/scripts/check_modules.py b/scripts/check_modules.py
index 3686ca00..416c3cc2 100644
--- a/scripts/check_modules.py
+++ b/scripts/check_modules.py
@@ -18,7 +18,7 @@ os_name = platform.system()
 modules_to_check = {
     "torch": ("1.11.0", "1.13.1", "2.0.0"),
     "torchvision": ("0.12.0", "0.14.1", "0.15.1"),
-    "sdkit": "1.0.98",
+    "sdkit": "1.0.101",
     "stable-diffusion-sdkit": "2.1.4",
     "rich": "12.6.0",
     "uvicorn": "0.19.0",
diff --git a/ui/easydiffusion/model_manager.py b/ui/easydiffusion/model_manager.py
index 2a8b57fd..458dae7f 100644
--- a/ui/easydiffusion/model_manager.py
+++ b/ui/easydiffusion/model_manager.py
@@ -13,6 +13,7 @@ KNOWN_MODEL_TYPES = [
     "gfpgan",
     "realesrgan",
     "lora",
+    "codeformer",
 ]
 MODEL_EXTENSIONS = {
     "stable-diffusion": [".ckpt", ".safetensors"],
@@ -21,6 +22,7 @@ MODEL_EXTENSIONS = {
     "gfpgan": [".pth"],
     "realesrgan": [".pth"],
     "lora": [".ckpt", ".safetensors"],
+    "codeformer": [".pth"],
 }
 DEFAULT_MODELS = {
     "stable-diffusion": [  # needed to support the legacy installations
@@ -133,6 +135,9 @@ def reload_models_if_necessary(context: Context, task_data: TaskData):
         if context.model_paths.get(model_type) != path
     }
 
+    if task_data.codeformer_upscale_faces and "realesrgan" not in models_to_reload.keys():
+        models_to_reload["realesrgan"] = resolve_model_to_use(DEFAULT_MODELS["realesrgan"][0], "realesrgan")
+
     if set_vram_optimizations(context) or set_clip_skip(context, task_data):  # reload SD
         models_to_reload["stable-diffusion"] = model_paths_in_req["stable-diffusion"]
 
@@ -159,7 +164,12 @@ def resolve_model_paths(task_data: TaskData):
     task_data.use_lora_model = resolve_model_to_use(task_data.use_lora_model, model_type="lora")
 
     if task_data.use_face_correction:
-        task_data.use_face_correction = resolve_model_to_use(task_data.use_face_correction, "gfpgan")
+        if "gfpgan" in task_data.use_face_correction.lower():
+            model_type = "gfpgan"
+        elif "codeformer" in task_data.use_face_correction.lower():
+            model_type = "codeformer"
+
+        task_data.use_face_correction = resolve_model_to_use(task_data.use_face_correction, model_type)
     if task_data.use_upscale and "realesrgan" in task_data.use_upscale.lower():
         task_data.use_upscale = resolve_model_to_use(task_data.use_upscale, "realesrgan")
 
@@ -240,17 +250,12 @@ def is_malicious_model(file_path):
 
 def getModels():
     models = {
-        "active": {
-            "stable-diffusion": "sd-v1-4",
-            "vae": "",
-            "hypernetwork": "",
-            "lora": "",
-        },
         "options": {
             "stable-diffusion": ["sd-v1-4"],
             "vae": [],
             "hypernetwork": [],
             "lora": [],
+            "codeformer": [],
         },
     }
 
@@ -307,6 +312,7 @@ def getModels():
     listModels(model_type="hypernetwork")
     listModels(model_type="gfpgan")
     listModels(model_type="lora")
+    listModels(model_type="codeformer")
 
     if models_scanned > 0:
         log.info(f"[green]Scanned {models_scanned} models. Nothing infected[/]")
diff --git a/ui/easydiffusion/renderer.py b/ui/easydiffusion/renderer.py
index 1ebd05ec..4be6e6bf 100644
--- a/ui/easydiffusion/renderer.py
+++ b/ui/easydiffusion/renderer.py
@@ -34,6 +34,7 @@ def init(device):
     context.temp_images = {}
     context.partial_x_samples = None
     context.model_load_errors = {}
+    context.enable_codeformer = True
 
     from easydiffusion import app
 
@@ -162,6 +163,8 @@ def filter_images(req: GenerateImageRequest, task_data: TaskData, images: list,
         filters_to_apply.append("nsfw_checker")
     if task_data.use_face_correction and "codeformer" in task_data.use_face_correction.lower():
         filters_to_apply.append("codeformer")
+
+        filter_params["upscale_faces"] = task_data.codeformer_upscale_faces
     elif task_data.use_face_correction and "gfpgan" in task_data.use_face_correction.lower():
         filters_to_apply.append("gfpgan")
     if task_data.use_upscale:
diff --git a/ui/easydiffusion/types.py b/ui/easydiffusion/types.py
index e4426714..4559a23a 100644
--- a/ui/easydiffusion/types.py
+++ b/ui/easydiffusion/types.py
@@ -23,7 +23,7 @@ class GenerateImageRequest(BaseModel):
     sampler_name: str = None  # "ddim", "plms", "heun", "euler", "euler_a", "dpm2", "dpm2_a", "lms"
     hypernetwork_strength: float = 0
     lora_alpha: float = 0
-    tiling: str = "none" # "none", "x", "y", "xy"
+    tiling: str = "none"  # "none", "x", "y", "xy"
 
 
 class TaskData(BaseModel):
@@ -51,6 +51,7 @@ class TaskData(BaseModel):
     stream_image_progress: bool = False
     stream_image_progress_interval: int = 5
     clip_skip: bool = False
+    codeformer_upscale_faces: bool = False
 
 
 class MergeRequest(BaseModel):
diff --git a/ui/index.html b/ui/index.html
index 21ec2550..4167a873 100644
--- a/ui/index.html
+++ b/ui/index.html
@@ -263,7 +263,12 @@
                     <div><ul>
                     <li><b class="settings-subheader">Render Settings</b></li>
                     <li class="pl-5"><input id="stream_image_progress" name="stream_image_progress" type="checkbox"> <label for="stream_image_progress">Show a live preview <small>(uses more VRAM, slower images)</small></label></li>
-                    <li class="pl-5"><input id="use_face_correction" name="use_face_correction" type="checkbox"> <label for="use_face_correction">Fix incorrect faces and eyes</label> <div style="display:inline-block;"><input id="gfpgan_model" type="text" spellcheck="false" autocomplete="off" class="model-filter" data-path="" /></div></li>
+                    <li class="pl-5">
+                        <input id="use_face_correction" name="use_face_correction" type="checkbox"> <label for="use_face_correction">Fix incorrect faces and eyes</label> <div style="display:inline-block;"><input id="gfpgan_model" type="text" spellcheck="false" autocomplete="off" class="model-filter" data-path="" /></div>
+                        <div id="codeformer_settings" class="displayNone sub-settings">
+                            <input id="codeformer_upscale_faces" name="codeformer_upscale_faces" type="checkbox"><label for="codeformer_upscale_faces">Upscale Faces <small>(improves the resolution of faces)</small></label>
+                        </div>
+                    </li>
                     <li class="pl-5">
                         <input id="use_upscale" name="use_upscale" type="checkbox"> <label for="use_upscale">Scale up by</label> 
                         <select id="upscale_amount" name="upscale_amount">
@@ -276,7 +281,7 @@
                             <option value="RealESRGAN_x4plus_anime_6B">RealESRGAN_x4plus_anime_6B</option>
                             <option value="latent_upscaler">Latent Upscaler 2x</option>
                         </select>
-                        <div id="latent_upscaler_settings" class="displayNone">
+                        <div id="latent_upscaler_settings" class="displayNone sub-settings">
                             <label for="latent_upscaler_steps_slider">Upscaling Steps:</label></td><td> <input id="latent_upscaler_steps_slider" name="latent_upscaler_steps_slider" class="editor-slider" value="10" type="range" min="1" max="50"> <input id="latent_upscaler_steps" name="latent_upscaler_steps" size="4" pattern="^[0-9\.]+$" onkeypress="preventNonNumericalInput(event)">
                         </div>
                     </li>
diff --git a/ui/media/css/main.css b/ui/media/css/main.css
index 8f4f49fa..e2d1f79a 100644
--- a/ui/media/css/main.css
+++ b/ui/media/css/main.css
@@ -1303,7 +1303,7 @@ body.wait-pause {
     display:none !important;
 }
 
-#latent_upscaler_settings {
+.sub-settings {
     padding-top: 3pt;
     padding-bottom: 3pt;
     padding-left: 5pt;
@@ -1322,7 +1322,7 @@ body.wait-pause {
     box-shadow: 0 0 10px rgba(0, 0, 0, 0.5);
     z-index: 9999;
     animation: slideInRight 0.5s ease forwards;
-    transition: bottom 0.5s ease; // Add a transition to smoothly reposition the toasts
+    transition: bottom 0.5s ease; /* Add a transition to smoothly reposition the toasts */
 }
 
 .toast-notification-error {
diff --git a/ui/media/js/main.js b/ui/media/js/main.js
index 8628732b..ef27cf1b 100644
--- a/ui/media/js/main.js
+++ b/ui/media/js/main.js
@@ -87,7 +87,7 @@ let promptStrengthField = document.querySelector("#prompt_strength")
 let samplerField = document.querySelector("#sampler_name")
 let samplerSelectionContainer = document.querySelector("#samplerSelection")
 let useFaceCorrectionField = document.querySelector("#use_face_correction")
-let gfpganModelField = new ModelDropdown(document.querySelector("#gfpgan_model"), "gfpgan")
+let gfpganModelField = new ModelDropdown(document.querySelector("#gfpgan_model"), ["codeformer", "gfpgan"])
 let useUpscalingField = document.querySelector("#use_upscale")
 let upscaleModelField = document.querySelector("#upscale_model")
 let upscaleAmountField = document.querySelector("#upscale_amount")
@@ -270,7 +270,9 @@ function shiftOrConfirm(e, prompt, fn) {
         confirm(
             '<small>Tip: To skip this dialog, use shift-click or disable the "Confirm dangerous actions" setting in the Settings tab.</small>',
             prompt,
-            () => { fn(e) }
+            () => {
+                fn(e)
+            }
         )
     }
 }
@@ -1261,6 +1263,10 @@ function getCurrentUserRequest() {
     }
     if (useFaceCorrectionField.checked) {
         newTask.reqBody.use_face_correction = gfpganModelField.value
+
+        if (gfpganModelField.value.includes("codeformer")) {
+            newTask.reqBody.codeformer_upscale_faces = document.querySelector("#codeformer_upscale_faces").checked
+        }
     }
     if (useUpscalingField.checked) {
         newTask.reqBody.use_upscale = upscaleModelField.value
@@ -1574,18 +1580,33 @@ metadataOutputFormatField.disabled = !saveToDiskField.checked
 gfpganModelField.disabled = !useFaceCorrectionField.checked
 useFaceCorrectionField.addEventListener("change", function(e) {
     gfpganModelField.disabled = !this.checked
+
+    onFixFaceModelChange()
 })
 
+function onFixFaceModelChange() {
+    let codeformerSettings = document.querySelector("#codeformer_settings")
+    if (gfpganModelField.value === "codeformer" && !gfpganModelField.disabled) {
+        codeformerSettings.classList.remove("displayNone")
+    } else {
+        codeformerSettings.classList.add("displayNone")
+    }
+}
+gfpganModelField.addEventListener("change", onFixFaceModelChange)
+onFixFaceModelChange()
+
 upscaleModelField.disabled = !useUpscalingField.checked
 upscaleAmountField.disabled = !useUpscalingField.checked
 useUpscalingField.addEventListener("change", function(e) {
     upscaleModelField.disabled = !this.checked
     upscaleAmountField.disabled = !this.checked
+
+    onUpscaleModelChange()
 })
 
 function onUpscaleModelChange() {
     let upscale4x = document.querySelector("#upscale_amount_4x")
-    if (upscaleModelField.value === "latent_upscaler") {
+    if (upscaleModelField.value === "latent_upscaler" && !upscaleModelField.disabled) {
         upscale4x.disabled = true
         upscaleAmountField.value = "2"
         latentUpscalerSettings.classList.remove("displayNone")
diff --git a/ui/media/js/searchable-models.js b/ui/media/js/searchable-models.js
index 9a723d15..f0d06b80 100644
--- a/ui/media/js/searchable-models.js
+++ b/ui/media/js/searchable-models.js
@@ -90,7 +90,12 @@ class ModelDropdown {
 
         if (modelsOptions !== undefined) {
             // reuse models from cache (only useful for plugins, which are loaded after models)
-            this.inputModels = modelsOptions[this.modelKey]
+            this.inputModels = []
+            let modelKeys = Array.isArray(this.modelKey) ? this.modelKey : [this.modelKey]
+            for (let i = 0; i < modelKeys.length; i++) {
+                let key = modelKeys[i]
+                this.inputModels.push(...modelsOptions[key])
+            }
             this.populateModels()
         }
         document.addEventListener(
@@ -98,6 +103,12 @@ class ModelDropdown {
             this.bind(function(e) {
                 // reload the models
                 this.inputModels = modelsOptions[this.modelKey]
+                this.inputModels = []
+                let modelKeys = Array.isArray(this.modelKey) ? this.modelKey : [this.modelKey]
+                for (let i = 0; i < modelKeys.length; i++) {
+                    let key = modelKeys[i]
+                    this.inputModels.push(...modelsOptions[key])
+                }
                 this.populateModels()
             }, this)
         )

From dd95df8f022a51bb8e5b097b94fca34e9fd5105a Mon Sep 17 00:00:00 2001
From: cmdr2 <secondary.cmdr2@gmail.com>
Date: Fri, 2 Jun 2023 16:34:29 +0530
Subject: [PATCH 6/6] Refactor the default model download code, remove
 check_models.py, don't check in legacy paths since that's already migrated
 during initialization; Download CodeFormer's model only when it's used for
 the first time

---
 scripts/check_models.py           | 105 --------------------------
 scripts/on_sd_start.bat           |   7 --
 scripts/on_sd_start.sh            |   6 --
 ui/easydiffusion/app.py           |  45 +++++++++--
 ui/easydiffusion/model_manager.py | 120 +++++++++++++++++++++---------
 5 files changed, 121 insertions(+), 162 deletions(-)
 delete mode 100644 scripts/check_models.py

diff --git a/scripts/check_models.py b/scripts/check_models.py
deleted file mode 100644
index a2186727..00000000
--- a/scripts/check_models.py
+++ /dev/null
@@ -1,105 +0,0 @@
-# this script runs inside the legacy "stable-diffusion" folder
-
-from sdkit.models import download_model, get_model_info_from_db
-from sdkit.utils import hash_file_quick
-
-import os
-import shutil
-from glob import glob
-import traceback
-
-models_base_dir = os.path.abspath(os.path.join("..", "models"))
-
-models_to_check = {
-    "stable-diffusion": [
-        {"file_name": "sd-v1-4.ckpt", "model_id": "1.4"},
-    ],
-    "gfpgan": [
-        {"file_name": "GFPGANv1.4.pth", "model_id": "1.4"},
-    ],
-    "realesrgan": [
-        {"file_name": "RealESRGAN_x4plus.pth", "model_id": "x4plus"},
-        {"file_name": "RealESRGAN_x4plus_anime_6B.pth", "model_id": "x4plus_anime_6"},
-    ],
-    "vae": [
-        {"file_name": "vae-ft-mse-840000-ema-pruned.ckpt", "model_id": "vae-ft-mse-840000-ema-pruned"},
-    ],
-    "codeformer": [
-        {"file_name": "codeformer.pth", "model_id": "codeformer-0.1.0"},
-    ],
-}
-MODEL_EXTENSIONS = {  # copied from easydiffusion/model_manager.py
-    "stable-diffusion": [".ckpt", ".safetensors"],
-    "vae": [".vae.pt", ".ckpt", ".safetensors"],
-    "hypernetwork": [".pt", ".safetensors"],
-    "gfpgan": [".pth"],
-    "realesrgan": [".pth"],
-    "lora": [".ckpt", ".safetensors"],
-    "codeformer": [".pth"],
-}
-
-
-def download_if_necessary(model_type: str, file_name: str, model_id: str):
-    model_path = os.path.join(models_base_dir, model_type, file_name)
-    expected_hash = get_model_info_from_db(model_type=model_type, model_id=model_id)["quick_hash"]
-
-    other_models_exist = any_model_exists(model_type)
-    known_model_exists = os.path.exists(model_path)
-    known_model_is_corrupt = known_model_exists and hash_file_quick(model_path) != expected_hash
-
-    if known_model_is_corrupt or (not other_models_exist and not known_model_exists):
-        print("> download", model_type, model_id)
-        download_model(model_type, model_id, download_base_dir=models_base_dir)
-
-
-def init():
-    migrate_legacy_model_location()
-
-    for model_type, models in models_to_check.items():
-        for model in models:
-            try:
-                download_if_necessary(model_type, model["file_name"], model["model_id"])
-            except:
-                traceback.print_exc()
-                fail(model_type)
-
-        print(model_type, "model(s) found.")
-
-
-### utilities
-def any_model_exists(model_type: str) -> bool:
-    extensions = MODEL_EXTENSIONS.get(model_type, [])
-    for ext in extensions:
-        if any(glob(f"{models_base_dir}/{model_type}/**/*{ext}", recursive=True)):
-            return True
-
-    return False
-
-
-def migrate_legacy_model_location():
-    'Move the models inside the legacy "stable-diffusion" folder, to their respective folders'
-
-    for model_type, models in models_to_check.items():
-        for model in models:
-            file_name = model["file_name"]
-            if os.path.exists(file_name):
-                dest_dir = os.path.join(models_base_dir, model_type)
-                os.makedirs(dest_dir, exist_ok=True)
-                shutil.move(file_name, os.path.join(dest_dir, file_name))
-
-
-def fail(model_name):
-    print(
-        f"""Error downloading the {model_name} model. Sorry about that, please try to:
-1. Run this installer again.
-2. If that doesn't fix it, please try to download the file manually. The address to download from, and the destination to save to are printed above this message.
-3. If those steps don't help, please copy *all* the error messages in this window, and ask the community at https://discord.com/invite/u9yhsFmEkB
-4. If that doesn't solve the problem, please file an issue at https://github.com/cmdr2/stable-diffusion-ui/issues
-Thanks!"""
-    )
-    exit(1)
-
-
-### start
-
-init()
diff --git a/scripts/on_sd_start.bat b/scripts/on_sd_start.bat
index ba205c9e..d2ef7321 100644
--- a/scripts/on_sd_start.bat
+++ b/scripts/on_sd_start.bat
@@ -79,13 +79,6 @@ call WHERE uvicorn > .tmp
     @echo conda_sd_ui_deps_installed >> ..\scripts\install_status.txt
 )
 
-@rem Download the required models
-call python ..\scripts\check_models.py
-if "%ERRORLEVEL%" NEQ "0" (
-    pause
-    exit /b
-)
-
 @>nul findstr /m "sd_install_complete" ..\scripts\install_status.txt
 @if "%ERRORLEVEL%" NEQ "0" (
     @echo sd_weights_downloaded >> ..\scripts\install_status.txt
diff --git a/scripts/on_sd_start.sh b/scripts/on_sd_start.sh
index 820c36ed..55f4da25 100755
--- a/scripts/on_sd_start.sh
+++ b/scripts/on_sd_start.sh
@@ -51,12 +51,6 @@ if ! command -v uvicorn &> /dev/null; then
     fail "UI packages not found!"
 fi
 
-# Download the required models
-if ! python ../scripts/check_models.py; then
-    read -p "Press any key to continue"
-    exit 1
-fi
-
 if [ `grep -c sd_install_complete ../scripts/install_status.txt` -gt "0" ]; then
     echo sd_weights_downloaded >> ../scripts/install_status.txt
     echo sd_install_complete >> ../scripts/install_status.txt
diff --git a/ui/easydiffusion/app.py b/ui/easydiffusion/app.py
index 3064e151..38e3392c 100644
--- a/ui/easydiffusion/app.py
+++ b/ui/easydiffusion/app.py
@@ -90,8 +90,8 @@ def init():
     os.makedirs(USER_SERVER_PLUGINS_DIR, exist_ok=True)
 
     # https://pytorch.org/docs/stable/storage.html
-    warnings.filterwarnings('ignore', category=UserWarning, message='TypedStorage is deprecated')
-    
+    warnings.filterwarnings("ignore", category=UserWarning, message="TypedStorage is deprecated")
+
     load_server_plugins()
 
     update_render_threads()
@@ -221,12 +221,41 @@ def open_browser():
 
         webbrowser.open(f"http://localhost:{port}")
 
-    Console().print(Panel(
-        "\n" +
-        "[white]Easy Diffusion is ready to serve requests.\n\n" +
-        "A new browser tab should have been opened by now.\n" +
-        f"If not, please open your web browser and navigate to [bold yellow underline]http://localhost:{port}/\n",
-        title="Easy Diffusion is ready", style="bold yellow on blue"))
+    Console().print(
+        Panel(
+            "\n"
+            + "[white]Easy Diffusion is ready to serve requests.\n\n"
+            + "A new browser tab should have been opened by now.\n"
+            + f"If not, please open your web browser and navigate to [bold yellow underline]http://localhost:{port}/\n",
+            title="Easy Diffusion is ready",
+            style="bold yellow on blue",
+        )
+    )
+
+
+def fail_and_die(fail_type: str, data: str):
+    suggestions = [
+        "Run this installer again.",
+        "If those steps don't help, please copy *all* the error messages in this window, and ask the community at https://discord.com/invite/u9yhsFmEkB",
+        "If that doesn't solve the problem, please file an issue at https://github.com/cmdr2/stable-diffusion-ui/issues",
+    ]
+
+    if fail_type == "model_download":
+        fail_label = f"Error downloading the {data} model"
+        suggestions.insert(
+            1,
+            "If that doesn't fix it, please try to download the file manually. The address to download from, and the destination to save to are printed above this message.",
+        )
+    else:
+        fail_label = "Error while installing Easy Diffusion"
+
+    msg = [f"{fail_label}. Sorry about that, please try to:"]
+    for i, suggestion in enumerate(suggestions):
+        msg.append(f"{i+1}. {suggestion}")
+    msg.append("Thanks!")
+
+    print("\n".join(msg))
+    exit(1)
 
 
 def get_image_modifiers():
diff --git a/ui/easydiffusion/model_manager.py b/ui/easydiffusion/model_manager.py
index 458dae7f..c4447033 100644
--- a/ui/easydiffusion/model_manager.py
+++ b/ui/easydiffusion/model_manager.py
@@ -1,10 +1,14 @@
 import os
+import shutil
+from glob import glob
+import traceback
 
 from easydiffusion import app
 from easydiffusion.types import TaskData
 from easydiffusion.utils import log
 from sdkit import Context
-from sdkit.models import load_model, scan_model, unload_model
+from sdkit.models import load_model, scan_model, unload_model, download_model, get_model_info_from_db
+from sdkit.utils import hash_file_quick
 
 KNOWN_MODEL_TYPES = [
     "stable-diffusion",
@@ -25,12 +29,19 @@ MODEL_EXTENSIONS = {
     "codeformer": [".pth"],
 }
 DEFAULT_MODELS = {
-    "stable-diffusion": [  # needed to support the legacy installations
-        "custom-model",  # only one custom model file was supported initially, creatively named 'custom-model'
-        "sd-v1-4",  # Default fallback.
+    "stable-diffusion": [
+        {"file_name": "sd-v1-4.ckpt", "model_id": "1.4"},
+    ],
+    "gfpgan": [
+        {"file_name": "GFPGANv1.4.pth", "model_id": "1.4"},
+    ],
+    "realesrgan": [
+        {"file_name": "RealESRGAN_x4plus.pth", "model_id": "x4plus"},
+        {"file_name": "RealESRGAN_x4plus_anime_6B.pth", "model_id": "x4plus_anime_6"},
+    ],
+    "vae": [
+        {"file_name": "vae-ft-mse-840000-ema-pruned.ckpt", "model_id": "vae-ft-mse-840000-ema-pruned"},
     ],
-    "gfpgan": ["GFPGANv1.3"],
-    "realesrgan": ["RealESRGAN_x4plus"],
 }
 MODELS_TO_LOAD_ON_START = ["stable-diffusion", "vae", "hypernetwork", "lora"]
 
@@ -39,6 +50,8 @@ known_models = {}
 
 def init():
     make_model_folders()
+    migrate_legacy_model_location()  # if necessary
+    download_default_models_if_necessary()
     getModels()  # run this once, to cache the picklescan results
 
 
@@ -77,7 +90,7 @@ def resolve_model_to_use(model_name: str = None, model_type: str = None):
     default_models = DEFAULT_MODELS.get(model_type, [])
     config = app.getConfig()
 
-    model_dirs = [os.path.join(app.MODELS_DIR, model_type), app.SD_DIR]
+    model_dir = os.path.join(app.MODELS_DIR, model_type)
     if not model_name:  # When None try user configured model.
         # config = getConfig()
         if "model" in config and model_type in config["model"]:
@@ -85,31 +98,25 @@ def resolve_model_to_use(model_name: str = None, model_type: str = None):
 
     if model_name:
         # Check models directory
-        models_dir_path = os.path.join(app.MODELS_DIR, model_type, model_name)
+        model_path = os.path.join(model_dir, model_name)
+        if os.path.exists(model_path):
+            return model_path
         for model_extension in model_extensions:
-            if os.path.exists(models_dir_path + model_extension):
-                return models_dir_path + model_extension
+            if os.path.exists(model_path + model_extension):
+                return model_path + model_extension
             if os.path.exists(model_name + model_extension):
                 return os.path.abspath(model_name + model_extension)
 
-    # Default locations
-    if model_name in default_models:
-        default_model_path = os.path.join(app.SD_DIR, model_name)
-        for model_extension in model_extensions:
-            if os.path.exists(default_model_path + model_extension):
-                return default_model_path + model_extension
-
     # Can't find requested model, check the default paths.
-    for default_model in default_models:
-        for model_dir in model_dirs:
-            default_model_path = os.path.join(model_dir, default_model)
-            for model_extension in model_extensions:
-                if os.path.exists(default_model_path + model_extension):
-                    if model_name is not None:
-                        log.warn(
-                            f"Could not find the configured custom model {model_name}{model_extension}. Using the default one: {default_model_path}{model_extension}"
-                        )
-                    return default_model_path + model_extension
+    if model_type == "stable-diffusion":
+        for default_model in default_models:
+            default_model_path = os.path.join(model_dir, default_model["file_name"])
+            if os.path.exists(default_model_path):
+                if model_name is not None:
+                    log.warn(
+                        f"Could not find the configured custom model {model_name}. Using the default one: {default_model_path}"
+                    )
+                return default_model_path
 
     return None
 
@@ -136,7 +143,9 @@ def reload_models_if_necessary(context: Context, task_data: TaskData):
     }
 
     if task_data.codeformer_upscale_faces and "realesrgan" not in models_to_reload.keys():
-        models_to_reload["realesrgan"] = resolve_model_to_use(DEFAULT_MODELS["realesrgan"][0], "realesrgan")
+        models_to_reload["realesrgan"] = resolve_model_to_use(
+            DEFAULT_MODELS["realesrgan"][0]["file_name"], "realesrgan"
+        )
 
     if set_vram_optimizations(context) or set_clip_skip(context, task_data):  # reload SD
         models_to_reload["stable-diffusion"] = model_paths_in_req["stable-diffusion"]
@@ -168,6 +177,7 @@ def resolve_model_paths(task_data: TaskData):
             model_type = "gfpgan"
         elif "codeformer" in task_data.use_face_correction.lower():
             model_type = "codeformer"
+            download_if_necessary("codeformer", "codeformer.pth", "codeformer-0.1.0")
 
         task_data.use_face_correction = resolve_model_to_use(task_data.use_face_correction, model_type)
     if task_data.use_upscale and "realesrgan" in task_data.use_upscale.lower():
@@ -179,7 +189,31 @@ def fail_if_models_did_not_load(context: Context):
         if model_type in context.model_load_errors:
             e = context.model_load_errors[model_type]
             raise Exception(f"Could not load the {model_type} model! Reason: " + e)
-            # concat 'e', don't use in format string (injection attack)
+
+
+def download_default_models_if_necessary():
+    for model_type, models in DEFAULT_MODELS.items():
+        for model in models:
+            try:
+                download_if_necessary(model_type, model["file_name"], model["model_id"])
+            except:
+                traceback.print_exc()
+                app.fail_and_die(fail_type="model_download", data=model_type)
+
+        print(model_type, "model(s) found.")
+
+
+def download_if_necessary(model_type: str, file_name: str, model_id: str):
+    model_path = os.path.join(app.MODELS_DIR, model_type, file_name)
+    expected_hash = get_model_info_from_db(model_type=model_type, model_id=model_id)["quick_hash"]
+
+    other_models_exist = any_model_exists(model_type)
+    known_model_exists = os.path.exists(model_path)
+    known_model_is_corrupt = known_model_exists and hash_file_quick(model_path) != expected_hash
+
+    if known_model_is_corrupt or (not other_models_exist and not known_model_exists):
+        print("> download", model_type, model_id)
+        download_model(model_type, model_id, download_base_dir=app.MODELS_DIR)
 
 
 def set_vram_optimizations(context: Context):
@@ -193,6 +227,26 @@ def set_vram_optimizations(context: Context):
     return False
 
 
+def migrate_legacy_model_location():
+    'Move the models inside the legacy "stable-diffusion" folder, to their respective folders'
+
+    for model_type, models in DEFAULT_MODELS.items():
+        for model in models:
+            file_name = model["file_name"]
+            legacy_path = os.path.join(app.SD_DIR, file_name)
+            if os.path.exists(legacy_path):
+                shutil.move(legacy_path, os.path.join(app.MODELS_DIR, model_type, file_name))
+
+
+def any_model_exists(model_type: str) -> bool:
+    extensions = MODEL_EXTENSIONS.get(model_type, [])
+    for ext in extensions:
+        if any(glob(f"{app.MODELS_DIR}/{model_type}/**/*{ext}", recursive=True)):
+            return True
+
+    return False
+
+
 def set_clip_skip(context: Context, task_data: TaskData):
     clip_skip = task_data.clip_skip
 
@@ -255,7 +309,7 @@ def getModels():
             "vae": [],
             "hypernetwork": [],
             "lora": [],
-            "codeformer": [],
+            "codeformer": ["codeformer"],
         },
     }
 
@@ -312,14 +366,8 @@ def getModels():
     listModels(model_type="hypernetwork")
     listModels(model_type="gfpgan")
     listModels(model_type="lora")
-    listModels(model_type="codeformer")
 
     if models_scanned > 0:
         log.info(f"[green]Scanned {models_scanned} models. Nothing infected[/]")
 
-    # legacy
-    custom_weight_path = os.path.join(app.SD_DIR, "custom-model.ckpt")
-    if os.path.exists(custom_weight_path):
-        models["options"]["stable-diffusion"].append("custom-model")
-
     return models