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',