mirror of
https://github.com/invoke-ai/InvokeAI
synced 2024-08-30 20:32:17 +00:00
Add prompt as output directory feature
Based on previous code by czyz
This commit is contained in:
parent
fd13526454
commit
4f8e7356b3
@ -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
|
||||||
@ -90,13 +91,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:
|
||||||
@ -167,6 +177,23 @@ 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:
|
||||||
|
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:
|
else:
|
||||||
current_outdir = outdir
|
current_outdir = outdir
|
||||||
|
|
||||||
@ -320,6 +347,14 @@ 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',
|
||||||
|
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
|
# GFPGAN related args
|
||||||
parser.add_argument(
|
parser.add_argument(
|
||||||
'--gfpgan_bg_upsampler',
|
'--gfpgan_bg_upsampler',
|
||||||
|
Loading…
x
Reference in New Issue
Block a user