From f9b6fc239cb447bf8772411fad30bd2db2da4cd7 Mon Sep 17 00:00:00 2001 From: paw37ski Date: Sat, 6 Dec 2025 15:12:13 +0100 Subject: [PATCH] Fixed TensorRT auto-install on first run --- dwpose/detection.py | 9 ++-- js/browse_files.js | 19 +++++---- js/dwpose_dlx.js | 37 ++++++++--------- main.py | 83 ++++++++++++++----------------------- nodes/custom_options.py | 3 ++ nodes/keypoint_printer.py | 6 ++- nodes/pose_interpolation.py | 19 ++++++++- 7 files changed, 91 insertions(+), 85 deletions(-) diff --git a/dwpose/detection.py b/dwpose/detection.py index 814d970..9d11d8f 100644 --- a/dwpose/detection.py +++ b/dwpose/detection.py @@ -105,12 +105,13 @@ def inference_detector(session, oriImg, model_type="ONNX", nms_threshold=0.45, s cudaStream = torch.cuda.current_stream().cuda_stream img_bchw = torch.from_numpy(img[None, :, :, :]) trt_output = session.infer({"images": img_bchw}, cudaStream) - output = trt_output['output'].cpu().numpy() - else: # ONNX + output_array = trt_output['output'].cpu().numpy() + else: # ONNX ort_inputs = {session.get_inputs()[0].name: img[None, :, :, :]} - output = session.run(None, ort_inputs) + onnx_output = session.run(None, ort_inputs) + output_array = onnx_output[0] - predictions = demo_postprocess(output[0], input_shape)[0] + predictions = demo_postprocess(output_array, input_shape)[0] boxes = predictions[:, :4] scores = predictions[:, 4:5] * predictions[:, 5:] diff --git a/js/browse_files.js b/js/browse_files.js index b3771e0..774cb6e 100644 --- a/js/browse_files.js +++ b/js/browse_files.js @@ -40,10 +40,11 @@ app.registerExtension({ app.graph.setDirtyCanvas(true, true); }; + // Helper function to upload files async function uploadFile(file, progressCallback) { try { const body = new FormData(); - body.append("image", file); + body.append("image", file); // ComfyUI's API uses 'image' for file uploads, even for JSON const url = api.apiURL("/upload/image"); const resp = await new Promise((resolve, reject) => { @@ -67,9 +68,10 @@ app.registerExtension({ // Add the "Upload keypoint dataset" button const uploadButton = node.addWidget("button", "Upload keypoint dataset", "upload", async () => { + // Create a hidden file input element const fileInput = document.createElement("input"); fileInput.type = "file"; - fileInput.accept = ".json"; + fileInput.accept = ".json"; // Accept only JSON files fileInput.style.display = "none"; document.body.appendChild(fileInput); @@ -77,24 +79,27 @@ app.registerExtension({ if (fileInput.files.length > 0) { try { const uploadedFile = fileInput.files[0]; - node.progress = 0; + // Show progress (optional) + node.progress = 0; // Initialize progress const resp = await uploadFile(uploadedFile, (p) => { node.progress = p; }); - node.progress = undefined; + node.progress = undefined; // Clear progress - const filename = JSON.parse(resp.responseText).name; + const filename = JSON.parse(resp.responseText).name; // Get the filename from the API response + // Refresh the file list and select the new file await updateFileList(filename); } catch (error) { console.error("Error during upload:", error); } } - document.body.removeChild(fileInput); + document.body.removeChild(fileInput); // Clean up the hidden input }; + // Programmatically click the hidden file input fileInput.click(); }); - uploadButton.options.serialize = false; + uploadButton.options.serialize = false; // Don't save this button in workflow // Initial file list load setTimeout(() => { diff --git a/js/dwpose_dlx.js b/js/dwpose_dlx.js index 985c2aa..8e900e2 100644 --- a/js/dwpose_dlx.js +++ b/js/dwpose_dlx.js @@ -63,7 +63,7 @@ app.registerExtension({ const providerType = providerWidget.value; const precision = precisionWidget.value; - let url = `/dwpose-dlx/get-models-list?provider_type=${providerType}`; + let url = `/dwpose_adv/get_model_list?provider_type=${providerType}`; if (providerType === "GPU" && precision) { url += `&precision=${precision}`; } @@ -72,28 +72,25 @@ app.registerExtension({ const response = await fetch(url); const data = await response.json(); - const handleWidgetUpdate = (widget, models) => { - if (!widget) return; - - if (models.length === 1 && models[0] === "No compatible GPU detected") { - widget.options.values = models; - widget.value = models[0]; - widget.disabled = true; + if (detectorModelWidget) { + const currentDetectorValue = detectorModelWidget.value; + detectorModelWidget.options.values = data.detector_models; + if (!data.detector_models.includes(currentDetectorValue)) { + detectorModelWidget.value = detectorModelWidget.options.values[0]; } else { - widget.disabled = false; - const currentValue = widget.value; - widget.options.values = models; - if (!models.includes(currentValue)) { - widget.value = widget.options.values[0]; - } else { - widget.value = currentValue; - } + detectorModelWidget.value = currentDetectorValue; } - }; - - handleWidgetUpdate(detectorModelWidget, data.detector_models); - handleWidgetUpdate(estimatorModelWidget, data.estimator_models); + } + if (estimatorModelWidget) { + const currentEstimatorValue = estimatorModelWidget.value; + estimatorModelWidget.options.values = data.estimator_models; + if (!data.estimator_models.includes(currentEstimatorValue)) { + estimatorModelWidget.value = estimatorModelWidget.options.values[0]; + } else { + estimatorModelWidget.value = currentEstimatorValue; + } + } } catch (error) { console.error("Error fetching model lists:", error); if (detectorModelWidget) detectorModelWidget.options.values = ["Error"]; diff --git a/main.py b/main.py index aa56452..1545aab 100644 --- a/main.py +++ b/main.py @@ -5,6 +5,7 @@ import cv2 import sys import time import json +import torch import numpy as np import folder_paths @@ -15,11 +16,6 @@ from . import HAS_TENSORRT from .scripts import logger from .nodes.custom_options import DWOPOSE_CUSTOM_OPTIONS_TYPE -try: - import torch - IS_GPU_AVAILABLE = torch.cuda.is_available() and torch.cuda.device_count() > 0 -except: - IS_GPU_AVAILABLE = False if HAS_TENSORRT: try: @@ -122,57 +118,42 @@ class DWposeDeluxeNode: } return inputs - @classmethod + @classmethod + def get_model_list(cls, provider_type, precision=None): + models_dir = os.path.join(folder_paths.models_dir, "dwpose") - def get_model_list(cls, provider_type, precision=None): + onnx_detector_models = [f for f in os.listdir(models_dir) if f.startswith("yolox_") and f.endswith(".onnx")] + onnx_estimator_models = [f for f in os.listdir(models_dir) if f.startswith("dw-ll_ucoco_") and f.endswith(".onnx")] - if provider_type == "GPU" and not IS_GPU_AVAILABLE: - return (["No compatible GPU detected"], ["No compatible GPU detected"]) + detector_models = [] + estimator_models = [] - models_dir = os.path.join(folder_paths.models_dir, "dwpose") - onnx_detector_models_base = [f for f in os.listdir(models_dir) if f.startswith("yolox_") and f.endswith(".onnx")] - onnx_estimator_models_base = [f for f in os.listdir(models_dir) if f.startswith("dw-ll_ucoco_") and f.endswith(".onnx")] - detector_models = [] - estimator_models = [] - - if provider_type == "CPU": - detector_models.extend(onnx_detector_models_base if onnx_detector_models_base else ["auto-download"]) - estimator_models.extend(onnx_estimator_models_base if onnx_estimator_models_base else ["auto-download"]) - elif provider_type == "GPU": - if precision == "fp32": - detector_models.extend(onnx_detector_models_base) - estimator_models.extend(onnx_estimator_models_base) - if not HAS_TENSORRT: - detector_models.append("TensorRT_Not_Installed") - estimator_models.append("TensorRT_Not_Installed") - elif precision: - if HAS_TENSORRT: - all_trt_files = [f for f in os.listdir(models_dir) if f.endswith(".trt")] - detector_prefix = f"yolox_l_{precision}" - matching_detector_models = [f for f in all_trt_files if f.startswith(detector_prefix) and f.endswith(".trt")] - detector_models.extend(sorted(matching_detector_models) if matching_detector_models else ["auto-build"]) - estimator_prefix = f"dw-ll_ucoco_384_{precision}" - matching_estimator_models = [f for f in all_trt_files if f.startswith(estimator_prefix) and f.endswith(".trt")] - estimator_models.extend(sorted(matching_estimator_models) if matching_estimator_models else ["auto-build"]) - else: - detector_models.append("TensorRT_Not_Installed") - estimator_models.append("TensorRT_Not_Installed") - else: - detector_models.append("select_precision") - estimator_models.append("select_precision") - if not detector_models: - detector_models.append("auto-download") + if provider_type == "CPU": + detector_models = onnx_detector_models if onnx_detector_models else ["auto-download"] + estimator_models = onnx_estimator_models if onnx_estimator_models else ["auto-download"] + elif provider_type == "GPU": + if not HAS_TENSORRT: + detector_models = ["TensorRT_Not_Installed"] + estimator_models = ["TensorRT_Not_Installed"] + elif precision: if HAS_TENSORRT: - detector_models.append("auto-build") + all_trt_files = [f for f in os.listdir(models_dir) if f.endswith(".trt")] + + detector_prefix = f"yolox_l_{precision}" + matching_detector_models = [f for f in all_trt_files if f.startswith(detector_prefix) and f.endswith(".trt")] + detector_models = sorted(matching_detector_models) if matching_detector_models else ["auto-build"] + + estimator_prefix = f"dw-ll_ucoco_384_{precision}" + matching_estimator_models = [f for f in all_trt_files if f.startswith(estimator_prefix) and f.endswith(".trt")] + estimator_models = sorted(matching_estimator_models) if matching_estimator_models else ["auto-build"] else: - detector_models.append("TensorRT_Not_Installed") - if not estimator_models: - estimator_models.append("auto-download") - if HAS_TENSORRT: - estimator_models.append("auto-build") - else: - estimator_models.append("TensorRT_Not_Installed") - return (sorted(list(set(detector_models))), sorted(list(set(estimator_models)))) + detector_models = ["TensorRT_Not_Installed"] + estimator_models = ["TensorRT_Not_Installed"] + else: + detector_models = ["select_precision"] + estimator_models = ["select_precision"] + + return (detector_models, estimator_models) @staticmethod def _get_model_list_api(provider_type, precision=None): diff --git a/nodes/custom_options.py b/nodes/custom_options.py index bc720a5..8f0b988 100644 --- a/nodes/custom_options.py +++ b/nodes/custom_options.py @@ -24,6 +24,9 @@ class CustomOptions: "face_threshold": ("FLOAT", {"default": 0.10, "min": 0.0, "max": 1.0, "step": 0.01}), "hand_threshold": ("FLOAT", {"default": 0.10, "min": 0.0, "max": 1.0, "step": 0.01}), "face_padding": ("FLOAT", {"default": 0.0, "min": 0.0, "max": 1.0, "step": 0.01}), + "neck_validity": ("FLOAT", {"default": 0.3, "min": 0.0, "max": 1.0, "step": 0.01}), + "nms_threshold": ("FLOAT", {"default": 0.45, "min": 0.0, "max": 1.0, "step": 0.01}), + "score_threshold": ("FLOAT", {"default": 0.1, "min": 0.0, "max": 1.0, "step": 0.01}), } } diff --git a/nodes/keypoint_printer.py b/nodes/keypoint_printer.py index 513548a..c5a18f9 100644 --- a/nodes/keypoint_printer.py +++ b/nodes/keypoint_printer.py @@ -35,10 +35,12 @@ class KeypointPrinter: logger.warning("Keypoints data is empty or invalid. Returning empty image.") return (torch.zeros((1, 64, 64, 3), dtype=torch.float32),) + # Coordinate Format Check for frame in keypoints: for person in frame.get("people", []): for key in ["pose_keypoints_2d", "face_keypoints_2d", "hand_left_keypoints_2d", "hand_right_keypoints_2d"]: kpts = person.get(key, []) + # Check only a few points for efficiency for i in range(0, min(len(kpts), 9), 3): if kpts[i] > 1.0 or kpts[i+1] > 1.0: logger.error("Keypoints are in absolute (pixel) format. Please use a Keypoint Converter node to normalize them before using the printer.") @@ -78,6 +80,7 @@ class KeypointPrinter: if poses_to_print != -1: people_to_process = people_to_process[:poses_to_print] + # Data Transformation current_candidate = [] current_subset = [] all_hand_peaks = [] @@ -87,7 +90,7 @@ class KeypointPrinter: for person in people_to_process: if show_body: pose_kpts_2d = np.array(person.get('pose_keypoints_2d', [])).reshape(-1, 3) - person_subset_row = [-1] * 25 + person_subset_row = [-1] * 25 # DWpose uses up to 25 keypoints for body/feet for kpt_idx, (x, y, conf) in enumerate(pose_kpts_2d): if kpt_idx < 25 and conf > 0: @@ -113,6 +116,7 @@ class KeypointPrinter: if normalized_face_kpts: all_lmks.append(normalized_face_kpts) + # Rendering if show_body and current_candidate and current_subset: canvas_np = dwpose_util.draw_bodypose(canvas_np, np.array(current_candidate), np.array(current_subset), show_feet, actual_options) diff --git a/nodes/pose_interpolation.py b/nodes/pose_interpolation.py index 6c21354..a5f5fd9 100644 --- a/nodes/pose_interpolation.py +++ b/nodes/pose_interpolation.py @@ -41,9 +41,10 @@ class PoseInterpolation: x_coords = pose_kpts_2d[:, 0] y_coords = pose_kpts_2d[:, 1] + # Filter out keypoints with 0 confidence before calculating bounds valid_kpts_indices = np.where(pose_kpts_2d[:, 2] > 0)[0] if len(valid_kpts_indices) == 0: - continue + continue # No valid keypoints for this person valid_x = x_coords[valid_kpts_indices] valid_y = y_coords[valid_kpts_indices] @@ -51,6 +52,8 @@ class PoseInterpolation: min_x, max_x = np.min(valid_x), np.max(valid_x) min_y, max_y = np.min(valid_y), np.max(valid_y) + # Bounding box coordinates are normalized, convert to pixels for area calculation + # The canvas width/height are available here bbox_width = (max_x - min_x) * canvas_width bbox_height = (max_y - min_y) * canvas_height @@ -65,7 +68,7 @@ class PoseInterpolation: return None, None, None return biggest_pose_data, canvas_width, canvas_height - return people[0], canvas_width, canvas_height + return people[0], canvas_width, canvas_height # Only one person, return it directly RETURN_TYPES = ("POSE_KEYPOINT",) RETURN_NAMES = ("keypoints",) @@ -96,6 +99,7 @@ class PoseInterpolation: logger.error("Could not extract valid start or end pose for interpolation.") raise InterruptProcessingException("Invalid start or end pose data for interpolation.") + # Part D (Output Canvas Size Handling): Check canvas dimensions and set target_width/height target_width = start_width target_height = start_height @@ -104,6 +108,8 @@ class PoseInterpolation: target_width = max(start_width, end_width) target_height = max(start_height, end_height) + # Part B: Check Compatibility + # Define all possible keypoint subsets keypoint_subsets = [ "pose_keypoints_2d", "face_keypoints_2d", @@ -117,6 +123,9 @@ class PoseInterpolation: end_kpts = end_pose_data.get(subset_name) if start_kpts is not None and end_kpts is not None: + # Ensure they have the same number of keypoints for interpolation + # Keypoints are stored as [x, y, conf, x, y, conf, ...] + # So the length of the list should be a multiple of 3 if len(start_kpts) == len(end_kpts) and len(start_kpts) % 3 == 0: common_subsets[subset_name] = (np.array(start_kpts).reshape(-1, 3), np.array(end_kpts).reshape(-1, 3)) else: @@ -128,15 +137,19 @@ class PoseInterpolation: logger.error("No common keypoint subsets found for interpolation.") raise InterruptProcessingException("Cannot interpolate without common keypoints.") + # Part C: Perform Linear Interpolation interpolated_frames = [] for i in range(frame_count): t = 0.0 if frame_count == 1 else i / (frame_count - 1) interpolated_person_data = {} + # Copy other non-keypoint data from start_pose_data if necessary + # For example, face_box if it exists and is not interpolated if "face_box" in start_pose_data and "face_box" in end_pose_data: start_box = np.array(start_pose_data["face_box"]) end_box = np.array(end_pose_data["face_box"]) interpolated_person_data["face_box"] = (start_box * (1 - t) + end_box * t).tolist() + # ... potentially other data to copy or interpolate ... for subset_name, (start_kpts_np, end_kpts_np) in common_subsets.items(): interpolated_kpts_np = (start_kpts_np * (1 - t) + end_kpts_np * t).reshape(-1).tolist() @@ -149,8 +162,10 @@ class PoseInterpolation: } interpolated_frames.append(interpolated_frame) + # Placeholder for Part D logic (return the interpolated_frames) logger.info(f"Linear interpolation completed. Generated {len(interpolated_frames)} frames.") + # Temporarily returning the interpolated frames return (interpolated_frames,) NODE_CLASS_MAPPINGS = {"PoseInterpolation": PoseInterpolation}