mirror of
https://github.com/vladmandic/automatic
synced 2026-09-19 09:14:35 +02:00
cleanup codeformer
This commit is contained in:
Vladimir Mandic
+7
-1
@@ -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
|
||||
|
||||
|
||||
@@ -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
|
||||
@@ -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])
|
||||
for method in urls.keys():
|
||||
download_pretrained_models(urls[method])
|
||||
|
||||
@@ -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])
|
||||
for method in urls.keys():
|
||||
download_pretrained_models(urls[method])
|
||||
|
||||
@@ -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"""<center><img src='https://user-images.githubusercontent.com/14334509/189166076-94bb2cac-4f4e-40fb-a69f-66709e3d98f5.png' alt='CodeFormer logo'></center>
|
||||
<b>Official Gradio demo</b> for <a href='https://github.com/sczhou/CodeFormer' target='_blank'><b>Towards Robust Blind Face Restoration with Codebook Lookup Transformer (NeurIPS 2022)</b></a>.<br>
|
||||
🔥 CodeFormer is a robust face restoration algorithm for old photos or AI-generated faces.<br>
|
||||
🤗 Try CodeFormer for improved stable-diffusion generation!<br>
|
||||
"""
|
||||
article = r"""
|
||||
If CodeFormer is helpful, please help to ⭐ the <a href='https://github.com/sczhou/CodeFormer' target='_blank'>Github Repo</a>. 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 <a rel="license" href="https://github.com/sczhou/CodeFormer/blob/master/LICENSE">S-Lab License 1.0</a>.
|
||||
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 <b>shangchenzhou@gmail.com</b>.
|
||||
|
||||
<div>
|
||||
🤗 Find Me:
|
||||
<a href="https://twitter.com/ShangchenZhou"><img style="margin-top:0.5em; margin-bottom:0.5em" src="https://img.shields.io/twitter/follow/ShangchenZhou?label=%40ShangchenZhou&style=social" alt="Twitter Follow"></a>
|
||||
<a href="https://github.com/sczhou"><img style="margin-top:0.5em; margin-bottom:2em" src="https://img.shields.io/github/followers/sczhou?style=social" alt="Github Follow"></a>
|
||||
</div>
|
||||
|
||||
<center><img src='https://visitor-badge-sczhou.glitch.me/badge?page_id=sczhou/CodeFormer' alt='visitors'></center>
|
||||
"""
|
||||
|
||||
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()
|
||||
@@ -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"
|
||||
@@ -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
|
||||
@@ -1,3 +0,0 @@
|
||||
# Weights
|
||||
|
||||
Put the downloaded pre-trained models to this folder.
|
||||
Reference in New Issue
Block a user