Xiabou/codeformer

1"""2This file is used for deploying hugging face demo:3https://huggingface.co/spaces/sczhou/CodeFormer4"""5 6import sys7sys.path.append('CodeFormer')8import os9import cv210import torch11import torch.nn.functional as F12import gradio as gr13 14from torchvision.transforms.functional import normalize15 16from basicsr.utils import imwrite, img2tensor, tensor2img17from basicsr.utils.download_util import load_file_from_url18from facelib.utils.face_restoration_helper import FaceRestoreHelper19from facelib.utils.misc import is_gray20from basicsr.archs.rrdbnet_arch import RRDBNet21from basicsr.utils.realesrgan_utils import RealESRGANer22 23from basicsr.utils.registry import ARCH_REGISTRY24 25 26os.system("pip freeze")27 28pretrain_model_url = {29    'codeformer': 'https://github.com/sczhou/CodeFormer/releases/download/v0.1.0/codeformer.pth',30    'detection': 'https://github.com/sczhou/CodeFormer/releases/download/v0.1.0/detection_Resnet50_Final.pth',31    'parsing': 'https://github.com/sczhou/CodeFormer/releases/download/v0.1.0/parsing_parsenet.pth',32    'realesrgan': 'https://github.com/sczhou/CodeFormer/releases/download/v0.1.0/RealESRGAN_x2plus.pth'33}34# download weights35if not os.path.exists('CodeFormer/weights/CodeFormer/codeformer.pth'):36    load_file_from_url(url=pretrain_model_url['codeformer'], model_dir='CodeFormer/weights/CodeFormer', progress=True, file_name=None)37if not os.path.exists('CodeFormer/weights/facelib/detection_Resnet50_Final.pth'):38    load_file_from_url(url=pretrain_model_url['detection'], model_dir='CodeFormer/weights/facelib', progress=True, file_name=None)39if not os.path.exists('CodeFormer/weights/facelib/parsing_parsenet.pth'):40    load_file_from_url(url=pretrain_model_url['parsing'], model_dir='CodeFormer/weights/facelib', progress=True, file_name=None)41if not os.path.exists('CodeFormer/weights/realesrgan/RealESRGAN_x2plus.pth'):42    load_file_from_url(url=pretrain_model_url['realesrgan'], model_dir='CodeFormer/weights/realesrgan', progress=True, file_name=None)43 44# download images45torch.hub.download_url_to_file(46    'https://replicate.com/api/models/sczhou/codeformer/files/fa3fe3d1-76b0-4ca8-ac0d-0a925cb0ff54/06.png',47    '01.png')48torch.hub.download_url_to_file(49    'https://replicate.com/api/models/sczhou/codeformer/files/a1daba8e-af14-4b00-86a4-69cec9619b53/04.jpg',50    '02.jpg')51torch.hub.download_url_to_file(52    'https://replicate.com/api/models/sczhou/codeformer/files/542d64f9-1712-4de7-85f7-3863009a7c3d/03.jpg',53    '03.jpg')54torch.hub.download_url_to_file(55    'https://replicate.com/api/models/sczhou/codeformer/files/a11098b0-a18a-4c02-a19a-9a7045d68426/010.jpg',56    '04.jpg')57torch.hub.download_url_to_file(58    'https://replicate.com/api/models/sczhou/codeformer/files/7cf19c2c-e0cf-4712-9af8-cf5bdbb8d0ee/012.jpg',59    '05.jpg')60 61def imread(img_path):62    img = cv2.imread(img_path)63    img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)64    return img65 66# set enhancer with RealESRGAN67def set_realesrgan():68    half = True if torch.cuda.is_available() else False69    model = RRDBNet(70        num_in_ch=3,71        num_out_ch=3,72        num_feat=64,73        num_block=23,74        num_grow_ch=32,75        scale=2,76    )77    upsampler = RealESRGANer(78        scale=2,79        model_path="CodeFormer/weights/realesrgan/RealESRGAN_x2plus.pth",80        model=model,81        tile=400,82        tile_pad=40,83        pre_pad=0,84        half=half,85    )86    return upsampler87 88upsampler = set_realesrgan()89device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')90codeformer_net = ARCH_REGISTRY.get("CodeFormer")(91    dim_embd=512,92    codebook_size=1024,93    n_head=8,94    n_layers=9,95    connect_list=["32", "64", "128", "256"],96).to(device)97ckpt_path = "CodeFormer/weights/CodeFormer/codeformer.pth"98checkpoint = torch.load(ckpt_path)["params_ema"]99codeformer_net.load_state_dict(checkpoint)100codeformer_net.eval()101 102os.makedirs('output', exist_ok=True)103 104def inference(image, background_enhance, face_upsample, upscale, codeformer_fidelity):105    """Run a single prediction on the model"""106    # take the default setting for the demo107    has_aligned = False108    only_center_face = False109    draw_box = False110    detection_model = "retinaface_resnet50"111 112    upscale = int(upscale) # covert type to int113    face_helper = FaceRestoreHelper(114        upscale,115        face_size=512,116        crop_ratio=(1, 1),117        det_model=detection_model,118        save_ext="png",119        use_parse=True,120        device=device,121    )122    bg_upsampler = upsampler if background_enhance else None123    face_upsampler = upsampler if face_upsample else None124 125    img = cv2.imread(str(image), cv2.IMREAD_COLOR)126 127    if has_aligned:128        # the input faces are already cropped and aligned129        img = cv2.resize(img, (512, 512), interpolation=cv2.INTER_LINEAR)130        face_helper.is_gray = is_gray(img, threshold=5)131        if face_helper.is_gray:132            print('Grayscale input: True')133        face_helper.cropped_faces = [img]134    else:135        face_helper.read_image(img)136        # get face landmarks for each face137        num_det_faces = face_helper.get_face_landmarks_5(138          only_center_face=only_center_face, resize=640, eye_dist_threshold=5139        )140        print(f"\tdetect {num_det_faces} faces")141        # align and warp each face142        face_helper.align_warp_face()143 144    # face restoration for each cropped face145    for idx, cropped_face in enumerate(face_helper.cropped_faces):146        # prepare data147        cropped_face_t = img2tensor(148            cropped_face / 255.0, bgr2rgb=True, float32=True149        )150        normalize(cropped_face_t, (0.5, 0.5, 0.5), (0.5, 0.5, 0.5), inplace=True)151        cropped_face_t = cropped_face_t.unsqueeze(0).to(device)152 153        try:154            with torch.no_grad():155                output = codeformer_net(156                    cropped_face_t, w=codeformer_fidelity, adain=True157                )[0]158                restored_face = tensor2img(output, rgb2bgr=True, min_max=(-1, 1))159            del output160            torch.cuda.empty_cache()161        except Exception as error:162            print(f"\tFailed inference for CodeFormer: {error}")163            restored_face = tensor2img(164                cropped_face_t, rgb2bgr=True, min_max=(-1, 1)165            )166 167        restored_face = restored_face.astype("uint8")168        face_helper.add_restored_face(restored_face)169 170    # paste_back171    if not has_aligned:172        # upsample the background173        if bg_upsampler is not None:174            # Now only support RealESRGAN for upsampling background175            bg_img = bg_upsampler.enhance(img, outscale=upscale)[0]176        else:177            bg_img = None178        face_helper.get_inverse_affine(None)179        # paste each restored face to the input image180        if face_upsample and face_upsampler is not None:181            restored_img = face_helper.paste_faces_to_input_image(182                upsample_img=bg_img,183                draw_box=draw_box,184                face_upsampler=face_upsampler,185            )186        else:187            restored_img = face_helper.paste_faces_to_input_image(188                upsample_img=bg_img, draw_box=draw_box189            )190 191    # save restored img192    save_path = f'output/out.png'193    imwrite(restored_img, str(save_path))194 195    restored_img = cv2.cvtColor(restored_img, cv2.COLOR_BGR2RGB)196    return restored_img, save_path197 198 199 200title = "CodeFormer: Robust Face Restoration and Enhancement Network"201description = r"""<center><img src='https://user-images.githubusercontent.com/14334509/189166076-94bb2cac-4f4e-40fb-a69f-66709e3d98f5.png' alt='CodeFormer logo'></center>202<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>203🔥 CodeFormer is a robust face restoration algorithm for old photos or AI-generated faces.<br>204🤗 Try CodeFormer for improved stable-diffusion generation!<br>205"""206article = r"""207If CodeFormer is helpful, please help to ⭐ the <a href='https://github.com/sczhou/CodeFormer' target='_blank'>Github Repo</a>. Thanks! 208[![GitHub Stars](https://img.shields.io/github/stars/sczhou/CodeFormer?style=social)](https://github.com/sczhou/CodeFormer)209 210---211 212📝 **Citation**213 214If our work is useful for your research, please consider citing:215```bibtex216@inproceedings{zhou2022codeformer,217    author = {Zhou, Shangchen and Chan, Kelvin C.K. and Li, Chongyi and Loy, Chen Change},218    title = {Towards Robust Blind Face Restoration with Codebook Lookup TransFormer},219    booktitle = {NeurIPS},220    year = {2022}221}222```223 224📋 **License**225 226This project is licensed under <a rel="license" href="https://github.com/sczhou/CodeFormer/blob/master/LICENSE">S-Lab License 1.0</a>. 227Redistribution and use for non-commercial purposes should follow this license.228 229📧 **Contact**230 231If you have any questions, please feel free to reach me out at <b>shangchenzhou@gmail.com</b>.232 233![visitors](https://visitor-badge.glitch.me/badge?page_id=sczhou/CodeFormer)234"""235 236demo = gr.Interface(237    inference, [238        gr.inputs.Image(type="filepath", label="Input"),239        gr.inputs.Checkbox(default=True, label="Background_Enhance"),240        gr.inputs.Checkbox(default=True, label="Face_Upsample"),241        gr.inputs.Number(default=2, label="Rescaling_Factor"),242        gr.Slider(0, 1, value=0.5, step=0.01, label='Codeformer_Fidelity: 0 for better quality, 1 for better identity')243    ], [244        gr.outputs.Image(type="numpy", label="Output"),245        gr.outputs.File(label="Download the output")246    ],247    title=title,248    description=description,249    article=article,       250    examples=[251        ['01.png', True, True, 2, 0.7],252        ['02.jpg', True, True, 2, 0.7],253        ['03.jpg', True, True, 2, 0.7],254        ['04.jpg', True, True, 2, 0.1],255        ['05.jpg', True, True, 2, 0.1]256      ]257    ).launch()258 259demo.queue(concurrency_count=4)260demo.launch()