Fix grid image saving, log to outdir path, display sampler options once

This commit is contained in:
Benjamin Warner 2022-08-28 19:34:55 -05:00
parent 373a2d9c32
commit 529040708b
2 changed files with 53 additions and 59 deletions

View File

@ -37,7 +37,7 @@ class PngWriter:
if not upscaled: if not upscaled:
self.files_written.append([self.filepath, seed]) self.files_written.append([self.filepath, seed])
def unique_filename(self, seed, upscaled, previouspath=None): def unique_filename(self, seed, upscaled=False, previouspath=None):
revision = 1 revision = 1
if previouspath is None: if previouspath is None:
@ -60,7 +60,7 @@ class PngWriter:
basename = os.path.basename(previouspath) basename = os.path.basename(previouspath)
x = re.match('^(\d+)\..*\.png', basename) x = re.match('^(\d+)\..*\.png', basename)
if not x: if not x:
return self.unique_filename(seed, previouspath) return self.unique_filename(seed, upscaled, previouspath)
basecount = int(x.groups()[0]) basecount = int(x.groups()[0])
series = 0 series = 0

View File

@ -86,16 +86,14 @@ def main():
"\n* Initialization done! Awaiting your command (-h for help, 'q' to quit)" "\n* Initialization done! Awaiting your command (-h for help, 'q' to quit)"
) )
log_path = os.path.join(opt.outdir, 'dream_log.txt') cmd_parser = create_cmd_parser()
with open(log_path, 'a') as log: if opt.web:
cmd_parser = create_cmd_parser() dream_server_loop(t2i)
if opt.web: else:
dream_server_loop(t2i) main_loop(t2i, opt.outdir, cmd_parser, infile)
else:
main_loop(t2i, opt.outdir, cmd_parser, log_path, infile)
log.close()
def main_loop(t2i, outdir, parser, log_path, infile):
def main_loop(t2i, outdir, parser, infile):
"""prompt/read/execute loop""" """prompt/read/execute loop"""
done = False done = False
last_seeds = [] last_seeds = []
@ -202,12 +200,13 @@ def main_loop(t2i, outdir, parser, log_path, infile):
continue continue
print('Outputs:') print('Outputs:')
write_log_message(t2i, normalized_prompt, results, log_path) log_path = os.path.join(current_outdir, 'dream_log.txt')
write_log_message(normalized_prompt, results, log_path)
print('goodbye!') print('goodbye!')
def get_next_command(infile=None) -> 'command string': def get_next_command(infile=None) -> str: #command string
if infile is None: if infile is None:
command = input('dream> ') command = input('dream> ')
else: else:
@ -239,9 +238,9 @@ def dream_server_loop(t2i):
dream_server.server_close() dream_server.server_close()
### the t2i variable doesn't seem to be necessary here. maybe remove it?
def write_log_message(t2i, prompt, results, log_path): def write_log_message(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""" """logs the name of the output image, prompt, and prompt args to the terminal and log file"""
log_lines = [f'{r[0]}: {prompt} -S{r[1]}\n' for r in results] log_lines = [f'{r[0]}: {prompt} -S{r[1]}\n' for r in results]
print(*log_lines, sep='') print(*log_lines, sep='')
@ -249,6 +248,17 @@ def write_log_message(t2i, prompt, results, log_path):
file.writelines(log_lines) file.writelines(log_lines)
SAMPLER_CHOICES=[
'ddim',
'k_dpm_2_a',
'k_dpm_2',
'k_euler_a',
'k_euler',
'k_heun',
'k_lms',
'plms',
]
def create_argv_parser(): def create_argv_parser():
parser = argparse.ArgumentParser( parser = argparse.ArgumentParser(
description="Parse script's command line args" description="Parse script's command line args"
@ -259,52 +269,44 @@ def create_argv_parser():
'-l', '-l',
dest='laion400m', dest='laion400m',
action='store_true', action='store_true',
help='fallback to the latent diffusion (laion400m) weights and config', help='Fallback to the latent diffusion (laion400m) weights and config',
) )
parser.add_argument( parser.add_argument(
'--from_file', '--from_file',
dest='infile', dest='infile',
type=str, type=str,
help='if specified, load prompts from this file', help='If specified, load prompts from this file',
) )
parser.add_argument( parser.add_argument(
'-n', '-n',
'--iterations', '--iterations',
type=int, type=int,
default=1, default=1,
help='number of images to generate', help='Number of images to generate',
) )
parser.add_argument( parser.add_argument(
'-F', '-F',
'--full_precision', '--full_precision',
dest='full_precision', dest='full_precision',
action='store_true', action='store_true',
help='use slower full precision math for calculations', help='Use slower full precision math for calculations',
) )
parser.add_argument( parser.add_argument(
'-A', '-A',
'-m', '-m',
'--sampler', '--sampler',
dest='sampler_name', dest='sampler_name',
choices=[ choices=SAMPLER_CHOICES,
'ddim', metavar='SAMPLER_NAME',
'k_dpm_2_a',
'k_dpm_2',
'k_euler_a',
'k_euler',
'k_heun',
'k_lms',
'plms',
],
default='k_lms', default='k_lms',
help='which sampler to use (k_lms)', help=f'Set the initial sampler. Default: k_lms. Supported samplers: {", ".join(SAMPLER_CHOICES)}',
) )
parser.add_argument( parser.add_argument(
'--outdir', '--outdir',
'-o', '-o',
type=str, type=str,
default='outputs/img-samples', default='outputs/img-samples',
help='directory in which to place generated images and a log of prompts and seeds (outputs/img-samples', help='Directory to save generated images and a log of prompts and seeds. Default: outputs/img-samples',
) )
parser.add_argument( parser.add_argument(
'--embedding_path', '--embedding_path',
@ -316,7 +318,7 @@ def create_argv_parser():
'-d', '-d',
type=str, type=str,
default='cuda', default='cuda',
help='device to run stable diffusion on. defaults to cuda `torch.cuda.current_device()` if avalible', help='Device to run Stable Diffusion on. Defaults to cuda `torch.cuda.current_device()` if avalible',
) )
# GFPGAN related args # GFPGAN related args
parser.add_argument( parser.add_argument(
@ -335,19 +337,19 @@ def create_argv_parser():
'--gfpgan_model_path', '--gfpgan_model_path',
type=str, type=str,
default='experiments/pretrained_models/GFPGANv1.3.pth', default='experiments/pretrained_models/GFPGANv1.3.pth',
help='indicates the path to the GFPGAN model, relative to --gfpgan_dir.', help='Indicates the path to the GFPGAN model, relative to --gfpgan_dir.',
) )
parser.add_argument( parser.add_argument(
'--gfpgan_dir', '--gfpgan_dir',
type=str, type=str,
default='../GFPGAN', default='../GFPGAN',
help='indicates the directory containing the GFPGAN code.', help='Indicates the directory containing the GFPGAN code.',
) )
parser.add_argument( parser.add_argument(
'--web', '--web',
dest='web', dest='web',
action='store_true', action='store_true',
help='start in web server mode.', help='Start in web server mode.',
) )
return parser return parser
@ -357,39 +359,39 @@ def create_cmd_parser():
description='Example: dream> a fantastic alien landscape -W1024 -H960 -s100 -n12' description='Example: dream> a fantastic alien landscape -W1024 -H960 -s100 -n12'
) )
parser.add_argument('prompt') parser.add_argument('prompt')
parser.add_argument('-s', '--steps', type=int, help='number of steps') parser.add_argument('-s', '--steps', type=int, help='Number of steps')
parser.add_argument( parser.add_argument(
'-S', '-S',
'--seed', '--seed',
type=int, type=int,
help='image seed; a +ve integer, or use -1 for the previous seed, -2 for the one before that, etc', help='Image seed; a +ve integer, or use -1 for the previous seed, -2 for the one before that, etc',
) )
parser.add_argument( parser.add_argument(
'-n', '-n',
'--iterations', '--iterations',
type=int, type=int,
default=1, default=1,
help='number of samplings to perform (slower, but will provide seeds for individual images)', help='Number of samplings to perform (slower, but will provide seeds for individual images)',
) )
parser.add_argument( parser.add_argument(
'-b', '-b',
'--batch_size', '--batch_size',
type=int, type=int,
default=1, default=1,
help='number of images to produce per sampling (will not provide seeds for individual images!)', help='Number of images to produce per sampling (will not provide seeds for individual images!)',
) )
parser.add_argument( parser.add_argument(
'-W', '--width', type=int, help='image width, multiple of 64' '-W', '--width', type=int, help='Image width, multiple of 64'
) )
parser.add_argument( parser.add_argument(
'-H', '--height', type=int, help='image height, multiple of 64' '-H', '--height', type=int, help='Image height, multiple of 64'
) )
parser.add_argument( parser.add_argument(
'-C', '-C',
'--cfg_scale', '--cfg_scale',
default=7.5, default=7.5,
type=float, type=float,
help='prompt configuration scale', help='Prompt configuration scale',
) )
parser.add_argument( parser.add_argument(
'-g', '--grid', action='store_true', help='generate a grid' '-g', '--grid', action='store_true', help='generate a grid'
@ -399,26 +401,26 @@ def create_cmd_parser():
'-o', '-o',
type=str, type=str,
default=None, default=None,
help='directory in which to place generated images and a log of prompts and seeds (outputs/img-samples', help='Directory to save generated images and a log of prompts and seeds',
) )
parser.add_argument( parser.add_argument(
'-i', '-i',
'--individual', '--individual',
action='store_true', action='store_true',
help='generate individual files (default)', help='Generate individual files (default)',
) )
parser.add_argument( parser.add_argument(
'-I', '-I',
'--init_img', '--init_img',
type=str, type=str,
help='path to input image for img2img mode (supersedes width and height)', help='Path to input image for img2img mode (supersedes width and height)',
) )
parser.add_argument( parser.add_argument(
'-f', '-f',
'--strength', '--strength',
default=0.75, default=0.75,
type=float, type=float,
help='strength for noising/unnoising. 0.0 preserves image exactly, 1.0 replaces it completely', help='Strength for noising/unnoising. 0.0 preserves image exactly, 1.0 replaces it completely',
) )
parser.add_argument( parser.add_argument(
'-G', '-G',
@ -447,7 +449,7 @@ def create_cmd_parser():
'-x', '-x',
'--skip_normalize', '--skip_normalize',
action='store_true', action='store_true',
help='skip subprompt weight normalization', help='Skip subprompt weight normalization',
) )
parser.add_argument( parser.add_argument(
'-A', '-A',
@ -456,17 +458,9 @@ def create_cmd_parser():
dest='sampler_name', dest='sampler_name',
default=None, default=None,
type=str, type=str,
choices=[ choices=SAMPLER_CHOICES,
'ddim', metavar='SAMPLER_NAME',
'k_dpm_2_a', help=f'Switch to a different sampler. Supported samplers: {", ".join(SAMPLER_CHOICES)}',
'k_dpm_2',
'k_euler_a',
'k_euler',
'k_heun',
'k_lms',
'plms',
],
help='Change to another supported sampler using this command',
) )
return parser return parser