Mirror from https://github.com/yuvraj108c/ComfyUI-Upscaler-Tensorrt

.gitattributes

@@ -33,3 +33,4 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
 *.zip filter=lfs diff=lfs merge=lfs -text
 *.zst filter=lfs diff=lfs merge=lfs -text
 *tfevents* filter=lfs diff=lfs merge=lfs -text
+assets/node_v3.png filter=lfs diff=lfs merge=lfs -text

.github/FUNDING.yml

@@ -0,0 +1,2 @@
+github: yuvraj108c
+custom: ["https://paypal.me/yuvraj108c", "https://buymeacoffee.com/yuvraj108cz"]

.gitignore

@@ -0,0 +1 @@
+__pycache__

LICENSE

@@ -0,0 +1,437 @@
+Attribution-NonCommercial-ShareAlike 4.0 International
+
+=======================================================================
+
+Creative Commons Corporation ("Creative Commons") is not a law firm and
+does not provide legal services or legal advice. Distribution of
+Creative Commons public licenses does not create a lawyer-client or
+other relationship. Creative Commons makes its licenses and related
+information available on an "as-is" basis. Creative Commons gives no
+warranties regarding its licenses, any material licensed under their
+terms and conditions, or any related information. Creative Commons
+disclaims all liability for damages resulting from their use to the
+fullest extent possible.
+
+Using Creative Commons Public Licenses
+
+Creative Commons public licenses provide a standard set of terms and
+conditions that creators and other rights holders may use to share
+original works of authorship and other material subject to copyright
+and certain other rights specified in the public license below. The
+following considerations are for informational purposes only, are not
+exhaustive, and do not form part of our licenses.
+
+     Considerations for licensors: Our public licenses are
+     intended for use by those authorized to give the public
+     permission to use material in ways otherwise restricted by
+     copyright and certain other rights. Our licenses are
+     irrevocable. Licensors should read and understand the terms
+     and conditions of the license they choose before applying it.
+     Licensors should also secure all rights necessary before
+     applying our licenses so that the public can reuse the
+     material as expected. Licensors should clearly mark any
+     material not subject to the license. This includes other CC-
+     licensed material, or material used under an exception or
+     limitation to copyright. More considerations for licensors:
+    wiki.creativecommons.org/Considerations_for_licensors
+
+     Considerations for the public: By using one of our public
+     licenses, a licensor grants the public permission to use the
+     licensed material under specified terms and conditions. If
+     the licensor's permission is not necessary for any reason--for
+     example, because of any applicable exception or limitation to
+     copyright--then that use is not regulated by the license. Our
+     licenses grant only permissions under copyright and certain
+     other rights that a licensor has authority to grant. Use of
+     the licensed material may still be restricted for other
+     reasons, including because others have copyright or other
+     rights in the material. A licensor may make special requests,
+     such as asking that all changes be marked or described.
+     Although not required by our licenses, you are encouraged to
+     respect those requests where reasonable. More considerations
+     for the public:
+    wiki.creativecommons.org/Considerations_for_licensees
+
+=======================================================================
+
+Creative Commons Attribution-NonCommercial-ShareAlike 4.0 International
+Public License
+
+By exercising the Licensed Rights (defined below), You accept and agree
+to be bound by the terms and conditions of this Creative Commons
+Attribution-NonCommercial-ShareAlike 4.0 International Public License
+("Public License"). To the extent this Public License may be
+interpreted as a contract, You are granted the Licensed Rights in
+consideration of Your acceptance of these terms and conditions, and the
+Licensor grants You such rights in consideration of benefits the
+Licensor receives from making the Licensed Material available under
+these terms and conditions.
+
+
+Section 1 -- Definitions.
+
+  a. Adapted Material means material subject to Copyright and Similar
+     Rights that is derived from or based upon the Licensed Material
+     and in which the Licensed Material is translated, altered,
+     arranged, transformed, or otherwise modified in a manner requiring
+     permission under the Copyright and Similar Rights held by the
+     Licensor. For purposes of this Public License, where the Licensed
+     Material is a musical work, performance, or sound recording,
+     Adapted Material is always produced where the Licensed Material is
+     synched in timed relation with a moving image.
+
+  b. Adapter's License means the license You apply to Your Copyright
+     and Similar Rights in Your contributions to Adapted Material in
+     accordance with the terms and conditions of this Public License.
+
+  c. BY-NC-SA Compatible License means a license listed at
+     creativecommons.org/compatiblelicenses, approved by Creative
+     Commons as essentially the equivalent of this Public License.
+
+  d. Copyright and Similar Rights means copyright and/or similar rights
+     closely related to copyright including, without limitation,
+     performance, broadcast, sound recording, and Sui Generis Database
+     Rights, without regard to how the rights are labeled or
+     categorized. For purposes of this Public License, the rights
+     specified in Section 2(b)(1)-(2) are not Copyright and Similar
+     Rights.
+
+  e. Effective Technological Measures means those measures that, in the
+     absence of proper authority, may not be circumvented under laws
+     fulfilling obligations under Article 11 of the WIPO Copyright
+     Treaty adopted on December 20, 1996, and/or similar international
+     agreements.
+
+  f. Exceptions and Limitations means fair use, fair dealing, and/or
+     any other exception or limitation to Copyright and Similar Rights
+     that applies to Your use of the Licensed Material.
+
+  g. License Elements means the license attributes listed in the name
+     of a Creative Commons Public License. The License Elements of this
+     Public License are Attribution, NonCommercial, and ShareAlike.
+
+  h. Licensed Material means the artistic or literary work, database,
+     or other material to which the Licensor applied this Public
+     License.
+
+  i. Licensed Rights means the rights granted to You subject to the
+     terms and conditions of this Public License, which are limited to
+     all Copyright and Similar Rights that apply to Your use of the
+     Licensed Material and that the Licensor has authority to license.
+
+  j. Licensor means the individual(s) or entity(ies) granting rights
+     under this Public License.
+
+  k. NonCommercial means not primarily intended for or directed towards
+     commercial advantage or monetary compensation. For purposes of
+     this Public License, the exchange of the Licensed Material for
+     other material subject to Copyright and Similar Rights by digital
+     file-sharing or similar means is NonCommercial provided there is
+     no payment of monetary compensation in connection with the
+     exchange.
+
+  l. Share means to provide material to the public by any means or
+     process that requires permission under the Licensed Rights, such
+     as reproduction, public display, public performance, distribution,
+     dissemination, communication, or importation, and to make material
+     available to the public including in ways that members of the
+     public may access the material from a place and at a time
+     individually chosen by them.
+
+  m. Sui Generis Database Rights means rights other than copyright
+     resulting from Directive 96/9/EC of the European Parliament and of
+     the Council of 11 March 1996 on the legal protection of databases,
+     as amended and/or succeeded, as well as other essentially
+     equivalent rights anywhere in the world.
+
+  n. You means the individual or entity exercising the Licensed Rights
+     under this Public License. Your has a corresponding meaning.
+
+
+Section 2 -- Scope.
+
+  a. License grant.
+
+       1. Subject to the terms and conditions of this Public License,
+          the Licensor hereby grants You a worldwide, royalty-free,
+          non-sublicensable, non-exclusive, irrevocable license to
+          exercise the Licensed Rights in the Licensed Material to:
+
+            a. reproduce and Share the Licensed Material, in whole or
+               in part, for NonCommercial purposes only; and
+
+            b. produce, reproduce, and Share Adapted Material for
+               NonCommercial purposes only.
+
+       2. Exceptions and Limitations. For the avoidance of doubt, where
+          Exceptions and Limitations apply to Your use, this Public
+          License does not apply, and You do not need to comply with
+          its terms and conditions.
+
+       3. Term. The term of this Public License is specified in Section
+          6(a).
+
+       4. Media and formats; technical modifications allowed. The
+          Licensor authorizes You to exercise the Licensed Rights in
+          all media and formats whether now known or hereafter created,
+          and to make technical modifications necessary to do so. The
+          Licensor waives and/or agrees not to assert any right or
+          authority to forbid You from making technical modifications
+          necessary to exercise the Licensed Rights, including
+          technical modifications necessary to circumvent Effective
+          Technological Measures. For purposes of this Public License,
+          simply making modifications authorized by this Section 2(a)
+          (4) never produces Adapted Material.
+
+       5. Downstream recipients.
+
+            a. Offer from the Licensor -- Licensed Material. Every
+               recipient of the Licensed Material automatically
+               receives an offer from the Licensor to exercise the
+               Licensed Rights under the terms and conditions of this
+               Public License.
+
+            b. Additional offer from the Licensor -- Adapted Material.
+               Every recipient of Adapted Material from You
+               automatically receives an offer from the Licensor to
+               exercise the Licensed Rights in the Adapted Material
+               under the conditions of the Adapter's License You apply.
+
+            c. No downstream restrictions. You may not offer or impose
+               any additional or different terms or conditions on, or
+               apply any Effective Technological Measures to, the
+               Licensed Material if doing so restricts exercise of the
+               Licensed Rights by any recipient of the Licensed
+               Material.
+
+       6. No endorsement. Nothing in this Public License constitutes or
+          may be construed as permission to assert or imply that You
+          are, or that Your use of the Licensed Material is, connected
+          with, or sponsored, endorsed, or granted official status by,
+          the Licensor or others designated to receive attribution as
+          provided in Section 3(a)(1)(A)(i).
+
+  b. Other rights.
+
+       1. Moral rights, such as the right of integrity, are not
+          licensed under this Public License, nor are publicity,
+          privacy, and/or other similar personality rights; however, to
+          the extent possible, the Licensor waives and/or agrees not to
+          assert any such rights held by the Licensor to the limited
+          extent necessary to allow You to exercise the Licensed
+          Rights, but not otherwise.
+
+       2. Patent and trademark rights are not licensed under this
+          Public License.
+
+       3. To the extent possible, the Licensor waives any right to
+          collect royalties from You for the exercise of the Licensed
+          Rights, whether directly or through a collecting society
+          under any voluntary or waivable statutory or compulsory
+          licensing scheme. In all other cases the Licensor expressly
+          reserves any right to collect such royalties, including when
+          the Licensed Material is used other than for NonCommercial
+          purposes.
+
+
+Section 3 -- License Conditions.
+
+Your exercise of the Licensed Rights is expressly made subject to the
+following conditions.
+
+  a. Attribution.
+
+       1. If You Share the Licensed Material (including in modified
+          form), You must:
+
+            a. retain the following if it is supplied by the Licensor
+               with the Licensed Material:
+
+                 i. identification of the creator(s) of the Licensed
+                    Material and any others designated to receive
+                    attribution, in any reasonable manner requested by
+                    the Licensor (including by pseudonym if
+                    designated);
+
+                ii. a copyright notice;
+
+               iii. a notice that refers to this Public License;
+
+                iv. a notice that refers to the disclaimer of
+                    warranties;
+
+                 v. a URI or hyperlink to the Licensed Material to the
+                    extent reasonably practicable;
+
+            b. indicate if You modified the Licensed Material and
+               retain an indication of any previous modifications; and
+
+            c. indicate the Licensed Material is licensed under this
+               Public License, and include the text of, or the URI or
+               hyperlink to, this Public License.
+
+       2. You may satisfy the conditions in Section 3(a)(1) in any
+          reasonable manner based on the medium, means, and context in
+          which You Share the Licensed Material. For example, it may be
+          reasonable to satisfy the conditions by providing a URI or
+          hyperlink to a resource that includes the required
+          information.
+       3. If requested by the Licensor, You must remove any of the
+          information required by Section 3(a)(1)(A) to the extent
+          reasonably practicable.
+
+  b. ShareAlike.
+
+     In addition to the conditions in Section 3(a), if You Share
+     Adapted Material You produce, the following conditions also apply.
+
+       1. The Adapter's License You apply must be a Creative Commons
+          license with the same License Elements, this version or
+          later, or a BY-NC-SA Compatible License.
+
+       2. You must include the text of, or the URI or hyperlink to, the
+          Adapter's License You apply. You may satisfy this condition
+          in any reasonable manner based on the medium, means, and
+          context in which You Share Adapted Material.
+
+       3. You may not offer or impose any additional or different terms
+          or conditions on, or apply any Effective Technological
+          Measures to, Adapted Material that restrict exercise of the
+          rights granted under the Adapter's License You apply.
+
+
+Section 4 -- Sui Generis Database Rights.
+
+Where the Licensed Rights include Sui Generis Database Rights that
+apply to Your use of the Licensed Material:
+
+  a. for the avoidance of doubt, Section 2(a)(1) grants You the right
+     to extract, reuse, reproduce, and Share all or a substantial
+     portion of the contents of the database for NonCommercial purposes
+     only;
+
+  b. if You include all or a substantial portion of the database
+     contents in a database in which You have Sui Generis Database
+     Rights, then the database in which You have Sui Generis Database
+     Rights (but not its individual contents) is Adapted Material,
+     including for purposes of Section 3(b); and
+
+  c. You must comply with the conditions in Section 3(a) if You Share
+     all or a substantial portion of the contents of the database.
+
+For the avoidance of doubt, this Section 4 supplements and does not
+replace Your obligations under this Public License where the Licensed
+Rights include other Copyright and Similar Rights.
+
+
+Section 5 -- Disclaimer of Warranties and Limitation of Liability.
+
+  a. UNLESS OTHERWISE SEPARATELY UNDERTAKEN BY THE LICENSOR, TO THE
+     EXTENT POSSIBLE, THE LICENSOR OFFERS THE LICENSED MATERIAL AS-IS
+     AND AS-AVAILABLE, AND MAKES NO REPRESENTATIONS OR WARRANTIES OF
+     ANY KIND CONCERNING THE LICENSED MATERIAL, WHETHER EXPRESS,
+     IMPLIED, STATUTORY, OR OTHER. THIS INCLUDES, WITHOUT LIMITATION,
+     WARRANTIES OF TITLE, MERCHANTABILITY, FITNESS FOR A PARTICULAR
+     PURPOSE, NON-INFRINGEMENT, ABSENCE OF LATENT OR OTHER DEFECTS,
+     ACCURACY, OR THE PRESENCE OR ABSENCE OF ERRORS, WHETHER OR NOT
+     KNOWN OR DISCOVERABLE. WHERE DISCLAIMERS OF WARRANTIES ARE NOT
+     ALLOWED IN FULL OR IN PART, THIS DISCLAIMER MAY NOT APPLY TO YOU.
+
+  b. TO THE EXTENT POSSIBLE, IN NO EVENT WILL THE LICENSOR BE LIABLE
+     TO YOU ON ANY LEGAL THEORY (INCLUDING, WITHOUT LIMITATION,
+     NEGLIGENCE) OR OTHERWISE FOR ANY DIRECT, SPECIAL, INDIRECT,
+     INCIDENTAL, CONSEQUENTIAL, PUNITIVE, EXEMPLARY, OR OTHER LOSSES,
+     COSTS, EXPENSES, OR DAMAGES ARISING OUT OF THIS PUBLIC LICENSE OR
+     USE OF THE LICENSED MATERIAL, EVEN IF THE LICENSOR HAS BEEN
+     ADVISED OF THE POSSIBILITY OF SUCH LOSSES, COSTS, EXPENSES, OR
+     DAMAGES. WHERE A LIMITATION OF LIABILITY IS NOT ALLOWED IN FULL OR
+     IN PART, THIS LIMITATION MAY NOT APPLY TO YOU.
+
+  c. The disclaimer of warranties and limitation of liability provided
+     above shall be interpreted in a manner that, to the extent
+     possible, most closely approximates an absolute disclaimer and
+     waiver of all liability.
+
+
+Section 6 -- Term and Termination.
+
+  a. This Public License applies for the term of the Copyright and
+     Similar Rights licensed here. However, if You fail to comply with
+     this Public License, then Your rights under this Public License
+     terminate automatically.
+
+  b. Where Your right to use the Licensed Material has terminated under
+     Section 6(a), it reinstates:
+
+       1. automatically as of the date the violation is cured, provided
+          it is cured within 30 days of Your discovery of the
+          violation; or
+
+       2. upon express reinstatement by the Licensor.
+
+     For the avoidance of doubt, this Section 6(b) does not affect any
+     right the Licensor may have to seek remedies for Your violations
+     of this Public License.
+
+  c. For the avoidance of doubt, the Licensor may also offer the
+     Licensed Material under separate terms or conditions or stop
+     distributing the Licensed Material at any time; however, doing so
+     will not terminate this Public License.
+
+  d. Sections 1, 5, 6, 7, and 8 survive termination of this Public
+     License.
+
+
+Section 7 -- Other Terms and Conditions.
+
+  a. The Licensor shall not be bound by any additional or different
+     terms or conditions communicated by You unless expressly agreed.
+
+  b. Any arrangements, understandings, or agreements regarding the
+     Licensed Material not stated herein are separate from and
+     independent of the terms and conditions of this Public License.
+
+
+Section 8 -- Interpretation.
+
+  a. For the avoidance of doubt, this Public License does not, and
+     shall not be interpreted to, reduce, limit, restrict, or impose
+     conditions on any use of the Licensed Material that could lawfully
+     be made without permission under this Public License.
+
+  b. To the extent possible, if any provision of this Public License is
+     deemed unenforceable, it shall be automatically reformed to the
+     minimum extent necessary to make it enforceable. If the provision
+     cannot be reformed, it shall be severed from this Public License
+     without affecting the enforceability of the remaining terms and
+     conditions.
+
+  c. No term or condition of this Public License will be waived and no
+     failure to comply consented to unless expressly agreed to by the
+     Licensor.
+
+  d. Nothing in this Public License constitutes or may be interpreted
+     as a limitation upon, or waiver of, any privileges and immunities
+     that apply to the Licensor or You, including from the legal
+     processes of any jurisdiction or authority.
+
+=======================================================================
+
+Creative Commons is not a party to its public
+licenses. Notwithstanding, Creative Commons may elect to apply one of
+its public licenses to material it publishes and in those instances
+will be considered the “Licensor.” The text of the Creative Commons
+public licenses is dedicated to the public domain under the CC0 Public
+Domain Dedication. Except for the limited purpose of indicating that
+material is shared under a Creative Commons public license or as
+otherwise permitted by the Creative Commons policies published at
+creativecommons.org/policies, Creative Commons does not authorize the
+use of the trademark "Creative Commons" or any other trademark or logo
+of Creative Commons without its prior written consent including,
+without limitation, in connection with any unauthorized modifications
+to any of its public licenses or any other arrangements,
+understandings, or agreements concerning use of licensed material. For
+the avoidance of doubt, this paragraph does not form part of the
+public licenses.
+
+Creative Commons may be contacted at creativecommons.org.

README.md

@@ -0,0 +1,122 @@
+<div align="center">
+
+# ComfyUI Upscaler TensorRT ⚡
+
+[![python](https://img.shields.io/badge/python-3.10.12-green)](https://www.python.org/downloads/release/python-31012/)
+[![cuda](https://img.shields.io/badge/cuda-12.7-green)](https://developer.nvidia.com/cuda-downloads)
+[![trt](https://img.shields.io/badge/TRT-10.9-green)](https://developer.nvidia.com/tensorrt)
+[![by-nc-sa/4.0](https://img.shields.io/badge/license-CC--BY--NC--SA--4.0-lightgrey)](https://creativecommons.org/licenses/by-nc-sa/4.0/deed.en)
+
+</div>
+
+This project provides a [Tensorrt](https://github.com/NVIDIA/TensorRT) implementation for fast image upscaling using models inside ComfyUI (2-4x faster)
+
+<p align="center">
+  <img src="assets/node_v3.png" style="height: 400px" />
+</p>
+
+## ⭐ Support
+If you like my projects and wish to see updates and new features, please consider supporting me. It helps a lot! 
+
+[![ComfyUI-Depth-Anything-Tensorrt](https://img.shields.io/badge/ComfyUI--Depth--Anything--Tensorrt-blue?style=flat-square)](https://github.com/yuvraj108c/ComfyUI-Depth-Anything-Tensorrt)
+[![ComfyUI-Upscaler-Tensorrt](https://img.shields.io/badge/ComfyUI--Upscaler--Tensorrt-blue?style=flat-square)](https://github.com/yuvraj108c/ComfyUI-Upscaler-Tensorrt)
+[![ComfyUI-Dwpose-Tensorrt](https://img.shields.io/badge/ComfyUI--Dwpose--Tensorrt-blue?style=flat-square)](https://github.com/yuvraj108c/ComfyUI-Dwpose-Tensorrt)
+[![ComfyUI-Rife-Tensorrt](https://img.shields.io/badge/ComfyUI--Rife--Tensorrt-blue?style=flat-square)](https://github.com/yuvraj108c/ComfyUI-Rife-Tensorrt)
+
+[![ComfyUI-Whisper](https://img.shields.io/badge/ComfyUI--Whisper-gray?style=flat-square)](https://github.com/yuvraj108c/ComfyUI-Whisper)
+[![ComfyUI_InvSR](https://img.shields.io/badge/ComfyUI__InvSR-gray?style=flat-square)](https://github.com/yuvraj108c/ComfyUI_InvSR)
+[![ComfyUI-Thera](https://img.shields.io/badge/ComfyUI--Thera-gray?style=flat-square)](https://github.com/yuvraj108c/ComfyUI-Thera)
+[![ComfyUI-Video-Depth-Anything](https://img.shields.io/badge/ComfyUI--Video--Depth--Anything-gray?style=flat-square)](https://github.com/yuvraj108c/ComfyUI-Video-Depth-Anything)
+[![ComfyUI-PiperTTS](https://img.shields.io/badge/ComfyUI--PiperTTS-gray?style=flat-square)](https://github.com/yuvraj108c/ComfyUI-PiperTTS)
+
+[![buy-me-coffees](https://i.imgur.com/3MDbAtw.png)](https://www.buymeacoffee.com/yuvraj108cZ)
+[![paypal-donation](https://i.imgur.com/w5jjubk.png)](https://paypal.me/yuvraj108c)
+---
+
+## ⏱️ Performance
+
+_Note: The following results were benchmarked on FP16 engines inside ComfyUI, using 100 identical frames_
+
+| Device |     Model     | Input Resolution (WxH) | Output Resolution (WxH) | FPS |
+| :----: | :-----------: | :--------------------: | :---------------------: | :-: |
+|  RTX5090  | 4x-UltraSharp |       512 x 512        |       2048 x 2048       |  12.7  |
+|  RTX5090  | 4x-UltraSharp |       1280 x 1280      |       5120 x 5120       |  2.0  |
+|  RTX4090  | 4x-UltraSharp |       512 x 512        |       2048 x 2048       |  6.7  |
+|  RTX4090  | 4x-UltraSharp |       1280 x 1280      |       5120 x 5120       |  1.1  |
+|  RTX3060  | 4x-UltraSharp |       512 x 512        |       2048 x 2048       |  2.2  |
+|  RTX3060  | 4x-UltraSharp |       1280 x 1280      |       5120 x 5120       |  0.35  |
+
+## 🚀 Installation
+- Install via the manager
+- Or, navigate to the `/ComfyUI/custom_nodes` directory
+
+```bash
+git clone https://github.com/yuvraj108c/ComfyUI-Upscaler-Tensorrt.git
+cd ./ComfyUI-Upscaler-Tensorrt
+pip install -r requirements.txt
+```
+
+## 🛠️ Supported Models
+
+- These upscaler models have been tested to work with Tensorrt. Onnx are available [here](https://huggingface.co/yuvraj108c/ComfyUI-Upscaler-Onnx/tree/main)
+- The exported tensorrt models support dynamic image resolutions from 256x256 to 1280x1280 px (e.g 960x540, 512x512, 1280x720 etc).
+
+   - [4x-AnimeSharp](https://openmodeldb.info/models/4x-AnimeSharp)
+   - [4x-UltraSharp](https://openmodeldb.info/models/4x-UltraSharp)
+   - [4x-WTP-UDS-Esrgan](https://openmodeldb.info/models/4x-WTP-UDS-Esrgan)
+   - [4x_NMKD-Siax_200k](https://openmodeldb.info/models/4x-NMKD-Siax-CX)
+   - [4x_RealisticRescaler_100000_G](https://openmodeldb.info/models/4x-RealisticRescaler)
+   - [4x_foolhardy_Remacri](https://openmodeldb.info/models/4x-Remacri)
+   - [RealESRGAN_x4](https://openmodeldb.info/models/4x-realesrgan-x4plus)
+   - [4xNomos2_otf_esrgan](https://openmodeldb.info/models/4x-Nomos2-otf-esrgan)
+   - [4x-ClearRealityV1](https://openmodeldb.info/models/4x-ClearRealityV1)
+   - [4x_UniversalUpscalerV2-Neutral_115000_swaG](https://openmodeldb.info/models/4x-UniversalUpscalerV2-Neutral)
+   - [4x-UltraSharpV2_Lite](https://huggingface.co/Kim2091/UltraSharpV2) 
+
+## ☀️ Usage
+
+- Load [example workflow](assets/tensorrt_upscaling_workflow.json) 
+- Choose the appropriate model from the dropdown
+- The tensorrt engine will be built automatically
+- Load an image of resolution between 256-1280px
+- Set `resize_to` to resize the upscaled images to fixed resolutions
+
+## 🔧 Custom Models
+- To export other ESRGAN models, you'll have to build the onnx model first, using [export_onnx.py](scripts/export_onnx.py) 
+- Place the onnx model in `/ComfyUI/models/onnx/YOUR_MODEL.onnx`
+- Then, add your model to this list [load_upscaler_config.json](load_upscaler_config.json)
+- Finally, run the same workflow and choose your model
+- If you've tested another working tensorrt model, let me know to add it officially to this node
+
+## 🚨 Updates
+### 27 Auguest 2025
+- Support 4x-UltraSharpV2_Lite, 4x_UniversalUpscalerV2-Neutral_115000_swaG, 4x-ClearRealityV1
+- Load models from config [PR#57](https://github.com/yuvraj108c/ComfyUI-Upscaler-Tensorrt/pull/57)
+
+### 30 April 2025
+- Merge https://github.com/yuvraj108c/ComfyUI-Upscaler-Tensorrt/pull/48 by @BiiirdPrograms to fix soft-lock by raising an error when input image dimensions unsupported
+### 4 March 2025 (breaking)
+- Automatic tensorrt engines are built from the workflow itself, to simplify the process for non-technical people
+- Separate model loading and tensorrt processing into different nodes
+- Optimise post processing
+- Update onnx export script
+
+## ⚠️ Known issues
+
+- If you upgrade tensorrt version, you'll have to rebuild the engines
+- Only models with ESRGAN architecture are currently working
+- High ram usage when exporting `.pth` to `.onnx`
+
+## 🤖 Environment tested
+
+- Ubuntu 22.04 LTS, Cuda 12.4, Tensorrt 10.8, Python 3.10, H100 GPU
+- Windows 11
+
+## 👏 Credits
+
+- [NVIDIA/Stable-Diffusion-WebUI-TensorRT](https://github.com/NVIDIA/Stable-Diffusion-WebUI-TensorRT)
+- [comfyanonymous/ComfyUI](https://github.com/comfyanonymous/ComfyUI)
+
+## License
+
+[Creative Commons Attribution-NonCommercial-ShareAlike 4.0 International (CC BY-NC-SA 4.0)](https://creativecommons.org/licenses/by-nc-sa/4.0/)

__init__.py

@@ -0,0 +1,203 @@
+import os
+import folder_paths
+import numpy as np
+import torch
+from comfy.utils import ProgressBar
+from .trt_utilities import Engine
+from .utilities import download_file, ColoredLogger, get_final_resolutions
+import comfy.model_management as mm
+import time
+import tensorrt
+import json # <--- Import json module
+
+logger = ColoredLogger("ComfyUI-Upscaler-Tensorrt")
+
+IMAGE_DIM_MIN = 256
+IMAGE_DIM_OPT = 512
+IMAGE_DIM_MAX = 1280
+
+# --- Function to load configuration ---
+def load_node_config(config_filename="load_upscaler_config.json"):
+    """Loads node configuration from a JSON file."""
+    current_dir = os.path.dirname(__file__)
+    config_path = os.path.join(current_dir, config_filename)
+    
+    default_config = { # Fallback in case file is missing or corrupt
+        "model": {
+            "options": ["4x-UltraSharp"],
+            "default": "4x-UltraSharp",
+            "tooltip": "Default model (fallback from code)"
+        },
+        "precision": {
+            "options": ["fp16", "fp32"],
+            "default": "fp16",
+            "tooltip": "Default precision (fallback from code)"
+        }
+    }
+
+    try:
+        with open(config_path, 'r') as f:
+            config = json.load(f)
+        logger.info(f"Successfully loaded configuration from {config_filename}")
+        return config
+    except FileNotFoundError:
+        logger.warning(f"Configuration file '{config_path}' not found. Using default fallback configuration.")
+        return default_config
+    except json.JSONDecodeError:
+        logger.error(f"Error decoding JSON from '{config_path}'. Using default fallback configuration.")
+        return default_config
+    except Exception as e:
+        logger.error(f"An unexpected error occurred while loading '{config_path}': {e}. Using default fallback.")
+        return default_config
+
+# --- Load the configuration once when the module is imported ---
+LOAD_UPSCALER_NODE_CONFIG = load_node_config()
+
+
+class UpscalerTensorrt:
+    @classmethod
+    def INPUT_TYPES(s):
+        return {
+            "required": {
+                "images": ("IMAGE", {"tooltip": f"Images to be upscaled. Resolution must be between {IMAGE_DIM_MIN} and {IMAGE_DIM_MAX} px"}),
+                "upscaler_trt_model": ("UPSCALER_TRT_MODEL", {"tooltip": "Tensorrt model built and loaded"}),
+                "resize_to": (["none", "HD", "FHD", "2k", "4k", "2x", "3x"],{"tooltip": "Resize the upscaled image to fixed resolutions, optional"}),
+            }
+        }
+    RETURN_NAMES = ("IMAGE",)
+    RETURN_TYPES = ("IMAGE",)
+    FUNCTION = "upscaler_tensorrt"
+    CATEGORY = "tensorrt"
+    DESCRIPTION = "Upscale images with tensorrt"
+
+    def upscaler_tensorrt(self, images, upscaler_trt_model, resize_to):
+        images_bchw = images.permute(0, 3, 1, 2)
+        B, C, H, W = images_bchw.shape
+
+        for dim in (H, W):
+            if dim > IMAGE_DIM_MAX or dim < IMAGE_DIM_MIN:
+                raise ValueError(f"Input image dimensions fall outside of the supported range: {IMAGE_DIM_MIN} to {IMAGE_DIM_MAX} px!\nImage dimensions: {W}px by {H}px")
+
+        final_width, final_height = get_final_resolutions(W, H, resize_to)
+        logger.info(f"Upscaling {B} images from H:{H}, W:{W} to H:{H*4}, W:{W*4} | Final resolution: H:{final_height}, W:{final_width} | resize_to: {resize_to}")
+
+        shape_dict = {
+            "input": {"shape": (1, 3, H, W)},
+            "output": {"shape": (1, 3, H*4, W*4)},
+        }
+        upscaler_trt_model.activate()
+        upscaler_trt_model.allocate_buffers(shape_dict=shape_dict)
+
+        cudaStream = torch.cuda.current_stream().cuda_stream
+        pbar = ProgressBar(B)
+        images_list = list(torch.split(images_bchw, split_size_or_sections=1))
+
+        upscaled_frames = torch.empty((B, C, final_height, final_width), dtype=torch.float32, device=mm.intermediate_device())
+        must_resize = W*4 != final_width or H*4 != final_height
+
+        for i, img in enumerate(images_list):
+            result = upscaler_trt_model.infer({"input": img}, cudaStream)
+            result = result["output"]
+
+            if must_resize:
+                result = torch.nn.functional.interpolate(
+                    result, 
+                    size=(final_height, final_width),
+                    mode='bicubic',
+                    antialias=True
+                )
+            upscaled_frames[i] = result.to(mm.intermediate_device())
+            pbar.update(1)
+
+        output = upscaled_frames.permute(0, 2, 3, 1)
+        upscaler_trt_model.reset()
+        mm.soft_empty_cache()
+
+        logger.info(f"Output shape: {output.shape}")
+        return (output,)
+
+class LoadUpscalerTensorrtModel:
+    @classmethod
+    def INPUT_TYPES(cls): # Changed 's' to 'cls' for convention
+        # Use the pre-loaded configuration
+        model_config = LOAD_UPSCALER_NODE_CONFIG.get("model", {})
+        precision_config = LOAD_UPSCALER_NODE_CONFIG.get("precision", {})
+        
+        # Provide sensible defaults if keys are missing in the config (though load_node_config handles this broadly)
+        model_options = model_config.get("options", ["4x-UltraSharp"])
+        model_default = model_config.get("default", "4x-UltraSharp")
+        model_tooltip = model_config.get("tooltip", "Select a model.")
+
+        precision_options = precision_config.get("options", ["fp16", "fp32"])
+        precision_default = precision_config.get("default", "fp16")
+        precision_tooltip = precision_config.get("tooltip", "Select precision.")
+
+        return {
+            "required": {
+                "model": (model_options, {"default": model_default, "tooltip": model_tooltip}),
+                "precision": (precision_options, {"default": precision_default, "tooltip": precision_tooltip}),
+            }
+        }
+    
+    RETURN_NAMES = ("upscaler_trt_model",)
+    RETURN_TYPES = ("UPSCALER_TRT_MODEL",)
+    # FUNCTION = "main" # This was duplicated, removing
+    CATEGORY = "tensorrt"
+    DESCRIPTION = "Load tensorrt models, they will be built automatically if not found."
+    FUNCTION = "load_upscaler_tensorrt_model" # This is the correct one
+    
+    def load_upscaler_tensorrt_model(self, model, precision):
+        tensorrt_models_dir = os.path.join(folder_paths.models_dir, "tensorrt", "upscaler")
+        onnx_models_dir = os.path.join(folder_paths.models_dir, "onnx")
+
+        os.makedirs(tensorrt_models_dir, exist_ok=True)
+        os.makedirs(onnx_models_dir, exist_ok=True)
+
+        onnx_model_path = os.path.join(onnx_models_dir, f"{model}.onnx")
+        
+        engine_channel = 3
+        engine_min_batch, engine_opt_batch, engine_max_batch = 1, 1, 1
+        engine_min_h, engine_opt_h, engine_max_h = IMAGE_DIM_MIN, IMAGE_DIM_OPT, IMAGE_DIM_MAX
+        engine_min_w, engine_opt_w, engine_max_w = IMAGE_DIM_MIN, IMAGE_DIM_OPT, IMAGE_DIM_MAX
+        tensorrt_model_path = os.path.join(tensorrt_models_dir, f"{model}_{precision}_{engine_min_batch}x{engine_channel}x{engine_min_h}x{engine_min_w}_{engine_opt_batch}x{engine_channel}x{engine_opt_h}x{engine_opt_w}_{engine_max_batch}x{engine_channel}x{engine_max_h}x{engine_max_w}_{tensorrt.__version__}.trt")
+
+        if not os.path.exists(tensorrt_model_path):
+            if not os.path.exists(onnx_model_path):
+                onnx_model_download_url = f"https://huggingface.co/yuvraj108c/ComfyUI-Upscaler-Onnx/resolve/main/{model}.onnx"
+                logger.info(f"Downloading {onnx_model_download_url}")
+                download_file(url=onnx_model_download_url, save_path=onnx_model_path)
+            else:
+                logger.info(f"Onnx model found at: {onnx_model_path}")
+
+            logger.info(f"Building TensorRT engine for {onnx_model_path}: {tensorrt_model_path}")
+            mm.soft_empty_cache()
+            s = time.time()
+            engine = Engine(tensorrt_model_path)
+            engine.build(
+                onnx_path=onnx_model_path,
+                fp16= True if precision == "fp16" else False,
+                input_profile=[
+                    {"input": [(engine_min_batch,engine_channel,engine_min_h,engine_min_w), (engine_opt_batch,engine_channel,engine_opt_h,engine_min_w), (engine_max_batch,engine_channel,engine_max_h,engine_max_w)]},
+                ],
+            )
+            e = time.time()
+            logger.info(f"Time taken to build: {(e-s)} seconds")
+
+        logger.info(f"Loading TensorRT engine: {tensorrt_model_path}")
+        mm.soft_empty_cache()
+        engine = Engine(tensorrt_model_path)
+        engine.load()
+
+        return (engine,)
+
+NODE_CLASS_MAPPINGS = {
+    "UpscalerTensorrt": UpscalerTensorrt,
+    "LoadUpscalerTensorrtModel": LoadUpscalerTensorrtModel,
+}
+
+NODE_DISPLAY_NAME_MAPPINGS = {
+    "UpscalerTensorrt": "Upscaler Tensorrt ⚡",
+    "LoadUpscalerTensorrtModel": "Load Upscale Tensorrt Model",
+}
+
+__all__ = ['NODE_CLASS_MAPPINGS', 'NODE_DISPLAY_NAME_MAPPINGS']

assets/tensorrt_upscaling_workflow.json

@@ -0,0 +1,184 @@
+{
+  "last_node_id": 13,
+  "last_link_id": 16,
+  "nodes": [
+    {
+      "id": 5,
+      "type": "LoadImage",
+      "pos": [
+        -301.260986328125,
+        16.96155548095703
+      ],
+      "size": [
+        315,
+        314
+      ],
+      "flags": {},
+      "order": 0,
+      "mode": 0,
+      "inputs": [],
+      "outputs": [
+        {
+          "name": "IMAGE",
+          "type": "IMAGE",
+          "links": [
+            16
+          ]
+        },
+        {
+          "name": "MASK",
+          "type": "MASK",
+          "links": null
+        }
+      ],
+      "properties": {
+        "Node name for S&R": "LoadImage"
+      },
+      "widgets_values": [
+        "example.png",
+        "image"
+      ]
+    },
+    {
+      "id": 2,
+      "type": "LoadUpscalerTensorrtModel",
+      "pos": [
+        -297.7986755371094,
+        450.532958984375
+      ],
+      "size": [
+        315,
+        82
+      ],
+      "flags": {},
+      "order": 1,
+      "mode": 0,
+      "inputs": [],
+      "outputs": [
+        {
+          "name": "upscaler_trt_model",
+          "type": "UPSCALER_TRT_MODEL",
+          "links": [
+            1
+          ],
+          "slot_index": 0
+        }
+      ],
+      "properties": {
+        "Node name for S&R": "LoadUpscalerTensorrtModel"
+      },
+      "widgets_values": [
+        "4x-UltraSharp",
+        "fp16"
+      ]
+    },
+    {
+      "id": 13,
+      "type": "PreviewImage",
+      "pos": [
+        519.0885009765625,
+        -51.63800048828125
+      ],
+      "size": [
+        706.2752075195312,
+        756.0552978515625
+      ],
+      "flags": {},
+      "order": 3,
+      "mode": 0,
+      "inputs": [
+        {
+          "name": "images",
+          "type": "IMAGE",
+          "link": 14
+        }
+      ],
+      "outputs": [],
+      "properties": {
+        "Node name for S&R": "PreviewImage"
+      },
+      "widgets_values": []
+    },
+    {
+      "id": 3,
+      "type": "UpscalerTensorrt",
+      "pos": [
+        111.23614501953125,
+        352.1241760253906
+      ],
+      "size": [
+        315,
+        78
+      ],
+      "flags": {},
+      "order": 2,
+      "mode": 0,
+      "inputs": [
+        {
+          "name": "images",
+          "type": "IMAGE",
+          "link": 16
+        },
+        {
+          "name": "upscaler_trt_model",
+          "type": "UPSCALER_TRT_MODEL",
+          "link": 1
+        }
+      ],
+      "outputs": [
+        {
+          "name": "IMAGE",
+          "type": "IMAGE",
+          "links": [
+            14
+          ],
+          "slot_index": 0
+        }
+      ],
+      "properties": {
+        "Node name for S&R": "UpscalerTensorrt"
+      },
+      "widgets_values": [
+        "none"
+      ]
+    }
+  ],
+  "links": [
+    [
+      1,
+      2,
+      0,
+      3,
+      1,
+      "UPSCALER_TRT_MODEL"
+    ],
+    [
+      14,
+      3,
+      0,
+      13,
+      0,
+      "IMAGE"
+    ],
+    [
+      16,
+      5,
+      0,
+      3,
+      0,
+      "IMAGE"
+    ]
+  ],
+  "groups": [],
+  "config": {},
+  "extra": {
+    "ds": {
+      "scale": 1,
+      "offset": [
+        500.6130318234485,
+        102.83383472565859
+      ]
+    }
+  },
+  "version": 0.4
+}

load_upscaler_config.json

@@ -0,0 +1,24 @@
+{
+  "model": {
+    "options": [
+      "4x-AnimeSharp",
+      "4x-UltraSharp",
+      "4x-WTP-UDS-Esrgan",
+      "4x_NMKD-Siax_200k",
+      "4x_RealisticRescaler_100000_G",
+      "4x_foolhardy_Remacri",
+      "RealESRGAN_x4",
+      "4xNomos2_otf_esrgan",
+      "4x_UniversalUpscalerV2-Neutral_115000_swaG",
+      "4x-ClearRealityV1",
+      "4x-UltraSharpV2_Lite"
+    ],
+    "default": "4x-UltraSharp",
+    "tooltip": "These models have been tested with tensorrt. Loaded from config."
+  },
+  "precision": {
+    "options": ["fp16", "fp32"],
+    "default": "fp16",
+    "tooltip": "Precision to build the tensorrt engines. Loaded from config."
+  }
+}

requirements.txt

@@ -0,0 +1,3 @@
+tensorrt<=10.12.0.36
+polygraphy
+requests

scripts/export_onnx.py

@@ -0,0 +1,79 @@
+# download the upscale models & place inside models/upscaler_models
+# edit model paths accordingly 
+
+import torch
+import folder_paths
+from spandrel import ModelLoader, ImageModelDescriptor
+
+model_name = "4xNomos2_otf_esrgan.pth"
+onnx_save_path = "./4xNomos2_otf_esrgan.onnx"
+
+model_path = folder_paths.get_full_path_or_raise("upscale_models", model_name)
+model = ModelLoader().load_from_file(model_path).model.eval().cuda()
+
+# Check dynamic shapes for esrgan 4x model
+def supports_dynamic_shapes_esrgan(model, scale=4):
+
+    input_shapes = [
+    (1, 3, 64, 64),
+    (1, 3, 128, 128),
+    (1, 3, 256, 192),
+    (1, 3, 512, 256),
+    (1, 3, 512, 512)
+    ]
+
+    all_passed = True
+
+    with torch.no_grad():
+        for shape in input_shapes:
+            try:
+                dummy_input = torch.randn(*shape).cuda()
+                output = model(dummy_input)
+
+                expected_h = shape[2] * scale
+                expected_w = shape[3] * scale
+
+                assert output.shape[0] == shape[0], "Batch size mismatch"
+                assert output.shape[1] == shape[1], "Channel mismatch"
+                assert output.shape[2] == expected_h, f"Height mismatch: expected {expected_h}, got {output.shape[2]}"
+                assert output.shape[3] == expected_w, f"Width mismatch: expected {expected_w}, got {output.shape[3]}"
+
+                print(f"Success: input {shape} → output {output.shape}")
+            except Exception as e:
+                all_passed = False
+                print(f"Failure: input {shape} → error: {e}")
+                torch.cuda.empty_cache()
+
+    if all_passed: print(f"Success: Dynamic shapes supported.")
+    if not all_passed: print(f"Failure: Dynamic shapes NOT supported.")
+    return all_passed
+
+# Use smaller dummy input if model supports
+if supports_dynamic_shapes_esrgan(model):
+    shape = (1, 3, 64, 64)
+    print(f"Using {shape} input (less VRAM usage)")
+else:
+    shape = (1, 3, 512, 512)
+    print(f"Using {shape} input (large VRAM usage)")
+
+x = torch.rand(*shape).cuda()
+
+dynamic_axes = {
+    "input": {0: "batch_size", 2: "width", 3: "height"},
+    "output": {0: "batch_size", 2: "width", 3: "height"},
+}
+
+with torch.no_grad():
+    torch.onnx.export(
+        model,
+        x,
+        onnx_save_path,
+        verbose=True,
+        input_names=['input'],
+        output_names=['output'],
+        opset_version=17,
+        export_params=True,
+        dynamic_axes=dynamic_axes,
+    )
+
+print("Saved onnx to:", onnx_save_path)

scripts/export_trt.py

@@ -0,0 +1,29 @@
+import torch
+import time
+from utilities import Engine
+
+def export_trt(trt_path=None, onnx_path=None, use_fp16=True):
+    if trt_path is None:
+        trt_path = input("Enter the path to save the TensorRT engine (e.g ./realesrgan.engine): ")
+    if onnx_path is None:
+        onnx_path = input("Enter the path to the ONNX model (e.g ./realesrgan.onnx): ")
+
+    engine = Engine(trt_path)
+
+    torch.cuda.empty_cache()
+
+    s = time.time()
+    ret = engine.build(
+        onnx_path,
+        use_fp16,
+        enable_preview=True,
+        input_profile=[
+            {"input": [(1,3,256,256), (1,3,512,512), (1,3,1280,1280)]}, # any sizes from 256x256 to 1280x1280
+        ],
+    )
+    e = time.time()
+    print(f"Time taken to build: {(e-s)} seconds")
+
+    return ret
+
+export_trt()

scripts/export_trt_from_directory.py

@@ -0,0 +1,62 @@
+import os
+import torch
+import time
+from utilities import Engine
+
+def export_trt(trt_path=None, onnx_path=None, use_fp16=True):
+    option = input("Choose an option:\n1. Convert a single ONNX file\n2. Convert all ONNX files in a directory\nEnter your choice (1 or 2): ")
+
+    if option == '1':
+        onnx_path = input("Enter the path to the ONNX model (e.g ./realesrgan.onnx): ")
+        onnx_files = [onnx_path]
+        trt_dir = input("Enter the path to save the TensorRT engine (e.g ./trt_engine/): ")
+    elif option == '2':
+        onnx_dir = input("Enter the directory path containing ONNX models (e.g ./onnx_models/): ")
+        onnx_files = [os.path.join(onnx_dir, file) for file in os.listdir(onnx_dir) if file.endswith('.onnx')]
+        if not onnx_files:
+            raise ValueError(f"No .onnx files found in directory: {onnx_dir}")
+        trt_dir = input("Enter the directory path to save the TensorRT engines (e.g ./trt_engine/): ")
+    else:
+        raise ValueError("Invalid option. Please choose either 1 or 2.")
+
+    # Check if trt_dir already exists as a directory
+    if not os.path.exists(trt_dir):
+        os.makedirs(trt_dir)
+        
+    #os.makedirs(trt_dir, exist_ok=True)
+    total_files = len(onnx_files)
+    for index, onnx_path in enumerate(onnx_files):
+        engine = Engine(trt_path)
+
+        torch.cuda.empty_cache()
+        base_name = os.path.splitext(os.path.basename(onnx_path))[0]
+        trt_path = os.path.join(trt_dir, f"{base_name}.engine")
+
+        print(f"Converting {onnx_path} to {trt_path}")
+
+        s = time.time()
+
+        # Initialize Engine with trt_path and clear CUDA cache
+        engine = Engine(trt_path)
+        torch.cuda.empty_cache()
+
+        ret = engine.build(
+        onnx_path,
+        use_fp16,
+        enable_preview=True,
+        input_profile=[
+            {"input": [(1,3,256,256), (1,3,512,512), (1,3,1280,1280)]}, # any sizes from 256x256 to 1280x1280
+        ],
+        )
+
+        e = time.time()
+        print(f"Time taken to build: {(e-s)} seconds")
+        if index < total_files - 1:
+            # Delay for 10 seconds
+            print("Delaying for 10 seconds...")
+            time.sleep(10)
+            print("Resuming operations after delay...")
+
+    return
+
+export_trt()

trt_utilities.py

@@ -0,0 +1,283 @@
+#
+# Copyright 2022 The HuggingFace Inc. team.
+# SPDX-FileCopyrightText: Copyright (c) 1993-2023 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+#
+import torch
+from torch.cuda import nvtx
+from collections import OrderedDict
+import numpy as np
+from polygraphy.backend.common import bytes_from_path
+from polygraphy import util
+from polygraphy.backend.trt import ModifyNetworkOutputs, Profile
+from polygraphy.backend.trt import (
+    engine_from_bytes,
+    engine_from_network,
+    network_from_onnx_path,
+    save_engine,
+)
+from polygraphy.logger import G_LOGGER
+import tensorrt as trt
+from logging import error, warning
+from tqdm import tqdm
+import copy
+
+TRT_LOGGER = trt.Logger(trt.Logger.ERROR)
+G_LOGGER.module_severity = G_LOGGER.ERROR
+
+# Map of numpy dtype -> torch dtype
+numpy_to_torch_dtype_dict = {
+    np.uint8: torch.uint8,
+    np.int8: torch.int8,
+    np.int16: torch.int16,
+    np.int32: torch.int32,
+    np.int64: torch.int64,
+    np.float16: torch.float16,
+    np.float32: torch.float32,
+    np.float64: torch.float64,
+    np.complex64: torch.complex64,
+    np.complex128: torch.complex128,
+}
+if np.version.full_version >= "1.24.0":
+    numpy_to_torch_dtype_dict[np.bool_] = torch.bool
+else:
+    numpy_to_torch_dtype_dict[np.bool] = torch.bool
+
+# Map of torch dtype -> numpy dtype
+torch_to_numpy_dtype_dict = {
+    value: key for (key, value) in numpy_to_torch_dtype_dict.items()
+}
+
+class TQDMProgressMonitor(trt.IProgressMonitor):
+    def __init__(self):
+        trt.IProgressMonitor.__init__(self)
+        self._active_phases = {}
+        self._step_result = True
+        self.max_indent = 5
+
+    def phase_start(self, phase_name, parent_phase, num_steps):
+        leave = False
+        try:
+            if parent_phase is not None:
+                nbIndents = (
+                    self._active_phases.get(parent_phase, {}).get(
+                        "nbIndents", self.max_indent
+                    )
+                    + 1
+                )
+                if nbIndents >= self.max_indent:
+                    return
+            else:
+                nbIndents = 0
+                leave = True
+            self._active_phases[phase_name] = {
+                "tq": tqdm(
+                    total=num_steps, desc=phase_name, leave=leave, position=nbIndents
+                ),
+                "nbIndents": nbIndents,
+                "parent_phase": parent_phase,
+            }
+        except KeyboardInterrupt:
+            # The phase_start callback cannot directly cancel the build, so request the cancellation from within step_complete.
+            _step_result = False
+
+    def phase_finish(self, phase_name):
+        try:
+            if phase_name in self._active_phases.keys():
+                self._active_phases[phase_name]["tq"].update(
+                    self._active_phases[phase_name]["tq"].total
+                    - self._active_phases[phase_name]["tq"].n
+                )
+
+                parent_phase = self._active_phases[phase_name].get("parent_phase", None)
+                while parent_phase is not None:
+                    self._active_phases[parent_phase]["tq"].refresh()
+                    parent_phase = self._active_phases[parent_phase].get(
+                        "parent_phase", None
+                    )
+                if (
+                    self._active_phases[phase_name]["parent_phase"]
+                    in self._active_phases.keys()
+                ):
+                    self._active_phases[
+                        self._active_phases[phase_name]["parent_phase"]
+                    ]["tq"].refresh()
+                del self._active_phases[phase_name]
+            pass
+        except KeyboardInterrupt:
+            _step_result = False
+
+    def step_complete(self, phase_name, step):
+        try:
+            if phase_name in self._active_phases.keys():
+                self._active_phases[phase_name]["tq"].update(
+                    step - self._active_phases[phase_name]["tq"].n
+                )
+            return self._step_result
+        except KeyboardInterrupt:
+            # There is no need to propagate this exception to TensorRT. We can simply cancel the build.
+            return False
+
+
+class Engine:
+    def __init__(
+        self,
+        engine_path,
+    ):
+        self.engine_path = engine_path
+        self.engine = None
+        self.context = None
+        self.buffers = OrderedDict()
+        self.tensors = OrderedDict()
+        self.cuda_graph_instance = None  # cuda graph
+
+    def __del__(self):
+        del self.engine
+        del self.context
+        del self.buffers
+        del self.tensors
+
+    def reset(self, engine_path=None):
+        # del self.engine
+        del self.context
+        del self.buffers
+        del self.tensors
+        # self.engine_path = engine_path
+
+        self.context = None
+        self.buffers = OrderedDict()
+        self.tensors = OrderedDict()
+        self.inputs = {}
+        self.outputs = {}
+
+    def build(
+        self,
+        onnx_path,
+        fp16,
+        input_profile=None,
+        enable_refit=False,
+        enable_preview=False,
+        enable_all_tactics=False,
+        timing_cache=None,
+        update_output_names=None,
+    ):
+        p = [Profile()]
+        if input_profile:
+            p = [Profile() for i in range(len(input_profile))]
+            for _p, i_profile in zip(p, input_profile):
+                for name, dims in i_profile.items():
+                    assert len(dims) == 3
+                    _p.add(name, min=dims[0], opt=dims[1], max=dims[2])
+
+        config_kwargs = {}
+        if not enable_all_tactics:
+            config_kwargs["tactic_sources"] = []
+
+        network = network_from_onnx_path(
+            onnx_path, flags=[trt.OnnxParserFlag.NATIVE_INSTANCENORM]
+        )
+        if update_output_names:
+            print(f"Updating network outputs to {update_output_names}")
+            network = ModifyNetworkOutputs(network, update_output_names)
+
+        builder = network[0]
+        config = builder.create_builder_config()
+        config.progress_monitor = TQDMProgressMonitor()
+
+        config.set_flag(trt.BuilderFlag.FP16) if fp16 else None
+        config.set_flag(trt.BuilderFlag.REFIT) if enable_refit else None
+
+        profiles = copy.deepcopy(p)
+        for profile in profiles:
+            # Last profile is used for set_calibration_profile.
+            calib_profile = profile.fill_defaults(network[1]).to_trt(
+                builder, network[1]
+            )
+            config.add_optimization_profile(calib_profile)
+
+        try:
+            engine = engine_from_network(
+                network,
+                config,
+            )
+        except Exception as e:
+            error(f"Failed to build engine: {e}")
+            return 1
+        try:
+            save_engine(engine, path=self.engine_path)
+        except Exception as e:
+            error(f"Failed to save engine: {e}")
+            return 1
+        return 0
+
+    def load(self):
+        self.engine = engine_from_bytes(bytes_from_path(self.engine_path))
+
+    def activate(self, reuse_device_memory=None):
+        if reuse_device_memory:
+            self.context = self.engine.create_execution_context_without_device_memory()
+        #    self.context.device_memory = reuse_device_memory
+        else:
+            self.context = self.engine.create_execution_context()
+
+    def allocate_buffers(self, shape_dict=None, device="cuda"):
+        nvtx.range_push("allocate_buffers")
+        for idx in range(self.engine.num_io_tensors):
+            name = self.engine.get_tensor_name(idx)
+            binding = self.engine[idx]
+            if shape_dict and binding in shape_dict:
+                shape = shape_dict[binding]["shape"]
+            else:
+                shape = self.context.get_tensor_shape(name)
+
+            dtype = trt.nptype(self.engine.get_tensor_dtype(name))
+            if self.engine.get_tensor_mode(name) == trt.TensorIOMode.INPUT:
+                self.context.set_input_shape(name, shape)
+            tensor = torch.empty(
+                tuple(shape), dtype=numpy_to_torch_dtype_dict[dtype]
+            ).to(device=device)
+            self.tensors[binding] = tensor
+        nvtx.range_pop()
+
+    def infer(self, feed_dict, stream, use_cuda_graph=False):
+        nvtx.range_push("set_tensors")
+        for name, buf in feed_dict.items():
+            self.tensors[name].copy_(buf)
+
+        for name, tensor in self.tensors.items():
+            self.context.set_tensor_address(name, tensor.data_ptr())
+        nvtx.range_pop()
+        nvtx.range_push("execute")
+        noerror = self.context.execute_async_v3(stream)
+        if not noerror:
+            raise ValueError("ERROR: inference failed.")
+        nvtx.range_pop()
+        return self.tensors
+
+    def __str__(self):
+        out = ""
+            
+        # When raising errors in the upscaler, this str() called by comfy's execution.py,
+        # but the engine won't have the attributes required for stringification
+        # If str() also raises an error, comfy gets soft-locked, not running prompts until restarted
+        if not hasattr(self.engine, "num_optimization_profiles") or not hasattr(self.engine, "num_bindings"):
+            return out
+        
+        for opt_profile in range(self.engine.num_optimization_profiles):
+            for binding_idx in range(self.engine.num_bindings):
+                name = self.engine.get_binding_name(binding_idx)
+                shape = self.engine.get_profile_shape(opt_profile, name)
+                out += f"\t{name} = {shape}\n"
+        return out

utilities.py

@@ -0,0 +1,139 @@
+import requests
+from tqdm import tqdm
+import logging
+import sys
+
+class ColoredLogger:
+    COLORS = {
+        'RED': '\033[91m',
+        'GREEN': '\033[92m',
+        'YELLOW': '\033[93m',
+        'BLUE': '\033[94m',
+        'MAGENTA': '\033[95m',
+        'RESET': '\033[0m'
+    }
+
+    LEVEL_COLORS = {
+        'DEBUG': COLORS['BLUE'],
+        'INFO': COLORS['GREEN'],
+        'WARNING': COLORS['YELLOW'],
+        'ERROR': COLORS['RED'],
+        'CRITICAL': COLORS['MAGENTA']
+    }
+
+    def __init__(self, name="MY-APP"):
+        self.logger = logging.getLogger(name)
+        self.logger.setLevel(logging.DEBUG)
+        self.app_name = name
+        
+        # Prevent message propagation to parent loggers
+        self.logger.propagate = False
+        
+        # Clear existing handlers
+        self.logger.handlers = []
+        
+        # Create console handler
+        handler = logging.StreamHandler(sys.stdout)
+        handler.setLevel(logging.DEBUG)
+        
+        # Custom formatter class to handle colored components
+        class ColoredFormatter(logging.Formatter):
+            def format(self, record):
+                # Color the level name according to severity
+                level_color = ColoredLogger.LEVEL_COLORS.get(record.levelname, '')
+                colored_levelname = f"{level_color}{record.levelname}{ColoredLogger.COLORS['RESET']}"
+                
+                # Color the logger name in blue
+                colored_name = f"{ColoredLogger.COLORS['BLUE']}{record.name}{ColoredLogger.COLORS['RESET']}"
+                
+                # Set the colored components
+                record.levelname = colored_levelname
+                record.name = colored_name
+                
+                return super().format(record)
+        
+        # Create formatter with the new format
+        formatter = ColoredFormatter('[%(name)s|%(levelname)s] - %(message)s')
+        handler.setFormatter(formatter)
+        
+        self.logger.addHandler(handler)
+
+
+    def debug(self, message):
+        self.logger.debug(f"{self.COLORS['BLUE']}{message}{self.COLORS['RESET']}")
+
+    def info(self, message):
+        self.logger.info(f"{self.COLORS['GREEN']}{message}{self.COLORS['RESET']}")
+
+    def warning(self, message):
+        self.logger.warning(f"{self.COLORS['YELLOW']}{message}{self.COLORS['RESET']}")
+
+    def error(self, message):
+        self.logger.error(f"{self.COLORS['RED']}{message}{self.COLORS['RESET']}")
+
+    def critical(self, message):
+        self.logger.critical(f"{self.COLORS['MAGENTA']}{message}{self.COLORS['RESET']}")
+
+def download_file(url, save_path):
+    """
+    Download a file from URL with progress bar
+    
+    Args:
+        url (str): URL of the file to download
+        save_path (str): Path to save the file as
+    """
+    GREEN = '\033[92m'
+    RESET = '\033[0m'
+    response = requests.get(url, stream=True)
+    total_size = int(response.headers.get('content-length', 0))
+    
+    with open(save_path, 'wb') as file, tqdm(
+        desc=save_path,
+        total=total_size,
+        unit='iB',
+        unit_scale=True,
+        unit_divisor=1024,
+        colour='green',
+        bar_format=f'{GREEN}{{l_bar}}{{bar}}{RESET}{GREEN}{{r_bar}}{RESET}' 
+    ) as progress_bar:
+        for data in response.iter_content(chunk_size=1024):
+            size = file.write(data)
+            progress_bar.update(size)
+
+def get_final_resolutions(width, height, resize_to):
+    final_width = None
+    final_height = None
+    aspect_ratio = float(width/height)
+
+    match resize_to:
+        case "HD":
+            final_width = 1280
+            final_height = 720
+        case "FHD":
+            final_width = 1920
+            final_height = 1080
+        case "2k":
+            final_width = 2560
+            final_height = 1440
+        case "4k":
+            final_width = 3840
+            final_height = 2160
+        case "none":
+            final_width = width*4
+            final_height = height*4
+        case "2x":
+            final_width = width*2
+            final_height = height*2
+        case "3x":
+            final_width = width*3
+            final_height = height*3
+
+    if aspect_ratio == 1.0:
+        final_width = final_height
+
+    if aspect_ratio < 1.0 and resize_to not in ("none", "2x", "3x"):
+        temp = final_width
+        final_width = final_height
+        final_height = temp
+
+    return (final_width, final_height)