286 lines
10 KiB
Python
286 lines
10 KiB
Python
import json
|
|
import os
|
|
|
|
import numpy as np
|
|
import torch
|
|
import torch.nn.functional as F
|
|
from PIL import Image
|
|
from PIL.PngImagePlugin import PngInfo
|
|
|
|
import folder_paths
|
|
|
|
|
|
def _resize_image(image, height, width):
|
|
if image.shape[1] == height and image.shape[2] == width:
|
|
return image
|
|
nchw = image.movedim(-1, 1)
|
|
resized = F.interpolate(nchw, size=(height, width), mode="bilinear", align_corners=False)
|
|
return resized.movedim(1, -1)
|
|
|
|
|
|
def _resize_mask(mask, height, width):
|
|
if mask.shape[-2] == height and mask.shape[-1] == width:
|
|
return mask
|
|
resized = F.interpolate(mask.unsqueeze(1), size=(height, width), mode="bilinear", align_corners=False)
|
|
return resized[:, 0, :, :]
|
|
|
|
|
|
def _match_batch(tensor, batch_size):
|
|
if tensor.shape[0] == batch_size:
|
|
return tensor
|
|
if tensor.shape[0] == 1:
|
|
return tensor.repeat(batch_size, *([1] * (tensor.dim() - 1)))
|
|
if tensor.shape[0] > batch_size:
|
|
return tensor[:batch_size]
|
|
|
|
pad_count = batch_size - tensor.shape[0]
|
|
last = tensor[-1:].repeat(pad_count, *([1] * (tensor.dim() - 1)))
|
|
return torch.cat([tensor, last], dim=0)
|
|
|
|
|
|
def _normalize_mask(mask, batch_size, height, width, device, dtype):
|
|
if mask is None:
|
|
return None
|
|
if mask.dim() == 2:
|
|
mask = mask.unsqueeze(0)
|
|
mask = mask.to(device=device, dtype=dtype).clamp(0.0, 1.0)
|
|
mask = _match_batch(mask, batch_size)
|
|
return _resize_mask(mask, height, width)
|
|
|
|
|
|
def _split_rgb_alpha(image, mask=None, mask_mode="comfy_transparency_mask"):
|
|
image = image.clamp(0.0, 1.0)
|
|
rgb = image[..., :3]
|
|
|
|
batch_size, height, width = image.shape[:3]
|
|
mask = _normalize_mask(mask, batch_size, height, width, image.device, image.dtype)
|
|
if mask is not None:
|
|
if mask_mode == "comfy_transparency_mask":
|
|
alpha = 1.0 - mask
|
|
else:
|
|
alpha = mask
|
|
elif image.shape[-1] >= 4:
|
|
alpha = image[..., 3]
|
|
else:
|
|
alpha = torch.ones((batch_size, height, width), device=image.device, dtype=image.dtype)
|
|
|
|
return rgb, alpha.clamp(0.0, 1.0)
|
|
|
|
|
|
def _canvas_size(mode, bottom_h, bottom_w, top_h, top_w, offset_x, offset_y, custom_w, custom_h):
|
|
if mode == "top_size":
|
|
return top_h, top_w, -offset_x, -offset_y, 0, 0
|
|
if mode == "union_size":
|
|
min_x = min(0, offset_x)
|
|
min_y = min(0, offset_y)
|
|
max_x = max(bottom_w, offset_x + top_w)
|
|
max_y = max(bottom_h, offset_y + top_h)
|
|
return max_y - min_y, max_x - min_x, -min_x, -min_y, offset_x - min_x, offset_y - min_y
|
|
if mode == "custom_size":
|
|
return custom_h, custom_w, 0, 0, offset_x, offset_y
|
|
return bottom_h, bottom_w, 0, 0, offset_x, offset_y
|
|
|
|
|
|
def _place_layer(rgb, alpha, canvas_h, canvas_w, x, y):
|
|
batch_size = rgb.shape[0]
|
|
out_rgb = torch.zeros((batch_size, canvas_h, canvas_w, 3), device=rgb.device, dtype=rgb.dtype)
|
|
out_alpha = torch.zeros((batch_size, canvas_h, canvas_w), device=rgb.device, dtype=rgb.dtype)
|
|
|
|
src_h, src_w = rgb.shape[1], rgb.shape[2]
|
|
dst_x0 = max(x, 0)
|
|
dst_y0 = max(y, 0)
|
|
dst_x1 = min(x + src_w, canvas_w)
|
|
dst_y1 = min(y + src_h, canvas_h)
|
|
if dst_x0 >= dst_x1 or dst_y0 >= dst_y1:
|
|
return out_rgb, out_alpha
|
|
|
|
src_x0 = dst_x0 - x
|
|
src_y0 = dst_y0 - y
|
|
src_x1 = src_x0 + (dst_x1 - dst_x0)
|
|
src_y1 = src_y0 + (dst_y1 - dst_y0)
|
|
|
|
out_rgb[:, dst_y0:dst_y1, dst_x0:dst_x1, :] = rgb[:, src_y0:src_y1, src_x0:src_x1, :]
|
|
out_alpha[:, dst_y0:dst_y1, dst_x0:dst_x1] = alpha[:, src_y0:src_y1, src_x0:src_x1]
|
|
return out_rgb, out_alpha
|
|
|
|
|
|
def _normal_composite(bottom_rgb, bottom_alpha, top_rgb, top_alpha, top_opacity):
|
|
top_alpha = (top_alpha * top_opacity).clamp(0.0, 1.0)
|
|
out_alpha = top_alpha + bottom_alpha * (1.0 - top_alpha)
|
|
premultiplied = (
|
|
top_rgb * top_alpha.unsqueeze(-1)
|
|
+ bottom_rgb * bottom_alpha.unsqueeze(-1) * (1.0 - top_alpha).unsqueeze(-1)
|
|
)
|
|
out_rgb = torch.where(
|
|
out_alpha.unsqueeze(-1) > 1e-6,
|
|
premultiplied / out_alpha.clamp(min=1e-6).unsqueeze(-1),
|
|
torch.zeros_like(premultiplied),
|
|
)
|
|
return out_rgb.clamp(0.0, 1.0), out_alpha.clamp(0.0, 1.0)
|
|
|
|
|
|
def _alpha_for_save(image, alpha, alpha_mode):
|
|
batch_size, height, width = image.shape[:3]
|
|
alpha = _normalize_mask(alpha, batch_size, height, width, image.device, image.dtype)
|
|
if alpha is not None:
|
|
if alpha_mode == "comfy_transparency_mask":
|
|
return (1.0 - alpha).clamp(0.0, 1.0)
|
|
return alpha.clamp(0.0, 1.0)
|
|
if image.shape[-1] >= 4:
|
|
return image[..., 3].clamp(0.0, 1.0)
|
|
return torch.ones((batch_size, height, width), device=image.device, dtype=image.dtype)
|
|
|
|
|
|
class PhotoshopNormalLayerComposite:
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
"required": {
|
|
"bottom_image": ("IMAGE",),
|
|
"top_image": ("IMAGE",),
|
|
"top_x": ("INT", {"default": 0, "min": -16384, "max": 16384, "step": 1}),
|
|
"top_y": ("INT", {"default": 0, "min": -16384, "max": 16384, "step": 1}),
|
|
"top_opacity": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 1.0, "step": 0.01}),
|
|
"canvas_mode": (["bottom_size", "top_size", "union_size", "custom_size"], {"default": "bottom_size"}),
|
|
"canvas_width": ("INT", {"default": 512, "min": 1, "max": 16384, "step": 1}),
|
|
"canvas_height": ("INT", {"default": 512, "min": 1, "max": 16384, "step": 1}),
|
|
"mask_mode": (
|
|
["comfy_transparency_mask", "alpha_mask"],
|
|
{"default": "comfy_transparency_mask"},
|
|
),
|
|
},
|
|
"optional": {
|
|
"bottom_alpha": ("MASK",),
|
|
"top_alpha": ("MASK",),
|
|
},
|
|
}
|
|
|
|
RETURN_TYPES = ("IMAGE", "IMAGE", "MASK", "MASK")
|
|
RETURN_NAMES = ("rgba_image", "rgb_preview", "alpha", "comfy_transparency_mask")
|
|
FUNCTION = "composite"
|
|
CATEGORY = "image/composite"
|
|
|
|
def composite(
|
|
self,
|
|
bottom_image,
|
|
top_image,
|
|
top_x,
|
|
top_y,
|
|
top_opacity,
|
|
canvas_mode,
|
|
canvas_width,
|
|
canvas_height,
|
|
mask_mode,
|
|
bottom_alpha=None,
|
|
top_alpha=None,
|
|
):
|
|
bottom_rgb, bottom_a = _split_rgb_alpha(bottom_image, bottom_alpha, mask_mode)
|
|
top_rgb, top_a = _split_rgb_alpha(top_image, top_alpha, mask_mode)
|
|
|
|
batch_size = max(bottom_rgb.shape[0], top_rgb.shape[0])
|
|
bottom_rgb = _match_batch(bottom_rgb, batch_size)
|
|
bottom_a = _match_batch(bottom_a, batch_size)
|
|
top_rgb = _match_batch(top_rgb, batch_size)
|
|
top_a = _match_batch(top_a, batch_size)
|
|
|
|
bottom_h, bottom_w = bottom_rgb.shape[1], bottom_rgb.shape[2]
|
|
top_h, top_w = top_rgb.shape[1], top_rgb.shape[2]
|
|
canvas_h, canvas_w, bottom_x, bottom_y, placed_top_x, placed_top_y = _canvas_size(
|
|
canvas_mode,
|
|
bottom_h,
|
|
bottom_w,
|
|
top_h,
|
|
top_w,
|
|
int(top_x),
|
|
int(top_y),
|
|
int(canvas_width),
|
|
int(canvas_height),
|
|
)
|
|
|
|
placed_bottom_rgb, placed_bottom_a = _place_layer(bottom_rgb, bottom_a, canvas_h, canvas_w, bottom_x, bottom_y)
|
|
placed_top_rgb, placed_top_a = _place_layer(top_rgb, top_a, canvas_h, canvas_w, placed_top_x, placed_top_y)
|
|
out_rgb, out_alpha = _normal_composite(
|
|
placed_bottom_rgb,
|
|
placed_bottom_a,
|
|
placed_top_rgb,
|
|
placed_top_a,
|
|
float(top_opacity),
|
|
)
|
|
|
|
rgba = torch.cat([out_rgb, out_alpha.unsqueeze(-1)], dim=-1)
|
|
comfy_mask = (1.0 - out_alpha).clamp(0.0, 1.0)
|
|
return (rgba, out_rgb, out_alpha, comfy_mask)
|
|
|
|
|
|
class SaveTransparentPNG:
|
|
def __init__(self):
|
|
self.output_dir = folder_paths.get_output_directory()
|
|
self.type = "output"
|
|
self.prefix_append = ""
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
"required": {
|
|
"image": ("IMAGE",),
|
|
"filename_prefix": ("STRING", {"default": "ComfyUI_transparent"}),
|
|
"alpha_mode": (
|
|
["auto_rgba_or_alpha_mask", "comfy_transparency_mask"],
|
|
{"default": "auto_rgba_or_alpha_mask"},
|
|
),
|
|
"compress_level": ("INT", {"default": 4, "min": 0, "max": 9, "step": 1}),
|
|
},
|
|
"optional": {
|
|
"alpha": ("MASK",),
|
|
},
|
|
"hidden": {"prompt": "PROMPT", "extra_pnginfo": "EXTRA_PNGINFO"},
|
|
}
|
|
|
|
RETURN_TYPES = ()
|
|
FUNCTION = "save"
|
|
OUTPUT_NODE = True
|
|
CATEGORY = "image/save"
|
|
|
|
def save(self, image, filename_prefix, alpha_mode, compress_level, alpha=None, prompt=None, extra_pnginfo=None):
|
|
image = image.clamp(0.0, 1.0)
|
|
alpha_mask_mode = "comfy_transparency_mask" if alpha_mode == "comfy_transparency_mask" else "alpha_mask"
|
|
out_alpha = _alpha_for_save(image, alpha, alpha_mask_mode)
|
|
rgb = image[..., :3]
|
|
rgba = torch.cat([rgb, out_alpha.unsqueeze(-1)], dim=-1)
|
|
|
|
full_output_folder, filename, counter, subfolder, filename_prefix = folder_paths.get_save_image_path(
|
|
filename_prefix + self.prefix_append,
|
|
self.output_dir,
|
|
rgba[0].shape[1],
|
|
rgba[0].shape[0],
|
|
)
|
|
|
|
metadata = PngInfo()
|
|
if prompt is not None:
|
|
metadata.add_text("prompt", json.dumps(prompt))
|
|
if extra_pnginfo is not None:
|
|
for key, value in extra_pnginfo.items():
|
|
metadata.add_text(key, json.dumps(value))
|
|
|
|
results = []
|
|
for batch_number, frame in enumerate(rgba):
|
|
pixels = (255.0 * frame.cpu().numpy()).clip(0, 255).astype(np.uint8)
|
|
img = Image.fromarray(pixels, mode="RGBA")
|
|
file = f"{filename}_{counter:05}_.png"
|
|
img.save(os.path.join(full_output_folder, file), pnginfo=metadata, compress_level=int(compress_level))
|
|
results.append({"filename": file, "subfolder": subfolder, "type": self.type})
|
|
counter += 1
|
|
|
|
return {"ui": {"images": results}}
|
|
|
|
|
|
NODE_CLASS_MAPPINGS = {
|
|
"PhotoshopNormalLayerComposite": PhotoshopNormalLayerComposite,
|
|
"SaveTransparentPNG": SaveTransparentPNG,
|
|
}
|
|
|
|
NODE_DISPLAY_NAME_MAPPINGS = {
|
|
"PhotoshopNormalLayerComposite": "PS Normal Layer Composite (RGBA)",
|
|
"SaveTransparentPNG": "Save Transparent PNG (RGBA)",
|
|
}
|