Merge pull request #4 from IBB666/copilot/fix-node-installation-issue

Restore ComfyUI node registration for IBB_POSE

ibb_pose_nodes.py

@@ -1,5 +1,7 @@
 import warnings
 import logging
+import colorsys
+import os
 
 # Suppress common warnings to reduce noise
 warnings.filterwarnings("ignore", category=UserWarning)
@@ -109,15 +111,12 @@ except ImportError:
             os.makedirs(full_output_folder, exist_ok=True)
             return full_output_folder, base_filename, 0, subfolder, filename_prefix
     folder_paths = MockFolderPaths()
-from huggingface_hub import snapshot_download
 import sys
 from pathlib import Path
 import glob
 import cv2
 import math
-import matplotlib.colors
 import tempfile
-from torchvision import transforms
 try:
     import model_management
 except ImportError:
@@ -157,6 +156,7 @@ except Exception as _e:
     )
 
 # --- Add custom folder paths to ComfyUI ---
+EMPTY_EMBED_DIR = str(Path(__file__).resolve().parent / "empty_text_encoder")
 SDPOSE_MODEL_DIR = os.path.join(folder_paths.models_dir, "IBB_POSE")
 YOLO_MODEL_DIR = os.path.join(folder_paths.models_dir, "yolo")
 
@@ -219,7 +219,14 @@ HAND_EDGES = [
 ]
 
 
+def _hand_edge_color(edge_index: int) -> tuple[int, int, int]:
+    color = colorsys.hsv_to_rgb(edge_index / float(len(HAND_EDGES)), 1.0, 1.0)
+    return tuple(int(channel * 255) for channel in color)
+
+
 def _download_file(url: str, dest: str) -> None:
+    import urllib.request
+
     os.makedirs(os.path.dirname(dest), exist_ok=True)
     print(f"IBB_POSE: Downloading {os.path.basename(dest)} ...")
     try:
@@ -427,8 +434,7 @@ def draw_wholebody_keypoints_openpose_style(canvas, keypoints, scores=None, thre
             if scores is not None and (scores[idx1] < threshold or scores[idx2] < threshold): continue
             x1, y1 = int(keypoints[idx1][0]), int(keypoints[idx1][1]); x2, y2 = int(keypoints[idx2][0]), int(keypoints[idx2][1])
             if x1 > 0.01 and y1 > 0.01 and x2 > 0.01 and y2 > 0.01 and 0 <= x1 < W and 0 <= y1 < H and 0 <= x2 < W and 0 <= y2 < H:
-                color = matplotlib.colors.hsv_to_rgb([ie / float(len(hand_edges)), 1.0, 1.0]) * 255
-                cv2.line(canvas, (x1, y1), (x2, y2), color, thickness=face_hand_line_width)
+                cv2.line(canvas, (x1, y1), (x2, y2), _hand_edge_color(ie), thickness=face_hand_line_width)
         for i in range(92, 113):
             if scores is not None and scores[i] < threshold: continue
             x, y = int(keypoints[i][0]), int(keypoints[i][1])
@@ -439,8 +445,7 @@ def draw_wholebody_keypoints_openpose_style(canvas, keypoints, scores=None, thre
             if scores is not None and (scores[idx1] < threshold or scores[idx2] < threshold): continue
             x1, y1 = int(keypoints[idx1][0]), int(keypoints[idx1][1]); x2, y2 = int(keypoints[idx2][0]), int(keypoints[idx2][1])
             if x1 > 0.01 and y1 > 0.01 and x2 > 0.01 and y2 > 0.01 and 0 <= x1 < W and 0 <= y1 < H and 0 <= x2 < W and 0 <= y2 < H:
-                color = matplotlib.colors.hsv_to_rgb([ie / float(len(hand_edges)), 1.0, 1.0]) * 255
-                cv2.line(canvas, (x1, y1), (x2, y2), color, thickness=face_hand_line_width)
+                cv2.line(canvas, (x1, y1), (x2, y2), _hand_edge_color(ie), thickness=face_hand_line_width)
         for i in range(113, 134):
             if scores is not None and i < len(scores) and scores[i] < threshold: continue
             x, y = int(keypoints[i][0]), int(keypoints[i][1])
@@ -540,9 +545,9 @@ def groundingdino_predict(dino_model_wrapper, image_pil, prompt, threshold):
 # --- Processing functions (adapted from SDPose_gradio.py) ---
 # (Functions detect_person_yolo, preprocess_image_for_sdpose, restore_keypoints_to_original, convert_to_openpose_json are omitted for brevity but would be pasted here)
 def detect_person_yolo(image, yolo_model_path, confidence_threshold=0.5):
-    if not YOLO_AVAILABLE: return [[0, 0, image.shape[1], image.shape[0]]], False
+    if not ULTRALYTICS_AVAILABLE: return [[0, 0, image.shape[1], image.shape[0]]], False
     try:
-        model = YOLO(yolo_model_path)
+        model = _YOLO(yolo_model_path)
         results = model(image, verbose=False)
         person_bboxes = []
         for result in results:
@@ -557,6 +562,8 @@ def detect_person_yolo(image, yolo_model_path, confidence_threshold=0.5):
         return [[0, 0, image.shape[1], image.shape[0]]], False
 
 def preprocess_image_for_sdpose(image, bbox=None, input_size=(768, 1024)):
+    from torchvision import transforms
+
     if isinstance(image, np.ndarray):
         pil_image = Image.fromarray(cv2.cvtColor(image, cv2.COLOR_BGR2RGB))
     else: pil_image = image
@@ -795,7 +802,7 @@ class IBBYOLOModelLoader:
     CATEGORY = "IBB_POSE"
 
     def load_yolo_model(self, model_name):
-        if not YOLO_AVAILABLE:
+        if not ULTRALYTICS_AVAILABLE:
             raise ImportError("ultralytics library is not available. Please install it to use YOLO models.")
         
         model_path = folder_paths.get_full_path("yolo", model_name)
@@ -803,7 +810,7 @@ class IBBYOLOModelLoader:
             raise FileNotFoundError(f"YOLO model not found: {model_name}")
             
         print(f"IBB_POSE: Loading YOLO model from {model_path}")
-        model = YOLO(model_path)
+        model = _YOLO(model_path)
         return (model,)
 
 
@@ -834,6 +841,13 @@ class IBBPoseModelLoader:
         return os.path.join(SDPOSE_MODEL_DIR, repo_name)
 
     def load_pose_model(self, model_type, unet_precision, device, unload_on_finish):
+        from diffusers import AutoencoderKL, DDPMScheduler, UNet2DConditionModel
+        from huggingface_hub import snapshot_download
+        from safetensors.torch import load_file
+
+        from .models.HeatmapHead import get_heatmap_head
+        from .models.ModifiedUNet import Modified_forward
+
         repo_id = {
             "Body": "teemosliang/SDPose-Body",
             "WholeBody": "teemosliang/SDPose-Wholebody"
@@ -844,8 +858,6 @@ class IBBPoseModelLoader:
         if not os.path.exists(os.path.join(model_path, "unet")):
             print(f"IBB_POSE: Downloading model from {repo_id} to {model_path}")
             snapshot_download(repo_id=repo_id, local_dir=model_path, local_dir_use_symlinks=False)
-
-    def load_model(self, skeleton_type, model_size, device, precision, auto_download):
         if not TORCH_AVAILABLE:
             raise ImportError(
                 "IBB_POSE: torch is required to run the node. "
@@ -857,32 +869,31 @@ class IBBPoseModelLoader:
         else:
             device = torch.device(device)
 
-        # --- Precision Handling ---
         dtype = {"fp32": torch.float32, "fp16": torch.float16, "bf16": torch.bfloat16}[unet_precision]
-        if device.type == 'cpu' and dtype == torch.float16:
-            print("IBB_POSE Warning: FP16 is not supported on CPU. Falling back to FP32.")
+        if device.type == "cpu" and dtype != torch.float32:
+            print("IBB_POSE Warning: FP16/BF16 is not supported on CPU. Falling back to FP32.")
             dtype = torch.float32
 
         print(f"IBB_POSE: Loading model on device: {device} with UNet precision: {unet_precision}")
-        
-        # --- Load Empty Embedding ---
+
         embed_path = os.path.join(EMPTY_EMBED_DIR, "empty_embedding.safetensors")
         if not os.path.exists(embed_path):
-            raise FileNotFoundError(f"Empty embedding not found at '{embed_path}'. Please run 'generate_empty_embedding.py' script first.")
+            raise FileNotFoundError(
+                f"Empty embedding not found at '{embed_path}'. Please run 'generate_empty_embedding.py' first."
+            )
         empty_text_embed = load_file(embed_path)["empty_text_embed"].to(device)
 
-        # --- Load Models ---
         unet = UNet2DConditionModel.from_pretrained(model_path, subfolder="unet", torch_dtype=dtype).to(device)
         unet = Modified_forward(unet, keypoint_scheme=keypoint_scheme)
-        vae = AutoencoderKL.from_pretrained(model_path, subfolder='vae').to(device) # VAE is small, keep fp32
-        
+        vae = AutoencoderKL.from_pretrained(model_path, subfolder="vae").to(device)
+
         dec_path = os.path.join(model_path, "decoder", "decoder.safetensors")
         hm_decoder = get_heatmap_head(mode=keypoint_scheme).to(device)
         if os.path.exists(dec_path):
             hm_decoder.load_state_dict(load_file(dec_path, device=str(device)), strict=True)
         hm_decoder = hm_decoder.to(dtype)
-        
-        noise_scheduler = DDPMScheduler.from_pretrained(model_path, subfolder='scheduler')
+
+        noise_scheduler = DDPMScheduler.from_pretrained(model_path, subfolder="scheduler")
 
         sdpose_model = {
             "unet": unet,
@@ -892,13 +903,12 @@ class IBBPoseModelLoader:
             "scheduler": noise_scheduler,
             "keypoint_scheme": keypoint_scheme,
             "device": device,
-            "unload_on_finish": unload_on_finish
+            "unload_on_finish": unload_on_finish,
         }
         return (sdpose_model,)
 
 
 # --- GroundingDINO Model Loader (Adapted from SAM2 Node) ---
-# (global groundingdino_model_list definition)
 groundingdino_model_dir_name = "grounding-dino"
 groundingdino_model_list = {
     "GroundingDINO_SwinT_OGC (694MB)": {
@@ -911,81 +921,41 @@ groundingdino_model_list = {
     },
 }
 
-        if skeleton_type in ("Body", "OpenPose"):
-            if not ULTRALYTICS_AVAILABLE:
-                raise ImportError(
-                    "IBB_POSE: ultralytics required for Body/OpenPose.\n"
-                    "  Install: pip install ultralytics"
-                )
-            model_path = _ensure_yolo_pose_model(model_size, auto_download)
-            print(f"IBB_POSE: Loading YOLO pose model -> {model_path}")
-            yolo = _YOLO(model_path)
-            model_info["backend"]   = "yolo_pose"
-            model_info["yolo_pose"] = yolo
-            try:
-                det_path = _ensure_yolo_det_model(auto_download)
-                model_info["yolo_det"] = _YOLO(det_path)
-            except Exception as exc:
-                print(f"IBB_POSE: YOLO detector load failed ({exc}). Multi-person via pose model.")
-                model_info["yolo_det"] = None
 
-        else:  # WholeBody
-            if not ONNXRUNTIME_AVAILABLE:
-                raise ImportError(
-                    "IBB_POSE: onnxruntime required for WholeBody.\n"
-                    "  Install: pip install onnxruntime  (or onnxruntime-gpu)"
-                )
-            det_path, pose_path = _ensure_dwpose_models(auto_download)
-
-            providers = ["CPUExecutionProvider"]
-            if dev.type != "cpu":
-                cuda_prov = ["CUDAExecutionProvider", "CPUExecutionProvider"]
-                try:
-                    available = set(ort.get_available_providers())
-                    if "CUDAExecutionProvider" in available:
-                        providers = cuda_prov
-                    else:
-                        raise RuntimeError("CUDAExecutionProvider unavailable")
-                except Exception:
-                    print("IBB_POSE: ONNX CUDAExecutionProvider unavailable – using CPU.")
-
-            print(f"IBB_POSE: Loading DWPose ONNX | providers: {providers}")
-            det_session  = ort.InferenceSession(det_path,  providers=providers)
-            pose_session = ort.InferenceSession(pose_path, providers=providers)
-            print("IBB_POSE: WholeBody detector backend -> YOLOX ONNX.")
-
-            model_info["backend"]      = "dwpose_onnx"
-            model_info["det_session"]  = det_session
-            model_info["pose_session"] = pose_session
-
-            if ULTRALYTICS_AVAILABLE:
-                try:
-                    det_yolo_path = _ensure_yolo_det_model(auto_download)
-                    model_info["yolo_det"] = _YOLO(det_yolo_path)
-                except Exception as exc:
-                    print(f"IBB_POSE: YOLO det unavailable ({exc}). Falling back to YOLOX.")
-                    model_info["yolo_det"] = None
-            else:
-                model_info["yolo_det"] = None
-
-        print(f"IBB_POSE: Model ready. skeleton={skeleton_type} device={dev} dtype={dtype}")
-        return (model_info,)
+def list_groundingdino_model():
+    return list(groundingdino_model_list.keys())
 
 
-    if dino_model_args.text_encoder_type == "bert-base-uncased":
-        dino_model_args.text_encoder_type = get_bert_base_uncased_model_path()
+def get_local_filepath(url, folder_name):
+    folder = os.path.join(folder_paths.models_dir, folder_name)
+    os.makedirs(folder, exist_ok=True)
+    filename = os.path.basename(url.split("?", 1)[0])
+    local_path = os.path.join(folder, filename)
+    if not os.path.exists(local_path):
+        _download_file(url, local_path)
+    return local_path
 
-    dino = local_groundingdino_build_model(dino_model_args)
-    checkpoint = torch.load(
-        get_local_filepath(
-            groundingdino_model_list[model_name]["model_url"],
-            groundingdino_model_dir_name,
-        ), map_location="cpu" # Load to CPU first
-    )
-    dino.load_state_dict(
-        local_groundingdino_clean_state_dict(checkpoint["model"]), strict=False
-    )
-    # Don't move to device here, let the Processor node handle it if needed
+
+def load_groundingdino_model(model_name):
+    try:
+        from groundingdino.models import build_model
+        from groundingdino.util.slconfig import SLConfig
+        from groundingdino.util.utils import clean_state_dict
+    except ImportError as exc:
+        raise ImportError(
+            "IBB_POSE: Failed to import 'groundingdino'. Please install it via pip: 'pip install groundingdino-py'"
+        ) from exc
+
+    if model_name not in groundingdino_model_list:
+        raise ValueError(f"IBB_POSE: Unknown GroundingDINO model: {model_name}")
+
+    config_path = get_local_filepath(groundingdino_model_list[model_name]["config_url"], groundingdino_model_dir_name)
+    model_path = get_local_filepath(groundingdino_model_list[model_name]["model_url"], groundingdino_model_dir_name)
+
+    dino_model_args = SLConfig.fromfile(config_path)
+    dino = build_model(dino_model_args)
+    checkpoint = torch.load(model_path, map_location="cpu")
+    dino.load_state_dict(clean_state_dict(checkpoint["model"]), strict=False)
     dino.eval()
     return dino
 
@@ -1190,7 +1160,7 @@ class IBBPoseProcessor:
                             detection_source = "GroundingDINO"
                     except Exception: pass
 
-            if not bboxes and yolo_model is not None and YOLO_AVAILABLE:
+            if not bboxes and yolo_model is not None and ULTRALYTICS_AVAILABLE:
                 try:
                     results = yolo_model(original_image_bgr, verbose=False)
                     yolo_bboxes = []