stable-diffusion-webui

1
from __future__ import annotations
2

3
import logging
4
import os
5

6
import torch
7

8
from modules import (
9
    devices,
10
    errors,
11
    face_restoration,
12
    face_restoration_utils,
13
    modelloader,
14
    shared,
15
)
16

17
logger = logging.getLogger(__name__)
18
model_url = "https://github.com/TencentARC/GFPGAN/releases/download/v1.3.0/GFPGANv1.4.pth"
19
model_download_name = "GFPGANv1.4.pth"
20
gfpgan_face_restorer: face_restoration.FaceRestoration | None = None
21

22

23
class FaceRestorerGFPGAN(face_restoration_utils.CommonFaceRestoration):
24
    def name(self):
25
        return "GFPGAN"
26

27
    def get_device(self):
28
        return devices.device_gfpgan
29

30
    def load_net(self) -> torch.Module:
31
        for model_path in modelloader.load_models(
32
            model_path=self.model_path,
33
            model_url=model_url,
34
            command_path=self.model_path,
35
            download_name=model_download_name,
36
            ext_filter=['.pth'],
37
        ):
38
            if 'GFPGAN' in os.path.basename(model_path):
39
                model = modelloader.load_spandrel_model(
40
                    model_path,
41
                    device=self.get_device(),
42
                    expected_architecture='GFPGAN',
43
                ).model
44
                model.different_w = True  # see https://github.com/chaiNNer-org/spandrel/pull/81
45
                return model
46
        raise ValueError("No GFPGAN model found")
47

48
    def restore(self, np_image):
49
        def restore_face(cropped_face_t):
50
            assert self.net is not None
51
            return self.net(cropped_face_t, return_rgb=False)[0]
52

53
        return self.restore_with_helper(np_image, restore_face)
54

55

56
def gfpgan_fix_faces(np_image):
57
    if gfpgan_face_restorer:
58
        return gfpgan_face_restorer.restore(np_image)
59
    logger.warning("GFPGAN face restorer not set up")
60
    return np_image
61

62

63
def setup_model(dirname: str) -> None:
64
    global gfpgan_face_restorer
65

66
    try:
67
        face_restoration_utils.patch_facexlib(dirname)
68
        gfpgan_face_restorer = FaceRestorerGFPGAN(model_path=dirname)
69
        shared.face_restorers.append(gfpgan_face_restorer)
70
    except Exception:
71
        errors.report("Error setting up GFPGAN", exc_info=True)
72