stable-diffusion-webui/test/test_face_restorers.py

import os
from test.conftest import test_files_path, test_outputs_path

import numpy as np
import pytest
from PIL import Image


@pytest.mark.usefixtures("initialize")
@pytest.mark.parametrize("restorer_name", ["gfpgan", "codeformer"])
def test_face_restorers(restorer_name):
    from modules import shared

    if restorer_name == "gfpgan":
        from modules import gfpgan_model
        gfpgan_model.setup_model(shared.cmd_opts.gfpgan_models_path)
        restorer = gfpgan_model.gfpgan_fix_faces
    elif restorer_name == "codeformer":
        from modules import codeformer_model
        codeformer_model.setup_model(shared.cmd_opts.codeformer_models_path)
        restorer = codeformer_model.codeformer.restore
    else:
        raise NotImplementedError("...")
    img = Image.open(os.path.join(test_files_path, "two-faces.jpg"))
    np_img = np.array(img, dtype=np.uint8)
    fixed_image = restorer(np_img)
    assert fixed_image.shape == np_img.shape
    assert not np.allclose(fixed_image, np_img)  # should have visibly changed
    Image.fromarray(fixed_image).save(os.path.join(test_outputs_path, f"{restorer_name}.png"))