Upload img_alpha_composite_coords.py

img_alpha_composite_coords.py

@@ -0,0 +1,179 @@
+# paste_rgba_at_xy.py
+# ComfyUI custom node: Paste a small RGBA image onto a larger RGBA canvas at (x, y)
+# Supports alpha compositing ("source-over") or hard replace.
+#
+# Usage:
+#  - canvas: big RGBA image (IMAGE)
+#  - overlay: small RGBA image (IMAGE)
+#  - x, y: top-left destination coordinate on the canvas
+#  - blend_mode: "alpha_over" (default) or "replace"
+#  - If either input is RGB (3ch), it will be upgraded to RGBA with alpha=1.0
+#  - Batching:
+#       * If one input has batch size 1 and the other >1, the single image is broadcast.
+#       * If both have batch >1, their batch sizes must match.
+
+import torch
+
+class PasteRGBAAtXY:
+    """
+    Paste a small RGBA image on a larger RGBA canvas at a specified (x, y) coordinate.
+
+    - Default blend mode is proper alpha compositing (SRC over).
+    - "replace" mode copies the overlay's RGBA pixels directly (no blending).
+    - Handles out-of-bounds and negative coordinates by clipping.
+    - Works with batches. Broadcasts a single overlay across a batch of canvases (and vice versa).
+    """
+
+    CATEGORY = "image/compose"
+    RETURN_TYPES = ("IMAGE",)
+    FUNCTION = "paste"
+
+    @classmethod
+    def INPUT_TYPES(cls):
+        return {
+            "required": {
+                "canvas": ("IMAGE",),
+                "overlay": ("IMAGE",),
+                "x": ("INT", {"default": 0, "min": -32768, "max": 32768, "step": 1}),
+                "y": ("INT", {"default": 0, "min": -32768, "max": 32768, "step": 1}),
+                "blend_mode": (["alpha_over", "replace"], {"default": "alpha_over"}),
+            }
+        }
+
+    @staticmethod
+    def _ensure_rgba(img: torch.Tensor) -> torch.Tensor:
+        """
+        Ensure an IMAGE tensor is RGBA. If RGB, append opaque alpha.
+        Shape convention in ComfyUI for IMAGE is [B, H, W, C] with float32 in [0, 1].
+        """
+        if img.ndim != 4:
+            raise ValueError(f"Expected image tensor with 4 dims [B,H,W,C], got shape {tuple(img.shape)}")
+        c = img.shape[-1]
+        if c == 4:
+            return img
+        if c == 3:
+            alpha = torch.ones((*img.shape[:-1], 1), dtype=img.dtype, device=img.device)
+            return torch.cat([img, alpha], dim=-1)
+        raise ValueError(f"Expected 3 or 4 channels, got {c}")
+
+    @staticmethod
+    def _pair_count(b1: int, b2: int) -> int:
+        """Return the output batch size if broadcasting is allowed."""
+        if b1 == b2:
+            return b1
+        if b1 == 1 or b2 == 1:
+            return max(b1, b2)
+        raise ValueError(f"Incompatible batch sizes: {b1} vs {b2}")
+
+    @staticmethod
+    def _get_batch(img: torch.Tensor, i: int) -> torch.Tensor:
+        """Fetch the i-th batch image with broadcasting if needed."""
+        b = img.shape[0]
+        return img[0] if b == 1 else img[i]
+
+    @staticmethod
+    def _alpha_over(dst_rgba: torch.Tensor, src_rgba: torch.Tensor, dx: int, dy: int) -> None:
+        """
+        Alpha-composite src onto dst in-place at integer offset (dx, dy).
+        Both tensors are HxWx4, float in [0,1]. Clips automatically when out of bounds.
+        """
+        Hc, Wc, _ = dst_rgba.shape
+        Ho, Wo, _ = src_rgba.shape
+
+        # Compute intersection rectangle on canvas (destination)
+        x0 = max(0, dx)
+        y0 = max(0, dy)
+        x1 = min(Wc, dx + Wo)
+        y1 = min(Hc, dy + Ho)
+
+        if x1 <= x0 or y1 <= y0:
+            return  # Nothing overlaps
+
+        # Corresponding source crop
+        sx0 = x0 - dx
+        sy0 = y0 - dy
+        w = x1 - x0
+        h = y1 - y0
+
+        dst_region = dst_rgba[y0:y0+h, x0:x0+w, :]
+        src_region = src_rgba[sy0:sy0+h, sx0:sx0+w, :]
+
+        # Split channels
+        cb = dst_region[..., :3]
+        ab = dst_region[..., 3:4]
+        co = src_region[..., :3]
+        ao = src_region[..., 3:4]
+
+        # Premultiply colors
+        cb_p = cb * ab
+        co_p = co * ao
+
+        # Source-over composition
+        out_a = ao + ab * (1.0 - ao)
+        out_c_p = co_p + cb_p * (1.0 - ao)
+
+        # Convert back to straight (guard against divide by zero)
+        eps = 1e-8
+        out_c = torch.where(out_a > eps, out_c_p / out_a.clamp_min(eps), torch.zeros_like(out_c_p))
+
+        # Write back (clamp just in case)
+        dst_region[..., :3] = out_c.clamp(0.0, 1.0)
+        dst_region[..., 3:4] = out_a.clamp(0.0, 1.0)
+
+        dst_rgba[y0:y0+h, x0:x0+w, :] = dst_region
+
+    @staticmethod
+    def _replace(dst_rgba: torch.Tensor, src_rgba: torch.Tensor, dx: int, dy: int) -> None:
+        """
+        Direct overwrite (copy) of src_rgba into dst_rgba at (dx, dy), clipped to bounds.
+        """
+        Hc, Wc, _ = dst_rgba.shape
+        Ho, Wo, _ = src_rgba.shape
+
+        x0 = max(0, dx)
+        y0 = max(0, dy)
+        x1 = min(Wc, dx + Wo)
+        y1 = min(Hc, dy + Ho)
+        if x1 <= x0 or y1 <= y0:
+            return
+
+        sx0 = x0 - dx
+        sy0 = y0 - dy
+        w = x1 - x0
+        h = y1 - y0
+
+        dst_rgba[y0:y0+h, x0:x0+w, :] = src_rgba[sy0:sy0+h, sx0:sx0+w, :]
+
+    def paste(self, canvas, overlay, x, y, blend_mode):
+        # Ensure RGBA
+        canvas = self._ensure_rgba(canvas)
+        overlay = self._ensure_rgba(overlay)
+
+        b_canvas = canvas.shape[0]
+        b_overlay = overlay.shape[0]
+        out_b = self._pair_count(b_canvas, b_overlay)
+
+        out_list = []
+        for i in range(out_b):
+            base = self._get_batch(canvas, i).clone()  # HxWx4
+            over = self._get_batch(overlay, i)         # HxWx4
+
+            if blend_mode == "alpha_over":
+                self._alpha_over(base, over, int(x), int(y))
+            else:  # "replace"
+                self._replace(base, over, int(x), int(y))
+
+            out_list.append(base)
+
+        out = torch.stack(out_list, dim=0)
+        return (out,)
+
+
+# --- ComfyUI registration ---
+NODE_CLASS_MAPPINGS = {
+    "PasteRGBAAtXY": PasteRGBAAtXY,
+}
+
+NODE_DISPLAY_NAME_MAPPINGS = {
+    "PasteRGBAAtXY": "Paste RGBA at (X,Y)",
+}