Implement CodeFormer Face Restoration (invoke-ai#669)

* Implement CodeFormer Face Restoration

* fix codeformer model destination path

Co-authored-by: Lincoln Stein <lincoln.stein@gmail.com>

Author: blessedcoolant

.gitignore

@@ -1,6 +1,7 @@
 # ignore default image save location and model symbolic link
 outputs/
 models/ldm/stable-diffusion-v1/model.ckpt
+ldm/restoration/codeformer/weights
 
 # ignore a directory which serves as a place for initial images
 inputs/

docs/features/UPSCALE.md

@@ -97,3 +97,39 @@ the base images.
 If you wish to stop during the image generation but want to upscale or face restore a particular
 generated image, pass it again with the same prompt and generated seed along with the `-U` and `-G`
 prompt arguments to perform those actions.
+
+## CodeFormer Support
+
+This repo also allows you to perform face restoration using
+[CodeFormer](https://github.com/sczhou/CodeFormer).
+
+In order to setup CodeFormer to work, you need to download the models like with GFPGAN. You can do
+this either by running `preload_models.py` or by manually downloading the
+[model file](https://github.com/sczhou/CodeFormer/releases/download/v0.1.0/codeformer.pth) and
+saving it to `ldm/restoration/codeformer/weights` folder.
+
+You can use `-ft` prompt argument to swap between CodeFormer and the default GFPGAN. The above
+mentioned `-G` prompt argument will allow you to control the strength of the restoration effect.
+
+### **Usage:**
+
+The following command will perform face restoration with CodeFormer instead of the default gfpgan.
+
+`<prompt> -G 0.8 -ft codeformer`
+
+**Other Options:**
+
+- `-cf` - cf or CodeFormer Fidelity takes values between `0` and `1`. 0 produces high quality
+  results but low accuracy and 1 produces lower quality results but higher accuacy to your original
+  face.
+
+The following command will perform face restoration with CodeFormer. CodeFormer will output a result
+that is closely matching to the input face.
+
+`<prompt> -G 1.0 -ft codeformer -cf 0.9`
+
+The following command will perform face restoration with CodeFormer. CodeFormer will output a result
+that is the best restoration possible. This may deviate slightly from the original face. This is an
+excellent option to use in situations when there is very little facial data to work with.
+
+`<prompt> -G 1.0 -ft codeformer -cf 0.1`

ldm/dream/args.py

@@ -516,13 +516,26 @@ def _create_dream_cmd_parser(self):
             help='Strength for noising/unnoising. 0.0 preserves image exactly, 1.0 replaces it completely',
             default=0.75,
         )
+        postprocessing_group.add_argument(
+            '-ft',
+            '--facetool',
+            type=str,
+            help='Select the face restoration AI to use: gfpgan, codeformer',
+        )
         postprocessing_group.add_argument(
             '-G',
             '--gfpgan_strength',
             type=float,
             help='The strength at which to apply the GFPGAN model to the result, in order to improve faces.',
             default=0,
         )
+        postprocessing_group.add_argument(
+            '-cf',
+            '--codeformer_fidelity',
+            type=float,
+            help='Takes values between 0 and 1. 0 produces high quality but low accuracy. 1 produces high accuracy but low quality.',
+            default=0.75
+        )
         postprocessing_group.add_argument(
             '-U',
             '--upscale',

ldm/generate.py

@@ -227,7 +227,9 @@ def prompt2image(
             embiggen       =    None,
             embiggen_tiles =    None,
             # these are specific to GFPGAN/ESRGAN
+            facetool         = None,
             gfpgan_strength  = 0,
+            codeformer_fidelity = None,
             save_original    = False,
             upscale          = None,
             # Set this True to handle KeyboardInterrupt internally
@@ -373,7 +375,9 @@ def process_image(image,seed):
             if upscale is not None or gfpgan_strength > 0:
                 self.upscale_and_reconstruct(results,
                                              upscale        = upscale,
+                                             facetool       = facetool,
                                              strength       = gfpgan_strength,
+                                             codeformer_fidelity = codeformer_fidelity,
                                              save_original  = save_original,
                                              image_callback = image_callback)
 
@@ -507,15 +511,20 @@ def correct_colors(self,
 
     def upscale_and_reconstruct(self,
                                 image_list,
+                                facetool      = 'gfpgan',
                                 upscale       = None,
                                 strength      =  0.0,
+                                codeformer_fidelity = 0.75,
                                 save_original = False,
                                 image_callback = None):
         try:
             if upscale is not None:
                 from ldm.gfpgan.gfpgan_tools import real_esrgan_upscale
             if strength > 0:
-                from ldm.gfpgan.gfpgan_tools import run_gfpgan
+                if facetool == 'codeformer':
+                    from ldm.restoration.codeformer.codeformer import CodeFormerRestoration
+                else:
+                    from ldm.gfpgan.gfpgan_tools import run_gfpgan
         except (ModuleNotFoundError, ImportError):
             print(traceback.format_exc(), file=sys.stderr)
             print('>> You may need to install the ESRGAN and/or GFPGAN modules')
@@ -534,9 +543,12 @@ def upscale_and_reconstruct(self,
                         seed,
                     )
                 if strength > 0:
-                    image = run_gfpgan(
-                        image, strength, seed, 1
-                    )
+                    if facetool == 'codeformer':
+                        image = CodeFormerRestoration().process(image=image, strength=strength, device=self.device, seed=seed, fidelity=codeformer_fidelity)
+                    else:
+                        image = run_gfpgan(
+                            image, strength, seed, 1
+                        )
             except Exception as e:
                 print(
                     f'>> Error running RealESRGAN or GFPGAN. Your image was not upscaled.\n{e}'

ldm/restoration/codeformer/codeformer.py

@@ -0,0 +1,76 @@
+import os
+import torch
+import numpy as np
+import warnings
+
+pretrained_model_url = 'https://github.com/sczhou/CodeFormer/releases/download/v0.1.0/codeformer.pth'
+
+class CodeFormerRestoration():
+    def __init__(self) -> None:
+        pass
+
+    def process(self, image, strength, device, seed=None, fidelity=0.75):
+        if seed is not None:
+            print(f'>> CodeFormer - Restoring Faces for image seed:{seed}')
+        with warnings.catch_warnings():
+            warnings.filterwarnings('ignore', category=DeprecationWarning)
+            warnings.filterwarnings('ignore', category=UserWarning)
+
+            from basicsr.utils.download_util import load_file_from_url
+            from basicsr.utils import img2tensor, tensor2img
+            from facexlib.utils.face_restoration_helper import FaceRestoreHelper
+            from ldm.restoration.codeformer.codeformer_arch import CodeFormer
+            from torchvision.transforms.functional import normalize
+            from PIL import Image
+            
+            cf_class = CodeFormer
+            
+            cf = cf_class(dim_embd=512, codebook_size=1024, n_head=8, n_layers=9, connect_list=['32', '64', '128', '256']).to(device)
+            
+            checkpoint_path = load_file_from_url(url=pretrained_model_url, model_dir=os.path.abspath('ldm/restoration/codeformer/weights'), progress=True)
+            checkpoint = torch.load(checkpoint_path)['params_ema']
+            cf.load_state_dict(checkpoint)
+            cf.eval()
+
+            image = image.convert('RGB')
+
+            face_helper = FaceRestoreHelper(upscale_factor=1, use_parse=True, device=device)
+            face_helper.clean_all()
+            face_helper.read_image(np.array(image, dtype=np.uint8))
+            face_helper.get_face_landmarks_5(resize=640, eye_dist_threshold=5)
+            face_helper.align_warp_face()
+
+            for idx, cropped_face in enumerate(face_helper.cropped_faces):
+                cropped_face_t = img2tensor(cropped_face / 255., 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 = cf(cropped_face_t, w=fidelity, adain=True)[0]
+                        restored_face = tensor2img(output.squeeze(0), rgb2bgr=True, min_max=(-1, 1))
+                    del output
+                    torch.cuda.empty_cache()
+                except RuntimeError as error:
+                    print(f'\tFailed inference for CodeFormer: {error}.')
+                    restored_face = cropped_face
+
+                restored_face = restored_face.astype('uint8')
+                face_helper.add_restored_face(restored_face)
+
+
+            face_helper.get_inverse_affine(None)
+
+            restored_img = face_helper.paste_faces_to_input_image()
+
+            res = Image.fromarray(restored_img)
+
+            if strength < 1.0:
+                # Resize the image to the new image if the sizes have changed
+                if restored_img.size != image.size:
+                    image = image.resize(res.size)
+                res = Image.blend(image, res, strength)
+
+            cf = None
+
+            return res