Back to Repositories

Testing Face Restoration Models in Stable-Diffusion-WebUI

This test suite validates face restoration functionality in the Stable Diffusion WebUI, specifically testing the GFPGAN and CodeFormer face restoration models. The tests ensure proper image processing and verification of restoration effects on facial features.

Test Coverage Overview

The test suite provides comprehensive coverage of face restoration functionality:

Tests both GFPGAN and CodeFormer restoration models
Validates image shape preservation after restoration
Verifies visible changes in restored images
Handles model initialization and setup
Tests with real image input (two-faces.jpg)

Implementation Analysis

The implementation uses pytest’s parametrization to test multiple face restorers with the same test logic. The approach dynamically loads appropriate model implementations, performs face restoration, and validates results through numpy array comparisons.

Key patterns include fixture usage for initialization, parametrized test cases, and image processing verification using numpy arrays.

Technical Details

Testing tools and configuration:

pytest for test framework and fixtures
PIL (Python Imaging Library) for image handling
NumPy for image array manipulation and comparison
Custom test paths configuration
Model-specific setup paths for GFPGAN and CodeFormer

Best Practices Demonstrated

The test implementation showcases several testing best practices:

Parametrized testing to reduce code duplication
Proper fixture usage for setup/teardown
Clear assertion messages and conditions
File handling with proper paths
Explicit type checking and shape validation
Output verification through both structural and content checks

automatic1111/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"))