From 581ee56802f3db25232e5820b5a40d3f986dab51 Mon Sep 17 00:00:00 2001 From: Rhys Yang Date: Fri, 22 Sep 2023 01:00:27 +0800 Subject: [PATCH] Add SD Prompt Saver --- nodes.py | 201 +++++++++++++++++++++++++++++++++++++++++++++++++++---- 1 file changed, 187 insertions(+), 14 deletions(-) diff --git a/nodes.py b/nodes.py index 66c822c..38f0265 100644 --- a/nodes.py +++ b/nodes.py @@ -5,6 +5,25 @@ @description: ComfyUI node version of SD Prompt Reader """ + +import os +import torch +import json +import numpy as np +from pathlib import Path +from PIL import Image, ImageOps +from PIL.PngImagePlugin import PngInfo + +import hashlib +import piexif +import piexif.helper + + +from nodes import MAX_RESOLUTION +from comfy.cli_args import args +import comfy.samplers +import folder_paths + from .stable_diffusion_prompt_reader.sd_prompt_reader.constants import ( SUPPORTED_FORMATS, MESSAGE, @@ -13,12 +32,6 @@ from .stable_diffusion_prompt_reader.sd_prompt_reader.image_data_reader import ( ImageDataReader, ) from .stable_diffusion_prompt_reader.sd_prompt_reader.__version__ import VERSION -import folder_paths -from pathlib import Path -import os -import torch -import numpy as np -from PIL import Image, ImageOps BLUE = "\033[1;34m" CYAN = "\033[36m" @@ -33,9 +46,6 @@ output_to_terminal("Reader core version: " + VERSION) class SDPromptReader: - def __init__(self): - pass - @classmethod def INPUT_TYPES(s): input_dir = folder_paths.get_input_directory() @@ -158,9 +168,172 @@ class SDPromptReader: return True -# A dictionary that contains all nodes you want to export with their names -# NOTE: names should be globally unique -NODE_CLASS_MAPPINGS = {"SDPromptReader": SDPromptReader} +class SDPromptSaver: + def __init__(self): + self.output_dir = folder_paths.get_output_directory() + self.type = "output" + self.prefix_append = "" -# A dictionary that contains the friendly/humanly readable titles for the nodes -NODE_DISPLAY_NAME_MAPPINGS = {"SDPromptReader": "SD Prompt Reader"} + @classmethod + def INPUT_TYPES(s): + return { + "required": { + "images": ("IMAGE",), + "filename_prefix": ("STRING", {"default": "ComfyUI"}), + "model_name": (folder_paths.get_filename_list("checkpoints"),), + "seed": ( + "INT", + { + "default": 0, + "min": 0, + "max": 0xFFFFFFFFFFFFFFFF, + }, + ), + "steps": ( + "INT", + {"default": 20, "min": 1, "max": 10000}, + ), + "cfg": ( + "FLOAT", + { + "default": 8.0, + "min": 0.0, + "max": 100.0, + "step": 0.5, + "round": 0.01, + }, + ), + "sampler_name": (comfy.samplers.KSampler.SAMPLERS,), + "scheduler": (comfy.samplers.KSampler.SCHEDULERS,), + "positive": ("STRING", {"default": "", "multiline": True}), + "negative": ("STRING", {"default": "", "multiline": True}), + "extension": (["png", "jpg", "webp"],), + }, + "optional": { + "width": ( + "INT", + {"default": 0, "min": 1, "max": MAX_RESOLUTION, "step": 8}, + ), + "height": ( + "INT", + {"default": 0, "min": 1, "max": MAX_RESOLUTION, "step": 8}, + ), + "lossless_webp": ("BOOLEAN", {"default": True}), + "jpg_webp_quality": ("INT", {"default": 100, "min": 1, "max": 100}), + }, + "hidden": {"prompt": "PROMPT", "extra_pnginfo": "EXTRA_PNGINFO"}, + } + + RETURN_TYPES = () + FUNCTION = "save_images" + + OUTPUT_NODE = True + + CATEGORY = "image" + + def save_images( + self, + images, + filename_prefix, + model_name: str = "", + seed: int = 0, + steps: int = 0, + cfg: float = 0.0, + sampler_name: str = "", + scheduler: str = "", + positive: str = "", + negative: str = "", + extension: str = "png", + width: int = 0, + height: int = 0, + lossless_webp: bool = True, + jpg_webp_quality: int = 100, + prompt=None, + extra_pnginfo=None, + ): + 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() + + for image in images: + i = 255.0 * image.cpu().numpy() + img = Image.fromarray(np.clip(i, 0, 255).astype(np.uint8)) + metadata = None + comment = ( + f"{positive}\n" + f"Negative prompt: {negative}\n" + f"Steps: {steps}, " + f"Sampler: {sampler_name}{''if scheduler == 'normal' else '_'+scheduler}, " + f"CFG scale: {cfg}, " + f"Seed: {seed}, " + f"Size: {img.width if width==0 else width}x{img.height if height==0 else height}, " + f"Model hash: {self.calculate_shorthash(model_name)}, " + f"Model: {Path(model_name).stem}, " + f"Version: ComfyUI" + ) + + file = Path(full_output_folder) / f"{filename}_{counter:05}_.{extension}" + if extension == "png": + if not args.disable_metadata: + metadata = PngInfo() + metadata.add_text("parameters", comment) + 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])) + + img.save( + file, + pnginfo=metadata, + compress_level=4, + ) + else: + img.save(file, quality=jpg_webp_quality, lossless=lossless_webp) + if not args.disable_metadata: + metadata = piexif.dump( + { + "Exif": { + piexif.ExifIFD.UserComment: piexif.helper.UserComment.dump( + comment, encoding="unicode" + ) + }, + } + ) + print(str(file)) + piexif.insert(metadata, str(file)) + + results.append( + {"filename": file.name, "subfolder": subfolder, "type": self.type} + ) + counter += 1 + + return {"ui": {"images": results}} + + @staticmethod + def calculate_shorthash(model_name): + hash_sha256 = hashlib.sha256() + blksize = 1024 * 1024 + file_name = folder_paths.get_full_path("checkpoints", model_name) + + with open(file_name, "rb") as f: + for chunk in iter(lambda: f.read(blksize), b""): + hash_sha256.update(chunk) + + return hash_sha256.hexdigest()[:10] + + +NODE_CLASS_MAPPINGS = {"SDPromptReader": SDPromptReader, "SDPromptSaver": SDPromptSaver} + +NODE_DISPLAY_NAME_MAPPINGS = { + "SDPromptReader": "SD Prompt Reader", + "SDPromptSaver": "SD Prompt Saver", +}