Add SD Prompt Saver
This commit is contained in:
Rhys Yang
@@ -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",
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user