| import numpy as np |
| import torch |
| import torch.nn.functional as F |
| from PIL import Image |
| import math |
|
|
| import comfy.utils |
|
|
|
|
| class Blend: |
| def __init__(self): |
| pass |
|
|
| @classmethod |
| def INPUT_TYPES(s): |
| return { |
| "required": { |
| "image1": ("IMAGE",), |
| "image2": ("IMAGE",), |
| "blend_factor": ("FLOAT", { |
| "default": 0.5, |
| "min": 0.0, |
| "max": 1.0, |
| "step": 0.01 |
| }), |
| "blend_mode": (["normal", "multiply", "screen", "overlay", "soft_light"],), |
| }, |
| } |
|
|
| RETURN_TYPES = ("IMAGE",) |
| FUNCTION = "blend_images" |
|
|
| CATEGORY = "image/postprocessing" |
|
|
| def blend_images(self, image1: torch.Tensor, image2: torch.Tensor, blend_factor: float, blend_mode: str): |
| if image1.shape != image2.shape: |
| image2 = image2.permute(0, 3, 1, 2) |
| image2 = comfy.utils.common_upscale(image2, image1.shape[2], image1.shape[1], upscale_method='bicubic', crop='center') |
| image2 = image2.permute(0, 2, 3, 1) |
|
|
| blended_image = self.blend_mode(image1, image2, blend_mode) |
| blended_image = image1 * (1 - blend_factor) + blended_image * blend_factor |
| blended_image = torch.clamp(blended_image, 0, 1) |
| return (blended_image,) |
|
|
| def blend_mode(self, img1, img2, mode): |
| if mode == "normal": |
| return img2 |
| elif mode == "multiply": |
| return img1 * img2 |
| elif mode == "screen": |
| return 1 - (1 - img1) * (1 - img2) |
| elif mode == "overlay": |
| return torch.where(img1 <= 0.5, 2 * img1 * img2, 1 - 2 * (1 - img1) * (1 - img2)) |
| elif mode == "soft_light": |
| return torch.where(img2 <= 0.5, img1 - (1 - 2 * img2) * img1 * (1 - img1), img1 + (2 * img2 - 1) * (self.g(img1) - img1)) |
| else: |
| raise ValueError(f"Unsupported blend mode: {mode}") |
|
|
| def g(self, x): |
| return torch.where(x <= 0.25, ((16 * x - 12) * x + 4) * x, torch.sqrt(x)) |
|
|
| def gaussian_kernel(kernel_size: int, sigma: float, device=None): |
| x, y = torch.meshgrid(torch.linspace(-1, 1, kernel_size, device=device), torch.linspace(-1, 1, kernel_size, device=device), indexing="ij") |
| d = torch.sqrt(x * x + y * y) |
| g = torch.exp(-(d * d) / (2.0 * sigma * sigma)) |
| return g / g.sum() |
|
|
| class Blur: |
| def __init__(self): |
| pass |
|
|
| @classmethod |
| def INPUT_TYPES(s): |
| return { |
| "required": { |
| "image": ("IMAGE",), |
| "blur_radius": ("INT", { |
| "default": 1, |
| "min": 1, |
| "max": 31, |
| "step": 1 |
| }), |
| "sigma": ("FLOAT", { |
| "default": 1.0, |
| "min": 0.1, |
| "max": 10.0, |
| "step": 0.1 |
| }), |
| }, |
| } |
|
|
| RETURN_TYPES = ("IMAGE",) |
| FUNCTION = "blur" |
|
|
| CATEGORY = "image/postprocessing" |
|
|
| def blur(self, image: torch.Tensor, blur_radius: int, sigma: float): |
| if blur_radius == 0: |
| return (image,) |
|
|
| batch_size, height, width, channels = image.shape |
|
|
| kernel_size = blur_radius * 2 + 1 |
| kernel = gaussian_kernel(kernel_size, sigma, device=image.device).repeat(channels, 1, 1).unsqueeze(1) |
|
|
| image = image.permute(0, 3, 1, 2) |
| padded_image = F.pad(image, (blur_radius,blur_radius,blur_radius,blur_radius), 'reflect') |
| blurred = F.conv2d(padded_image, kernel, padding=kernel_size // 2, groups=channels)[:,:,blur_radius:-blur_radius, blur_radius:-blur_radius] |
| blurred = blurred.permute(0, 2, 3, 1) |
|
|
| return (blurred,) |
|
|
| class Quantize: |
| def __init__(self): |
| pass |
|
|
| @classmethod |
| def INPUT_TYPES(s): |
| return { |
| "required": { |
| "image": ("IMAGE",), |
| "colors": ("INT", { |
| "default": 256, |
| "min": 1, |
| "max": 256, |
| "step": 1 |
| }), |
| "dither": (["none", "floyd-steinberg"],), |
| }, |
| } |
|
|
| RETURN_TYPES = ("IMAGE",) |
| FUNCTION = "quantize" |
|
|
| CATEGORY = "image/postprocessing" |
|
|
| def quantize(self, image: torch.Tensor, colors: int = 256, dither: str = "FLOYDSTEINBERG"): |
| batch_size, height, width, _ = image.shape |
| result = torch.zeros_like(image) |
|
|
| dither_option = Image.Dither.FLOYDSTEINBERG if dither == "floyd-steinberg" else Image.Dither.NONE |
|
|
| for b in range(batch_size): |
| tensor_image = image[b] |
| img = (tensor_image * 255).to(torch.uint8).numpy() |
| pil_image = Image.fromarray(img, mode='RGB') |
|
|
| palette = pil_image.quantize(colors=colors) |
| quantized_image = pil_image.quantize(colors=colors, palette=palette, dither=dither_option) |
|
|
| quantized_array = torch.tensor(np.array(quantized_image.convert("RGB"))).float() / 255 |
| result[b] = quantized_array |
|
|
| return (result,) |
|
|
| class Sharpen: |
| def __init__(self): |
| pass |
|
|
| @classmethod |
| def INPUT_TYPES(s): |
| return { |
| "required": { |
| "image": ("IMAGE",), |
| "sharpen_radius": ("INT", { |
| "default": 1, |
| "min": 1, |
| "max": 31, |
| "step": 1 |
| }), |
| "sigma": ("FLOAT", { |
| "default": 1.0, |
| "min": 0.1, |
| "max": 10.0, |
| "step": 0.1 |
| }), |
| "alpha": ("FLOAT", { |
| "default": 1.0, |
| "min": 0.0, |
| "max": 5.0, |
| "step": 0.1 |
| }), |
| }, |
| } |
|
|
| RETURN_TYPES = ("IMAGE",) |
| FUNCTION = "sharpen" |
|
|
| CATEGORY = "image/postprocessing" |
|
|
| def sharpen(self, image: torch.Tensor, sharpen_radius: int, sigma:float, alpha: float): |
| if sharpen_radius == 0: |
| return (image,) |
|
|
| batch_size, height, width, channels = image.shape |
|
|
| kernel_size = sharpen_radius * 2 + 1 |
| kernel = gaussian_kernel(kernel_size, sigma) * -(alpha*10) |
| center = kernel_size // 2 |
| kernel[center, center] = kernel[center, center] - kernel.sum() + 1.0 |
| kernel = kernel.repeat(channels, 1, 1).unsqueeze(1) |
|
|
| tensor_image = image.permute(0, 3, 1, 2) |
| tensor_image = F.pad(tensor_image, (sharpen_radius,sharpen_radius,sharpen_radius,sharpen_radius), 'reflect') |
| sharpened = F.conv2d(tensor_image, kernel, padding=center, groups=channels)[:,:,sharpen_radius:-sharpen_radius, sharpen_radius:-sharpen_radius] |
| sharpened = sharpened.permute(0, 2, 3, 1) |
|
|
| result = torch.clamp(sharpened, 0, 1) |
|
|
| return (result,) |
|
|
| class ImageScaleToTotalPixels: |
| upscale_methods = ["nearest-exact", "bilinear", "area", "bicubic", "lanczos"] |
| crop_methods = ["disabled", "center"] |
|
|
| @classmethod |
| def INPUT_TYPES(s): |
| return {"required": { "image": ("IMAGE",), "upscale_method": (s.upscale_methods,), |
| "megapixels": ("FLOAT", {"default": 1.0, "min": 0.01, "max": 16.0, "step": 0.01}), |
| }} |
| RETURN_TYPES = ("IMAGE",) |
| FUNCTION = "upscale" |
|
|
| CATEGORY = "image/upscaling" |
|
|
| def upscale(self, image, upscale_method, megapixels): |
| samples = image.movedim(-1,1) |
| total = int(megapixels * 1024 * 1024) |
|
|
| scale_by = math.sqrt(total / (samples.shape[3] * samples.shape[2])) |
| width = round(samples.shape[3] * scale_by) |
| height = round(samples.shape[2] * scale_by) |
|
|
| s = comfy.utils.common_upscale(samples, width, height, upscale_method, "disabled") |
| s = s.movedim(1,-1) |
| return (s,) |
|
|
| NODE_CLASS_MAPPINGS = { |
| "ImageBlend": Blend, |
| "ImageBlur": Blur, |
| "ImageQuantize": Quantize, |
| "ImageSharpen": Sharpen, |
| "ImageScaleToTotalPixels": ImageScaleToTotalPixels, |
| } |
|
|
