Upload Salia_UltralyticsDetectorProvider2.py

Salia_UltralyticsDetectorProvider2.py

@@ -0,0 +1,741 @@
+"""
+Salia Ultralytics Detector Provider (ComfyUI custom node)
+
+Goal:
+- Provide the same outputs as Impact-Subpack's `UltralyticsDetectorProvider`:
+    - BBOX_DETECTOR
+    - SEGM_DETECTOR
+- But packaged so you can drop it into your own custom node folder (your Salia_* environment)
+  without requiring ComfyUI-Impact-Subpack.
+
+Notes:
+- This file intentionally keeps dependencies minimal and self-contained.
+- It uses `ultralytics.YOLO` to run `.pt` models directly (no TensorRT build step).
+- For PyTorch >= 2.6, `torch.load` defaults to `weights_only=True` which can break
+  legacy `.pt` checkpoints. This file adds an OPTIONAL whitelist-based fallback
+  to `weights_only=False` (unsafe) for specifically trusted model filenames.
+"""
+
+from __future__ import annotations
+
+import os
+import logging
+import pickle
+from datetime import datetime
+from contextlib import contextmanager
+from collections import namedtuple
+
+import folder_paths
+
+from PIL import Image
+import numpy as np
+import torch
+import torch.nn.functional as F
+
+try:
+    import cv2  # opencv-python or opencv-python-headless
+except Exception:
+    cv2 = None
+
+
+# ---------------------------
+# Model folders (same layout as Impact Subpack)
+# ---------------------------
+
+_SUPPORTED_PT_EXTS = getattr(folder_paths, "supported_pt_extensions", [".pt", ".pth", ".ckpt", ".safetensors"])
+
+
+def _add_folder_path_and_extensions(folder_name: str, paths: list[str], extensions: list[str] | tuple[str, ...]):
+    """Add/merge a folder_paths entry without depending on Impact-Pack helpers."""
+    if folder_name in folder_paths.folder_names_and_paths:
+        existing_paths, existing_exts = folder_paths.folder_names_and_paths[folder_name]
+        merged_paths = list(existing_paths)
+        for p in paths:
+            if p not in merged_paths:
+                merged_paths.append(p)
+        merged_exts = list(existing_exts)
+        for ext in extensions:
+            if ext not in merged_exts:
+                merged_exts.append(ext)
+        folder_paths.folder_names_and_paths[folder_name] = (merged_paths, tuple(merged_exts))
+    else:
+        folder_paths.folder_names_and_paths[folder_name] = (list(paths), tuple(extensions))
+
+
+def _update_model_paths(base_path: str):
+    """Register standard Impact-Subpack ultralytics model locations."""
+    _add_folder_path_and_extensions(
+        "ultralytics_bbox",
+        [os.path.join(base_path, "ultralytics", "bbox")],
+        _SUPPORTED_PT_EXTS,
+    )
+    _add_folder_path_and_extensions(
+        "ultralytics_segm",
+        [os.path.join(base_path, "ultralytics", "segm")],
+        _SUPPORTED_PT_EXTS,
+    )
+    _add_folder_path_and_extensions(
+        "ultralytics",
+        [os.path.join(base_path, "ultralytics")],
+        _SUPPORTED_PT_EXTS,
+    )
+
+
+# Register common folders (models_dir + ComfyUI-Manager download_model_base)
+_update_model_paths(folder_paths.models_dir)
+if "download_model_base" in folder_paths.folder_names_and_paths:
+    try:
+        _update_model_paths(folder_paths.get_folder_paths("download_model_base")[0])
+    except Exception:
+        pass
+
+# Also register local folder(s) inside THIS custom-node extension, so you can keep
+# models next to your Salia_*.py files if you want.
+_THIS_DIR = os.path.dirname(os.path.abspath(__file__))
+for local_dir in [
+    os.path.join(_THIS_DIR, "nodes"),
+    os.path.join(_THIS_DIR, "models"),
+    _THIS_DIR,
+]:
+    if os.path.isdir(local_dir):
+        _add_folder_path_and_extensions("ultralytics_bbox", [local_dir], _SUPPORTED_PT_EXTS)
+        _add_folder_path_and_extensions("ultralytics_segm", [local_dir], _SUPPORTED_PT_EXTS)
+        _add_folder_path_and_extensions("ultralytics", [local_dir], _SUPPORTED_PT_EXTS)
+
+
+# ---------------------------
+# Optional safe-load fallback (PyTorch >= 2.6)
+# ---------------------------
+
+_ORIG_TORCH_LOAD = torch.load
+
+
+def _get_whitelist_file() -> str | None:
+    """Create/return the whitelist file path under ComfyUI's user directory."""
+    try:
+        user_dir = folder_paths.get_user_directory()
+    except Exception:
+        user_dir = None
+
+    if not user_dir or not os.path.isdir(user_dir):
+        return None
+
+    wl_dir = os.path.join(user_dir, "default", "ComfyUI-Salia-Ultralytics")
+    wl_file = os.path.join(wl_dir, "model-whitelist.txt")
+    try:
+        os.makedirs(wl_dir, exist_ok=True)
+        if not os.path.exists(wl_file):
+            with open(wl_file, "w", encoding="utf-8") as f:
+                f.write("# Add base filenames of trusted legacy models here (one per line).\n")
+                f.write("# Example: eyes.pt\n")
+                f.write("# These will be allowed to load with weights_only=False if safe loading fails.\n")
+                f.write("# WARNING: Only add models you trust.\n")
+    except Exception:
+        return None
+
+    return wl_file
+
+
+_WHITELIST_PATH = _get_whitelist_file()
+
+
+# ---------------------------
+# Model path logging (requested)
+# ---------------------------
+
+def _get_model_load_log_file() -> str:
+    """
+    Log file path used to record which ultralytics model file was actually loaded.
+    Prefer the same ComfyUI user dir used for the whitelist (if available).
+    """
+    # If whitelist exists, put log next to it (same directory).
+    if _WHITELIST_PATH:
+        base_dir = os.path.dirname(_WHITELIST_PATH)
+        return os.path.join(base_dir, "model-load-log.txt")
+
+    # Fallback: try ComfyUI user directory
+    try:
+        user_dir = folder_paths.get_user_directory()
+    except Exception:
+        user_dir = None
+
+    if user_dir and os.path.isdir(user_dir):
+        base_dir = os.path.join(user_dir, "default", "ComfyUI-Salia-Ultralytics")
+        try:
+            os.makedirs(base_dir, exist_ok=True)
+        except Exception:
+            pass
+        return os.path.join(base_dir, "model-load-log.txt")
+
+    # Last resort: next to this python file
+    return os.path.join(_THIS_DIR, "model-load-log.txt")
+
+
+_MODEL_LOAD_LOG_PATH = _get_model_load_log_file()
+
+
+def _find_all_model_paths(model_name: str) -> list[str]:
+    """
+    Find all possible on-disk matches across the registered ultralytics folders.
+    Useful if the same filename exists in multiple locations.
+    """
+    matches: list[str] = []
+
+    try:
+        ultra_roots = folder_paths.get_folder_paths("ultralytics")
+    except Exception:
+        ultra_roots = []
+
+    try:
+        bbox_roots = folder_paths.get_folder_paths("ultralytics_bbox")
+    except Exception:
+        bbox_roots = []
+
+    try:
+        segm_roots = folder_paths.get_folder_paths("ultralytics_segm")
+    except Exception:
+        segm_roots = []
+
+    def add_if_exists(root: str, rel: str):
+        p = os.path.join(root, rel)
+        if os.path.exists(p):
+            matches.append(os.path.abspath(p))
+
+    # model_name might be "bbox/foo.pt" or "segm/foo.pt" (includes subfolder)
+    for r in ultra_roots:
+        add_if_exists(r, model_name)
+
+    # Also search the specialized bbox/segm roots with the prefix stripped
+    if model_name.startswith("bbox/"):
+        rel = model_name[5:]
+        for r in bbox_roots:
+            add_if_exists(r, rel)
+    elif model_name.startswith("segm/"):
+        rel = model_name[5:]
+        for r in segm_roots:
+            add_if_exists(r, rel)
+
+    # De-dupe preserving order
+    out: list[str] = []
+    seen = set()
+    for p in matches:
+        if p not in seen:
+            seen.add(p)
+            out.append(p)
+    return out
+
+
+def _log_selected_model(model_name: str, model_path: str, matches: list[str] | None = None):
+    """
+    Prints the resolved model path to console AND appends it to a log file.
+    """
+    # 1) Console output
+    print(f"[Salia Ultralytics] Selected model_name: {model_name}")
+    print(f"[Salia Ultralytics] Resolved model_path: {model_path}")
+    if matches and len(matches) > 1:
+        print("[Salia Ultralytics] Multiple matches found (first one is used by get_full_path):")
+        for p in matches:
+            print(f"  - {p}")
+    print(f"[Salia Ultralytics] Model load log file: {_MODEL_LOAD_LOG_PATH}")
+
+    # Also emit to python logging (ComfyUI typically captures this)
+    logging.info("[Salia Ultralytics] Selected model_name: %s", model_name)
+    logging.info("[Salia Ultralytics] Resolved model_path: %s", model_path)
+    if matches and len(matches) > 1:
+        logging.warning("[Salia Ultralytics] Multiple matches found (first one is used by get_full_path):")
+        for p in matches:
+            logging.warning("  - %s", p)
+    logging.info("[Salia Ultralytics] Model load log file: %s", _MODEL_LOAD_LOG_PATH)
+
+    # 2) File append
+    try:
+        ts = datetime.now().isoformat(timespec="seconds")
+        exists = os.path.isfile(model_path)
+        size = os.path.getsize(model_path) if exists else -1
+
+        log_dir = os.path.dirname(_MODEL_LOAD_LOG_PATH)
+        if log_dir:
+            os.makedirs(log_dir, exist_ok=True)
+
+        with open(_MODEL_LOAD_LOG_PATH, "a", encoding="utf-8") as f:
+            f.write(f"{ts}\t{model_name}\t{model_path}\texists={exists}\tsize={size}\n")
+            if matches and len(matches) > 1:
+                for p in matches:
+                    f.write(f"{ts}\tmatch\t{p}\n")
+    except Exception as e:
+        logging.warning("[Salia Ultralytics] Failed to write model-load log to %s: %s", _MODEL_LOAD_LOG_PATH, e)
+
+
+def _load_whitelist(filepath: str | None) -> set[str]:
+    if not filepath:
+        return set()
+    try:
+        approved: set[str] = set()
+        with open(filepath, "r", encoding="utf-8") as f:
+            for line in f:
+                line = line.strip()
+                if line and not line.startswith("#"):
+                    approved.add(os.path.basename(line))
+        return approved
+    except Exception:
+        return set()
+
+
+_MODEL_WHITELIST = _load_whitelist(_WHITELIST_PATH)
+
+
+def _torch_load_wrapper(*args, **kwargs):
+    """Try safe load first; if it fails due to weights-only restrictions, allow fallback if whitelisted."""
+    filename = None
+    if args and isinstance(args[0], str):
+        filename = os.path.basename(args[0])
+    elif isinstance(kwargs.get("f"), str):
+        filename = os.path.basename(kwargs["f"])
+
+    try:
+        return _ORIG_TORCH_LOAD(*args, **kwargs)
+    except pickle.UnpicklingError as e:
+        msg = str(e)
+        # Heuristic: this is the common PyTorch >=2.6 safe-load failure mode.
+        maybe_weights_only_error = (
+            "Weights only load failed" in msg
+            or "Unsupported global" in msg
+            or "disallowed" in msg
+            or "not allowed" in msg
+            or "getattr" in msg
+        )
+
+        if not maybe_weights_only_error:
+            raise
+
+        # Refresh whitelist from disk (so users can edit without restarting, sometimes)
+        global _MODEL_WHITELIST
+        _MODEL_WHITELIST = _load_whitelist(_WHITELIST_PATH)
+
+        if filename and filename in _MODEL_WHITELIST:
+            logging.warning(
+                "[Salia Ultralytics] Safe torch.load failed for '%s'. Retrying with weights_only=False because it's whitelisted (%s).",
+                filename,
+                _WHITELIST_PATH,
+            )
+            retry_kwargs = dict(kwargs)
+            retry_kwargs["weights_only"] = False
+            return _ORIG_TORCH_LOAD(*args, **retry_kwargs)
+
+        logging.error(
+            "[Salia Ultralytics] Blocked unsafe model load for '%s'.\n"
+            "Safe loading failed and the file is not whitelisted.\n"
+            "If you TRUST this model, add its base name to: %s",
+            filename or "[unknown]",
+            _WHITELIST_PATH or "[whitelist path unavailable]",
+        )
+        raise
+
+
+@contextmanager
+def _patched_torch_load_for_ultralytics():
+    """Patch torch.load only while ultralytics loads a checkpoint."""
+    # If PyTorch doesn't even have the safe-loader feature, don't patch.
+    if not hasattr(torch.serialization, "safe_globals"):
+        yield
+        return
+
+    prev = torch.load
+    torch.load = _torch_load_wrapper
+    try:
+        yield
+    finally:
+        torch.load = prev
+
+
+def _load_yolo(model_path: str):
+    """Load an Ultralytics YOLO model (with optional safe-load fallback)."""
+    try:
+        from ultralytics import YOLO  # lazy import
+    except Exception as e:
+        raise ImportError(
+            "[Salia Ultralytics] ultralytics is not installed. Install it in your ComfyUI env, e.g.:\n"
+            "pip install ultralytics"
+        ) from e
+
+    with _patched_torch_load_for_ultralytics():
+        return YOLO(model_path)
+
+
+# ---------------------------
+# Minimal Impact-compatible utilities (self-contained)
+# ---------------------------
+
+def _tensor2np_rgb(image: torch.Tensor) -> np.ndarray:
+    """Convert a ComfyUI IMAGE tensor to a uint8 RGB numpy image."""
+    # ComfyUI image is usually: (B,H,W,C) float in [0,1]
+    if not isinstance(image, torch.Tensor):
+        raise TypeError(f"Expected torch.Tensor, got {type(image)}")
+
+    if image.dim() == 4:
+        img = image[0]
+    else:
+        img = image
+
+    img = img.detach()
+    if img.is_cuda:
+        img = img.cpu()
+
+    img = img.clamp(0, 1).numpy()
+    if img.shape[-1] == 1:
+        img = np.repeat(img, 3, axis=-1)
+
+    img_u8 = (img * 255.0).round().astype(np.uint8)
+    return img_u8
+
+
+def tensor2pil(image: torch.Tensor) -> Image.Image:
+    return Image.fromarray(_tensor2np_rgb(image))
+
+
+def make_crop_region(w: int, h: int, bbox_xyxy, crop_factor: float, crop_min_size: int | None = None):
+    x1, y1, x2, y2 = [float(v) for v in bbox_xyxy]
+    bbox_w = max(1.0, x2 - x1)
+    bbox_h = max(1.0, y2 - y1)
+
+    crop_w = bbox_w * float(crop_factor)
+    crop_h = bbox_h * float(crop_factor)
+
+    if crop_min_size is not None:
+        crop_w = max(crop_w, float(crop_min_size))
+        crop_h = max(crop_h, float(crop_min_size))
+
+    cx = (x1 + x2) / 2.0
+    cy = (y1 + y2) / 2.0
+
+    rx1 = int(round(cx - crop_w / 2.0))
+    ry1 = int(round(cy - crop_h / 2.0))
+    rx2 = int(round(cx + crop_w / 2.0))
+    ry2 = int(round(cy + crop_h / 2.0))
+
+    rx1 = max(0, min(w - 1, rx1))
+    ry1 = max(0, min(h - 1, ry1))
+    rx2 = max(rx1 + 1, min(w, rx2))
+    ry2 = max(ry1 + 1, min(h, ry2))
+
+    return (rx1, ry1, rx2, ry2)
+
+
+def crop_image(image: torch.Tensor, crop_region):
+    x1, y1, x2, y2 = crop_region
+    x1, y1, x2, y2 = int(x1), int(y1), int(x2), int(y2)
+    if image.dim() == 4:
+        return image[:, y1:y2, x1:x2, :]
+    if image.dim() == 3:
+        return image[y1:y2, x1:x2, :]
+    raise ValueError(f"Unexpected image tensor shape: {tuple(image.shape)}")
+
+
+def crop_ndarray2(arr: np.ndarray, crop_region):
+    x1, y1, x2, y2 = crop_region
+    x1, y1, x2, y2 = int(x1), int(y1), int(x2), int(y2)
+    return arr[y1:y2, x1:x2]
+
+
+def dilate_masks(segmasks, dilation: int):
+    if dilation <= 0:
+        return segmasks
+    if cv2 is None:
+        raise ImportError(
+            "[Salia Ultralytics] opencv-python is required for mask dilation but cv2 could not be imported.\n"
+            "Install: pip install opencv-python-headless"
+        )
+
+    k = int(dilation)
+    ksize = k * 2 + 1
+    kernel = cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (ksize, ksize))
+
+    out = []
+    for bbox, mask, conf in segmasks:
+        m = (mask > 0.5).astype(np.uint8) * 255
+        m = cv2.dilate(m, kernel, iterations=1)
+        out.append((bbox, (m > 0).astype(np.float32), conf))
+    return out
+
+
+def combine_masks(segmasks, out_shape_hw: tuple[int, int] | None = None) -> torch.Tensor:
+    if not segmasks:
+        if out_shape_hw is None:
+            return torch.zeros((1, 1, 1), dtype=torch.float32)
+        h, w = out_shape_hw
+        return torch.zeros((1, h, w), dtype=torch.float32)
+
+    base = segmasks[0][1]
+    combined = np.zeros_like(base, dtype=np.float32)
+    for _, m, _ in segmasks:
+        combined = np.maximum(combined, m.astype(np.float32))
+    return torch.from_numpy(combined).unsqueeze(0)
+
+
+# ---------------------------
+# Impact-compatible detector wrapper objects
+# ---------------------------
+
+SEG = namedtuple(
+    "SEG",
+    [
+        "cropped_image",
+        "cropped_mask",
+        "confidence",
+        "crop_region",
+        "bbox",
+        "label",
+        "control_net_wrapper",
+    ],
+    defaults=[None],
+)
+
+
+class NO_BBOX_DETECTOR:
+    pass
+
+
+class NO_SEGM_DETECTOR:
+    pass
+
+
+def _create_segmasks(results):
+    # results = [labels, bboxes_xyxy, segms, confs]
+    bboxes = results[1]
+    segms = results[2]
+    confs = results[3]
+
+    out = []
+    for i in range(len(segms)):
+        out.append((bboxes[i], segms[i].astype(np.float32), confs[i]))
+    return out
+
+
+def _inference_bbox(model, image_pil: Image.Image, confidence: float = 0.3, device: str = ""):
+    pred = model(image_pil, conf=confidence, device=device)
+
+    bboxes = pred[0].boxes.xyxy.cpu().numpy()  # xyxy
+    if bboxes.shape[0] == 0:
+        return [[], [], [], []]
+
+    # Make simple rectangle masks for each bbox
+    np_img = np.array(image_pil)
+    if np_img.ndim == 2:
+        h, w = np_img.shape
+    else:
+        h, w = np_img.shape[0], np_img.shape[1]
+
+    segms = []
+    for x0, y0, x1, y1 in bboxes:
+        m = np.zeros((h, w), dtype=np.uint8)
+        x0i, y0i, x1i, y1i = int(x0), int(y0), int(x1), int(y1)
+        x0i = max(0, min(w - 1, x0i))
+        x1i = max(0, min(w, x1i))
+        y0i = max(0, min(h - 1, y0i))
+        y1i = max(0, min(h, y1i))
+        if cv2 is not None:
+            cv2.rectangle(m, (x0i, y0i), (x1i, y1i), 255, -1)
+        else:
+            m[y0i:y1i, x0i:x1i] = 255
+        segms.append((m > 0))
+
+    labels = []
+    confs = []
+    for i in range(len(bboxes)):
+        labels.append(pred[0].names[int(pred[0].boxes[i].cls.item())])
+        confs.append(pred[0].boxes[i].conf.detach().cpu().numpy())
+
+    return [labels, list(bboxes), segms, confs]
+
+
+def _inference_segm(model, image_pil: Image.Image, confidence: float = 0.3, device: str = ""):
+    pred = model(image_pil, conf=confidence, device=device)
+
+    bboxes = pred[0].boxes.xyxy.cpu().numpy()  # xyxy
+    if bboxes.shape[0] == 0:
+        return [[], [], [], []]
+
+    if pred[0].masks is None or pred[0].masks.data is None:
+        # fallback: no masks, treat like bbox
+        return _inference_bbox(model, image_pil, confidence=confidence, device=device)
+
+    segms = pred[0].masks.data.detach().cpu().numpy()  # (n, h, w) in model-space
+
+    # Resize masks back to original image size
+    h_orig = image_pil.size[1]
+    w_orig = image_pil.size[0]
+
+    results = [[], [], [], []]
+
+    for i in range(len(bboxes)):
+        results[0].append(pred[0].names[int(pred[0].boxes[i].cls.item())])
+        results[1].append(bboxes[i])
+
+        mask = torch.from_numpy(segms[i]).float()
+        mask = F.interpolate(mask.unsqueeze(0).unsqueeze(0), size=(h_orig, w_orig), mode="bilinear", align_corners=False)
+        mask = mask.squeeze(0).squeeze(0)
+
+        results[2].append(mask.numpy())
+        results[3].append(pred[0].boxes[i].conf.detach().cpu().numpy())
+
+    return results
+
+
+class SaliaUltraBBoxDetector:
+    def __init__(self, model):
+        self.model = model
+
+    def detect(self, image, threshold, dilation, crop_factor, drop_size=1, detailer_hook=None):
+        drop_size = max(int(drop_size), 1)
+        detected = _inference_bbox(self.model, tensor2pil(image), confidence=float(threshold))
+        segmasks = _create_segmasks(detected)
+
+        if int(dilation) > 0:
+            segmasks = dilate_masks(segmasks, int(dilation))
+
+        items = []
+        h = image.shape[1]
+        w = image.shape[2]
+
+        for (bbox, mask, conf), label in zip(segmasks, detected[0]):
+            x1, y1, x2, y2 = bbox
+            if (x2 - x1) > drop_size and (y2 - y1) > drop_size:
+                crop_region = make_crop_region(w, h, bbox, float(crop_factor))
+
+                if detailer_hook is not None and hasattr(detailer_hook, "post_crop_region"):
+                    crop_region = detailer_hook.post_crop_region(w, h, bbox, crop_region)
+
+                cropped_image = crop_image(image, crop_region)
+                cropped_mask = crop_ndarray2(mask, crop_region)
+
+                items.append(SEG(cropped_image, cropped_mask, conf, crop_region, bbox, label, None))
+
+        segs = (image.shape[1], image.shape[2]), items
+
+        if detailer_hook is not None and hasattr(detailer_hook, "post_detection"):
+            segs = detailer_hook.post_detection(segs)
+
+        return segs
+
+    def detect_combined(self, image, threshold, dilation):
+        detected = _inference_bbox(self.model, tensor2pil(image), confidence=float(threshold))
+        segmasks = _create_segmasks(detected)
+        if int(dilation) > 0:
+            segmasks = dilate_masks(segmasks, int(dilation))
+        return combine_masks(segmasks, out_shape_hw=(image.shape[1], image.shape[2]))
+
+    def setAux(self, x):
+        pass
+
+
+class SaliaUltraSegmDetector:
+    def __init__(self, model):
+        self.model = model
+
+    def detect(self, image, threshold, dilation, crop_factor, drop_size=1, detailer_hook=None):
+        drop_size = max(int(drop_size), 1)
+        detected = _inference_segm(self.model, tensor2pil(image), confidence=float(threshold))
+        segmasks = _create_segmasks(detected)
+
+        if int(dilation) > 0:
+            segmasks = dilate_masks(segmasks, int(dilation))
+
+        items = []
+        h = image.shape[1]
+        w = image.shape[2]
+
+        for (bbox, mask, conf), label in zip(segmasks, detected[0]):
+            x1, y1, x2, y2 = bbox
+            if (x2 - x1) > drop_size and (y2 - y1) > drop_size:
+                crop_region = make_crop_region(w, h, bbox, float(crop_factor))
+
+                if detailer_hook is not None and hasattr(detailer_hook, "post_crop_region"):
+                    crop_region = detailer_hook.post_crop_region(w, h, bbox, crop_region)
+
+                cropped_image = crop_image(image, crop_region)
+                cropped_mask = crop_ndarray2(mask, crop_region)
+
+                items.append(SEG(cropped_image, cropped_mask, conf, crop_region, bbox, label, None))
+
+        segs = (image.shape[1], image.shape[2]), items
+
+        if detailer_hook is not None and hasattr(detailer_hook, "post_detection"):
+            segs = detailer_hook.post_detection(segs)
+
+        return segs
+
+    def detect_combined(self, image, threshold, dilation):
+        detected = _inference_segm(self.model, tensor2pil(image), confidence=float(threshold))
+        segmasks = _create_segmasks(detected)
+        if int(dilation) > 0:
+            segmasks = dilate_masks(segmasks, int(dilation))
+        return combine_masks(segmasks, out_shape_hw=(image.shape[1], image.shape[2]))
+
+    def setAux(self, x):
+        pass
+
+
+# ---------------------------
+# The actual ComfyUI Node
+# ---------------------------
+
+class SaliaUltralyticsDetectorProvider2:
+    """Load an Ultralytics `.pt` model and provide Impact-compatible detectors."""
+
+    @classmethod
+    def INPUT_TYPES(cls):
+        bboxs = ["bbox/" + x for x in folder_paths.get_filename_list("ultralytics_bbox")]
+        segms = ["segm/" + x for x in folder_paths.get_filename_list("ultralytics_segm")]
+        return {"required": {"model_name": (bboxs + segms,)}}
+
+    RETURN_TYPES = ("BBOX_DETECTOR", "SEGM_DETECTOR")
+    FUNCTION = "doit"
+    CATEGORY = "Salia/Detectors"
+
+    def doit(self, model_name: str):
+        # First, allow selecting a file like "bbox/foo.pt" that lives under models/ultralytics/bbox
+        model_path = folder_paths.get_full_path("ultralytics", model_name)
+
+        if model_path is None:
+            if model_name.startswith("bbox/"):
+                model_path = folder_paths.get_full_path("ultralytics_bbox", model_name[5:])
+            elif model_name.startswith("segm/"):
+                model_path = folder_paths.get_full_path("ultralytics_segm", model_name[5:])
+
+        if model_path is None:
+            cands = []
+            try:
+                cands.extend(folder_paths.get_folder_paths("ultralytics"))
+                if model_name.startswith("bbox/"):
+                    cands.extend(folder_paths.get_folder_paths("ultralytics_bbox"))
+                elif model_name.startswith("segm/"):
+                    cands.extend(folder_paths.get_folder_paths("ultralytics_segm"))
+            except Exception:
+                pass
+
+            formatted = "\n\t".join(cands)
+            raise ValueError(
+                f"[Salia Ultralytics] model file '{model_name}' was not found.\n"
+                f"Searched these folders:\n\t{formatted}\n"
+                f"Tip: put bbox models in 'models/ultralytics/bbox' or segm models in 'models/ultralytics/segm'."
+            )
+
+        # NEW: print + log the resolved on-disk path (and any duplicates)
+        matches = _find_all_model_paths(model_name)
+        _log_selected_model(model_name, os.path.abspath(model_path), matches)
+
+        model = _load_yolo(model_path)
+
+        if model_name.startswith("bbox/"):
+            return SaliaUltraBBoxDetector(model), NO_SEGM_DETECTOR()
+        else:
+            return SaliaUltraBBoxDetector(model), SaliaUltraSegmDetector(model)
+
+
+NODE_CLASS_MAPPINGS = {
+    "SaliaUltralyticsDetectorProvider2": SaliaUltralyticsDetectorProvider2,
+}
+
+NODE_DISPLAY_NAME_MAPPINGS = {
+    "SaliaUltralyticsDetectorProvider2": "Salia Ultralytics Detector 2 (Salia)",
+}