diff --git a/ui/easydiffusion/model_manager.py b/ui/easydiffusion/model_manager.py index 2bb4852d..a0b2489a 100644 --- a/ui/easydiffusion/model_manager.py +++ b/ui/easydiffusion/model_manager.py @@ -107,12 +107,18 @@ def resolve_model_to_use(model_name: str = None, model_type: str = None): def reload_models_if_necessary(context: Context, task_data: TaskData): - face_correction_model = "codeformer" if "codeformer" in task_data.use_face_correction.lower() else "gfpgan" + if hasattr(task_data, 'use_face_correction') and task_data.use_face_correction: + face_correction_model = "codeformer" if "codeformer" in task_data.use_face_correction.lower() else "gfpgan" + face_correction_value = task_data.use_face_correction + else: + face_correction_model = "gfpgan" + face_correction_value = None + model_paths_in_req = { "stable-diffusion": task_data.use_stable_diffusion_model, "vae": task_data.use_vae_model, "hypernetwork": task_data.use_hypernetwork_model, - face_correction_model: task_data.use_face_correction, + face_correction_model: face_correction_value, "realesrgan": task_data.use_upscale, "nsfw_checker": True if task_data.block_nsfw else None, "lora": task_data.use_lora_model,