|
|
|
@@ -26,19 +26,18 @@ original_locale = locale.setlocale(locale.LC_TIME, '')
|
|
|
|
|
|
|
|
|
|
# class SaveImageExtended -------------------------------------------------------------------------------
|
|
|
|
|
class SaveImageExtended:
|
|
|
|
|
#Version: 2.43
|
|
|
|
|
version = 2.43
|
|
|
|
|
version = 2.44
|
|
|
|
|
type = 'output'
|
|
|
|
|
|
|
|
|
|
png_compress_level = 9
|
|
|
|
|
avif_quality = 100
|
|
|
|
|
avif_quality = 91
|
|
|
|
|
webp_quality = 91
|
|
|
|
|
jpeg_quality = 91
|
|
|
|
|
|
|
|
|
|
filename_prefix = 'myFile'
|
|
|
|
|
filename_prefix = 'ComfyUI'
|
|
|
|
|
filename_keys = 'sampler_name, scheduler, cfg, steps'
|
|
|
|
|
foldername_prefix = 'myFolder'
|
|
|
|
|
foldername_keys = 'ckpt_name, ./exampleSubfolder'
|
|
|
|
|
foldername_prefix = ''
|
|
|
|
|
foldername_keys = 'ckpt_name, ./subfolder'
|
|
|
|
|
delimiter = '_'
|
|
|
|
|
save_job_data = 'basic, models, sampler, prompt'
|
|
|
|
|
job_data_per_image = False
|
|
|
|
@@ -49,13 +48,13 @@ class SaveImageExtended:
|
|
|
|
|
counter_positions = ['last', 'first']
|
|
|
|
|
one_counter_per_folder = True
|
|
|
|
|
image_preview = True
|
|
|
|
|
extToRemove = ['.safetensors', '.ckpt', '.pt']
|
|
|
|
|
extToRemove = ['.safetensors', '.ckpt', '.pt', '.bin', '.pth']
|
|
|
|
|
|
|
|
|
|
print(f"\033[92m[save_image_extended]\033[0m version: {version}\033[0m")
|
|
|
|
|
if pillow_avif not in sys.modules:
|
|
|
|
|
output_ext = '.avif'
|
|
|
|
|
output_exts = ['.avif', '.png', '.webp', '.jpg']
|
|
|
|
|
print(f"\033[92m[save_image_extended] AVIF is supported! Woohoo!\033[0m")
|
|
|
|
|
print(f"\033[92m[save_image_extended] AVIF is supported! Woohoo!\033[0m\n")
|
|
|
|
|
else:
|
|
|
|
|
output_ext = '.png'
|
|
|
|
|
output_exts = ['.png', '.webp', '.jpg']
|
|
|
|
@@ -147,9 +146,9 @@ class SaveImageExtended:
|
|
|
|
|
if counter_position not in self.counter_positions: counter_position = self.counter_position
|
|
|
|
|
if counter_position == 'last':
|
|
|
|
|
# BUG: this works only if extension is 3 letters like png, this will break with webp and avif:
|
|
|
|
|
counters = [int(file[-(extLen + counter_digits):-extLen]) if file[-(extLen + counter_digits):-extLen].isdigit() else 0 for file in files if one_counter_per_folder or file.startswith(filename_prefix)]
|
|
|
|
|
counters = [int(file[-(extLen + counter_digits):-extLen]) if file[-(extLen + counter_digits):-extLen].isdecimal() else 0 for file in files if one_counter_per_folder or file.startswith(filename_prefix)]
|
|
|
|
|
else:
|
|
|
|
|
counters = [int(file[:counter_digits]) if file[:counter_digits].isdigit() else 0 for file in files if one_counter_per_folder or file[counter_digits +1:].startswith(filename_prefix)]
|
|
|
|
|
counters = [int(file[:counter_digits]) if file[:counter_digits].isdecimal() else 0 for file in files if one_counter_per_folder or file[counter_digits +1:].startswith(filename_prefix)]
|
|
|
|
|
|
|
|
|
|
if counters:
|
|
|
|
|
counter = max(counters) + 1
|
|
|
|
@@ -160,7 +159,8 @@ class SaveImageExtended:
|
|
|
|
|
return counter
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
def find_keys_recursively(self, prompt, keys_to_find, found_values):
|
|
|
|
|
# find_keys_recursively is a self-updating recursive method, that will update the dict found_values
|
|
|
|
|
def find_keys_recursively(self, prompt={}, keys_to_find=[], found_values={}):
|
|
|
|
|
for key, value in prompt.items():
|
|
|
|
|
if key in keys_to_find:
|
|
|
|
|
found_values[key] = value
|
|
|
|
@@ -168,15 +168,17 @@ class SaveImageExtended:
|
|
|
|
|
self.find_keys_recursively(value, keys_to_find, found_values)
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
def cleanup_fileName(self, value):
|
|
|
|
|
if isinstance(value, str):
|
|
|
|
|
def cleanup_fileName(self, file='', extToRemove=extToRemove):
|
|
|
|
|
if isinstance(file, str):
|
|
|
|
|
# takes care of all the possible safetensor extensions under the sun
|
|
|
|
|
value = os.path.splitext(os.path.basename(value))[0]
|
|
|
|
|
return value
|
|
|
|
|
# cannot do that... maybe the user want a string.string fixed value to use, that does not end with extToRemove
|
|
|
|
|
# file = os.path.splitext(os.path.basename(file))[0]
|
|
|
|
|
for ext in extToRemove: file = file.removesuffix(ext)
|
|
|
|
|
return file
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
# this method pretty much does the same as find_keys_recursively, except it's for job.json export
|
|
|
|
|
def find_parameter_values(self, target_keys, prompt, found_values={}):
|
|
|
|
|
def find_parameter_values(self, target_keys, prompt={}, found_values={}):
|
|
|
|
|
loras_string = ''
|
|
|
|
|
for key, value in prompt.items():
|
|
|
|
|
# print(f"debug find_parameter_values: key={key} value={value}")
|
|
|
|
@@ -204,9 +206,15 @@ class SaveImageExtended:
|
|
|
|
|
return found_values
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
def generate_custom_name(self, keys_to_extract, prefix, delimiter, resolution, prompt):
|
|
|
|
|
# String Type Example isdecimal() isdigit() isnumeric()
|
|
|
|
|
# --------------------------- --------- ----------- --------- -----------
|
|
|
|
|
# Base 10 Numbers '0123' True True True
|
|
|
|
|
# Fractions and Superscripts '⅔','2²' False True True
|
|
|
|
|
# Roman Numerals 'ↁ' False False True
|
|
|
|
|
# --------------------------- --------- ----------- --------- -----------
|
|
|
|
|
def generate_custom_name(self, keys_to_extract, prefix, delimiter, prompt):
|
|
|
|
|
custom_name = prefix
|
|
|
|
|
if prompt is not None and len(keys_to_extract) > 0:
|
|
|
|
|
if prompt is not None and keys_to_extract != ['']:
|
|
|
|
|
found_values = {}
|
|
|
|
|
# print(f"debug generate_custom_name: --prefix: {prefix}")
|
|
|
|
|
# print(f"debug generate_custom_name: --keys_to_extract: {keys_to_extract}")
|
|
|
|
@@ -221,51 +229,82 @@ class SaveImageExtended:
|
|
|
|
|
# 'inputs': {'cfg': 1.6,
|
|
|
|
|
# 'denoise': 1.0, ...
|
|
|
|
|
for key in keys_to_extract:
|
|
|
|
|
value = None
|
|
|
|
|
node, nodeKey = None, None
|
|
|
|
|
|
|
|
|
|
# check if this is a subfolder: starts with ./ or /, can also end with /
|
|
|
|
|
if '/' in key:
|
|
|
|
|
# key is a fsubfolder
|
|
|
|
|
value = key
|
|
|
|
|
else:
|
|
|
|
|
splitKey = key.split('.')
|
|
|
|
|
# we also exclude cases like "string." or ".string" or "string.string"
|
|
|
|
|
if len(splitKey) > 1:
|
|
|
|
|
# key has the form string.string
|
|
|
|
|
if '' not in splitKey:
|
|
|
|
|
# key has the form string.string
|
|
|
|
|
if splitKey[0].isdecimal():
|
|
|
|
|
# key has the form num.widget_name like 123.widget_name, we will then look for widget_name value in node #123
|
|
|
|
|
node, nodeKey = splitKey[0], splitKey[1]
|
|
|
|
|
if node in prompt:
|
|
|
|
|
# print(f"debug generate_custom_name: --node.nodeKey = {node}.{nodeKey}")
|
|
|
|
|
# splitKey[0] = #node number found in prompt, we will recurse only in that node:
|
|
|
|
|
value = self.find_keys_recursively(prompt[node], [nodeKey], found_values)
|
|
|
|
|
self.find_keys_recursively(prompt[node], [nodeKey], found_values)
|
|
|
|
|
else:
|
|
|
|
|
# if splitKey[0] = #node number not found in prompt, we will just inform the user
|
|
|
|
|
# if splitKey[0] = #num node not found in prompt; #num could have changed or user made a typo
|
|
|
|
|
print(f"SaveImageExtended info: node #{node} not found")
|
|
|
|
|
self.find_keys_recursively(prompt, [nodeKey], found_values)
|
|
|
|
|
else:
|
|
|
|
|
# from now on we will wotk with nodeKey, that will save us tons of if then else
|
|
|
|
|
# key is in the form string.string = fixed string
|
|
|
|
|
value = self.cleanup_fileName(key)
|
|
|
|
|
else:
|
|
|
|
|
# key is in the form ".string" or "string." or "." - we won't clean that up and keep as is
|
|
|
|
|
value = key
|
|
|
|
|
else:
|
|
|
|
|
# nodeKey is not a folder, has no dot, could be a valid key to find, could be a fixed string - keep as is
|
|
|
|
|
nodeKey = key
|
|
|
|
|
# we just try and find the last value for that key, whichever node it's in:
|
|
|
|
|
value = self.find_keys_recursively(prompt, [nodeKey], found_values)
|
|
|
|
|
self.find_keys_recursively(prompt, [nodeKey], found_values)
|
|
|
|
|
# is key num.widget_name
|
|
|
|
|
# is key subfolder
|
|
|
|
|
|
|
|
|
|
# at this point we have a nodeKey but maybe no value
|
|
|
|
|
# now we analyze each value found and format them accordingly:
|
|
|
|
|
value = found_values[nodeKey]
|
|
|
|
|
# at this point we have a nodeKey, or a value, or both
|
|
|
|
|
# print(f"debug generate_custom_name: ----key: {nodeKey}")
|
|
|
|
|
# print(f"debug generate_custom_name: ----value: {value}")
|
|
|
|
|
|
|
|
|
|
if value is None:
|
|
|
|
|
# key not found = it's a fixed string
|
|
|
|
|
# print(f"debug generate_custom_name: ------value=key: {nodeKey}")
|
|
|
|
|
if nodeKey is not None:
|
|
|
|
|
if nodeKey in found_values: value = found_values[nodeKey]
|
|
|
|
|
if value is None:
|
|
|
|
|
value = nodeKey
|
|
|
|
|
else:
|
|
|
|
|
value = self.cleanup_fileName(value)
|
|
|
|
|
|
|
|
|
|
# at this point, value is not None anymore
|
|
|
|
|
# now we analyze each value found and format them accordingly:
|
|
|
|
|
# print(f"debug generate_custom_name: ----value: {value}")
|
|
|
|
|
delim = delimiter
|
|
|
|
|
|
|
|
|
|
# now we build the custom_name:
|
|
|
|
|
if isinstance(value, str):
|
|
|
|
|
# prefix and keys can very well be subfolders ending or starting with a /
|
|
|
|
|
if custom_name.endswith('/'):
|
|
|
|
|
# for subfolders, do not start filename with a delimiter...
|
|
|
|
|
delim = ''
|
|
|
|
|
else:
|
|
|
|
|
# for subfolders in keys, do not clean the filename...
|
|
|
|
|
if '/' in value and not value.endswith('/'):
|
|
|
|
|
# print(f"debug generate_custom_name: ---------: folder")
|
|
|
|
|
delim= ''
|
|
|
|
|
# ".string" case
|
|
|
|
|
if value.startswith('.'):
|
|
|
|
|
delim = ''
|
|
|
|
|
|
|
|
|
|
if isinstance(value, float):
|
|
|
|
|
value = round(float(value), 1)
|
|
|
|
|
|
|
|
|
|
# now we build the custom_name:
|
|
|
|
|
if isinstance(value, str):
|
|
|
|
|
value = self.cleanup_fileName(value)
|
|
|
|
|
# prefix and keys can very well be subfolders ending or starting with a /
|
|
|
|
|
if (value.startswith('./') or value.startswith('/') or custom_name.endswith('/')):
|
|
|
|
|
# for subfolders, do not start filename with a delimiter...
|
|
|
|
|
custom_name += f"{value}"
|
|
|
|
|
else:
|
|
|
|
|
custom_name += f"{delimiter}{value}"
|
|
|
|
|
else:
|
|
|
|
|
# could be int or float, can't be anything else
|
|
|
|
|
custom_name += f"{delimiter}{value}"
|
|
|
|
|
custom_name += f"{delim}{value}"
|
|
|
|
|
# print(f"debug generate_custom_name: ------custom_name: {custom_name}")
|
|
|
|
|
return custom_name.strip(delimiter)
|
|
|
|
|
# for each key
|
|
|
|
|
return custom_name.strip(delimiter).strip('.').strip('/').strip(delimiter)
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
def save_job_to_json(self, save_job_data, prompt, filename_prefix, positive_text_opt, negative_text_opt, job_custom_text, resolution, output_path, filename):
|
|
|
|
@@ -459,39 +498,53 @@ class SaveImageExtended:
|
|
|
|
|
# class SaveImageExtended -------------------------------------------------------------------------------
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
# node will never return None values, except for optional input. Impossible.
|
|
|
|
|
def save_images(self,
|
|
|
|
|
counter_digits,
|
|
|
|
|
counter_position,
|
|
|
|
|
one_counter_per_folder,
|
|
|
|
|
delimiter,
|
|
|
|
|
filename_keys,
|
|
|
|
|
foldername_keys,
|
|
|
|
|
images,
|
|
|
|
|
image_preview,
|
|
|
|
|
filename_prefix,
|
|
|
|
|
filename_keys,
|
|
|
|
|
foldername_prefix,
|
|
|
|
|
foldername_keys,
|
|
|
|
|
delimiter,
|
|
|
|
|
save_job_data,
|
|
|
|
|
job_data_per_image,
|
|
|
|
|
job_custom_text,
|
|
|
|
|
save_metadata,
|
|
|
|
|
filename_prefix=filename_prefix,
|
|
|
|
|
foldername_prefix=foldername_prefix,
|
|
|
|
|
extra_pnginfo=None,
|
|
|
|
|
counter_digits,
|
|
|
|
|
counter_position,
|
|
|
|
|
one_counter_per_folder,
|
|
|
|
|
image_preview,
|
|
|
|
|
output_ext,
|
|
|
|
|
negative_text_opt=None,
|
|
|
|
|
positive_text_opt=None,
|
|
|
|
|
extra_pnginfo=None,
|
|
|
|
|
prompt=None,
|
|
|
|
|
output_ext=output_ext
|
|
|
|
|
):
|
|
|
|
|
|
|
|
|
|
# print(f"filename_prefix = x{filename_prefix}x")
|
|
|
|
|
# print(f"filename_keys = x{filename_keys}x")
|
|
|
|
|
# print(f"foldername_prefix = x{foldername_prefix}x")
|
|
|
|
|
# print(f"foldername_keys = x{foldername_keys}x")
|
|
|
|
|
# print(f"delimiter = x{delimiter}x")
|
|
|
|
|
# print(f"save_job_data = x{save_job_data}x")
|
|
|
|
|
# print(f"job_data_per_image = x{job_data_per_image}x")
|
|
|
|
|
# print(f"output_ext = x{output_ext}x")
|
|
|
|
|
|
|
|
|
|
# apply default values: we replicate the default save image box
|
|
|
|
|
if not filename_prefix and not filename_keys: filename_prefix=self.filename_prefix
|
|
|
|
|
if delimiter: delimiter = delimiter[0]
|
|
|
|
|
|
|
|
|
|
filename_keys_to_extract = [item.strip() for item in filename_keys.split(',')]
|
|
|
|
|
foldername_keys_to_extract = [item.strip() for item in foldername_keys.split(',')]
|
|
|
|
|
|
|
|
|
|
custom_filename = self.generate_custom_name(filename_keys_to_extract, filename_prefix, delimiter, prompt)
|
|
|
|
|
custom_foldername = self.generate_custom_name(foldername_keys_to_extract, foldername_prefix, delimiter, prompt)
|
|
|
|
|
|
|
|
|
|
# Get set resolution value
|
|
|
|
|
i = 255. * images[0].cpu().numpy()
|
|
|
|
|
img = Image.fromarray(np.clip(i, 0, 255).astype(np.uint8))
|
|
|
|
|
resolution = f'{img.width}x{img.height}'
|
|
|
|
|
|
|
|
|
|
delimiter = delimiter[0]
|
|
|
|
|
filename_keys_to_extract = [item.strip() for item in filename_keys.split(',')]
|
|
|
|
|
foldername_keys_to_extract = [item.strip() for item in foldername_keys.split(',')]
|
|
|
|
|
custom_filename = self.generate_custom_name(filename_keys_to_extract, filename_prefix, delimiter, resolution, prompt)
|
|
|
|
|
custom_foldername = self.generate_custom_name(foldername_keys_to_extract, foldername_prefix, delimiter, resolution, prompt)
|
|
|
|
|
|
|
|
|
|
# Create folders, count images, save images
|
|
|
|
|
try:
|
|
|
|
|
full_output_folder, filename, _, _, custom_filename = folder_paths.get_save_image_path(custom_filename, self.output_dir, images[0].shape[1], images[0].shape[0])
|
|
|
|
|