From a5a1d3358918626e2860db1effcabea3787fee2b Mon Sep 17 00:00:00 2001
From: patriceac <48073125+patriceac@users.noreply.github.com>
Date: Sun, 21 May 2023 18:32:48 -0700
Subject: [PATCH] Fix face restoration model selection

---
 ui/easydiffusion/model_manager.py | 10 ++++++++--
 1 file changed, 8 insertions(+), 2 deletions(-)

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,