Update Salia_Croppytools.py

Salia_Croppytools.py

@@ -441,7 +441,7 @@ class apply_segment:
             "required": {
                 "mask": ("MASK",),
                 "image": (choices, {}),          # dropdown asset (used ONLY for loaded mask)
-                "img": ("IMAGE",),               # the image to receive final_mask as alpha (overlay source)
+                "img": ("IMAGE",),               # the image to receive alpha_mask as alpha (overlay source)
                 "canvas": ("IMAGE",),            # destination
                 "x": ("INT", {"default": 0, "min": -100000, "max": 100000, "step": 1}),
                 "y": ("INT", {"default": 0, "min": -100000, "max": 100000, "step": 1}),
@@ -458,53 +458,63 @@ class apply_segment:
 
         combiner = _AILab_MaskCombiner_Exact()
 
-        # Load asset mask (do NOT invert)
-        _img_asset, loaded_mask = load_image_from_assets(image)
+        # -------------------
+        # 1) invert input mask -> inverse_mask
+        # -------------------
+        inverse_mask = 1.0 - mask
+        inverse_mask = torch.clamp(inverse_mask, 0.0, 1.0)
 
-        # Invert input mask, then combine with loaded mask (RMBG exact combine => maximum)
-        inv_mask = 1.0 - mask
-        final_mask, = combiner.combine_masks(inv_mask, mode="combine", mask_2=loaded_mask)
+        # Ensure CPU so we don't ever hit CPU/GPU device mismatches with loaded masks/PIL resizing
+        inverse_mask = inverse_mask.detach().cpu()
 
-        # Apply final_mask as alpha to input image -> final_overlay (RGBA)
+        # -------------------
+        # 2) alpha_mask = combine_masks_with_loaded(inverse_mask)
+        # i.e. max(inverse_mask, 1 - loaded_mask) using RMBG exact combiner
+        # -------------------
+        _img_asset, loaded_mask = load_image_from_assets(image)
+        loaded_mask = loaded_mask.detach().cpu()
+
+        alpha_mask, = combiner.combine_masks(
+            inverse_mask,
+            mode="combine",
+            mask_2=(1.0 - loaded_mask)
+        )
+        alpha_mask = torch.clamp(alpha_mask, 0.0, 1.0)
+
+        # -------------------
+        # 3) join image with alpha using alpha_mask -> overlay (RGBA)
+        #    IMPORTANT: replace alpha, don't max() with existing alpha
+        # -------------------
         img = _as_image(img)
         B, H, W, C = img.shape
 
-        # Resize final_mask to match img H/W if needed (uses RMBG exact resize helper)
-        # (target_shape must look like a mask shape [B,H,W], but resize keeps its own batch count)
-        final_mask_resized = combiner._resize_if_needed(final_mask, (final_mask.shape[0], H, W))
+        # Resize alpha_mask to match img H/W if needed (same resize logic style as RMBG)
+        alpha_mask_resized = combiner._resize_if_needed(alpha_mask, (alpha_mask.shape[0], H, W))
 
-        # Batch match (simple 1->N expansion only)
-        if final_mask_resized.shape[0] != B:
-            if final_mask_resized.shape[0] == 1 and B > 1:
-                final_mask_resized = final_mask_resized.expand(B, H, W)
-            elif B == 1 and final_mask_resized.shape[0] > 1:
-                img = img.expand(final_mask_resized.shape[0], *img.shape[1:])
+        # Batch match (allow 1 -> N expansion)
+        if alpha_mask_resized.shape[0] != B:
+            if alpha_mask_resized.shape[0] == 1 and B > 1:
+                alpha_mask_resized = alpha_mask_resized.expand(B, H, W)
+            elif B == 1 and alpha_mask_resized.shape[0] > 1:
+                img = img.expand(alpha_mask_resized.shape[0], *img.shape[1:])
                 B = img.shape[0]
             else:
-                raise ValueError(f"Batch mismatch: img batch={B}, final_mask batch={final_mask_resized.shape[0]}")
-
-        if C == 3:
-            # RGB -> RGBA with alpha = final_mask
-            alpha = final_mask_resized.to(device=img.device, dtype=img.dtype)
-            final_overlay = torch.cat([img, alpha.unsqueeze(-1)], dim=-1)
-        else:
-            # RGBA: combine existing alpha with final_mask using RMBG combine (maximum)
-            rgb = img[..., :3]
-            alpha_img = img[..., 3]  # [B,H,W]
+                raise ValueError(
+                    f"Batch mismatch: img batch={B}, alpha_mask batch={alpha_mask_resized.shape[0]}"
+                )
 
-            # RMBG combine uses PIL-resize sometimes, so keep combine inputs on CPU
-            a1 = alpha_img.detach().cpu()
-            a2 = final_mask_resized.detach().cpu()
-            combined_alpha, = combiner.combine_masks(a1, mode="combine", mask_2=a2)
+        alpha_mask_resized = alpha_mask_resized.to(device=img.device, dtype=img.dtype).clamp(0.0, 1.0)
 
-            combined_alpha = combined_alpha.to(device=img.device, dtype=img.dtype)
-            final_overlay = torch.cat([rgb, combined_alpha.unsqueeze(-1)], dim=-1)
+        rgb = img[..., :3]
+        overlay = torch.cat([rgb, alpha_mask_resized.unsqueeze(-1)], dim=-1)  # RGBA
 
-        # Paste final_overlay onto canvas at (x,y)
+        # -------------------
+        # 4) paste overlay onto canvas at (x,y) -> final output
+        # -------------------
         canvas = _as_image(canvas)
-        final_overlay = final_overlay.to(device=canvas.device, dtype=canvas.dtype)
+        overlay = overlay.to(device=canvas.device, dtype=canvas.dtype)
 
-        out = _alpha_over_region(final_overlay, canvas, x, y)
+        out = _alpha_over_region(overlay, canvas, x, y)
         return (out,)
 
     @classmethod
@@ -525,7 +535,6 @@ class apply_segment:
             return f"File not found in assets/images: {image}"
         return True
 
-
 # -----------------------------
 # Node mappings
 # -----------------------------