Upload Inspyrenet_Rembg2.py

Inspyrenet_Rembg2.py

@@ -0,0 +1,197 @@
+from PIL import Image
+import os
+import urllib.request
+
+import torch
+import numpy as np
+from transparent_background import Remover
+from tqdm import tqdm
+
+
+CKPT_PATH = "/root/.transparent-background/ckpt_base.pth"
+CKPT_URL = "https://huggingface.co/saliacoel/x/resolve/main/ckpt_base.pth"
+
+
+def _ensure_ckpt_base():
+    """
+    1) Check /root/.transparent-background/ckpt_base.pth
+       - if exists: do nothing
+       - else: download from CKPT_URL
+    """
+    try:
+        if os.path.isfile(CKPT_PATH) and os.path.getsize(CKPT_PATH) > 0:
+            return
+    except Exception:
+        # If getsize fails for any reason, fall through to download attempt.
+        pass
+
+    os.makedirs(os.path.dirname(CKPT_PATH), exist_ok=True)
+    tmp_path = CKPT_PATH + ".tmp"
+
+    try:
+        with urllib.request.urlopen(CKPT_URL) as resp:
+            total = resp.headers.get("Content-Length")
+            total = int(total) if total is not None else None
+
+            with open(tmp_path, "wb") as f:
+                if total:
+                    with tqdm(
+                        total=total,
+                        unit="B",
+                        unit_scale=True,
+                        desc="Downloading ckpt_base.pth",
+                    ) as pbar:
+                        while True:
+                            chunk = resp.read(1024 * 1024)
+                            if not chunk:
+                                break
+                            f.write(chunk)
+                            pbar.update(len(chunk))
+                else:
+                    while True:
+                        chunk = resp.read(1024 * 1024)
+                        if not chunk:
+                            break
+                        f.write(chunk)
+
+        os.replace(tmp_path, CKPT_PATH)
+
+    finally:
+        # Clean up partial download if something went wrong
+        if os.path.isfile(tmp_path):
+            try:
+                os.remove(tmp_path)
+            except Exception:
+                pass
+
+
+# Tensor to PIL
+def tensor2pil(image: torch.Tensor) -> Image.Image:
+    arr = image.detach().cpu().numpy()
+    # Handle accidental singleton batch dim
+    if arr.ndim == 4 and arr.shape[0] == 1:
+        arr = arr[0]
+    arr = np.clip(255.0 * arr, 0, 255).astype(np.uint8)
+    return Image.fromarray(arr)
+
+
+# Convert PIL to Tensor
+def pil2tensor(image: Image.Image) -> torch.Tensor:
+    return torch.from_numpy(np.array(image).astype(np.float32) / 255.0).unsqueeze(0)
+
+
+def _rgba_to_rgb_on_white(pil_img: Image.Image) -> Image.Image:
+    """
+    5) If input is RGBA:
+         - alpha composite over WHITE background
+         - convert to RGB (drop alpha)
+       If input is RGB:
+         - carry on
+    """
+    if pil_img.mode == "RGBA":
+        bg = Image.new("RGBA", pil_img.size, (255, 255, 255, 255))
+        composited = Image.alpha_composite(bg, pil_img)
+        return composited.convert("RGB")
+
+    if pil_img.mode != "RGB":
+        return pil_img.convert("RGB")
+
+    return pil_img
+
+
+class InspyrenetRembg2:
+    """
+    Original node kept (unchanged behavior/output), except it now ensures ckpt exists.
+    """
+    def __init__(self):
+        pass
+
+    @classmethod
+    def INPUT_TYPES(s):
+        return {
+            "required": {
+                "image": ("IMAGE",),
+                "torchscript_jit": (["default", "on"],)
+            },
+        }
+
+    RETURN_TYPES = ("IMAGE", "MASK")
+    FUNCTION = "remove_background"
+    CATEGORY = "image"
+
+    def remove_background(self, image, torchscript_jit):
+        _ensure_ckpt_base()
+
+        if (torchscript_jit == "default"):
+            remover = Remover()
+        else:
+            remover = Remover(jit=True)
+
+        img_list = []
+        for img in tqdm(image, "Inspyrenet Rembg"):
+            mid = remover.process(tensor2pil(img), type='rgba')
+            out = pil2tensor(mid)
+            img_list.append(out)
+
+        img_stack = torch.cat(img_list, dim=0)
+        mask = img_stack[:, :, :, 3]
+        return (img_stack, mask)
+
+
+class InspyrenetRembg3:
+    """
+    New node per requested changes:
+      - ensures ckpt_base.pth exists (downloads if missing)
+      - torchscript_jit hardcoded to "default" (no input, no JIT)
+      - NO MASK output (IMAGE only)
+      - if input is RGBA: composite over white, convert to RGB, then run remover
+      - output remains RGBA (type='rgba')
+    """
+    def __init__(self):
+        pass
+
+    @classmethod
+    def INPUT_TYPES(s):
+        return {
+            "required": {
+                "image": ("IMAGE",),
+            },
+        }
+
+    RETURN_TYPES = ("IMAGE",)
+    FUNCTION = "remove_background"
+    CATEGORY = "image"
+
+    def remove_background(self, image):
+        _ensure_ckpt_base()
+
+        # 3) hardcode torchscript_jit == "default"
+        remover = Remover()
+
+        img_list = []
+        for img in tqdm(image, "Inspyrenet Rembg3"):
+            pil_in = tensor2pil(img)
+
+            # 5) normalize input to RGB for the model:
+            #    - if RGBA -> alpha composite on white -> RGB
+            #    - if RGB  -> keep
+            pil_rgb = _rgba_to_rgb_on_white(pil_in)
+
+            # do functionality as usual, output RGBA
+            mid = remover.process(pil_rgb, type="rgba")
+            out = pil2tensor(mid)
+            img_list.append(out)
+
+        img_stack = torch.cat(img_list, dim=0)
+        return (img_stack,)
+
+
+NODE_CLASS_MAPPINGS = {
+    "InspyrenetRembg2": InspyrenetRembg2,
+    "InspyrenetRembg3": InspyrenetRembg3,
+}
+
+NODE_DISPLAY_NAME_MAPPINGS = {
+    "InspyrenetRembg2": "Inspyrenet Rembg2",
+    "InspyrenetRembg3": "Inspyrenet Rembg3",
+}