holm/easydiffusion
1
0
Fork 0
forked from extern/easydiffusion
Code Issues Pull Requests Packages Projects Releases Wiki Activity

sdkit 1.0.101 - CodeFormer as an option to improve faces

Browse Source
This commit is contained in:
cmdr2 2023-06-01 16:50:01 +05:30
parent 32c4f10626
commit 0860e35d17
9 changed files with 68 additions and 17 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

4
scripts/check_models.py
View File

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

2
scripts/check_modules.py
View File

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

20
ui/easydiffusion/model_manager.py
View File

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

3
ui/easydiffusion/renderer.py
View File

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

3
ui/easydiffusion/types.py
View File

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

9
ui/index.html
View File

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

4
ui/media/css/main.css
View File

@ -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 {

27
ui/media/js/main.js
View File

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

13
ui/media/js/searchable-models.js
View File

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