diff --git a/.gitignore b/.gitignore index 313b49c..fe2369e 100644 --- a/.gitignore +++ b/.gitignore @@ -1,3 +1,4 @@ +models __pycache__ -RIFE_RTX_ENGINE.cache -*.bak \ No newline at end of file +.vscode +.bak \ No newline at end of file diff --git a/.gitignore.bak b/.gitignore.bak new file mode 100644 index 0000000..236f609 --- /dev/null +++ b/.gitignore.bak @@ -0,0 +1,3 @@ +models +__pycache__ +.vscode \ No newline at end of file diff --git a/LICENSE b/LICENSE index 7a4a3ea..bfef380 100644 --- a/LICENSE +++ b/LICENSE @@ -1,202 +1,437 @@ +Attribution-NonCommercial-ShareAlike 4.0 International - Apache License - Version 2.0, January 2004 - http://www.apache.org/licenses/ +======================================================================= - TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION +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. - 1. Definitions. +Using Creative Commons Public Licenses - "License" shall mean the terms and conditions for use, reproduction, - and distribution as defined by Sections 1 through 9 of this document. +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. - "Licensor" shall mean the copyright owner or entity authorized by - the copyright owner that is granting the License. + 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 - "Legal Entity" shall mean the union of the acting entity and all - other entities that control, are controlled by, or are under common - control with that entity. For the purposes of this definition, - "control" means (i) the power, direct or indirect, to cause the - direction or management of such entity, whether by contract or - otherwise, or (ii) ownership of fifty percent (50%) or more of the - outstanding shares, or (iii) beneficial ownership of such entity. + 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 - "You" (or "Your") shall mean an individual or Legal Entity - exercising permissions granted by this License. +======================================================================= - "Source" form shall mean the preferred form for making modifications, - including but not limited to software source code, documentation - source, and configuration files. +Creative Commons Attribution-NonCommercial-ShareAlike 4.0 International +Public License - "Object" form shall mean any form resulting from mechanical - transformation or translation of a Source form, including but - not limited to compiled object code, generated documentation, - and conversions to other media types. +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. - "Work" shall mean the work of authorship, whether in Source or - Object form, made available under the License, as indicated by a - copyright notice that is included in or attached to the work - (an example is provided in the Appendix below). - "Derivative Works" shall mean any work, whether in Source or Object - form, that is based on (or derived from) the Work and for which the - editorial revisions, annotations, elaborations, or other modifications - represent, as a whole, an original work of authorship. For the purposes - of this License, Derivative Works shall not include works that remain - separable from, or merely link (or bind by name) to the interfaces of, - the Work and Derivative Works thereof. +Section 1 -- Definitions. - "Contribution" shall mean any work of authorship, including - the original version of the Work and any modifications or additions - to that Work or Derivative Works thereof, that is intentionally - submitted to Licensor for inclusion in the Work by the copyright owner - or by an individual or Legal Entity authorized to submit on behalf of - the copyright owner. For the purposes of this definition, "submitted" - means any form of electronic, verbal, or written communication sent - to the Licensor or its representatives, including but not limited to - communication on electronic mailing lists, source code control systems, - and issue tracking systems that are managed by, or on behalf of, the - Licensor for the purpose of discussing and improving the Work, but - excluding communication that is conspicuously marked or otherwise - designated in writing by the copyright owner as "Not a Contribution." + 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. - "Contributor" shall mean Licensor and any individual or Legal Entity - on behalf of whom a Contribution has been received by Licensor and - subsequently incorporated within the Work. + 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. - 2. Grant of Copyright License. Subject to the terms and conditions of - this License, each Contributor hereby grants to You a perpetual, - worldwide, non-exclusive, no-charge, royalty-free, irrevocable - copyright license to reproduce, prepare Derivative Works of, - publicly display, publicly perform, sublicense, and distribute the - Work and such Derivative Works in Source or Object form. + 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. - 3. Grant of Patent License. Subject to the terms and conditions of - this License, each Contributor hereby grants to You a perpetual, - worldwide, non-exclusive, no-charge, royalty-free, irrevocable - (except as stated in this section) patent license to make, have made, - use, offer to sell, sell, import, and otherwise transfer the Work, - where such license applies only to those patent claims licensable - by such Contributor that are necessarily infringed by their - Contribution(s) alone or by combination of their Contribution(s) - with the Work to which such Contribution(s) was submitted. If You - institute patent litigation against any entity (including a - cross-claim or counterclaim in a lawsuit) alleging that the Work - or a Contribution incorporated within the Work constitutes direct - or contributory patent infringement, then any patent licenses - granted to You under this License for that Work shall terminate - as of the date such litigation is filed. + 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. - 4. Redistribution. You may reproduce and distribute copies of the - Work or Derivative Works thereof in any medium, with or without - modifications, and in Source or Object form, provided that You - meet the following conditions: + 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. - (a) You must give any other recipients of the Work or - Derivative Works a copy of this License; and + 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. - (b) You must cause any modified files to carry prominent notices - stating that You changed the files; and + 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. - (c) You must retain, in the Source form of any Derivative Works - that You distribute, all copyright, patent, trademark, and - attribution notices from the Source form of the Work, - excluding those notices that do not pertain to any part of - the Derivative Works; and + h. Licensed Material means the artistic or literary work, database, + or other material to which the Licensor applied this Public + License. - (d) If the Work includes a "NOTICE" text file as part of its - distribution, then any Derivative Works that You distribute must - include a readable copy of the attribution notices contained - within such NOTICE file, excluding those notices that do not - pertain to any part of the Derivative Works, in at least one - of the following places: within a NOTICE text file distributed - as part of the Derivative Works; within the Source form or - documentation, if provided along with the Derivative Works; or, - within a display generated by the Derivative Works, if and - wherever such third-party notices normally appear. The contents - of the NOTICE file are for informational purposes only and - do not modify the License. You may add Your own attribution - notices within Derivative Works that You distribute, alongside - or as an addendum to the NOTICE text from the Work, provided - that such additional attribution notices cannot be construed - as modifying the 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. - You may add Your own copyright statement to Your modifications and - may provide additional or different license terms and conditions - for use, reproduction, or distribution of Your modifications, or - for any such Derivative Works as a whole, provided Your use, - reproduction, and distribution of the Work otherwise complies with - the conditions stated in this License. + j. Licensor means the individual(s) or entity(ies) granting rights + under this Public License. - 5. Submission of Contributions. Unless You explicitly state otherwise, - any Contribution intentionally submitted for inclusion in the Work - by You to the Licensor shall be under the terms and conditions of - this License, without any additional terms or conditions. - Notwithstanding the above, nothing herein shall supersede or modify - the terms of any separate license agreement you may have executed - with Licensor regarding such Contributions. + 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. - 6. Trademarks. This License does not grant permission to use the trade - names, trademarks, service marks, or product names of the Licensor, - except as required for reasonable and customary use in describing the - origin of the Work and reproducing the content of the NOTICE file. + 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. - 7. Disclaimer of Warranty. Unless required by applicable law or - agreed to in writing, Licensor provides the Work (and each - Contributor provides its Contributions) on an "AS IS" BASIS, - WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or - implied, including, without limitation, any warranties or conditions - of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A - PARTICULAR PURPOSE. You are solely responsible for determining the - appropriateness of using or redistributing the Work and assume any - risks associated with Your exercise of permissions under this License. + 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. - 8. Limitation of Liability. In no event and under no legal theory, - whether in tort (including negligence), contract, or otherwise, - unless required by applicable law (such as deliberate and grossly - negligent acts) or agreed to in writing, shall any Contributor be - liable to You for damages, including any direct, indirect, special, - incidental, or consequential damages of any character arising as a - result of this License or out of the use or inability to use the - Work (including but not limited to damages for loss of goodwill, - work stoppage, computer failure or malfunction, or any and all - other commercial damages or losses), even if such Contributor - has been advised of the possibility of such damages. + n. You means the individual or entity exercising the Licensed Rights + under this Public License. Your has a corresponding meaning. - 9. Accepting Warranty or Additional Liability. While redistributing - the Work or Derivative Works thereof, You may choose to offer, - and charge a fee for, acceptance of support, warranty, indemnity, - or other liability obligations and/or rights consistent with this - License. However, in accepting such obligations, You may act only - on Your own behalf and on Your sole responsibility, not on behalf - of any other Contributor, and only if You agree to indemnify, - defend, and hold each Contributor harmless for any liability - incurred by, or claims asserted against, such Contributor by reason - of your accepting any such warranty or additional liability. - END OF TERMS AND CONDITIONS +Section 2 -- Scope. - APPENDIX: How to apply the Apache License to your work. + a. License grant. - To apply the Apache License to your work, attach the following - boilerplate notice, with the fields enclosed by brackets "[]" - replaced with your own identifying information. (Don't include - the brackets!) The text should be enclosed in the appropriate - comment syntax for the file format. We also recommend that a - file or class name and description of purpose be included on the - same "printed page" as the copyright notice for easier - identification within third-party archives. + 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: - Copyright [yyyy] [name of copyright owner] + a. reproduce and Share the Licensed Material, in whole or + in part, for NonCommercial purposes only; and - 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 + b. produce, reproduce, and Share Adapted Material for + NonCommercial purposes only. - http://www.apache.org/licenses/LICENSE-2.0 + 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. - 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. \ No newline at end of file + 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. \ No newline at end of file diff --git a/README.md b/README.md index 2fbaf6e..8709c42 100644 --- a/README.md +++ b/README.md @@ -1,76 +1,115 @@ -# ComfyUI RIFE TensorRT-RTX +
-[![python](https://img.shields.io/badge/python-3.13.9-green)](https://www.python.org/downloads/release/python-31313/) -[![cuda](https://img.shields.io/badge/cuda-13.0-green)](https://developer.nvidia.com/cuda-downloads) -[![trt](https://img.shields.io/badge/TRT_RTX-1.4.0.76-green)](https://developer.nvidia.com/tensorrt) +# ComfyUI Rife TensorRT ⚡ + +[![python](https://img.shields.io/badge/python-3.12-green)](https://www.python.org/downloads/) +[![cuda](https://img.shields.io/badge/cuda-13.0-green)](https://developer.nvidia.com/cuda-13-0-2-download-archive) +[![trt](https://img.shields.io/badge/TRT-10.14.1.48-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) -## Overview - -This project is a fork of the original [ComfyUI_RIFE_TensorRT_Auto](https://github.com/silveroxides/ComfyUI_RIFE_TensorRT_Auto) by [silveroxides](https://github.com/silveroxides). -This project uses TensorRT-RTX to do RIFE interpolation - -## Features -- **NVIDIA TensorRT Acceleration**: Leverages Tensor Cores for up to 30x faster frame interpolation compared to standard local methods. -- **Dynamic Shape Support**: Configure minimum, optimal, and maximum resolutions to handle varying input sizes without rebuilding the engine for every image. -- **Automatic Engine Building**: Automatically compiles specialized `.trt` files from your `.onnx` models if a matching engine isn't found. -- **Persistent Timing Cache**: Saves hardware-specific optimization data to a local cache to speed up subsequent engine builds. -- **Memory Efficient**: Optimized for VRAM management with automatic garbage collection and cache clearing after inference. +![node](https://github.com/user-attachments/assets/5fd6d529-300c-42a5-b9cf-46e031f0bcb5) -## Examples +
-![Simple Example](assets/example.png) +This project provides a [TensorRT](https://github.com/NVIDIA/TensorRT) implementation of [RIFE](https://github.com/hzwer/ECCV2022-RIFE) for ultra fast frame interpolation inside ComfyUI -## Installation +This project is licensed under [CC BY-NC-SA](https://creativecommons.org/licenses/by-nc-sa/4.0/), everyone is FREE to access, use, modify and redistribute with the same license. -### ComfyUI Manager +If you like the project, please give me a star! ⭐ -This project has not been submitted to the ComfyUI Manager registry yet. But you can still install it this way: +--- -1. Open ComfyUI Manager. -2. Click the `Custom Nodes Manager` to open the custom nodes manager page. -3. On the bottom right corner, click the `Install via Git URL` button. -4. Enter the URL of this repository: `https://github.com/ThreadsOfFate/ComfyUI-RIFE-TensorRT-RTX` -5. Click "Confirm". -6. Restart ComfyUI. +## ⏱️ Performance -### Manual +_Note: The following results were benchmarked on FP16 engines inside ComfyUI, using 2000 frames consisting of 2 alternating similar frames, averaged 2-3 times_ -1. On the github page, click on the green `<> Code` button and then "Download ZIP". -2. Extract the root folder within the downloaded ZIP file to your ComfyUI `custom_nodes` directory. -3. Using the same python environment that runs ComfyUI, install the required dependencies: `python -m pip install -r custom_nodes/ComfyUI-Upscaler-TensorRT-Advanced/requirements.txt`. -4. Restart ComfyUI. +| Device | Rife Engine | Resolution| Multiplier | FPS | +| :----: | :-: | :-: | :-: | :-: | +| H100 | rife49_ensemble_True_scale_1_sim | 512 x 512 | 2 | 45 | +| H100 | rife49_ensemble_True_scale_1_sim | 512 x 512 | 4 | 57 | +| H100 | rife49_ensemble_True_scale_1_sim | 1280 x 1280 | 2 | 21 | +## 🚀 Installation -### Requirements +Navigate to the ComfyUI `/custom_nodes` directory + +```bash +git clone https://github.com/yuvraj108c/ComfyUI-Rife-Tensorrt +cd ./ComfyUI-Rife-Tensorrt pip install -r requirements.txt +``` +### ⚠️ CUDA Version Selection -### CUDA Toolkit Required -Please download NVIDIA CUDA Toolkit [https://developer.nvidia.com/cuda/toolkit](https://developer.nvidia.com/cuda/toolkit) +This node defaults to **CUDA 13** (RTX 50 series, driver 580+). -###Nodes Included -🟦 RIFE RTX Engine Builder (RIFERTXEngineLoader) - Loads a TensorRT-RTX engine. If the engine file for your specific GPU and resolution settings does not exist, it will automatically build one from the source ONNX model. -🟦 RIFE RTX Frame Interpolation (RIFERTXFrameInterpolation) - The core execution node. It takes an image and a loaded engine to perform the upscale. It supports an optional resize input for custom final dimensions. -🟦 RIFE RTX Engine Config (RIFERTXEngineConfig) - Allows you to define the resolution bounds (min/opt/max) for the engine. - Min/Max: The range of resolutions the engine can handle. - Opt: The resolution the engine is most optimized for. - - -GPU: NVIDIA RTX Series (TensorRT-RTX requires Tensor Cores). -Software: NVIDIA Drivers and a compatible version of TensorRT (handled via requirements.txt). +**For CUDA 12 (RTX 30/40 series):** +```bash +pip install -r requirements_cu12.txt +``` -###Troubleshooting -First Build Time: The first time you load a model with new shape settings, it may take several minutes to compile the engine. Check the console for progress via the TQDM bar. -VRAM Errors: If the build fails, try reducing the width_max or height_max in the Dynamic Shape Config. +**For CUDA 13 (Default):** +```bash +pip install -r requirements.txt +``` +### 📦 CUDA Toolkit Required -## 🛠️ Supported Models +The node automatically detects your CUDA installation via `CUDA_PATH` or `CUDA_HOME` environment variables. -rife47_ensemble_True_scale_1_sim -rife48_ensemble_True_scale_1_sim -rife49_ensemble_True_scale_1_sim +If CUDA is not detected, download from: https://developer.nvidia.com/cuda-13-0-2-download-archive + +### 🎯 Resolution Profiles + +The node supports resolution profiles to optimize VRAM usage: +- **small**: 480-896px (recommended for most video generation) +- **medium**: 720-1280px (for higher resolution videos) +- **custom**: Connect a "RIFE Custom Resolution Config" node for manual control + +The following RIFE models are supported and will be automatically downloaded and built: + - **rife49_ensemble_True_scale_1_sim** (default) - Latest and most accurate + - **rife48_ensemble_True_scale_1_sim** - Good balance of speed and quality + - **rife47_ensemble_True_scale_1_sim** - Fastest option + +Models are automatically downloaded from [HuggingFace](https://huggingface.co/yuvraj108c/rife-onnx) and TensorRT engines are built on first use. + +## ☀️ Usage + +1. **Load Model**: Insert `Right Click -> Add Node -> tensorrt -> Load Rife Tensorrt Model` + - Choose your preferred RIFE model (rife47, rife48, or rife49) + - Select precision (fp16 recommended for speed, fp32 for maximum accuracy) + - Select resolution profile (small, medium, or custom) + - The model will be automatically downloaded and TensorRT engine built on first use + +2. **Process Frames**: Insert `Right Click -> Add Node -> tensorrt -> Rife Tensorrt` + - Connect the loaded model from step 1 + - Input your video frames + - Configure interpolation settings (multiplier, CUDA graph, etc.) + +## 🤖 Environment tested + +- Windows 11, CUDA 13.0, TensorRT 10.14.1.48, Python 3.12, RTX 5070 Ti +- WSL Ubuntu 24.04.03 LTS, CUDA 12.9, TensorRT 10.13.3.9, Python 3.12.11, RTX 5080 + +## 🚨 Updates + +### January 2026 +- **CUDA 13 Default**: Updated to CUDA 13.0 and TensorRT 10.14.1.48 +- **Auto CUDA Detection**: Automatically finds CUDA toolkit and DLL paths +- **Resolution Profiles**: Added small/medium/custom profiles to reduce VRAM usage + +### December 2025 +- **Automatic Model Management**: No more manual downloads! Models are automatically downloaded from HuggingFace and TensorRT engines are built on demand +- **Improved Workflow**: New two-node system with `Load Rife Tensorrt Model` + `Rife Tensorrt` for better organization +- **Updated Dependencies**: TensorRT updated to 10.13.3.9 for better performance and compatibility + +## 👏 Credits + +- https://github.com/styler00dollar/VSGAN-tensorrt-docker +- https://github.com/Fannovel16/ComfyUI-Frame-Interpolation +- https://github.com/hzwer/ECCV2022-RIFE + +## License + +[Creative Commons Attribution-NonCommercial-ShareAlike 4.0 International (CC BY-NC-SA 4.0)](https://creativecommons.org/licenses/by-nc-sa/4.0/) diff --git a/__init__.py b/__init__.py index a6d53a9..cbb82a2 100644 --- a/__init__.py +++ b/__init__.py @@ -1,20 +1,324 @@ -"""ComfyUI-Rife-TensorRT Extension Initialization. +import torch +import os +import sys +from pathlib import Path +from comfy.model_management import get_torch_device +from .vfi_utilities import preprocess_frames, postprocess_frames, generate_frames_rife, logger +from .trt_utilities import Engine +from .utilities import download_file, ColoredLogger +import folder_paths +import time -This module acts as the entry point for ComfyUI, exporting the node classes -and display names required to register the TensorRT-RTX RIFE implementation -into the ComfyUI node graph. +# Auto-detect CUDA toolkit and add DLL path before importing polygraphy +def _setup_cuda_dll_path(): + """Auto-detect CUDA toolkit and add cudart64 DLL path on Windows.""" + if not sys.platform.startswith("win"): + return + + cuda_root = None + + # Check for CUDA_PATH or CUDA_HOME environment variables + cuda_root = os.environ.get("CUDA_PATH") or os.environ.get("CUDA_HOME") + + if not cuda_root: + # Try default Windows install location + program_files = os.environ.get("PROGRAMFILES") + if program_files: + cuda_base = Path(program_files) / "NVIDIA GPU Computing Toolkit" / "CUDA" + if cuda_base.exists(): + # Find highest version directory + versions = sorted([d for d in cuda_base.iterdir() if d.is_dir()], reverse=True) + if versions: + cuda_root = str(versions[0]) + + if cuda_root: + cuda_path = Path(cuda_root) + # CUDA 13.0+ puts cudart64 in bin/x64 subdirectory + cuda_bin_x64 = cuda_path / "bin" / "x64" + if cuda_bin_x64.exists() and any(cuda_bin_x64.glob("cudart64*.dll")): + os.add_dll_directory(str(cuda_bin_x64)) + return + # Fallback to regular bin directory for older CUDA versions + cuda_bin = cuda_path / "bin" + if cuda_bin.exists() and any(cuda_bin.glob("cudart64*.dll")): + os.add_dll_directory(str(cuda_bin)) + return + + # CUDA toolkit not found - print warning with download link + print("[ComfyUI-Rife-TensorRT] WARNING: CUDA toolkit not found.") + print(" Set CUDA_PATH environment variable or install CUDA toolkit.") + print(" Download: https://developer.nvidia.com/cuda-13-0-2-download-archive") -Attributes: - NODE_CLASS_MAPPINGS (dict): Maps internal lookup strings to node classes. - NODE_DISPLAY_NAME_MAPPINGS (dict): Maps internal lookup strings to - human-readable UI names. - WEB_DIRECTORY (str): Path to the directory containing frontend - javascript extensions. -""" +_setup_cuda_dll_path() -from .nodes import NODE_CLASS_MAPPINGS, NODE_DISPLAY_NAME_MAPPINGS +from polygraphy import cuda +import comfy.model_management as mm +import tensorrt +import json -WEB_DIRECTORY = "./js" +ENGINE_DIR = os.path.join(folder_paths.models_dir, "tensorrt", "rife") + +# Default resolution profiles (fallback if config is missing) +DEFAULT_RESOLUTION_PROFILES = { + "small": {"min": 384, "opt": 720, "max": 1080}, + "medium": {"min": 672, "opt": 1080, "max": 1312} +} + +# Logger for this module +rife_logger = ColoredLogger("ComfyUI-Rife-Tensorrt") + +# Function to load configuration +def load_node_config(config_filename="load_rife_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 = { + "model": { + "options": ["rife49_ensemble_True_scale_1_sim"], + "default": "rife49_ensemble_True_scale_1_sim", + "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) + rife_logger.info(f"Successfully loaded configuration from {config_filename}") + return config + except FileNotFoundError: + rife_logger.warning(f"Configuration file '{config_path}' not found. Using default fallback configuration.") + return default_config + except json.JSONDecodeError: + rife_logger.error(f"Error decoding JSON from '{config_path}'. Using default fallback configuration.") + return default_config + except Exception as e: + rife_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_RIFE_NODE_CONFIG = load_node_config() + + +class CustomResolutionConfig: + """Node to configure custom resolution dimensions for TensorRT engine building.""" + @classmethod + def INPUT_TYPES(s): + return { + "required": { + "min_dim": ("INT", {"default": 384, "min": 64, "max": 4096, "step": 8, "tooltip": "Minimum resolution dimension"}), + "opt_dim": ("INT", {"default": 720, "min": 64, "max": 4096, "step": 8, "tooltip": "Optimal resolution dimension (most common)"}), + "max_dim": ("INT", {"default": 1312, "min": 64, "max": 4096, "step": 8, "tooltip": "Maximum resolution dimension"}), + } + } + + RETURN_TYPES = ("RIFE_RESOLUTION_CONFIG",) + RETURN_NAMES = ("resolution_config",) + FUNCTION = "configure" + CATEGORY = "tensorrt" + DESCRIPTION = "Configure custom resolution dimensions for RIFE TensorRT engine." + + def configure(self, min_dim, opt_dim, max_dim): + config = { + "min": min_dim, + "opt": opt_dim, + "max": max_dim, + } + return (config,) + + +class AutoLoadRifeTensorrtModel: + @classmethod + def INPUT_TYPES(cls): + # Use the pre-loaded configuration + model_config = LOAD_RIFE_NODE_CONFIG.get("model", {}) + precision_config = LOAD_RIFE_NODE_CONFIG.get("precision", {}) + + # Provide sensible defaults if keys are missing in the config + model_options = model_config.get("options", ["rife49_ensemble_True_scale_1_sim"]) + model_default = model_config.get("default", "rife49_ensemble_True_scale_1_sim") + model_tooltip = model_config.get("tooltip", "Select a RIFE model.") + + precision_options = precision_config.get("options", ["fp16", "fp32"]) + precision_default = precision_config.get("default", "fp16") + precision_tooltip = precision_config.get("tooltip", "Select precision.") + + # Resolution profile configuration + profile_config = LOAD_RIFE_NODE_CONFIG.get("resolution_profile", {}) + profile_options = profile_config.get("options", ["small", "medium"]) + # Ensure 'custom' is always available + if "custom" not in profile_options: + profile_options = profile_options + ["custom"] + profile_default = profile_config.get("default", "small") + profile_tooltip = profile_config.get("tooltip", "Resolution range for TensorRT engine. Use 'custom' with the INT inputs below.") + + return { + "required": { + "model": (model_options, {"default": model_default, "tooltip": model_tooltip}), + "precision": (precision_options, {"default": precision_default, "tooltip": precision_tooltip}), + "resolution_profile": (profile_options, {"default": profile_default, "tooltip": profile_tooltip}), + }, + "optional": { + "custom_config": ("RIFE_RESOLUTION_CONFIG", {"tooltip": "Custom resolution config (used when profile='custom')"}), + } + } + + RETURN_NAMES = ("rife_trt_model",) + RETURN_TYPES = ("RIFE_TRT_MODEL",) + CATEGORY = "tensorrt" + DESCRIPTION = "Load RIFE tensorrt models, they will be built automatically if not found." + FUNCTION = "load_rife_tensorrt_model" + + def load_rife_tensorrt_model(self, model, precision, resolution_profile, custom_config=None): + tensorrt_models_dir = os.path.join(folder_paths.models_dir, "tensorrt", "rife") + 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") + + # Get resolution dimensions based on profile + if resolution_profile == "custom": + if custom_config is None: + rife_logger.warning("Custom profile selected but no custom_config provided. Using defaults that cover both small and medium ranges.") + dim_min, dim_opt, dim_max = 384, 720, 1312 + else: + dim_min = custom_config.get("min", 384) + dim_opt = custom_config.get("opt", 720) + dim_max = custom_config.get("max", 1312) + # Use dimensions in profile name for custom engines + profile_name = f"custom_{dim_min}_{dim_opt}_{dim_max}" + else: + profiles = LOAD_RIFE_NODE_CONFIG.get("resolution_profiles", DEFAULT_RESOLUTION_PROFILES) + profile = profiles.get(resolution_profile, DEFAULT_RESOLUTION_PROFILES["small"]) + dim_min = profile.get("min", 384) + dim_opt = profile.get("opt", 720) + dim_max = profile.get("max", 1080) + profile_name = resolution_profile + rife_logger.info(f"Using resolution profile '{profile_name}': min={dim_min}, opt={dim_opt}, max={dim_max}") + + # Build tensorrt model path with detailed naming (includes profile) + engine_channel = 3 + engine_min_batch, engine_opt_batch, engine_max_batch = 1, 1, 1 + engine_min_h, engine_opt_h, engine_max_h = dim_min, dim_opt, dim_max + engine_min_w, engine_opt_w, engine_max_w = dim_min, dim_opt, dim_max + tensorrt_model_path = os.path.join(tensorrt_models_dir, f"{model}_{precision}_{profile_name}_{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/rife-onnx/resolve/main/{model}.onnx" + rife_logger.info(f"Downloading {onnx_model_download_url}") + download_file(url=onnx_model_download_url, save_path=onnx_model_path) + else: + rife_logger.info(f"ONNX model found at: {onnx_model_path}") + + rife_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=[ + { + "img0": [(engine_min_batch, engine_channel, engine_min_h, engine_min_w), (engine_opt_batch, engine_channel, engine_opt_h, engine_opt_w), (engine_max_batch, engine_channel, engine_max_h, engine_max_w)], + "img1": [(engine_min_batch, engine_channel, engine_min_h, engine_min_w), (engine_opt_batch, engine_channel, engine_opt_h, engine_opt_w), (engine_max_batch, engine_channel, engine_max_h, engine_max_w)], + } + ], + ) + e = time.time() + rife_logger.info(f"Time taken to build: {(e-s)} seconds") + + rife_logger.info(f"Loading TensorRT engine: {tensorrt_model_path}") + mm.soft_empty_cache() + engine = Engine(tensorrt_model_path) + engine.load() + + return (engine,) + +class AutoRifeTensorrt: + @classmethod + def INPUT_TYPES(s): + return { + "required": { + "frames": ("IMAGE", {"tooltip": "Input frames for video frame interpolation"}), + "rife_trt_model": ("RIFE_TRT_MODEL", {"tooltip": "Tensorrt model built and loaded"}), + "clear_cache_after_n_frames": ("INT", {"default": 100, "min": 1, "max": 1000, "tooltip": "Clear CUDA cache after processing this many frames"}), + "multiplier": ("INT", {"default": 2, "min": 1, "tooltip": "Frame interpolation multiplier"}), + "use_cuda_graph": ("BOOLEAN", {"default": True, "tooltip": "Use CUDA graph for better performance"}), + "keep_model_loaded": ("BOOLEAN", {"default": False, "tooltip": "Keep model loaded in memory after processing"}), + }, + } + + RETURN_TYPES = ("IMAGE", ) + FUNCTION = "vfi" + CATEGORY = "tensorrt" + OUTPUT_NODE=True + + def vfi( + self, + frames, + rife_trt_model, + clear_cache_after_n_frames=100, + multiplier=2, + use_cuda_graph=True, + keep_model_loaded=False, + ): + B, H, W, C = frames.shape + shape_dict = { + "img0": {"shape": (1, 3, H, W)}, + "img1": {"shape": (1, 3, H, W)}, + "output": {"shape": (1, 3, H, W)}, + } + + cudaStream = cuda.Stream() + + # Use the provided model directly + engine = rife_trt_model + logger(f"Using loaded TensorRT engine") + + # Activate and allocate buffers for the engine + engine.activate() + engine.allocate_buffers(shape_dict=shape_dict) + + frames = preprocess_frames(frames) + + def return_middle_frame(frame_0, frame_1, timestep): + timestep_t = torch.tensor([timestep], dtype=torch.float32).to(get_torch_device()) + # s = time.time() + output = engine.infer({"img0": frame_0, "img1": frame_1, "timestep": timestep_t}, cudaStream, use_cuda_graph) + # e = time.time() + # print(f"Time taken to infer: {(e-s)*1000} ms") + + result = output['output'] + return result + + result = generate_frames_rife(frames, clear_cache_after_n_frames, multiplier, return_middle_frame) + out = postprocess_frames(result) + + if not keep_model_loaded: + engine.reset() + + return (out,) + + +NODE_CLASS_MAPPINGS = { + "AutoRifeTensorrt": AutoRifeTensorrt, + "AutoLoadRifeTensorrtModel": AutoLoadRifeTensorrtModel, + "CustomResolutionConfig": CustomResolutionConfig, +} + +NODE_DISPLAY_NAME_MAPPINGS = { + "AutoRifeTensorrt": "Auto RIFE TensorRT", + "AutoLoadRifeTensorrtModel": "(Down)load RIFE TensorRT Model", + "CustomResolutionConfig": "RIFE Custom Resolution Config", +} __all__ = ['NODE_CLASS_MAPPINGS', 'NODE_DISPLAY_NAME_MAPPINGS'] diff --git a/scripts/export_onnx.py b/export_onnx.py similarity index 100% rename from scripts/export_onnx.py rename to export_onnx.py diff --git a/scripts/export_trt.py b/export_trt.py similarity index 100% rename from scripts/export_trt.py rename to export_trt.py diff --git a/load_rife_config.json b/load_rife_config.json new file mode 100644 index 0000000..3db8bf6 --- /dev/null +++ b/load_rife_config.json @@ -0,0 +1,25 @@ +{ + "model": { + "options": [ + "rife47_ensemble_True_scale_1_sim", + "rife48_ensemble_True_scale_1_sim", + "rife49_ensemble_True_scale_1_sim" + ], + "default": "rife49_ensemble_True_scale_1_sim", + "tooltip": "RIFE models for video frame interpolation. 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." + }, + "resolution_profile": { + "options": ["small", "medium"], + "default": "small", + "tooltip": "Resolution range for TensorRT engine. Small: 384-1080px, Medium: 672-1312px. Smaller ranges use less VRAM." + }, + "resolution_profiles": { + "small": { "min": 384, "opt": 720, "max": 1080 }, + "medium": { "min": 672, "opt": 1080, "max": 1312 } + } +} \ No newline at end of file diff --git a/nodes.py b/nodes.py deleted file mode 100644 index bb0ddcc..0000000 --- a/nodes.py +++ /dev/null @@ -1,311 +0,0 @@ -"""TensorRT-RTX RIFE Implementation for ComfyUI. - -This module provides custom nodes for video frame interpolation (VFI) using -NVIDIA's TensorRT-RTX. It includes nodes for configuring engine resolutions, -building/loading optimized .trt engines, and performing high-speed inference. -""" - -import json -import os -import sys - -import folder_paths -import logging -import tensorrt_rtx as trt -import time -import torch - -import comfy.model_management as mm - -from pathlib import Path - -# Use polygraphy TensorRTX backend -os.environ["POLYGRAPHY_USE_TENSORRT_RTX"] = "1" - -from polygraphy import cuda -from .vfi_utilities import preprocess_frames, postprocess_frames, generate_frames_rife -from .trt_utilities import Engine - -def update_folder_names_and_paths(key, targets=[]): - """Registers specific folder paths within ComfyUI's path manager. - - Ensures that custom subdirectories (like onnx_rife) are recognized by - ComfyUI's internal file system and UI loaders. - - Args: - key (str): The primary key to register in folder_names_and_paths. - targets (list): Fallback keys to inherit paths from if the primary - key is missing. - """ - # check for existing key - base = folder_paths.folder_names_and_paths.get(key, ([], {})) - base = base[0] if isinstance(base[0], (list, set, tuple)) else [] - # find base key & add w/ fallback, sanity check + warning - target = next((x for x in targets if x in folder_paths.folder_names_and_paths), targets[0]) - orig, _ = folder_paths.folder_names_and_paths.get(target, ([], {})) - folder_paths.folder_names_and_paths[key] = (orig or base, {".onnx"}) - if base and base != orig: - logging.warning(f"Unknown file list already present on key {key}: {base}") - - -# Auto-detect CUDA toolkit and add DLL path before importing polygraphy -def _setup_cuda_dll_path(): - """Auto-detects the CUDA toolkit and adds cudart64 DLL path on Windows. - - This is necessary for TensorRT-RTX to correctly locate runtime libraries - when running in a portable environment or on systems where CUDA is not - globally in the system PATH. - """ - if not sys.platform.startswith("win"): - return - - cuda_root = None - - # Check for CUDA_PATH or CUDA_HOME environment variables - cuda_root = os.environ.get("CUDA_PATH") or os.environ.get("CUDA_HOME") - - if not cuda_root: - # Try default Windows install location - program_files = os.environ.get("PROGRAMFILES") - if program_files: - cuda_base = Path(program_files) / "NVIDIA GPU Computing Toolkit" / "CUDA" - if cuda_base.exists(): - # Find highest version directory - versions = sorted([d for d in cuda_base.iterdir() if d.is_dir()], reverse=True) - if versions: - cuda_root = str(versions[0]) - - if cuda_root: - cuda_path = Path(cuda_root) - # CUDA 13.0+ puts cudart64 in bin/x64 subdirectory - cuda_bin_x64 = cuda_path / "bin" / "x64" - if cuda_bin_x64.exists() and any(cuda_bin_x64.glob("cudart64*.dll")): - os.add_dll_directory(str(cuda_bin_x64)) - return - # Fallback to regular bin directory for older CUDA versions - cuda_bin = cuda_path / "bin" - if cuda_bin.exists() and any(cuda_bin.glob("cudart64*.dll")): - os.add_dll_directory(str(cuda_bin)) - return - - # CUDA toolkit not found - print warning with download link - print("[ComfyUI-Rife-TensorRT] WARNING: CUDA toolkit not found.") - print(" Set CUDA_PATH environment variable or install CUDA toolkit.") - print(" Download: https://developer.nvidia.com/cuda-13-0-2-download-archive") - -# Initialize Environment -_setup_cuda_dll_path() -os.environ["POLYGRAPHY_USE_TENSORRT_RTX"] = "1" -update_folder_names_and_paths("onnx_rife", ["onnx"]) - -class RIFERTXEngineOptions: - """Data container for RIFE shape boundaries. - - Attributes: - min_dim (int): Minimum supported resolution. - opt_dim (int): Targeted optimal resolution. - max_dim (int): Maximum supported resolution. - """ - min_dim: int = 256 - opt_dim: int = 512 - max_dim: int = 1280 - -class RIFERTXEngineConfig: - """ComfyUI Node for configuring resolution profiles. - - Provides a UI interface to set the dynamic shape profiles required - during the TensorRT engine build process. - """ - @classmethod - def INPUT_TYPES(s): - return { - "required": { - "min_dim": ("INT", {"default": 256, "min": 64, "max": 4096, "step": 8, "tooltip": "Minimum resolution dimension"}), - "opt_dim": ("INT", {"default": 512, "min": 64, "max": 4096, "step": 8, "tooltip": "Optimal resolution dimension (most common)"}), - "max_dim": ("INT", {"default": 1280, "min": 64, "max": 4096, "step": 8, "tooltip": "Maximum resolution dimension"}), - } - } - - RETURN_TYPES = ("RIFE_RTX_ENGINE_CONFIG",) - RETURN_NAMES = ("RIFE_ENGINE_CONFIG",) - FUNCTION = "configure" - CATEGORY = "TensorRT_RTX/RIFE" - DESCRIPTION = "Configure custom resolution dimensions for RIFE TensorRT engine." - - def configure(self, min_dim, opt_dim, max_dim): - """Wraps dimensions into a configuration dictionary. - - Args: - min_dim (int): Minimum resolution. - opt_dim (int): Optimal resolution. - max_dim (int): Maximum resolution. - - Returns: - tuple: Contains the configuration dictionary. - """ - config = {"min_dim": min_dim, "opt_dim": opt_dim, "max_dim": max_dim,} - return (config,) - - -class RIFERTXEngineLoader: - """ComfyUI Node for loading or building RIFE TensorRT engines. - - Checks if a .trt engine matching the requested configuration exists. - If not, it compiles the ONNX model into a TensorRT engine. - """ - @classmethod - def INPUT_TYPES(cls): - onnx_names = [x for x in folder_paths.get_filename_list("onnx_rife") if "rife" in x.lower()] - return { - "required": { "onnx_name": (onnx_names,)}, - "optional": { "config": ("RIFE_RTX_ENGINE_CONFIG", {"tooltip": "Options for building the TensorRT engine"}),} - } - - RETURN_NAMES = ("RIFE_ENGINE",) - RETURN_TYPES = ("RIFE_RTX_ENGINE",) - CATEGORY = "TensorRT_RTX/RIFE" - DESCRIPTION = "Load RIFE tensorrt_rtx models, they will be built automatically if not found." - FUNCTION = "load_rife_tensorrt_rtx_model" - - def load_rife_tensorrt_rtx_model(self, onnx_name, config=None): - """Loads or builds the RIFE TensorRT engine. - - Args: - onnx_name (str): Name of the source ONNX file. - config (dict, optional): Resolution profile. Defaults to RIFERTXEngineOptions. - - Returns: - tuple: The loaded Engine object. - """ - if config is None: - config = RIFERTXEngineOptions() - - tensorrt_models_dir = os.path.join(folder_paths.models_dir, "tensorrt", "rife") - 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"{onnx_name}") - - # Build tensorrt model path with detailed naming (includes profile) - engine_channel = 3 - e_min, e_opt, e_max = config.min_dim, config.opt_dim, config.max_dim - tensorrt_model_path = os.path.join(tensorrt_models_dir, f"{onnx_name}_1x{engine_channel}x{e_min}x{e_min}_1x{engine_channel}x{e_opt}x{e_opt}_1x{engine_channel}x{e_max}x{e_max}_{trt.__version__}.trt") - - engine = Engine(tensorrt_model_path) - - if not os.path.exists(tensorrt_model_path): - logging.info(f"Building TensorRT_RTX engine for {onnx_model_path}: {tensorrt_model_path}") - - mm.soft_empty_cache() - s = time.time() - - result = engine.build( - onnx_path=onnx_model_path, - input_profile=[ - { - "img0": [(1, engine_channel, e_min, e_min), (1, engine_channel, e_opt, e_opt), (1, engine_channel, e_max, e_max)], - "img1": [(1, engine_channel, e_min, e_min), (1, engine_channel, e_opt, e_opt), (1, engine_channel, e_max, e_max)], - } - ], - ) - - if result != 0: - raise ValueError(f"Failed to build the engine") - e = time.time() - logging.info(f"Time taken to build: {(e-s)} seconds") - - logging.info(f"Loading TensorRT_RTX engine: {tensorrt_model_path}") - mm.soft_empty_cache() - engine.load() - - return (engine,) - -class RIFERTXFrameInterpolation: - """ComfyUI Node for performing VFI inference. - - Uses the loaded TensorRT engine to interpolate frames in a video sequence. - """ - @classmethod - def INPUT_TYPES(s): - return { - "required": { - "frames": ("IMAGE", {"tooltip": "Input frames for video frame interpolation"}), - "engine": ("RIFE_RTX_ENGINE", {"tooltip": "Tensorrt_RTX model built and loaded"}), - "clear_cache_after_n_frames": ("INT", {"default": 100, "min": 1, "max": 1000, "tooltip": "Clear CUDA cache after processing this many frames"}), - "multiplier": ("INT", {"default": 2, "min": 1, "tooltip": "Frame interpolation multiplier"}), - "use_cuda_graph": ("BOOLEAN", {"default": True, "tooltip": "Use CUDA graph for better performance"}), - "keep_model_loaded": ("BOOLEAN", {"default": False, "tooltip": "Keep model loaded in memory after processing"}), - }, - } - - RETURN_TYPES = ("IMAGE", ) - FUNCTION = "vfi" - CATEGORY = "TensorRT_RTX/RIFE" - OUTPUT_NODE=True - - def vfi( - self, - frames, - engine, - clear_cache_after_n_frames=100, - multiplier=2, - use_cuda_graph=True, - keep_model_loaded=False, - ): - """Interpolates frames using the TensorRT engine. - - Args: - frames (torch.Tensor): Input image batch. - engine (Engine): The active TensorRT engine. - clear_cache_after_n_frames (int): Frequency of CUDA cache clearing. - multiplier (int): Target frame rate multiplier. - use_cuda_graph (bool): Whether to use CUDA graph optimization. - keep_model_loaded (bool): If True, keeps engine in VRAM after use. - - Returns: - tuple: Interpolated image batch. - """ - B, H, W, C = frames.shape - shape_dict = { - "img0": {"shape": (1, 3, H, W)}, - "img1": {"shape": (1, 3, H, W)}, - "output": {"shape": (1, 3, H, W)}, - } - - cudaStream = cuda.Stream() - - # Activate and allocate buffers for the engine - engine.activate() - engine.allocate_buffers(shape_dict=shape_dict) - - frames = preprocess_frames(frames) - - def return_middle_frame(frame_0, frame_1, timestep): - timestep_t = torch.tensor([timestep], dtype=torch.float32).to(mm.get_torch_device()) - output = engine.infer({"img0": frame_0, "img1": frame_1, "timestep": timestep_t}, cudaStream, use_cuda_graph) - result = output['output'] - - return result - - result = generate_frames_rife(frames, clear_cache_after_n_frames, multiplier, return_middle_frame) - out = postprocess_frames(result) - - if not keep_model_loaded: - engine.reset() - - return (out,) - - -NODE_CLASS_MAPPINGS = { - "RIFERTXFrameInterpolation": RIFERTXFrameInterpolation, - "RIFERTXEngineLoader": RIFERTXEngineLoader, - "RIFERTXEngineConfig": RIFERTXEngineConfig, -} - -NODE_DISPLAY_NAME_MAPPINGS = { - "RIFERTXFrameInterpolation": "RIFE RTX Frame Interpolation", - "RIFERTXEngineLoader": "RIFE RTX Engine Loader", - "RIFERTXEngineConfig": "RIFE RTX Engine Config", -} \ No newline at end of file diff --git a/pyproject.toml b/pyproject.toml new file mode 100644 index 0000000..c7ce2ae --- /dev/null +++ b/pyproject.toml @@ -0,0 +1,14 @@ +[project] +name = "ComfyUI_RIFE_TensorRT_Auto" +description = "Ultra fast frame interpolation using RIFE with TensorRT with Automatic downloading" +version = "1.1.1" +license = { file = "LICENSE" } + +[project.urls] +Repository = "https://github.com/silveroxides/ComfyUI_RIFE_TensorRT_Auto" +# Used by Comfy Registry https://comfyregistry.org + +[tool.comfy] +PublisherId = "silveroxides" +DisplayName = "ComfyUI_RIFE_TensorRT_Auto" +Icon = "" diff --git a/requirements.txt b/requirements.txt new file mode 100644 index 0000000..a171e1c --- /dev/null +++ b/requirements.txt @@ -0,0 +1,12 @@ +# Default: CUDA 13 (RTX 50 series, CUDA 13.0+) +# For CUDA 12, use requirements_cu12.txt instead +tensorrt_cu13<10.16.0 +tensorrt_cu13_bindings<10.16.0 +tensorrt_cu13_libs<10.16.0 +cuda-toolkit +einops +colored +polygraphy +cuda-python +requests +tqdm \ No newline at end of file diff --git a/requirements.txt.bak b/requirements.txt.bak new file mode 100644 index 0000000..3426b0b --- /dev/null +++ b/requirements.txt.bak @@ -0,0 +1,12 @@ +# Default: CUDA 13 (RTX 50 series, CUDA 13.0+) +# For CUDA 12, use requirements_cu12.txt instead +tensorrt_cu13==10.14.1.48 +tensorrt_cu13_bindings==10.14.1.48 +tensorrt_cu13_libs==10.14.1.48 +cuda-toolkit>=13.0.0,<13.1.0 +einops +colored +polygraphy +cuda-python +requests +tqdm \ No newline at end of file diff --git a/requirements_cu12.txt b/requirements_cu12.txt new file mode 100644 index 0000000..cbf070d --- /dev/null +++ b/requirements_cu12.txt @@ -0,0 +1,12 @@ +# Legacy: CUDA 12 (RTX 30/40 series, CUDA 12.x) +# For CUDA 13, use requirements.txt instead +tensorrt-cu12==10.13.3.9 +tensorrt-cu12-libs==10.13.3.9 +tensorrt-cu12-bindings==10.13.3.9 +cuda-toolkit>=12.8.0,<13.0.0 +einops +colored +polygraphy +cuda-python +requests +tqdm diff --git a/trt_utilities.py b/trt_utilities.py index eb409e2..4eac528 100644 --- a/trt_utilities.py +++ b/trt_utilities.py @@ -1,43 +1,27 @@ -"""TensorRT Acceleration Utilities for RIFE. - -This module provides the low-level interface for NVIDIA TensorRT, handling -engine compilation, timing caches, memory allocation, and optimized inference -using CUDA Graphs. -""" - -import copy -import io -import os - -import cuda.bindings.runtime as cudart -import onnx -import numpy as np -import tensorrt_rtx as trt import torch - +from torch.cuda import nvtx from collections import OrderedDict -from logging import error, warning -from onnxconverter_common import float16 -from onnxsim import simplify - -# Use polygraphy TensorRTX backend -os.environ["POLYGRAPHY_USE_TENSORRT_RTX"] = "1" - +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, - network_from_onnx_bytes, - save_engine, ) + +import modelopt.onnx.autocast as autocast +import onnx + +from onnx import TensorProto + from polygraphy.logger import G_LOGGER - -from torch.cuda import nvtx +import tensorrt as trt +from logging import error, warning from tqdm import tqdm +import copy +import cuda.bindings.runtime as cudart +TRT_LOGGER = trt.Logger(trt.Logger.ERROR) G_LOGGER.module_severity = G_LOGGER.ERROR # Map of numpy dtype -> torch dtype @@ -63,47 +47,8 @@ torch_to_numpy_dtype_dict = { value: key for (key, value) in numpy_to_torch_dtype_dict.items() } -# Path for your persistent hardware profile -node_dir = os.path.dirname(os.path.realpath(__file__)) -MASTER_CACHE_FILE = f"{node_dir}\\RIFE_RTX_ENGINE.cache" - -def get_timing_cache(config): - """Retrieves or creates a TensorRT timing cache. - - Args: - config (trt.IBuilderConfig): The builder configuration object. - - Returns: - trt.ITimingCache: The loaded or new timing cache. - """ - if os.path.exists(MASTER_CACHE_FILE): - with open(MASTER_CACHE_FILE, "rb") as f: - cache_data = f.read() - return config.create_timing_cache(cache_data) - return config.create_timing_cache(b"") - -def save_timing_cache(timing_cache): - """Saves the current timing cache to disk for future builds. - - Args: - timing_cache (trt.ITimingCache): The cache to serialize. - """ - with open(MASTER_CACHE_FILE, "wb") as f: - f.write(timing_cache.serialize()) - +# https://github.com/Jeff-LiangF/streamv2v/blob/18c1a3bd56ff348d54a3300605936980bb13b03c/src/streamv2v/acceleration/tensorrt/utilities.py def CUASSERT(cuda_ret): - """Asserts that a CUDA call was successful. - https://github.com/Jeff-LiangF/streamv2v/blob/18c1a3bd56ff348d54a3300605936980bb13b03c/src/streamv2v/acceleration/tensorrt/utilities.py - - Args: - cuda_ret (tuple): The return value from a cuda-python call. - - Returns: - Any: The result of the CUDA call if successful. - - Raises: - RuntimeError: If the CUDA call returned an error. - """ err = cuda_ret[0] if err != cudart.cudaError_t.cudaSuccess: raise RuntimeError( @@ -112,10 +57,28 @@ def CUASSERT(cuda_ret): if len(cuda_ret) > 1: return cuda_ret[1] return None + + +def is_already_fp16(onnx_path: str) -> bool: + model = onnx.load(onnx_path) + graph = model.graph + + # Check all weight initializers + for initializer in graph.initializer: + if initializer.data_type not in (TensorProto.FLOAT16, TensorProto.INT8, + TensorProto.INT32, TensorProto.INT64, + TensorProto.BOOL): + return False + + # Check graph inputs that have a known element type + for inp in graph.input: + t = inp.type.tensor_type + if t.HasField("elem_type") and t.elem_type == TensorProto.FLOAT: + return False + + return True class TQDMProgressMonitor(trt.IProgressMonitor): - """Progress monitor for TensorRT engine builds using TQDM bars.""" - def __init__(self): trt.IProgressMonitor.__init__(self) self._active_phases = {} @@ -123,7 +86,6 @@ class TQDMProgressMonitor(trt.IProgressMonitor): self.max_indent = 5 def phase_start(self, phase_name, parent_phase, num_steps): - """Called when a build phase starts.""" leave = False try: if parent_phase is not None: @@ -150,7 +112,6 @@ class TQDMProgressMonitor(trt.IProgressMonitor): _step_result = False def phase_finish(self, phase_name): - """Called when a build phase finishes.""" try: if phase_name in self._active_phases.keys(): self._active_phases[phase_name]["tq"].update( @@ -177,7 +138,6 @@ class TQDMProgressMonitor(trt.IProgressMonitor): _step_result = False def step_complete(self, phase_name, step): - """Called when a step within a phase completes.""" try: if phase_name in self._active_phases.keys(): self._active_phases[phase_name]["tq"].update( @@ -189,15 +149,6 @@ class TQDMProgressMonitor(trt.IProgressMonitor): return False class Engine: - """Wrapper for TensorRT engine management. - - Handles loading, building, and executing inference with optional - CUDA Graph support. - - Attributes: - engine_path (str): File path to the .trt engine. - """ - def __init__( self, engine_path, @@ -211,7 +162,7 @@ class Engine: self.graph = None def __del__(self): - """Ensures CUDA resources are freed on deletion.""" + # Clean up CUDA graph resources if hasattr(self, 'cuda_graph_instance') and self.cuda_graph_instance is not None: try: cudart.cudaGraphDestroy(self.cuda_graph_instance) @@ -229,11 +180,7 @@ class Engine: del self.tensors def reset(self, engine_path=None): - """Cleans up engine and CUDA resources. - - Args: - engine_path (str, optional): New path to set for the engine. - """ + # Clean up CUDA graph resources first if hasattr(self, 'cuda_graph_instance') and self.cuda_graph_instance is not None: try: cudart.cudaGraphDestroy(self.cuda_graph_instance) @@ -268,6 +215,7 @@ class Engine: def build( self, onnx_path, + fp16, input_profile=None, enable_refit=False, enable_preview=False, @@ -275,16 +223,39 @@ class Engine: timing_cache=None, update_output_names=None, ): - """Compiles an ONNX model into a TensorRT engine. + print(f"Building TensorRT engine for {onnx_path}: {self.engine_path}") + + if not is_already_fp16(onnx_path): + print(f"File not FP16... Converting") + converted = autocast.convert_to_mixed_precision( + onnx_path=onnx_path, + low_precision_type="fp16", + keep_io_types=True, + ) + onnx.save(converted, onnx_path) # overwrite in-place + + + logger = trt.Logger(trt.Logger.WARNING) + builder = trt.Builder(logger) + network = builder.create_network() # strong typing is always on in 11.x + config = builder.create_builder_config() + + config.progress_monitor = TQDMProgressMonitor() + + if enable_refit: + config.set_flag(trt.BuilderFlag.REFIT) + + parser = trt.OnnxParser(network, logger) + with open(onnx_path, "rb") as f: + if not parser.parse(f.read()): + for i in range(parser.num_errors): + print(parser.get_error(i)) + raise RuntimeError("ONNX parse failed") + + if update_output_names: + print(f"Updating network outputs to {update_output_names}") + network = ModifyNetworkOutputs(network, update_output_names) - Args: - onnx_path (str): Path to source ONNX file. - input_profile (list): List of dicts defining min/opt/max shapes. - **kwargs: Configuration flags (enable_refit, etc). - - Returns: - int: 0 if success, 1 if failure. - """ p = [Profile()] if input_profile: p = [Profile() for i in range(len(input_profile))] @@ -292,79 +263,30 @@ class Engine: 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"] = [] - - # Load ONNX file - onnx_bytes = onnx.load(onnx_path) - onnx_precision = onnx_bytes.graph.input[0].type.tensor_type.elem_type - onnx_simp, check = simplify(onnx_bytes) - - if not check: - error(f"Simplified ONNX model: {onnx_path} could not be validated") - return 1 - - # Keep simplified model in memory to not overwrite the onnx file in models directory - # We dont want to be destructive - onnx_buffer = io.BytesIO() - onnx_buffer.write(onnx_simp.SerializeToString()) - onnx_buffer.seek(0) - onnx_buffer.name = "onnx_simp_mem_rife.onnx" - - #network = network_from_onnx_path(onnx_path, flags=[trt.OnnxParserFlag.NATIVE_INSTANCENORM]) - network = network_from_onnx_bytes(onnx_buffer.getvalue(), 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 - - # Limit workspace (e.g., to 2GB or 4GB depending on your GPU) - # 1 << 30 is 1GB - #config.set_memory_pool_limit(trt.MemoryPoolType.WORKSPACE, 4 * (1 << 30)) - config.builder_optimization_level = 5 - - # Use Master Cache - timing_cache = get_timing_cache(config) - config.set_timing_cache(timing_cache, False) - + 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] - ) + calib_profile = profile.fill_defaults(network).to_trt(builder, network) config.add_optimization_profile(calib_profile) - - result = 0 - + try: - engine = engine_from_network(network, config) - save_engine(engine, path=self.engine_path) + engine = builder.build_serialized_network(network, config) except Exception as e: - error(f"Failed to build engine: {e}") - result = 1 - - # Update Master Cache with new findings - save_timing_cache(timing_cache) - print(f"Updated Master Cache.") - - return result + error(f"Failed to build engine: {e}") + return 1 + try: + with open(self.engine_path, "wb") as f: + f.write(bytes(engine)) + except Exception as e: + error(f"Failed to save engine: {e}") + return 1 + return 0 def load(self): - """Loads the serialized engine from disk.""" self.engine = engine_from_bytes(bytes_from_path(self.engine_path)) def activate(self, reuse_device_memory=None): - """Initializes the execution context.""" + # If engine was reset, reload it if self.engine is None: self.load() @@ -375,12 +297,7 @@ class Engine: self.context = self.engine.create_execution_context() def allocate_buffers(self, shape_dict=None, device="cuda"): - """Allocates GPU memory for input and output tensors. - - Args: - shape_dict (dict, optional): Specific shapes for tensors. - device (str): Device to allocate on (default "cuda"). - """ + # Clean up CUDA graph resources since tensors will be recreated if hasattr(self, 'cuda_graph_instance') and self.cuda_graph_instance is not None: try: cudart.cudaGraphDestroy(self.cuda_graph_instance) @@ -413,16 +330,6 @@ class Engine: nvtx.range_pop() def infer(self, feed_dict, stream, use_cuda_graph=False): - """Performs inference on the provided data. - - Args: - feed_dict (dict): Input tensors. - stream (cuda.Stream): CUDA stream for execution. - use_cuda_graph (bool): If True, uses/captures CUDA graphs. - - Returns: - OrderedDict: The output tensors. - """ for name, buf in feed_dict.items(): self.tensors[name].copy_(buf) diff --git a/trt_utilities.py.bak b/trt_utilities.py.bak new file mode 100644 index 0000000..c20f56a --- /dev/null +++ b/trt_utilities.py.bak @@ -0,0 +1,364 @@ +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, +) + +import modelopt.onnx.autocast as autocast +import onnx + +from onnx import TensorProto + +from polygraphy.logger import G_LOGGER +import tensorrt as trt +from logging import error, warning +from tqdm import tqdm +import copy +import cuda.bindings.runtime as cudart + +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() +} + +# https://github.com/Jeff-LiangF/streamv2v/blob/18c1a3bd56ff348d54a3300605936980bb13b03c/src/streamv2v/acceleration/tensorrt/utilities.py +def CUASSERT(cuda_ret): + err = cuda_ret[0] + if err != cudart.cudaError_t.cudaSuccess: + raise RuntimeError( + f"CUDA ERROR: {err}, error code reference: https://nvidia.github.io/cuda-python/module/cudart.html#cuda.cudart.cudaError_t" + ) + if len(cuda_ret) > 1: + return cuda_ret[1] + return None + + +def is_already_fp16(onnx_path: str) -> bool: + model = onnx.load(onnx_path) + graph = model.graph + + # Check all weight initializers + for initializer in graph.initializer: + if initializer.data_type not in (TensorProto.FLOAT16, TensorProto.INT8, + TensorProto.INT32, TensorProto.INT64, + TensorProto.BOOL): + return False + + # Check graph inputs that have a known element type + for inp in graph.input: + t = inp.type.tensor_type + if t.HasField("elem_type") and t.elem_type == TensorProto.FLOAT: + return False + + return True + +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 + self.graph = None + + def __del__(self): + # Clean up CUDA graph resources + if hasattr(self, 'cuda_graph_instance') and self.cuda_graph_instance is not None: + try: + cudart.cudaGraphDestroy(self.cuda_graph_instance) + except: + pass + if hasattr(self, 'graph') and self.graph is not None: + try: + cudart.cudaGraphDestroy(self.graph) + except: + pass + + del self.engine + del self.context + del self.buffers + del self.tensors + + def reset(self, engine_path=None): + # Clean up CUDA graph resources first + if hasattr(self, 'cuda_graph_instance') and self.cuda_graph_instance is not None: + try: + cudart.cudaGraphDestroy(self.cuda_graph_instance) + except: + pass + self.cuda_graph_instance = None + if hasattr(self, 'graph') and self.graph is not None: + try: + cudart.cudaGraphDestroy(self.graph) + except: + pass + self.graph = None + + if hasattr(self, 'engine') and self.engine is not None: + del self.engine + if hasattr(self, 'context') and self.context is not None: + del self.context + if hasattr(self, 'buffers'): + del self.buffers + if hasattr(self, 'tensors'): + del self.tensors + + self.engine = None + self.context = None + self.engine_path = engine_path if engine_path else self.engine_path + + 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, + ): + print(f"Building TensorRT engine for {onnx_path}: {self.engine_path}") + + if not is_already_fp16(onnx_path): + print(f"File not FP16... Converting") + converted = autocast.convert_to_mixed_precision( + onnx_path=onnx_path, + low_precision_type="fp16", + keep_io_types=True, + ) + onnx.save(converted, onnx_path) # overwrite in-place + + + logger = trt.Logger(trt.Logger.WARNING) + builder = trt.Builder(logger) + network = builder.create_network() # strong typing is always on in 11.x + config = builder.create_builder_config() + + config.progress_monitor = TQDMProgressMonitor() + + if enable_refit: + config.set_flag(trt.BuilderFlag.REFIT) + + parser = trt.OnnxParser(network, logger) + with open(onnx_path, "rb") as f: + if not parser.parse(f.read()): + for i in range(parser.num_errors): + print(parser.get_error(i)) + raise RuntimeError("ONNX parse failed") + + if update_output_names: + print(f"Updating network outputs to {update_output_names}") + network = ModifyNetworkOutputs(network, update_output_names) + + 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]) + + profiles = copy.deepcopy(p) + for profile in profiles: + calib_profile = profile.fill_defaults(network).to_trt(builder, network) + config.add_optimization_profile(calib_profile) + + try: + engine = builder.build_serialized_network(network, config) + except Exception as e: + error(f"Failed to build engine: {e}") + return 1 + try: + with open(self.engine_path, "wb") as f: + f.write(bytes(engine)) + 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 engine was reset, reload it + if self.engine is None: + self.load() + + 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"): + # Clean up CUDA graph resources since tensors will be recreated + if hasattr(self, 'cuda_graph_instance') and self.cuda_graph_instance is not None: + try: + cudart.cudaGraphDestroy(self.cuda_graph_instance) + except: + pass + self.cuda_graph_instance = None + if hasattr(self, 'graph') and self.graph is not None: + try: + cudart.cudaGraphDestroy(self.graph) + except: + pass + self.graph = None + + 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): + 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()) + + if use_cuda_graph: + if self.cuda_graph_instance is not None: + CUASSERT(cudart.cudaGraphLaunch(self.cuda_graph_instance, stream.ptr)) + CUASSERT(cudart.cudaStreamSynchronize(stream.ptr)) + else: + # do inference before CUDA graph capture + noerror = self.context.execute_async_v3(stream.ptr) + if not noerror: + raise ValueError("ERROR: inference failed.") + # capture cuda graph + CUASSERT( + cudart.cudaStreamBeginCapture(stream.ptr, cudart.cudaStreamCaptureMode.cudaStreamCaptureModeGlobal) + ) + self.context.execute_async_v3(stream.ptr) + self.graph = CUASSERT(cudart.cudaStreamEndCapture(stream.ptr)) + self.cuda_graph_instance = CUASSERT(cudart.cudaGraphInstantiate(self.graph, 0)) + else: + noerror = self.context.execute_async_v3(stream.ptr) + if not noerror: + raise ValueError("ERROR: inference failed.") + + return self.tensors + diff --git a/utilities.py b/utilities.py new file mode 100644 index 0000000..fd4a248 --- /dev/null +++ b/utilities.py @@ -0,0 +1,101 @@ +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) \ No newline at end of file diff --git a/vfi_utilities.py b/vfi_utilities.py index 69ef2f0..0acdb6b 100644 --- a/vfi_utilities.py +++ b/vfi_utilities.py @@ -1,50 +1,37 @@ -"""Utility functions for RIFE Video Frame Interpolation (VFI). +# https://github.com/Fannovel16/ComfyUI-Frame-Interpolation/blob/main/vfi_utils.py -This module contains preprocessing, postprocessing, and the core interpolation -loop logic used to generate intermediate frames between a sequence of images. - -https://github.com/Fannovel16/ComfyUI-Frame-Interpolation/blob/main/vfi_utils.py -""" - -import einops import os import torch import typing - -import logging -import numpy as np - +import einops from comfy.model_management import soft_empty_cache, get_torch_device +import numpy as np from comfy.utils import ProgressBar +from colored import Fore, Back, Style DEVICE = get_torch_device() -def preprocess_frames(frames: torch.Tensor) -> torch.Tensor: - """Prepares ComfyUI image batches for model inference. +def load_file_from_github_release(model_type, ckpt_name): + error_strs = [] + for i, base_model_download_url in enumerate(BASE_MODEL_DOWNLOAD_URLS): + try: + return load_file_from_url(base_model_download_url + ckpt_name, get_ckpt_container_path(model_type)) + except Exception: + traceback_str = traceback.format_exc() + if i < len(BASE_MODEL_DOWNLOAD_URLS) - 1: + print("Failed! Trying another endpoint.") + error_strs.append(f"Error when downloading from: {base_model_download_url + ckpt_name}\n\n{traceback_str}") - Converts images from (Batch, Height, Width, Channels) to - (Batch, Channels, Height, Width) and ensures only RGB channels are present. + error_str = '\n\n'.join(error_strs) + raise Exception(f"Tried all GitHub base urls to download {ckpt_name} but no suceess. Below is the error log:\n\n{error_str}") - Args: - frames (torch.Tensor): Input tensor in BHWC format. +def logger(msg): + print(f'{Style.reset}{Fore.cyan}⚡ [Rife Tensorrt] - {msg}{Style.reset}') - Returns: - torch.Tensor: Preprocessed tensor in BCHW format. - """ +def preprocess_frames(frames): return einops.rearrange(frames[..., :3], "n h w c -> n c h w") -def postprocess_frames(frames: torch.Tensor) -> torch.Tensor: - """Converts model output back to ComfyUI standard format. - - Converts images from (Batch, Channels, Height, Width) back to - (Batch, Height, Width, Channels) and moves them to CPU memory. - - Args: - frames (torch.Tensor): Output tensor from the model in BCHW format. - - Returns: - torch.Tensor: Postprocessed tensor in BHWC format on CPU. - """ +def postprocess_frames(frames): return einops.rearrange(frames, "n c h w -> n h w c")[..., :3].cpu() def generate_frames_rife( @@ -53,40 +40,20 @@ def generate_frames_rife( multiplier, return_middle_frame_function ): - """Core interpolation loop for generating intermediate frames. - Iterates through a sequence of frames and calls the provided inference - function to generate 'n' intermediate frames based on the multiplier. - - Args: - frames (torch.Tensor): The preprocessed frame sequence (BCHW). - clear_cache_after_n_frames (int): Interval for VRAM cache clearing. - multiplier (int): The factor by which to increase the frame count. - return_middle_frame_function (Callable): Function that takes two - frames and a timestep and returns the interpolated frame. - - Returns: - torch.Tensor: The complete sequence including original and - interpolated frames. - """ output_frames = torch.zeros(multiplier*frames.shape[0], *frames.shape[1:], device="cpu") out_len = 0 - cache_counter = 0 - - pbar = ProgressBar(len(frames)) - - logging.info(f"Interpolating frames") - for frame_itr in range(len(frames) - 1): + number_of_frames_processed_since_last_cleared_cuda_cache = 0 + # pbar = ProgressBar(len(frames)) + + for frame_itr in range(len(frames) - 1): # Skip the final frame since there are no frames after it frame_0 = frames[frame_itr:frame_itr+1] frame_1 = frames[frame_itr+1:frame_itr+2] - - # Store the current original frame - output_frames[out_len] = frame_0 + output_frames[out_len] = frame_0 # Start with first frame out_len += 1 - # Generate intermediate frames for middle_i in range(1, multiplier): timestep = middle_i/multiplier middle_frame = return_middle_frame_function(frame_0, frame_1, timestep).detach().cpu() @@ -95,22 +62,25 @@ def generate_frames_rife( output_frames[out_len] = middle_frame out_len +=1 - # VRAM Management - cache_counter += 1 - if cache_counter >= clear_cache_after_n_frames: + # Try to avoid a memory overflow by clearing cuda cache regularly + number_of_frames_processed_since_last_cleared_cuda_cache += 1 + if number_of_frames_processed_since_last_cleared_cuda_cache >= clear_cache_after_n_frames: soft_empty_cache() - logging.info(f"Clearing cache...") - cache_counter = 0 + number_of_frames_processed_since_last_cleared_cuda_cache = 0 + logger("Clearing cache...") + + # pbar.update(1) - pbar.update(1) # Append final frame output_frames[out_len] = frames[-1:] - logging.info(f"done! - {(len(frames) -1) * (multiplier-1)} new frames generated at resolution: {output_frames[0].shape}") + logger(f"done! - {(len(frames) -1) * (multiplier-1)} new frames generated at resolution: {output_frames[0].shape}") out_len += 1 # clear cache for courtesy soft_empty_cache() - - # Slice to actual length in case of rounding or logic offsets - return output_frames[:out_len] + logger("Final clearing cache done ...") +# + res = output_frames[:out_len] + + return res