sczhou/CodeFormer
2.4k
1"""2This file is used for deploying hugging face demo:3https://huggingface.co/spaces/sczhou/CodeFormer4"""5import spaces6import sys7sys.path.append('CodeFormer')8import os9import cv210import numpy as np11import torch12import torch.nn.functional as F13import uuid, threading, time, glob14import gradio as gr15 16from torchvision.transforms.functional import normalize17 18from basicsr.utils import imwrite, img2tensor, tensor2img19from basicsr.utils.download_util import load_file_from_url20from facelib.utils.face_restoration_helper import FaceRestoreHelper21from basicsr.archs.rrdbnet_arch import RRDBNet22from basicsr.utils.realesrgan_utils import RealESRGANer23from facelib.utils.misc import is_gray24 25from basicsr.utils.registry import ARCH_REGISTRY26 27 28os.system("pip freeze")29 30pretrain_model_url = {31 'codeformer': 'https://github.com/sczhou/CodeFormer/releases/download/v0.1.0/codeformer.pth',32 'detection': 'https://github.com/sczhou/CodeFormer/releases/download/v0.1.0/detection_Resnet50_Final.pth',33 'parsing': 'https://github.com/sczhou/CodeFormer/releases/download/v0.1.0/parsing_parsenet.pth',34 'realesrgan': 'https://github.com/sczhou/CodeFormer/releases/download/v0.1.0/RealESRGAN_x2plus.pth'35}36# download weights37if not os.path.exists('CodeFormer/weights/CodeFormer/codeformer.pth'):38 load_file_from_url(url=pretrain_model_url['codeformer'], model_dir='CodeFormer/weights/CodeFormer', progress=True, file_name=None)39if not os.path.exists('CodeFormer/weights/facelib/detection_Resnet50_Final.pth'):40 load_file_from_url(url=pretrain_model_url['detection'], model_dir='CodeFormer/weights/facelib', progress=True, file_name=None)41if not os.path.exists('CodeFormer/weights/facelib/parsing_parsenet.pth'):42 load_file_from_url(url=pretrain_model_url['parsing'], model_dir='CodeFormer/weights/facelib', progress=True, file_name=None)43if not os.path.exists('CodeFormer/weights/realesrgan/RealESRGAN_x2plus.pth'):44 load_file_from_url(url=pretrain_model_url['realesrgan'], model_dir='CodeFormer/weights/realesrgan', progress=True, file_name=None)45 46# download images47torch.hub.download_url_to_file(48 'https://replicate.com/api/models/sczhou/codeformer/files/fa3fe3d1-76b0-4ca8-ac0d-0a925cb0ff54/06.png',49 '01.png')50torch.hub.download_url_to_file(51 'https://replicate.com/api/models/sczhou/codeformer/files/a1daba8e-af14-4b00-86a4-69cec9619b53/04.jpg',52 '02.jpg')53torch.hub.download_url_to_file(54 'https://replicate.com/api/models/sczhou/codeformer/files/542d64f9-1712-4de7-85f7-3863009a7c3d/03.jpg',55 '03.jpg')56torch.hub.download_url_to_file(57 'https://replicate.com/api/models/sczhou/codeformer/files/a11098b0-a18a-4c02-a19a-9a7045d68426/010.jpg',58 '04.jpg')59torch.hub.download_url_to_file(60 'https://replicate.com/api/models/sczhou/codeformer/files/7cf19c2c-e0cf-4712-9af8-cf5bdbb8d0ee/012.jpg',61 '05.jpg')62torch.hub.download_url_to_file(63 'https://raw.githubusercontent.com/sczhou/CodeFormer/master/inputs/cropped_faces/0729.png',64 '06.png')65 66 67def imread_unicode_safe(path):68 with open(path, "rb") as f:69 data = np.frombuffer(f.read(), dtype=np.uint8)70 return cv2.imdecode(data, cv2.IMREAD_COLOR)71 72def delayed_remove(path, delay=60):73 time.sleep(delay)74 try:75 if os.path.exists(path):76 os.remove(path)77 print(f"[CLEANUP] removed: {path}")78 else:79 print(f"[CLEANUP] already gone: {path}")80 except Exception as e:81 print(f"[CLEANUP] failed: {path} | {e}")82 83# set enhancer with RealESRGAN84def set_realesrgan():85 half = True if torch.cuda.is_available() else False86 model = RRDBNet(87 num_in_ch=3,88 num_out_ch=3,89 num_feat=64,90 num_block=23,91 num_grow_ch=32,92 scale=2,93 )94 upsampler = RealESRGANer(95 scale=2,96 model_path="CodeFormer/weights/realesrgan/RealESRGAN_x2plus.pth",97 model=model,98 tile=400,99 tile_pad=40,100 pre_pad=0,101 half=half,102 )103 return upsampler104 105upsampler = set_realesrgan()106device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')107codeformer_net = ARCH_REGISTRY.get("CodeFormer")(108 dim_embd=512,109 codebook_size=1024,110 n_head=8,111 n_layers=9,112 connect_list=["32", "64", "128", "256"],113).to(device)114ckpt_path = "CodeFormer/weights/CodeFormer/codeformer.pth"115checkpoint = torch.load(ckpt_path)["params_ema"]116codeformer_net.load_state_dict(checkpoint)117codeformer_net.eval()118 119os.makedirs('output', exist_ok=True)120 121@spaces.GPU122def inference(image, face_align, background_enhance, face_upsample, upscale, codeformer_fidelity):123 """Run a single prediction on the model"""124 try: # global try125 # take the default setting for the demo126 only_center_face = False127 draw_box = False128 detection_model = "retinaface_resnet50"129 130 face_align = face_align if face_align is not None else True131 background_enhance = background_enhance if background_enhance is not None else True132 face_upsample = face_upsample if face_upsample is not None else True133 upscale = upscale if (upscale is not None and upscale > 0) else 2134 135 has_aligned = not face_align136 upscale = 1 if has_aligned else upscale137 138 if isinstance(image, dict):139 image_path = image.get("name")140 elif isinstance(image, str):141 image_path = image142 else:143 image_path = None144 raise gr.Error("Invalid input image.")145 146 if not os.path.exists(image_path):147 raise gr.Error("Invalid input image.")148 149 print('Inp:', image_path, background_enhance, face_upsample, upscale, codeformer_fidelity)150 151 img = imread_unicode_safe(image_path)152 153 if img is None:154 raise gr.Error("Failed to read input image.")155 156 print('\timage size:', img.shape)157 158 upscale = int(upscale) # convert type to int159 if upscale > 4: # avoid memory exceeded due to too large upscale160 upscale = 4 161 if upscale > 2 and max(img.shape[:2])>1000: # avoid memory exceeded due to too large img resolution162 upscale = 2 163 if min(img.shape[:2]) > 1100 or max(img.shape[:2])>1500: # avoid memory exceeded due to too large img resolution164 upscale = 1165 background_enhance = False166 face_upsample = False167 168 h, w = img.shape[:2]169 if h * w > 4_000_000: # avoid memory exceeded due to too large img resolution170 raise gr.Error(171 "Image resolution is too large (>4 megapixels). "172 "To keep the demo responsive and avoid long queue times, this case is skipped. "173 "For such inputs, please deploy this demo locally and remove this limit."174 )175 176 face_helper = FaceRestoreHelper(177 upscale,178 face_size=512,179 crop_ratio=(1, 1),180 det_model=detection_model,181 save_ext="png",182 use_parse=True,183 device=device,184 )185 bg_upsampler = upsampler if background_enhance else None186 face_upsampler = upsampler if face_upsample else None187 188 if has_aligned:189 # the input faces are already cropped and aligned190 img = cv2.resize(img, (512, 512), interpolation=cv2.INTER_LINEAR)191 face_helper.is_gray = is_gray(img, threshold=5)192 if face_helper.is_gray:193 print('\tgrayscale input: True')194 face_helper.cropped_faces = [img]195 else:196 face_helper.read_image(img)197 # get face landmarks for each face198 num_det_faces = face_helper.get_face_landmarks_5(199 only_center_face=only_center_face, resize=640, eye_dist_threshold=5200 )201 print(f'\tdetect {num_det_faces} faces')202 # align and warp each face203 face_helper.align_warp_face()204 205 if min(img.shape[:2]) > 1000 and num_det_faces > 15:206 raise gr.Error(207 "Too many faces detected (>15) in a high-resolution image. "208 "To keep the demo responsive and avoid long queue times, this case is skipped. "209 "For such inputs, please deploy this demo locally and remove this limit."210 )211 212 213 # face restoration for each cropped face214 for idx, cropped_face in enumerate(face_helper.cropped_faces):215 # prepare data216 cropped_face_t = img2tensor(cropped_face / 255., bgr2rgb=True, float32=True)217 normalize(cropped_face_t, (0.5, 0.5, 0.5), (0.5, 0.5, 0.5), inplace=True)218 cropped_face_t = cropped_face_t.unsqueeze(0).to(device)219 220 try:221 with torch.no_grad():222 output = codeformer_net(223 cropped_face_t, w=codeformer_fidelity, adain=True224 )[0]225 restored_face = tensor2img(output, rgb2bgr=True, min_max=(-1, 1))226 del output227 torch.cuda.empty_cache()228 except RuntimeError as error:229 print(f"Failed inference for CodeFormer: {error}")230 restored_face = tensor2img(cropped_face_t, rgb2bgr=True, min_max=(-1, 1))231 232 restored_face = restored_face.astype("uint8")233 face_helper.add_restored_face(restored_face, cropped_face)234 235 # paste_back236 if not has_aligned:237 # upsample the background238 if bg_upsampler is not None:239 # Now only support RealESRGAN for upsampling background240 bg_img = bg_upsampler.enhance(img, outscale=upscale)[0]241 else:242 bg_img = None243 face_helper.get_inverse_affine(None)244 # paste each restored face to the input image245 if face_upsample and face_upsampler is not None:246 restored_img = face_helper.paste_faces_to_input_image(247 upsample_img=bg_img,248 draw_box=draw_box,249 face_upsampler=face_upsampler,250 )251 else:252 restored_img = face_helper.paste_faces_to_input_image(253 upsample_img=bg_img, draw_box=draw_box254 )255 else:256 restored_img = restored_face257 258 # save restored img259 # save_path = f'output/out.png'260 # imwrite(restored_img, str(save_path))261 262 # restored_img = cv2.cvtColor(restored_img, cv2.COLOR_BGR2RGB)263 # return restored_img264 265 266 #save restored img267 save_path = f"output/{uuid.uuid4().hex}.png"268 imwrite(restored_img, save_path)269 print(f"[SAVE] path={save_path} outputs={len(glob.glob('output/*.png'))}")270 271 threading.Thread(272 target=delayed_remove,273 args=(save_path,30),274 daemon=True275 ).start()276 277 return save_path, None278 279 except gr.Error:280 raise281 282 except Exception as error:283 print('[UNEXPECTED ERROR]', error)284 raise gr.Error("Unexpected error. Please try another image.")285 286 287title = "CodeFormer: Robust Face Restoration and Enhancement Network"288 289description = r"""<center><img src='https://user-images.githubusercontent.com/14334509/189166076-94bb2cac-4f4e-40fb-a69f-66709e3d98f5.png' alt='CodeFormer logo'></center>290<br>291<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>292🔥 CodeFormer is a robust face restoration algorithm for old photos or AI-generated faces.<br>293🤗 Try CodeFormer for improved stable-diffusion generation!<br>294"""295 296article = r"""297If CodeFormer is helpful, please help to ⭐ the <a href='https://github.com/sczhou/CodeFormer' target='_blank'>Github Repo</a>. Thanks! 298[](https://github.com/sczhou/CodeFormer)299 300---301 302📝 **Citation**303 304If our work is useful for your research, please consider citing:305```bibtex306@inproceedings{zhou2022codeformer,307 author = {Zhou, Shangchen and Chan, Kelvin C.K. and Li, Chongyi and Loy, Chen Change},308 title = {Towards Robust Blind Face Restoration with Codebook Lookup TransFormer},309 booktitle = {NeurIPS},310 year = {2022}311}312```313 314📋 **License**315 316This project is licensed under <a rel="license" href="https://github.com/sczhou/CodeFormer/blob/master/LICENSE">S-Lab License 1.0</a>. 317Redistribution and use for non-commercial purposes should follow this license.318 319📧 **Contact**320 321If you have any questions, please feel free to reach me out at <b>shangchenzhou@gmail.com</b>.322 323🤗 **Find Me:**324<style type="text/css">325td {326 padding-right: 0px !important;327}328 329.gradio-container-4-37-2 .prose table, .gradio-container-4-37-2 .prose tr, .gradio-container-4-37-2 .prose td, .gradio-container-4-37-2 .prose th {330 border: 0px solid #ffffff;331 border-bottom: 0px solid #ffffff;332}333 334</style>335 336<table>337<tr>338 <td><a href="https://github.com/sczhou"><img style="margin:-0.8em 0 2em 0" src="https://img.shields.io/github/followers/sczhou?style=social" alt="Github Follow"></a></td>339 <td><a href="https://twitter.com/ShangchenZhou"><img style="margin:-0.8em 0 2em 0" src="https://img.shields.io/twitter/follow/ShangchenZhou?label=%40ShangchenZhou&style=social" alt="Twitter Follow"></a></td>340</tr>341</table>342 343<center><img src='https://api.infinitescript.com/badgen/count?name=sczhou/CodeFormer<ext=Visitors&color=6dc9aa' alt='visitors'></center>344"""345 346with gr.Blocks() as demo:347 gr.Markdown(title)348 gr.Markdown(description)349 with gr.Row():350 with gr.Column():351 input_img = gr.Image(type="filepath", label="Input")352 face_align = gr.Checkbox(value=True, label="Pre_Face_Align")353 background_enhance = gr.Checkbox(value=True, label="Background_Enhance")354 face_enhance = gr.Checkbox(value=True, label="Face_Upsample")355 upscale_factor = gr.Number(value=2, label="Rescaling_Factor (up to 4)")356 codeformer_fidelity = gr.Slider(0, 1, value=0.5, step=0.01, label='Codeformer_Fidelity (0 for better quality, 1 for better identity)')357 submit = gr.Button('Enhance Image')358 with gr.Column():359 output_img = gr.Image(type="filepath", label="Output")360 note = gr.Markdown("**Please download the output within 30 seconds.**")361 362 inps = [input_img, face_align, background_enhance, face_enhance, upscale_factor, codeformer_fidelity]363 outs = [output_img, note]364 submit.click(fn=inference, inputs=inps, outputs=outs)365 366 ex = gr.Examples([367 ['01.png', True, True, True, 2, 0.7],368 ['02.jpg', True, True, True, 2, 0.7],369 ['03.jpg', True, True, True, 2, 0.7],370 ['04.jpg', True, True, True, 2, 0.1],371 ['05.jpg', True, True, True, 2, 0.1],372 ['06.png', False, True, True, 1, 0.5]373 ],374 inputs=inps,375 cache_examples=False)376 377 gr.Markdown(article)378 379DEBUG = os.getenv('DEBUG') == '1'380demo.queue(api_open=False, max_size=10, default_concurrency_limit=2)381demo.launch(debug=DEBUG)