├── .gitignore
├── README.md
├── __init__.py
├── assets
    ├── example.png
    └── prompts.png
└── save_image_extended.py


/.gitignore:
--------------------------------------------------------------------------------
1 | __pycache__/


--------------------------------------------------------------------------------
/README.md:
--------------------------------------------------------------------------------
 1 | # Save Image Extended for ComfyUI
 2 | 
 3 | <p align="center">
 4 |  <img src="assets/example.png" />
 5 | </p>
 6 | 
 7 |  Customize the information saved in file- and folder names. Use the values of sampler parameters as part of file or folder names. <br>Save data about the generated job (sampler, prompts, models) as entries in a `json` (text) file, in each folder.
 8 | 
 9 |  ## Installation
10 | 1. Open a terminal inside the 'custom_nodes' folder located in your ComfyUI installation dir
11 | 2. Use the `git clone` command to clone the [save-image-extended-comfyui](https://github.com/thedyze/save-image-extended-comfyui) repo.
12 | ```
13 | git clone https://github.com/thedyze/save-image-extended-comfyui
14 | ```
15 | 
16 | ## Parameters / Usage
17 | 
18 | - `filename_prefix` -  String prefix added to files.
19 | - `filename_keys` - Comma separated string with sampler parameters to add to filename. E.g: `sampler_name, scheduler, cfg, denoise` Added to filename in written order. `resolution`  also works. `vae_name` `model_name` (upscale model), `ckpt_name` (checkpoint) are others that should work. Here you can try any parameter name of any node. As long as the parameter has the same variable name defined in the `prompt` object they should work. The same applies to `foldername_keys`. 
20 | - `foldername_prefix` - String prefix added to folders.
21 | - `foldername_keys` - Comma separated string with sampler parameters to add to foldername.
22 | - `delimiter` - Delimiter character, either `underscore`, `dot`, or `comma`.
23 | - `save_job_data` - If enabled, saves information about each job as entries in a `jobs.json` text file, inside the generated folder. Mulitple options for saving `prompt`, `basic data`, `sampler settings`, `loaded models`.
24 | - `job_data_per_image` - When enabled, saves individual job data files for each image.
25 | - `job_custom_text` - Custom string to save along with the job data. Right click the node and convert to input to connect with another node.
26 | - `save_metadata` - Saves metadata into the image.
27 | - `counter_digits` - Number of digits used for the image counter. `3` = image_001.png. Will adjust the counter if files are deleted. Looks for the highest number in the folder, does not fill gaps.
28 | - `counter_position` - Image counter first or last in the filename.
29 | - `one_counter_per_folder` - Toggles the counter. Either one counter per folder, or resets when a parameter/prompt changes.
30 | - `image_preview` - Turns the image preview on and off.
31 | 
32 | ## Node inputs
33 | 
34 | - `images` - The generated images.
35 | - `positive_text_opt` - Optional string input for when using custom nodes for positive prompt text.
36 | - `negative_text_opt` - Optional string input for when using custom nodes for negative prompt text.
37 | 
38 | ## Automatic folder names and date/time in names
39 | 
40 | Convert the 'prefix' parameters to inputs (right click in the node and select e.g 'convert foldername_prefix to input'. Then attach the 'Get Date Time String' custom node from JPS to these inputs. This way a new folder name can be automatically generated each time generate is pressed.
41 | #
42 | Disclaimer: Does not check for illegal characters entered in file or folder names. May not be compatible with every other custom node, depending on changes in the `prompt` object. Tested and working with default samplers, Efficiency nodes and UltimateSDUpscale.
43 | #
44 | <br>
45 | 
46 | <p align="center">
47 |  <img src="assets/prompts.png" />
48 | <br><br>
49 |  Happy saving!
50 | </p>
51 | 
52 | 
53 | 
54 | 


--------------------------------------------------------------------------------
/__init__.py:
--------------------------------------------------------------------------------
1 | from .save_image_extended import NODE_CLASS_MAPPINGS, NODE_DISPLAY_NAME_MAPPINGS
2 | 
3 | __all__ = ['NODE_CLASS_MAPPINGS', 'NODE_DISPLAY_NAME_MAPPINGS']
4 | 


--------------------------------------------------------------------------------
/assets/example.png:
--------------------------------------------------------------------------------
https://raw.githubusercontent.com/thedyze/save-image-extended-comfyui/5f104fddfc8281d9fd85d03e63c5f70454db6701/assets/example.png


--------------------------------------------------------------------------------
/assets/prompts.png:
--------------------------------------------------------------------------------
https://raw.githubusercontent.com/thedyze/save-image-extended-comfyui/5f104fddfc8281d9fd85d03e63c5f70454db6701/assets/prompts.png


--------------------------------------------------------------------------------
/save_image_extended.py:
--------------------------------------------------------------------------------
  1 | import os
  2 | import re
  3 | import sys
  4 | import json
  5 | from PIL import Image
  6 | from PIL.PngImagePlugin import PngInfo
  7 | import numpy as np
  8 | import locale
  9 | from datetime import datetime
 10 | from pathlib import Path
 11 | 
 12 | sys.path.insert(0, os.path.join(os.path.dirname(os.path.realpath(__file__)), 'comfy'))
 13 | original_locale = locale.setlocale(locale.LC_TIME, '')
 14 | 
 15 | import folder_paths
 16 | 
 17 | class SaveImageExtended:
 18 | 	def __init__(self):
 19 | 		self.output_dir = folder_paths.get_output_directory()
 20 | 		self.type = 'output'
 21 | 		self.prefix_append = ''
 22 | 
 23 | 	@classmethod
 24 | 	def INPUT_TYPES(s):
 25 | 		return {
 26 | 			'required': {
 27 | 				'images': ('IMAGE', ),
 28 | 				'filename_prefix': ('STRING', {'default': 'myFile'}),
 29 | 				'filename_keys': ('STRING', {'default': 'steps, cfg', 'multiline': False}),
 30 | 				'foldername_prefix': ('STRING', {'default': 'myPix'}),
 31 | 				'foldername_keys': ('STRING', {'default': 'sampler_name, scheduler', 'multiline': False}),
 32 | 				'delimiter': (['underscore','dot', 'comma'], {'default': 'underscore'}),
 33 | 				'save_job_data': (['disabled', 'prompt', 'basic, prompt', 'basic, sampler, prompt', 'basic, models, sampler, prompt'],{'default': 'basic, prompt'}),
 34 | 				'job_data_per_image': (['disabled', 'enabled'],{'default': 'disabled'}),
 35 | 				'job_custom_text': ('STRING', {'default': '', 'multiline': False}),
 36 | 				'save_metadata': (['disabled', 'enabled'], {'default': 'enabled'}),
 37 | 				'counter_digits': ([2, 3, 4, 5, 6], {'default': 3}),
 38 | 				'counter_position': (['first', 'last'], {'default': 'last'}),
 39 | 				'one_counter_per_folder': (['disabled', 'enabled'], {'default': 'disabled'}),
 40 | 				'image_preview': (['disabled', 'enabled'], {'default': 'enabled'}),
 41 | 			},
 42 | 			"optional": {
 43 |                     "positive_text_opt": ("STRING", {"forceInput": True}),
 44 | 					"negative_text_opt": ("STRING", {"forceInput": True}),
 45 |                     },
 46 | 			'hidden': {'prompt': 'PROMPT', 'extra_pnginfo': 'EXTRA_PNGINFO'},
 47 | 		}
 48 | 
 49 | 	RETURN_TYPES = ()
 50 | 	FUNCTION = 'save_images'
 51 | 	OUTPUT_NODE = True
 52 | 	CATEGORY = 'image'
 53 | 
 54 | 	def get_subfolder_path(self, image_path, output_path):
 55 | 		image_path = Path(image_path).resolve()
 56 | 		output_path = Path(output_path).resolve()
 57 | 		relative_path = image_path.relative_to(output_path)
 58 | 		subfolder_path = relative_path.parent
 59 | 
 60 | 		return str(subfolder_path)
 61 | 
 62 | 	# Get current counter number from file names
 63 | 	def get_latest_counter(self, one_counter_per_folder, folder_path, filename_prefix, counter_digits, counter_position='last'):
 64 | 		counter = 1
 65 | 		if not os.path.exists(folder_path):
 66 | 			print(f"Folder {folder_path} does not exist, starting counter at 1.")
 67 | 			return counter
 68 | 
 69 | 		try:
 70 | 			files = [f for f in os.listdir(folder_path) if f.endswith('.png')]
 71 | 			if files:
 72 | 				if counter_position == 'last':
 73 | 					counters = [int(f[-(4 + counter_digits):-4]) if f[-(4 + counter_digits):-4].isdigit() else 0 for f in files if one_counter_per_folder == 'enabled' or f.startswith(filename_prefix)]
 74 | 				elif counter_position == 'first':
 75 | 					counters = [int(f[:counter_digits]) if f[:counter_digits].isdigit() else 0 for f in files if one_counter_per_folder == 'enabled' or f[counter_digits +1:].startswith(filename_prefix)]
 76 | 				else:
 77 | 					print("Invalid counter_position. Using 'last' as default.")
 78 | 					counters = [int(f[-(4 + counter_digits):-4]) if f[-(4 + counter_digits):-4].isdigit() else 0 for f in files if one_counter_per_folder == 'enabled' or f.startswith(filename_prefix)]
 79 | 
 80 | 				if counters:
 81 | 					counter = max(counters) + 1
 82 | 
 83 | 		except Exception as e:
 84 | 			print(f"An error occurred while finding the latest counter: {e}")
 85 | 
 86 | 		return counter
 87 | 
 88 | 	@staticmethod
 89 | 	def find_keys_recursively(d, keys_to_find, found_values):
 90 | 		for key, value in d.items():
 91 | 			if key in keys_to_find:
 92 | 				found_values[key] = value
 93 | 			if isinstance(value, dict):
 94 | 				SaveImageExtended.find_keys_recursively(value, keys_to_find, found_values)
 95 | 
 96 | 	@staticmethod
 97 | 	def remove_file_extension(value):
 98 | 		if isinstance(value, str) and value.endswith('.safetensors'):
 99 | 			base_value = os.path.basename(value)
100 | 			value = base_value[:-12]
101 | 		if isinstance(value, str) and value.endswith('.pt'):
102 | 			base_value = os.path.basename(value)
103 | 			value = base_value[:-3]
104 | 
105 | 		return value
106 | 
107 | 	@staticmethod
108 | 	def find_parameter_values(target_keys, obj, found_values=None):
109 | 		if found_values is None:
110 | 			found_values = {}
111 | 
112 | 		if not isinstance(target_keys, list):
113 | 			target_keys = [target_keys]
114 | 
115 | 		loras_string = ''
116 | 		for key, value in obj.items():
117 | 			if 'loras' in target_keys:
118 | 				# Match both formats: lora_xx and lora_name_x
119 | 				if re.match(r'lora(_name)?(_\d+)?', key):
120 | 					if value.endswith('.safetensors'):
121 | 						value = SaveImageExtended.remove_file_extension(value)
122 | 					if value != 'None':
123 | 						loras_string += f'{value}, '
124 | 
125 | 			if key in target_keys:
126 | 				if (isinstance(value, str) and value.endswith('.safetensors')) or (isinstance(value, str) and value.endswith('.pt')):
127 | 					value = SaveImageExtended.remove_file_extension(value)
128 | 				found_values[key] = value
129 | 
130 | 			if isinstance(value, dict):
131 | 				SaveImageExtended.find_parameter_values(target_keys, value, found_values)
132 | 
133 | 		if 'loras' in target_keys and loras_string:
134 | 			found_values['loras'] = loras_string.strip(', ')
135 | 
136 | 		if len(target_keys) == 1:
137 | 			return found_values.get(target_keys[0], None)
138 | 
139 | 		return found_values
140 | 
141 | 	@staticmethod
142 | 	def generate_custom_name(keys_to_extract, prefix, delimiter_char, resolution, prompt):
143 | 		custom_name = prefix
144 | 
145 | 		if prompt is not None and len(keys_to_extract) > 0:
146 | 			found_values = {'resolution': resolution}
147 | 			SaveImageExtended.find_keys_recursively(prompt, keys_to_extract, found_values)
148 | 			for key in keys_to_extract:
149 | 				value = found_values.get(key)
150 | 				if value is not None:
151 | 					if key == 'cfg' or key =='denoise':
152 | 						try:
153 | 							value = round(float(value), 1)
154 | 						except ValueError:
155 | 							pass
156 | 
157 | 					if (isinstance(value, str) and value.endswith('.safetensors')) or (isinstance(value, str) and value.endswith('.pt')):
158 | 						value = SaveImageExtended.remove_file_extension(value)
159 | 
160 | 					custom_name += f'{delimiter_char}{value}'
161 | 
162 | 		return custom_name.strip(delimiter_char)
163 | 
164 | 	@staticmethod
165 | 	def save_job_to_json(save_job_data, prompt, filename_prefix, positive_text_opt, negative_text_opt, job_custom_text, resolution, output_path, filename):
166 | 		prompt_keys_to_save = {}
167 | 		if 'basic' in save_job_data:
168 | 			if len(filename_prefix) > 0:
169 | 				prompt_keys_to_save['filename_prefix'] = filename_prefix
170 | 			prompt_keys_to_save['resolution'] = resolution
171 | 		if len(job_custom_text) > 0:
172 | 			prompt_keys_to_save['custom_text'] = job_custom_text
173 | 
174 | 		if 'models' in save_job_data:
175 | 			models = SaveImageExtended.find_parameter_values(['ckpt_name', 'loras', 'vae_name', 'model_name'], prompt)
176 | 			if models.get('ckpt_name'):
177 | 				prompt_keys_to_save['checkpoint'] = models['ckpt_name']
178 | 			if models.get('loras'):
179 | 				prompt_keys_to_save['loras'] = models['loras']
180 | 			if models.get('vae_name'):
181 | 				prompt_keys_to_save['vae'] = models['vae_name']
182 | 			if models.get('model_name'):
183 | 				prompt_keys_to_save['upscale_model'] = models['model_name']
184 | 
185 | 
186 | 
187 | 		if 'sampler' in save_job_data:
188 | 			prompt_keys_to_save['sampler_parameters'] = SaveImageExtended.find_parameter_values(['seed', 'steps', 'cfg', 'sampler_name', 'scheduler', 'denoise'], prompt)
189 | 
190 | 		if 'prompt' in save_job_data:
191 | 			if positive_text_opt is not None:
192 | 				if not (isinstance(positive_text_opt, list) and
193 | 						len(positive_text_opt) == 2 and
194 | 						isinstance(positive_text_opt[0], str) and
195 | 						len(positive_text_opt[0]) < 6 and
196 | 						isinstance(positive_text_opt[1], (int, float))):
197 | 					prompt_keys_to_save['positive_prompt'] = positive_text_opt
198 | 
199 | 			if negative_text_opt is not None:
200 | 				if not (isinstance(positive_text_opt, list) and len(negative_text_opt) == 2 and isinstance(negative_text_opt[0], str) and isinstance(negative_text_opt[1], (int, float))):
201 | 					prompt_keys_to_save['negative_prompt'] = negative_text_opt
202 | 
203 | 			#If no user input for prompts
204 | 			if positive_text_opt is None and negative_text_opt is None:
205 | 				if prompt is not None:
206 | 					for key in prompt:
207 | 						class_type = prompt[key].get('class_type', None)
208 | 						inputs = prompt[key].get('inputs', {})
209 | 
210 | 						# Efficiency Loaders prompt structure
211 | 						if class_type == 'Efficient Loader' or class_type == 'Eff. Loader SDXL':
212 | 							if 'positive' in inputs and 'negative' in inputs:
213 | 								prompt_keys_to_save['positive_prompt'] = inputs.get('positive')
214 | 								prompt_keys_to_save['negative_prompt'] = inputs.get('negative')
215 | 
216 | 						# KSampler/UltimateSDUpscale prompt structure
217 | 						elif class_type == 'KSampler' or class_type == 'KSamplerAdvanced' or class_type == 'UltimateSDUpscale':
218 | 							positive_ref = inputs.get('positive', [])[0] if 'positive' in inputs else None
219 | 							negative_ref = inputs.get('negative', [])[0] if 'negative' in inputs else None
220 | 
221 | 							positive_text = prompt.get(str(positive_ref), {}).get('inputs', {}).get('text', None)
222 | 							negative_text = prompt.get(str(negative_ref), {}).get('inputs', {}).get('text', None)
223 | 
224 | 							# If we get non text inputs
225 | 							if positive_text is not None:
226 | 								if isinstance(positive_text, list):
227 | 									if len(positive_text) == 2:
228 | 										if isinstance(positive_text[0], str) and len(positive_text[0]) < 6:
229 | 											if isinstance(positive_text[1], (int, float)):
230 | 												continue
231 | 								prompt_keys_to_save['positive_prompt'] = positive_text
232 | 
233 | 							if negative_text is not None:
234 | 								if isinstance(negative_text, list):
235 | 									if len(negative_text) == 2:
236 | 										if isinstance(negative_text[0], str) and len(negative_text[0]) < 6:
237 | 											if isinstance(negative_text[1], (int, float)):
238 | 												continue
239 | 								prompt_keys_to_save['negative_prompt'] = negative_text
240 | 
241 | 		# Append data and save
242 | 		json_file_path = os.path.join(output_path, filename)
243 | 		existing_data = {}
244 | 		if os.path.exists(json_file_path):
245 | 			try:
246 | 				with open(json_file_path, 'r') as f:
247 | 					existing_data = json.load(f)
248 | 			except json.JSONDecodeError:
249 | 				print(f"The file {json_file_path} is empty or malformed. Initializing with empty data.")
250 | 				existing_data = {}
251 | 
252 | 		timestamp = datetime.now().strftime('%c')
253 | 		new_entry = {}
254 | 		new_entry[timestamp] = prompt_keys_to_save
255 | 		existing_data.update(new_entry)
256 | 
257 | 		with open(json_file_path, 'w') as f:
258 | 			json.dump(existing_data, f, indent=4)
259 | 
260 | 
261 | 	def save_images(self,
262 | 				 counter_digits,
263 | 				 counter_position,
264 | 				 one_counter_per_folder,
265 | 				 delimiter,
266 | 				 filename_keys,
267 | 				 foldername_keys,
268 | 				 images,
269 | 				 image_preview,
270 | 				 save_job_data,
271 | 				 job_data_per_image,
272 | 				 job_custom_text,
273 | 				 save_metadata,
274 | 				 filename_prefix='',
275 | 				 foldername_prefix='',
276 | 				 extra_pnginfo=None,
277 | 				 negative_text_opt=None,
278 | 				 positive_text_opt=None,
279 | 				 prompt=None
280 | 				):
281 | 
282 | 		delimiter_char = "_" if delimiter =='underscore' else '.' if delimiter =='dot' else ','
283 | 
284 | 		# Get set resolution value
285 | 		i = 255. * images[0].cpu().numpy()
286 | 		img = Image.fromarray(np.clip(i, 0, 255).astype(np.uint8))
287 | 		resolution = f'{img.width}x{img.height}'
288 | 
289 | 		filename_keys_to_extract = [item.strip() for item in filename_keys.split(',')]
290 | 		foldername_keys_to_extract = [item.strip() for item in foldername_keys.split(',')]
291 | 		custom_filename = SaveImageExtended.generate_custom_name(filename_keys_to_extract, filename_prefix, delimiter_char, resolution, prompt)
292 | 		custom_foldername = SaveImageExtended.generate_custom_name(foldername_keys_to_extract, foldername_prefix, delimiter_char, resolution, prompt)
293 | 
294 | 		# Create and save images
295 | 		try:
296 | 			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])
297 | 			output_path = os.path.join(full_output_folder, custom_foldername)
298 | 			os.makedirs(output_path, exist_ok=True)
299 | 			counter = self.get_latest_counter(one_counter_per_folder, output_path, filename, counter_digits, counter_position)
300 | 
301 | 			results = list()
302 | 			for image in images:
303 | 				i = 255. * image.cpu().numpy()
304 | 				img = Image.fromarray(np.clip(i, 0, 255).astype(np.uint8))
305 | 				metadata = None
306 | 				if save_metadata == 'enabled':
307 | 					metadata = PngInfo()
308 | 					if prompt is not None:
309 | 						metadata.add_text('prompt', json.dumps(prompt))
310 | 					if extra_pnginfo is not None:
311 | 						for x in extra_pnginfo:
312 | 							metadata.add_text(x, json.dumps(extra_pnginfo[x]))
313 | 
314 | 				if counter_position == 'last':
315 | 					file = f'{filename}{delimiter_char}{counter:0{counter_digits}}.png'
316 | 				else:
317 | 					file = f'{counter:0{counter_digits}}{delimiter_char}{filename}.png'
318 | 
319 | 				image_path = os.path.join(output_path, file)
320 | 				img.save(image_path, pnginfo=metadata, compress_level=4)
321 | 
322 | 				if save_job_data != 'disabled' and job_data_per_image =='enabled':
323 | 					SaveImageExtended.save_job_to_json(save_job_data, prompt, filename_prefix, positive_text_opt, negative_text_opt, job_custom_text, resolution, output_path, f'{file.strip(".png")}.json')
324 | 
325 | 				subfolder = self.get_subfolder_path(image_path, self.output_dir)
326 | 				results.append({ 'filename': file, 'subfolder': subfolder, 'type': self.type})
327 | 				counter += 1
328 | 
329 | 			if save_job_data != 'disabled' and job_data_per_image =='disabled':
330 | 				SaveImageExtended.save_job_to_json(save_job_data, prompt, filename_prefix, positive_text_opt, negative_text_opt, job_custom_text, resolution, output_path, 'jobs.json')
331 | 
332 | 		except OSError as e:
333 | 			print(f'An error occurred while creating the subfolder or saving the image: {e}')
334 | 		else:
335 | 			if image_preview == 'disabled':
336 | 				results = list()
337 | 			return { 'ui': { 'images': results } }
338 | 
339 | NODE_CLASS_MAPPINGS = {
340 |     'SaveImageExtended': SaveImageExtended,
341 | }
342 | 
343 | NODE_DISPLAY_NAME_MAPPINGS = {
344 |     'SaveImageExtended': 'Save Image Extended',
345 | }
346 | 


--------------------------------------------------------------------------------