MdONeilsl-ComfyUI-MdColorMod/py/color_balance.py

import torch
from typing import Tuple, Optional

@torch.no_grad()                     # 1) no gradient tracking → saves VRAM & time
def exec_operation_color_balance(
    image: torch.Tensor,             # shape: (B, H, W, C) with C = 3 or 4
    mask: Optional[torch.Tensor],    # can be (H,W), (B,H,W), (B,H,W,1), ...
    shadow_adj: Tuple[float, float, float],
    midtones_adj: Tuple[float, float, float],
    highlights_adj: Tuple[float, float, float],
    preserve_luminosity: bool
) -> torch.Tensor:
    """
    Apply GIMP‑style color balance.
    Each adjustment tuple = (Cyan‑Red, Magenta‑Green, Yellow‑Blue) in [-100, 100].

    Improvements over original:
      - Robust batch‑mask handling (2D mask → broadcast to whole batch)
      - Full RGB / RGBA support (alpha channel is never color‑shifted)
      - Input validation with clear error messages
      - `torch.no_grad()` + single memory allocation → VRAM optimised
      - Works on any device (CPU / CUDA / MPS)
    """

    # ----------------------------------------------------------------------
    # 1. Input validation & preparation
    # ----------------------------------------------------------------------
    if image.dim() != 4:
        raise ValueError(f"image must be 4D (B,H,W,C), got {image.dim()}D")
    B, H, W, C = image.shape
    if C not in (3, 4):
        raise ValueError(f"only 3 (RGB) or 4 (RGBA) channels allowed, got {C}")
    for adj in (shadow_adj, midtones_adj, highlights_adj):
        if len(adj) != 3:
            raise ValueError("each adjustment must be a sequence of 3 floats")
    device = image.device

    # Work on a float copy (same device)
    img = image.to(dtype=torch.float32, copy=True)   # single allocation
    rgb = img[..., :3]                               # view – no copy

    # ----------------------------------------------------------------------
    # 2. Prepare mask  →  broadcastable (B, H, W, 1)
    # ----------------------------------------------------------------------
    mask_tensor: Optional[torch.Tensor] = None
    if mask is not None:
        m = mask.to(device=device, dtype=torch.float32)
        # unify shapes
        if m.dim() == 2:                     # (H, W)  → apply to all images
            m = m[None, ..., None]           # (1, H, W, 1)
            m = m.expand(B, -1, -1, -1)      # (B, H, W, 1)
        elif m.dim() == 3:
            if m.shape[-1] == 1:             # (H, W, 1) or (B, H, W, 1)?
                if m.shape[0] == B:
                    m = m                   # (B, H, W, 1) – already fine
                else:
                    m = m[None]             # (1, H, W, 1)
                    m = m.expand(B, -1, -1, -1)
            else:                            # (B, H, W)
                m = m.unsqueeze(-1)          # (B, H, W, 1)
        elif m.dim() == 4:
            if m.shape[0] == 1 and B > 1:
                m = m.expand(B, -1, -1, -1) # broadcast batch dimension
        else:
            raise ValueError(f"mask shape {mask.shape} not supported")
        # safety check
        if m.shape[:3] != (B, H, W):
            raise ValueError(
                f"mask spatial dimensions must match image. "
                f"Expected (B,H,W)=({B},{H},{W}), got {m.shape[:3]}"
            )
        mask_tensor = m

    # ----------------------------------------------------------------------
    # 3. Luma & transfer‑function weights  (Rec. 709)
    # ----------------------------------------------------------------------
    luma_w = torch.tensor([0.299, 0.587, 0.114], device=device, dtype=torch.float32)
    luma = (rgb * luma_w).sum(dim=-1, keepdim=True)          # (B, H, W, 1)

    def smoothstep(edge0, edge1, x):
        t = torch.clamp((x - edge0) / (edge1 - edge0), 0.0, 1.0)
        return t * t * (3.0 - 2.0 * t)

    shadow_w     = smoothstep(0.5, 0.0, luma)                # 1 at 0, 0 at 0.5
    midtones_w   = 4.0 * luma * (1.0 - luma)                 # peak 1 at 0.5
    highlights_w = smoothstep(0.5, 1.0, luma)                # 0 at 0.5, 1 at 1.0

    # ----------------------------------------------------------------------
    # 4. Normalised adjustments → (1,1,1,3)
    # ----------------------------------------------------------------------
    def to_tensor(adj):
        return torch.tensor(adj, device=device, dtype=torch.float32).view(1,1,1,3) / 100.0

    delta = (
        shadow_w     * to_tensor(shadow_adj) +
        midtones_w   * to_tensor(midtones_adj) +
        highlights_w * to_tensor(highlights_adj)
    )                                                         # (B, H, W, 3)

    # ----------------------------------------------------------------------
    # 5. Apply delta (with or without mask)
    # ----------------------------------------------------------------------
    if mask_tensor is not None:
        blended = rgb + delta * mask_tensor                    # only RGB
    else:
        blended = rgb + delta

    # Build output (keep alpha untouched)
    if C == 4:
        alpha = img[..., 3:4]
        out = torch.cat([blended, alpha], dim=-1)
    else:
        out = blended

    # ----------------------------------------------------------------------
    # 6. Optional luminance preservation (RGB only)
    # ----------------------------------------------------------------------
    if preserve_luminosity:
        new_luma = (out[..., :3] * luma_w).sum(dim=-1, keepdim=True)
        eps = 1e-8
        scale = luma / (new_luma + eps)
        out[..., :3] = out[..., :3] * scale

    # ----------------------------------------------------------------------
    # 7. Final clamp (entire image, including alpha, exactly like original)
    # ----------------------------------------------------------------------
    out.clamp_(0.0, 1.0)

    return out