nodes/hobinrude-ComfyUI_DWposeDeluxe
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

Fixed TensorRT auto-install on first run

Browse Source
This commit is contained in:
paw37ski
2025-12-06 15:12:13 +01:00
parent 0158d1daef
commit f9b6fc239c
7 changed files with 91 additions and 85 deletions
Download Patch File Download Diff File Expand all files Collapse all files
dwpose/detection.py
+5 -4
View File
@@ -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:]
js/browse_files.js
+12 -7
View File
@@ -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(() => {
js/dwpose_dlx.js
+17 -20
View File
@@ -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"];
main.py
+32 -51
View File
@@ -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):
nodes/custom_options.py
+3
View File
@@ -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}),
}
}
nodes/keypoint_printer.py
+5 -1
View File
@@ -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)
nodes/pose_interpolation.py
+17 -2
View File
@@ -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}