Mirrored_Repos/InvokeAI
1
0
Fork 0
mirror of https://github.com/invoke-ai/InvokeAI synced 2024-08-30 20:32:17 +00:00
Code Issues Packages Projects Releases Wiki Activity

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

Browse Source
This commit is contained in:
coreco 2023-02-08 22:36:35 -06:00
parent 4785a1cd05
commit 1f76b30e54
6 changed files with 29 additions and 11 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

1
.gitignore vendored
View File

@ -1,4 +1,5 @@
# ignore default image save location and model symbolic link
.idea/
embeddings/
outputs/
models/ldm/stable-diffusion-v1/model.ckpt

2
ldm/invoke/CLI.py
View File

@ -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:

6
ldm/invoke/args.py
View File

@ -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,

8
ldm/invoke/config/invokeai_configure.py
View File

@ -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():

4
ldm/invoke/restoration/base.py
View File

@ -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;

11
ldm/invoke/restoration/realesrgan.py
View File

@ -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")