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