MdONeilsl-ComfyUI-MdColorMod/py/color_posterize.py

import torch
from typing import Tuple, Optional

from .libs.color import srgb_to_linear_vec, linear_to_srgb_vec
from .libs.process import compute_chunk_size

# ============================================================
# GEGL-style posterization core
# ============================================================

def _posterize_channel_gegl(x: torch.Tensor, levels: int) -> torch.Tensor:
    """
    GEGL-style uniform quantizer.
    Equivalent to:
        floor(x * (levels - 1) + 0.5) / (levels - 1)
    """
    levels = max(2, int(levels))
    scale = levels - 1
    return torch.floor(x * scale + 0.5) / scale


# ============================================================
# Main operator – improved with batch chunking, alpha preservation,
# error handling, and early exit.
# ============================================================

def exec_operator_color_posterize(
    image: torch.Tensor,
    levels: int,
    mask: Optional[torch.Tensor] = None,
    linear: bool = False,   # GEGL toggle: operate in linear light
) -> Tuple[torch.Tensor]:
    """
    Apply GEGL‑style posterisation to a batch of images.

    The input *image* is a float tensor of shape ``(B, H, W, C)`` with
    ``C == 3`` (RGB) or ``C == 4`` (RGBA).  *levels* controls the number
    of discrete levels per colour channel (minimum 2).  An optional *mask*
    tensor of shape ``(H, W)`` (single mask) or ``(B, H, W)`` (per‑image
    mask) defines the per‑pixel strength of the effect: ``0.0`` leaves the
    pixel unchanged, ``1.0`` applies full posterisation.

    When the input contains an alpha channel (``C == 4``), the alpha
    channel is **preserved unchanged**; only the RGB channels are
    posterised and blended with the mask.

    The function validates input shapes, performs an early exit for empty
    batches, and processes the batch in GPU‑friendly chunks as determined
    by ``compute_chunk_size``.  No spatial tiling is performed.

    Parameters
    ----------
    image : torch.Tensor
        Batch of images (B,H,W,C), float.
    levels : int
        Number of posterisation levels (2‑256).
    mask : torch.Tensor or None
        Optional strength mask (H,W) or (B,H,W).
    linear : bool
        If ``True``, the effect is applied in linear light
        (sRGB ↔ linear conversion is used).

    Returns
    -------
    Tuple[torch.Tensor]
        A 1‑tuple containing the posterised image batch.
    """

    # ---------- input validation ----------
    if not isinstance(image, torch.Tensor):
        raise TypeError(f"image must be a torch.Tensor, got {type(image)}")
    if image.dim() != 4:
        raise ValueError(f"image must be 4D (B,H,W,C), got shape {image.shape}")
    C = image.shape[-1]
    if C not in (1, 3, 4):
        raise ValueError(f"Unsupported number of channels: {C}. Only 1, 3, or 4 are supported")
    if image.dtype not in (torch.float32, torch.float64, torch.float16):
        raise TypeError(f"image dtype must be float, got {image.dtype}")

    levels = max(2, int(levels))   # ensure >= 2

    if mask is not None:
        if not isinstance(mask, torch.Tensor):
            raise TypeError(f"mask must be a torch.Tensor or None, got {type(mask)}")
        if mask.dim() == 2:
            pass  # (H,W) will be broadcast
        elif mask.dim() == 3:
            if mask.shape[0] != image.shape[0]:
                raise ValueError(
                    f"Batch dimension mismatch: image batch {image.shape[0]}, "
                    f"mask batch {mask.shape[0]}"
                )
        else:
            raise ValueError(f"mask must be 2D (H,W) or 3D (B,H,W), got shape {mask.shape}")
        if mask.device != image.device:
            raise ValueError("image and mask must reside on the same device")

    # ---------- early exit for empty batches ----------
    if image.shape[0] == 0:
        return (image,)

    # ---------- alpha channel handling ----------
    has_alpha = (C == 4)
    if has_alpha:
        rgb = image[..., :3]
        alpha = image[..., 3:4]          # keep as (B,H,W,1)
    else:
        rgb = image
        alpha = None

    B = rgb.shape[0]

    # ---------- determine chunk size (VRAM‑aware) ----------
    # The helper considers both image and mask sizes.
    chunk_size = compute_chunk_size(image, mask)
    chunk_size = max(1, min(chunk_size, B))   # clamp to valid range

    out_chunks = []

    # ---------- chunked processing ----------
    for start in range(0, B, chunk_size):
        end = min(start + chunk_size, B)

        chunk_rgb = rgb[start:end].clone()   # clone to avoid in‑place issues
        # Prepare mask chunk
        chunk_mask = None
        if mask is not None:
            if mask.ndim == 3:          # per‑image masks
                chunk_mask = mask[start:end]
            else:                        # single (H,W) mask, reuse for all images
                chunk_mask = mask

        # ---- posterisation (per‑channel) ----
        x = torch.clamp(chunk_rgb, 0.0, 1.0)

        if linear:
            x = srgb_to_linear_vec(x)

        x_p = _posterize_channel_gegl(x, levels)

        if linear:
            x_p = linear_to_srgb_vec(x_p)

        x_p = torch.clamp(x_p, 0.0, 1.0)

        # ---- mask blending (only on RGB) ----
        if chunk_mask is not None:
            # Ensure mask has shape (B',H,W,1) for broadcasting
            if chunk_mask.ndim == 2:
                chunk_mask = chunk_mask.unsqueeze(0).unsqueeze(-1)  # (1,H,W,1)
            else:
                chunk_mask = chunk_mask.unsqueeze(-1)               # (B',H,W,1)
            chunk_mask = chunk_mask.to(dtype=x.dtype, device=x.device)
            x_p = chunk_rgb * (1.0 - chunk_mask) + x_p * chunk_mask

        # ---- reassemble alpha if present ----
        if has_alpha:
            chunk_out = torch.cat([x_p, alpha[start:end]], dim=-1)
        else:
            chunk_out = x_p

        out_chunks.append(chunk_out)

    # ---------- concatenate chunks ----------
    final_image = torch.cat(out_chunks, dim=0)

    return (final_image,)