| |
|
|
| 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) |
| |
| 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') |
|
|
| |
| 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) |
| |
| 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') |
|
|
| |
| 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)" |
| } |
