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

Merge branch 'development' into development

Browse Source
This commit is contained in:
Peter Baylies 2022-09-23 15:41:06 -04:00 committed by GitHub
commit 880142708d
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
4 changed files with 38 additions and 25 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(): class Restoration():
def __init__(self, gfpgan_dir='./src/gfpgan', gfpgan_model_path='experiments/pretrained_models/GFPGANv1.3.pth', esrgan_bg_tile=400) -> None: def __init__(self) -> None:
self.gfpgan_dir = gfpgan_dir pass
self.gfpgan_model_path = gfpgan_model_path
self.esrgan_bg_tile = esrgan_bg_tile
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 # Load GFPGAN
gfpgan = self.load_gfpgan() gfpgan = self.load_gfpgan(gfpgan_dir, gfpgan_model_path)
if gfpgan.gfpgan_model_exists: if gfpgan.gfpgan_model_exists:
print('>> GFPGAN Initialized') print('>> GFPGAN Initialized')
else:
print('>> GFPGAN Disabled')
gfpgan = None
# Load CodeFormer # Load CodeFormer
codeformer = self.load_codeformer() codeformer = self.load_codeformer()
if codeformer.codeformer_model_exists: if codeformer.codeformer_model_exists:
print('>> CodeFormer Initialized') print('>> CodeFormer Initialized')
else:
print('>> CodeFormer Disabled')
codeformer = None
return gfpgan, codeformer return gfpgan, codeformer
# Face Restore Models # Face Restore Models
def load_gfpgan(self): def load_gfpgan(self, gfpgan_dir, gfpgan_model_path):
from ldm.dream.restoration.gfpgan import GFPGAN 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): def load_codeformer(self):
from ldm.dream.restoration.codeformer import CodeFormerRestoration from ldm.dream.restoration.codeformer import CodeFormerRestoration
return CodeFormerRestoration() return CodeFormerRestoration()
# Upscale Models # Upscale Models
def load_esrgan(self): def load_esrgan(self, esrgan_bg_tile=400):
from ldm.dream.restoration.realesrgan import ESRGAN from ldm.dream.restoration.realesrgan import ESRGAN
esrgan = ESRGAN(self.esrgan_bg_tile) esrgan = ESRGAN(esrgan_bg_tile)
print('>> ESRGAN Initialized') print('>> ESRGAN Initialized')
return esrgan; 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) self.gfpgan_model_exists = os.path.isfile(self.model_path)
if not self.gfpgan_model_exists: if not self.gfpgan_model_exists:
raise Exception( print('## NOT FOUND: GFPGAN model not found at ' + self.model_path)
'GFPGAN model not found at path ' + self.model_path) return None
sys.path.append(os.path.abspath(gfpgan_dir)) sys.path.append(os.path.abspath(gfpgan_dir))
def model_exists(self): def model_exists(self):

14
ldm/generate.py
View File

@ -736,11 +736,17 @@ class Generate:
else: else:
print(">> ESRGAN is disabled. Image not upscaled.") print(">> ESRGAN is disabled. Image not upscaled.")
if strength > 0: if strength > 0:
if self.gfpgan is not None and self.codeformer is not None: if self.gfpgan is not None or self.codeformer is not None:
if facetool == 'codeformer': if self.gfpgan is None:
image = self.codeformer.process(image=image, strength=strength, device=self.device, seed=seed, fidelity=codeformer_fidelity) 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: else:
image = self.gfpgan.process(image, strength, seed) image = self.gfpgan.process(image, strength, seed)
else: else:
print(">> Face Restoration is disabled.") print(">> Face Restoration is disabled.")
except Exception as e: except Exception as e:

21
scripts/dream.py Executable file → Normal file
View File

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