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

Merge branch 'prompt_as_dir' of https://github.com/BlueAmulet/stable-diffusion into BlueAmulet-prompt_as_dir

Browse Source
This commit is contained in:
Lincoln Stein 2022-08-29 22:13:37 -04:00
commit d48299e56c
1 changed files with 32 additions and 2 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

34
scripts/dream.py
View File

@ -4,6 +4,7 @@
import argparse import argparse
import shlex import shlex
import os import os
import re
import sys import sys
import copy import copy
import warnings import warnings
@ -91,13 +92,22 @@ def main():
if opt.web: if opt.web:
dream_server_loop(t2i) dream_server_loop(t2i)
else: 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""" """prompt/read/execute loop"""
done = False done = False
last_seeds = [] 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: while not done:
try: try:
@ -170,6 +180,20 @@ def main_loop(t2i, outdir, parser, infile):
if not os.path.exists(opt.outdir): if not os.path.exists(opt.outdir):
os.makedirs(opt.outdir) os.makedirs(opt.outdir)
current_outdir = opt.outdir current_outdir = opt.outdir
elif prompt_as_dir:
# 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 ('Writing files to directory: "' + current_outdir + '"')
# make sure the output directory exists
if not os.path.exists(current_outdir):
os.makedirs(current_outdir)
else: else:
current_outdir = outdir current_outdir = outdir
@ -335,6 +359,12 @@ def create_argv_parser():
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',
) )
parser.add_argument(
'--prompt_as_dir',
'-p',
action='store_true',
help='Place images in subdirectories named after the prompt.',
)
# GFPGAN related args # GFPGAN related args
parser.add_argument( parser.add_argument(
'--gfpgan_bg_upsampler', '--gfpgan_bg_upsampler',