# layerstyle advance import os.path import shutil from PIL.PngImagePlugin import PngInfo import datetime from .imagefunc import * class SaveImagePlus: def __init__(self): self.NODE_NAME = 'SaveImagePlus' self.output_dir = folder_paths.get_output_directory() self.type = "output" self.prefix_append = "" self.compress_level = 4 @classmethod def INPUT_TYPES(s): return {"required": {"images": ("IMAGE", ), "custom_path": ("STRING", {"default": ""}), "filename_prefix": ("STRING", {"default": "comfyui"}), "timestamp": (["None", "second", "millisecond"],), "format": (["png", "jpg"],), "quality": ("INT", {"default": 80, "min": 10, "max": 100, "step": 1}), "meta_data": ("BOOLEAN", {"default": False}), "blind_watermark": ("STRING", {"default": ""}), "save_workflow_as_json": ("BOOLEAN", {"default": False}), "preview": ("BOOLEAN", {"default": True}), }, "hidden": {"prompt": "PROMPT", "extra_pnginfo": "EXTRA_PNGINFO"}, } RETURN_TYPES = () FUNCTION = "save_image_plus" OUTPUT_NODE = True CATEGORY = '😺dzNodes/LayerUtility/SystemIO' def save_image_plus(self, images, custom_path, filename_prefix, timestamp, format, quality, meta_data, blind_watermark, preview, save_workflow_as_json, prompt=None, extra_pnginfo=None): now = datetime.datetime.now() custom_path = custom_path.replace("%date", now.strftime("%Y-%m-%d")) custom_path = custom_path.replace("%time", now.strftime("%H-%M-%S")) filename_prefix = filename_prefix.replace("%date", now.strftime("%Y-%m-%d")) filename_prefix = filename_prefix.replace("%time", now.strftime("%H-%M-%S")) filename_prefix += self.prefix_append full_output_folder, filename, counter, subfolder, filename_prefix = folder_paths.get_save_image_path(filename_prefix, self.output_dir, images[0].shape[1], images[0].shape[0]) results = list() temp_sub_dir = generate_random_name('_savepreview_', '_temp', 16) temp_dir = os.path.join(folder_paths.get_temp_directory(), temp_sub_dir) for image in images: i = 255. * image.cpu().numpy() img = Image.fromarray(np.clip(i, 0, 255).astype(np.uint8)) if blind_watermark != "": img_mode = img.mode wm_size = watermark_image_size(img) import qrcode qr = qrcode.QRCode( version=1, error_correction=qrcode.constants.ERROR_CORRECT_H, box_size=20, border=1, ) qr.add_data(blind_watermark.encode('utf-8')) qr.make(fit=True) qr_image = qr.make_image(fill_color="black", back_color="white") qr_image = qr_image.resize((wm_size, wm_size), Image.BICUBIC).convert("L") y, u, v, _ = image_channel_split(img, mode='YCbCr') _u = add_invisibal_watermark(u, qr_image) wm_img = image_channel_merge((y, _u, v), mode='YCbCr') if img.mode == "RGBA": img = RGB2RGBA(wm_img, img.split()[-1]) else: img = wm_img.convert(img_mode) metadata = None if meta_data: metadata = PngInfo() if prompt is not None: metadata.add_text("prompt", json.dumps(prompt)) if extra_pnginfo is not None: for x in extra_pnginfo: metadata.add_text(x, json.dumps(extra_pnginfo[x])) if timestamp == "millisecond": file = f'{filename}_{now.strftime("%Y-%m-%d_%H-%M-%S-%f")[:-3]}' elif timestamp == "second": file = f'{filename}_{now.strftime("%Y-%m-%d_%H-%M-%S")}' else: file = f'{filename}_{counter:05}' preview_filename = "" if custom_path != "": if not os.path.exists(custom_path): try: os.makedirs(custom_path) except Exception as e: log(f"Error: {self.NODE_NAME} skipped, because unable to create temporary folder.", message_type='warning') raise FileNotFoundError(f"cannot create custom_path {custom_path}, {e}") full_output_folder = os.path.normpath(custom_path) # save preview image to temp_dir if os.path.isdir(temp_dir): shutil.rmtree(temp_dir) try: os.makedirs(temp_dir) except Exception as e: print(e) log(f"Error: {self.NODE_NAME} skipped, because unable to create temporary folder.", message_type='warning') try: preview_filename = os.path.join(generate_random_name('saveimage_preview_', '_temp', 16) + '.png') img.save(os.path.join(temp_dir, preview_filename)) except Exception as e: print(e) log(f"Error: {self.NODE_NAME} skipped, because unable to create temporary file.", message_type='warning') # check if file exists, change filename while os.path.isfile(os.path.join(full_output_folder, f"{file}.{format}")): counter += 1 if timestamp == "millisecond": file = f'{filename}_{now.strftime("%Y-%m-%d_%H-%M-%S-%f")[:-3]}_{counter:05}' elif timestamp == "second": file = f'{filename}_{now.strftime("%Y-%m-%d_%H-%M-%S")}_{counter:05}' else: file = f"{filename}_{counter:05}" image_file_name = os.path.join(full_output_folder, f"{file}.{format}") json_file_name = os.path.join(full_output_folder, f"{file}.json") if format == "png": img.save(image_file_name, pnginfo=metadata, compress_level= (100 - quality) // 10) else: if img.mode == "RGBA": img = img.convert("RGB") img.save(image_file_name, quality=quality) log(f"{self.NODE_NAME} -> Saving image to {image_file_name}") if save_workflow_as_json: try: workflow = (extra_pnginfo or {}).get('workflow') if workflow is None: log('No workflow found, skipping saving of JSON') with open(f'{json_file_name}', 'w', encoding='utf-8') as workflow_file: json.dump(workflow, workflow_file, ensure_ascii=False, indent=4) log(f'Saved workflow to {json_file_name}') except Exception as e: log( f'Failed to save workflow as json due to: {e}, proceeding with the remainder of saving execution', message_type="warning") if preview: if custom_path == "": results.append({ "filename": f"{file}.{format}", "subfolder": subfolder, "type": self.type }) else: results.append({ "filename": preview_filename, "subfolder": temp_sub_dir, "type": "temp" }) counter += 1 return { "ui": { "images": results } } class SaveImagePlusV2: def __init__(self): self.NODE_NAME = 'SaveImagePlusV2' self.output_dir = folder_paths.get_output_directory() self.type = "output" self.prefix_append = "" self.compress_level = 4 @classmethod def INPUT_TYPES(s): return {"required": {"images": ("IMAGE", ), "custom_path": ("STRING", {"default": ""}), "custom_filename": ("STRING", {"default": ""}), "filename_prefix": ("STRING", {"default": "comfyui"}), "timestamp": (["None", "second", "millisecond"],), "dpi": ("INT", {"default": 96, "min": 1, "max": 10000, "step": 1}), "format": (["png", "jpg"],), "quality": ("INT", {"default": 80, "min": 10, "max": 100, "step": 1}), "meta_data": ("BOOLEAN", {"default": False}), "blind_watermark": ("STRING", {"default": ""}), "save_workflow_as_json": ("BOOLEAN", {"default": False}), "preview": ("BOOLEAN", {"default": True}), }, "hidden": {"prompt": "PROMPT", "extra_pnginfo": "EXTRA_PNGINFO"}, } RETURN_TYPES = () FUNCTION = "save_image_plus_v2" OUTPUT_NODE = True CATEGORY = '😺dzNodes/LayerUtility/SystemIO' def save_image_plus_v2(self, images, custom_path, custom_filename, filename_prefix, timestamp, dpi, format, quality, meta_data, blind_watermark, preview, save_workflow_as_json, prompt=None, extra_pnginfo=None): now = datetime.datetime.now() custom_path = custom_path.replace("%date", now.strftime("%Y-%m-%d")) custom_path = custom_path.replace("%time", now.strftime("%H-%M-%S")) filename_prefix = filename_prefix.replace("%date", now.strftime("%Y-%m-%d")) filename_prefix = filename_prefix.replace("%time", now.strftime("%H-%M-%S")) filename_prefix += self.prefix_append full_output_folder, filename, counter, subfolder, filename_prefix = folder_paths.get_save_image_path(filename_prefix, self.output_dir, images[0].shape[1], images[0].shape[0]) results = list() temp_sub_dir = generate_random_name('_savepreview_', '_temp', 16) temp_dir = os.path.join(folder_paths.get_temp_directory(), temp_sub_dir) for image in images: i = 255. * image.cpu().numpy() img = Image.fromarray(np.clip(i, 0, 255).astype(np.uint8)) if blind_watermark != "": img_mode = img.mode wm_size = watermark_image_size(img) import qrcode qr = qrcode.QRCode( version=1, error_correction=qrcode.constants.ERROR_CORRECT_H, box_size=20, border=1, ) qr.add_data(blind_watermark.encode('utf-8')) qr.make(fit=True) qr_image = qr.make_image(fill_color="black", back_color="white") qr_image = qr_image.resize((wm_size, wm_size), Image.BICUBIC).convert("L") y, u, v, _ = image_channel_split(img, mode='YCbCr') _u = add_invisibal_watermark(u, qr_image) wm_img = image_channel_merge((y, _u, v), mode='YCbCr') if img.mode == "RGBA": img = RGB2RGBA(wm_img, img.split()[-1]) else: img = wm_img.convert(img_mode) metadata = None if meta_data: metadata = PngInfo() if prompt is not None: metadata.add_text("prompt", json.dumps(prompt)) if extra_pnginfo is not None: for x in extra_pnginfo: metadata.add_text(x, json.dumps(extra_pnginfo[x])) if custom_filename == "": if timestamp == "millisecond": file = f'{filename}_{now.strftime("%Y-%m-%d_%H-%M-%S-%f")[:-3]}' elif timestamp == "second": file = f'{filename}_{now.strftime("%Y-%m-%d_%H-%M-%S")}' else: file = f'{filename}_{counter:05}' else: custom_filename = custom_filename.replace("%date", now.strftime("%Y-%m-%d")) custom_filename = custom_filename.replace("%time", now.strftime("%H-%M-%S")) filename = custom_filename file = filename preview_filename = "" if custom_path != "": if not os.path.exists(custom_path): try: os.makedirs(custom_path) except Exception as e: log(f"Error: {self.NODE_NAME} skipped, because unable to create temporary folder.", message_type='warning') raise FileNotFoundError(f"cannot create custom_path {custom_path}, {e}") full_output_folder = os.path.normpath(custom_path) # save preview image to temp_dir if os.path.isdir(temp_dir): shutil.rmtree(temp_dir) try: os.makedirs(temp_dir) except Exception as e: print(e) log(f"Error: {self.NODE_NAME} skipped, because unable to create temporary folder.", message_type='warning') try: preview_filename = os.path.join(generate_random_name('saveimage_preview_', '_temp', 16) + '.png') img.save(os.path.join(temp_dir, preview_filename), pnginfo=metadata, compress_level=(100 - quality) // 10, dpi=(dpi, dpi)) except Exception as e: print(e) log(f"Error: {self.NODE_NAME} skipped, because unable to create temporary file.", message_type='warning') # check if file exists, change filename if custom_filename == "": while os.path.isfile(os.path.join(full_output_folder, f"{file}.{format}")): counter += 1 if timestamp == "millisecond": file = f'{filename}_{now.strftime("%Y-%m-%d_%H-%M-%S-%f")[:-3]}_{counter:05}' elif timestamp == "second": file = f'{filename}_{now.strftime("%Y-%m-%d_%H-%M-%S")}_{counter:05}' else: file = f"{filename}_{counter:05}" image_file_name = os.path.join(full_output_folder, f"{file}.{format}") json_file_name = os.path.join(full_output_folder, f"{file}.json") if format == "png": img.save(image_file_name, pnginfo=metadata, compress_level=(100 - quality) // 10, dpi=(dpi, dpi)) else: if img.mode == "RGBA": img = img.convert("RGB") img.save(image_file_name, quality=quality, dpi=(dpi, dpi)) log(f"{self.NODE_NAME} -> Saving image to {image_file_name}") if save_workflow_as_json: try: workflow = (extra_pnginfo or {}).get('workflow') if workflow is None: log('No workflow found, skipping saving of JSON') with open(f'{json_file_name}', 'w', encoding='utf-8') as workflow_file: json.dump(workflow, workflow_file, ensure_ascii=False, indent=4) log(f'Saved workflow to {json_file_name}') except Exception as e: log( f'Failed to save workflow as json due to: {e}, proceeding with the remainder of saving execution', message_type="warning") if preview: if custom_path == "": results.append({ "filename": f"{file}.{format}", "subfolder": subfolder, "type": self.type }) else: results.append({ "filename": preview_filename, "subfolder": temp_sub_dir, "type": "temp" }) counter += 1 return { "ui": { "images": results } } NODE_CLASS_MAPPINGS = { "LayerUtility: SaveImagePlus": SaveImagePlus, "LayerUtility: SaveImagePlusV2": SaveImagePlusV2 } NODE_DISPLAY_NAME_MAPPINGS = { "LayerUtility: SaveImagePlus": "LayerUtility: SaveImage Plus(Advance)", "LayerUtility: SaveImagePlusV2": "LayerUtility: SaveImagePlus V2(Advance)" }