diff --git a/.gitignore b/.gitignore index b4f9a2bb0f..9adb0be85a 100644 --- a/.gitignore +++ b/.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 \ No newline at end of file +models diff --git a/ldm/invoke/CLI.py b/ldm/invoke/CLI.py index fd61c7c8bf..e365ab2d5e 100644 --- a/ldm/invoke/CLI.py +++ b/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: diff --git a/ldm/invoke/args.py b/ldm/invoke/args.py index ace0786154..8e5c9ff18c 100644 --- a/ldm/invoke/args.py +++ b/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, diff --git a/ldm/invoke/config/invokeai_configure.py b/ldm/invoke/config/invokeai_configure.py index 3a53799a13..ef45a023d6 100755 --- a/ldm/invoke/config/invokeai_configure.py +++ b/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(): diff --git a/ldm/invoke/restoration/base.py b/ldm/invoke/restoration/base.py index 5b4bc483c2..ea0b02bd59 100644 --- a/ldm/invoke/restoration/base.py +++ b/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; diff --git a/ldm/invoke/restoration/realesrgan.py b/ldm/invoke/restoration/realesrgan.py index 132a5e55e2..def55589af 100644 --- a/ldm/invoke/restoration/realesrgan.py +++ b/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")