Commit 5ed5e95f authored by high_byte's avatar high_byte
Browse files

add face restoration option to xyz_grid

parent 27e319dc
Loading
Loading
Loading
Loading
+15 −0
Original line number Diff line number Diff line
@@ -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),
]