diff --git a/CHANGELOG.md b/CHANGELOG.md
index 9767f1d7d..1e88ec824 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -1,12 +1,18 @@
# Change Log for SD.Next
-## Update for 2023-11-07
+## Update for 2023-11-08
- **Extra networks**
- Use multi-threading for 5x load speedup
+- **General**:
+ - Reworked parser when pasting previously generated images/prompts
- **Diffusers**
- Fix DPM SDE scheduler
+- **Fixes**
- Fix inpaint
+ - More uniform models paths
+ - Improve extension compatibility
+ - Improve BF16 support
## Update for 2023-11-06
diff --git a/repositories/codeformer/facelib/detection/__init__.py b/repositories/codeformer/facelib/detection/__init__.py
index 5d1f8fc21..1c021d410 100644
--- a/repositories/codeformer/facelib/detection/__init__.py
+++ b/repositories/codeformer/facelib/detection/__init__.py
@@ -1,14 +1,16 @@
import os
+from copy import deepcopy
import torch
from torch import nn
-from copy import deepcopy
-
from facelib.utils import load_file_from_url
from facelib.utils import download_pretrained_models
from facelib.detection.yolov5face.models.common import Conv
-
from .retinaface.retinaface import RetinaFace
from .yolov5face.face_detector import YoloDetector
+from modules import paths
+
+
+model_dir = os.path.join(paths.models_path, 'Codeformer')
def init_detection_model(model_name, half=False, device='cuda'):
@@ -32,7 +34,7 @@ def init_retinaface_model(model_name, half=False, device='cuda'):
else:
raise NotImplementedError(f'{model_name} is not implemented.')
- model_path = load_file_from_url(url=model_url, model_dir='weights/facelib', progress=True, file_name=None)
+ model_path = load_file_from_url(url=model_url, model_dir=model_dir, progress=True, file_name=None)
load_net = torch.load(model_path, map_location=lambda storage, loc: storage)
# remove unnecessary 'module.'
for k, v in deepcopy(load_net).items():
@@ -55,8 +57,8 @@ def init_yolov5face_model(model_name, device='cuda'):
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 = load_file_from_url(url=model_url, model_dir='weights/facelib', progress=True, file_name=None)
+
+ model_path = load_file_from_url(url=model_url, model_dir=model_dir, 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()
@@ -69,32 +71,3 @@ def init_yolov5face_model(model_name, device='cuda'):
m._non_persistent_buffers_set = set() # pytorch 1.6.0 compatibility
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
\ No newline at end of file
diff --git a/repositories/codeformer/scripts/download_pretrained_models.py b/repositories/codeformer/scripts/download_pretrained_models.py
index daa6e8ca1..16ff6e170 100644
--- a/repositories/codeformer/scripts/download_pretrained_models.py
+++ b/repositories/codeformer/scripts/download_pretrained_models.py
@@ -1,40 +1,26 @@
-import argparse
import os
-from os import path as osp
-
+from modules import paths
from basicsr.utils.download_util import load_file_from_url
-def download_pretrained_models(method, file_urls):
- save_path_root = f'./weights/{method}'
- os.makedirs(save_path_root, exist_ok=True)
+urls = {
+ 'CodeFormer': {
+ 'codeformer.pth': 'https://github.com/sczhou/CodeFormer/releases/download/v0.1.0/codeformer.pth'
+ },
+ 'facelib': {
+ # 'yolov5l-face.pth': 'https://github.com/sczhou/CodeFormer/releases/download/v0.1.0/yolov5l-face.pth',
+ 'detection_Resnet50_Final.pth': 'https://github.com/sczhou/CodeFormer/releases/download/v0.1.0/detection_Resnet50_Final.pth',
+ 'parsing_parsenet.pth': 'https://github.com/sczhou/CodeFormer/releases/download/v0.1.0/parsing_parsenet.pth'
+ }
+}
+
+def download_pretrained_models(file_urls):
+ model_dir = os.path.join(paths.models_path, 'Codeformer')
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)
+ load_file_from_url(url=file_url, model_dir=model_dir, progress=True, file_name=file_name)
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_urls = {
- 'CodeFormer': {
- 'codeformer.pth': 'https://github.com/sczhou/CodeFormer/releases/download/v0.1.0/codeformer.pth'
- },
- 'facelib': {
- # 'yolov5l-face.pth': 'https://github.com/sczhou/CodeFormer/releases/download/v0.1.0/yolov5l-face.pth',
- 'detection_Resnet50_Final.pth': 'https://github.com/sczhou/CodeFormer/releases/download/v0.1.0/detection_Resnet50_Final.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_urls.keys():
- download_pretrained_models(method, file_urls[method])
- else:
- download_pretrained_models(args.method, file_urls[args.method])
\ No newline at end of file
+ for method in urls.keys():
+ download_pretrained_models(urls[method])
diff --git a/repositories/codeformer/scripts/download_pretrained_models_from_gdrive.py b/repositories/codeformer/scripts/download_pretrained_models_from_gdrive.py
index 7df5be6fc..5a5c6bd44 100644
--- a/repositories/codeformer/scripts/download_pretrained_models_from_gdrive.py
+++ b/repositories/codeformer/scripts/download_pretrained_models_from_gdrive.py
@@ -1,18 +1,16 @@
-import argparse
import os
+from modules import paths
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)
+model_dir = os.path.join(paths.models_path, 'Codeformer')
+
+def download_pretrained_models(file_ids):
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))
+ save_path = osp.abspath(osp.join(model_dir, 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':
@@ -29,21 +27,13 @@ def download_pretrained_models(method, file_ids):
# 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 = {
+ urls = {
'CodeFormer': {
'codeformer.pth': '1v_E_vZvP-dQPF55Kc5SRCjaKTQXDz-JB'
},
@@ -52,9 +42,5 @@ if __name__ == '__main__':
'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])
\ No newline at end of file
+ for method in urls.keys():
+ download_pretrained_models(urls[method])
diff --git a/repositories/codeformer/web-demos/hugging_face/app.py b/repositories/codeformer/web-demos/hugging_face/app.py
deleted file mode 100644
index 7da0fc947..000000000
--- a/repositories/codeformer/web-demos/hugging_face/app.py
+++ /dev/null
@@ -1,280 +0,0 @@
-"""
-This file is used for deploying hugging face demo:
-https://huggingface.co/spaces/sczhou/CodeFormer
-"""
-
-import sys
-sys.path.append('CodeFormer')
-import os
-import cv2
-import torch
-import torch.nn.functional as F
-import gradio as gr
-
-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
-from facelib.utils.misc import is_gray
-from basicsr.archs.rrdbnet_arch import RRDBNet
-from basicsr.utils.realesrgan_utils import RealESRGANer
-
-from basicsr.utils.registry import ARCH_REGISTRY
-
-
-os.system("pip freeze")
-
-pretrain_model_url = {
- 'codeformer': 'https://github.com/sczhou/CodeFormer/releases/download/v0.1.0/codeformer.pth',
- 'detection': 'https://github.com/sczhou/CodeFormer/releases/download/v0.1.0/detection_Resnet50_Final.pth',
- 'parsing': 'https://github.com/sczhou/CodeFormer/releases/download/v0.1.0/parsing_parsenet.pth',
- 'realesrgan': 'https://github.com/sczhou/CodeFormer/releases/download/v0.1.0/RealESRGAN_x2plus.pth'
-}
-# download weights
-if not os.path.exists('CodeFormer/weights/CodeFormer/codeformer.pth'):
- load_file_from_url(url=pretrain_model_url['codeformer'], model_dir='CodeFormer/weights/CodeFormer', progress=True, file_name=None)
-if not os.path.exists('CodeFormer/weights/facelib/detection_Resnet50_Final.pth'):
- load_file_from_url(url=pretrain_model_url['detection'], model_dir='CodeFormer/weights/facelib', progress=True, file_name=None)
-if not os.path.exists('CodeFormer/weights/facelib/parsing_parsenet.pth'):
- load_file_from_url(url=pretrain_model_url['parsing'], model_dir='CodeFormer/weights/facelib', progress=True, file_name=None)
-if not os.path.exists('CodeFormer/weights/realesrgan/RealESRGAN_x2plus.pth'):
- load_file_from_url(url=pretrain_model_url['realesrgan'], model_dir='CodeFormer/weights/realesrgan', progress=True, file_name=None)
-
-# download images
-torch.hub.download_url_to_file(
- 'https://replicate.com/api/models/sczhou/codeformer/files/fa3fe3d1-76b0-4ca8-ac0d-0a925cb0ff54/06.png',
- '01.png')
-torch.hub.download_url_to_file(
- 'https://replicate.com/api/models/sczhou/codeformer/files/a1daba8e-af14-4b00-86a4-69cec9619b53/04.jpg',
- '02.jpg')
-torch.hub.download_url_to_file(
- 'https://replicate.com/api/models/sczhou/codeformer/files/542d64f9-1712-4de7-85f7-3863009a7c3d/03.jpg',
- '03.jpg')
-torch.hub.download_url_to_file(
- 'https://replicate.com/api/models/sczhou/codeformer/files/a11098b0-a18a-4c02-a19a-9a7045d68426/010.jpg',
- '04.jpg')
-torch.hub.download_url_to_file(
- 'https://replicate.com/api/models/sczhou/codeformer/files/7cf19c2c-e0cf-4712-9af8-cf5bdbb8d0ee/012.jpg',
- '05.jpg')
-
-def imread(img_path):
- img = cv2.imread(img_path)
- img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
- return img
-
-# set enhancer with RealESRGAN
-def set_realesrgan():
- half = True if torch.cuda.is_available() else False
- model = RRDBNet(
- num_in_ch=3,
- num_out_ch=3,
- num_feat=64,
- num_block=23,
- num_grow_ch=32,
- scale=2,
- )
- upsampler = RealESRGANer(
- scale=2,
- model_path="CodeFormer/weights/realesrgan/RealESRGAN_x2plus.pth",
- model=model,
- tile=400,
- tile_pad=40,
- pre_pad=0,
- half=half,
- )
- return upsampler
-
-upsampler = set_realesrgan()
-device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
-codeformer_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 = "CodeFormer/weights/CodeFormer/codeformer.pth"
-checkpoint = torch.load(ckpt_path)["params_ema"]
-codeformer_net.load_state_dict(checkpoint)
-codeformer_net.eval()
-
-os.makedirs('output', exist_ok=True)
-
-def inference(image, background_enhance, face_upsample, upscale, codeformer_fidelity):
- """Run a single prediction on the model"""
- try: # global try
- # take the default setting for the demo
- has_aligned = False
- only_center_face = False
- draw_box = False
- detection_model = "retinaface_resnet50"
- print('Inp:', image, background_enhance, face_upsample, upscale, codeformer_fidelity)
-
- img = cv2.imread(str(image), cv2.IMREAD_COLOR)
- print('\timage size:', img.shape)
-
- upscale = int(upscale) # convert type to int
- if upscale > 4: # avoid memory exceeded due to too large upscale
- upscale = 4
- if upscale > 2 and max(img.shape[:2])>1000: # avoid memory exceeded due to too large img resolution
- upscale = 2
- if max(img.shape[:2]) > 1500: # avoid memory exceeded due to too large img resolution
- upscale = 1
- background_enhance = False
- face_upsample = False
-
- face_helper = FaceRestoreHelper(
- upscale,
- face_size=512,
- crop_ratio=(1, 1),
- det_model=detection_model,
- save_ext="png",
- use_parse=True,
- device=device,
- )
- bg_upsampler = upsampler if background_enhance else None
- face_upsampler = upsampler if face_upsample else None
-
- if has_aligned:
- # the input faces are already cropped and aligned
- img = cv2.resize(img, (512, 512), interpolation=cv2.INTER_LINEAR)
- face_helper.is_gray = is_gray(img, threshold=5)
- if face_helper.is_gray:
- print('\tgrayscale input: True')
- face_helper.cropped_faces = [img]
- else:
- face_helper.read_image(img)
- # get face landmarks for each face
- num_det_faces = face_helper.get_face_landmarks_5(
- only_center_face=only_center_face, resize=640, eye_dist_threshold=5
- )
- print(f'\tdetect {num_det_faces} faces')
- # align and warp each face
- face_helper.align_warp_face()
-
- # face restoration for each cropped face
- for idx, cropped_face in enumerate(face_helper.cropped_faces):
- # prepare data
- cropped_face_t = img2tensor(
- cropped_face / 255.0, 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(device)
-
- try:
- with torch.no_grad():
- output = codeformer_net(
- cropped_face_t, w=codeformer_fidelity, adain=True
- )[0]
- restored_face = tensor2img(output, rgb2bgr=True, min_max=(-1, 1))
- del output
- torch.cuda.empty_cache()
- except RuntimeError as error:
- print(f"Failed inference for CodeFormer: {error}")
- restored_face = tensor2img(
- cropped_face_t, rgb2bgr=True, min_max=(-1, 1)
- )
-
- restored_face = restored_face.astype("uint8")
- face_helper.add_restored_face(restored_face)
-
- # paste_back
- if not has_aligned:
- # upsample the background
- if bg_upsampler is not None:
- # Now only support RealESRGAN for upsampling background
- bg_img = bg_upsampler.enhance(img, outscale=upscale)[0]
- else:
- bg_img = None
- face_helper.get_inverse_affine(None)
- # paste each restored face to the input image
- if face_upsample and face_upsampler is not None:
- restored_img = face_helper.paste_faces_to_input_image(
- upsample_img=bg_img,
- draw_box=draw_box,
- face_upsampler=face_upsampler,
- )
- else:
- restored_img = face_helper.paste_faces_to_input_image(
- upsample_img=bg_img, draw_box=draw_box
- )
-
- # save restored img
- save_path = f'output/out.png'
- imwrite(restored_img, str(save_path))
-
- restored_img = cv2.cvtColor(restored_img, cv2.COLOR_BGR2RGB)
- return restored_img, save_path
- except Exception as error:
- print('Global exception', error)
- return None, None
-
-
-title = "CodeFormer: Robust Face Restoration and Enhancement Network"
-description = r"""
-Official Gradio demo for Towards Robust Blind Face Restoration with Codebook Lookup Transformer (NeurIPS 2022).
-🔥 CodeFormer is a robust face restoration algorithm for old photos or AI-generated faces.
-🤗 Try CodeFormer for improved stable-diffusion generation!
-"""
-article = r"""
-If CodeFormer is helpful, please help to ⭐ the Github Repo. Thanks!
-[](https://github.com/sczhou/CodeFormer)
-
----
-
-📝 **Citation**
-
-If our work is useful for your research, please consider citing:
-```bibtex
-@inproceedings{zhou2022codeformer,
- author = {Zhou, Shangchen and Chan, Kelvin C.K. and Li, Chongyi and Loy, Chen Change},
- title = {Towards Robust Blind Face Restoration with Codebook Lookup TransFormer},
- booktitle = {NeurIPS},
- year = {2022}
-}
-```
-
-📋 **License**
-
-This project is licensed under S-Lab License 1.0.
-Redistribution and use for non-commercial purposes should follow this license.
-
-📧 **Contact**
-
-If you have any questions, please feel free to reach me out at shangchenzhou@gmail.com.
-
-
- 🤗 Find Me:
-
-
-
-"""
-
-demo = gr.Interface(
- inference, [
- gr.inputs.Image(type="filepath", label="Input"),
- gr.inputs.Checkbox(default=True, label="Background_Enhance"),
- gr.inputs.Checkbox(default=True, label="Face_Upsample"),
- gr.inputs.Number(default=2, label="Rescaling_Factor (up to 4)"),
- gr.Slider(0, 1, value=0.5, step=0.01, label='Codeformer_Fidelity (0 for better quality, 1 for better identity)')
- ], [
- gr.outputs.Image(type="numpy", label="Output"),
- gr.outputs.File(label="Download the output")
- ],
- title=title,
- description=description,
- article=article,
- examples=[
- ['01.png', True, True, 2, 0.7],
- ['02.jpg', True, True, 2, 0.7],
- ['03.jpg', True, True, 2, 0.7],
- ['04.jpg', True, True, 2, 0.1],
- ['05.jpg', True, True, 2, 0.1]
- ]
- )
-
-demo.queue(concurrency_count=2)
-demo.launch()
\ No newline at end of file
diff --git a/repositories/codeformer/web-demos/replicate/cog.yaml b/repositories/codeformer/web-demos/replicate/cog.yaml
deleted file mode 100644
index 3f4589690..000000000
--- a/repositories/codeformer/web-demos/replicate/cog.yaml
+++ /dev/null
@@ -1,30 +0,0 @@
-"""
-This file is used for deploying replicate demo:
-https://replicate.com/sczhou/codeformer
-"""
-
-build:
- gpu: true
- cuda: "11.3"
- python_version: "3.8"
- system_packages:
- - "libgl1-mesa-glx"
- - "libglib2.0-0"
- python_packages:
- - "ipython==8.4.0"
- - "future==0.18.2"
- - "lmdb==1.3.0"
- - "scikit-image==0.19.3"
- - "torch==1.11.0 --extra-index-url=https://download.pytorch.org/whl/cu113"
- - "torchvision==0.12.0 --extra-index-url=https://download.pytorch.org/whl/cu113"
- - "scipy==1.9.0"
- - "gdown==4.5.1"
- - "pyyaml==6.0"
- - "tb-nightly==2.11.0a20220906"
- - "tqdm==4.64.1"
- - "yapf==0.32.0"
- - "lpips==0.1.4"
- - "Pillow==9.2.0"
- - "opencv-python==4.6.0.66"
-
-predict: "predict.py:Predictor"
diff --git a/repositories/codeformer/web-demos/replicate/predict.py b/repositories/codeformer/web-demos/replicate/predict.py
deleted file mode 100644
index 61935e9e7..000000000
--- a/repositories/codeformer/web-demos/replicate/predict.py
+++ /dev/null
@@ -1,189 +0,0 @@
-"""
-This file is used for deploying replicate demo:
-https://replicate.com/sczhou/codeformer
-running: cog predict -i image=@inputs/whole_imgs/04.jpg -i codeformer_fidelity=0.5 -i upscale=2
-push: cog push r8.im/sczhou/codeformer
-"""
-
-import tempfile
-import cv2
-import torch
-from torchvision.transforms.functional import normalize
-try:
- from cog import BasePredictor, Input, Path
-except Exception:
- print('please install cog package')
-
-from basicsr.utils import imwrite, img2tensor, tensor2img
-from basicsr.archs.rrdbnet_arch import RRDBNet
-from basicsr.utils.realesrgan_utils import RealESRGANer
-from basicsr.utils.registry import ARCH_REGISTRY
-from facelib.utils.face_restoration_helper import FaceRestoreHelper
-
-
-class Predictor(BasePredictor):
- def setup(self):
- """Load the model into memory to make running multiple predictions efficient"""
- self.device = "cuda:0"
- self.upsampler = set_realesrgan()
- self.net = ARCH_REGISTRY.get("CodeFormer")(
- dim_embd=512,
- codebook_size=1024,
- n_head=8,
- n_layers=9,
- connect_list=["32", "64", "128", "256"],
- ).to(self.device)
- ckpt_path = "weights/CodeFormer/codeformer.pth"
- checkpoint = torch.load(ckpt_path)[
- "params_ema"
- ] # update file permission if cannot load
- self.net.load_state_dict(checkpoint)
- self.net.eval()
-
- def predict(
- self,
- image: Path = Input(description="Input image"),
- codeformer_fidelity: float = Input(
- default=0.5,
- ge=0,
- le=1,
- description="Balance the quality (lower number) and fidelity (higher number).",
- ),
- background_enhance: bool = Input(
- description="Enhance background image with Real-ESRGAN", default=True
- ),
- face_upsample: bool = Input(
- description="Upsample restored faces for high-resolution AI-created images",
- default=True,
- ),
- upscale: int = Input(
- description="The final upsampling scale of the image",
- default=2,
- ),
- ) -> Path:
- """Run a single prediction on the model"""
-
- # take the default setting for the demo
- has_aligned = False
- only_center_face = False
- draw_box = False
- detection_model = "retinaface_resnet50"
-
- self.face_helper = FaceRestoreHelper(
- upscale,
- face_size=512,
- crop_ratio=(1, 1),
- det_model=detection_model,
- save_ext="png",
- use_parse=True,
- device=self.device,
- )
-
- bg_upsampler = self.upsampler if background_enhance else None
- face_upsampler = self.upsampler if face_upsample else None
-
- img = cv2.imread(str(image), cv2.IMREAD_COLOR)
-
- if has_aligned:
- # the input faces are already cropped and aligned
- img = cv2.resize(img, (512, 512), interpolation=cv2.INTER_LINEAR)
- self.face_helper.cropped_faces = [img]
- else:
- self.face_helper.read_image(img)
- # get face landmarks for each face
- num_det_faces = self.face_helper.get_face_landmarks_5(
- only_center_face=only_center_face, resize=640, eye_dist_threshold=5
- )
- print(f"\tdetect {num_det_faces} faces")
- # align and warp each face
- self.face_helper.align_warp_face()
-
- # face restoration for each cropped face
- for idx, cropped_face in enumerate(self.face_helper.cropped_faces):
- # prepare data
- cropped_face_t = img2tensor(
- cropped_face / 255.0, 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(self.device)
-
- try:
- with torch.no_grad():
- output = self.net(
- cropped_face_t, w=codeformer_fidelity, adain=True
- )[0]
- restored_face = tensor2img(output, rgb2bgr=True, min_max=(-1, 1))
- del output
- torch.cuda.empty_cache()
- except Exception as error:
- print(f"\tFailed inference for CodeFormer: {error}")
- restored_face = tensor2img(
- cropped_face_t, rgb2bgr=True, min_max=(-1, 1)
- )
-
- restored_face = restored_face.astype("uint8")
- self.face_helper.add_restored_face(restored_face)
-
- # paste_back
- if not has_aligned:
- # upsample the background
- if bg_upsampler is not None:
- # Now only support RealESRGAN for upsampling background
- bg_img = bg_upsampler.enhance(img, outscale=upscale)[0]
- else:
- bg_img = None
- self.face_helper.get_inverse_affine(None)
- # paste each restored face to the input image
- if face_upsample and face_upsampler is not None:
- restored_img = self.face_helper.paste_faces_to_input_image(
- upsample_img=bg_img,
- draw_box=draw_box,
- face_upsampler=face_upsampler,
- )
- else:
- restored_img = self.face_helper.paste_faces_to_input_image(
- upsample_img=bg_img, draw_box=draw_box
- )
-
- # save restored img
- out_path = Path(tempfile.mkdtemp()) / 'output.png'
- imwrite(restored_img, str(out_path))
-
- return out_path
-
-
-def imread(img_path):
- img = cv2.imread(img_path)
- img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
- return img
-
-
-def set_realesrgan():
- if not torch.cuda.is_available(): # CPU
- import warnings
-
- warnings.warn(
- "The unoptimized RealESRGAN is slow on CPU. We do not use it. "
- "If you really want to use it, please modify the corresponding codes.",
- category=RuntimeWarning,
- )
- upsampler = None
- else:
- model = RRDBNet(
- num_in_ch=3,
- num_out_ch=3,
- num_feat=64,
- num_block=23,
- num_grow_ch=32,
- scale=2,
- )
- upsampler = RealESRGANer(
- scale=2,
- model_path="./weights/realesrgan/RealESRGAN_x2plus.pth",
- model=model,
- tile=400,
- tile_pad=40,
- pre_pad=0,
- half=True,
- )
- return upsampler
diff --git a/repositories/codeformer/weights/CodeFormer/.gitkeep b/repositories/codeformer/weights/CodeFormer/.gitkeep
deleted file mode 100644
index e69de29bb..000000000
diff --git a/repositories/codeformer/weights/README.md b/repositories/codeformer/weights/README.md
deleted file mode 100644
index 67ad334bd..000000000
--- a/repositories/codeformer/weights/README.md
+++ /dev/null
@@ -1,3 +0,0 @@
-# Weights
-
-Put the downloaded pre-trained models to this folder.
\ No newline at end of file
diff --git a/repositories/codeformer/weights/facelib/.gitkeep b/repositories/codeformer/weights/facelib/.gitkeep
deleted file mode 100644
index e69de29bb..000000000