Restoration Modules are now optional

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

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)