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

cosmetic fixup to how the outputs are reported

Browse Source
This commit is contained in:
Lincoln Stein
2022-08-28 17:06:33 -04:00
parent 5d67986997 e0de42bd03
commit ea2ee33be8
6 changed files with 479 additions and 346 deletions
Download Patch File Download Diff File Expand all files Collapse all files

147
scripts/dream.py
View File

@ -11,6 +11,7 @@ import ldm.dream.readline
from ldm.dream.pngwriter import PngWriter, PromptFormatter
from ldm.dream.server import DreamServer, ThreadingDreamServer
def main():
"""Initialize command-line parsers and the diffusion model"""
arg_parser = create_argv_parser()
@ -67,52 +68,22 @@ def main():
if opt.infile:
try:
if os.path.isfile(opt.infile):
infile = open(opt.infile,'r')
elif opt.infile=='-': # stdin
infile = open(opt.infile, 'r')
elif opt.infile == '-': # stdin
infile = sys.stdin
else:
raise FileNotFoundError(f'{opt.infile} not found.')
except (FileNotFoundError,IOError) as e:
except (FileNotFoundError, IOError) as e:
print(f'{e}. Aborting.')
sys.exit(-1)
# preload the model
t2i.load_model()
# load GFPGAN if requested
if opt.use_gfpgan:
print('\n* --gfpgan was specified, loading gfpgan...')
with warnings.catch_warnings():
warnings.filterwarnings('ignore', category=DeprecationWarning)
try:
model_path = os.path.join(
opt.gfpgan_dir, opt.gfpgan_model_path
)
if not os.path.isfile(model_path):
raise Exception(
'GFPGAN model not found at path ' + model_path
)
sys.path.append(os.path.abspath(opt.gfpgan_dir))
from gfpgan import GFPGANer
bg_upsampler = load_gfpgan_bg_upsampler(
opt.gfpgan_bg_upsampler, opt.gfpgan_bg_tile
)
t2i.gfpgan = GFPGANer(
model_path=model_path,
upscale=opt.gfpgan_upscale,
arch='clean',
channel_multiplier=2,
bg_upsampler=bg_upsampler,
)
except Exception:
import traceback
print('Error loading GFPGAN:', file=sys.stderr)
print(traceback.format_exc(), file=sys.stderr)
if not infile:
print(
"\n* Initialization done! Awaiting your command (-h for help, 'q' to quit)"
)
log_path = os.path.join(opt.outdir, 'dream_log.txt')
with open(log_path, 'a') as log:
@ -237,9 +208,10 @@ def main_loop(t2i, outdir, parser, log_path, infile):
print('goodbye!')
def get_next_command(infile=None) -> 'command string':
if infile is None:
command = input("dream> ")
command = input('dream> ')
else:
command = infile.readline()
if not command:
@ -269,78 +241,13 @@ def dream_server_loop(t2i):
dream_server.server_close()
def load_gfpgan_bg_upsampler(bg_upsampler, bg_tile=400):
import torch
if bg_upsampler == 'realesrgan':
if not torch.cuda.is_available(): # CPU
import warnings
warnings.warn(
'The unoptimized RealESRGAN is slow on CPU. We do not use it. '
'If you really want to use it, please modify the corresponding codes.'
)
bg_upsampler = None
else:
from basicsr.archs.rrdbnet_arch import RRDBNet
from realesrgan import RealESRGANer
model = RRDBNet(
num_in_ch=3,
num_out_ch=3,
num_feat=64,
num_block=23,
num_grow_ch=32,
scale=2,
)
bg_upsampler = RealESRGANer(
scale=2,
model_path='https://github.com/xinntao/Real-ESRGAN/releases/download/v0.2.1/RealESRGAN_x2plus.pth',
model=model,
tile=bg_tile,
tile_pad=10,
pre_pad=0,
half=True,
) # need to set False in CPU mode
else:
bg_upsampler = None
return bg_upsampler
# variant generation is going to be superseded by a generalized
# "prompt-morph" functionality
# def generate_variants(t2i,outdir,opt,previous_gens):
# variants = []
# print(f"Generating {opt.variants} variant(s)...")
# newopt = copy.deepcopy(opt)
# newopt.iterations = 1
# newopt.variants = None
# for r in previous_gens:
# newopt.init_img = r[0]
# prompt = PromptFormatter(t2i,newopt).normalize_prompt()
# print(f"] generating variant for {newopt.init_img}")
# for j in range(0,opt.variants):
# try:
# file_writer = PngWriter(outdir,prompt,newopt.batch_size)
# callback = file_writer.write_image
# t2i.prompt2image(image_callback=callback,**vars(newopt))
# results = file_writer.files_written
# variants.append([prompt,results])
# except AssertionError as e:
# print(e)
# continue
# print(f'{opt.variants} variants generated')
# return variants
### the t2i variable doesn't seem to be necessary here. maybe remove it?
def write_log_message(t2i, prompt, results, log_path):
"""logs the name of the output image, its prompt and seed to the terminal, log file, and a Dream text chunk in the PNG metadata"""
log_lines = [f"{r[0]}: {prompt} -S{r[1]}\n" for r in results]
print(*log_lines, sep="")
log_lines = [f'{r[0]}: {prompt} -S{r[1]}\n' for r in results]
print(*log_lines, sep='')
with open(log_path, "a") as file:
with open(log_path, 'a') as file:
file.writelines(log_lines)
@ -413,18 +320,6 @@ def create_argv_parser():
help='device to run stable diffusion on. defaults to cuda `torch.cuda.current_device()` if avalible',
)
# GFPGAN related args
parser.add_argument(
'--gfpgan',
dest='use_gfpgan',
action='store_true',
help='load gfpgan for use in the dreambot. Note: Enabling GFPGAN will require more GPU memory',
)
parser.add_argument(
'--gfpgan_upscale',
type=int,
default=2,
help='The final upsampling scale of the image. Default: 2. Only used if --gfpgan is specified',
)
parser.add_argument(
'--gfpgan_bg_upsampler',
type=str,
@ -529,10 +424,24 @@ def create_cmd_parser():
parser.add_argument(
'-G',
'--gfpgan_strength',
default=None,
default=0,
type=float,
help='The strength at which to apply the GFPGAN model to the result, in order to improve faces.',
)
parser.add_argument(
'-U',
'--upscale',
nargs=2,
default=None,
type=float,
help='Scale factor for Real-ESRGAN. Either use 2 or 4.',
)
parser.add_argument(
'-save_orig',
'--save_original',
action='store_true',
help='Save original. Use it when upscaling to save both versions.',
)
# variants is going to be superseded by a generalized "prompt-morph" function
# parser.add_argument('-v','--variants',type=int,help="in img2img mode, the first generated image will get passed back to img2img to generate the requested number of variants")
parser.add_argument(

20
scripts/preload_models.py
View File

@ -3,6 +3,9 @@
# Before running stable-diffusion on an internet-isolated machine,
# run this script from one with internet connectivity. The
# two machines must share a common .cache directory.
from transformers import CLIPTokenizer, CLIPTextModel
import clip
from transformers import BertTokenizerFast
import sys
import transformers
import os
@ -12,7 +15,6 @@ transformers.logging.set_verbosity_error()
# this will preload the Bert tokenizer fles
print('preloading bert tokenizer...')
from transformers import BertTokenizerFast
tokenizer = BertTokenizerFast.from_pretrained('bert-base-uncased')
print('...success')
@ -28,8 +30,6 @@ version = 'openai/clip-vit-large-patch14'
print('preloading CLIP model (Ignore the deprecation warnings)...')
sys.stdout.flush()
import clip
from transformers import CLIPTokenizer, CLIPTextModel
tokenizer = CLIPTokenizer.from_pretrained(version)
transformer = CLIPTextModel.from_pretrained(version)
@ -63,6 +63,20 @@ if gfpgan:
scale=2,
),
)
RealESRGANer(
scale=4,
model_path='https://github.com/xinntao/Real-ESRGAN/releases/download/v0.1.0/RealESRGAN_x4plus.pth',
model=RRDBNet(
num_in_ch=3,
num_out_ch=3,
num_feat=64,
num_block=23,
num_grow_ch=32,
scale=4,
),
)
FaceRestoreHelper(1, det_model='retinaface_resnet50')
print('...success')
except Exception: