adding support for ESRGAN denoising strength, which allows for improved detail retention when upscaling photorelistic faces

.gitignore

@@ -1,4 +1,5 @@
 # ignore default image save location and model symbolic link
+.idea/
 embeddings/
 outputs/
 models/ldm/stable-diffusion-v1/model.ckpt
@@ -232,4 +233,4 @@ installer/update.bat
 installer/update.sh
 
 # no longer stored in source directory
-models
+models

ldm/invoke/CLI.py

@@ -1057,7 +1057,7 @@ def load_face_restoration(opt):
             else:
                 print('>> Face restoration disabled')
             if opt.esrgan:
-                esrgan = restoration.load_esrgan(opt.esrgan_bg_tile)
+                esrgan = restoration.load_esrgan(opt.esrgan_bg_tile, opt.esrgan_denoise_str)
             else:
                 print('>> Upscaling disabled')
         else:

ldm/invoke/args.py

@@ -671,6 +671,12 @@ class Args(object):
             default=400,
             help='Tile size for background sampler, 0 for no tile during testing. Default: 400.',
         )
+        postprocessing_group.add_argument(
+            '--esrgan_denoise_str',
+            type=float,
+            default=0.9,
+            help='esrgan denoise str. 0 is no denoise, 1 is max denoise.  Default: 0.9',
+        )
         postprocessing_group.add_argument(
             '--gfpgan_model_path',
             type=str,

ldm/invoke/config/invokeai_configure.py

@@ -128,7 +128,7 @@ script do it for you. Manual installation is described at:
 
 https://invoke-ai.github.io/InvokeAI/installation/020_INSTALL_MANUAL/
 
-You may download the recommended models (about 15GB total), install all models (40 GB!!) 
+You may download the recommended models (about 15GB total), install all models (40 GB!!)
 select a customized set, or completely skip this step.
 """
     )
@@ -583,7 +583,7 @@ def new_config_file_contents(successfully_downloaded: dict, config_file: Path, o
         # model is a diffusers (indicated with a path)
         if conf.get(model) and Path(successfully_downloaded[model]).is_dir():
             offer_to_delete_weights(model, conf[model], opt.yes_to_all)
-            
+
         stanza = {}
         mod = Datasets[model]
         stanza["description"] = mod["description"]
@@ -635,7 +635,7 @@ def offer_to_delete_weights(model_name: str, conf_stanza: dict, yes_to_all: bool
             weights.unlink()
         except OSError as e:
             print(str(e))
-    
+
 # ---------------------------------------------
 # this will preload the Bert tokenizer fles
 def download_bert():
@@ -683,10 +683,18 @@ def download_clip():
 def download_realesrgan():
     print("Installing models from RealESRGAN...", file=sys.stderr)
     model_url = "https://github.com/xinntao/Real-ESRGAN/releases/download/v0.2.5.0/realesr-general-x4v3.pth"
+    wdn_model_url = "https://github.com/xinntao/Real-ESRGAN/releases/download/v0.2.5.0/realesr-general-wdn-x4v3.pth"
+
     model_dest = os.path.join(
         Globals.root, "models/realesrgan/realesr-general-x4v3.pth"
     )
+
+    wdn_model_dest = os.path.join(
+        Globals.root, "models/realesrgan/realesr-general-wdn-x4v3.pth"
+    )
+
     download_with_progress_bar(model_url, model_dest, "RealESRGAN")
+    download_with_progress_bar(wdn_model_url, wdn_model_dest, "RealESRGANwdn")
 
 
 def download_gfpgan():

ldm/invoke/restoration/base.py

@@ -31,8 +31,8 @@ class Restoration():
         return CodeFormerRestoration()
 
     # Upscale Models
-    def load_esrgan(self, esrgan_bg_tile=400):
+    def load_esrgan(self, esrgan_bg_tile=400, denoise_str=0.9):
         from ldm.invoke.restoration.realesrgan import ESRGAN
-        esrgan = ESRGAN(esrgan_bg_tile)
+        esrgan = ESRGAN(esrgan_bg_tile, denoise_str)
         print('>> ESRGAN Initialized')
         return esrgan;

ldm/invoke/restoration/realesrgan.py

@@ -8,8 +8,9 @@ from PIL import Image
 from PIL.Image import Image as ImageType
 
 class ESRGAN():
-    def __init__(self, bg_tile_size=400) -> None:
+    def __init__(self, bg_tile_size=400, denoise_str=0.9) -> None:
         self.bg_tile_size = bg_tile_size
+        self.denoise_str=denoise_str
 
         if not torch.cuda.is_available():  # CPU or MPS on M1
             use_half_precision = False
@@ -26,14 +27,16 @@ class ESRGAN():
         from realesrgan import RealESRGANer
 
         model = SRVGGNetCompact(num_in_ch=3, num_out_ch=3, num_feat=64, num_conv=32, upscale=4, act_type='prelu')
-        model_path = os.path.join(Globals.root,'models/realesrgan/realesr-general-x4v3.pth')
+        model_path = os.path.join(Globals.root, 'models/realesrgan/realesr-general-x4v3.pth')
+        wdn_model_path = os.path.join(Globals.root, 'models/realesrgan/realesr-general-wdn-x4v3.pth')
         scale = 4
 
         bg_upsampler = RealESRGANer(
             scale=scale,
-            model_path=model_path,
+            model_path=[model_path, wdn_model_path],
             model=model,
             tile=self.bg_tile_size,
+            dni_weight=[self.denoise_str, 1 - self.denoise_str],
             tile_pad=10,
             pre_pad=0,
             half=use_half_precision,
@@ -60,7 +63,7 @@ class ESRGAN():
 
         if seed is not None:
             print(
-                f'>> Real-ESRGAN Upscaling seed:{seed} : scale:{upsampler_scale}x'
+                f'>> Real-ESRGAN Upscaling seed:{seed}, scale:{upsampler_scale}x, tile:{self.bg_tile_size}, denoise:{self.denoise_str}'
             )
         # ESRGAN outputs images with partial transparency if given RGBA images; convert to RGB
         image = image.convert("RGB")