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 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""" # take the default setting for the demo has_aligned = False only_center_face = False draw_box = False detection_model = "retinaface_resnet50" 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 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) 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 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") 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 title = "CodeFormer: Robust Face Restoration and Enhancement Network" description = r"""
CodeFormer logo
Official Gradio demo for Towards Robust Blind Face Restoration with Codebook Lookup Transformer.
🔥 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! [![GitHub Stars](https://img.shields.io/github/stars/sczhou/CodeFormer?style=social)](https://github.com/sczhou/CodeFormer) --- 📝 Citation If our work is useful for your research, please consider citing: ```bibtex @article{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}, journal = {arXiv preprint arXiv:2206.11253}, year = {2022} } ``` 📧 Contact If you have any questions, please feel free to reach me out at shangchenzhou@gmail.com. ![visitors](https://visitor-badge.glitch.me/badge?page_id=sczhou/CodeFormer) """ 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"), 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"), ], 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] ] ).launch()