add hugging_face demo.

README.md

@@ -20,7 +20,8 @@ S-Lab, Nanyang Technological University
 
 ### Update
 
-- **2022.09.09**: Integrated to [Replicate](https://replicate.com/). Try out online demo! [![Replicate](https://img.shields.io/badge/Demo-%F0%9F%9A%80%20Replicate-blue)](https://replicate.com/sczhou/codeformer)
+- **2022.09.14**: Integrated to :hugs: [Hugging Face](https://replicate.com/). Try out online demo! [![Hugging Face](https://img.shields.io/badge/Demo-%F0%9F%A4%97%20Hugging%20Face-blue)](https://replicate.com/sczhou/codeformer)
+- **2022.09.09**: Integrated to :rocket: [Replicate](https://replicate.com/). Try out online demo! [![Replicate](https://img.shields.io/badge/Demo-%F0%9F%9A%80%20Replicate-blue)](https://replicate.com/sczhou/codeformer)
 - **2022.09.04**: Add face upsampling `--face_upsample` for high-resolution AI-created face enhancement.
 - **2022.08.23**: Some modifications on face detection and fusion for better AI-created face enhancement.
 - **2022.08.07**: Integrate [Real-ESRGAN](https://github.com/xinntao/Real-ESRGAN) to support background image enhancement.

app.py

@@ -0,0 +1,244 @@
+"""
+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 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"""<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</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!
+[![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 <b>[email protected]</b>.
+
+![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()

cog.yaml

@@ -1,3 +1,8 @@
+"""
+This file is used for deploying replicate demo:
+https://replicate.com/sczhou/codeformer
+"""
+
 build:
   gpu: true
   cuda: "11.3"

facelib/detection/__init__.py

@@ -25,10 +25,10 @@ def init_detection_model(model_name, half=False, device='cuda'):
 def init_retinaface_model(model_name, half=False, device='cuda'):
     if model_name == 'retinaface_resnet50':
         model = RetinaFace(network_name='resnet50', half=half)
-        model_url = 'https://github.com/xinntao/facexlib/releases/download/v0.1.0/detection_Resnet50_Final.pth'
+        model_url = 'https://github.com/sczhou/CodeFormer/releases/download/v0.1.0/detection_Resnet50_Final.pth'
     elif model_name == 'retinaface_mobile0.25':
         model = RetinaFace(network_name='mobile0.25', half=half)
-        model_url = 'https://github.com/xinntao/facexlib/releases/download/v0.1.0/detection_mobilenet0.25_Final.pth'
+        model_url = 'https://github.com/sczhou/CodeFormer/releases/download/v0.1.0/detection_mobilenet0.25_Final.pth'
     else:
         raise NotImplementedError(f'{model_name} is not implemented.')
 

facelib/utils/face_restoration_helper.py

@@ -59,7 +59,7 @@ class FaceRestoreHelper(object):
                  use_parse=False,
                  device=None):
         self.template_3points = template_3points  # improve robustness
-        self.upscale_factor = upscale_factor
+        self.upscale_factor = int(upscale_factor)
         # the cropped face ratio based on the square face
         self.crop_ratio = crop_ratio  # (h, w)
         assert (self.crop_ratio[0] >= 1 and self.crop_ratio[1] >= 1), 'crop ration only supports >=1'

predict.py

@@ -1,15 +1,18 @@
 """
-download checkpoints to ./weights beforehand 
-python scripts/download_pretrained_models.py facelib
-python scripts/download_pretrained_models.py CodeFormer
-wget 'https://github.com/xinntao/Real-ESRGAN/releases/download/v0.2.1/RealESRGAN_x2plus.pth'
+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
-from cog import BasePredictor, Input, Path
+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
@@ -22,7 +25,7 @@ class Predictor(BasePredictor):
     def setup(self):
         """Load the model into memory to make running multiple predictions efficient"""
         self.device = "cuda:0"
-        self.bg_upsampler = set_realesrgan()
+        self.upsampler = set_realesrgan()
         self.net = ARCH_REGISTRY.get("CodeFormer")(
             dim_embd=512,
             codebook_size=1024,
@@ -76,8 +79,8 @@ class Predictor(BasePredictor):
             device=self.device,
         )
 
-        bg_upsampler = self.bg_upsampler if background_enhance else None
-        face_upsampler = self.bg_upsampler if face_upsample else None
+        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)
 
@@ -143,10 +146,8 @@ class Predictor(BasePredictor):
                 )
 
         # save restored img
-        out_path = Path(tempfile.mkdtemp()) / "output.png"
-
-        if not has_aligned and restored_img is not None:
-            imwrite(restored_img, str(out_path))
+        out_path = Path(tempfile.mkdtemp()) / 'output.png'
+        imwrite(restored_img, str(out_path))
 
         return out_path
 
@@ -166,7 +167,7 @@ def set_realesrgan():
             "If you really want to use it, please modify the corresponding codes.",
             category=RuntimeWarning,
         )
-        bg_upsampler = None
+        upsampler = None
     else:
         model = RRDBNet(
             num_in_ch=3,
@@ -176,13 +177,13 @@ def set_realesrgan():
             num_grow_ch=32,
             scale=2,
         )
-        bg_upsampler = RealESRGANer(
+        upsampler = RealESRGANer(
             scale=2,
-            model_path="./weights/RealESRGAN_x2plus.pth",
+            model_path="./weights/realesrgan/RealESRGAN_x2plus.pth",
             model=model,
             tile=400,
             tile_pad=40,
             pre_pad=0,
             half=True,
         )
-    return bg_upsampler
+    return upsampler