V0.1.0 release.

README.md

@@ -6,7 +6,8 @@
 
 [Paper](https://arxiv.org/abs/2206.11253) | [Project Page](https://shangchenzhou.com/projects/CodeFormer/) | [Video](https://youtu.be/d3VDpkXlueI)
 
-<a href="https://colab.research.google.com/drive/1m52PNveE4PBhYrecj34cnpEeiHcC5LTb?usp=sharing"><img src="https://colab.research.google.com/assets/colab-badge.svg" alt="google colab logo"></a>
+<a href="https://colab.research.google.com/drive/1m52PNveE4PBhYrecj34cnpEeiHcC5LTb?usp=sharing"><img src="https://colab.research.google.com/assets/colab-badge.svg" alt="google colab logo"></a> ![visitors](https://visitor-badge.glitch.me/badge?page_id=sczhou/CodeFormer)
+
 
 [Shangchen Zhou](https://shangchenzhou.com/), [Kelvin C.K. Chan](https://ckkelvinchan.github.io/), [Chongyi Li](https://li-chongyi.github.io/), [Chen Change Loy](https://www.mmlab-ntu.com/person/ccloy/) 
 

facelib/detection/__init__.py

@@ -49,17 +49,14 @@ def init_retinaface_model(model_name, half=False, device='cuda'):
 def init_yolov5face_model(model_name, device='cuda'):
     if model_name == 'YOLOv5l':
         model = YoloDetector(config_name='facelib/detection/yolov5face/models/yolov5l.yaml', device=device)
-        f_id = {'yolov5l-face.pth': '131578zMA6B2x8VQHyHfa6GEPtulMCNzV'}
+        model_url = 'https://github.com/sczhou/CodeFormer/releases/download/v0.1.0/yolov5l-face.pth'
     elif model_name == 'YOLOv5n':
         model = YoloDetector(config_name='facelib/detection/yolov5face/models/yolov5n.yaml', device=device)
-        f_id = {'yolov5n-face.pth': '1fhcpFvWZqghpGXjYPIne2sw1Fy4yhw6o'}
+        model_url = 'https://github.com/sczhou/CodeFormer/releases/download/v0.1.0/yolov5n-face.pth'
     else:
         raise NotImplementedError(f'{model_name} is not implemented.')
-
-    model_path = os.path.join('weights/facelib', list(f_id.keys())[0])
-    if not os.path.exists(model_path):
-        download_pretrained_models(file_ids=f_id, save_path_root='weights/facelib')
-
+    
+    model_path = load_file_from_url(url=model_url, model_dir='weights/facelib', progress=True, file_name=None)
     load_net = torch.load(model_path, map_location=lambda storage, loc: storage)
     model.detector.load_state_dict(load_net, strict=True)
     model.detector.eval()
@@ -71,4 +68,33 @@ def init_yolov5face_model(model_name, device='cuda'):
         elif isinstance(m, Conv):
             m._non_persistent_buffers_set = set()  # pytorch 1.6.0 compatibility
 
-    return model
+    return model
+
+
+# Download from Google Drive
+# def init_yolov5face_model(model_name, device='cuda'):
+#     if model_name == 'YOLOv5l':
+#         model = YoloDetector(config_name='facelib/detection/yolov5face/models/yolov5l.yaml', device=device)
+#         f_id = {'yolov5l-face.pth': '131578zMA6B2x8VQHyHfa6GEPtulMCNzV'}
+#     elif model_name == 'YOLOv5n':
+#         model = YoloDetector(config_name='facelib/detection/yolov5face/models/yolov5n.yaml', device=device)
+#         f_id = {'yolov5n-face.pth': '1fhcpFvWZqghpGXjYPIne2sw1Fy4yhw6o'}
+#     else:
+#         raise NotImplementedError(f'{model_name} is not implemented.')
+
+#     model_path = os.path.join('weights/facelib', list(f_id.keys())[0])
+#     if not os.path.exists(model_path):
+#         download_pretrained_models(file_ids=f_id, save_path_root='weights/facelib')
+
+#     load_net = torch.load(model_path, map_location=lambda storage, loc: storage)
+#     model.detector.load_state_dict(load_net, strict=True)
+#     model.detector.eval()
+#     model.detector = model.detector.to(device).float()
+
+#     for m in model.detector.modules():
+#         if type(m) in [nn.Hardswish, nn.LeakyReLU, nn.ReLU, nn.ReLU6, nn.SiLU]:
+#             m.inplace = True  # pytorch 1.7.0 compatibility
+#         elif isinstance(m, Conv):
+#             m._non_persistent_buffers_set = set()  # pytorch 1.6.0 compatibility
+
+#     return model

facelib/parsing/__init__.py

@@ -8,10 +8,10 @@ from .parsenet import ParseNet
 def init_parsing_model(model_name='bisenet', half=False, device='cuda'):
     if model_name == 'bisenet':
         model = BiSeNet(num_class=19)
-        model_url = 'https://github.com/xinntao/facexlib/releases/download/v0.2.0/parsing_bisenet.pth'
+        model_url = 'https://github.com/sczhou/CodeFormer/releases/download/v0.1.0/parsing_bisenet.pth'
     elif model_name == 'parsenet':
         model = ParseNet(in_size=512, out_size=512, parsing_ch=19)
-        model_url = 'https://github.com/xinntao/facexlib/releases/download/v0.2.2/parsing_parsenet.pth'
+        model_url = 'https://github.com/sczhou/CodeFormer/releases/download/v0.1.0/parsing_parsenet.pth'
     else:
         raise NotImplementedError(f'{model_name} is not implemented.')
 

inference_codeformer.py

@@ -6,11 +6,16 @@ import glob
 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 facelib.utils.face_restoration_helper import FaceRestoreHelper
 import torch.nn.functional as F
 
 from basicsr.utils.registry import ARCH_REGISTRY
 
+pretrain_model_url = {
+    'restoration': 'https://github.com/sczhou/CodeFormer/releases/download/v0.1.0/codeformer.pth',
+}
+
 if __name__ == '__main__':
     device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
     parser = argparse.ArgumentParser()
@@ -59,8 +64,10 @@ if __name__ == '__main__':
     # ------------------ set up CodeFormer restorer -------------------
     net = ARCH_REGISTRY.get('CodeFormer')(dim_embd=512, codebook_size=1024, n_head=8, n_layers=9, 
                                             connect_list=['32', '64', '128', '256']).to(device)
-
-    ckpt_path = 'weights/CodeFormer/codeformer.pth'
+    
+    # ckpt_path = 'weights/CodeFormer/codeformer.pth'
+    ckpt_path = load_file_from_url(url=pretrain_model_url['restoration'], 
+                                    model_dir='weights/CodeFormer', progress=True, file_name=None)
     checkpoint = torch.load(ckpt_path)['params_ema']
     net.load_state_dict(checkpoint)
     net.eval()

requirements.txt

@@ -14,7 +14,7 @@ torchvision
 tqdm
 yapf
 lpips
-gdown # supports downloading the large file from Google Drive
+# gdown # supports downloading the large file from Google Drive
 # cmake
 # dlib
 # conda install -c conda-forge dlib

scripts/download_pretrained_models.py

@@ -2,31 +2,16 @@ import argparse
 import os
 from os import path as osp
 
-# from basicsr.utils.download_util import download_file_from_google_drive
-import gdown
+from basicsr.utils.download_util import load_file_from_url
 
 
-def download_pretrained_models(method, file_ids):
+def download_pretrained_models(method, file_urls):
     save_path_root = f'./weights/{method}'
     os.makedirs(save_path_root, exist_ok=True)
 
-    for file_name, file_id in file_ids.items():
-        file_url = 'https://drive.google.com/uc?id='+file_id
-        save_path = osp.abspath(osp.join(save_path_root, file_name))
-        if osp.exists(save_path):
-            user_response = input(f'{file_name} already exist. Do you want to cover it? Y/N\n')
-            if user_response.lower() == 'y':
-                print(f'Covering {file_name} to {save_path}')
-                gdown.download(file_url, save_path, quiet=False)
-                # download_file_from_google_drive(file_id, save_path)
-            elif user_response.lower() == 'n':
-                print(f'Skipping {file_name}')
-            else:
-                raise ValueError('Wrong input. Only accepts Y/N.')
-        else:
-            print(f'Downloading {file_name} to {save_path}')
-            gdown.download(file_url, save_path, quiet=False)
-            # download_file_from_google_drive(file_id, save_path)
+    for file_name, file_url in file_urls.items():
+        save_path = load_file_from_url(url=file_url, model_dir=save_path_root, progress=True, file_name=file_name)
+
 
 if __name__ == '__main__':
     parser = argparse.ArgumentParser()
@@ -37,24 +22,18 @@ if __name__ == '__main__':
         help=("Options: 'CodeFormer' 'facelib'. Set to 'all' to download all the models."))
     args = parser.parse_args()
 
-    # file name: file id
-    # 'dlib': {
-    #     'mmod_human_face_detector-4cb19393.dat': '1qD-OqY8M6j4PWUP_FtqfwUPFPRMu6ubX',
-    #     'shape_predictor_5_face_landmarks-c4b1e980.dat': '1vF3WBUApw4662v9Pw6wke3uk1qxnmLdg',
-    #     'shape_predictor_68_face_landmarks-fbdc2cb8.dat': '1tJyIVdCHaU6IDMDx86BZCxLGZfsWB8yq'
-    # }
-    file_ids = {
+    file_urls = {
         'CodeFormer': {
-            'codeformer.pth': '1v_E_vZvP-dQPF55Kc5SRCjaKTQXDz-JB'
+            'codeformer.pth': 'https://github.com/sczhou/CodeFormer/releases/download/v0.1.0/codeformer.pth'
         },
         'facelib': {
-            'yolov5l-face.pth': '131578zMA6B2x8VQHyHfa6GEPtulMCNzV',
-            'parsing_parsenet.pth': '16pkohyZZ8ViHGBk3QtVqxLZKzdo466bK'
+            'yolov5l-face.pth': 'https://github.com/sczhou/CodeFormer/releases/download/v0.1.0/yolov5l-face.pth',
+            'parsing_parsenet.pth': 'https://github.com/sczhou/CodeFormer/releases/download/v0.1.0/parsing_parsenet.pth'
         }
     }
 
     if args.method == 'all':
-        for method in file_ids.keys():
-            download_pretrained_models(method, file_ids[method])
+        for method in file_urls.keys():
+            download_pretrained_models(method, file_urls[method])
     else:
-        download_pretrained_models(args.method, file_ids[args.method])
+        download_pretrained_models(args.method, file_urls[args.method])

scripts/download_pretrained_models_from_gdrive.py

@@ -0,0 +1,60 @@
+import argparse
+import os
+from os import path as osp
+
+# from basicsr.utils.download_util import download_file_from_google_drive
+import gdown
+
+
+def download_pretrained_models(method, file_ids):
+    save_path_root = f'./weights/{method}'
+    os.makedirs(save_path_root, exist_ok=True)
+
+    for file_name, file_id in file_ids.items():
+        file_url = 'https://drive.google.com/uc?id='+file_id
+        save_path = osp.abspath(osp.join(save_path_root, file_name))
+        if osp.exists(save_path):
+            user_response = input(f'{file_name} already exist. Do you want to cover it? Y/N\n')
+            if user_response.lower() == 'y':
+                print(f'Covering {file_name} to {save_path}')
+                gdown.download(file_url, save_path, quiet=False)
+                # download_file_from_google_drive(file_id, save_path)
+            elif user_response.lower() == 'n':
+                print(f'Skipping {file_name}')
+            else:
+                raise ValueError('Wrong input. Only accepts Y/N.')
+        else:
+            print(f'Downloading {file_name} to {save_path}')
+            gdown.download(file_url, save_path, quiet=False)
+            # download_file_from_google_drive(file_id, save_path)
+
+if __name__ == '__main__':
+    parser = argparse.ArgumentParser()
+
+    parser.add_argument(
+        'method',
+        type=str,
+        help=("Options: 'CodeFormer' 'facelib'. Set to 'all' to download all the models."))
+    args = parser.parse_args()
+
+    # file name: file id
+    # 'dlib': {
+    #     'mmod_human_face_detector-4cb19393.dat': '1qD-OqY8M6j4PWUP_FtqfwUPFPRMu6ubX',
+    #     'shape_predictor_5_face_landmarks-c4b1e980.dat': '1vF3WBUApw4662v9Pw6wke3uk1qxnmLdg',
+    #     'shape_predictor_68_face_landmarks-fbdc2cb8.dat': '1tJyIVdCHaU6IDMDx86BZCxLGZfsWB8yq'
+    # }
+    file_ids = {
+        'CodeFormer': {
+            'codeformer.pth': '1v_E_vZvP-dQPF55Kc5SRCjaKTQXDz-JB'
+        },
+        'facelib': {
+            'yolov5l-face.pth': '131578zMA6B2x8VQHyHfa6GEPtulMCNzV',
+            'parsing_parsenet.pth': '16pkohyZZ8ViHGBk3QtVqxLZKzdo466bK'
+        }
+    }
+
+    if args.method == 'all':
+        for method in file_ids.keys():
+            download_pretrained_models(method, file_ids[method])
+    else:
+        download_pretrained_models(args.method, file_ids[args.method])