Skip to content

add engine builder node #6

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
wants to merge 1 commit into
base: master
Choose a base branch
from

Conversation

ryanontheinside
Copy link

Add EngineBuilder Node for YoloNasPose-Tensorrt

Description

This PR adds a new EngineBuilder node to the ComfyUI-YoloNasPose-Tensorrt extension, allowing users to easily build TensorRT engines from ONNX models directly in ComfyUI.

Changes

  • Created new engine builder node

Features

  • Model Selection: Choose from the supported ONNX models
  • Custom Engine Names: Option to specify custom engine file names
  • Custom ONNX Models: Support for using custom ONNX models
  • Automatic Downloads: Auto download onnx from HF, and build engine
  • User Feedback: Clear status messages for all operations

Usage

  1. Add the "YoloNasPose Engine Builder" node to your workflow
  2. Select the desired model size
  3. (Optional) Configure custom engine name and FP16 settings
  4. Run the workflow to build the TensorRT engine

Technical Details

  • Models are downloaded to models/onnx/yolo-nas-pose/
  • Engines are built in models/tensorrt/yolo-nas-pose/
  • Supports both predefined and custom ONNX models
  • Includes proper error handling and status reporting

Testing

  • Tested with some models sizes in ComfyUI (method 1)
  • Tested with some model sizes by running the script directly (method 2)

for convenience
@ryanontheinside ryanontheinside marked this pull request as ready for review May 21, 2025 17:21
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Labels
None yet
Projects
None yet
Development

Successfully merging this pull request may close these issues.

1 participant