""" @author: receyuki @title: SD Prompt Reader @nickname: SD Prompt Reader @description: ComfyUI node version of the 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, ) from .stable_diffusion_prompt_reader.sd_prompt_reader.image_data_reader import ( ImageDataReader, ) from .__version__ import VERSION as NODE_VERSION from .stable_diffusion_prompt_reader.sd_prompt_reader.__version__ import ( VERSION as CORE_VERSION, ) BLUE = "\033[1;34m" CYAN = "\033[36m" RESET = "\033[0m" def output_to_terminal(text: str): print(f"{RESET+BLUE}" f"[SD Prompt Reader] " f"{CYAN+text+RESET}") output_to_terminal("Node version: " + NODE_VERSION) output_to_terminal("Core version: " + CORE_VERSION) class SDPromptReader: @classmethod def INPUT_TYPES(s): input_dir = folder_paths.get_input_directory() files = [ f for f in os.listdir(input_dir) if os.path.isfile(os.path.join(input_dir, f)) ] return { "required": { "image": (sorted(files), {"image_upload": True}), "data_index": ( "INT", {"default": 0, "min": 0, "max": 255, "step": 1}, ), }, } RETURN_TYPES = ( "IMAGE", "MASK", "STRING", "STRING", "INT", "INT", "FLOAT", "INT", "INT", "STRING", ) RETURN_NAMES = ( "IMAGE", "MASK", "POSITIVE", "NEGATIVE", "SEED", "STEPS", "CFG", "WIDTH", "HEIGHT", "SETTING", ) FUNCTION = "load_image" CATEGORY = "SD Prompt Reader" OUTPUT_NODE = True def load_image(self, image, data_index): image_path = folder_paths.get_annotated_filepath(image) i = Image.open(image_path) i = ImageOps.exif_transpose(i) image = i.convert("RGB") image = np.array(image).astype(np.float32) / 255.0 image = torch.from_numpy(image)[None,] if "A" in i.getbands(): mask = np.array(i.getchannel("A")).astype(np.float32) / 255.0 mask = 1.0 - torch.from_numpy(mask) else: mask = torch.zeros((64, 64), dtype=torch.float32, device="cpu") if Path(image_path).suffix not in SUPPORTED_FORMATS: output_to_terminal(MESSAGE["suffix_error"][1]) raise ValueError(MESSAGE["suffix_error"][1]) with open(Path(image_path), "rb") as f: image_data = ImageDataReader(f) if not image_data.tool: output_to_terminal(MESSAGE["format_error"][1]) raise ValueError(MESSAGE["format_error"][1]) seed = int( self.param_parser(image_data.parameter.get("seed"), data_index) or 0 ) steps = int( self.param_parser(image_data.parameter.get("steps"), data_index) or 0 ) cfg = float( self.param_parser(image_data.parameter.get("cfg"), data_index) or 0 ) width = int(image_data.width or 0) height = int(image_data.height or 0) output_to_terminal("Positive: \n" + image_data.positive) output_to_terminal("Negative: \n" + image_data.negative) output_to_terminal("Setting: \n" + image_data.setting) return { "ui": { "text": (image_data.positive, image_data.negative, image_data.setting) }, "result": ( image, mask, image_data.positive, image_data.negative, seed, steps, cfg, width, height, image_data.setting, ), } @staticmethod def param_parser(data: str, index: int): data_list = data.strip("()").split(",") return data_list[0] if len(data_list) == 1 else data_list[index] @classmethod def IS_CHANGED(s, image, data_index): image_path = folder_paths.get_annotated_filepath(image) with open(Path(image_path), "rb") as f: image_data = ImageDataReader(f) return image_data.props @classmethod def VALIDATE_INPUTS(s, image, data_index): if not folder_paths.exists_annotated_filepath(image): return "Invalid image file: {}".format(image) return True class SDPromptSaver: def __init__(self): self.output_dir = folder_paths.get_output_directory() self.type = "output" self.prefix_append = "" @classmethod def INPUT_TYPES(s): return { "required": { "images": ("IMAGE",), "filename_prefix": ("STRING", {"default": "ComfyUI"}), }, "optional": { "model_name": (folder_paths.get_filename_list("checkpoints"),), "model_name_str": ("STRING", {"default": ""}), "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,), "sampler_name_str": ("STRING", {"default": ""}), "scheduler": (comfy.samplers.KSampler.SCHEDULERS,), "scheduler_str": ("STRING", {"default": ""}), "positive": ("STRING", {"default": "", "multiline": True}), "negative": ("STRING", {"default": "", "multiline": True}), "extension": (["png", "jpg", "webp"],), "width": ( "INT", {"default": 0, "min": 1, "max": MAX_RESOLUTION, "step": 8}, ), "height": ( "INT", {"default": 0, "min": 1, "max": MAX_RESOLUTION, "step": 8}, ), "calculate_model_hash": ("BOOLEAN", {"default": False}), "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 = "SD Prompt Reader" def save_images( self, images, filename_prefix, model_name: str = "", model_name_str: str = "", seed: int = 0, steps: int = 0, cfg: float = 0.0, sampler_name: str = "", sampler_name_str: str = "", scheduler: str = "", scheduler_str: str = "", positive: str = "", negative: str = "", extension: str = "png", width: int = 0, height: int = 0, calculate_model_hash: bool = False, 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() model_name_real = model_name_str if model_name_str else model_name sampler_name_real = sampler_name_str if sampler_name_str else sampler_name scheduler_real = scheduler_str if scheduler_str else scheduler for image in images: i = 255.0 * image.cpu().numpy() img = Image.fromarray(np.clip(i, 0, 255).astype(np.uint8)) metadata = None model_hash = ( f"Model hash: {self.calculate_model_hash(model_name_real)}, " if calculate_model_hash else "" ) comment = ( f"{positive}\n" f"Negative prompt: {negative}\n" f"Steps: {steps}, " f"Sampler: {sampler_name_real}{''if scheduler_real == 'normal' else '_'+scheduler_real}, " 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}" f"Model: {Path(model_name_real).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" ) }, } ) piexif.insert(metadata, str(file)) results.append( {"filename": file.name, "subfolder": subfolder, "type": self.type} ) counter += 1 return {"ui": {"images": results}} @staticmethod def calculate_model_hash(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] class SDParameterGenerator: @classmethod def INPUT_TYPES(s): return { "required": { "ckpt_name": (folder_paths.get_filename_list("checkpoints"),), "config_name": ( ["disable"] + folder_paths.get_filename_list("configs"), {"default": "disable"}, ), }, "optional": { "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,), "width": ( "INT", {"default": 512, "min": 1, "max": MAX_RESOLUTION, "step": 8}, ), "height": ( "INT", {"default": 512, "min": 1, "max": MAX_RESOLUTION, "step": 8}, ), }, } RETURN_TYPES = ( "MODEL", "CLIP", "VAE", folder_paths.get_filename_list("checkpoints"), "INT", "INT", "FLOAT", comfy.samplers.KSampler.SAMPLERS, comfy.samplers.KSampler.SCHEDULERS, "INT", "INT", ) RETURN_NAMES = ( "MODEL", "CLIP", "VAE", "MODEL_NAME", "SEED", "STEPS", "CFG", "SAMPLER_NAME", "SCHEDULER", "WIDTH", "HEIGHT", ) FUNCTION = "generate_parameter" CATEGORY = "SD Prompt Reader" def generate_parameter( self, ckpt_name, config_name, seed, steps, cfg, sampler_name, scheduler, width, height, output_vae=True, output_clip=True, ): ckpt_path = folder_paths.get_full_path("checkpoints", ckpt_name) if config_name != "disable": config_path = folder_paths.get_full_path("configs", config_name) checkpoint = comfy.sd.load_checkpoint( config_path, ckpt_path, output_vae=True, output_clip=True, embedding_directory=folder_paths.get_folder_paths("embeddings"), ) else: checkpoint = comfy.sd.load_checkpoint_guess_config( ckpt_path, output_vae=True, output_clip=True, embedding_directory=folder_paths.get_folder_paths("embeddings"), )[:3] return checkpoint + ( ckpt_name, seed, steps, cfg, sampler_name, scheduler, width, height, ) class SDPromptMerger: @classmethod def INPUT_TYPES(s): return { "required": { "text_g": ( "STRING", {"default": "", "multiline": True, "forceInput": True}, ), "text_l": ( "STRING", {"default": "", "multiline": True, "forceInput": True}, ), }, } RETURN_TYPES = ("STRING",) FUNCTION = "merge_prompt" CATEGORY = "SD Prompt Reader" def merge_prompt(self, text_g, text_l): if text_l == "": return text_g return (text_g + "\n" + text_l,) class SDTypeConverter: @classmethod def INPUT_TYPES(s): return { "required": {}, "optional": { "model_name": (folder_paths.get_filename_list("checkpoints"),), "sampler_name": (comfy.samplers.KSampler.SAMPLERS,), "scheduler": (comfy.samplers.KSampler.SCHEDULERS,), }, } RETURN_TYPES = ( "STRING", "STRING", "STRING", ) RETURN_NAMES = ( "MODEL_NAME_STR", "SAMPLER_NAME_STR", "SCHEDULER_STR", ) FUNCTION = "convert_string" CATEGORY = "SD Prompt Reader" def convert_string( self, model_name: str = "", sampler_name: str = "", scheduler: str = "" ): return ( model_name, sampler_name, scheduler, ) NODE_CLASS_MAPPINGS = { "SDPromptReader": SDPromptReader, "SDPromptSaver": SDPromptSaver, "SDParameterGenerator": SDParameterGenerator, "SDPromptMerger": SDPromptMerger, "SDTypeConverter": SDTypeConverter, } NODE_DISPLAY_NAME_MAPPINGS = { "SDPromptReader": "SD Prompt Reader", "SDPromptSaver": "SD Prompt Saver", "SDParameterGenerator": "SD Parameter Generator", "SDPromptMerger": "SD Prompt Merger", "SDTypeConverter": "SD Type Converter", }