Fix a bug where the realesrgan model would get unloaded after the first request in a batch while using Codeformer with upscaling of faces

This commit is contained in:
cmdr2 2023-06-05 15:08:57 +05:30
parent 8a2c09c6de
commit a10aa92634
2 changed files with 7 additions and 5 deletions

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.101",
"sdkit": "1.0.102",
"stable-diffusion-sdkit": "2.1.4",
"rich": "12.6.0",
"uvicorn": "0.19.0",

View File

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