stack86/CodeFormer
0
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')60torch.hub.download_url_to_file(61 'https://raw.githubusercontent.com/sczhou/CodeFormer/master/inputs/cropped_faces/0729.png',62 '06.png')63 64def imread(img_path):65 img = cv2.imread(img_path)66 img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)67 return img68 69# set enhancer with RealESRGAN70def set_realesrgan():71 half = True if torch.cuda.is_available() else False72 model = RRDBNet(73 num_in_ch=3,74 num_out_ch=3,75 num_feat=64,76 num_block=23,77 num_grow_ch=32,78 scale=2,79 )80 upsampler = RealESRGANer(81 scale=2,82 model_path="CodeFormer/weights/realesrgan/RealESRGAN_x2plus.pth",83 model=model,84 tile=400,85 tile_pad=40,86 pre_pad=0,87 half=half,88 )89 return upsampler90 91upsampler = set_realesrgan()92device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')93codeformer_net = ARCH_REGISTRY.get("CodeFormer")(94 dim_embd=512,95 codebook_size=1024,96 n_head=8,97 n_layers=9,98 connect_list=["32", "64", "128", "256"],99).to(device)100ckpt_path = "CodeFormer/weights/CodeFormer/codeformer.pth"101checkpoint = torch.load(ckpt_path)["params_ema"]102codeformer_net.load_state_dict(checkpoint)103codeformer_net.eval()104 105os.makedirs('output', exist_ok=True)106 107def inference(image, face_align, background_enhance, face_upsample, upscale, codeformer_fidelity):108 """Run a single prediction on the model"""109 try: # global try110 # take the default setting for the demo111 only_center_face = False112 draw_box = False113 detection_model = "retinaface_resnet50"114 115 print('Inp:', image, background_enhance, face_upsample, upscale, codeformer_fidelity)116 face_align = face_align if face_align is not None else True117 background_enhance = background_enhance if background_enhance is not None else True118 face_upsample = face_upsample if face_upsample is not None else True119 upscale = upscale if (upscale is not None and upscale > 0) else 2120 121 has_aligned = not face_align122 upscale = 1 if has_aligned else upscale123 124 img = cv2.imread(str(image), cv2.IMREAD_COLOR)125 print('\timage size:', img.shape)126 127 upscale = int(upscale) # convert type to int128 if upscale > 4: # avoid memory exceeded due to too large upscale129 upscale = 4 130 if upscale > 2 and max(img.shape[:2])>1000: # avoid memory exceeded due to too large img resolution131 upscale = 2 132 if max(img.shape[:2]) > 1500: # avoid memory exceeded due to too large img resolution133 upscale = 1134 background_enhance = False135 face_upsample = False136 137 face_helper = FaceRestoreHelper(138 upscale,139 face_size=512,140 crop_ratio=(1, 1),141 det_model=detection_model,142 save_ext="png",143 use_parse=True,144 device=device,145 )146 bg_upsampler = upsampler if background_enhance else None147 face_upsampler = upsampler if face_upsample else None148 149 if has_aligned:150 # the input faces are already cropped and aligned151 img = cv2.resize(img, (512, 512), interpolation=cv2.INTER_LINEAR)152 face_helper.is_gray = is_gray(img, threshold=5)153 if face_helper.is_gray:154 print('\tgrayscale input: True')155 face_helper.cropped_faces = [img]156 else:157 face_helper.read_image(img)158 # get face landmarks for each face159 num_det_faces = face_helper.get_face_landmarks_5(160 only_center_face=only_center_face, resize=640, eye_dist_threshold=5161 )162 print(f'\tdetect {num_det_faces} faces')163 # align and warp each face164 face_helper.align_warp_face()165 166 # face restoration for each cropped face167 for idx, cropped_face in enumerate(face_helper.cropped_faces):168 # prepare data169 cropped_face_t = img2tensor(170 cropped_face / 255.0, bgr2rgb=True, float32=True171 )172 normalize(cropped_face_t, (0.5, 0.5, 0.5), (0.5, 0.5, 0.5), inplace=True)173 cropped_face_t = cropped_face_t.unsqueeze(0).to(device)174 175 try:176 with torch.no_grad():177 output = codeformer_net(178 cropped_face_t, w=codeformer_fidelity, adain=True179 )[0]180 restored_face = tensor2img(output, rgb2bgr=True, min_max=(-1, 1))181 del output182 torch.cuda.empty_cache()183 except RuntimeError as error:184 print(f"Failed inference for CodeFormer: {error}")185 restored_face = tensor2img(186 cropped_face_t, rgb2bgr=True, min_max=(-1, 1)187 )188 189 restored_face = restored_face.astype("uint8")190 face_helper.add_restored_face(restored_face)191 192 # paste_back193 if not has_aligned:194 # upsample the background195 if bg_upsampler is not None:196 # Now only support RealESRGAN for upsampling background197 bg_img = bg_upsampler.enhance(img, outscale=upscale)[0]198 else:199 bg_img = None200 face_helper.get_inverse_affine(None)201 # paste each restored face to the input image202 if face_upsample and face_upsampler is not None:203 restored_img = face_helper.paste_faces_to_input_image(204 upsample_img=bg_img,205 draw_box=draw_box,206 face_upsampler=face_upsampler,207 )208 else:209 restored_img = face_helper.paste_faces_to_input_image(210 upsample_img=bg_img, draw_box=draw_box211 )212 else:213 restored_img = restored_face214 215 # save restored img216 save_path = f'output/out.png'217 imwrite(restored_img, str(save_path))218 219 restored_img = cv2.cvtColor(restored_img, cv2.COLOR_BGR2RGB)220 return restored_img221 except Exception as error:222 print('Global exception', error)223 return None, None224 225 226title = "人脸恢复与增强"227 228description = r"""229"""230 231article = r"""232"""233 234demo = gr.Interface(235 inference, [236 gr.Image(type="filepath", label="Input"),237 gr.Checkbox(value=True, label="Pre_Face_Align"),238 gr.Checkbox(value=True, label="Background_Enhance"),239 gr.Checkbox(value=True, label="Face_Upsample"),240 gr.Number(value=2, label="Rescaling_Factor (up to 4)"),241 gr.Slider(0, 1, value=0.5, step=0.01, label='Codeformer_Fidelity (0 for better quality, 1 for better identity)')242 ], [243 gr.Image(type="numpy", label="Output").style(height='auto')244 ],245 title=title,246 description=description,247 article=article, 248 examples=[249 ['01.png', True, True, True, 2, 0.7],250 ['02.jpg', True, True, True, 2, 0.7],251 ['03.jpg', True, True, True, 2, 0.7],252 ['04.jpg', True, True, True, 2, 0.1],253 ['05.jpg', True, True, True, 2, 0.1],254 ['06.png', False, True, True, 1, 0.5]255 ])256 257DEBUG = os.getenv('DEBUG') == '1'258demo.queue(api_open=False, concurrency_count=2, max_size=10)259demo.launch(debug=DEBUG)260# demo.launch(debug=DEBUG, share=True)