MdONeilsl
b6a2d1dd79
Introduce a complete GIMP/GEGL-inspired color module for ComfyUI. Adds: node declarations (nodes/*) and matching operator implementations (py/*) for many color/image operations (balance, temperature, hue/chroma, hue/saturation, saturation, exposure, levels, invert, auto adjustments, maps, posterize, desaturate, tone mapping, component extract/mix, etc.). Includes support libs (py/libs: color, process, blur), layer compositing (layers_blending), VRAM‑aware chunking, mask handling and alpha preservation. Adds web UI extension and widgets (web/*), package metadata, dynamic __init__.py loader, README update and .gitignore tweaks.
186 lines
6.6 KiB
Python
186 lines
6.6 KiB
Python
import torch
|
||
from typing import Tuple, Optional
|
||
import logging
|
||
|
||
from .libs.color import rgb_to_hsv, hsv_to_rgb, rgb_to_lch, lch_to_rgb, rgb_to_yuv, yuv_to_rgb
|
||
from .libs.process import compute_chunk_size
|
||
|
||
logger = logging.getLogger(__name__)
|
||
|
||
# ----------------------------------------------------------------------
|
||
# Core saturation algorithm – processes a single batch on one device
|
||
# ----------------------------------------------------------------------
|
||
def _saturation_core(
|
||
image: torch.Tensor, # (B, H, W, C) C = 3 or 4
|
||
scale: float,
|
||
color_space: str,
|
||
mask: Optional[torch.Tensor] # (B, H, W, C) or None
|
||
) -> torch.Tensor:
|
||
"""
|
||
Apply saturation scaling to an image batch (RGB or RGBA).
|
||
Handles exactly one device; no chunking, no fallback.
|
||
"""
|
||
|
||
# Early exit when no change is requested
|
||
if scale == 1.0:
|
||
return image
|
||
|
||
# Clamp input to valid colour range
|
||
img = image.clamp(0.0, 1.0)
|
||
|
||
# Separate alpha if present
|
||
if img.shape[-1] == 4:
|
||
alpha = img[..., 3:4]
|
||
rgb = img[..., :3]
|
||
else:
|
||
alpha = None
|
||
rgb = img
|
||
|
||
# ----- colour space conversion & scaling -----
|
||
if color_space == "Native":
|
||
h, s, v = rgb_to_hsv(rgb)
|
||
s = torch.clamp(s * scale, 0.0, 1.0)
|
||
result_rgb = hsv_to_rgb(h, s, v)
|
||
elif color_space == "CIE Lab/Lch":
|
||
l, c, h = rgb_to_lch(rgb)
|
||
c = c * scale
|
||
result_rgb = lch_to_rgb(l, c, h)
|
||
else: # CIE Yuv
|
||
y, u, v = rgb_to_yuv(rgb)
|
||
u = u * scale
|
||
v = v * scale
|
||
result_rgb = yuv_to_rgb(y, u, v)
|
||
|
||
result_rgb = torch.clamp(result_rgb, 0.0, 1.0)
|
||
|
||
# Reassemble RGBA if needed
|
||
if alpha is not None:
|
||
result = torch.cat([result_rgb, alpha], dim=-1)
|
||
else:
|
||
result = result_rgb
|
||
|
||
# Mask blending (applied to all channels, including alpha)
|
||
if mask is not None:
|
||
result = img * (1.0 - mask) + result * mask
|
||
|
||
return result
|
||
|
||
|
||
# ----------------------------------------------------------------------
|
||
# Public entry point – batch, device, VRAM management, error handling
|
||
# ----------------------------------------------------------------------
|
||
def exec_operator_color_saturation(
|
||
image: torch.Tensor,
|
||
scale: float,
|
||
color_space: str,
|
||
mask: Optional[torch.Tensor] = None
|
||
) -> Tuple[torch.Tensor]:
|
||
"""
|
||
Adjust the saturation of an image or batch of images.
|
||
|
||
Parameters
|
||
----------
|
||
image : torch.Tensor
|
||
Input image(s) in RGB or RGBA format.
|
||
Shape: (H, W, 3) or (H, W, 4) for a single image,
|
||
(B, H, W, 3) or (B, H, W, 4) for a batch.
|
||
scale : float
|
||
Saturation multiplier. 1.0 = no change, <1 reduces saturation,
|
||
>1 increases saturation.
|
||
color_space : str
|
||
Colour space in which saturation is scaled.
|
||
One of ``"Native"``, ``"CIE Lab/Lch"``, ``"CIE Yuv"``.
|
||
mask : torch.Tensor, optional
|
||
Optional mask defining where the adjustment is applied.
|
||
Shape must match the image batch size and spatial dimensions.
|
||
(B, H, W) or (B, 1, H, W) or (B, H, W, 1).
|
||
|
||
Returns
|
||
-------
|
||
Tuple[torch.Tensor]
|
||
A tuple containing the saturation‑adjusted image(s) in the same
|
||
format and shape as the input.
|
||
"""
|
||
# --- validate input shapes and types ---
|
||
if not isinstance(image, torch.Tensor):
|
||
raise TypeError("image must be a torch.Tensor")
|
||
if image.dim() not in (3, 4):
|
||
raise ValueError(f"image must be 3D (H,W,C) or 4D (B,H,W,C), got {image.dim()}D")
|
||
if image.shape[-1] not in (3, 4):
|
||
raise ValueError(f"image channels must be 3 or 4, got {image.shape[-1]}")
|
||
|
||
if not isinstance(scale, (int, float)):
|
||
raise TypeError("scale must be a float")
|
||
if color_space not in ("Native", "CIE Lab/Lch", "CIE Yuv"):
|
||
raise ValueError(f"unknown color_space '{color_space}'")
|
||
|
||
# Normalise to batch format (B, H, W, C)
|
||
single_input = image.dim() == 3
|
||
if single_input:
|
||
image = image.unsqueeze(0) # (1, H, W, C)
|
||
if mask is not None:
|
||
mask = mask.unsqueeze(0) # (1, …) whatever dim
|
||
|
||
batch_size, h, w, channels = image.shape
|
||
device = image.device
|
||
|
||
# --- mask handling (ensure (B, H, W, C) compatible) ---
|
||
if mask is not None:
|
||
mask = mask.to(device)
|
||
# bring to (B, H, W, 1)
|
||
if mask.dim() == 2: # (H, W)
|
||
mask = mask.unsqueeze(0).unsqueeze(-1)
|
||
elif mask.dim() == 3: # (B, H, W)
|
||
mask = mask.unsqueeze(-1)
|
||
elif mask.dim() == 4:
|
||
if mask.shape[1] == 1: # (B, 1, H, W)
|
||
mask = mask.permute(0, 2, 3, 1) # (B, H, W, 1)
|
||
elif mask.shape[-1] != 1:
|
||
raise ValueError("mask must have a single channel")
|
||
else:
|
||
raise ValueError("mask must be 2D, 3D or 4D")
|
||
|
||
if mask.shape[0] != batch_size:
|
||
raise ValueError(f"batch size mismatch: image {batch_size}, mask {mask.shape[0]}")
|
||
# Expand mask to full channel dimension
|
||
mask = mask.expand(-1, -1, -1, channels)
|
||
|
||
# --- VRAM optimised chunked processing (GPU only) ---
|
||
VRAM_CHUNK_SIZE = compute_chunk_size(image) # process 4 images at once on GPU to limit peak memory
|
||
|
||
def _process_chunked(img_batch, msk_batch):
|
||
"""Process a whole batch, potentially in chunks."""
|
||
if device.type == 'cuda' and batch_size > VRAM_CHUNK_SIZE:
|
||
results = []
|
||
for i in range(0, batch_size, VRAM_CHUNK_SIZE):
|
||
end = min(i + VRAM_CHUNK_SIZE, batch_size)
|
||
sub_img = img_batch[i:end]
|
||
sub_msk = msk_batch[i:end] if msk_batch is not None else None
|
||
results.append(_saturation_core(sub_img, scale, color_space, sub_msk))
|
||
# free GPU memory peaks
|
||
torch.cuda.empty_cache()
|
||
return torch.cat(results, dim=0)
|
||
else:
|
||
return _saturation_core(img_batch, scale, color_space, msk_batch)
|
||
|
||
# --- execute with GPU OOM fallback ---
|
||
try:
|
||
result = _process_chunked(image, mask)
|
||
except RuntimeError as e:
|
||
if 'out of memory' in str(e) and device.type == 'cuda':
|
||
logger.warning("CUDA out of memory – falling back to CPU")
|
||
torch.cuda.empty_cache()
|
||
cpu_image = image.cpu()
|
||
cpu_mask = mask.cpu() if mask is not None else None
|
||
result_cpu = _saturation_core(cpu_image, scale, color_space, cpu_mask)
|
||
result = result_cpu.to(device)
|
||
else:
|
||
raise
|
||
|
||
# Remove batch dimension if input was a single image
|
||
if single_input:
|
||
result = result.squeeze(0)
|
||
|
||
return (result,)
|
||
|