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.
77 lines
2.2 KiB
Python
77 lines
2.2 KiB
Python
import torch
|
|
from typing import Optional
|
|
|
|
def compute_chunk_size(
|
|
image: torch.Tensor,
|
|
mask: Optional[torch.Tensor] = None,
|
|
) -> int:
|
|
"""
|
|
Determine a safe batch-size chunk for GPU processing based on the
|
|
currently available CUDA memory.
|
|
|
|
The estimation accounts for the memory required by the input image,
|
|
an optional mask tensor, the output image, and temporary working
|
|
tensors created during color conversion.
|
|
|
|
Args:
|
|
image (torch.Tensor):
|
|
Input image tensor of shape (B, H, W, C).
|
|
|
|
mask (torch.Tensor, optional):
|
|
Optional mask tensor associated with the image batch.
|
|
|
|
Returns:
|
|
int:
|
|
Maximum number of images that can be processed safely in a
|
|
single chunk. If the image resides on the CPU, the complete
|
|
batch size is returned.
|
|
"""
|
|
if image.device.type != "cuda":
|
|
return image.shape[0]
|
|
|
|
free_mem, _ = torch.cuda.mem_get_info(image.device)
|
|
|
|
image_batch = max(image.shape[0], 1)
|
|
|
|
# Bytes required by one image sample
|
|
image_sample_bytes = (
|
|
image[0].numel() * image.element_size()
|
|
)
|
|
|
|
# Bytes required by one mask sample
|
|
mask_sample_bytes = 0
|
|
|
|
if mask is not None:
|
|
if mask.device != image.device:
|
|
mask = mask.to(image.device)
|
|
|
|
mask_batch = max(mask.shape[0], 1)
|
|
|
|
if mask_batch == image_batch:
|
|
mask_sample_bytes = mask[0].numel() * mask.element_size()
|
|
elif mask_batch == 1:
|
|
# Single mask broadcast across the batch
|
|
mask_sample_bytes = mask[0].numel() * mask.element_size()
|
|
else:
|
|
raise ValueError(
|
|
f"Mask batch size ({mask_batch}) does not match image batch size ({image_batch})."
|
|
)
|
|
|
|
# Working memory:
|
|
# input image
|
|
# output image
|
|
# temporary color conversion tensors
|
|
# optional mask
|
|
working_factor = 4
|
|
|
|
required_per_sample = (
|
|
image_sample_bytes * working_factor +
|
|
mask_sample_bytes
|
|
)
|
|
|
|
# Keep a safety margin so other CUDA allocations have room.
|
|
usable_memory = int(free_mem * 0.90)
|
|
|
|
return max(1, min(image_batch, usable_memory // max(required_per_sample, 1)))
|
|
|