Merge pull request #294 from EliasOenal/master

Fixes for mps/Metal: use of seeds, img2img, CodeFormer

1 files changed, 5 insertions, 3 deletions

diff --git a/modules/codeformer_model.py b/modules/codeformer_model.py
index fd1da692..6cd29c83 100644
--- a/modules/codeformer_model.py
+++ b/modules/codeformer_model.py
@@ -47,6 +47,8 @@ def setup_codeformer():
             def __init__(self):

                 self.net = None

                 self.face_helper = None

+                if shared.device.type == 'mps': # CodeFormer currently does not support mps backend

+                    shared.device_codeformer = torch.device('cpu')

 

             def create_models(self):

 

@@ -54,13 +56,13 @@ def setup_codeformer():
                     self.net.to(shared.device)

                     return self.net, self.face_helper

 

-                net = net_class(dim_embd=512, codebook_size=1024, n_head=8, n_layers=9, connect_list=['32', '64', '128', '256']).to(shared.device)

+                net = net_class(dim_embd=512, codebook_size=1024, n_head=8, n_layers=9, connect_list=['32', '64', '128', '256']).to(shared.device_codeformer)

                 ckpt_path = load_file_from_url(url=pretrain_model_url, model_dir=os.path.join(path, 'weights/CodeFormer'), progress=True)

                 checkpoint = torch.load(ckpt_path)['params_ema']

                 net.load_state_dict(checkpoint)

                 net.eval()

 

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

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

 

                 self.net = net

                 self.face_helper = face_helper

@@ -82,7 +84,7 @@ def setup_codeformer():
                 for idx, cropped_face in enumerate(self.face_helper.cropped_faces):

                     cropped_face_t = img2tensor(cropped_face / 255., bgr2rgb=True, float32=True)

                     normalize(cropped_face_t, (0.5, 0.5, 0.5), (0.5, 0.5, 0.5), inplace=True)

-                    cropped_face_t = cropped_face_t.unsqueeze(0).to(shared.device)

+                    cropped_face_t = cropped_face_t.unsqueeze(0).to(shared.device_codeformer)

 

                     try:

                         with torch.no_grad():