add face restoration option to xyz_grid

This commit is contained in:
high_byte 2023-03-12 03:29:07 +02:00
parent 27e319dc4f
commit 5ed5e95fb8
1 changed files with 15 additions and 0 deletions

View File

@ -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),
]