Unverified Commit 151e2cc6 authored by AUTOMATIC1111's avatar AUTOMATIC1111 Committed by GitHub
Browse files

Merge pull request #4461 from brkirch/face-restoration-device-fix 

Fix setting device for GFPGAN and CodeFormer
parents cc90dcc9 f4a488f5

modules/codeformer_model.py

+3 −0
Original line number Diff line number Diff line
@@ -36,6 +36,7 @@ def setup_model(dirname): 
        from basicsr.utils.download_util import load_file_from_url

        from basicsr.utils import imwrite, img2tensor, tensor2img

        from facelib.utils.face_restoration_helper import FaceRestoreHelper

        from facelib.detection.retinaface import retinaface

        from modules.shared import cmd_opts



        net_class = CodeFormer
@@ -65,6 +66,8 @@ def setup_model(dirname): 
                net.load_state_dict(checkpoint)

                net.eval()



                if hasattr(retinaface, 'device'):

                    retinaface.device = devices.device_codeformer

                face_helper = FaceRestoreHelper(1, face_size=512, crop_ratio=(1, 1), det_model='retinaface_resnet50', save_ext='png', use_parse=True, device=devices.device_codeformer)



                self.net = net

modules/gfpgan_model.py

+3 −1
Original line number Diff line number Diff line
@@ -36,7 +36,9 @@ def gfpgann(): 
    else:

        print("Unable to load gfpgan model!")

        return None

    model = gfpgan_constructor(model_path=model_file, upscale=1, arch='clean', channel_multiplier=2, bg_upsampler=None)

    if hasattr(facexlib.detection.retinaface, 'device'):

        facexlib.detection.retinaface.device = devices.device_gfpgan

    model = gfpgan_constructor(model_path=model_file, upscale=1, arch='clean', channel_multiplier=2, bg_upsampler=None, device=devices.device_gfpgan)

    loaded_gfpgan_model = model



    return model

CRA Git | Maintained and supported by SUSTech CRA and CCSE