sdkit 1.0.101 - CodeFormer as an option to improve faces

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"],
 }
 
 

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

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[/]")

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:

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):

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>

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 {

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

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