Optimize GFPGAN and Improve User Experience

TODO.txt

@@ -1,6 +1,7 @@
 Feature requests:
 
-1. "gobig" mode - split image into strips, scale up, add detail using
+
+1. "gobig" mode - split image into strips, scale up, add detail using  - DONE!
    img2img and reassemble with feathering. Issue #66.
    See https://github.com/jquesnelle/txt2imghd
 
@@ -15,18 +16,20 @@ Feature requests:
 6. Support for loading variations of the stable-diffusion
    weights #49
 
-7. Support for klms and other non-ddim samplers in img2img() #36
+7. Support for klms and other non-ddim samplers in img2img() #36 - DONE!
 
 8. Pass a shell command to open up an image viewer on the last
    batch of images generated #29.
 
+9. Change sampler and outdir after initialization #115
+
 Code Refactorization:
 
-1. Move the PNG file generation code out of simplet2i and into
+1. Move the PNG file generation code out of simplet2i and into - DONE!
    separate module. txt2img() and img2img() should return Image
    objects, and parent code is responsible for filenaming logic.
 
-2. Refactor redundant code that is shared between txt2img() and
+2. Refactor redundant code that is shared between txt2img() and - DONE!
    img2img().
 
-3. Experiment with replacing CompViz code with HuggingFace.
+3. Experiment with replacing CompViz code with HuggingFace. - NOT WORTH IT!

ldm/dream/pngwriter.py

@@ -23,9 +23,9 @@ def __init__(self, outdir, prompt=None, batch_size=1):
         self.files_written = []
         os.makedirs(outdir, exist_ok=True)
 
-    def write_image(self, image, seed):
+    def write_image(self, image, seed, upscaled=False):
         self.filepath = self.unique_filename(
-            seed, self.filepath
+            seed, upscaled, self.filepath
         )   # will increment name in some sensible way
         try:
             prompt = f'{self.prompt} -S{seed}'
@@ -34,7 +34,7 @@ def write_image(self, image, seed):
             print(e)
         self.files_written.append([self.filepath, seed])
 
-    def unique_filename(self, seed, previouspath=None):
+    def unique_filename(self, seed, upscaled, previouspath=None):
         revision = 1
 
         if previouspath is None:
@@ -68,6 +68,8 @@ def unique_filename(self, seed, previouspath=None):
                 if self.batch_size > 1 or os.path.exists(
                     os.path.join(self.outdir, filename)
                 ):
+                    if upscaled:
+                      break
                     filename = f'{basecount:06}.{seed}.{series:02}.png'
                 finished = not os.path.exists(
                     os.path.join(self.outdir, filename)

ldm/gfpgan/gfpgan_tools.py

@@ -0,0 +1,132 @@
+import torch
+import warnings
+import os
+import sys
+import numpy as np
+
+from PIL import Image
+from scripts.dream import create_argv_parser
+
+arg_parser = create_argv_parser()
+opt = arg_parser.parse_args()
+
+
+def _run_gfpgan(image, strength, prompt, seed, upsampler_scale=4):
+    print(
+        f"\n* GFPGAN - Restoring Faces: {prompt} : seed:{seed}")
+    with warnings.catch_warnings():
+        warnings.filterwarnings("ignore", category=DeprecationWarning)
+        warnings.filterwarnings("ignore", category=UserWarning)
+
+        try:
+            model_path = os.path.join(
+                opt.gfpgan_dir, opt.gfpgan_model_path)
+            if not os.path.isfile(model_path):
+                raise Exception(
+                    "GFPGAN model not found at path "+model_path)
+
+            sys.path.append(os.path.abspath(opt.gfpgan_dir))
+            from gfpgan import GFPGANer
+
+            bg_upsampler = _load_gfpgan_bg_upsampler(
+                opt.gfpgan_bg_upsampler, upsampler_scale, opt.gfpgan_bg_tile)
+
+            gfpgan = GFPGANer(model_path=model_path, upscale=upsampler_scale,
+                              arch='clean', channel_multiplier=2, bg_upsampler=bg_upsampler)
+        except Exception:
+            import traceback
+            print("Error loading GFPGAN:", file=sys.stderr)
+            print(traceback.format_exc(), file=sys.stderr)
+
+    if (gfpgan is None):
+        print(f"GFPGAN not initialized, it must be loaded via the --gfpgan argument")
+        return image
+
+    image = image.convert("RGB")
+
+    cropped_faces, restored_faces, restored_img = gfpgan.enhance(np.array(
+        image, dtype=np.uint8), has_aligned=False, only_center_face=False, paste_back=True)
+    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)
+
+    if torch.cuda.is_available():
+        torch.cuda.empty_cache()
+    gfpgan = None
+
+    return res
+
+
+def _load_gfpgan_bg_upsampler(bg_upsampler, upsampler_scale, bg_tile=400):
+    if bg_upsampler == 'realesrgan':
+        if not torch.cuda.is_available():  # CPU
+            warnings.warn('The unoptimized RealESRGAN is slow on CPU. We do not use it. '
+                          'If you really want to use it, please modify the corresponding codes.')
+            bg_upsampler = None
+        else:
+            model_path = {2: 'https://github.com/xinntao/Real-ESRGAN/releases/download/v0.2.1/RealESRGAN_x2plus.pth',
+                          4: 'https://github.com/xinntao/Real-ESRGAN/releases/download/v0.1.0/RealESRGAN_x4plus.pth'}
+
+            if upsampler_scale not in model_path:
+                return None
+
+            from basicsr.archs.rrdbnet_arch import RRDBNet
+            from realesrgan import RealESRGANer
+
+            if (upsampler_scale == 4):
+                model = RRDBNet(num_in_ch=3, num_out_ch=3, num_feat=64,
+                                num_block=23, num_grow_ch=32, scale=4)
+            if (upsampler_scale == 2):
+                model = RRDBNet(num_in_ch=3, num_out_ch=3, num_feat=64,
+                                num_block=23, num_grow_ch=32, scale=2)
+
+            bg_upsampler = RealESRGANer(
+                scale=upsampler_scale,
+                model_path=model_path[upsampler_scale],
+                model=model,
+                tile=bg_tile,
+                tile_pad=10,
+                pre_pad=0,
+                half=True)  # need to set False in CPU mode
+    else:
+        bg_upsampler = None
+
+    return bg_upsampler
+
+
+def real_esrgan_upscale(image, strength, upsampler_scale, prompt, seed):
+    print(
+        f"\n* Real-ESRGAN Upscaling: {prompt} : seed:{seed} : scale:{upsampler_scale}x")
+
+    with warnings.catch_warnings():
+        warnings.filterwarnings("ignore", category=DeprecationWarning)
+        warnings.filterwarnings("ignore", category=UserWarning)
+
+        try:
+            upsampler = _load_gfpgan_bg_upsampler(
+                opt.gfpgan_bg_upsampler, upsampler_scale, opt.gfpgan_bg_tile)
+        except Exception:
+            import traceback
+            print("Error loading Real-ESRGAN:", file=sys.stderr)
+            print(traceback.format_exc(), file=sys.stderr)
+
+    output, img_mode = upsampler.enhance(np.array(
+        image, dtype=np.uint8), outscale=upsampler_scale, alpha_upsampler=opt.gfpgan_bg_upsampler)
+
+    res = Image.fromarray(output)
+
+    if strength < 1.0:
+        # Resize the image to the new image if the sizes have changed
+        if output.size != image.size:
+            image = image.resize(res.size)
+        res = Image.blend(image, res, strength)
+
+    if torch.cuda.is_available():
+        torch.cuda.empty_cache()
+    upsampler = None
+
+    return res