Upload salia_detailer_ezpz_gated.py

salia_detailer_ezpz_gated.py

@@ -0,0 +1,890 @@
+import hashlib
+import shutil
+import threading
+import urllib.request
+from pathlib import Path
+from typing import Any, Dict, Tuple, Optional
+
+import numpy as np
+import torch
+from PIL import Image, ImageOps
+
+import folder_paths
+import comfy.model_management as model_management
+
+
+# transformers is required for depth-estimation pipeline
+try:
+    from transformers import pipeline
+except Exception as e:
+    pipeline = None
+    _TRANSFORMERS_IMPORT_ERROR = e
+
+
+# -------------------------------------------------------------------------------------
+# Global caches (checkpoint + controlnet) so using the node multiple times won't reload
+# -------------------------------------------------------------------------------------
+
+_CKPT_CACHE: Dict[str, Tuple[Any, Any, Any]] = {}
+_CN_CACHE: Dict[str, Any] = {}
+_CKPT_LOCK = threading.Lock()
+_CN_LOCK = threading.Lock()
+
+
+# -------------------------------------------------------------------------------------
+# Plugin root detection (works whether file is in plugin root or nodes/)
+# -------------------------------------------------------------------------------------
+
+def _find_plugin_root() -> Path:
+    """
+    Walk upwards from this file until we find an 'assets' folder.
+    Robust against hyphen/underscore package naming and different file placement.
+    """
+    here = Path(__file__).resolve()
+    for parent in [here.parent] + list(here.parents)[:12]:
+        if (parent / "assets").is_dir():
+            return parent
+    # fallback: typical nodes/<file>.py
+    return here.parent.parent
+
+
+PLUGIN_ROOT = _find_plugin_root()
+
+
+# -------------------------------------------------------------------------------------
+# PIL helpers (Lanczos resize for IMAGE and MASK)
+# -------------------------------------------------------------------------------------
+
+def _pil_lanczos():
+    if hasattr(Image, "Resampling"):
+        return Image.Resampling.LANCZOS
+    return Image.LANCZOS
+
+
+def _image_tensor_to_pil(img: torch.Tensor) -> Image.Image:
+    """
+    Comfy IMAGE: [B,H,W,C] or [H,W,C], float [0..1] -> PIL RGB/RGBA
+    """
+    if img.ndim == 4:
+        img = img[0]
+    img = img.detach().cpu().float().clamp(0, 1)
+    arr = (img.numpy() * 255.0).round().astype(np.uint8)
+    if arr.shape[-1] == 4:
+        return Image.fromarray(arr, mode="RGBA")
+    return Image.fromarray(arr, mode="RGB")
+
+
+def _pil_to_image_tensor(pil: Image.Image) -> torch.Tensor:
+    """
+    PIL RGB/RGBA -> Comfy IMAGE [1,H,W,C], float [0..1]
+    """
+    if pil.mode not in ("RGB", "RGBA"):
+        pil = pil.convert("RGBA") if "A" in pil.getbands() else pil.convert("RGB")
+    arr = np.array(pil).astype(np.float32) / 255.0
+    t = torch.from_numpy(arr)  # [H,W,C]
+    return t.unsqueeze(0)
+
+
+def _mask_tensor_to_pil(mask: torch.Tensor) -> Image.Image:
+    """
+    Comfy MASK: [B,H,W] or [H,W], float [0..1] -> PIL L
+    """
+    if mask.ndim == 3:
+        mask = mask[0]
+    mask = mask.detach().cpu().float().clamp(0, 1)
+    arr = (mask.numpy() * 255.0).round().astype(np.uint8)
+    return Image.fromarray(arr, mode="L")
+
+
+def _pil_to_mask_tensor(pil_l: Image.Image) -> torch.Tensor:
+    """
+    PIL L -> Comfy MASK [1,H,W], float [0..1]
+    """
+    if pil_l.mode != "L":
+        pil_l = pil_l.convert("L")
+    arr = np.array(pil_l).astype(np.float32) / 255.0
+    t = torch.from_numpy(arr)  # [H,W]
+    return t.unsqueeze(0)
+
+
+def _resize_image_lanczos(img: torch.Tensor, w: int, h: int) -> torch.Tensor:
+    """
+    Resize Comfy IMAGE [B,H,W,C] with Lanczos via PIL.
+    """
+    if img.ndim != 4:
+        raise ValueError("Expected IMAGE tensor with shape [B,H,W,C].")
+    outs = []
+    for i in range(img.shape[0]):
+        pil = _image_tensor_to_pil(img[i].unsqueeze(0))
+        pil = pil.resize((int(w), int(h)), resample=_pil_lanczos())
+        outs.append(_pil_to_image_tensor(pil))
+    return torch.cat(outs, dim=0)
+
+
+def _resize_mask_lanczos(mask: torch.Tensor, w: int, h: int) -> torch.Tensor:
+    """
+    Resize Comfy MASK [B,H,W] with Lanczos via PIL.
+    """
+    if mask.ndim != 3:
+        raise ValueError("Expected MASK tensor with shape [B,H,W].")
+    outs = []
+    for i in range(mask.shape[0]):
+        pil = _mask_tensor_to_pil(mask[i].unsqueeze(0))
+        pil = pil.resize((int(w), int(h)), resample=_pil_lanczos())
+        outs.append(_pil_to_mask_tensor(pil))
+    return torch.cat(outs, dim=0)
+
+
+# -------------------------------------------------------------------------------------
+# ✅ ComfyUI 0.5.1 FIX: Manual JoinImageWithAlpha equivalent
+# -------------------------------------------------------------------------------------
+
+def _rgb_to_rgba_with_comfy_mask(rgb: torch.Tensor, mask: torch.Tensor) -> torch.Tensor:
+    """
+    Make RGBA from:
+      rgb:  IMAGE [B,H,W,3] float [0..1]
+      mask: MASK  [B,H,W]   float [0..1]  (Comfy convention: 1=masked/transparent)
+    Output:
+      rgba: IMAGE [B,H,W,4] where alpha = 1 - mask  (1=opaque, 0=transparent)
+    """
+    if rgb.ndim == 3:
+        rgb = rgb.unsqueeze(0)
+    if mask.ndim == 2:
+        mask = mask.unsqueeze(0)
+
+    if rgb.ndim != 4 or rgb.shape[-1] != 3:
+        raise ValueError(f"rgb must be [B,H,W,3], got {tuple(rgb.shape)}")
+    if mask.ndim != 3:
+        raise ValueError(f"mask must be [B,H,W], got {tuple(mask.shape)}")
+
+    # Batch match
+    if mask.shape[0] != rgb.shape[0]:
+        if mask.shape[0] == 1 and rgb.shape[0] > 1:
+            mask = mask.expand(rgb.shape[0], -1, -1)
+        else:
+            raise ValueError("Batch mismatch between rgb and mask.")
+
+    # Size match
+    if mask.shape[1] != rgb.shape[1] or mask.shape[2] != rgb.shape[2]:
+        raise ValueError(
+            f"Mask size mismatch. rgb={rgb.shape[2]}x{rgb.shape[1]} mask={mask.shape[2]}x{mask.shape[1]}"
+        )
+
+    mask = mask.to(device=rgb.device, dtype=rgb.dtype).clamp(0, 1)
+    alpha = (1.0 - mask).unsqueeze(-1).clamp(0, 1)  # [B,H,W,1]
+
+    rgba = torch.cat([rgb.clamp(0, 1), alpha], dim=-1)  # [B,H,W,4]
+    return rgba
+
+
+# -------------------------------------------------------------------------------------
+# Core lazy loaders (checkpoint + controlnet), cached globally
+# -------------------------------------------------------------------------------------
+
+def _load_checkpoint_cached(ckpt_name: str):
+    """
+    Mirrors comfy-core CheckpointLoaderSimple, but cached to avoid double-loads.
+    Returns: (model, clip, vae)
+    """
+    with _CKPT_LOCK:
+        if ckpt_name in _CKPT_CACHE:
+            return _CKPT_CACHE[ckpt_name]
+
+        import nodes
+        loader = nodes.CheckpointLoaderSimple()
+        fn = getattr(loader, loader.FUNCTION)
+        model, clip, vae = fn(ckpt_name=ckpt_name)
+
+        _CKPT_CACHE[ckpt_name] = (model, clip, vae)
+        return model, clip, vae
+
+
+def _load_controlnet_cached(control_net_name: str):
+    """
+    Mirrors comfy-core ControlNetLoader, but cached to avoid double-loads.
+    Returns: controlnet
+    """
+    with _CN_LOCK:
+        if control_net_name in _CN_CACHE:
+            return _CN_CACHE[control_net_name]
+
+        import nodes
+        loader = nodes.ControlNetLoader()
+        fn = getattr(loader, loader.FUNCTION)
+        (cn,) = fn(control_net_name=control_net_name)
+
+        _CN_CACHE[control_net_name] = cn
+        return cn
+
+
+# -------------------------------------------------------------------------------------
+# Assets/images dropdown + loader (inlined; no LoadImage_SaliaOnline_Assets dependency)
+# -------------------------------------------------------------------------------------
+
+def _assets_images_dir() -> Path:
+    return PLUGIN_ROOT / "assets" / "images"
+
+
+def _list_asset_pngs() -> list:
+    img_dir = _assets_images_dir()
+    if not img_dir.is_dir():
+        return []
+    files = []
+    for p in img_dir.rglob("*"):
+        if p.is_file() and p.suffix.lower() == ".png":
+            files.append(p.relative_to(img_dir).as_posix())
+    files.sort()
+    return files
+
+
+def _safe_asset_path(asset_rel_path: str) -> Path:
+    img_dir = _assets_images_dir()
+    if not img_dir.is_dir():
+        raise FileNotFoundError(f"assets/images folder not found: {img_dir}")
+
+    base = img_dir.resolve()
+    rel = Path(asset_rel_path)
+
+    if rel.is_absolute():
+        raise ValueError("Absolute paths are not allowed for asset_image.")
+
+    full = (base / rel).resolve()
+
+    # path traversal protection
+    if base != full and base not in full.parents:
+        raise ValueError(f"Invalid asset path (path traversal blocked): {asset_rel_path}")
+
+    if not full.is_file():
+        raise FileNotFoundError(f"Asset PNG not found in assets/images: {asset_rel_path}")
+    if full.suffix.lower() != ".png":
+        raise ValueError(f"Asset is not a PNG: {asset_rel_path}")
+
+    return full
+
+
+def _load_asset_image_and_mask(asset_rel_path: str) -> Tuple[torch.Tensor, torch.Tensor]:
+    """
+    Returns (IMAGE, MASK) in ComfyUI formats.
+
+    Mask semantics: match ComfyUI core LoadImage:
+      - alpha is RGBA alpha channel normalized to [0..1]
+      - mask = 1 - alpha
+    """
+    p = _safe_asset_path(asset_rel_path)
+
+    im = Image.open(p)
+    im = ImageOps.exif_transpose(im)
+
+    rgba = im.convert("RGBA")
+    rgb = rgba.convert("RGB")
+
+    rgb_arr = np.array(rgb).astype(np.float32) / 255.0  # [H,W,3]
+    img_t = torch.from_numpy(rgb_arr)[None, ...]
+
+    alpha = np.array(rgba.getchannel("A")).astype(np.float32) / 255.0  # [H,W]
+    mask = 1.0 - alpha  # Comfy MASK convention
+
+    mask_t = torch.from_numpy(mask)[None, ...]
+    return img_t, mask_t
+
+
+# -------------------------------------------------------------------------------------
+# Salia_Depth (INLINED, no imports from other files)
+# -------------------------------------------------------------------------------------
+
+MODEL_DIR = PLUGIN_ROOT / "assets" / "depth"
+MODEL_DIR.mkdir(parents=True, exist_ok=True)
+
+REQUIRED_FILES = {
+    "config.json": "https://huggingface.co/saliacoel/depth/resolve/main/config.json",
+    "model.safetensors": "https://huggingface.co/saliacoel/depth/resolve/main/model.safetensors",
+    "preprocessor_config.json": "https://huggingface.co/saliacoel/depth/resolve/main/preprocessor_config.json",
+}
+
+ZOE_FALLBACK_REPO_ID = "Intel/zoedepth-nyu-kitti"
+
+_PIPE_CACHE: Dict[Tuple[str, str], Any] = {}  # (model_source, device_str) -> pipeline
+_PIPE_LOCK = threading.Lock()
+
+
+def _have_required_files() -> bool:
+    return all((MODEL_DIR / name).exists() for name in REQUIRED_FILES.keys())
+
+
+def _download_url_to_file(url: str, dst: Path, timeout: int = 180) -> None:
+    dst.parent.mkdir(parents=True, exist_ok=True)
+    tmp = dst.with_suffix(dst.suffix + ".tmp")
+
+    if tmp.exists():
+        try:
+            tmp.unlink()
+        except Exception:
+            pass
+
+    req = urllib.request.Request(url, headers={"User-Agent": "ComfyUI-SaliaDepth/1.1"})
+    with urllib.request.urlopen(req, timeout=timeout) as r, open(tmp, "wb") as f:
+        shutil.copyfileobj(r, f)
+
+    tmp.replace(dst)
+
+
+def ensure_local_model_files() -> bool:
+    if _have_required_files():
+        return True
+    try:
+        for fname, url in REQUIRED_FILES.items():
+            fpath = MODEL_DIR / fname
+            if fpath.exists():
+                continue
+            _download_url_to_file(url, fpath)
+        return _have_required_files()
+    except Exception:
+        return False
+
+
+def HWC3(x: np.ndarray) -> np.ndarray:
+    assert x.dtype == np.uint8
+    if x.ndim == 2:
+        x = x[:, :, None]
+    assert x.ndim == 3
+    H, W, C = x.shape
+    assert C == 1 or C == 3 or C == 4
+    if C == 3:
+        return x
+    if C == 1:
+        return np.concatenate([x, x, x], axis=2)
+    # C == 4
+    color = x[:, :, 0:3].astype(np.float32)
+    alpha = x[:, :, 3:4].astype(np.float32) / 255.0
+    y = color * alpha + 255.0 * (1.0 - alpha)  # white background
+    y = y.clip(0, 255).astype(np.uint8)
+    return y
+
+
+def pad64(x: int) -> int:
+    return int(np.ceil(float(x) / 64.0) * 64 - x)
+
+
+def safer_memory(x: np.ndarray) -> np.ndarray:
+    return np.ascontiguousarray(x.copy()).copy()
+
+
+def resize_image_with_pad_min_side(
+    input_image: np.ndarray,
+    resolution: int,
+    upscale_method: str = "INTER_CUBIC",
+    skip_hwc3: bool = False,
+    mode: str = "edge",
+) -> Tuple[np.ndarray, Any]:
+    cv2 = None
+    try:
+        import cv2 as _cv2
+        cv2 = _cv2
+    except Exception:
+        cv2 = None
+
+    img = input_image if skip_hwc3 else HWC3(input_image)
+
+    H_raw, W_raw, _ = img.shape
+    if resolution <= 0:
+        return img, (lambda x: x)
+
+    k = float(resolution) / float(min(H_raw, W_raw))
+    H_target = int(np.round(float(H_raw) * k))
+    W_target = int(np.round(float(W_raw) * k))
+
+    if cv2 is not None:
+        upscale_methods = {
+            "INTER_NEAREST": cv2.INTER_NEAREST,
+            "INTER_LINEAR": cv2.INTER_LINEAR,
+            "INTER_AREA": cv2.INTER_AREA,
+            "INTER_CUBIC": cv2.INTER_CUBIC,
+            "INTER_LANCZOS4": cv2.INTER_LANCZOS4,
+        }
+        method = upscale_methods.get(upscale_method, cv2.INTER_CUBIC)
+        img = cv2.resize(img, (W_target, H_target), interpolation=method if k > 1 else cv2.INTER_AREA)
+    else:
+        pil = Image.fromarray(img)
+        resample = Image.BICUBIC if k > 1 else Image.LANCZOS
+        pil = pil.resize((W_target, H_target), resample=resample)
+        img = np.array(pil, dtype=np.uint8)
+
+    H_pad, W_pad = pad64(H_target), pad64(W_target)
+    img_padded = np.pad(img, [[0, H_pad], [0, W_pad], [0, 0]], mode=mode)
+
+    def remove_pad(x: np.ndarray) -> np.ndarray:
+        return safer_memory(x[:H_target, :W_target, ...])
+
+    return safer_memory(img_padded), remove_pad
+
+
+def pad_only_to_64(img_u8: np.ndarray, mode: str = "edge") -> Tuple[np.ndarray, Any]:
+    img = HWC3(img_u8)
+    H_raw, W_raw, _ = img.shape
+    H_pad, W_pad = pad64(H_raw), pad64(W_raw)
+    img_padded = np.pad(img, [[0, H_pad], [0, W_pad], [0, 0]], mode=mode)
+
+    def remove_pad(x: np.ndarray) -> np.ndarray:
+        return safer_memory(x[:H_raw, :W_raw, ...])
+
+    return safer_memory(img_padded), remove_pad
+
+
+def composite_rgba_over_white_keep_alpha(inp_u8: np.ndarray) -> Tuple[np.ndarray, Optional[np.ndarray]]:
+    if inp_u8.ndim == 3 and inp_u8.shape[2] == 4:
+        rgba = inp_u8.astype(np.uint8)
+        rgb = rgba[:, :, 0:3].astype(np.float32)
+        a = (rgba[:, :, 3:4].astype(np.float32) / 255.0)
+        rgb_white = (rgb * a + 255.0 * (1.0 - a)).clip(0, 255).astype(np.uint8)
+        alpha_u8 = rgba[:, :, 3].copy()
+        return rgb_white, alpha_u8
+    return HWC3(inp_u8), None
+
+
+def apply_alpha_then_black_background(depth_rgb_u8: np.ndarray, alpha_u8: np.ndarray) -> np.ndarray:
+    depth_rgb_u8 = HWC3(depth_rgb_u8)
+    a = (alpha_u8.astype(np.float32) / 255.0)[:, :, None]
+    out = (depth_rgb_u8.astype(np.float32) * a).clip(0, 255).astype(np.uint8)
+    return out
+
+
+def comfy_tensor_to_u8(img: torch.Tensor) -> np.ndarray:
+    if img.ndim == 4:
+        img = img[0]
+    arr = img.detach().cpu().float().clamp(0, 1).numpy()
+    u8 = (arr * 255.0).round().astype(np.uint8)
+    return u8
+
+
+def u8_to_comfy_tensor(img_u8: np.ndarray) -> torch.Tensor:
+    img_u8 = HWC3(img_u8)
+    t = torch.from_numpy(img_u8.astype(np.float32) / 255.0)
+    return t.unsqueeze(0)  # [1,H,W,C]
+
+
+def _try_load_pipeline(model_source: str, device: torch.device):
+    if pipeline is None:
+        raise RuntimeError(f"transformers import failed: {_TRANSFORMERS_IMPORT_ERROR}")
+
+    key = (model_source, str(device))
+    with _PIPE_LOCK:
+        if key in _PIPE_CACHE:
+            return _PIPE_CACHE[key]
+
+        p = pipeline(task="depth-estimation", model=model_source)
+        try:
+            p.model = p.model.to(device)
+            p.device = device
+        except Exception:
+            pass
+
+        _PIPE_CACHE[key] = p
+        return p
+
+
+def get_depth_pipeline(device: torch.device):
+    if ensure_local_model_files():
+        try:
+            return _try_load_pipeline(str(MODEL_DIR), device)
+        except Exception:
+            pass
+    try:
+        return _try_load_pipeline(ZOE_FALLBACK_REPO_ID, device)
+    except Exception:
+        return None
+
+
+def depth_estimate_zoe_style(
+    pipe,
+    input_rgb_u8: np.ndarray,
+    detect_resolution: int,
+    upscale_method: str = "INTER_CUBIC",
+) -> np.ndarray:
+    if detect_resolution == -1:
+        work_img, remove_pad = pad_only_to_64(input_rgb_u8, mode="edge")
+    else:
+        work_img, remove_pad = resize_image_with_pad_min_side(
+            input_rgb_u8,
+            int(detect_resolution),
+            upscale_method=upscale_method,
+            skip_hwc3=False,
+            mode="edge",
+        )
+
+    pil_image = Image.fromarray(work_img)
+
+    with torch.no_grad():
+        result = pipe(pil_image)
+        depth = result["depth"]
+
+        if isinstance(depth, Image.Image):
+            depth_array = np.array(depth, dtype=np.float32)
+        else:
+            depth_array = np.array(depth, dtype=np.float32)
+
+        vmin = float(np.percentile(depth_array, 2))
+        vmax = float(np.percentile(depth_array, 85))
+
+        depth_array = depth_array - vmin
+        denom = (vmax - vmin)
+        if abs(denom) < 1e-12:
+            denom = 1e-6
+        depth_array = depth_array / denom
+
+        depth_array = 1.0 - depth_array
+        depth_image = (depth_array * 255.0).clip(0, 255).astype(np.uint8)
+
+    detected_map = remove_pad(HWC3(depth_image))
+    return detected_map
+
+
+def resize_to_original(depth_rgb_u8: np.ndarray, w0: int, h0: int) -> np.ndarray:
+    try:
+        import cv2
+        out = cv2.resize(depth_rgb_u8, (w0, h0), interpolation=cv2.INTER_LINEAR)
+        return out.astype(np.uint8)
+    except Exception:
+        pil = Image.fromarray(depth_rgb_u8)
+        pil = pil.resize((w0, h0), resample=Image.BILINEAR)
+        return np.array(pil, dtype=np.uint8)
+
+
+def _salia_depth_execute(image: torch.Tensor, resolution: int = -1) -> torch.Tensor:
+    """
+    Internal callable version of Salia_Depth:
+      input: IMAGE [B,H,W,3 or 4]
+      output: IMAGE [B,H,W,3]
+    """
+    try:
+        device = model_management.get_torch_device()
+    except Exception:
+        device = torch.device("cpu")
+
+    pipe_obj = None
+    try:
+        pipe_obj = get_depth_pipeline(device)
+    except Exception:
+        pipe_obj = None
+
+    if pipe_obj is None:
+        return image
+
+    if image.ndim == 3:
+        image = image.unsqueeze(0)
+
+    outs = []
+    for i in range(image.shape[0]):
+        try:
+            h0 = int(image[i].shape[0])
+            w0 = int(image[i].shape[1])
+
+            inp_u8 = comfy_tensor_to_u8(image[i])
+
+            rgb_for_depth, alpha_u8 = composite_rgba_over_white_keep_alpha(inp_u8)
+            had_rgba = alpha_u8 is not None
+
+            depth_rgb = depth_estimate_zoe_style(
+                pipe=pipe_obj,
+                input_rgb_u8=rgb_for_depth,
+                detect_resolution=int(resolution),
+                upscale_method="INTER_CUBIC",
+            )
+
+            depth_rgb = resize_to_original(depth_rgb, w0=w0, h0=h0)
+
+            if had_rgba:
+                if alpha_u8.shape[0] != h0 or alpha_u8.shape[1] != w0:
+                    try:
+                        import cv2
+                        alpha_u8 = cv2.resize(alpha_u8, (w0, h0), interpolation=cv2.INTER_LINEAR).astype(np.uint8)
+                    except Exception:
+                        pil_a = Image.fromarray(alpha_u8)
+                        pil_a = pil_a.resize((w0, h0), resample=Image.BILINEAR)
+                        alpha_u8 = np.array(pil_a, dtype=np.uint8)
+
+                depth_rgb = apply_alpha_then_black_background(depth_rgb, alpha_u8)
+
+            outs.append(u8_to_comfy_tensor(depth_rgb))
+        except Exception:
+            outs.append(image[i].unsqueeze(0))
+
+    return torch.cat(outs, dim=0)
+
+
+# -------------------------------------------------------------------------------------
+# Alpha-over paste (RGBA square onto base at X,Y)
+# -------------------------------------------------------------------------------------
+
+def _alpha_over_region(base: torch.Tensor, overlay_rgba: torch.Tensor, x: int, y: int) -> torch.Tensor:
+    """
+    base: [B,H,W,C] where C is 3 or 4, float [0..1]
+    overlay_rgba: [B,s,s,4] float [0..1]
+    """
+    if base.ndim != 4 or overlay_rgba.ndim != 4:
+        raise ValueError("base and overlay must be [B,H,W,C].")
+
+    B, H, W, C = base.shape
+    b2, sH, sW, c2 = overlay_rgba.shape
+    if c2 != 4:
+        raise ValueError("overlay_rgba must have 4 channels (RGBA).")
+    if sH != sW:
+        raise ValueError("overlay must be square.")
+    s = sH
+
+    if x < 0 or y < 0 or x + s > W or y + s > H:
+        raise ValueError(f"Square paste out of bounds. base={W}x{H}, paste at ({x},{y}) size={s}")
+
+    if b2 != B:
+        if b2 == 1 and B > 1:
+            overlay_rgba = overlay_rgba.expand(B, -1, -1, -1)
+        else:
+            raise ValueError("Batch mismatch between base and overlay.")
+
+    out = base.clone()
+
+    overlay_rgb = overlay_rgba[..., 0:3].clamp(0, 1)
+    overlay_a = overlay_rgba[..., 3:4].clamp(0, 1)
+
+    base_rgb = out[:, y:y + s, x:x + s, 0:3]
+    comp_rgb = overlay_rgb * overlay_a + base_rgb * (1.0 - overlay_a)
+    out[:, y:y + s, x:x + s, 0:3] = comp_rgb
+
+    if C == 4:
+        base_a = out[:, y:y + s, x:x + s, 3:4].clamp(0, 1)
+        comp_a = overlay_a + base_a * (1.0 - overlay_a)
+        out[:, y:y + s, x:x + s, 3:4] = comp_a
+
+    return out.clamp(0, 1)
+
+
+# -------------------------------------------------------------------------------------
+# The One-Node Workflow
+# -------------------------------------------------------------------------------------
+
+class Salia_ezpz_gated:
+    CATEGORY = "image/salia"
+    RETURN_TYPES = ("IMAGE",)
+    RETURN_NAMES = ("image",)
+    FUNCTION = "run"
+
+    @classmethod
+    def INPUT_TYPES(cls):
+        ckpts = folder_paths.get_filename_list("checkpoints") or ["<no checkpoints found>"]
+        cns = folder_paths.get_filename_list("controlnet") or ["<no controlnets found>"]
+        assets = _list_asset_pngs() or ["<no pngs found>"]
+
+        try:
+            import comfy.samplers
+            sampler_names = comfy.samplers.KSampler.SAMPLERS
+            scheduler_names = comfy.samplers.KSampler.SCHEDULERS
+        except Exception:
+            sampler_names = ["euler"]
+            scheduler_names = ["karras"]
+
+        upscale_choices = ["1", "2", "4", "6", "8", "10", "12", "14", "16"]
+
+        return {
+            "required": {
+                "image": ("IMAGE",),
+                "trigger_string": ("STRING", {"default": ""}),
+
+                "X_coord": ("INT", {"default": 0, "min": 0, "max": 16384, "step": 1}),
+                "Y_coord": ("INT", {"default": 0, "min": 0, "max": 16384, "step": 1}),
+                "square_size": ("INT", {"default": 384, "min": 8, "max": 8192, "step": 1}),
+
+                "positive_prompt": ("STRING", {"default": "", "multiline": True}),
+                "negative_prompt": ("STRING", {"default": "", "multiline": True}),
+
+                "upscale_factor": (upscale_choices, {"default": "4"}),
+
+                "ckpt_name": (ckpts, {}),
+                "control_net_name": (cns, {}),
+                "asset_image": (assets, {}),
+
+                "controlnet_strength": ("FLOAT", {"default": 0.33, "min": 0.00, "max": 10.00, "step": 0.01}),
+                "controlnet_start_percent": ("FLOAT", {"default": 0.00, "min": 0.00, "max": 1.00, "step": 0.01}),
+                "controlnet_end_percent": ("FLOAT", {"default": 1.00, "min": 0.00, "max": 1.00, "step": 0.01}),
+
+                "steps": ("INT", {"default": 30, "min": 1, "max": 200, "step": 1}),
+                "cfg": ("FLOAT", {"default": 2.6, "min": 0.00, "max": 10.00, "step": 0.05}),
+                "sampler_name": (sampler_names, {"default": "euler"} if "euler" in sampler_names else {}),
+                "scheduler": (scheduler_names, {"default": "karras"} if "karras" in scheduler_names else {}),
+                "denoise": ("FLOAT", {"default": 0.35, "min": 0.00, "max": 1.00, "step": 0.01}),
+            }
+        }
+
+    def run(
+        self,
+        image: torch.Tensor,
+        trigger_string: str = "",
+        X_coord: int = 0,
+        Y_coord: int = 0,
+        square_size: int = 384,
+        positive_prompt: str = "",
+        negative_prompt: str = "",
+        upscale_factor: str = "4",
+        ckpt_name: str = "",
+        control_net_name: str = "",
+        asset_image: str = "",
+        controlnet_strength: float = 0.33,
+        controlnet_start_percent: float = 0.0,
+        controlnet_end_percent: float = 1.0,
+        steps: int = 30,
+        cfg: float = 2.6,
+        sampler_name: str = "euler",
+        scheduler: str = "karras",
+        denoise: float = 0.35,
+    ):
+        # If trigger_string is exactly empty, bypass everything and return input unchanged.
+        if trigger_string == "":
+            return (image,)
+
+        # Normalize input to [B,H,W,C]
+        if image.ndim == 3:
+            image = image.unsqueeze(0)
+        if image.ndim != 4:
+            raise ValueError("Input image must be [B,H,W,C].")
+
+        B, H, W, C = image.shape
+        if C not in (3, 4):
+            raise ValueError("Input image must have 3 (RGB) or 4 (RGBA) channels.")
+
+        x = int(X_coord)
+        y = int(Y_coord)
+        s = int(square_size)
+
+        up = int(upscale_factor)
+        if up not in (1, 2, 4, 6, 8, 10, 12, 14, 16):
+            raise ValueError("upscale_factor must be one of: 1,2,4,6,8,10,12,14,16")
+
+        if s <= 0:
+            raise ValueError("square_size must be > 0")
+        if x < 0 or y < 0 or x + s > W or y + s > H:
+            raise ValueError(f"Crop out of bounds. image={W}x{H}, crop at ({x},{y}) size={s}")
+
+        up_w = s * up
+        up_h = s * up
+
+        # VAE/UNet path likes multiples of 8
+        if (up_w % 8) != 0 or (up_h % 8) != 0:
+            raise ValueError("square_size * upscale_factor must be divisible by 8 (required by VAE pipeline).")
+
+        start_p = float(max(0.0, min(1.0, controlnet_start_percent)))
+        end_p = float(max(0.0, min(1.0, controlnet_end_percent)))
+        if end_p < start_p:
+            start_p, end_p = end_p, start_p
+
+        # 1) Crop square
+        crop = image[:, y:y + s, x:x + s, :]
+        crop_rgb = crop[:, :, :, 0:3].contiguous()
+
+        # 2) Depth (inline Salia_Depth) then upscale with Lanczos
+        depth_small = _salia_depth_execute(crop_rgb, resolution=s)
+        depth_up = _resize_image_lanczos(depth_small, up_w, up_h)
+
+        # 3) Upscale crop for VAE encode
+        crop_up = _resize_image_lanczos(crop_rgb, up_w, up_h)
+
+        # 4) Load asset mask and resize
+        if asset_image == "<no pngs found>":
+            raise FileNotFoundError("No PNGs found in assets/images for this plugin.")
+        _asset_img_unused, asset_mask = _load_asset_image_and_mask(asset_image)
+
+        if asset_mask.ndim == 2:
+            asset_mask = asset_mask.unsqueeze(0)
+        if asset_mask.ndim != 3:
+            raise ValueError("Asset mask must be [B,H,W].")
+
+        if asset_mask.shape[0] != B:
+            if asset_mask.shape[0] == 1 and B > 1:
+                asset_mask = asset_mask.expand(B, -1, -1)
+            else:
+                raise ValueError("Batch mismatch for asset mask.")
+
+        asset_mask_up = _resize_mask_lanczos(asset_mask, up_w, up_h)
+
+        # 5) Load checkpoint + controlnet (cached)
+        if ckpt_name == "<no checkpoints found>":
+            raise FileNotFoundError("No checkpoints found in models/checkpoints.")
+        if control_net_name == "<no controlnets found>":
+            raise FileNotFoundError("No controlnets found in models/controlnet.")
+
+        model, clip, vae = _load_checkpoint_cached(ckpt_name)
+        controlnet = _load_controlnet_cached(control_net_name)
+
+        import nodes
+
+        # 6) CLIP encodes
+        pos_enc = nodes.CLIPTextEncode()
+        neg_enc = nodes.CLIPTextEncode()
+        pos_fn = getattr(pos_enc, pos_enc.FUNCTION)
+        neg_fn = getattr(neg_enc, neg_enc.FUNCTION)
+        (pos_cond,) = pos_fn(text=str(positive_prompt), clip=clip)
+        (neg_cond,) = neg_fn(text=str(negative_prompt), clip=clip)
+
+        # 7) Apply ControlNet
+        cn_apply = nodes.ControlNetApplyAdvanced()
+        cn_fn = getattr(cn_apply, cn_apply.FUNCTION)
+        pos_cn, neg_cn = cn_fn(
+            strength=float(controlnet_strength),
+            start_percent=float(start_p),
+            end_percent=float(end_p),
+            positive=pos_cond,
+            negative=neg_cond,
+            control_net=controlnet,
+            image=depth_up,
+            vae=vae,
+        )
+
+        # 8) VAE Encode
+        vae_enc = nodes.VAEEncode()
+        vae_enc_fn = getattr(vae_enc, vae_enc.FUNCTION)
+        (latent,) = vae_enc_fn(pixels=crop_up, vae=vae)
+
+        # 9) KSampler (deterministic seed derived from inputs)
+        seed_material = (
+            f"{ckpt_name}|{control_net_name}|{asset_image}|{x}|{y}|{s}|{up}|"
+            f"{steps}|{cfg}|{sampler_name}|{scheduler}|{denoise}|"
+            f"{controlnet_strength}|{start_p}|{end_p}|"
+            f"{positive_prompt}|{negative_prompt}"
+        ).encode("utf-8", errors="ignore")
+        seed64 = int(hashlib.sha256(seed_material).hexdigest()[:16], 16)
+
+        ksampler = nodes.KSampler()
+        k_fn = getattr(ksampler, ksampler.FUNCTION)
+        (sampled_latent,) = k_fn(
+            seed=seed64,
+            steps=int(steps),
+            cfg=float(cfg),
+            sampler_name=str(sampler_name),
+            scheduler=str(scheduler),
+            denoise=float(denoise),
+            model=model,
+            positive=pos_cn,
+            negative=neg_cn,
+            latent_image=latent,
+        )
+
+        # 10) VAE Decode -> RGB
+        vae_dec = nodes.VAEDecode()
+        vae_dec_fn = getattr(vae_dec, vae_dec.FUNCTION)
+        (decoded_rgb,) = vae_dec_fn(samples=sampled_latent, vae=vae)
+
+        # 11) ✅ Manual "JoinImageWithAlpha"
+        rgba_up = _rgb_to_rgba_with_comfy_mask(decoded_rgb, asset_mask_up)
+
+        # 12) Downscale RGBA back to crop size
+        rgba_square = _resize_image_lanczos(rgba_up, s, s)
+
+        # 13) Paste back onto original at X,Y (alpha-over)
+        out = _alpha_over_region(image, rgba_square, x=x, y=y)
+        return (out,)
+
+
+NODE_CLASS_MAPPINGS = {
+    "Salia_ezpz_gated": Salia_ezpz_gated,
+}
+
+NODE_DISPLAY_NAME_MAPPINGS = {
+    "Salia_ezpz_gated": "Salia EZPZ Gated",
+}