support MPS (Mac M1) device; Happy Lunar New Year!

README.md

@@ -146,4 +146,4 @@ This project is licensed under <a rel="license" href="https://github.com/sczhou/
 This project is based on [BasicSR](https://github.com/XPixelGroup/BasicSR). Some codes are brought from [Unleashing Transformers](https://github.com/samb-t/unleashing-transformers), [YOLOv5-face](https://github.com/deepcam-cn/yolov5-face), and [FaceXLib](https://github.com/xinntao/facexlib). We also adopt [Real-ESRGAN](https://github.com/xinntao/Real-ESRGAN) to support background image enhancement. Thanks for their awesome works.
 
 ### Contact
-If you have any question, please feel free to reach me out at `shangchenzhou@gmail.com`.
+If you have any questions, please feel free to reach me out at `shangchenzhou@gmail.com`.

basicsr/setup.py

@@ -6,8 +6,8 @@ import os
 import subprocess
 import sys
 import time
-import torch
 from torch.utils.cpp_extension import BuildExtension, CppExtension, CUDAExtension
+from utils.misc import gpu_is_available
 
 version_file = './basicsr/version.py'
 
@@ -87,7 +87,8 @@ def make_cuda_ext(name, module, sources, sources_cuda=None):
     define_macros = []
     extra_compile_args = {'cxx': []}
 
-    if torch.cuda.is_available() or os.getenv('FORCE_CUDA', '0') == '1':
+    # if torch.cuda.is_available() or os.getenv('FORCE_CUDA', '0') == '1':
+    if gpu_is_available or os.getenv('FORCE_CUDA', '0') == '1':
         define_macros += [('WITH_CUDA', None)]
         extension = CUDAExtension
         extra_compile_args['nvcc'] = [

basicsr/utils/misc.py

@@ -1,13 +1,36 @@
-import numpy as np
 import os
+import re
 import random
 import time
 import torch
+import numpy as np
 from os import path as osp
 
 from .dist_util import master_only
 from .logger import get_root_logger
 
+IS_HIGH_VERSION = [int(m) for m in list(re.findall(r"^([0-9]+)\.([0-9]+)\.([0-9]+)([^0-9][a-zA-Z0-9]*)?(\+git.*)?$",\
+    torch.__version__)[0][:3])] >= [1, 12, 0]
+
+def gpu_is_available():
+    if IS_HIGH_VERSION:
+        if torch.backends.mps.is_available():
+            return True
+    return True if torch.cuda.is_available() and torch.backends.cudnn.is_available() else False
+
+def get_device(gpu_id=None):
+    if gpu_id is None:
+        gpu_str = ''
+    elif isinstance(gpu_id, int):
+        gpu_str = f':{gpu_id}'
+    else:
+        raise TypeError('Input should be int value.')
+
+    if IS_HIGH_VERSION:
+        if torch.backends.mps.is_available():
+            return torch.device('mps'+gpu_str)
+    return torch.device('cuda'+gpu_str if torch.cuda.is_available() and torch.backends.cudnn.is_available() else 'cpu')
+
 
 def set_random_seed(seed):
     """Set random seeds."""

basicsr/utils/realesrgan_utils.py

@@ -5,12 +5,12 @@ import os
 import queue
 import threading
 import torch
-from basicsr.utils.download_util import load_file_from_url
 from torch.nn import functional as F
+from basicsr.utils.download_util import load_file_from_url
+from basicsr.utils.misc import get_device
 
 # ROOT_DIR = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
 
-
 class RealESRGANer():
     """A helper class for upsampling images with RealESRGAN.
 
@@ -44,11 +44,14 @@ class RealESRGANer():
         self.half = half
 
         # initialize model
-        if gpu_id:
-            self.device = torch.device(
-                f'cuda:{gpu_id}' if torch.cuda.is_available() else 'cpu') if device is None else device
-        else:
-            self.device = torch.device('cuda' if torch.cuda.is_available() else 'cpu') if device is None else device
+        # if gpu_id:
+        #     self.device = torch.device(
+        #         f'cuda:{gpu_id}' if torch.cuda.is_available() else 'cpu') if device is None else device
+        # else:
+        #     self.device = torch.device('cuda' if torch.cuda.is_available() else 'cpu') if device is None else device
+
+        self.device = get_device(gpu_id) if device is None else device
+        
         # if the model_path starts with https, it will first download models to the folder: realesrgan/weights
         if model_path.startswith('https://'):
             model_path = load_file_from_url(

facelib/detection/retinaface/retinaface.py

@@ -11,11 +11,13 @@ from facelib.detection.retinaface.retinaface_net import FPN, SSH, MobileNetV1, m
 from facelib.detection.retinaface.retinaface_utils import (PriorBox, batched_decode, batched_decode_landm, decode, decode_landm,
                                                  py_cpu_nms)
 
-device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
+from basicsr.utils.misc import get_device
+# device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
+device = get_device()
 
 
 def generate_config(network_name):
-
+    
     cfg_mnet = {
         'name': 'mobilenet0.25',
         'min_sizes': [[16, 32], [64, 128], [256, 512]],

facelib/detection/yolov5face/face_detector.py

@@ -1,13 +1,10 @@
-import copy
-import os
-from pathlib import Path
-
 import cv2
-import numpy as np
+import copy
+import re
 import torch
-from torch import nn
+import numpy as np
 
-from facelib.detection.yolov5face.models.common import Conv
+from pathlib import Path
 from facelib.detection.yolov5face.models.yolo import Model
 from facelib.detection.yolov5face.utils.datasets import letterbox
 from facelib.detection.yolov5face.utils.general import (
@@ -17,7 +14,9 @@ from facelib.detection.yolov5face.utils.general import (
     scale_coords_landmarks,
 )
 
-IS_HIGH_VERSION = tuple(map(int, torch.__version__.split('+')[0].split('.')[:2])) >= (1, 9, 0)
+# IS_HIGH_VERSION = tuple(map(int, torch.__version__.split('+')[0].split('.')[:2])) >= (1, 9)
+IS_HIGH_VERSION = [int(m) for m in list(re.findall(r"^([0-9]+)\.([0-9]+)\.([0-9]+)([^0-9][a-zA-Z0-9]*)?(\+git.*)?$",\
+    torch.__version__)[0][:3])] >= [1, 9, 0]
 
 
 def isListempty(inList):

facelib/utils/face_restoration_helper.py

@@ -7,6 +7,7 @@ from torchvision.transforms.functional import normalize
 from facelib.detection import init_detection_model
 from facelib.parsing import init_parsing_model
 from facelib.utils.misc import img2tensor, imwrite, is_gray, bgr2gray
+from basicsr.utils.misc import get_device
 
 
 def get_largest_face(det_faces, h, w):
@@ -97,7 +98,8 @@ class FaceRestoreHelper(object):
         self.pad_input_imgs = []
 
         if device is None:
-            self.device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
+            # self.device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
+            self.device = get_device()
         else:
             self.device = device
 

inference_codeformer.py

@@ -6,9 +6,9 @@ import torch
 from torchvision.transforms.functional import normalize
 from basicsr.utils import imwrite, img2tensor, tensor2img
 from basicsr.utils.download_util import load_file_from_url
+from basicsr.utils.misc import gpu_is_available, get_device
 from facelib.utils.face_restoration_helper import FaceRestoreHelper
 from facelib.utils.misc import is_gray
-import torch.nn.functional as F
 
 from basicsr.utils.registry import ARCH_REGISTRY
 
@@ -19,9 +19,7 @@ pretrain_model_url = {
 def set_realesrgan():
     from basicsr.archs.rrdbnet_arch import RRDBNet
     from basicsr.utils.realesrgan_utils import RealESRGANer
-
-    cuda_is_available = torch.cuda.is_available()
-    half = True if cuda_is_available else False
+    
     model = RRDBNet(
         num_in_ch=3,
         num_out_ch=3,
@@ -37,10 +35,10 @@ def set_realesrgan():
         tile=args.bg_tile,
         tile_pad=40,
         pre_pad=0,
-        half=half, # need to set False in CPU mode
+        half=torch.cuda.is_available(), # need to set False in CPU/MPS mode
     )
 
-    if not cuda_is_available:  # CPU
+    if not gpu_is_available():  # CPU
         import warnings
         warnings.warn('Running on CPU now! Make sure your PyTorch version matches your CUDA.'
                         'The unoptimized RealESRGAN is slow on CPU. '
@@ -49,7 +47,8 @@ def set_realesrgan():
     return upsampler
 
 if __name__ == '__main__':
-    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
+    # device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
+    device = get_device()
     parser = argparse.ArgumentParser()
 
     parser.add_argument('-i', '--input_path', type=str, default='./inputs/whole_imgs', 
@@ -79,10 +78,10 @@ if __name__ == '__main__':
     # ------------------------ input & output ------------------------
     w = args.fidelity_weight
     input_video = False
-    if args.input_path.endswith(('jpg', 'png')): # input single img path
+    if args.input_path.endswith(('jpg', 'jpeg', 'png', 'JPG', 'JPEG', 'PNG')): # input single img path
         input_img_list = [args.input_path]
         result_root = f'results/test_img_{w}'
-    elif args.input_path.endswith(('mp4', 'mov', 'avi')): # input video path
+    elif args.input_path.endswith(('mp4', 'mov', 'avi', 'MP4', 'MOV', 'AVI')): # input video path
         from basicsr.utils.video_util import VideoReader, VideoWriter
         input_img_list = []
         vidreader = VideoReader(args.input_path)
@@ -100,7 +99,7 @@ if __name__ == '__main__':
         if args.input_path.endswith('/'):  # solve when path ends with /
             args.input_path = args.input_path[:-1]
         # scan all the jpg and png images
-        input_img_list = sorted(glob.glob(os.path.join(args.input_path, '*.[jp][pn]g')))
+        input_img_list = sorted(glob.glob(os.path.join(args.input_path, '*.[jpJP][pnPN]*[gG]')))
         result_root = f'results/{os.path.basename(args.input_path)}_{w}'
 
     if not args.output_path is None: # set output path

web-demos/hugging_face/app.py

@@ -13,15 +13,16 @@ import gradio as gr
 
 from torchvision.transforms.functional import normalize
 
+from basicsr.archs.rrdbnet_arch import RRDBNet
 from basicsr.utils import imwrite, img2tensor, tensor2img
 from basicsr.utils.download_util import load_file_from_url
-from facelib.utils.face_restoration_helper import FaceRestoreHelper
-from facelib.utils.misc import is_gray
-from basicsr.archs.rrdbnet_arch import RRDBNet
+from basicsr.utils.misc import gpu_is_available, get_device
 from basicsr.utils.realesrgan_utils import RealESRGANer
-
 from basicsr.utils.registry import ARCH_REGISTRY
 
+from facelib.utils.face_restoration_helper import FaceRestoreHelper
+from facelib.utils.misc import is_gray
+
 
 os.system("pip freeze")
 
@@ -65,7 +66,8 @@ def imread(img_path):
 
 # set enhancer with RealESRGAN
 def set_realesrgan():
-    half = True if torch.cuda.is_available() else False
+    # half = True if torch.cuda.is_available() else False
+    half = True if gpu_is_available() else False
     model = RRDBNet(
         num_in_ch=3,
         num_out_ch=3,
@@ -86,7 +88,8 @@ def set_realesrgan():
     return upsampler
 
 upsampler = set_realesrgan()
-device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
+# device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
+device = get_device()
 codeformer_net = ARCH_REGISTRY.get("CodeFormer")(
     dim_embd=512,
     codebook_size=1024,

web-demos/replicate/predict.py

@@ -14,12 +14,13 @@ try:
 except Exception:
     print('please install cog package')
 
-from basicsr.utils import imwrite, img2tensor, tensor2img
 from basicsr.archs.rrdbnet_arch import RRDBNet
+from basicsr.utils import imwrite, img2tensor, tensor2img
 from basicsr.utils.realesrgan_utils import RealESRGANer
+from basicsr.utils.misc import gpu_is_available
 from basicsr.utils.registry import ARCH_REGISTRY
-from facelib.utils.face_restoration_helper import FaceRestoreHelper
 
+from facelib.utils.face_restoration_helper import FaceRestoreHelper
 
 class Predictor(BasePredictor):
     def setup(self):
@@ -159,7 +160,8 @@ def imread(img_path):
 
 
 def set_realesrgan():
-    if not torch.cuda.is_available():  # CPU
+    # if not torch.cuda.is_available():  # CPU
+    if not gpu_is_available():  # CPU
         import warnings
 
         warnings.warn(