add an option to skip adding number to filenames when saving.

rework filename pattern function go through the pattern once and not calculate any of replacements until they are actually encountered in the pattern.
This commit is contained in:
AUTOMATIC 2022-10-24 14:03:58 +03:00
parent eb007e5884
commit 8da1bd48bf
2 changed files with 128 additions and 116 deletions

View File

@ -1,4 +1,7 @@
import datetime import datetime
import sys
import traceback
import pytz import pytz
import io import io
import math import math
@ -274,10 +277,15 @@ invalid_filename_chars = '<>:"/\\|?*\n'
invalid_filename_prefix = ' ' invalid_filename_prefix = ' '
invalid_filename_postfix = ' .' invalid_filename_postfix = ' .'
re_nonletters = re.compile(r'[\s' + string.punctuation + ']+') re_nonletters = re.compile(r'[\s' + string.punctuation + ']+')
re_pattern = re.compile(r"([^\[\]]+|\[([^]]+)]|[\[\]]*)")
re_pattern_arg = re.compile(r"(.*)<([^>]*)>$")
max_filename_part_length = 128 max_filename_part_length = 128
def sanitize_filename_part(text, replace_spaces=True): def sanitize_filename_part(text, replace_spaces=True):
if text is None:
return None
if replace_spaces: if replace_spaces:
text = text.replace(' ', '_') text = text.replace(' ', '_')
@ -287,49 +295,103 @@ def sanitize_filename_part(text, replace_spaces=True):
return text return text
def apply_filename_pattern(x, p, seed, prompt): class FilenameGenerator:
max_prompt_words = opts.directories_max_prompt_words replacements = {
'seed': lambda self: self.seed if self.seed is not None else '',
'steps': lambda self: self.p and self.p.steps,
'cfg': lambda self: self.p and self.p.cfg_scale,
'width': lambda self: self.p and self.p.width,
'height': lambda self: self.p and self.p.height,
'styles': lambda self: self.p and sanitize_filename_part(", ".join([style for style in self.p.styles if not style == "None"]) or "None", replace_spaces=False),
'sampler': lambda self: self.p and sanitize_filename_part(sd_samplers.samplers[self.p.sampler_index].name, replace_spaces=False),
'model_hash': lambda self: getattr(self.p, "sd_model_hash", shared.sd_model.sd_model_hash),
'date': lambda self: datetime.datetime.now().strftime('%Y-%m-%d'),
'datetime': lambda self, *args: self.datetime(*args), # accepts formats: [datetime], [datetime<Format>], [datetime<Format><Time Zone>]
'job_timestamp': lambda self: getattr(self.p, "job_timestamp", shared.state.job_timestamp),
'prompt': lambda self: sanitize_filename_part(self.prompt),
'prompt_no_styles': lambda self: self.prompt_no_style(),
'prompt_spaces': lambda self: sanitize_filename_part(self.prompt, replace_spaces=False),
'prompt_words': lambda self: self.prompt_words(),
}
default_time_format = '%Y%m%d%H%M%S'
if seed is not None: def __init__(self, p, seed, prompt):
x = re.sub(r'\[seed]', str(seed), x, flags=re.IGNORECASE) self.p = p
self.seed = seed
self.prompt = prompt
if p is not None: def prompt_no_style(self):
x = re.sub(r'\[steps]', str(p.steps), x, flags=re.IGNORECASE) if self.p is None or self.prompt is None:
x = re.sub(r'\[cfg]', str(p.cfg_scale), x, flags=re.IGNORECASE) return None
x = re.sub(r'\[width]', str(p.width), x, flags=re.IGNORECASE)
x = re.sub(r'\[height]', str(p.height), x, flags=re.IGNORECASE)
x = re.sub(r'\[styles]', sanitize_filename_part(", ".join([x for x in p.styles if not x == "None"]) or "None", replace_spaces=False), x, flags=re.IGNORECASE)
x = re.sub(r'\[sampler]', sanitize_filename_part(sd_samplers.samplers[p.sampler_index].name, replace_spaces=False), x, flags=re.IGNORECASE)
x = re.sub(r'\[model_hash]', getattr(p, "sd_model_hash", shared.sd_model.sd_model_hash), x, flags=re.IGNORECASE) prompt_no_style = self.prompt
current_time = datetime.datetime.now() for style in shared.prompt_styles.get_style_prompts(self.p.styles):
x = re.sub(r'\[date]', current_time.strftime('%Y-%m-%d'), x, flags=re.IGNORECASE)
x = replace_datetime(x, current_time) # replace [datetime], [datetime<Format>], [datetime<Format><Time Zone>]
x = re.sub(r'\[job_timestamp]', getattr(p, "job_timestamp", shared.state.job_timestamp), x, flags=re.IGNORECASE)
# Apply [prompt] at last. Because it may contain any replacement word.^M
if prompt is not None:
x = re.sub(r'\[prompt]', sanitize_filename_part(prompt), x, flags=re.IGNORECASE)
if re.search(r'\[prompt_no_styles]', x, re.IGNORECASE):
prompt_no_style = prompt
for style in shared.prompt_styles.get_style_prompts(p.styles):
if len(style) > 0: if len(style) > 0:
style_parts = [y for y in style.split("{prompt}")] for part in style.split("{prompt}"):
for part in style_parts:
prompt_no_style = prompt_no_style.replace(part, "").replace(", ,", ",").strip().strip(',') prompt_no_style = prompt_no_style.replace(part, "").replace(", ,", ",").strip().strip(',')
prompt_no_style = prompt_no_style.replace(style, "").strip().strip(',').strip()
x = re.sub(r'\[prompt_no_styles]', sanitize_filename_part(prompt_no_style, replace_spaces=False), x, flags=re.IGNORECASE)
x = re.sub(r'\[prompt_spaces]', sanitize_filename_part(prompt, replace_spaces=False), x, flags=re.IGNORECASE) prompt_no_style = prompt_no_style.replace(style, "").strip().strip(',').strip()
if re.search(r'\[prompt_words]', x, re.IGNORECASE):
words = [x for x in re_nonletters.split(prompt or "") if len(x) > 0] return sanitize_filename_part(prompt_no_style, replace_spaces=False)
def prompt_words(self):
words = [x for x in re_nonletters.split(self.prompt or "") if len(x) > 0]
if len(words) == 0: if len(words) == 0:
words = ["empty"] words = ["empty"]
x = re.sub(r'\[prompt_words]', sanitize_filename_part(" ".join(words[0:max_prompt_words]), replace_spaces=False), x, flags=re.IGNORECASE) return sanitize_filename_part(" ".join(words[0:opts.directories_max_prompt_words]), replace_spaces=False)
if cmd_opts.hide_ui_dir_config: def datetime(self, *args):
x = re.sub(r'^[\\/]+|\.{2,}[\\/]+|[\\/]+\.{2,}', '', x) time_datetime = datetime.datetime.now()
return x time_format = args[0] if len(args) > 0 else self.default_time_format
time_zone = pytz.timezone(args[1]) if len(args) > 1 else None
time_zone_time = time_datetime.astimezone(time_zone)
try:
formatted_time = time_zone_time.strftime(time_format)
except (ValueError, TypeError) as _:
formatted_time = time_zone_time.strftime(self.default_time_format)
return sanitize_filename_part(formatted_time, replace_spaces=False)
def apply(self, x):
res = ''
for m in re_pattern.finditer(x):
text, pattern = m.groups()
if pattern is None:
res += text
continue
pattern_args = []
while True:
m = re_pattern_arg.match(pattern)
if m is None:
break
pattern, arg = m.groups()
pattern_args.insert(0, arg)
fun = self.replacements.get(pattern.lower())
if fun is not None:
try:
replacement = fun(self, *pattern_args)
except Exception:
replacement = None
print(f"Error adding [{pattern}] to filename", file=sys.stderr)
print(traceback.format_exc(), file=sys.stderr)
if replacement is None:
res += f'[{pattern}]'
else:
res += str(replacement)
continue
res += f'[{pattern}]'
return res
def get_next_sequence_number(path, basename): def get_next_sequence_number(path, basename):
@ -354,66 +416,8 @@ def get_next_sequence_number(path, basename):
return result + 1 return result + 1
def replace_datetime(input_str: str, time_datetime: datetime.datetime = None):
"""
Args:
input_str (`str`):
the String to be Formatted
time_datetime (`datetime.datetime`)
the time to be used, if None, use datetime.datetime.now()
Formats sub_string of input_str with formatted datetime with time zone support.
accepts sub_string format: [datetime], [datetime<Format>], [datetime<Format><Time Zone>]
case insensitive
e.g.
input: "___[Datetime<%Y_%m_%d %H-%M-%S><Asia/Tokyo>]___"
return: "___2022_10_22 20-40-14___"
handles invalid Formats and Time Zones
time format reference:
https://docs.python.org/3/library/datetime.html#strftime-and-strptime-format-codes
valid time zones
print(pytz.all_timezones)
https://pytz.sourceforge.net/
"""
default_time_format = '%Y%m%d%H%M%S'
if time_datetime is None:
time_datetime = datetime.datetime.now()
# match all datetime to be replace
match_itr = re.finditer(r'\[datetime(?:<([^>]*)>(?:<([^>]*)>)?)?]', input_str, re.IGNORECASE)
for match in reversed(list(match_itr)):
# extract format
time_format = match.group(1)
if time_format == '':
# if time_format is blank use default YYYYMMDDHHMMSS
time_format = default_time_format
# extract timezone
try:
time_zone = pytz.timezone(match.group(2))
except pytz.exceptions.UnknownTimeZoneError as _:
# if no time_zone or invalid, use system time
time_zone = None
# generate time string
time_zone_time = time_datetime.astimezone(time_zone)
try:
formatted_time = time_zone_time.strftime(time_format)
except (ValueError, TypeError) as _:
# if format error then use default_time_format
formatted_time = time_zone_time.strftime(default_time_format)
formatted_time = sanitize_filename_part(formatted_time, replace_spaces=False)
input_str = input_str[:match.start()] + formatted_time + input_str[match.end():]
return input_str
def save_image(image, path, basename, seed=None, prompt=None, extension='png', info=None, short_filename=False, no_prompt=False, grid=False, pnginfo_section_name='parameters', p=None, existing_info=None, forced_filename=None, suffix="", save_to_dirs=None): def save_image(image, path, basename, seed=None, prompt=None, extension='png', info=None, short_filename=False, no_prompt=False, grid=False, pnginfo_section_name='parameters', p=None, existing_info=None, forced_filename=None, suffix="", save_to_dirs=None):
'''Save an image. """Save an image.
Args: Args:
image (`PIL.Image`): image (`PIL.Image`):
@ -444,7 +448,9 @@ def save_image(image, path, basename, seed=None, prompt=None, extension='png', i
The full path of the saved imaged. The full path of the saved imaged.
txt_fullfn (`str` or None): txt_fullfn (`str` or None):
If a text file is saved for this image, this will be its full path. Otherwise None. If a text file is saved for this image, this will be its full path. Otherwise None.
''' """
namegen = FilenameGenerator(p, seed, prompt)
if extension == 'png' and opts.enable_pnginfo and info is not None: if extension == 'png' and opts.enable_pnginfo and info is not None:
pnginfo = PngImagePlugin.PngInfo() pnginfo = PngImagePlugin.PngInfo()
@ -460,24 +466,25 @@ def save_image(image, path, basename, seed=None, prompt=None, extension='png', i
save_to_dirs = (grid and opts.grid_save_to_dirs) or (not grid and opts.save_to_dirs and not no_prompt) save_to_dirs = (grid and opts.grid_save_to_dirs) or (not grid and opts.save_to_dirs and not no_prompt)
if save_to_dirs: if save_to_dirs:
dirname = apply_filename_pattern(opts.directories_filename_pattern or "[prompt_words]", p, seed, prompt).strip('\\ /') dirname = namegen.apply(opts.directories_filename_pattern or "[prompt_words]").lstrip(' ').rstrip('\\ /')
path = os.path.join(path, dirname) path = os.path.join(path, dirname)
os.makedirs(path, exist_ok=True) os.makedirs(path, exist_ok=True)
if forced_filename is None: if forced_filename is None:
if short_filename or prompt is None or seed is None: if short_filename or seed is None:
file_decoration = "" file_decoration = ""
elif opts.save_to_dirs:
file_decoration = opts.samples_filename_pattern or "[seed]"
else: else:
file_decoration = opts.samples_filename_pattern or "[seed]-[prompt_spaces]" file_decoration = opts.samples_filename_pattern or "[seed]"
if file_decoration != "": add_number = opts.save_images_add_number or file_decoration == ''
if file_decoration != "" and add_number:
file_decoration = "-" + file_decoration file_decoration = "-" + file_decoration
file_decoration = apply_filename_pattern(file_decoration, p, seed, prompt) + suffix file_decoration = namegen.apply(file_decoration) + suffix
if add_number:
basecount = get_next_sequence_number(path, basename) basecount = get_next_sequence_number(path, basename)
fullfn = None fullfn = None
fullfn_without_extension = None fullfn_without_extension = None
@ -487,6 +494,9 @@ def save_image(image, path, basename, seed=None, prompt=None, extension='png', i
fullfn_without_extension = os.path.join(path, f"{fn}{file_decoration}") fullfn_without_extension = os.path.join(path, f"{fn}{file_decoration}")
if not os.path.exists(fullfn): if not os.path.exists(fullfn):
break break
else:
fullfn = os.path.join(path, f"{file_decoration}.{extension}")
fullfn_without_extension = os.path.join(path, file_decoration)
else: else:
fullfn = os.path.join(path, f"{forced_filename}.{extension}") fullfn = os.path.join(path, f"{forced_filename}.{extension}")
fullfn_without_extension = os.path.join(path, forced_filename) fullfn_without_extension = os.path.join(path, forced_filename)

View File

@ -86,6 +86,7 @@ parser.add_argument("--device-id", type=str, help="Select the default CUDA devic
cmd_opts = parser.parse_args() cmd_opts = parser.parse_args()
restricted_opts = [ restricted_opts = [
"samples_filename_pattern", "samples_filename_pattern",
"directories_filename_pattern",
"outdir_samples", "outdir_samples",
"outdir_txt2img_samples", "outdir_txt2img_samples",
"outdir_img2img_samples", "outdir_img2img_samples",
@ -190,7 +191,8 @@ options_templates = {}
options_templates.update(options_section(('saving-images', "Saving images/grids"), { options_templates.update(options_section(('saving-images', "Saving images/grids"), {
"samples_save": OptionInfo(True, "Always save all generated images"), "samples_save": OptionInfo(True, "Always save all generated images"),
"samples_format": OptionInfo('png', 'File format for images'), "samples_format": OptionInfo('png', 'File format for images'),
"samples_filename_pattern": OptionInfo("", "Images filename pattern"), "samples_filename_pattern": OptionInfo("", "Images filename pattern", component_args=hide_dirs),
"save_images_add_number": OptionInfo(True, "Add number to filename when saving", component_args=hide_dirs),
"grid_save": OptionInfo(True, "Always save all generated image grids"), "grid_save": OptionInfo(True, "Always save all generated image grids"),
"grid_format": OptionInfo('png', 'File format for grids'), "grid_format": OptionInfo('png', 'File format for grids'),
@ -225,8 +227,8 @@ options_templates.update(options_section(('saving-to-dirs', "Saving to a directo
"save_to_dirs": OptionInfo(False, "Save images to a subdirectory"), "save_to_dirs": OptionInfo(False, "Save images to a subdirectory"),
"grid_save_to_dirs": OptionInfo(False, "Save grids to a subdirectory"), "grid_save_to_dirs": OptionInfo(False, "Save grids to a subdirectory"),
"use_save_to_dirs_for_ui": OptionInfo(False, "When using \"Save\" button, save images to a subdirectory"), "use_save_to_dirs_for_ui": OptionInfo(False, "When using \"Save\" button, save images to a subdirectory"),
"directories_filename_pattern": OptionInfo("", "Directory name pattern"), "directories_filename_pattern": OptionInfo("", "Directory name pattern", component_args=hide_dirs),
"directories_max_prompt_words": OptionInfo(8, "Max prompt words for [prompt_words] pattern", gr.Slider, {"minimum": 1, "maximum": 20, "step": 1}), "directories_max_prompt_words": OptionInfo(8, "Max prompt words for [prompt_words] pattern", gr.Slider, {"minimum": 1, "maximum": 20, "step": 1, **hide_dirs}),
})) }))
options_templates.update(options_section(('upscaling', "Upscaling"), { options_templates.update(options_section(('upscaling', "Upscaling"), {