Update salia_depth.py

salia_depth.py

@@ -1,611 +1,611 @@
-import os
-import shutil
-import urllib.request
-from pathlib import Path
-from typing import Dict, Tuple, Any, Optional, List
-
-import numpy as np
-import torch
-from PIL import Image
-
-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
-
-
-# --------------------------------------------------------------------------------------
-# Paths / sources
-# --------------------------------------------------------------------------------------
-
-# This file: comfyui-salia_online/nodes/Salia_Depth.py
-# Plugin root: comfyui-salia_online/
-PLUGIN_ROOT = Path(__file__).resolve().parent.parent
-
-# Requested local path: assets/depth
-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-path" fallback
-ZOE_FALLBACK_REPO_ID = "Intel/zoedepth-nyu-kitti"
-
-
-# --------------------------------------------------------------------------------------
-# Logging helpers
-# --------------------------------------------------------------------------------------
-
-def _make_logger() -> Tuple[List[str], Any]:
-    lines: List[str] = []
-
-    def log(msg: str):
-        # console
-        try:
-            print(msg)
-        except Exception:
-            pass
-        # UI string
-        lines.append(str(msg))
-
-    return lines, log
-
-
-def _fmt_bytes(n: Optional[int]) -> str:
-    if n is None:
-        return "?"
-    # simple readable
-    for unit in ["B", "KB", "MB", "GB", "TB"]:
-        if n < 1024:
-            return f"{n:.0f}{unit}"
-        n /= 1024.0
-    return f"{n:.1f}PB"
-
-
-def _file_size(path: Path) -> Optional[int]:
-    try:
-        return path.stat().st_size
-    except Exception:
-        return None
-
-
-def _hf_cache_info() -> Dict[str, str]:
-    info: Dict[str, str] = {}
-    info["env.HF_HOME"] = os.environ.get("HF_HOME", "")
-    info["env.HF_HUB_CACHE"] = os.environ.get("HF_HUB_CACHE", "")
-    info["env.TRANSFORMERS_CACHE"] = os.environ.get("TRANSFORMERS_CACHE", "")
-    info["env.HUGGINGFACE_HUB_CACHE"] = os.environ.get("HUGGINGFACE_HUB_CACHE", "")
-
-    try:
-        from huggingface_hub import constants as hf_constants
-        # These exist in most hub versions:
-        info["huggingface_hub.constants.HF_HOME"] = str(getattr(hf_constants, "HF_HOME", ""))
-        info["huggingface_hub.constants.HF_HUB_CACHE"] = str(getattr(hf_constants, "HF_HUB_CACHE", ""))
-    except Exception:
-        pass
-
-    return info
-
-
-# --------------------------------------------------------------------------------------
-# Download helpers
-# --------------------------------------------------------------------------------------
-
-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:
-    """
-    Download with atomic temp rename.
-    """
-    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(log) -> bool:
-    """
-    Ensure assets/depth contains the 3 files.
-    Returns True if present or downloaded successfully, else False.
-    """
-    # Always log expected locations + URLs, even if we don't download.
-    log("[SaliaDepth] ===== Local model file check =====")
-    log(f"[SaliaDepth] Plugin root: {PLUGIN_ROOT}")
-    log(f"[SaliaDepth] Local model dir (on drive): {MODEL_DIR}")
-
-    for fname, url in REQUIRED_FILES.items():
-        fpath = MODEL_DIR / fname
-        exists = fpath.exists()
-        size = _file_size(fpath) if exists else None
-        log(f"[SaliaDepth]   - {fname}")
-        log(f"[SaliaDepth]       local path: {fpath}  exists={exists}  size={_fmt_bytes(size)}")
-        log(f"[SaliaDepth]       remote url : {url}")
-
-    if _have_required_files():
-        log("[SaliaDepth] All required local files already exist. No download needed.")
-        return True
-
-    log("[SaliaDepth] One or more local files missing. Attempting download...")
-
-    try:
-        for fname, url in REQUIRED_FILES.items():
-            fpath = MODEL_DIR / fname
-            if fpath.exists():
-                continue
-            log(f"[SaliaDepth] Downloading '{fname}' -> '{fpath}'")
-            _download_url_to_file(url, fpath)
-            log(f"[SaliaDepth] Downloaded '{fname}' size={_fmt_bytes(_file_size(fpath))}")
-
-        ok = _have_required_files()
-        log(f"[SaliaDepth] Download finished. ok={ok}")
-        return ok
-    except Exception as e:
-        log(f"[SaliaDepth] Download failed with error: {repr(e)}")
-        return False
-
-
-# --------------------------------------------------------------------------------------
-# Exact Zoe-style preprocessing helpers (copied/adapted from your snippet)
-# --------------------------------------------------------------------------------------
-
-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",
-    log=None
-) -> Tuple[np.ndarray, Any]:
-    """
-    EXACT behavior like your zoe.transformers.py:
-      k = resolution / min(H,W)
-      resize to (W_target, H_target)
-      pad to multiple of 64
-      return padded image and remove_pad() closure
-    """
-    # prefer cv2 like original for matching results
-    cv2 = None
-    try:
-        import cv2 as _cv2
-        cv2 = _cv2
-    except Exception:
-        cv2 = None
-        if log:
-            log("[SaliaDepth] WARN: cv2 not available; resizing will use PIL fallback (may change results).")
-
-    if skip_hwc3:
-        img = input_image
-    else:
-        img = HWC3(input_image)
-
-    H_raw, W_raw, _ = img.shape
-    if resolution <= 0:
-        # keep original, but still pad to 64 (we will handle padding separately for -1 path)
-        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 fallback
-        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]:
-    """
-    For resolution == -1: keep original resolution but still pad to multiples of 64,
-    then provide remove_pad that returns original size.
-    """
-    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
-
-
-# --------------------------------------------------------------------------------------
-# RGBA rules (as you requested)
-# --------------------------------------------------------------------------------------
-
-def composite_rgba_over_white_keep_alpha(inp_u8: np.ndarray) -> Tuple[np.ndarray, Optional[np.ndarray]]:
-    """
-    If RGBA: return RGB composited over WHITE + alpha_u8 kept separately.
-    If RGB: return input RGB + None alpha.
-    """
-    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
-    # force to RGB
-    return HWC3(inp_u8), None
-
-
-def apply_alpha_then_black_background(depth_rgb_u8: np.ndarray, alpha_u8: np.ndarray) -> np.ndarray:
-    """
-    Requested output rule:
-      - attach alpha to depth (conceptually RGBA)
-      - composite over BLACK
-      - output RGB
-    That is equivalent to depth_rgb * alpha.
-    """
-    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
-
-
-# --------------------------------------------------------------------------------------
-# ComfyUI conversion helpers
-# --------------------------------------------------------------------------------------
-
-def comfy_tensor_to_u8(img: torch.Tensor) -> np.ndarray:
-    """
-    Comfy IMAGE: float [0..1], shape [H,W,C] or [B,H,W,C]
-    Convert to uint8 HWC.
-    """
-    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]
-
-
-# --------------------------------------------------------------------------------------
-# Pipeline loading (local-first, then zoe fallback)
-# --------------------------------------------------------------------------------------
-
-_PIPE_CACHE: Dict[Tuple[str, str], Any] = {}  # (model_source, device_str) -> pipeline
-
-
-def _try_load_pipeline(model_source: str, device: torch.device, log):
-    """
-    Use transformers.pipeline like Zoe code does.
-    We intentionally do NOT pass device=... here, and instead move model like Zoe node.
-    """
-    if pipeline is None:
-        raise RuntimeError(f"transformers import failed: {_TRANSFORMERS_IMPORT_ERROR}")
-
-    key = (model_source, str(device))
-    if key in _PIPE_CACHE:
-        log(f"[SaliaDepth] Using cached pipeline for source='{model_source}' device='{device}'")
-        return _PIPE_CACHE[key]
-
-    log(f"[SaliaDepth] Creating pipeline(task='depth-estimation', model='{model_source}')")
-    p = pipeline(task="depth-estimation", model=model_source)
-
-    # Try to move model to torch device, like ZoeDetector.to()
-    try:
-        p.model = p.model.to(device)
-        p.device = device  # Zoe code sets this; newer transformers uses torch.device internally
-        log(f"[SaliaDepth] Moved pipeline model to device: {device}")
-    except Exception as e:
-        log(f"[SaliaDepth] WARN: Could not move pipeline model to device {device}: {repr(e)}")
-
-    # Log config info for debugging
-    try:
-        cfg = p.model.config
-        log(f"[SaliaDepth] Model class: {p.model.__class__.__name__}")
-        log(f"[SaliaDepth] Config class: {cfg.__class__.__name__}")
-        log(f"[SaliaDepth] Config model_type: {getattr(cfg, 'model_type', '')}")
-        log(f"[SaliaDepth] Config _name_or_path: {getattr(cfg, '_name_or_path', '')}")
-    except Exception as e:
-        log(f"[SaliaDepth] WARN: Could not log model config: {repr(e)}")
-
-    _PIPE_CACHE[key] = p
-    return p
-
-
-def get_depth_pipeline(device: torch.device, log):
-    """
-    1) Ensure assets/depth files exist (download if missing)
-    2) Try load local dir
-    3) Fallback to Intel/zoedepth-nyu-kitti
-    4) If both fail -> None
-    """
-    # Always log HF cache info (helps locate where fallback downloads go)
-    log("[SaliaDepth] ===== Hugging Face cache info (fallback path) =====")
-    for k, v in _hf_cache_info().items():
-        if v:
-            log(f"[SaliaDepth] {k} = {v}")
-    log(f"[SaliaDepth] Zoe fallback repo id: {ZOE_FALLBACK_REPO_ID}")
-
-    # Local-first
-    local_ok = ensure_local_model_files(log)
-    if local_ok:
-        try:
-            log(f"[SaliaDepth] Trying LOCAL model from directory: {MODEL_DIR}")
-            return _try_load_pipeline(str(MODEL_DIR), device, log)
-        except Exception as e:
-            log(f"[SaliaDepth] Local model load FAILED: {repr(e)}")
-
-    # Fallback
-    try:
-        log(f"[SaliaDepth] Trying ZOE fallback model: {ZOE_FALLBACK_REPO_ID}")
-        return _try_load_pipeline(ZOE_FALLBACK_REPO_ID, device, log)
-    except Exception as e:
-        log(f"[SaliaDepth] Zoe fallback load FAILED: {repr(e)}")
-
-    return None
-
-
-# --------------------------------------------------------------------------------------
-# Depth inference (Zoe-style)
-# --------------------------------------------------------------------------------------
-
-def depth_estimate_zoe_style(
-    pipe,
-    input_rgb_u8: np.ndarray,
-    detect_resolution: int,
-    log,
-    upscale_method: str = "INTER_CUBIC"
-) -> np.ndarray:
-    """
-    Matches your ZoeDetector.__call__ logic very closely.
-    Returns uint8 RGB depth map.
-    """
-    # detect_resolution:
-    #   - if -1: keep original but pad-to-64
-    #   - else: min-side resize to detect_resolution, then pad-to-64
-    if detect_resolution == -1:
-        work_img, remove_pad = pad_only_to_64(input_rgb_u8, mode="edge")
-        log(f"[SaliaDepth] Preprocess: resolution=-1 (no resize), padded to 64. work={work_img.shape}")
-    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",
-            log=log
-        )
-        log(f"[SaliaDepth] Preprocess: min-side resized to {detect_resolution}, padded to 64. work={work_img.shape}")
-
-    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)
-
-        # EXACT normalization like your Zoe code
-        vmin = float(np.percentile(depth_array, 2))
-        vmax = float(np.percentile(depth_array, 85))
-
-        log(f"[SaliaDepth] Depth raw stats: shape={depth_array.shape} vmin(p2)={vmin:.6f} vmax(p85)={vmax:.6f} mean={float(depth_array.mean()):.6f}")
-
-        depth_array = depth_array - vmin
-        denom = (vmax - vmin)
-        if abs(denom) < 1e-12:
-            # avoid division by zero; log it
-            log("[SaliaDepth] WARN: vmax==vmin; forcing denom epsilon to avoid NaNs.")
-            denom = 1e-6
-        depth_array = depth_array / denom
-
-        # EXACT invert like your Zoe code
-        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))
-    log(f"[SaliaDepth] Output (post-remove_pad): {detected_map.shape} dtype={detected_map.dtype}")
-    return detected_map
-
-
-def resize_to_original(depth_rgb_u8: np.ndarray, w0: int, h0: int, log) -> np.ndarray:
-    """
-    Resize depth output back to original input size.
-    Use cv2 if available, else PIL.
-    """
-    try:
-        import cv2
-        out = cv2.resize(depth_rgb_u8, (w0, h0), interpolation=cv2.INTER_LINEAR)
-        return out.astype(np.uint8)
-    except Exception as e:
-        log(f"[SaliaDepth] WARN: cv2 resize failed ({repr(e)}); using PIL.")
-        pil = Image.fromarray(depth_rgb_u8)
-        pil = pil.resize((w0, h0), resample=Image.BILINEAR)
-        return np.array(pil, dtype=np.uint8)
-
-
-# --------------------------------------------------------------------------------------
-# ComfyUI Node
-# --------------------------------------------------------------------------------------
-
-class Salia_Depth_Preprocessor:
-    @classmethod
-    def INPUT_TYPES(cls):
-        return {
-            "required": {
-                "image": ("IMAGE",),
-                # note: default -1, min -1
-                "resolution": ("INT", {"default": -1, "min": -1, "max": 8192, "step": 1}),
-            }
-        }
-
-    # 2 outputs: image + log string
-    RETURN_TYPES = ("IMAGE", "STRING")
-    FUNCTION = "execute"
-    CATEGORY = "ControlNet Preprocessors/Normal and Depth Estimators"
-
-    def execute(self, image, resolution=-1):
-        lines, log = _make_logger()
-        log("[SaliaDepth] ==================================================")
-        log("[SaliaDepth] SaliaDepthPreprocessor starting")
-        log(f"[SaliaDepth] resolution input = {resolution}")
-
-        # Get torch device
-        try:
-            device = model_management.get_torch_device()
-        except Exception as e:
-            device = torch.device("cpu")
-            log(f"[SaliaDepth] WARN: model_management.get_torch_device failed: {repr(e)} -> using CPU")
-
-        log(f"[SaliaDepth] torch device = {device}")
-
-        # Load pipeline
-        pipe = None
-        try:
-            pipe = get_depth_pipeline(device, log)
-        except Exception as e:
-            log(f"[SaliaDepth] ERROR: get_depth_pipeline crashed: {repr(e)}")
-            pipe = None
-
-        if pipe is None:
-            log("[SaliaDepth] FATAL: No pipeline available. Returning input image unchanged.")
-            return (image, "\n".join(lines))
-
-        # Batch support
-        if image.ndim == 3:
-            image = image.unsqueeze(0)
-
-        outs = []
-        for i in range(image.shape[0]):
-            try:
-                # Original dimensions
-                h0 = int(image[i].shape[0])
-                w0 = int(image[i].shape[1])
-                c0 = int(image[i].shape[2])
-                log(f"[SaliaDepth] ---- Batch index {i} input shape = ({h0},{w0},{c0}) ----")
-
-                inp_u8 = comfy_tensor_to_u8(image[i])
-
-                # RGBA rule (pre)
-                rgb_for_depth, alpha_u8 = composite_rgba_over_white_keep_alpha(inp_u8)
-                had_rgba = alpha_u8 is not None
-                log(f"[SaliaDepth] had_rgba={had_rgba}")
-
-                # Run depth (Zoe-style)
-                depth_rgb = depth_estimate_zoe_style(
-                    pipe=pipe,
-                    input_rgb_u8=rgb_for_depth,
-                    detect_resolution=int(resolution),
-                    log=log,
-                    upscale_method="INTER_CUBIC"
-                )
-
-                # Resize back to original input size
-                depth_rgb = resize_to_original(depth_rgb, w0=w0, h0=h0, log=log)
-
-                # RGBA rule (post)
-                if had_rgba:
-                    # Use original alpha at original size.
-                    # If alpha size differs, resize alpha to match.
-                    if alpha_u8.shape[0] != h0 or alpha_u8.shape[1] != w0:
-                        log("[SaliaDepth] Alpha size mismatch; resizing alpha to original size.")
-                        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)
-
-                    # "Put alpha on RGB turning it into RGBA, then put BLACK background behind it, then back to RGB"
-                    depth_rgb = apply_alpha_then_black_background(depth_rgb, alpha_u8)
-                    log("[SaliaDepth] Applied RGBA post-step (alpha + black background).")
-
-                outs.append(u8_to_comfy_tensor(depth_rgb))
-
-            except Exception as e:
-                log(f"[SaliaDepth] ERROR: Inference failed at batch index {i}: {repr(e)}")
-                log("[SaliaDepth] Passing through original input image for this batch item.")
-                outs.append(image[i].unsqueeze(0))
-
-        out = torch.cat(outs, dim=0)
-        log("[SaliaDepth] Done.")
-        return (out, "\n".join(lines))
-
-
-NODE_CLASS_MAPPINGS = {
-    "SaliaDepthPreprocessor": Salia_Depth_Preprocessor
-}
-
-NODE_DISPLAY_NAME_MAPPINGS = {
-    "SaliaDepthPreprocessor": "Salia Depth (local assets/depth + logs)"
+import os
+import shutil
+import urllib.request
+from pathlib import Path
+from typing import Dict, Tuple, Any, Optional, List
+
+import numpy as np
+import torch
+from PIL import Image
+
+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
+
+
+# --------------------------------------------------------------------------------------
+# Paths / sources
+# --------------------------------------------------------------------------------------
+
+# This file: comfyui-salia_online/nodes/Salia_Depth.py
+# Plugin root: comfyui-salia_online/
+PLUGIN_ROOT = Path(__file__).resolve().parent.parent
+
+# Requested local path: assets/depth
+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-path" fallback
+ZOE_FALLBACK_REPO_ID = "Intel/zoedepth-nyu-kitti"
+
+
+# --------------------------------------------------------------------------------------
+# Logging helpers
+# --------------------------------------------------------------------------------------
+
+def _make_logger() -> Tuple[List[str], Any]:
+    lines: List[str] = []
+
+    def log(msg: str):
+        # console
+        try:
+            print(msg)
+        except Exception:
+            pass
+        # UI string
+        lines.append(str(msg))
+
+    return lines, log
+
+
+def _fmt_bytes(n: Optional[int]) -> str:
+    if n is None:
+        return "?"
+    # simple readable
+    for unit in ["B", "KB", "MB", "GB", "TB"]:
+        if n < 1024:
+            return f"{n:.0f}{unit}"
+        n /= 1024.0
+    return f"{n:.1f}PB"
+
+
+def _file_size(path: Path) -> Optional[int]:
+    try:
+        return path.stat().st_size
+    except Exception:
+        return None
+
+
+def _hf_cache_info() -> Dict[str, str]:
+    info: Dict[str, str] = {}
+    info["env.HF_HOME"] = os.environ.get("HF_HOME", "")
+    info["env.HF_HUB_CACHE"] = os.environ.get("HF_HUB_CACHE", "")
+    info["env.TRANSFORMERS_CACHE"] = os.environ.get("TRANSFORMERS_CACHE", "")
+    info["env.HUGGINGFACE_HUB_CACHE"] = os.environ.get("HUGGINGFACE_HUB_CACHE", "")
+
+    try:
+        from huggingface_hub import constants as hf_constants
+        # These exist in most hub versions:
+        info["huggingface_hub.constants.HF_HOME"] = str(getattr(hf_constants, "HF_HOME", ""))
+        info["huggingface_hub.constants.HF_HUB_CACHE"] = str(getattr(hf_constants, "HF_HUB_CACHE", ""))
+    except Exception:
+        pass
+
+    return info
+
+
+# --------------------------------------------------------------------------------------
+# Download helpers
+# --------------------------------------------------------------------------------------
+
+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:
+    """
+    Download with atomic temp rename.
+    """
+    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(log) -> bool:
+    """
+    Ensure assets/depth contains the 3 files.
+    Returns True if present or downloaded successfully, else False.
+    """
+    # Always log expected locations + URLs, even if we don't download.
+    log("[SaliaDepth] ===== Local model file check =====")
+    log(f"[SaliaDepth] Plugin root: {PLUGIN_ROOT}")
+    log(f"[SaliaDepth] Local model dir (on drive): {MODEL_DIR}")
+
+    for fname, url in REQUIRED_FILES.items():
+        fpath = MODEL_DIR / fname
+        exists = fpath.exists()
+        size = _file_size(fpath) if exists else None
+        log(f"[SaliaDepth]   - {fname}")
+        log(f"[SaliaDepth]       local path: {fpath}  exists={exists}  size={_fmt_bytes(size)}")
+        log(f"[SaliaDepth]       remote url : {url}")
+
+    if _have_required_files():
+        log("[SaliaDepth] All required local files already exist. No download needed.")
+        return True
+
+    log("[SaliaDepth] One or more local files missing. Attempting download...")
+
+    try:
+        for fname, url in REQUIRED_FILES.items():
+            fpath = MODEL_DIR / fname
+            if fpath.exists():
+                continue
+            log(f"[SaliaDepth] Downloading '{fname}' -> '{fpath}'")
+            _download_url_to_file(url, fpath)
+            log(f"[SaliaDepth] Downloaded '{fname}' size={_fmt_bytes(_file_size(fpath))}")
+
+        ok = _have_required_files()
+        log(f"[SaliaDepth] Download finished. ok={ok}")
+        return ok
+    except Exception as e:
+        log(f"[SaliaDepth] Download failed with error: {repr(e)}")
+        return False
+
+
+# --------------------------------------------------------------------------------------
+# Exact Zoe-style preprocessing helpers (copied/adapted from your snippet)
+# --------------------------------------------------------------------------------------
+
+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",
+    log=None
+) -> Tuple[np.ndarray, Any]:
+    """
+    EXACT behavior like your zoe.transformers.py:
+      k = resolution / min(H,W)
+      resize to (W_target, H_target)
+      pad to multiple of 64
+      return padded image and remove_pad() closure
+    """
+    # prefer cv2 like original for matching results
+    cv2 = None
+    try:
+        import cv2 as _cv2
+        cv2 = _cv2
+    except Exception:
+        cv2 = None
+        if log:
+            log("[SaliaDepth] WARN: cv2 not available; resizing will use PIL fallback (may change results).")
+
+    if skip_hwc3:
+        img = input_image
+    else:
+        img = HWC3(input_image)
+
+    H_raw, W_raw, _ = img.shape
+    if resolution <= 0:
+        # keep original, but still pad to 64 (we will handle padding separately for -1 path)
+        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 fallback
+        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]:
+    """
+    For resolution == -1: keep original resolution but still pad to multiples of 64,
+    then provide remove_pad that returns original size.
+    """
+    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
+
+
+# --------------------------------------------------------------------------------------
+# RGBA rules (as you requested)
+# --------------------------------------------------------------------------------------
+
+def composite_rgba_over_white_keep_alpha(inp_u8: np.ndarray) -> Tuple[np.ndarray, Optional[np.ndarray]]:
+    """
+    If RGBA: return RGB composited over WHITE + alpha_u8 kept separately.
+    If RGB: return input RGB + None alpha.
+    """
+    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
+    # force to RGB
+    return HWC3(inp_u8), None
+
+
+def apply_alpha_then_black_background(depth_rgb_u8: np.ndarray, alpha_u8: np.ndarray) -> np.ndarray:
+    """
+    Requested output rule:
+      - attach alpha to depth (conceptually RGBA)
+      - composite over BLACK
+      - output RGB
+    That is equivalent to depth_rgb * alpha.
+    """
+    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
+
+
+# --------------------------------------------------------------------------------------
+# ComfyUI conversion helpers
+# --------------------------------------------------------------------------------------
+
+def comfy_tensor_to_u8(img: torch.Tensor) -> np.ndarray:
+    """
+    Comfy IMAGE: float [0..1], shape [H,W,C] or [B,H,W,C]
+    Convert to uint8 HWC.
+    """
+    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]
+
+
+# --------------------------------------------------------------------------------------
+# Pipeline loading (local-first, then zoe fallback)
+# --------------------------------------------------------------------------------------
+
+_PIPE_CACHE: Dict[Tuple[str, str], Any] = {}  # (model_source, device_str) -> pipeline
+
+
+def _try_load_pipeline(model_source: str, device: torch.device, log):
+    """
+    Use transformers.pipeline like Zoe code does.
+    We intentionally do NOT pass device=... here, and instead move model like Zoe node.
+    """
+    if pipeline is None:
+        raise RuntimeError(f"transformers import failed: {_TRANSFORMERS_IMPORT_ERROR}")
+
+    key = (model_source, str(device))
+    if key in _PIPE_CACHE:
+        log(f"[SaliaDepth] Using cached pipeline for source='{model_source}' device='{device}'")
+        return _PIPE_CACHE[key]
+
+    log(f"[SaliaDepth] Creating pipeline(task='depth-estimation', model='{model_source}')")
+    p = pipeline(task="depth-estimation", model=model_source)
+
+    # Try to move model to torch device, like ZoeDetector.to()
+    try:
+        p.model = p.model.to(device)
+        p.device = device  # Zoe code sets this; newer transformers uses torch.device internally
+        log(f"[SaliaDepth] Moved pipeline model to device: {device}")
+    except Exception as e:
+        log(f"[SaliaDepth] WARN: Could not move pipeline model to device {device}: {repr(e)}")
+
+    # Log config info for debugging
+    try:
+        cfg = p.model.config
+        log(f"[SaliaDepth] Model class: {p.model.__class__.__name__}")
+        log(f"[SaliaDepth] Config class: {cfg.__class__.__name__}")
+        log(f"[SaliaDepth] Config model_type: {getattr(cfg, 'model_type', '')}")
+        log(f"[SaliaDepth] Config _name_or_path: {getattr(cfg, '_name_or_path', '')}")
+    except Exception as e:
+        log(f"[SaliaDepth] WARN: Could not log model config: {repr(e)}")
+
+    _PIPE_CACHE[key] = p
+    return p
+
+
+def get_depth_pipeline(device: torch.device, log):
+    """
+    1) Ensure assets/depth files exist (download if missing)
+    2) Try load local dir
+    3) Fallback to Intel/zoedepth-nyu-kitti
+    4) If both fail -> None
+    """
+    # Always log HF cache info (helps locate where fallback downloads go)
+    log("[SaliaDepth] ===== Hugging Face cache info (fallback path) =====")
+    for k, v in _hf_cache_info().items():
+        if v:
+            log(f"[SaliaDepth] {k} = {v}")
+    log(f"[SaliaDepth] Zoe fallback repo id: {ZOE_FALLBACK_REPO_ID}")
+
+    # Local-first
+    local_ok = ensure_local_model_files(log)
+    if local_ok:
+        try:
+            log(f"[SaliaDepth] Trying LOCAL model from directory: {MODEL_DIR}")
+            return _try_load_pipeline(str(MODEL_DIR), device, log)
+        except Exception as e:
+            log(f"[SaliaDepth] Local model load FAILED: {repr(e)}")
+
+    # Fallback
+    try:
+        log(f"[SaliaDepth] Trying ZOE fallback model: {ZOE_FALLBACK_REPO_ID}")
+        return _try_load_pipeline(ZOE_FALLBACK_REPO_ID, device, log)
+    except Exception as e:
+        log(f"[SaliaDepth] Zoe fallback load FAILED: {repr(e)}")
+
+    return None
+
+
+# --------------------------------------------------------------------------------------
+# Depth inference (Zoe-style)
+# --------------------------------------------------------------------------------------
+
+def depth_estimate_zoe_style(
+    pipe,
+    input_rgb_u8: np.ndarray,
+    detect_resolution: int,
+    log,
+    upscale_method: str = "INTER_CUBIC"
+) -> np.ndarray:
+    """
+    Matches your ZoeDetector.__call__ logic very closely.
+    Returns uint8 RGB depth map.
+    """
+    # detect_resolution:
+    #   - if -1: keep original but pad-to-64
+    #   - else: min-side resize to detect_resolution, then pad-to-64
+    if detect_resolution == -1:
+        work_img, remove_pad = pad_only_to_64(input_rgb_u8, mode="edge")
+        log(f"[SaliaDepth] Preprocess: resolution=-1 (no resize), padded to 64. work={work_img.shape}")
+    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",
+            log=log
+        )
+        log(f"[SaliaDepth] Preprocess: min-side resized to {detect_resolution}, padded to 64. work={work_img.shape}")
+
+    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)
+
+        # EXACT normalization like your Zoe code
+        vmin = float(np.percentile(depth_array, 2))
+        vmax = float(np.percentile(depth_array, 85))
+
+        log(f"[SaliaDepth] Depth raw stats: shape={depth_array.shape} vmin(p2)={vmin:.6f} vmax(p85)={vmax:.6f} mean={float(depth_array.mean()):.6f}")
+
+        depth_array = depth_array - vmin
+        denom = (vmax - vmin)
+        if abs(denom) < 1e-12:
+            # avoid division by zero; log it
+            log("[SaliaDepth] WARN: vmax==vmin; forcing denom epsilon to avoid NaNs.")
+            denom = 1e-6
+        depth_array = depth_array / denom
+
+        # EXACT invert like your Zoe code
+        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))
+    log(f"[SaliaDepth] Output (post-remove_pad): {detected_map.shape} dtype={detected_map.dtype}")
+    return detected_map
+
+
+def resize_to_original(depth_rgb_u8: np.ndarray, w0: int, h0: int, log) -> np.ndarray:
+    """
+    Resize depth output back to original input size.
+    Use cv2 if available, else PIL.
+    """
+    try:
+        import cv2
+        out = cv2.resize(depth_rgb_u8, (w0, h0), interpolation=cv2.INTER_LINEAR)
+        return out.astype(np.uint8)
+    except Exception as e:
+        log(f"[SaliaDepth] WARN: cv2 resize failed ({repr(e)}); using PIL.")
+        pil = Image.fromarray(depth_rgb_u8)
+        pil = pil.resize((w0, h0), resample=Image.BILINEAR)
+        return np.array(pil, dtype=np.uint8)
+
+
+# --------------------------------------------------------------------------------------
+# ComfyUI Node
+# --------------------------------------------------------------------------------------
+
+class Salia_Depth_Preprocessor:
+    @classmethod
+    def INPUT_TYPES(cls):
+        return {
+            "required": {
+                "image": ("IMAGE",),
+                # note: default -1, min -1
+                "resolution": ("INT", {"default": -1, "min": -1, "max": 8192, "step": 1}),
+            }
+        }
+
+    # 2 outputs: image + log string
+    RETURN_TYPES = ("IMAGE", "STRING")
+    FUNCTION = "execute"
+    CATEGORY = "ControlNet Preprocessors/Normal and Depth Estimators"
+
+    def execute(self, image, resolution=-1):
+        lines, log = _make_logger()
+        log("[SaliaDepth] ==================================================")
+        log("[SaliaDepth] SaliaDepthPreprocessor starting")
+        log(f"[SaliaDepth] resolution input = {resolution}")
+
+        # Get torch device
+        try:
+            device = model_management.get_torch_device()
+        except Exception as e:
+            device = torch.device("cpu")
+            log(f"[SaliaDepth] WARN: model_management.get_torch_device failed: {repr(e)} -> using CPU")
+
+        log(f"[SaliaDepth] torch device = {device}")
+
+        # Load pipeline
+        pipe = None
+        try:
+            pipe = get_depth_pipeline(device, log)
+        except Exception as e:
+            log(f"[SaliaDepth] ERROR: get_depth_pipeline crashed: {repr(e)}")
+            pipe = None
+
+        if pipe is None:
+            log("[SaliaDepth] FATAL: No pipeline available. Returning input image unchanged.")
+            return (image, "\n".join(lines))
+
+        # Batch support
+        if image.ndim == 3:
+            image = image.unsqueeze(0)
+
+        outs = []
+        for i in range(image.shape[0]):
+            try:
+                # Original dimensions
+                h0 = int(image[i].shape[0])
+                w0 = int(image[i].shape[1])
+                c0 = int(image[i].shape[2])
+                log(f"[SaliaDepth] ---- Batch index {i} input shape = ({h0},{w0},{c0}) ----")
+
+                inp_u8 = comfy_tensor_to_u8(image[i])
+
+                # RGBA rule (pre)
+                rgb_for_depth, alpha_u8 = composite_rgba_over_white_keep_alpha(inp_u8)
+                had_rgba = alpha_u8 is not None
+                log(f"[SaliaDepth] had_rgba={had_rgba}")
+
+                # Run depth (Zoe-style)
+                depth_rgb = depth_estimate_zoe_style(
+                    pipe=pipe,
+                    input_rgb_u8=rgb_for_depth,
+                    detect_resolution=int(resolution),
+                    log=log,
+                    upscale_method="INTER_CUBIC"
+                )
+
+                # Resize back to original input size
+                depth_rgb = resize_to_original(depth_rgb, w0=w0, h0=h0, log=log)
+
+                # RGBA rule (post)
+                if had_rgba:
+                    # Use original alpha at original size.
+                    # If alpha size differs, resize alpha to match.
+                    if alpha_u8.shape[0] != h0 or alpha_u8.shape[1] != w0:
+                        log("[SaliaDepth] Alpha size mismatch; resizing alpha to original size.")
+                        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)
+
+                    # "Put alpha on RGB turning it into RGBA, then put BLACK background behind it, then back to RGB"
+                    depth_rgb = apply_alpha_then_black_background(depth_rgb, alpha_u8)
+                    log("[SaliaDepth] Applied RGBA post-step (alpha + black background).")
+
+                outs.append(u8_to_comfy_tensor(depth_rgb))
+
+            except Exception as e:
+                log(f"[SaliaDepth] ERROR: Inference failed at batch index {i}: {repr(e)}")
+                log("[SaliaDepth] Passing through original input image for this batch item.")
+                outs.append(image[i].unsqueeze(0))
+
+        out = torch.cat(outs, dim=0)
+        log("[SaliaDepth] Done.")
+        return (out, "\n".join(lines))
+
+
+NODE_CLASS_MAPPINGS = {
+    "SaliaDepthPreprocessor": Salia_Depth_Preprocessor
+}
+
+NODE_DISPLAY_NAME_MAPPINGS = {
+    "SaliaDepthPreprocessor": "Salia Depth (local assets/depth + logs)"
 }