MdONeilsl-ComfyUI-MdColorMod/py/color_exposure.py

import torch

from .libs.process import compute_chunk_size

def exec_operator_color_exposure(
    image: torch.Tensor,
    black_level: float,
    exposure: float,
    mask: torch.Tensor = None,
) -> torch.Tensor:
    """
    Apply a GEGL‑style exposure adjustment to an image batch, with optional
    mask blending and full alpha‑channel preservation.

    Exposure is modeled as:

        white = 2^(-exposure)
        gain  = 1 / (white - black_level)     (clamped for numerical stability)
        result = (image - black_level) * gain

    When a mask is provided, the adjusted image is blended with the original
    using the mask as the per‑pixel mixing weight:

        output = image * (1 - mask) + result * mask

    Args:
        image (torch.Tensor): Input image tensor of shape ``(B, H, W, C)``,
                              with C ∈ {1, 3, 4}.
        black_level (float): Black level offset (clamped to avoid division
                             by zero).
        exposure (float): Exposure value (stops). 0 → no change, >0 brighter,
                          <0 darker.
        mask (torch.Tensor, optional): Blend mask of shape ``(B, H, W)`` or
                                       ``(1, H, W)``. If omitted, no blending.

    Returns:
        torch.Tensor: Exposure‑adjusted image tensor with the same shape and
                      dtype as the input, but always float32 internally.
    """
    # --- Input validation ----------------------------------------------------
    if not isinstance(image, torch.Tensor):
        raise ValueError("`image` must be a torch.Tensor")
    if image.dim() != 4:
        raise ValueError(
            f"`image` must have 4 dimensions (B,H,W,C), got {image.dim()}"
        )
    if image.shape[-1] not in (1, 3, 4):
        raise ValueError(
            f"Unsupported number of channels: {image.shape[-1]}. "
            "Expected 1, 3, or 4."
        )
    if mask is not None:
        if not isinstance(mask, torch.Tensor):
            raise ValueError("`mask` must be a torch.Tensor")
        if mask.dim() not in (3, 4):
            raise ValueError(
                f"`mask` must have 3 or 4 dimensions, got {mask.dim()}"
            )
        if mask.shape[0] not in (1, image.shape[0]):
            raise ValueError(
                f"Batch size of mask ({mask.shape[0]}) must be 1 or match "
                f"image batch size ({image.shape[0]})"
            )

    # --- Early exit (identity transform) -------------------------------------
    if exposure == 0.0 and black_level == 0.0:
        # Even with a mask, the blend would be identity → return original
        return image

    # --- Device & precision setup --------------------------------------------
    image = image.to(dtype=torch.float32)
    device = image.device

    # --- Alpha channel separation (if any) -----------------------------------
    has_alpha = image.shape[-1] == 4
    if has_alpha:
        rgb = image[..., :3]          # (B, H, W, 3)
        alpha = image[..., 3:4]       # (B, H, W, 1)
    else:
        rgb = image                  # operate on all channels
        alpha = None

    B, H, W, C_rgb = rgb.shape

    # --- Mask preparation (batched or single) --------------------------------
    mask_tensor = None
    if mask is not None:
        mask_tensor = mask.to(device)
        # Normalise mask shape to (B_mask, H, W, 1)
        if mask_tensor.dim() == 3:
            mask_tensor = mask_tensor.unsqueeze(-1)   # (B, H, W, 1)
        elif mask_tensor.dim() == 4 and mask_tensor.shape[1] == 1:
            # ComfyUI sometimes uses (B, 1, H, W)
            mask_tensor = mask_tensor.permute(0, 2, 3, 1)  # (B, H, W, 1)

        # If a single mask is given, broadcast across the batch later
        single_mask = mask_tensor.shape[0] == 1

    # --- VRAM‑aware chunked processing ---------------------------------------
    chunk_size = compute_chunk_size(image)  # based on full image (includes alpha if present)
    results = []

    for start in range(0, B, chunk_size):
        end = min(start + chunk_size, B)
        rgb_chunk = rgb[start:end]          # (chunk_B, H, W, C_rgb)

        # --- Apply GEGL exposure transform (colour channels only) ------------
        white = torch.exp2(torch.tensor(-exposure, device=device))
        diff = torch.clamp(white - black_level, min=1e-6)
        gain = 1.0 / diff
        transformed = (rgb_chunk - black_level) * gain

        # --- Optional mask blending on colour channels -----------------------
        if mask_tensor is not None:
            if single_mask:
                # Use the single mask for every chunk
                mask_chunk = mask_tensor.expand(end - start, -1, -1, -1)  # (chunk_B, H, W, 1)
            else:
                mask_chunk = mask_tensor[start:end]  # (chunk_B, H, W, 1)

            # Expand to match rgb_chunk shape (broadcast over channels)
            mask_expanded = mask_chunk.expand_as(rgb_chunk)
            result_chunk = rgb_chunk * (1.0 - mask_expanded) + transformed * mask_expanded
        else:
            result_chunk = transformed

        # --- Reattach alpha channel if present -------------------------------
        if has_alpha:
            alpha_chunk = alpha[start:end]   # (chunk_B, H, W, 1)
            result_chunk = torch.cat([result_chunk, alpha_chunk], dim=-1)

        results.append(result_chunk)

    return torch.cat(results, dim=0)