Add support for Apple Silicon (M1, M2, M3, …)

.gitignore

@@ -127,3 +127,5 @@ dmypy.json
 
 # Pyre type checker
 .pyre/
+
+weights/*

README.md

@@ -32,6 +32,8 @@ _<sup>2</sup>[Department of Computing, The Hong Kong Polytechnic University](htt
 <img src="figs/seg2face_00.jpg" width="390px"/> <img src="figs/seg2face_01.jpg" width="390px"/>
 
 ## News
+(2023-12-16) GPEN can run on Apple Silicon GPU now by using `--use_mps`.
+
 (2023-02-15) **GPEN-BFR-1024** and **GPEN-BFR-2048** are now publicly available. Please download them via \[[ModelScope2](https://www.modelscope.cn/models/damo/cv_gpen_image-portrait-enhancement-hires/summary)\].
 
 (2023-02-15) We provide online demos via \[[ModelScope1](https://www.modelscope.cn/models/damo/cv_gpen_image-portrait-enhancement/summary)\] and \[[ModelScope2](https://www.modelscope.cn/models/damo/cv_gpen_image-portrait-enhancement-hires/summary)\]. 

demo.py

@@ -78,6 +78,7 @@ def generate_mask(H, W, img=None):
     parser.add_argument('--alpha', type=float, default=1, help='blending the results')
     parser.add_argument('--use_sr', action='store_true', help='use sr or not')
     parser.add_argument('--use_cuda', action='store_true', help='use cuda or not')
+    parser.add_argument('--use_mps', action='store_true', help='use Apple Silicon or not')
     parser.add_argument('--save_face', action='store_true', help='save face or not')
     parser.add_argument('--aligned', action='store_true', help='input are aligned faces or not')
     parser.add_argument('--sr_model', type=str, default='realesrnet', help='SR model')
@@ -93,14 +94,21 @@ def generate_mask(H, W, img=None):
 
     os.makedirs(args.outdir, exist_ok=True)
 
+    if args.use_cuda:
+        device = 'cuda'
+    elif args.use_mps:
+        device = 'mps'
+    else:
+        device = 'cpu'
+
     if args.task == 'FaceEnhancement': 
-        processer = FaceEnhancement(args, in_size=args.in_size, model=args.model, use_sr=args.use_sr, device='cuda' if args.use_cuda else 'cpu')
+        processer = FaceEnhancement(args, in_size=args.in_size, model=args.model, use_sr=args.use_sr, device=device)
     elif args.task == 'FaceColorization':
-        processer = FaceColorization(in_size=args.in_size, model=args.model, device='cuda' if args.use_cuda else 'cpu')
+        processer = FaceColorization(in_size=args.in_size, model=args.model, device=device)
     elif args.task == 'FaceInpainting':
-        processer = FaceInpainting(in_size=args.in_size, model=args.model, device='cuda' if args.use_cuda else 'cpu')
+        processer = FaceInpainting(in_size=args.in_size, model=args.model, device=device)
     elif args.task == 'Segmentation2Face':
-        processer = Segmentation2Face(in_size=args.in_size, model=args.model, is_norm=False, device='cuda' if args.use_cuda else 'cpu')
+        processer = Segmentation2Face(in_size=args.in_size, model=args.model, is_norm=False, device=device)
 
 
     files = sorted(glob.glob(os.path.join(args.indir, '*.*g')))

face_parse/face_parsing.py

@@ -58,7 +58,7 @@ def process_tensor(self, imt):
     def img2tensor(self, img):
         img = img[..., ::-1]
         img = img / 255. * 2 - 1
-        img_tensor = torch.from_numpy(img.transpose(2, 0, 1)).unsqueeze(0).to(self.device)
+        img_tensor = torch.from_numpy(img.transpose(2, 0, 1)).unsqueeze(0).to(self.device, dtype=torch.float32)
         return img_tensor.float()
 
     def tenor2mask(self, tensor):