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

Restoration Modules are now optional

Browse Source
This commit is contained in:
blessedcoolant 2022-09-23 19:40:31 +12:00 committed by Lincoln Stein
parent 413064cf45
commit cd8be1d0e9
4 changed files with 37 additions and 26 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

24
ldm/dream/restoration/base.py
View File

@ -1,34 +1,38 @@
class Restoration():
def __init__(self, gfpgan_dir='./src/gfpgan', gfpgan_model_path='experiments/pretrained_models/GFPGANv1.3.pth', esrgan_bg_tile=400) -> None:
self.gfpgan_dir = gfpgan_dir
self.gfpgan_model_path = gfpgan_model_path
self.esrgan_bg_tile = esrgan_bg_tile
def __init__(self) -> None:
pass
def load_face_restore_models(self):
def load_face_restore_models(self, gfpgan_dir='./src/gfpgan', gfpgan_model_path='experiments/pretrained_models/GFPGANv1.3.pth'):
# Load GFPGAN
gfpgan = self.load_gfpgan()
gfpgan = self.load_gfpgan(gfpgan_dir, gfpgan_model_path)
if gfpgan.gfpgan_model_exists:
print('>> GFPGAN Initialized')
else:
print('>> GFPGAN Disabled')
gfpgan = None
# Load CodeFormer
codeformer = self.load_codeformer()
if codeformer.codeformer_model_exists:
print('>> CodeFormer Initialized')
else:
print('>> CodeFormer Disabled')
codeformer = None
return gfpgan, codeformer
# Face Restore Models
def load_gfpgan(self):
def load_gfpgan(self, gfpgan_dir, gfpgan_model_path):
from ldm.dream.restoration.gfpgan import GFPGAN
return GFPGAN(self.gfpgan_dir, self.gfpgan_model_path)
return GFPGAN(gfpgan_dir, gfpgan_model_path)
def load_codeformer(self):
from ldm.dream.restoration.codeformer import CodeFormerRestoration
return CodeFormerRestoration()
# Upscale Models
def load_esrgan(self):
def load_esrgan(self, esrgan_bg_tile=400):
from ldm.dream.restoration.realesrgan import ESRGAN
esrgan = ESRGAN(self.esrgan_bg_tile)
esrgan = ESRGAN(esrgan_bg_tile)
print('>> ESRGAN Initialized')
return esrgan;

4
ldm/dream/restoration/gfpgan.py
View File

@ -17,8 +17,8 @@ class GFPGAN():
self.gfpgan_model_exists = os.path.isfile(self.model_path)
if not self.gfpgan_model_exists:
raise Exception(
'GFPGAN model not found at path ' + self.model_path)
print('## NOT FOUND: GFPGAN model not found at ' + self.model_path)
return None
sys.path.append(os.path.abspath(gfpgan_dir))
def model_exists(self):

14
ldm/generate.py
View File

@ -726,11 +726,17 @@ class Generate:
else:
print(">> ESRGAN is disabled. Image not upscaled.")
if strength > 0:
if self.gfpgan is not None and self.codeformer is not None:
if facetool == 'codeformer':
image = self.codeformer.process(image=image, strength=strength, device=self.device, seed=seed, fidelity=codeformer_fidelity)
if self.gfpgan is not None or self.codeformer is not None:
if self.gfpgan is None:
if facetool == 'codeformer':
if self.codeformer is not None:
image = self.codeformer.process(image=image, strength=strength, device=self.device, seed=seed, fidelity=codeformer_fidelity)
else:
print('>> CodeFormer not found. Face restoration is disabled.')
else:
print('>> GFPGAN not found. Face restoration is disabled.')
else:
image = self.gfpgan.process(image, strength, seed)
image = self.gfpgan.process(image, strength, seed)
else:
print(">> Face Restoration is disabled.")
except Exception as e:

21
scripts/dream.py
View File

@ -47,16 +47,17 @@ def main():
# Loading Face Restoration and ESRGAN Modules
try:
gfpgan, codeformer, esrgan = None, None, None
from ldm.dream.restoration import Restoration
restoration = Restoration(opt.gfpgan_dir, opt.gfpgan_model_path, opt.esrgan_bg_tile)
if opt.restore:
gfpgan, codeformer = restoration.load_face_restore_models()
else:
print('>> Face restoration disabled')
if opt.esrgan:
esrgan = restoration.load_esrgan()
else:
print('>> Upscaling disabled')
if opt.restore or opt.esrgan:
from ldm.dream.restoration import Restoration
restoration = Restoration()
if opt.restore:
gfpgan, codeformer = restoration.load_face_restore_models(opt.gfpgan_dir, opt.gfpgan_model_path)
else:
print('>> Face restoration disabled')
if opt.esrgan:
esrgan = restoration.load_esrgan(opt.esrgan_bg_tile)
else:
print('>> Upscaling disabled')
except (ModuleNotFoundError, ImportError):
import traceback
print(traceback.format_exc(), file=sys.stderr)