Upload salia_detailer_ezpz.py

salia_detailer_ezpz.py

@@ -4,7 +4,7 @@ from typing import Any, Dict, Tuple, Optional
 
 import torch
 import numpy as np
-from PIL import Image
+from PIL import Image, ImageOps
 
 import folder_paths
 
@@ -149,45 +149,131 @@ def _load_controlnet_cached(control_net_name: str):
 
 
 # -------------------------------------------------------------------------------------
-# Asset dropdown support (from comfyui-salia_online assets/images)
-# (We still lazy-call the user's LoadImage_SaliaOnline_Assets for consistent mask behavior.)
+# Assets/images dropdown + loader (INLINED, no LoadImage_SaliaOnline_Assets dependency)
 # -------------------------------------------------------------------------------------
 
-def _list_asset_pngs_fallback():
-    # Fallback scanner (if utils import fails)
+_ASSETS_DIR_CACHE: Optional["object"] = None
+_ASSETS_DIR_LOCK = threading.Lock()
+
+
+def _find_assets_images_dir():
+    """
+    Find the plugin's assets/images folder by walking upward from this file.
+    This is robust even if Comfy imports modules in weird ways.
+    """
+    from pathlib import Path
+
+    here = Path(__file__).resolve()
+    # check a few levels up; plugin root should be near
+    for parent in [here.parent] + list(here.parents)[:8]:
+        candidate = parent / "assets" / "images"
+        if candidate.is_dir():
+            return candidate
+    return None
+
+
+def _assets_images_dir():
+    global _ASSETS_DIR_CACHE
+    with _ASSETS_DIR_LOCK:
+        if _ASSETS_DIR_CACHE is not None:
+            # If it was found once, reuse.
+            try:
+                if _ASSETS_DIR_CACHE.is_dir():
+                    return _ASSETS_DIR_CACHE
+            except Exception:
+                pass
+
+        found = _find_assets_images_dir()
+        _ASSETS_DIR_CACHE = found
+        return found
+
+
+def _list_asset_pngs():
+    """
+    List PNGs inside assets/images (recursive), returning paths relative to assets/images.
+    """
+    img_dir = _assets_images_dir()
+    if img_dir is None:
+        return []
+
+    files = []
     try:
-        from pathlib import Path
-        plugin_root = Path(__file__).resolve().parent.parent
-        img_dir = plugin_root / "assets" / "images"
-        if not img_dir.exists():
-            return []
-        files = sorted([p.name for p in img_dir.glob("*.png")])
+        for p in img_dir.rglob("*"):
+            if p.is_file() and p.suffix.lower() == ".png":
+                rel = p.relative_to(img_dir).as_posix()
+                files.append(rel)
+        files.sort()
         return files
     except Exception:
         return []
 
 
-def _list_asset_pngs():
-    try:
-        # Prefer your plugin's own list function (same dropdown as your node)
-        from ..utils.io import list_pngs  # type: ignore
-        return list_pngs() or []
-    except Exception:
-        return _list_asset_pngs_fallback()
+def _safe_asset_path(asset_rel_path: str):
+    """
+    Resolve a selected dropdown entry to an actual file path inside assets/images.
+    Prevents path traversal.
+    """
+    from pathlib import Path
+
+    img_dir = _assets_images_dir()
+    if img_dir is None:
+        raise FileNotFoundError("assets/images folder not found (could not locate plugin assets).")
+
+    base = img_dir.resolve()
+    rel = Path(asset_rel_path)
+
+    if rel.is_absolute():
+        raise ValueError("Absolute paths are not allowed for asset_image.")
+
+    # Resolve and verify containment
+    full = (base / rel).resolve()
+    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_mask(asset_name: str):
+def _load_asset_image_and_mask(asset_rel_path: str):
     """
-    Lazy-import and run your LoadImage_SaliaOnline_Assets node.
-    Returns: MASK
+    Load PNG from assets/images and return (IMAGE, MASK) in ComfyUI formats.
+
+    IMPORTANT: Mask semantics match ComfyUI core LoadImage:
+      - If PNG has alpha: mask = 1 - alpha
+      - If no alpha: mask = 0 (opaque)
     """
-    # NOTE: Keep this lazy so importing the plugin doesn't force-load anything.
-    from .salia_loadimage_assets import LoadImage_SaliaOnline_Assets  # lazy-ish (light)
+    p = _safe_asset_path(asset_rel_path)
+
+    im = Image.open(p)
+    im = ImageOps.exif_transpose(im)
+
+    # Ensure we can extract alpha if present
+    had_alpha = ("A" in im.getbands())
+    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, ...]
 
-    loader = LoadImage_SaliaOnline_Assets()
-    img, mask = loader.run(asset_name)
-    return mask
+    if had_alpha:
+        alpha = np.array(rgba.getchannel("A")).astype(np.float32) / 255.0  # [H,W], 1=opaque
+        mask = 1.0 - alpha  # Comfy MASK convention
+    else:
+        h, w = rgb.size[1], rgb.size[0]
+        mask = np.zeros((h, w), dtype=np.float32)
 
+    mask_t = torch.from_numpy(mask)[None, ...]
+    return img_t, mask_t
+
+
+# -------------------------------------------------------------------------------------
+# Salia_Depth (still lazy import, unchanged)
+# -------------------------------------------------------------------------------------
 
 def _run_salia_depth(image: torch.Tensor, resolution: int) -> torch.Tensor:
     """
@@ -394,12 +480,12 @@ class Salia_Detailer_EZPZ:
         crop_up = _resize_image_lanczos(crop_rgb, up_w, up_h)
 
         # -------------------------
-        # 4) Load asset mask (dropdown) and resize it to match upscaled resolution
+        # 4) Load asset mask (INLINE assets loader) and resize to match upscaled resolution
         # -------------------------
         if asset_image == "<no pngs found>":
-            raise FileNotFoundError("No PNGs found in comfyui-salia_online/assets/images")
+            raise FileNotFoundError("No PNGs found in assets/images for this plugin.")
 
-        asset_mask = _load_asset_mask(asset_image)  # MASK
+        _asset_img_unused, asset_mask = _load_asset_image_and_mask(asset_image)  # MASK is what we need
         if asset_mask.ndim == 2:
             asset_mask = asset_mask.unsqueeze(0)
         if asset_mask.ndim != 3:
@@ -466,7 +552,6 @@ class Salia_Detailer_EZPZ:
         # -------------------------
         # 9) KSampler
         # -------------------------
-        # No seed input requested: derive a stable seed from inputs so changing anything changes seed.
         seed_material = (
             f"{ckpt_name}|{control_net_name}|{asset_image}|{x}|{y}|{s}|{up}|"
             f"{steps}|{cfg}|{sampler_name}|{scheduler}|{denoise}|"
@@ -503,7 +588,6 @@ class Salia_Detailer_EZPZ:
         join = nodes.JoinImageWithAlpha()
         join_fn = getattr(join, join.FUNCTION)
 
-        # Some Comfy versions name the mask input "alpha", others "mask".
         try:
             (rgba_up,) = join_fn(image=decoded_rgb, alpha=asset_mask_up)
         except TypeError: