Commit f4a488f5 authored by brkirch's avatar brkirch
Browse files

Set device for facelib/facexlib and gfpgan 

* FaceXLib/FaceLib doesn't pass the device argument to RetinaFace but instead chooses one itself and sets it to a global - in order to use a device other than its internally chosen default it is necessary to manually replace the default value
* The GFPGAN constructor needs the device argument to work with MPS or a CUDA device ID that differs from the default
parent 804d9fb8

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