From 4f8e7356b34fa5570ad5eef9f850d9d1235041f6 Mon Sep 17 00:00:00 2001 From: BlueAmulet <43395286+BlueAmulet@users.noreply.github.com> Date: Mon, 29 Aug 2022 10:54:40 -0600 Subject: [PATCH 1/3] Add prompt as output directory feature Based on previous code by czyz --- scripts/dream.py | 39 +++++++++++++++++++++++++++++++++++++-- 1 file changed, 37 insertions(+), 2 deletions(-) diff --git a/scripts/dream.py b/scripts/dream.py index 5c51c6e68b..8c3bd6553e 100755 --- a/scripts/dream.py +++ b/scripts/dream.py @@ -4,6 +4,7 @@ import argparse import shlex import os +import re import sys import copy import warnings @@ -90,13 +91,22 @@ def main(): if opt.web: dream_server_loop(t2i) else: - main_loop(t2i, opt.outdir, cmd_parser, infile) + main_loop(t2i, opt.outdir, opt.prompt_as_dir, cmd_parser, infile) -def main_loop(t2i, outdir, parser, infile): +def main_loop(t2i, outdir, prompt_as_dir, parser, infile): """prompt/read/execute loop""" done = False last_seeds = [] + path_filter = re.compile(r'[<>:"/\\|?*]') + + # os.pathconf is not available on Windows + if hasattr(os, 'pathconf'): + path_max = os.pathconf(outdir, 'PC_PATH_MAX') + name_max = os.pathconf(outdir, 'PC_NAME_MAX') + else: + path_max = 260 + name_max = 255 while not done: try: @@ -167,6 +177,23 @@ def main_loop(t2i, outdir, parser, infile): if not os.path.exists(opt.outdir): os.makedirs(opt.outdir) current_outdir = opt.outdir + elif prompt_as_dir: + if prompt_as_dir == 'simple': + # this line will sanitize the entire prompt as one folder + subdir = path_filter.sub('_', opt.prompt)[:name_max].rstrip(' .') + else: + # this line will allow the prompt to create subdirectories + subdir = os.path.join(*[path_filter.sub('_', x)[:name_max].rstrip(' .') for x in opt.prompt.replace('\\', '/').split('/') if x.rstrip(' .') != '']) + + # truncate path to maximum allowed length + subdir = subdir[:(path_max - 27 - len(os.path.abspath(outdir)))] + current_outdir = os.path.join(outdir, subdir) + + print ('\nWriting files to directory: "' + current_outdir + '"\n') + + # make sure the output directory exists + if not os.path.exists(current_outdir): + os.makedirs(current_outdir) else: current_outdir = outdir @@ -320,6 +347,14 @@ def create_argv_parser(): default='cuda', help='Device to run Stable Diffusion on. Defaults to cuda `torch.cuda.current_device()` if avalible', ) + parser.add_argument( + '--prompt_as_dir', + '-p', + type=str, + nargs='?', + const='y', + help='Output images are placed in subdirectories named using the prompt string. Only used if --prompt_as_dir or -p is specified.\nIf set to \"simple\" then all non-alphanumeric characters in the prompt will be replaced by underscores for the folder name.', + ) # GFPGAN related args parser.add_argument( '--gfpgan_bg_upsampler', From 2e22d9ecf17ffce878490b737244498023be4fc3 Mon Sep 17 00:00:00 2001 From: BlueAmulet <43395286+BlueAmulet@users.noreply.github.com> Date: Mon, 29 Aug 2022 18:10:15 -0600 Subject: [PATCH 2/3] Address bakkot review --- scripts/dream.py | 17 ++++++----------- 1 file changed, 6 insertions(+), 11 deletions(-) diff --git a/scripts/dream.py b/scripts/dream.py index 8c3bd6553e..962d2e9fa9 100755 --- a/scripts/dream.py +++ b/scripts/dream.py @@ -178,18 +178,15 @@ def main_loop(t2i, outdir, prompt_as_dir, parser, infile): os.makedirs(opt.outdir) current_outdir = opt.outdir elif prompt_as_dir: - if prompt_as_dir == 'simple': - # this line will sanitize the entire prompt as one folder - subdir = path_filter.sub('_', opt.prompt)[:name_max].rstrip(' .') - else: - # this line will allow the prompt to create subdirectories - subdir = os.path.join(*[path_filter.sub('_', x)[:name_max].rstrip(' .') for x in opt.prompt.replace('\\', '/').split('/') if x.rstrip(' .') != '']) + # sanitize the prompt to a valid folder name + subdir = path_filter.sub('_', opt.prompt)[:name_max].rstrip(' .') # truncate path to maximum allowed length + # 27 is the length of '######.##########.##.png', plus two separators and a NUL subdir = subdir[:(path_max - 27 - len(os.path.abspath(outdir)))] current_outdir = os.path.join(outdir, subdir) - print ('\nWriting files to directory: "' + current_outdir + '"\n') + print ('Writing files to directory: "' + current_outdir + '"') # make sure the output directory exists if not os.path.exists(current_outdir): @@ -350,10 +347,8 @@ def create_argv_parser(): parser.add_argument( '--prompt_as_dir', '-p', - type=str, - nargs='?', - const='y', - help='Output images are placed in subdirectories named using the prompt string. Only used if --prompt_as_dir or -p is specified.\nIf set to \"simple\" then all non-alphanumeric characters in the prompt will be replaced by underscores for the folder name.', + action='store_true', + help='Place images in subdirectories named after the prompt.', ) # GFPGAN related args parser.add_argument( From 42ffcd7204d761a3752f79b53f97a07181ee071c Mon Sep 17 00:00:00 2001 From: Lincoln Stein Date: Mon, 29 Aug 2022 22:34:09 -0400 Subject: [PATCH 3/3] add the recently added commands to the readline command-line-completion list; fix command-line documentation bug, closing issue #188 --- ldm/dream/readline.py | 47 +++++++++++++++++++------------------------ scripts/dream.py | 2 +- 2 files changed, 22 insertions(+), 27 deletions(-) diff --git a/ldm/dream/readline.py b/ldm/dream/readline.py index f40fe83316..5cf99523fc 100644 --- a/ldm/dream/readline.py +++ b/ldm/dream/readline.py @@ -86,32 +86,27 @@ if readline_available: readline.set_completer( Completer( [ - 'cd', - 'pwd', - '--steps', - '-s', - '--seed', - '-S', - '--iterations', - '-n', - '--batch_size', - '-b', - '--width', - '-W', - '--height', - '-H', - '--cfg_scale', - '-C', - '--grid', - '-g', - '--individual', - '-i', - '--init_img', - '-I', - '--strength', - '-f', - '-v', - '--variants', + '--steps','-s', + '--seed','-S', + '--iterations','-n', + '--batch_size','-b', + '--width','-W','--height','-H', + '--cfg_scale','-C', + '--grid','-g', + '--individual','-i', + '--init_img','-I', + '--strength','-f', + '--variants','-v', + '--outdir','-o', + '--sampler','-A','-m', + '--embedding_path', + '--device', + '--grid','-g', + '--gfpgan_strength','-G', + '--upscale','-U', + '-save_orig','--save_original', + '--skip_normalize','-x', + '--log_tokenization','t', ] ).complete ) diff --git a/scripts/dream.py b/scripts/dream.py index ae79b76c17..0f4dbdae06 100755 --- a/scripts/dream.py +++ b/scripts/dream.py @@ -437,7 +437,7 @@ def create_cmd_parser(): '--cfg_scale', default=7.5, type=float, - help='Prompt configuration scale', + help='Classifier free guidance (CFG) scale - higher numbers cause generator to "try" harder.', ) parser.add_argument( '-g', '--grid', action='store_true', help='generate a grid'