add face restoration option to xyz_grid

1 files changed, 15 insertions, 0 deletions

diff --git a/scripts/xyz_grid.py b/scripts/xyz_grid.py
index 9a0678fa..ce584981 100644
--- a/scripts/xyz_grid.py
+++ b/scripts/xyz_grid.py
@@ -132,6 +132,20 @@ def apply_uni_pc_order(p, x, xs):
     opts.data["uni_pc_order"] = min(x, p.steps - 1)

 

 

+def apply_face_restore(p, opt, x):

+    opt = opt.lower()

+    if opt == 'codeformer':

+        is_active = True

+        p.face_restoration_model = 'CodeFormer'

+    elif opt == 'gfpgan':

+        is_active = True

+        p.face_restoration_model = 'GFPGAN'

+    else:

+        is_active = opt in ('true', 'yes', 'y', '1')

+

+    p.restore_faces = is_active

+

+

 def format_value_add_label(p, opt, x):

     if type(x) == float:

         x = round(x, 8)

@@ -210,6 +224,7 @@ axis_options = [
     AxisOption("VAE", str, apply_vae, cost=0.7, choices=lambda: list(sd_vae.vae_dict)),

     AxisOption("Styles", str, apply_styles, choices=lambda: list(shared.prompt_styles.styles)),

     AxisOption("UniPC Order", int, apply_uni_pc_order, cost=0.5),

+    AxisOption("Face restore", str, apply_face_restore, format_value=format_value),

 ]