136 lines
6.0 KiB
Python
136 lines
6.0 KiB
Python
|
|
import os, json, io, base64
|
|
|
|
import numpy as np
|
|
from PIL import Image
|
|
from PIL.PngImagePlugin import PngInfo
|
|
|
|
import folder_paths
|
|
from comfy.cli_args import args
|
|
|
|
from .helpers.imgutil import scale_tensor_image
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
# ==============================================================================
|
|
# adpated from ComfyUI SaveImage node
|
|
# ==============================================================================
|
|
class PseudoSaveImageWithEmbeddedMasks:
|
|
def __init__(self):
|
|
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": {
|
|
"imgs": ("IMAGE", {"tooltip": "The images to save."}),
|
|
"filename_prefix": ("STRING", {"default": "pseudocomfy", "tooltip": "The prefix for the file to save. This may include formatting information such as %date:yyyy-MM-dd% or %Empty Latent Image.width% to include values from nodes."})
|
|
},
|
|
"optional": {
|
|
"msks": ("MASK", {"tooltip": "The masks to embed in the image."}), # List of mask tensors ([1, H, W])
|
|
},
|
|
"hidden": {
|
|
"prompt": "PROMPT", "extra_pnginfo": "EXTRA_PNGINFO"
|
|
},
|
|
}
|
|
|
|
INPUT_IS_LIST = True
|
|
|
|
RETURN_TYPES = ()
|
|
FUNCTION = "save_images"
|
|
|
|
OUTPUT_NODE = True
|
|
|
|
CATEGORY = "Pseudocomfy/IO"
|
|
DESCRIPTION = "Saves the input images with embedded masks as extra channels."
|
|
|
|
def save_images(self, imgs, filename_prefix="pseudocomfy", msks=[], prompt=None, extra_pnginfo=None):
|
|
# given that INPUT_IS_LIST, msks is a list of masks that should be embedded in each image in each batch
|
|
# given that INPUT_IS_LIST, imgs is a list of image batches
|
|
image_batches = imgs
|
|
typical_image_shape = image_batches[0][0].shape
|
|
|
|
# if inputs that are meant to be singletons are a list, use the first element
|
|
if isinstance(filename_prefix, list) and len(filename_prefix)>0: filename_prefix = filename_prefix[0]
|
|
if isinstance(prompt, list) and len(prompt)>0: prompt = prompt[0]
|
|
if isinstance(extra_pnginfo, list) and len(extra_pnginfo)>0: extra_pnginfo = extra_pnginfo[0]
|
|
|
|
print(f"[pseudocomfy] SaveImageWithEmbeddedMasks")
|
|
print(f"\tfilename_prefix = {filename_prefix}")
|
|
print(f"\ttypical_image_shape = {tuple(typical_image_shape)}")
|
|
print(f"\timage_batches count = {len(image_batches)}")
|
|
for i, img_batch in enumerate(image_batches):
|
|
print(f"\timg_batch[{i}] shape = {tuple(img_batch.shape)}")
|
|
print(f"\tmsks count = {len(msks)}")
|
|
for i, msk in enumerate(msks):
|
|
print(f"\tmsks[{i}] shape = {tuple(msk.shape)}")
|
|
|
|
|
|
# Ensure all masks have shape [1, width, height]
|
|
|
|
for i in range(len(msks)):
|
|
mask = msks[i]
|
|
# Check mask has three dimensions and the first dimension is 1
|
|
if not (isinstance(mask.shape, tuple) and len(mask.shape) == 3 and mask.shape[0] == 1):
|
|
raise ValueError(f"Mask at index {i} must have shape [1, width, height], got {mask.shape}")
|
|
|
|
# scale the mask to the desired width and height if necessary
|
|
if mask.shape[-2] != typical_image_shape[0] or mask.shape[-1] != typical_image_shape[1]:
|
|
msks[i] = scale_tensor_image(mask, typical_image_shape[1], typical_image_shape[0]) # note that for scale_tensor_image, the order of width and height is reversed from order of tensor shape
|
|
|
|
|
|
print(f"\tmsks count = {len(msks)}")
|
|
for i, msk in enumerate(msks):
|
|
print(f"\tmsks[{i}] shape = {tuple(msk.shape)}")
|
|
|
|
|
|
filename_prefix += self.prefix_append
|
|
full_output_folder, filename, counter, subfolder, filename_prefix = folder_paths.get_save_image_path(filename_prefix, self.output_dir, typical_image_shape[1], typical_image_shape[0])
|
|
results = list()
|
|
|
|
for img_batch in image_batches:
|
|
for (batch_number, image) in enumerate(img_batch):
|
|
i = 255. * image.cpu().numpy()
|
|
image = Image.fromarray(np.clip(i, 0, 255).astype(np.uint8))
|
|
metadata = None
|
|
if not args.disable_metadata:
|
|
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]))
|
|
|
|
# Embed ALL masks as base64 PNGs in metadata
|
|
for mask_idx, mask in enumerate(msks):
|
|
mask_np = mask.squeeze().cpu().numpy()
|
|
mask_img = Image.fromarray(np.clip(255 * mask_np, 0, 255).astype(np.uint8), mode="L")
|
|
buf = io.BytesIO()
|
|
mask_img.save(buf, format="PNG")
|
|
b64 = base64.b64encode(buf.getvalue()).decode("utf-8")
|
|
mask_json = {
|
|
"format": "PNG",
|
|
"encoding": "base64",
|
|
"size": list(mask_img.size),
|
|
"data": b64
|
|
}
|
|
metadata.add_text(f"pseudocomfy_mask_{mask_idx:02d}", json.dumps(mask_json))
|
|
|
|
filename_with_batch_num = filename.replace("%batch_num%", str(batch_number))
|
|
file = f"{filename_with_batch_num}_{counter:05}_.png"
|
|
image.save(os.path.join(full_output_folder, file), pnginfo=metadata, compress_level=self.compress_level)
|
|
results.append({
|
|
"filename": file,
|
|
"subfolder": subfolder,
|
|
"type": self.type
|
|
})
|
|
counter += 1
|
|
|
|
return { "ui": { "images": results } } |