Upload Salia_Croppytools.py

Salia_Croppytools.py

@@ -0,0 +1,460 @@
+import os
+from typing import Tuple
+
+import torch
+import torch.nn.functional as F
+
+# Salia utils (same style as your loader node)
+try:
+    from ..utils.io import list_pngs, load_image_from_assets, file_hash, safe_path
+except Exception:
+    # Fallback if you placed this file in a different package depth
+    try:
+        from .utils.io import list_pngs, load_image_from_assets, file_hash, safe_path
+    except Exception as e:
+        _UTILS_IMPORT_ERR = e
+
+        def _missing(*args, **kwargs):
+            raise ImportError(
+                "Could not import Salia utils (list_pngs/load_image_from_assets/file_hash/safe_path). "
+                "Place this node file in the same package layout as your other Salia nodes.\n"
+                f"Original import error: {_UTILS_IMPORT_ERR}"
+            )
+
+        list_pngs = _missing
+        load_image_from_assets = _missing
+        file_hash = _missing
+        safe_path = _missing
+
+
+# -----------------------------
+# Helpers
+# -----------------------------
+
+def _as_image(img: torch.Tensor) -> torch.Tensor:
+    # ComfyUI IMAGE is usually [B,H,W,C]
+    if not isinstance(img, torch.Tensor):
+        raise TypeError("IMAGE must be a torch.Tensor")
+    if img.dim() != 4:
+        raise ValueError(f"Expected IMAGE shape [B,H,W,C], got {tuple(img.shape)}")
+    if img.shape[-1] not in (3, 4):
+        raise ValueError(f"Expected IMAGE channels 3 (RGB) or 4 (RGBA), got C={img.shape[-1]}")
+    return img
+
+
+def _as_mask(msk: torch.Tensor) -> torch.Tensor:
+    # ComfyUI MASK is usually [B,H,W] float 0..1
+    if not isinstance(msk, torch.Tensor):
+        raise TypeError("MASK must be a torch.Tensor")
+    if msk.dim() == 2:
+        msk = msk.unsqueeze(0)
+    if msk.dim() != 3:
+        raise ValueError(f"Expected MASK shape [B,H,W] (or [H,W]), got {tuple(msk.shape)}")
+    return msk
+
+
+def _match_batch(a: torch.Tensor, b: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]:
+    ba = a.shape[0]
+    bb = b.shape[0]
+    if ba == bb:
+        return a, b
+    if ba == 1 and bb > 1:
+        return a.expand(bb, *a.shape[1:]), b
+    if bb == 1 and ba > 1:
+        return a, b.expand(ba, *b.shape[1:])
+    raise ValueError(f"Batch mismatch: A has batch {ba}, B has batch {bb} (and neither is 1).")
+
+
+def _resize_mask_to(msk: torch.Tensor, target_h: int, target_w: int) -> torch.Tensor:
+    # msk: [B,H,W] -> resize to [B,target_h,target_w]
+    if msk.shape[1] == target_h and msk.shape[2] == target_w:
+        return msk
+    x = msk.unsqueeze(1)  # [B,1,H,W]
+    x = F.interpolate(x, size=(target_h, target_w), mode="bilinear", align_corners=False)
+    return x.squeeze(1)
+
+
+def _combine_alpha_union(a: torch.Tensor, b: torch.Tensor) -> torch.Tensor:
+    """
+    "Alpha combine" (union) like standard alpha coverage:
+        out = 1 - (1-a)*(1-b)
+    """
+    a = a.clamp(0.0, 1.0)
+    b = b.clamp(0.0, 1.0)
+    return (1.0 - (1.0 - a) * (1.0 - b)).clamp(0.0, 1.0)
+
+
+def _crop_with_padding(image: torch.Tensor, x: int, y: int, w: int, h: int) -> torch.Tensor:
+    """
+    Crops [x,y] top-left, size w*h. If out of bounds, pads with zeros.
+    image: [B,H,W,C]
+    returns: [B,h,w,C]
+    """
+    image = _as_image(image)
+    B, H, W, C = image.shape
+    w = max(1, int(w))
+    h = max(1, int(h))
+    x = int(x)
+    y = int(y)
+
+    out = torch.zeros((B, h, w, C), device=image.device, dtype=image.dtype)
+
+    # intersection in source
+    x0s = max(0, x)
+    y0s = max(0, y)
+    x1s = min(W, x + w)
+    y1s = min(H, y + h)
+
+    if x1s <= x0s or y1s <= y0s:
+        return out
+
+    # destination offsets
+    x0d = x0s - x
+    y0d = y0s - y
+    x1d = x0d + (x1s - x0s)
+    y1d = y0d + (y1s - y0s)
+
+    out[:, y0d:y1d, x0d:x1d, :] = image[:, y0s:y1s, x0s:x1s, :]
+    return out
+
+
+def _ensure_rgba(img: torch.Tensor) -> torch.Tensor:
+    """
+    img: [B,H,W,C] where C is 3 or 4
+    returns RGBA [B,H,W,4]
+    """
+    img = _as_image(img)
+    if img.shape[-1] == 4:
+        return img
+    # RGB -> RGBA with alpha=1
+    B, H, W, _ = img.shape
+    alpha = torch.ones((B, H, W, 1), device=img.device, dtype=img.dtype)
+    return torch.cat([img, alpha], dim=-1)
+
+
+def _alpha_over_region(overlay: torch.Tensor, canvas: torch.Tensor, x: int, y: int) -> torch.Tensor:
+    """
+    Places overlay at canvas pixel position (x,y) top-left corner.
+    Supports RGB/RGBA for both. Uses alpha-over if overlay has alpha or canvas has alpha.
+    Returns same channel count as canvas (3->3, 4->4).
+    """
+    overlay = _as_image(overlay)
+    canvas = _as_image(canvas)
+
+    overlay, canvas = _match_batch(overlay, canvas)
+
+    B, Hc, Wc, Cc = canvas.shape
+    Bo, Ho, Wo, Co = overlay.shape
+
+    x = int(x)
+    y = int(y)
+
+    out = canvas.clone()
+
+    # intersection on canvas
+    x0c = max(0, x)
+    y0c = max(0, y)
+    x1c = min(Wc, x + Wo)
+    y1c = min(Hc, y + Ho)
+
+    if x1c <= x0c or y1c <= y0c:
+        return out
+
+    # corresponding region on overlay
+    x0o = x0c - x
+    y0o = y0c - y
+    x1o = x0o + (x1c - x0c)
+    y1o = y0o + (y1c - y0c)
+
+    canvas_region = out[:, y0c:y1c, x0c:x1c, :]
+    overlay_region = overlay[:, y0o:y1o, x0o:x1o, :]
+
+    # Convert both regions to RGBA for compositing
+    canvas_rgba = _ensure_rgba(canvas_region)
+    overlay_rgba = _ensure_rgba(overlay_region)
+
+    over_rgb = overlay_rgba[..., :3].clamp(0.0, 1.0)
+    over_a = overlay_rgba[..., 3:4].clamp(0.0, 1.0)
+
+    under_rgb = canvas_rgba[..., :3].clamp(0.0, 1.0)
+    under_a = canvas_rgba[..., 3:4].clamp(0.0, 1.0)
+
+    # Premultiplied alpha composite: out = over + under*(1-over_a)
+    over_pm = over_rgb * over_a
+    under_pm = under_rgb * under_a
+
+    out_a = over_a + under_a * (1.0 - over_a)
+    out_pm = over_pm + under_pm * (1.0 - over_a)
+
+    eps = 1e-6
+    out_rgb = torch.where(out_a > eps, out_pm / (out_a + eps), torch.zeros_like(out_pm))
+    out_rgb = out_rgb.clamp(0.0, 1.0)
+    out_a = out_a.clamp(0.0, 1.0)
+
+    if Cc == 3:
+        out[:, y0c:y1c, x0c:x1c, :] = out_rgb
+    else:
+        out[:, y0c:y1c, x0c:x1c, :] = torch.cat([out_rgb, out_a], dim=-1)
+
+    return out
+
+
+# -----------------------------
+# 1) Cropout_Square_From_IMG
+# -----------------------------
+
+class Cropout_Square_From_IMG:
+    CATEGORY = "image/salia"
+
+    @classmethod
+    def INPUT_TYPES(cls):
+        return {
+            "required": {
+                "img": ("IMAGE",),
+                "x": ("INT", {"default": 0, "min": -100000, "max": 100000, "step": 1}),
+                "y": ("INT", {"default": 0, "min": -100000, "max": 100000, "step": 1}),
+                "square_size": ("INT", {"default": 512, "min": 1, "max": 16384, "step": 1}),
+            }
+        }
+
+    RETURN_TYPES = ("IMAGE",)
+    RETURN_NAMES = ("image",)
+    FUNCTION = "run"
+
+    def run(self, img, x, y, square_size):
+        cropped = _crop_with_padding(img, x, y, square_size, square_size)
+        return (cropped,)
+
+
+# -----------------------------
+# 2) Cropout_Rect_From_IMG
+# -----------------------------
+
+class Cropout_Rect_From_IMG:
+    CATEGORY = "image/salia"
+
+    @classmethod
+    def INPUT_TYPES(cls):
+        return {
+            "required": {
+                "img": ("IMAGE",),
+                "x": ("INT", {"default": 0, "min": -100000, "max": 100000, "step": 1}),
+                "y": ("INT", {"default": 0, "min": -100000, "max": 100000, "step": 1}),
+                "width": ("INT", {"default": 512, "min": 1, "max": 16384, "step": 1}),
+                "height": ("INT", {"default": 512, "min": 1, "max": 16384, "step": 1}),
+            }
+        }
+
+    RETURN_TYPES = ("IMAGE",)
+    RETURN_NAMES = ("image",)
+    FUNCTION = "run"
+
+    def run(self, img, x, y, width, height):
+        cropped = _crop_with_padding(img, x, y, width, height)
+        return (cropped,)
+
+
+# -----------------------------
+# 3) Paste_rect_to_img
+# -----------------------------
+
+class Paste_rect_to_img:
+    CATEGORY = "image/salia"
+
+    @classmethod
+    def INPUT_TYPES(cls):
+        return {
+            "required": {
+                "overlay": ("IMAGE",),
+                "canvas": ("IMAGE",),
+                "x": ("INT", {"default": 0, "min": -100000, "max": 100000, "step": 1}),
+                "y": ("INT", {"default": 0, "min": -100000, "max": 100000, "step": 1}),
+            }
+        }
+
+    RETURN_TYPES = ("IMAGE",)
+    RETURN_NAMES = ("image",)
+    FUNCTION = "run"
+
+    def run(self, overlay, canvas, x, y):
+        out = _alpha_over_region(overlay, canvas, x, y)
+        return (out,)
+
+
+# -----------------------------
+# 4) Combine_2_masks
+# -----------------------------
+
+class Combine_2_masks:
+    CATEGORY = "mask/salia"
+
+    @classmethod
+    def INPUT_TYPES(cls):
+        return {
+            "required": {
+                "maskA": ("MASK",),
+                "maskB": ("MASK",),
+            }
+        }
+
+    RETURN_TYPES = ("MASK",)
+    RETURN_NAMES = ("mask",)
+    FUNCTION = "run"
+
+    def run(self, maskA, maskB):
+        a = _as_mask(maskA)
+        b = _as_mask(maskB)
+
+        a, b = _match_batch(a, b)
+        b = _resize_mask_to(b, a.shape[1], a.shape[2])
+
+        out = _combine_alpha_union(a, b)
+        return (out,)
+
+
+# -----------------------------
+# 5) Combine_2_masks_invert_1
+# -----------------------------
+
+class Combine_2_masks_invert_1:
+    CATEGORY = "mask/salia"
+
+    @classmethod
+    def INPUT_TYPES(cls):
+        return {
+            "required": {
+                "maskA": ("MASK",),
+                "maskB": ("MASK",),
+            }
+        }
+
+    RETURN_TYPES = ("MASK",)
+    RETURN_NAMES = ("mask",)
+    FUNCTION = "run"
+
+    def run(self, maskA, maskB):
+        a = _as_mask(maskA)
+        b = _as_mask(maskB)
+
+        a, b = _match_batch(a, b)
+        b = _resize_mask_to(b, a.shape[1], a.shape[2])
+
+        a_inv = (1.0 - a).clamp(0.0, 1.0)
+        out = _combine_alpha_union(a_inv, b)
+        return (out,)
+
+
+# -----------------------------
+# 6) Combine_2_masks_inverse
+# -----------------------------
+
+class Combine_2_masks_inverse:
+    CATEGORY = "mask/salia"
+
+    @classmethod
+    def INPUT_TYPES(cls):
+        return {
+            "required": {
+                "maskA": ("MASK",),
+                "maskB": ("MASK",),
+            }
+        }
+
+    RETURN_TYPES = ("MASK",)
+    RETURN_NAMES = ("mask",)
+    FUNCTION = "run"
+
+    def run(self, maskA, maskB):
+        a = _as_mask(maskA)
+        b = _as_mask(maskB)
+
+        a, b = _match_batch(a, b)
+        b = _resize_mask_to(b, a.shape[1], a.shape[2])
+
+        a_inv = (1.0 - a).clamp(0.0, 1.0)
+        b_inv = (1.0 - b).clamp(0.0, 1.0)
+
+        combined_inv = _combine_alpha_union(a_inv, b_inv)
+        out = (1.0 - combined_inv).clamp(0.0, 1.0)  # == a*b (intersection)
+        return (out,)
+
+
+# -----------------------------
+# 7) combine_masks_with_loaded
+# -----------------------------
+
+class combine_masks_with_loaded:
+    CATEGORY = "mask/salia"
+
+    @classmethod
+    def INPUT_TYPES(cls):
+        choices = list_pngs() or ["<no pngs found>"]
+        return {
+            "required": {
+                "mask": ("MASK",),
+                "image": (choices, {}),
+            }
+        }
+
+    RETURN_TYPES = ("MASK",)
+    RETURN_NAMES = ("mask",)
+    FUNCTION = "run"
+
+    def run(self, mask, image):
+        if image == "<no pngs found>":
+            raise FileNotFoundError("No PNGs in assets/images")
+
+        base = _as_mask(mask)
+
+        # Load image+mask from assets (Salia util)
+        _img, loaded_mask = load_image_from_assets(image)
+        loaded = _as_mask(loaded_mask)
+
+        base, loaded = _match_batch(base, loaded)
+        loaded = _resize_mask_to(loaded, base.shape[1], base.shape[2])
+
+        out = _combine_alpha_union(base, loaded)
+        return (out,)
+
+    @classmethod
+    def IS_CHANGED(cls, mask, image):
+        if image == "<no pngs found>":
+            return image
+        return file_hash(image)
+
+    @classmethod
+    def VALIDATE_INPUTS(cls, mask, image):
+        if image == "<no pngs found>":
+            return "No PNGs in assets/images"
+        try:
+            path = safe_path(image)
+        except Exception as e:
+            return str(e)
+        if not os.path.isfile(path):
+            return f"File not found in assets/images: {image}"
+        return True
+
+
+# -----------------------------
+# Node mappings
+# -----------------------------
+
+NODE_CLASS_MAPPINGS = {
+    "Cropout_Square_From_IMG": Cropout_Square_From_IMG,
+    "Cropout_Rect_From_IMG": Cropout_Rect_From_IMG,
+    "Paste_rect_to_img": Paste_rect_to_img,
+    "Combine_2_masks": Combine_2_masks,
+    "Combine_2_masks_invert_1": Combine_2_masks_invert_1,
+    "Combine_2_masks_inverse": Combine_2_masks_inverse,
+    "combine_masks_with_loaded": combine_masks_with_loaded,
+}
+
+NODE_DISPLAY_NAME_MAPPINGS = {
+    "Cropout_Square_From_IMG": "Cropout_Square_From_IMG",
+    "Cropout_Rect_From_IMG": "Cropout_Rect_From_IMG",
+    "Paste_rect_to_img": "Paste_rect_to_img",
+    "Combine_2_masks": "Combine_2_masks",
+    "Combine_2_masks_invert_1": "Combine_2_masks_invert_1",
+    "Combine_2_masks_inverse": "Combine_2_masks_inverse",
+    "combine_masks_with_loaded": "combine_masks_with_loaded",
+}