Upload dn04.py

dn04.py

@@ -0,0 +1,606 @@
+# FaceDetailerStandalone_MIN_FIXED_FAST_EMBEDDED_SAM.py
+# One-node Face Detailer (image-only) with fixed settings + embedded Ultralytics bbox detector + embedded SAM loader.
+# - Output parity with Impact Pack Face Detailer at the same settings
+# - No separate bbox-detector node; detector is cached/constructed internally
+# - No separate SAM loader node; SAM is cached/constructed internally
+# - Lightweight runtime overhead (cached imports, inference_mode, fused layers, TF32, FP16 on CUDA)
+
+import os
+from dataclasses import dataclass
+from typing import List, Tuple, Optional
+
+import warnings
+warnings.filterwarnings("ignore")
+
+# Silence OpenCV before importing it (env var) and after (setLogLevel)
+os.environ["OPENCV_LOG_LEVEL"] = "ERROR"
+
+import numpy as np
+import torch
+import comfy
+from PIL import Image
+import cv2
+
+try:
+    if hasattr(cv2, "setLogLevel"):
+        try:
+            lvl = cv2.LOG_LEVEL_ERROR if hasattr(cv2, "LOG_LEVEL_ERROR") else 3  # 3 == error
+            cv2.setLogLevel(lvl)
+        except Exception:
+            pass
+except Exception:
+    pass
+
+# ---------------- Fixed FaceDetailer settings (do not expose in UI) ----------------
+# GUIDE_SIZE = 512
+# GUIDE_SIZE_FOR_BBOX = True
+# MAX_SIZE = 1024
+# STEPS = 30
+# CFG = 7.0
+# SCHEDULER = "simple"
+# DENOISE = 0.5
+# FEATHER = 5
+# NOISE_MASK = True
+# FORCE_INPAINT = True
+# BBOX_THRESHOLD = 0.5
+# BBOX_DILATION = 10
+# BBOX_CROP_FACTOR = 3.0
+# DROP_SIZE = 10
+# SAM_DETECTION_HINT = "center-1"
+# SAM_DILATION = 0
+# SAM_THRESHOLD = 0.93
+# SAM_BBOX_EXPANSION = 0
+# SAM_MASK_HINT_THRESHOLD = 0.7
+# SAM_MASK_HINT_USE_NEGATIVE = "False"
+# WILDCARD = ""
+# CYCLE = 1
+# INPAINT_MODEL = False
+# NOISE_MASK_FEATHER = 20
+# TILED_ENCODE = False
+# TILED_DECODE = False
+# ---------------------------------------------------------------------
+
+# ---------------- Ultralytics / YOLO detector integration (embedded) ----------------
+
+# Torch runtime perf switches
+torch.backends.cudnn.benchmark = True  # autotune best conv algorithms
+if torch.cuda.is_available():
+    torch.backends.cuda.matmul.allow_tf32 = True
+    torch.backends.cudnn.allow_tf32 = True
+    try:
+        torch.set_float32_matmul_precision("high")  # PyTorch 2.x
+    except Exception:
+        pass
+
+# Optional Impact Pack interop (SEG type)
+try:
+    # If Impact Pack is installed, use its SEG to be perfectly compatible.
+    from impact.core import SEG as _IMPACT_SEG  # type: ignore
+    _USE_IMPACT_SEG = True
+except Exception:
+    _USE_IMPACT_SEG = False
+
+    @dataclass
+    class _LocalSEG:
+        cropped_image: Optional[torch.Tensor]
+        cropped_mask: np.ndarray  # 2D float32 [0..1]
+        confidence: float
+        crop_region: Tuple[int, int, int, int]  # (x1,y1,x2,y2)
+        bbox: Tuple[int, int, int, int]         # (x1,y1,x2,y2)
+        label: str
+        control_net_wrapper: Optional[object] = None
+
+SEG = _IMPACT_SEG if _USE_IMPACT_SEG else _LocalSEG
+
+# ---------------------------------------------------------------------
+# LOCAL ASSET PATHS (no hardcoded absolute paths)
+# ---------------------------------------------------------------------
+
+# Base directory of this node file (cross-platform, works on RunPod/ComfyUI)
+BASE_DIR = os.path.dirname(os.path.abspath(__file__))
+
+# Local YOLO model path inside this custom node folder
+YOLO_MODEL_PATH = os.path.join(BASE_DIR, "assets", "face_yolov8m_salia.pt")
+YOLO_IMGSZ = 640
+
+# Local SAM checkpoint path inside this custom node folder
+SAM_CKPT_PATH = os.path.join(BASE_DIR, "assets", "sam_vit_b_01ec64_salia.pth")
+
+# Cached instances (process-local)
+_CACHED_YOLO_MODEL = None
+_CACHED_ULTRA_DETECTOR = None
+
+
+def _tensor_to_pil(image: torch.Tensor) -> Image.Image:
+    # image: [1, H, W, 3], float(0..1)
+    img = image[0].detach().cpu().clamp(0, 1).numpy()
+    img = (img * 255.0).round().astype(np.uint8)  # (H, W, 3) RGB
+    return Image.fromarray(img, mode="RGB")
+
+
+def _make_crop_region(w: int, h: int, bbox_xyxy, crop_factor: float) -> Tuple[int, int, int, int]:
+    x1, y1, x2, y2 = map(int, bbox_xyxy)
+    cx = (x1 + x2) * 0.5
+    cy = (y1 + y2) * 0.5
+    bw = (x2 - x1)
+    bh = (y2 - y1)
+    new_w = max(1, int(bw * crop_factor))
+    new_h = max(1, int(bh * crop_factor))
+    # center to image
+    nx1 = int(max(0, round(cx - new_w * 0.5)))
+    ny1 = int(max(0, round(cy - new_h * 0.5)))
+    nx2 = int(min(w, nx1 + new_w))
+    ny2 = int(min(h, ny1 + new_h))
+    # clamp again
+    nx1 = max(0, min(nx1, w - 1))
+    ny1 = max(0, min(ny1, h - 1))
+    nx2 = max(nx1 + 1, min(nx2, w))
+    ny2 = max(ny1 + 1, min(ny2, h))
+    return (nx1, ny1, nx2, ny2)
+
+
+def _crop_tensor_image(image: torch.Tensor, crop: Tuple[int, int, int, int]) -> torch.Tensor:
+    # image: [1,H,W,3]; crop: (x1,y1,x2,y2)
+    x1, y1, xb, yb = crop
+    return image[:, y1:yb, x1:xb, :].contiguous()
+
+
+def _crop_ndarray(mask: np.ndarray, crop: Tuple[int, int, int, int]) -> np.ndarray:
+    # mask: [H,W] float/bool/uint8; crop: (x1,y1,x2,y2)
+    x1, y1, xb, yb = crop
+    return mask[int(y1):int(yb), int(x1):int(xb)]
+
+
+def _dilate_masks(segmasks: List[Tuple[np.ndarray, np.ndarray, float]], factor: int):
+    if factor == 0 or not segmasks:
+        return segmasks
+    k = abs(int(factor))
+    if k < 1:
+        return segmasks
+    kernel = cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (k, k))
+    do_dilate = factor > 0
+    out = []
+    for (bbox, m, conf) in segmasks:
+        u8 = (m * 255.0).astype(np.uint8) if m.dtype != np.uint8 else m
+        d = cv2.dilate(u8, kernel, iterations=1) if do_dilate else cv2.erode(u8, kernel, iterations=1)
+        out.append((bbox, d.astype(np.float32) / 255.0, conf))
+    return out
+
+
+def _combine_masks(segmasks: List[Tuple[np.ndarray, np.ndarray, float]]) -> Optional[torch.Tensor]:
+    if not segmasks:
+        return None
+    h = segmasks[0][1].shape[0]
+    w = segmasks[0][1].shape[1]
+    acc = np.zeros((h, w), dtype=np.uint8)
+    for _, m, _ in segmasks:
+        u8 = (m * 255.0).astype(np.uint8) if m.dtype != np.uint8 else m
+        acc = cv2.bitwise_or(acc, u8)
+    return torch.from_numpy(acc.astype(np.float32) / 255.0)  # [H,W], float32 0..1 CPU
+
+
+def _pick_device_str(user_device: str = "") -> str:
+    if user_device:
+        return user_device
+    return "cuda" if torch.cuda.is_available() else "cpu"
+
+
+@torch.inference_mode()
+def _inference_bbox(model, image_pil: Image.Image, confidence: float = 0.3, device: str = ""):
+    """
+    Returns results = [labels(str), bboxes(xyxy), segms(full-image bool masks), conf(float)]
+    For bbox models, segm "masks" are rectangles from the boxes (Subpack parity).
+    """
+    pred = model(
+        image_pil,
+        conf=confidence,
+        device=_pick_device_str(device),
+        verbose=False,
+        imgsz=YOLO_IMGSZ,  # fixed size can be faster
+    )
+
+    p0 = pred[0]
+    boxes = p0.boxes
+    bboxes = boxes.xyxy.detach().cpu().numpy()  # (N,4) float, xyxy
+
+    W, H = image_pil.size
+    segms = []
+    for x0, y0, x1, y1 in bboxes:
+        m = np.zeros((H, W), np.uint8)
+        cv2.rectangle(m, (int(x0), int(y0)), (int(x1), int(y1)), 255, -1)
+        segms.append(m.astype(bool))
+
+    if bboxes.shape[0] == 0:
+        return [[], [], [], []]
+
+    results = [[], [], [], []]
+    names = p0.names
+    for i, (bbox, segm) in enumerate(zip(bboxes, segms)):
+        cls_i = int(boxes.cls[i].item())
+        results[0].append(names[cls_i])
+        results[1].append(bbox)
+        results[2].append(segm)
+        results[3].append(float(boxes.conf[i].item()))
+    return results
+
+
+def _create_segmasks(results):
+    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
+
+
+class UltraBBoxDetector:
+    def __init__(self, yolo_model):
+        self.bbox_model = yolo_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.bbox_model, _tensor_to_pil(image), threshold)
+        segmasks = _create_segmasks(detected)
+        if int(dilation) != 0:
+            segmasks = _dilate_masks(segmasks, int(dilation))
+
+        H = int(image.shape[1])
+        W = int(image.shape[2])
+        items = []
+        for (bbox_xyxy, full_mask, conf), label in zip(segmasks, detected[0]):
+            x1, y1, x2, y2 = map(int, bbox_xyxy)
+            if (x2 - x1) > drop_size and (y2 - y1) > drop_size:
+                crop_region = _make_crop_region(W, H, (x1, y1, x2, y2), 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, (x1, y1, x2, y2), crop_region)
+
+                cropped_image = _crop_tensor_image(image, crop_region)
+                cropped_mask = _crop_ndarray(full_mask, crop_region).astype(np.float32)
+                items.append(SEG(cropped_image, cropped_mask, float(conf), crop_region, (x1, y1, x2, y2), str(label), None))
+
+        segs = ((H, W), 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.bbox_model, _tensor_to_pil(image), threshold)
+        segmasks = _create_segmasks(detected)
+        if int(dilation) != 0:
+            segmasks = _dilate_masks(segmasks, int(dilation))
+        return _combine_masks(segmasks)
+
+    def setAux(self, x):
+        # kept for signature parity
+        pass
+
+
+def _load_ultralytics_model(model_path: str):
+    # Import here so that module import doesn't hard-fail if ultralytics is missing
+    try:
+        from ultralytics import YOLO
+    except Exception as e:
+        raise RuntimeError(
+            "[FaceDetailerStandalone] The 'ultralytics' package is required for the embedded bbox detector.\n"
+            "Install in your ComfyUI python:  python -m pip install --upgrade ultralytics"
+        ) from e
+
+    if not os.path.isfile(model_path):
+        raise FileNotFoundError(
+            "[FaceDetailerStandalone] Embedded YOLO model file not found.\n"
+            f"Expected at: {model_path}\n"
+            "Please place 'face_yolov8m_salia.pt' in the 'assets' folder next to this node."
+        )
+
+    yolo = YOLO(model_path)
+
+    # One-time graph/model optimizations
+    try:
+        dev = _pick_device_str()
+        try:
+            yolo.to(dev)  # newer Ultralytics
+        except Exception:
+            yolo.model.to(dev)  # older versions
+    except Exception:
+        pass
+
+    # Fuse Conv+BN where possible (small speedup)
+    try:
+        yolo.fuse()
+    except Exception:
+        pass
+
+    # Use half precision weights on CUDA (big win; safe for inference)
+    try:
+        if torch.cuda.is_available():
+            yolo.model.half()
+    except Exception:
+        pass
+
+    return yolo
+
+
+def _get_embedded_detector():
+    global _CACHED_YOLO_MODEL, _CACHED_ULTRA_DETECTOR
+    if _CACHED_ULTRA_DETECTOR is not None:
+        return _CACHED_ULTRA_DETECTOR
+    if _CACHED_YOLO_MODEL is None:
+        _CACHED_YOLO_MODEL = _load_ultralytics_model(YOLO_MODEL_PATH)
+    _CACHED_ULTRA_DETECTOR = UltraBBoxDetector(_CACHED_YOLO_MODEL)
+    return _CACHED_ULTRA_DETECTOR
+
+# ---------------- Embedded SAM loader (GPU-only, hardcoded path, reuse one predictor) ----------------
+# Matches your SAMLoaderStandalone design, but embedded + cached.
+
+
+def _to_numpy_rgb(image_tensor):
+    """
+    Comfy 'IMAGE' is NHWC in [0..1]. Convert to uint8 HxWx3 RGB numpy.
+    Accepts torch.Tensor (NHWC) or numpy already in HWC.
+    """
+    if isinstance(image_tensor, torch.Tensor):
+        img = image_tensor
+        if img.dim() == 4 and img.shape[0] == 1:
+            img = img[0]
+        img = (img.clamp(0, 1) * 255.0).to(torch.uint8).cpu().numpy()  # HWC
+        return img
+    elif isinstance(image_tensor, np.ndarray):
+        if image_tensor.dtype != np.uint8:
+            img = np.clip(image_tensor, 0, 255).astype(np.uint8)
+        else:
+            img = image_tensor
+        return img
+    else:
+        raise TypeError(f"Unsupported image type for SAM: {type(image_tensor)}")
+
+
+class _SAMWrapperGPUOnlyFast:
+    """
+    FaceDetailer-compatible wrapper:
+      - Stays on CUDA
+      - Reuses a single SamPredictor
+      - predict(image, points, plabs, bbox, threshold) -> list[HxW float32 CPU masks]
+    """
+    def __init__(self, model):
+        self.model = model
+        dev = comfy.model_management.get_torch_device()
+        if "cuda" not in str(dev).lower():
+            raise RuntimeError(
+                f"[FaceDetailerStandalone] GPU-only SAM: CUDA device not available (got '{dev}')."
+            )
+        self._device = dev
+        self.model.to(self._device).eval()
+        # Lazy import for segment_anything predictor
+        from segment_anything import SamPredictor  # type: ignore
+        # Reuse one predictor instance (cheaper than re-creating every call)
+        self._predictor = SamPredictor(self.model)
+
+    def prepare_device(self):
+        if "cuda" not in str(self._device).lower():
+            raise RuntimeError("[FaceDetailerStandalone] CUDA device lost/unavailable for SAM.")
+
+    def release_device(self):
+        # GPU-only; keep on GPU (no-op)
+        pass
+
+    @torch.inference_mode()
+    def predict(self, image, points, plabs, bbox, threshold: float):
+        """
+        image: Comfy IMAGE (NHWC, [0..1]) or numpy
+        points: list[[x,y], ...] or None
+        plabs:  list[int] (1=fg, 0=bg) or None
+        bbox:   [x1,y1,x2,y2] or None
+        threshold: float in [0..1]
+        returns: list of HxW float32 CPU masks (0/1)
+        """
+        self.prepare_device()
+
+        np_img = _to_numpy_rgb(image)
+        # Some builds call set_image(img, "RGB"); accept both signatures.
+        try:
+            self._predictor.set_image(np_img, "RGB")
+        except TypeError:
+            self._predictor.set_image(np_img)
+
+        pc = np.array(points, dtype=np.float32) if points else None
+        pl = np.array(plabs, dtype=np.int32) if plabs else None
+        bx = np.array(bbox, dtype=np.float32) if bbox is not None else None
+
+        # Keep provided behavior: multimask_output=False
+        masks, scores, _ = self._predictor.predict(
+            point_coords=pc,
+            point_labels=pl,
+            box=bx,
+            multimask_output=False
+        )
+
+        out = []
+        if masks is not None and scores is not None:
+            for m, s in zip(masks, scores):
+                if float(s) >= float(threshold):
+                    if isinstance(m, torch.Tensor):
+                        t = m.to(torch.float32).cpu()
+                    else:
+                        t = torch.from_numpy(m.astype(np.float32)).cpu()
+                    out.append(t)
+        return out
+
+
+# Cache for SAM
+_CACHED_SAM_MODEL = None
+
+
+def _get_embedded_sam():
+    """Load SAM vit_b from SAM_CKPT_PATH and attach GPU-only fast wrapper, cached."""
+    global _CACHED_SAM_MODEL
+    if _CACHED_SAM_MODEL is not None:
+        return _CACHED_SAM_MODEL
+
+    if not os.path.isfile(SAM_CKPT_PATH):
+        raise FileNotFoundError(
+            f"[FaceDetailerStandalone] SAM checkpoint not found:\n  {SAM_CKPT_PATH}\n"
+            f"Place 'sam_vit_b_01ec64_salia.pth' in the 'assets' folder next to this node."
+        )
+
+    # Import here to avoid module import failure at file load time
+    try:
+        from segment_anything import sam_model_registry  # type: ignore
+    except Exception as e:
+        raise RuntimeError(
+            "[FaceDetailerStandalone] 'segment_anything' is not installed for embedded SAM. "
+            "Install in your Comfy python, e.g.:  python -m pip install "
+            "git+https://github.com/facebookresearch/segment-anything"
+        ) from e
+
+    # Fixed to vit_b (matches 'sam_vit_b_01ec64' weights)
+    sam = sam_model_registry['vit_b'](checkpoint=SAM_CKPT_PATH)
+    sam.eval()  # ensure eval mode
+
+    # Attach GPU-only, faster wrapper
+    wrapper = _SAMWrapperGPUOnlyFast(sam)
+    sam.sam_wrapper = wrapper
+
+    _CACHED_SAM_MODEL = sam
+    return _CACHED_SAM_MODEL
+
+# ---------------- Impact Pack Face Detailer binding ----------------
+_ENHANCE_FACE = None
+_IMPORT_ERR = None
+try:
+    from impact.impact_pack import FaceDetailer as _FD
+    _ENHANCE_FACE = _FD.enhance_face
+except Exception as _e:
+    _IMPORT_ERR = _e
+    _ENHANCE_FACE = None
+
+# ---------------- Single public node ----------------
+class dn_04:
+    @classmethod
+    def INPUT_TYPES(cls):
+        # Only essential, connectable parts remain editable. (No bbox or SAM inputs.)
+        return {
+            "required": {
+                "image": ("IMAGE",),
+                "model": ("MODEL", {"tooltip": "If `ImpactDummyInput` is connected to model, inference is skipped."}),
+                "clip": ("CLIP",),
+                "vae": ("VAE",),
+
+                # Keep sampler selectable; all other knobs are fixed
+                "sampler_name": (comfy.samplers.KSampler.SAMPLERS,),
+
+                # Conditioning stays connectable
+                "positive": ("CONDITIONING",),
+                "negative": ("CONDITIONING",),
+
+                # Keep seed editable but fixed after generate for reproducibility
+                "seed": ("INT", {
+                    "default": 0,
+                    "min": 0,
+                    "max": 0xffffffffffffffff,
+                    "step": 1,
+                    "control_after_generate": "fixed",
+                }),
+            },
+            "optional": {
+                # No external SAM input; embedded
+            }
+        }
+
+    RETURN_TYPES = ("IMAGE",)
+    RETURN_NAMES = ("image",)
+    FUNCTION = "doit"
+    CATEGORY = "ImpactPack/Standalone"
+    DESCRIPTION = (
+        "Face Detailer with requested parameters hardcoded (non-editable), "
+        "and embedded Ultralytics face bbox detector + embedded SAM (no external input nodes). "
+        "Optimized call path (cached imports + inference_mode) for lower overhead; "
+        "results identical to Impact Pack Face Detailer at the same settings."
+    )
+
+    def doit(
+        self,
+        image, model, clip, vae,
+        sampler_name,
+        positive, negative,
+        seed,
+    ):
+        if _ENHANCE_FACE is None:
+            raise RuntimeError(
+                "ComfyUI-Impact-Pack is required for Face Detailer logic. "
+                "Please install/enable ComfyUI-Impact-Pack."
+            ) from _IMPORT_ERR
+
+        # Embedded detector & SAM (cached)
+        bbox_detector = _get_embedded_detector()
+        sam_model_opt = _get_embedded_sam()
+
+        enhance = _ENHANCE_FACE
+
+        # Determine batch size safely
+        B = image.shape[0] if (hasattr(image, "shape") and image.ndim == 4) else 1
+
+        # No autograd, faster kernel choices, identical math for inference
+        with torch.inference_mode():
+            if B == 1:
+                # Fast-path for single image (avoid list + cat)
+                single = image[0] if image.ndim == 4 else image  # [H,W,C]
+                enhanced_img, _, _, _, _ = enhance(
+                    single.unsqueeze(0),  # -> [1,H,W,C]
+                    model, clip, vae,
+                    512, True, 1024,          # guide_size, guide_for_bbox, max_size
+                    seed, 38, 7.0,            # steps, cfg
+                    sampler_name, "simple",   # scheduler name
+                    positive, negative,
+                    0.4, 5, True, True,       # denoise, feather, noise_mask, force_inpaint
+                    0.5, 10, 3.0,             # bbox_threshold, bbox_dilation, bbox_crop_factor
+                    "center-1", 0, 0.93, 0,   # sam_detection_hint, sam_dilation, sam_threshold, sam_bbox_expansion
+                    0.7, "False",             # sam_mask_hint_threshold, sam_mask_hint_use_negative
+                    10, bbox_detector,        # drop_size, bbox_detector
+                    # Internals not exposed (kept fixed/None)
+                    segm_detector=None, sam_model_opt=sam_model_opt,
+                    wildcard_opt="", detailer_hook=None,
+                    refiner_ratio=None, refiner_model=None, refiner_clip=None,
+                    refiner_positive=None, refiner_negative=None,
+                    cycle=1, inpaint_model=False,
+                    noise_mask_feather=20,
+                    scheduler_func_opt=None,
+                    tiled_encode=False, tiled_decode=False,
+                )
+                return (enhanced_img,)
+
+            # Batch of images; per-frame process with seed+i
+            out_imgs = []
+            for i, single in enumerate(image.unbind(0)):
+                enhanced_img, _, _, _, _ = enhance(
+                    single.unsqueeze(0),  # [1,H,W,C]
+                    model, clip, vae,
+                    512, True, 1024,
+                    seed + i, 30, 7.0,
+                    sampler_name, "simple",
+                    positive, negative,
+                    0.5, 5, True, True,
+                    0.5, 10, 3.0,
+                    "center-1", 0, 0.93, 0,
+                    0.7, "False",
+                    10, bbox_detector,
+                    segm_detector=None, sam_model_opt=sam_model_opt,
+                    wildcard_opt="", detailer_hook=None,
+                    refiner_ratio=None, refiner_model=None, refiner_clip=None,
+                    refiner_positive=None, refiner_negative=None,
+                    cycle=1, inpaint_model=False,
+                    noise_mask_feather=20,
+                    scheduler_func_opt=None,
+                    tiled_encode=False, tiled_decode=False,
+                )
+                out_imgs.append(enhanced_img)
+
+            return (torch.cat(out_imgs, dim=0),)
+
+
+NODE_CLASS_MAPPINGS = {
+    "dn_04": dn_04,
+}
+NODE_DISPLAY_NAME_MAPPINGS = {
+    "dn_04": "dn_04",
+}