From a10aa92634f5743319a7ad1e2faf3f68c875fdf8 Mon Sep 17 00:00:00 2001 From: cmdr2 Date: Mon, 5 Jun 2023 15:08:57 +0530 Subject: [PATCH] Fix a bug where the realesrgan model would get unloaded after the first request in a batch while using Codeformer with upscaling of faces --- scripts/check_modules.py | 2 +- ui/easydiffusion/model_manager.py | 10 ++++++---- 2 files changed, 7 insertions(+), 5 deletions(-) diff --git a/scripts/check_modules.py b/scripts/check_modules.py index 416c3cc2..cbd6d0c1 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.101", + "sdkit": "1.0.102", "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 c4447033..29cf529b 100644 --- a/ui/easydiffusion/model_manager.py +++ b/ui/easydiffusion/model_manager.py @@ -142,10 +142,12 @@ 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]["file_name"], "realesrgan" - ) + if task_data.codeformer_upscale_faces: + if "realesrgan" not in models_to_reload and "realesrgan" not in context.models: + default_realesrgan = DEFAULT_MODELS["realesrgan"][0]["file_name"] + models_to_reload["realesrgan"] = resolve_model_to_use(default_realesrgan, "realesrgan") + elif "realesrgan" in models_to_reload and models_to_reload["realesrgan"] is None: + del models_to_reload["realesrgan"] # don't unload 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"]