commit cdd8db862505806ddb08d87953bc109b9dc4cb82 Author: wehub-resource-sync Date: Mon Jul 13 13:37:51 2026 +0800 chore: import upstream snapshot with attribution diff --git a/.github/ISSUE_TEMPLATE/PULL_REQUEST_TEMPLATE.md b/.github/ISSUE_TEMPLATE/PULL_REQUEST_TEMPLATE.md new file mode 100644 index 0000000..65d04d3 --- /dev/null +++ b/.github/ISSUE_TEMPLATE/PULL_REQUEST_TEMPLATE.md @@ -0,0 +1,28 @@ +# Contribution Guide + +We welcome your contributions to this repository. To ensure elegant code style and better code quality, we have prepared the following contribution guidelines. + +## What We Accept + ++ This PR fixes a typo or improves the documentation (if this is the case, you may skip the other checks). ++ This PR fixes a specific issue — please reference the issue number in the PR description. Make sure your code strictly follows the coding standards below. ++ This PR introduces a new feature — please clearly explain the necessity and implementation of the feature. Make sure your code strictly follows the coding standards below. + +## Code Style Guide + +Good code style is an art. We have prepared a `pyproject.toml` and a `pre-commit` hook to enforce consistent code formatting across the project. You can clean up your code following the steps below: + +1. Install the required dependencies: +```shell + pip install ruff pre-commit +``` +2. Then, run the following command: +```shell + pre-commit run --all-files +``` +If your code complies with the standards, you should not see any errors. + +## Naming Conventions + +- Please use **English** for naming; do not use Pinyin or other languages. All comments should also be in English. +- Follow **PEP8** naming conventions strictly, and use underscores to separate words. Avoid meaningless names such as `a`, `b`, `c`. diff --git a/.github/ISSUE_TEMPLATE/bug_report.yaml b/.github/ISSUE_TEMPLATE/bug_report.yaml new file mode 100644 index 0000000..4f1f0dc --- /dev/null +++ b/.github/ISSUE_TEMPLATE/bug_report.yaml @@ -0,0 +1,51 @@ +name: "\U0001F41B Bug Report" +description: Submit a bug report to help us improve CogVideoX / 提交一个 Bug 问题报告来帮助我们改进 CogVideoX 开源模型 +body: + - type: textarea + id: system-info + attributes: + label: System Info / 系統信息 + description: Your operating environment / 您的运行环境信息 + placeholder: Includes Cuda version, Diffusers version, Python version, operating system, hardware information (if you suspect a hardware problem)... / 包括Cuda版本，Diffusers，Python版本，操作系统，硬件信息(如果您怀疑是硬件方面的问题)... + validations: + required: true + + - type: checkboxes + id: information-scripts-examples + attributes: + label: Information / 问题信息 + description: 'The problem arises when using: / 问题出现在' + options: + - label: "The official example scripts / 官方的示例脚本" + - label: "My own modified scripts / 我自己修改的脚本和任务" + + - type: textarea + id: reproduction + validations: + required: true + attributes: + label: Reproduction / 复现过程 + description: | + Please provide a code example that reproduces the problem you encountered, preferably with a minimal reproduction unit. + If you have code snippets, error messages, stack traces, please provide them here as well. + Please format your code correctly using code tags. See https://help.github.com/en/github/writing-on-github/creating-and-highlighting-code-blocks#syntax-highlighting + Do not use screenshots, as they are difficult to read and (more importantly) do not allow others to copy and paste your code. + + 请提供能重现您遇到的问题的代码示例,最好是最小复现单元。 + 如果您有代码片段、错误信息、堆栈跟踪，也请在此提供。 + 请使用代码标签正确格式化您的代码。请参见 https://help.github.com/en/github/writing-on-github/creating-and-highlighting-code-blocks#syntax-highlighting + 请勿使用截图，因为截图难以阅读，而且（更重要的是）不允许他人复制粘贴您的代码。 + placeholder: | + Steps to reproduce the behavior/复现Bug的步骤: + + 1. + 2. + 3. + + - type: textarea + id: expected-behavior + validations: + required: true + attributes: + label: Expected behavior / 期待表现 + description: "A clear and concise description of what you would expect to happen. /简单描述您期望发生的事情。" diff --git a/.github/ISSUE_TEMPLATE/feature-request.yaml b/.github/ISSUE_TEMPLATE/feature-request.yaml new file mode 100644 index 0000000..513e96c --- /dev/null +++ b/.github/ISSUE_TEMPLATE/feature-request.yaml @@ -0,0 +1,34 @@ +name: "\U0001F680 Feature request" +description: Submit a request for a new CogVideoX feature / 提交一个新的 CogVideoX开源模型的功能建议 +labels: [ "feature" ] +body: + - type: textarea + id: feature-request + validations: + required: true + attributes: + label: Feature request / 功能建议 + description: | + A brief description of the functional proposal. Links to corresponding papers and code are desirable. + 对功能建议的简述。最好提供对应的论文和代码链接。 + + - type: textarea + id: motivation + validations: + required: true + attributes: + label: Motivation / 动机 + description: | + Your motivation for making the suggestion. If that motivation is related to another GitHub issue, link to it here. + 您提出建议的动机。如果该动机与另一个 GitHub 问题有关，请在此处提供对应的链接。 + + - type: textarea + id: contribution + validations: + required: true + attributes: + label: Your contribution / 您的贡献 + description: | + + Your PR link or any other link you can help with. + 您的PR链接或者其他您能提供帮助的链接。 diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml new file mode 100644 index 0000000..fff4c90 --- /dev/null +++ b/.pre-commit-config.yaml @@ -0,0 +1,19 @@ +repos: + - repo: https://github.com/astral-sh/ruff-pre-commit + rev: v0.4.5 + hooks: + - id: ruff + args: [--fix, --respect-gitignore, --config=pyproject.toml] + - id: ruff-format + args: [--config=pyproject.toml] + + - repo: https://github.com/pre-commit/pre-commit-hooks + rev: v4.5.0 + hooks: + - id: trailing-whitespace + - id: end-of-file-fixer + - id: check-yaml + - id: check-toml + - id: check-case-conflict + - id: check-merge-conflict + - id: debug-statements diff --git a/LICENSE b/LICENSE new file mode 100644 index 0000000..f98e413 --- /dev/null +++ b/LICENSE @@ -0,0 +1,201 @@ + Apache License + Version 2.0, January 2004 + http://www.apache.org/licenses/ + + TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION + + 1. Definitions. + + "License" shall mean the terms and conditions for use, reproduction, + and distribution as defined by Sections 1 through 9 of this document. + + "Licensor" shall mean the copyright owner or entity authorized by + the copyright owner that is granting the License. + + "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. + + "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. + + "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. + + "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. + + "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." + + "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. + + 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. + + 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. + + 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: + + (a) You must give any other recipients of the Work or + Derivative Works a copy of this License; and + + (b) You must cause any modified files to carry prominent notices + stating that You changed the files; and + + (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 + + (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. + + 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. + + 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. + + 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. + + 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. + + 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. + + 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 + + APPENDIX: How to apply the Apache License to your work. + + 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. + + Copyright 2024 CogVideo Model Team @ Zhipu AI + + Licensed under the Apache License, Version 2.0 (the "License"); + you may not use this file except in compliance with the License. + You may obtain a copy of the License at + + http://www.apache.org/licenses/LICENSE-2.0 + + Unless required by applicable law or agreed to in writing, software + distributed under the License is distributed on an "AS IS" BASIS, + WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + See the License for the specific language governing permissions and + limitations under the License. diff --git a/MODEL_LICENSE b/MODEL_LICENSE new file mode 100644 index 0000000..7f8da3b --- /dev/null +++ b/MODEL_LICENSE @@ -0,0 +1,71 @@ +The CogVideoX License + +1. Definitions + +“Licensor” means the CogVideoX Model Team that distributes its Software. + +“Software” means the CogVideoX model parameters made available under this license. + +2. License Grant + +Under the terms and conditions of this license, the licensor hereby grants you a non-exclusive, worldwide, non-transferable, non-sublicensable, revocable, royalty-free copyright license. The intellectual property rights of the generated content belong to the user to the extent permitted by applicable local laws. +This license allows you to freely use all open-source models in this repository for academic research. Users who wish to use the models for commercial purposes must register and obtain a basic commercial license in https://open.bigmodel.cn/mla/form . +Users who have registered and obtained the basic commercial license can use the models for commercial activities for free, but must comply with all terms and conditions of this license. Additionally, the number of service users (visits) for your commercial activities must not exceed 1 million visits per month. +If the number of service users (visits) for your commercial activities exceeds 1 million visits per month, you need to contact our business team to obtain more commercial licenses. +The above copyright statement and this license statement should be included in all copies or significant portions of this software. + +3. Restriction + +You will not use, copy, modify, merge, publish, distribute, reproduce, or create derivative works of the Software, in whole or in part, for any military, or illegal purposes. + +You will not use the Software for any act that may undermine China's national security and national unity, harm the public interest of society, or infringe upon the rights and interests of human beings. + +4. Disclaimer + +THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE. + +5. Limitation of Liability + +EXCEPT TO THE EXTENT PROHIBITED BY APPLICABLE LAW, IN NO EVENT AND UNDER NO LEGAL THEORY, WHETHER BASED IN TORT, NEGLIGENCE, CONTRACT, LIABILITY, OR OTHERWISE WILL ANY LICENSOR BE LIABLE TO YOU FOR ANY DIRECT, INDIRECT, SPECIAL, INCIDENTAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES, OR ANY OTHER COMMERCIAL LOSSES, EVEN IF THE LICENSOR HAS BEEN ADVISED OF THE POSSIBILITY OF SUCH DAMAGES. + +6. Dispute Resolution + +This license shall be governed and construed in accordance with the laws of People’s Republic of China. Any dispute arising from or in connection with this License shall be submitted to Haidian District People's Court in Beijing. + +Note that the license is subject to update to a more comprehensive version. For any questions related to the license and copyright, please contact us at license@zhipuai.cn. + +1. 定义 + +“许可方”是指分发其软件的 CogVideoX 模型团队。 + +“软件”是指根据本许可提供的 CogVideoX 模型参数。 + +2. 许可授予 + +根据本许可的条款和条件，许可方特此授予您非排他性、全球性、不可转让、不可再许可、可撤销、免版税的版权许可。生成内容的知识产权所属，可根据适用当地法律的规定，在法律允许的范围内由用户享有生成内容的知识产权或其他权利。 +本许可允许您免费使用本仓库中的所有开源模型进行学术研究。对于希望将模型用于商业目的的用户，需在 https://open.bigmodel.cn/mla/form 完成登记并获得基础商用授权。 + +经过登记并获得基础商用授权的用户可以免费使用本模型进行商业活动，但必须遵守本许可的所有条款和条件。 +在本许可证下，您的商业活动的服务用户数量（访问量）不得超过100万人次访问 / 每月。如果超过，您需要与我们的商业团队联系以获得更多的商业许可。 +上述版权声明和本许可声明应包含在本软件的所有副本或重要部分中。 + +3.限制 + +您不得出于任何军事或非法目的使用、复制、修改、合并、发布、分发、复制或创建本软件的全部或部分衍生作品。 + +您不得利用本软件从事任何危害国家安全和国家统一、危害社会公共利益、侵犯人身权益的行为。 + +4.免责声明 + +本软件“按原样”提供，不提供任何明示或暗示的保证，包括但不限于对适销性、特定用途的适用性和非侵权性的保证。 +在任何情况下，作者或版权持有人均不对任何索赔、损害或其他责任负责，无论是在合同诉讼、侵权行为还是其他方面，由软件或软件的使用或其他交易引起、由软件引起或与之相关 软件。 + +5. 责任限制 + +除适用法律禁止的范围外，在任何情况下且根据任何法律理论，无论是基于侵权行为、疏忽、合同、责任或其他原因，任何许可方均不对您承担任何直接、间接、特殊、偶然、示范性、 或间接损害，或任何其他商业损失，即使许可人已被告知此类损害的可能性。 + +6.争议解决 + +本许可受中华人民共和国法律管辖并按其解释。 因本许可引起的或与本许可有关的任何争议应提交北京市海淀区人民法院。 + +请注意，许可证可能会更新到更全面的版本。 有关许可和版权的任何问题，请通过 license@zhipuai.cn 与我们联系。 diff --git a/README.md b/README.md new file mode 100644 index 0000000..45440ca --- /dev/null +++ b/README.md @@ -0,0 +1,458 @@ +# CogVideo & CogVideoX + +[中文阅读](./README_zh.md) + +[日本語で読む](./README_ja.md) + +
+ +
+

+Experience the CogVideoX-5B model online at 🤗 Huggingface Space or 🤖 ModelScope Space +

+

+📚 View the paper and user guide +

+

+ 👋 Join our WeChat and Discord +

+

+📍 Visit QingYing and API Platform to experience larger-scale commercial video generation models. +

+ +## Project Updates + +- 🔥🔥 **News**: ```2025/03/24```: We have launched [CogKit](https://github.com/THUDM/CogKit), a fine-tuning and inference framework for the **CogView4** and **CogVideoX** series. This toolkit allows you to fully explore and utilize our multimodal generation models. +- 🔥 **News**: ```2025/02/28```: DDIM Inverse is now supported in `CogVideoX-5B` and `CogVideoX1.5-5B`. Check [here](inference/ddim_inversion.py). +- 🔥 **News**: ```2025/01/08```: We have updated the code for `Lora` fine-tuning based on the `diffusers` version model, which uses less GPU memory. For more details, please see [here](finetune/README.md). +- 🔥 **News**: ```2024/11/15```: We released the `CogVideoX1.5` model in the diffusers version. Only minor parameter adjustments are needed to continue using previous code. +- 🔥 **News**: ```2024/11/08```: We have released the CogVideoX1.5 model. CogVideoX1.5 is an upgraded version of the open-source model CogVideoX. +The CogVideoX1.5-5B series supports 10-second videos with higher resolution, and CogVideoX1.5-5B-I2V supports video generation at any resolution. +The SAT code has already been updated, while the diffusers version is still under adaptation. Download the SAT version code [here](https://huggingface.co/THUDM/CogVideoX1.5-5B-SAT). +- 🔥 **News**: ```2024/10/13```: A more cost-effective fine-tuning framework for `CogVideoX-5B` that works with a single + 4090 GPU, [cogvideox-factory](https://github.com/a-r-r-o-w/cogvideox-factory), has been released. It supports + fine-tuning with multiple resolutions. Feel free to use it! +- 🔥 **News**: ```2024/10/10```: We have updated our technical report. Please + click [here](https://arxiv.org/pdf/2408.06072) to view it. More training details and a demo have been added. To see + the demo, click [here](https://yzy-thu.github.io/CogVideoX-demo/).- 🔥 **News**: ```2024/10/09```: We have publicly + released the [technical documentation](https://zhipu-ai.feishu.cn/wiki/DHCjw1TrJiTyeukfc9RceoSRnCh) for CogVideoX + fine-tuning on Feishu, further increasing distribution flexibility. All examples in the public documentation can be + fully reproduced. +- 🔥 **News**: ```2024/9/19```: We have open-sourced the CogVideoX series image-to-video model **CogVideoX-5B-I2V**. + This model can take an image as a background input and generate a video combined with prompt words, offering greater + controllability. With this, the CogVideoX series models now support three tasks: text-to-video generation, video + continuation, and image-to-video generation. Welcome to try it online + at [Experience](https://huggingface.co/spaces/THUDM/CogVideoX-5B-Space). +- 🔥 ```2024/9/19```: The Caption + model [CogVLM2-Caption](https://huggingface.co/THUDM/cogvlm2-llama3-caption), used in the training process of + CogVideoX to convert video data into text descriptions, has been open-sourced. Welcome to download and use it. +- 🔥 ```2024/8/27```: We have open-sourced a larger model in the CogVideoX series, **CogVideoX-5B**. We have + significantly optimized the model's inference performance, greatly lowering the inference threshold. + You can run **CogVideoX-2B** on older GPUs like `GTX 1080TI`, and **CogVideoX-5B** on desktop GPUs like `RTX 3060`. Please strictly + follow the [requirements](requirements.txt) to update and install dependencies, and refer + to [cli_demo](inference/cli_demo.py) for inference code. Additionally, the open-source license for + the **CogVideoX-2B** model has been changed to the **Apache 2.0 License**. +- 🔥 ```2024/8/6```: We have open-sourced **3D Causal VAE**, used for **CogVideoX-2B**, which can reconstruct videos with + almost no loss. +- 🔥 ```2024/8/6```: We have open-sourced the first model of the CogVideoX series video generation models, **CogVideoX-2B + **. +- 🌱 **Source**: ```2022/5/19```: We have open-sourced the CogVideo video generation model (now you can see it in + the `CogVideo` branch). This is the first open-source large Transformer-based text-to-video generation model. You can + access the [ICLR'23 paper](https://arxiv.org/abs/2205.15868) for technical details. + +## Table of Contents + +Jump to a specific section: + +- [Quick Start](#quick-start) + - [Prompt Optimization](#prompt-optimization) + - [SAT](#sat) + - [Diffusers](#diffusers) +- [Gallery](#gallery) + - [CogVideoX-5B](#cogvideox-5b) + - [CogVideoX-2B](#cogvideox-2b) +- [Model Introduction](#model-introduction) +- [Friendly Links](#friendly-links) +- [Project Structure](#project-structure) + - [Quick Start with Colab](#quick-start-with-colab) + - [Inference](#inference) + - [finetune](#finetune) + - [sat](#sat-1) + - [Tools](#tools) +- [CogVideo(ICLR'23)](#cogvideoiclr23) +- [Citation](#citation) +- [Model-License](#model-license) + +## Quick Start + +### Prompt Optimization + +Before running the model, please refer to [this guide](inference/convert_demo.py) to see how we use large models like +GLM-4 (or other comparable products, such as GPT-4) to optimize the model. This is crucial because the model is trained +with long prompts, and a good prompt directly impacts the quality of the video generation. + +### SAT + +**Please make sure your Python version is between 3.10 and 3.12, inclusive of both 3.10 and 3.12.** + +Follow instructions in [sat_demo](sat/README.md): Contains the inference code and fine-tuning code of SAT weights. It is +recommended to improve based on the CogVideoX model structure. Innovative researchers use this code to better perform +rapid stacking and development. + +### Diffusers + +**Please make sure your Python version is between 3.10 and 3.12, inclusive of both 3.10 and 3.12.** + +``` +pip install -r requirements.txt +``` + +Then follow [diffusers_demo](inference/cli_demo.py): A more detailed explanation of the inference code, mentioning the +significance of common parameters. + +For more details on quantized inference, please refer +to [diffusers-torchao](https://github.com/sayakpaul/diffusers-torchao/). With Diffusers and TorchAO, quantized inference +is also possible leading to memory-efficient inference as well as speedup in some cases when compiled. A full list of +memory and time benchmarks with various settings on A100 and H100 has been published +at [diffusers-torchao](https://github.com/sayakpaul/diffusers-torchao). + +## Gallery + +### CogVideoX-5B + + + + + + + + + + + + + + +
+ + + + + + + +
+ + + + + + + +
+ +### CogVideoX-2B + + + + + + + + +
+ + + + + + + +
+ +To view the corresponding prompt words for the gallery, please click [here](resources/galary_prompt.md) + +## Model Introduction + +CogVideoX is an open-source version of the video generation model originating +from [QingYing](https://chatglm.cn/video?lang=en?fr=osm_cogvideo). The table below displays the list of video generation +models we currently offer, along with their foundational information. + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + +
Model NameCogVideoX1.5-5B (Latest)CogVideoX1.5-5B-I2V (Latest)CogVideoX-2BCogVideoX-5BCogVideoX-5B-I2V
Release DateNovember 8, 2024November 8, 2024August 6, 2024August 27, 2024September 19, 2024
Video Resolution1360 * 768 Min(W, H) = 768
768 ≤ Max(W, H) ≤ 1360
Max(W, H) % 16 = 0 		720 * 480
Number of FramesShould be 16N + 1 where N <= 10 (default 81)Should be 8N + 1 where N <= 6 (default 49)
Inference PrecisionBF16 (Recommended), FP16, FP32, FP8*, INT8, Not supported: INT4FP16*(Recommended), BF16, FP32, FP8*, INT8, Not supported: INT4BF16 (Recommended), FP16, FP32, FP8*, INT8, Not supported: INT4
Single GPU Memory Usage
SAT BF16: 76GB
diffusers BF16: from 10GB*
diffusers INT8(torchao): from 7GB*		SAT FP16: 18GB
diffusers FP16: 4GB minimum*
diffusers INT8 (torchao): 3.6GB minimum*		SAT BF16: 26GB
diffusers BF16 : 5GB minimum*
diffusers INT8 (torchao): 4.4GB minimum*
Multi-GPU Memory UsageBF16: 24GB* using diffusers
FP16: 10GB* using diffusers
BF16: 15GB* using diffusers
Inference Speed
(Step = 50, FP/BF16)		Single A100: ~1000 seconds (5-second video)
Single H100: ~550 seconds (5-second video)		Single A100: ~90 seconds
Single H100: ~45 seconds		Single A100: ~180 seconds
Single H100: ~90 seconds
Prompt LanguageEnglish*
Prompt Token Limit224 Tokens226 Tokens
Video Length5 seconds or 10 seconds6 seconds
Frame Rate16 frames / second 8 frames / second
Position Encoding3d_rope_pos_embed3d_sincos_pos_embed3d_rope_pos_embed3d_rope_pos_embed + learnable_pos_embed
Download Link (Diffusers)🤗 HuggingFace
🤖 ModelScope
🟣 WiseModel		🤗 HuggingFace
🤖 ModelScope
🟣 WiseModel		🤗 HuggingFace
🤖 ModelScope
🟣 WiseModel		🤗 HuggingFace
🤖 ModelScope
🟣 WiseModel		🤗 HuggingFace
🤖 ModelScope
🟣 WiseModel
Download Link (SAT)🤗 HuggingFace
🤖 ModelScope
🟣 WiseModel		SAT
+ +**Data Explanation** + ++ While testing using the diffusers library, all optimizations included in the diffusers library were enabled. This + scheme has not been tested for actual memory usage on devices outside of **NVIDIA A100 / H100** architectures. + Generally, this scheme can be adapted to all **NVIDIA Ampere architecture** and above devices. If optimizations are + disabled, memory consumption will multiply, with peak memory usage being about 3 times the value in the table. + However, speed will increase by about 3-4 times. You can selectively disable some optimizations, including: + +``` +pipe.enable_sequential_cpu_offload() +pipe.vae.enable_slicing() +pipe.vae.enable_tiling() +``` + ++ For multi-GPU inference, the `enable_sequential_cpu_offload()` optimization needs to be disabled. ++ Using INT8 models will slow down inference, which is done to accommodate lower-memory GPUs while maintaining minimal + video quality loss, though inference speed will significantly decrease. ++ The CogVideoX-2B model was trained in `FP16` precision, and all CogVideoX-5B models were trained in `BF16` precision. + We recommend using the precision in which the model was trained for inference. ++ [PytorchAO](https://github.com/pytorch/ao) and [Optimum-quanto](https://github.com/huggingface/optimum-quanto/) can be + used to quantize the text encoder, transformer, and VAE modules to reduce the memory requirements of CogVideoX. This + allows the model to run on free T4 Colabs or GPUs with smaller memory! Also, note that TorchAO quantization is fully + compatible with `torch.compile`, which can significantly improve inference speed. FP8 precision must be used on + devices with NVIDIA H100 and above, requiring source installation of `torch`, `torchao` Python packages. CUDA 12.4 is recommended. ++ The inference speed tests also used the above memory optimization scheme. Without memory optimization, inference speed + increases by about 10%. Only the `diffusers` version of the model supports quantization. ++ The model only supports English input; other languages can be translated into English for use via large model + refinement. + + +## Friendly Links + +We highly welcome contributions from the community and actively contribute to the open-source community. The following +works have already been adapted for CogVideoX, and we invite everyone to use them: + ++ [LeMiCa](https://unicomai.github.io/LeMiCa/): a diffusion model inference acceleration solution developed by China Unicom Data Science and Artificial Intelligence Research Institute. By leveraging cache-based techniques and global denoising path optimization, LeMiCa provides efficient inference support for CogVideoX, achieving nearly 2.5x lossless acceleration while maintaining visual consistency and quality. ++ [RIFLEx-CogVideoX](https://github.com/thu-ml/RIFLEx): + RIFLEx extends the video with just one line of code: `freq[k-1]=(2np.pi)/(Ls)`. The framework not only supports training-free inference, but also offers models fine-tuned based on CogVideoX. By fine-tuning the model for just 1,000 steps on original-length videos, RIFLEx significantly enhances its length extrapolation capability. ++ [CogVideoX-Fun](https://github.com/aigc-apps/CogVideoX-Fun): CogVideoX-Fun is a modified pipeline based on the + CogVideoX architecture, supporting flexible resolutions and multiple launch methods. ++ [CogStudio](https://github.com/pinokiofactory/cogstudio): A separate repository for CogVideo's Gradio Web UI, which + supports more functional Web UIs. ++ [Xorbits Inference](https://github.com/xorbitsai/inference): A powerful and comprehensive distributed inference + framework, allowing you to easily deploy your own models or the latest cutting-edge open-source models with just one + click. ++ [ComfyUI-CogVideoXWrapper](https://github.com/kijai/ComfyUI-CogVideoXWrapper) Use the ComfyUI framework to integrate + CogVideoX into your workflow. ++ [VideoSys](https://github.com/NUS-HPC-AI-Lab/VideoSys): VideoSys provides a user-friendly, high-performance + infrastructure for video generation, with full pipeline support and continuous integration of the latest models and + techniques. ++ [AutoDL Space](https://www.codewithgpu.com/i/THUDM/CogVideo/CogVideoX-5b-demo): A one-click deployment Huggingface + Space image provided by community members. ++ [Interior Design Fine-Tuning Model](https://huggingface.co/collections/bertjiazheng/koolcogvideox-66e4762f53287b7f39f8f3ba): + is a fine-tuned model based on CogVideoX, specifically designed for interior design. ++ [xDiT](https://github.com/xdit-project/xDiT): xDiT is a scalable inference engine for Diffusion Transformers (DiTs) + on multiple GPU Clusters. xDiT supports real-time image and video generations services. + [cogvideox-factory](https://github.com/a-r-r-o-w/cogvideox-factory): A cost-effective + fine-tuning framework for CogVideoX, compatible with the `diffusers` version model. Supports more resolutions, and + fine-tuning CogVideoX-5B can be done with a single 4090 GPU. ++ [CogVideoX-Interpolation](https://github.com/feizc/CogvideX-Interpolation): A pipeline based on the modified CogVideoX + structure, aimed at providing greater flexibility for keyframe interpolation generation. ++ [DiffSynth-Studio](https://github.com/modelscope/DiffSynth-Studio): DiffSynth Studio is a diffusion engine. It has + restructured the architecture, including text encoders, UNet, VAE, etc., enhancing computational performance while + maintaining compatibility with open-source community models. The framework has been adapted for CogVideoX. ++ [CogVideoX-Controlnet](https://github.com/TheDenk/cogvideox-controlnet): A simple ControlNet module code that includes the CogVideoX model. ++ [VideoTuna](https://github.com/VideoVerses/VideoTuna): VideoTuna is the first repo that integrates multiple AI video generation models for text-to-video, image-to-video, text-to-image generation. ++ [ConsisID](https://github.com/PKU-YuanGroup/ConsisID): An identity-preserving text-to-video generation model, bases on CogVideoX-5B, which keep the face consistent in the generated video by frequency decomposition. ++ [A Step by Step Tutorial](https://www.youtube.com/watch?v=5UCkMzP2VLE&ab_channel=SECourses): A step-by-step guide on installing and optimizing the CogVideoX1.5-5B-I2V model in Windows and cloud environments. Special thanks to the [FurkanGozukara](https://github.com/FurkanGozukara) for his effort and support! + +## Project Structure + +This open-source repository will guide developers to quickly get started with the basic usage and fine-tuning examples +of the **CogVideoX** open-source model. + +### Quick Start with Colab + +Here provide three projects that can be run directly on free Colab T4 instances: + ++ [CogVideoX-5B-T2V-Colab.ipynb](https://colab.research.google.com/drive/1pCe5s0bC_xuXbBlpvIH1z0kfdTLQPzCS?usp=sharing): + CogVideoX-5B Text-to-Video Colab code. ++ [CogVideoX-5B-T2V-Int8-Colab.ipynb](https://colab.research.google.com/drive/1DUffhcjrU-uz7_cpuJO3E_D4BaJT7OPa?usp=sharing): + CogVideoX-5B Quantized Text-to-Video Inference Colab code, which takes about 30 minutes per run. ++ [CogVideoX-5B-I2V-Colab.ipynb](https://colab.research.google.com/drive/17CqYCqSwz39nZAX2YyonDxosVKUZGzcX?usp=sharing): + CogVideoX-5B Image-to-Video Colab code. ++ [CogVideoX-5B-V2V-Colab.ipynb](https://colab.research.google.com/drive/1comfGAUJnChl5NwPuO8Ox5_6WCy4kbNN?usp=sharing): + CogVideoX-5B Video-to-Video Colab code. + +### Inference + ++ [dcli_demo](inference/cli_demo.py): A more detailed inference code explanation, including the significance of + common parameters. All of this is covered here. ++ [cli_demo_quantization](inference/cli_demo_quantization.py): + Quantized model inference code that can run on devices with lower memory. You can also modify this code to support + running CogVideoX models in FP8 precision. ++ [diffusers_vae_demo](inference/cli_vae_demo.py): Code for running VAE inference separately. ++ [space demo](inference/gradio_composite_demo): The same GUI code as used in the Huggingface Space, with frame + interpolation and super-resolution tools integrated. + +
+ +
+ ++ [convert_demo](inference/convert_demo.py): How to convert user input into long-form input suitable for CogVideoX. + Since CogVideoX is trained on long texts, we need to transform the input text distribution to match the training data + using an LLM. The script defaults to using GLM-4, but it can be replaced with GPT, Gemini, or any other large language + model. ++ [gradio_web_demo](inference/gradio_composite_demo): A simple Gradio web application demonstrating how to use the + CogVideoX-2B / 5B model to generate videos. Similar to our Huggingface Space, you can use this script to run a simple + web application for video generation. + +### finetune + ++ [finetune_demo](finetune/README.md): Fine-tuning scheme and details of the diffusers version of the CogVideoX model. + +### sat + ++ [sat_demo](sat/README.md): Contains the inference code and fine-tuning code of SAT weights. It is recommended to + improve based on the CogVideoX model structure. Innovative researchers use this code to better perform rapid stacking + and development. + +### Tools + +This folder contains some tools for model conversion / caption generation, etc. + ++ [convert_weight_sat2hf](tools/convert_weight_sat2hf.py): Converts SAT model weights to Huggingface model weights. ++ [caption_demo](tools/caption/README.md): Caption tool, a model that understands videos and outputs descriptions in + text. ++ [export_sat_lora_weight](tools/export_sat_lora_weight.py): SAT fine-tuning model export tool, exports the SAT Lora + Adapter in diffusers format. ++ [load_cogvideox_lora](tools/load_cogvideox_lora.py): Tool code for loading the diffusers version of fine-tuned Lora + Adapter. ++ [llm_flux_cogvideox](tools/llm_flux_cogvideox/llm_flux_cogvideox.py): Automatically generate videos using an + open-source local large language model + Flux + CogVideoX. ++ [parallel_inference_xdit](tools/parallel_inference/parallel_inference_xdit.py): + Supported by [xDiT](https://github.com/xdit-project/xDiT), parallelize the + video generation process on multiple GPUs. + +## CogVideo(ICLR'23) + +The official repo for the +paper: [CogVideo: Large-scale Pretraining for Text-to-Video Generation via Transformers](https://arxiv.org/abs/2205.15868) +is on the [CogVideo branch](https://github.com/THUDM/CogVideo/tree/CogVideo) + +**CogVideo is able to generate relatively high-frame-rate videos.** +A 4-second clip of 32 frames is shown below. + +![High-frame-rate sample](https://raw.githubusercontent.com/THUDM/CogVideo/CogVideo/assets/appendix-sample-highframerate.png) + +![Intro images](https://raw.githubusercontent.com/THUDM/CogVideo/CogVideo/assets/intro-image.png) +
+ +
+ + +The demo for CogVideo is at [https://models.aminer.cn/cogvideo](https://models.aminer.cn/cogvideo/), where you can get +hands-on practice on text-to-video generation. *The original input is in Chinese.* + +## Citation + +🌟 If you find our work helpful, please leave us a star and cite our paper. + +``` +@article{yang2024cogvideox, + title={CogVideoX: Text-to-Video Diffusion Models with An Expert Transformer}, + author={Yang, Zhuoyi and Teng, Jiayan and Zheng, Wendi and Ding, Ming and Huang, Shiyu and Xu, Jiazheng and Yang, Yuanming and Hong, Wenyi and Zhang, Xiaohan and Feng, Guanyu and others}, + journal={arXiv preprint arXiv:2408.06072}, + year={2024} +} +@article{hong2022cogvideo, + title={CogVideo: Large-scale Pretraining for Text-to-Video Generation via Transformers}, + author={Hong, Wenyi and Ding, Ming and Zheng, Wendi and Liu, Xinghan and Tang, Jie}, + journal={arXiv preprint arXiv:2205.15868}, + year={2022} +} +``` + +## Model-License + +The code in this repository is released under the [Apache 2.0 License](LICENSE). + +The CogVideoX-2B model (including its corresponding Transformers module and VAE module) is released under +the [Apache 2.0 License](LICENSE). + +The CogVideoX-5B model (Transformers module, include I2V and T2V) is released under +the [CogVideoX LICENSE](https://huggingface.co/THUDM/CogVideoX-5b/blob/main/LICENSE). diff --git a/README.wehub.md b/README.wehub.md new file mode 100644 index 0000000..486c96c --- /dev/null +++ b/README.wehub.md @@ -0,0 +1,7 @@ +# WeHub 来源说明 + +- 原始项目：`zai-org/CogVideo` +- 原始仓库：https://github.com/zai-org/CogVideo +- 导入方式：上游默认分支的最新快照 +- 原作者、版权和许可证信息以原始仓库及本仓库 LICENSE 为准 +- 本文件仅用于记录来源，不代表 WeHub 是原项目作者 diff --git a/README_ja.md b/README_ja.md new file mode 100644 index 0000000..afca6db --- /dev/null +++ b/README_ja.md @@ -0,0 +1,432 @@ +# CogVideo & CogVideoX + +[Read this in English](./README.md) + +[中文阅读](./README_zh.md) + +
+ +
+

+ 🤗 Huggingface Space または 🤖 ModelScope Space で CogVideoX-5B モデルをオンラインで体験してください +

+

+📚 論文と使用ドキュメントを表示します。 +

+

+ 👋 WeChat と Discord に参加 +

+

+📍 清影 と APIプラットフォーム を訪問して、より大規模な商用ビデオ生成モデルを体験. +

+ +## 更新とニュース +- 🔥🔥 ```2025/03/24```: [CogKit](https://github.com/THUDM/CogKit) は **CogView4** および **CogVideoX** シリーズの微調整と推論のためのフレームワークです。このツールキットを活用することで、私たちのマルチモーダル生成モデルを最大限に活用できます。 +- **ニュース**: ```2025/02/28```: DDIM Inverse が `CogVideoX-5B` と `CogVideoX1.5-5B` でサポートされました。詳細は [こちら](inference/ddim_inversion.py) をご覧ください。 +- **ニュース**: ```2025/01/08```: 私たちは`diffusers`バージョンのモデルをベースにした`Lora`微調整用のコードを更新しました。より少ないVRAM（ビデオメモリ）で動作します。詳細については[こちら](finetune/README_ja.md)をご覧ください。 +- **ニュース**: ```2024/11/15```: `CogVideoX1.5` モデルのdiffusersバージョンをリリースしました。わずかなパラメータ調整で以前のコードをそのまま利用可能です。 +- **ニュース**: ```2024/11/08```: `CogVideoX1.5` モデルをリリースしました。CogVideoX1.5 は CogVideoX オープンソースモデルのアップグレードバージョンです。 +CogVideoX1.5-5B シリーズモデルは、10秒 長の動画とより高い解像度をサポートしており、`CogVideoX1.5-5B-I2V` は任意の解像度での動画生成に対応しています。 +SAT コードはすでに更新されており、`diffusers` バージョンは現在適応中です。 +SAT バージョンのコードは [こちら](https://huggingface.co/THUDM/CogVideoX1.5-5B-SAT) からダウンロードできます。 +- 🔥 **ニュース**: ```2024/10/13```: コスト削減のため、単一の4090 GPUで`CogVideoX-5B` + を微調整できるフレームワーク [cogvideox-factory](https://github.com/a-r-r-o-w/cogvideox-factory) + がリリースされました。複数の解像度での微調整に対応しています。ぜひご利用ください！ +- 🔥**ニュース**: ```2024/10/10```: + 技術報告書を更新し、より詳細なトレーニング情報とデモを追加しました。 +- 🔥 **ニュース**: ```2024/10/10```: 技術報告書を更新しました。[こちら](https://arxiv.org/pdf/2408.06072) + をクリックしてご覧ください。さらにトレーニングの詳細とデモを追加しました。デモを見るには[こちら](https://yzy-thu.github.io/CogVideoX-demo/) + をクリックしてください。 +- 🔥**ニュース**: ```2024/10/09```: 飛書の[技術ドキュメント](https://zhipu-ai.feishu.cn/wiki/DHCjw1TrJiTyeukfc9RceoSRnCh) + でCogVideoXの微調整ガイドを公開しています。分配の自由度をさらに高めるため、公開されているドキュメント内のすべての例が完全に再現可能です。 +- 🔥**ニュース**: ```2024/9/19```: CogVideoXシリーズの画像生成ビデオモデル **CogVideoX-5B-I2V** + をオープンソース化しました。このモデルは、画像を背景入力として使用し、プロンプトワードと組み合わせてビデオを生成することができ、より高い制御性を提供します。これにより、CogVideoXシリーズのモデルは、テキストからビデオ生成、ビデオの継続、画像からビデオ生成の3つのタスクをサポートするようになりました。オンラインでの[体験](https://huggingface.co/spaces/THUDM/CogVideoX-5B-Space) + をお楽しみください。 +- 🔥 **ニュース**: ```2024/9/19```: + CogVideoXのトレーニングプロセスでビデオデータをテキスト記述に変換するために使用されるキャプションモデル [CogVLM2-Caption](https://huggingface.co/THUDM/cogvlm2-llama3-caption) + をオープンソース化しました。ダウンロードしてご利用ください。 +- 🔥 ```2024/8/27```: CogVideoXシリーズのより大きなモデル **CogVideoX-5B** + をオープンソース化しました。モデルの推論性能を大幅に最適化し、推論のハードルを大幅に下げました。`GTX 1080TI` などの旧型GPUで + **CogVideoX-2B** を、`RTX 3060` などのデスクトップGPUで **CogVideoX-5B** + モデルを実行できます。依存関係を更新・インストールするために、[要件](requirements.txt) + を厳守し、推論コードは [cli_demo](inference/cli_demo.py) を参照してください。さらに、**CogVideoX-2B** モデルのオープンソースライセンスが + **Apache 2.0 ライセンス** に変更されました。 +- 🔥 ```2024/8/6```: **CogVideoX-2B** 用の **3D Causal VAE** をオープンソース化しました。これにより、ビデオをほぼ無損失で再構築することができます。 +- 🔥 ```2024/8/6```: CogVideoXシリーズのビデオ生成モデルの最初のモデル、**CogVideoX-2B** をオープンソース化しました。 +- 🌱 **ソース**: ```2022/5/19```: CogVideoビデオ生成モデルをオープンソース化しました（現在、`CogVideo` + ブランチで確認できます）。これは、トランスフォーマーに基づく初のオープンソース大規模テキスト生成ビデオモデルです。技術的な詳細については、[ICLR'23論文](https://arxiv.org/abs/2205.15868) + をご覧ください。 + +**より強力なモデルが、より大きなパラメータサイズで登場予定です。お楽しみに！** + +## 目次 + +特定のセクションにジャンプ： + +- [クイックスタート](#クイックスタート) + - [プロンプトの最適化](#プロンプトの最適化) + - [SAT](#sat) + - [Diffusers](#diffusers) +- [Gallery](#gallery) + - [CogVideoX-5B](#cogvideox-5b) + - [CogVideoX-2B](#cogvideox-2b) +- [モデル紹介](#モデル紹介) +- [友好的リンク](#友好的リンク) +- [プロジェクト構造](#プロジェクト構造) + - [Colabでのクイックスタート](#colabでのクイックスタート) + - [Inference](#inference) + - [finetune](#finetune) + - [sat](#sat-1) + - [ツール](#ツール) +- [CogVideo(ICLR'23)](#cogvideoiclr23) +- [引用](#引用) +- [ライセンス契約](#ライセンス契約) + +## クイックスタート + +### プロンプトの最適化 + +モデルを実行する前に、[こちら](inference/convert_demo.py) +を参考にして、GLM-4（または同等の製品、例えばGPT-4）の大規模モデルを使用してどのようにモデルを最適化するかをご確認ください。これは非常に重要です。モデルは長いプロンプトでトレーニングされているため、良いプロンプトがビデオ生成の品質に直接影響を与えます。 + +### SAT + +[sat_demo](sat/README.md) の指示に従ってください: +SATウェイトの推論コードと微調整コードが含まれています。CogVideoXモデル構造に基づいて改善することをお勧めします。革新的な研究者は、このコードを使用して迅速なスタッキングと開発を行うことができます。 + +### Diffusers + +``` +pip install -r requirements.txt +``` + +次に [diffusers_demo](inference/cli_demo.py) を参照してください: 推論コードの詳細な説明が含まれており、一般的なパラメータの意味についても言及しています。 + +量子化推論の詳細については、[diffusers-torchao](https://github.com/sayakpaul/diffusers-torchao/) を参照してください。Diffusers +と TorchAO を使用することで、量子化推論も可能となり、メモリ効率の良い推論や、コンパイル時に場合によっては速度の向上が期待できます。A100 +および H100 +上でのさまざまな設定におけるメモリおよび時間のベンチマークの完全なリストは、[diffusers-torchao](https://github.com/sayakpaul/diffusers-torchao) +に公開されています。 + +## Gallery + +### CogVideoX-5B + + + + + + + + + + + + + + +
+ + + + + + + +
+ + + + + + + +
+ +### CogVideoX-2B + + + + + + + + +
+ + + + + + + +
+ +ギャラリーの対応するプロンプトワードを表示するには、[こちら](resources/galary_prompt.md)をクリックしてください + +## モデル紹介 + +CogVideoXは、[清影](https://chatglm.cn/video?fr=osm_cogvideox) と同源のオープンソース版ビデオ生成モデルです。 +以下の表に、提供しているビデオ生成モデルの基本情報を示します: + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + +
モデル名CogVideoX1.5-5B (最新)CogVideoX1.5-5B-I2V (最新)CogVideoX-2BCogVideoX-5BCogVideoX-5B-I2V
公開日2024年11月8日2024年11月8日2024年8月6日2024年8月27日2024年9月19日
ビデオ解像度1360 * 768 Min(W, H) = 768
768 ≤ Max(W, H) ≤ 1360
Max(W, H) % 16 = 0 		720 * 480
フレーム数16N + 1 (N <= 10) である必要があります (デフォルト 81)8N + 1 (N <= 6) である必要があります (デフォルト 49)
推論精度BF16(推奨), FP16, FP32，FP8*，INT8，INT4非対応FP16*(推奨), BF16, FP32，FP8*，INT8，INT4非対応BF16(推奨), FP16, FP32，FP8*，INT8，INT4非対応
単一GPUメモリ消費量
SAT BF16: 76GB
diffusers BF16：10GBから*
diffusers INT8(torchao)：7GBから*		SAT FP16: 18GB
diffusers FP16: 4GB以上*
diffusers INT8(torchao): 3.6GB以上*		SAT BF16: 26GB
diffusers BF16 : 5GB以上*
diffusers INT8(torchao): 4.4GB以上*
複数GPU推論メモリ消費量BF16: 24GB* using diffusers
FP16: 10GB* diffusers使用
BF16: 15GB* diffusers使用
推論速度
(Step = 50, FP/BF16)		シングルA100: ~1000秒(5秒ビデオ)
シングルH100: ~550秒(5秒ビデオ)		シングルA100: ~90秒
シングルH100: ~45秒		シングルA100: ~180秒
シングルH100: ~90秒
プロンプト言語英語*
プロンプト長さの上限224トークン226トークン
ビデオ長さ5秒または10秒6秒
フレームレート16フレーム/秒8フレーム/秒
位置エンコーディング3d_rope_pos_embed3d_sincos_pos_embed3d_rope_pos_embed3d_rope_pos_embed + learnable_pos_embed
ダウンロードリンク (Diffusers)🤗 HuggingFace
🤖 ModelScope
🟣 WiseModel		🤗 HuggingFace
🤖 ModelScope
🟣 WiseModel		🤗 HuggingFace
🤖 ModelScope
🟣 WiseModel		🤗 HuggingFace
🤖 ModelScope
🟣 WiseModel		🤗 HuggingFace
🤖 ModelScope
🟣 WiseModel
ダウンロードリンク (SAT)🤗 HuggingFace
🤖 ModelScope
🟣 WiseModel		SAT
+ +**データ解説** + ++ diffusersライブラリを使用してテストする際には、`diffusers`ライブラリが提供する全ての最適化が有効になっています。この方法は + **NVIDIA A100 / H100**以外のデバイスでのメモリ/メモリ消費のテストは行っていません。通常、この方法は**NVIDIA + Ampereアーキテクチャ** + 以上の全てのデバイスに適応できます。最適化を無効にすると、メモリ消費は倍増し、ピークメモリ使用量は表の3倍になりますが、速度は約3〜4倍向上します。以下の最適化を部分的に無効にすることが可能です: + +``` +pipe.enable_sequential_cpu_offload() +pipe.vae.enable_slicing() +pipe.vae.enable_tiling() +``` + ++ マルチGPUで推論する場合、`enable_sequential_cpu_offload()`最適化を無効にする必要があります。 ++ INT8モデルを使用すると推論速度が低下しますが、これはメモリの少ないGPUで正常に推論を行い、ビデオ品質の損失を最小限に抑えるための措置です。推論速度は大幅に低下します。 ++ CogVideoX-2Bモデルは`FP16`精度でトレーニングされており、CogVideoX-5Bモデルは`BF16` + 精度でトレーニングされています。推論時にはモデルがトレーニングされた精度を使用することをお勧めします。 ++ [PytorchAO](https://github.com/pytorch/ao)および[Optimum-quanto](https://github.com/huggingface/optimum-quanto/) + は、CogVideoXのメモリ要件を削減するためにテキストエンコーダ、トランスフォーマ、およびVAEモジュールを量子化するために使用できます。これにより、無料のT4 + Colabやより少ないメモリのGPUでモデルを実行することが可能になります。同様に重要なのは、TorchAOの量子化は`torch.compile` + と完全に互換性があり、推論速度を大幅に向上させることができる点です。`NVIDIA H100`およびそれ以上のデバイスでは`FP8` + 精度を使用する必要があります。これには、`torch`、`torchao` Pythonパッケージのソースコードからのインストールが必要です。`CUDA 12.4`の使用をお勧めします。 ++ 推論速度テストも同様に、上記のメモリ最適化方法を使用しています。メモリ最適化を使用しない場合、推論速度は約10％向上します。 + `diffusers`バージョンのモデルのみが量子化をサポートしています。 ++ モデルは英語入力のみをサポートしており、他の言語は大規模モデルの改善を通じて英語に翻訳できます。 + + +## 友好的リンク + +コミュニティからの貢献を大歓迎し、私たちもオープンソースコミュニティに積極的に貢献しています。以下の作品はすでにCogVideoXに対応しており、ぜひご利用ください： + ++ [LeMiCa](https://unicomai.github.io/LeMiCa/): 中国聯通データサイエンス・人工知能研究所が開発した拡散モデル推論加速ソリューション。LeMiCaは、キャッシュベースの技術とグローバルノイズ除去パス最適化を活用することで、CogVideoXの効率的な推論サポートを提供し、視覚的な一貫性と品質を維持しながら、約2.5倍のロスレス加速を実現します。 ++ [RIFLEx-CogVideoX](https://github.com/thu-ml/RIFLEx): + RIFLExは動画の長さを外挿する手法で、たった1行のコードで動画の長さを元の2倍に延長できます。RIFLExはトレーニング不要の推論をサポートするだけでなく、CogVideoXをベースにファインチューニングしたモデルも提供しています。元の長さの動画でわずか1000ステップのファインチューニングを行うだけで、長さ外挿能力を大幅に向上させることができます。 ++ [CogVideoX-Fun](https://github.com/aigc-apps/CogVideoX-Fun): + CogVideoX-Funは、CogVideoXアーキテクチャを基にした改良パイプラインで、自由な解像度と複数の起動方法をサポートしています。 ++ [CogStudio](https://github.com/pinokiofactory/cogstudio): CogVideo の Gradio Web UI の別のリポジトリ。より高機能な Web + UI をサポートします。 ++ [Xorbits Inference](https://github.com/xorbitsai/inference): + 強力で包括的な分散推論フレームワークであり、ワンクリックで独自のモデルや最新のオープンソースモデルを簡単にデプロイできます。 ++ [ComfyUI-CogVideoXWrapper](https://github.com/kijai/ComfyUI-CogVideoXWrapper) + ComfyUIフレームワークを使用して、CogVideoXをワークフローに統合します。 ++ [VideoSys](https://github.com/NUS-HPC-AI-Lab/VideoSys): VideoSysは、使いやすく高性能なビデオ生成インフラを提供し、最新のモデルや技術を継続的に統合しています。 ++ [AutoDLイメージ](https://www.codewithgpu.com/i/THUDM/CogVideo/CogVideoX-5b-demo): コミュニティメンバーが提供するHuggingface + Spaceイメージのワンクリックデプロイメント。 ++ [インテリアデザイン微調整モデル](https://huggingface.co/collections/bertjiazheng/koolcogvideox-66e4762f53287b7f39f8f3ba): + は、CogVideoXを基盤にした微調整モデルで、インテリアデザイン専用に設計されています。 ++ [xDiT](https://github.com/xdit-project/xDiT): + xDiTは、複数のGPUクラスター上でDiTsを並列推論するためのエンジンです。xDiTはリアルタイムの画像およびビデオ生成サービスをサポートしています。 ++ [CogVideoX-Interpolation](https://github.com/feizc/CogvideX-Interpolation): + キーフレーム補間生成において、より大きな柔軟性を提供することを目的とした、CogVideoX構造を基にした修正版のパイプライン。 ++ [DiffSynth-Studio](https://github.com/modelscope/DiffSynth-Studio): DiffSynth + Studioは、拡散エンジンです。テキストエンコーダー、UNet、VAEなどを含むアーキテクチャを再構築し、オープンソースコミュニティモデルとの互換性を維持しつつ、計算性能を向上させました。このフレームワークはCogVideoXに適応しています。 ++ [CogVideoX-Controlnet](https://github.com/TheDenk/cogvideox-controlnet): CogVideoXモデルを含むシンプルなControlNetモジュールのコード。 ++ [VideoTuna](https://github.com/VideoVerses/VideoTuna): VideoTuna は、テキストからビデオ、画像からビデオ、テキストから画像生成のための複数のAIビデオ生成モデルを統合した最初のリポジトリです。 ++ [ConsisID](https://github.com/PKU-YuanGroup/ConsisID): 一貫性のある顔を保持するために、周波数分解を使用するCogVideoX-5Bに基づいたアイデンティティ保持型テキストから動画生成モデル。 ++ [ステップバイステップチュートリアル](https://www.youtube.com/watch?v=5UCkMzP2VLE&ab_channel=SECourses): WindowsおよびクラウドでのCogVideoX1.5-5B-I2Vモデルのインストールと最適化に関するステップバイステップガイド。[FurkanGozukara](https://github.com/FurkanGozukara)氏の尽力とサポートに感謝いたします！ + +## プロジェクト構造 + +このオープンソースリポジトリは、**CogVideoX** オープンソースモデルの基本的な使用方法と微調整の例を迅速に開始するためのガイドです。 + +### Colabでのクイックスタート + +無料のColab T4上で直接実行できる3つのプロジェクトを提供しています。 + ++ [CogVideoX-5B-T2V-Colab.ipynb](https://colab.research.google.com/drive/1pCe5s0bC_xuXbBlpvIH1z0kfdTLQPzCS?usp=sharing): + CogVideoX-5B テキストからビデオへの生成用Colabコード。 ++ [CogVideoX-5B-T2V-Int8-Colab.ipynb](https://colab.research.google.com/drive/1DUffhcjrU-uz7_cpuJO3E_D4BaJT7OPa?usp=sharing): + CogVideoX-5B テキストからビデオへの量子化推論用Colabコード。1回の実行に約30分かかります。 ++ [CogVideoX-5B-I2V-Colab.ipynb](https://colab.research.google.com/drive/17CqYCqSwz39nZAX2YyonDxosVKUZGzcX?usp=sharing): + CogVideoX-5B 画像からビデオへの生成用Colabコード。 ++ [CogVideoX-5B-V2V-Colab.ipynb](https://colab.research.google.com/drive/1comfGAUJnChl5NwPuO8Ox5_6WCy4kbNN?usp=sharing): + CogVideoX-5B ビデオからビデオへの生成用Colabコード。 + +### Inference + ++ [cli_demo](inference/cli_demo.py): 推論コードの詳細な説明が含まれており、一般的なパラメータの意味についても言及しています。 ++ [cli_demo_quantization](inference/cli_demo_quantization.py): + 量子化モデル推論コードで、低メモリのデバイスでも実行可能です。また、このコードを変更して、FP8 精度の CogVideoX + モデルの実行をサポートすることもできます。 ++ [diffusers_vae_demo](inference/cli_vae_demo.py): VAE推論コードの実行には現在71GBのメモリが必要ですが、将来的には最適化される予定です。 ++ [space demo](inference/gradio_composite_demo): Huggingface Spaceと同じGUIコードで、フレーム補間や超解像ツールが組み込まれています。 + +
+ +
+ ++ [convert_demo](inference/convert_demo.py): + ユーザー入力をCogVideoXに適した形式に変換する方法。CogVideoXは長いキャプションでトレーニングされているため、入力テキストをLLMを使用してトレーニング分布と一致させる必要があります。デフォルトではGLM-4を使用しますが、GPT、Geminiなどの他のLLMに置き換えることもできます。 ++ [gradio_web_demo](inference/gradio_web_demo.py): CogVideoX-2B / 5B モデルを使用して動画を生成する方法を示す、シンプルな + Gradio Web UI デモです。私たちの Huggingface Space と同様に、このスクリプトを使用して Web デモを起動することができます。 + +### finetune + ++ [train_cogvideox_lora](finetune/README_ja.md): CogVideoX diffusers 微調整方法の詳細な説明が含まれています。このコードを使用して、自分のデータセットで + CogVideoX を微調整することができます。 + +### sat + ++ [sat_demo](sat/README.md): + SATウェイトの推論コードと微調整コードが含まれています。CogVideoXモデル構造に基づいて改善することをお勧めします。革新的な研究者は、このコードを使用して迅速なスタッキングと開発を行うことができます。 + +### ツール + +このフォルダには、モデル変換/キャプション生成などのツールが含まれています。 + ++ [convert_weight_sat2hf](tools/convert_weight_sat2hf.py): SAT モデルの重みを Huggingface モデルの重みに変換します。 ++ [caption_demo](tools/caption/README_ja.md): Caption ツール、ビデオを理解してテキストで出力するモデル。 ++ [export_sat_lora_weight](tools/export_sat_lora_weight.py): SAT ファインチューニングモデルのエクスポートツール、SAT Lora + Adapter を diffusers 形式でエクスポートします。 ++ [load_cogvideox_lora](tools/load_cogvideox_lora.py): diffusers 版のファインチューニングされた Lora Adapter + をロードするためのツールコード。 ++ [llm_flux_cogvideox](tools/llm_flux_cogvideox/llm_flux_cogvideox.py): オープンソースのローカル大規模言語モデル + + Flux + CogVideoX を使用して自動的に動画を生成します。 ++ [parallel_inference_xdit](tools/parallel_inference/parallel_inference_xdit.py)： + [xDiT](https://github.com/xdit-project/xDiT) + によってサポートされ、ビデオ生成プロセスを複数の GPU で並列化します。 ++ [cogvideox-factory](https://github.com/a-r-r-o-w/cogvideox-factory): CogVideoXの低コスト微調整フレームワークで、 + `diffusers`バージョンのモデルに適応しています。より多くの解像度に対応し、単一の4090 GPUでCogVideoX-5Bの微調整が可能です。 + +## CogVideo(ICLR'23) + +論文の公式リポジトリ: [CogVideo: Large-scale Pretraining for Text-to-Video Generation via Transformers](https://arxiv.org/abs/2205.15868) +は [CogVideo branch](https://github.com/THUDM/CogVideo/tree/CogVideo) にあります。 + +**CogVideoは比較的高フレームレートのビデオを生成することができます。** +32フレームの4秒間のクリップが以下に示されています。 + +![High-frame-rate sample](https://raw.githubusercontent.com/THUDM/CogVideo/CogVideo/assets/appendix-sample-highframerate.png) + +![Intro images](https://raw.githubusercontent.com/THUDM/CogVideo/CogVideo/assets/intro-image.png) +
+ +
+ + +CogVideoのデモは [https://models.aminer.cn/cogvideo](https://models.aminer.cn/cogvideo/) で体験できます。 +*元の入力は中国語です。* + +## 引用 + +🌟 私たちの仕事が役立つと思われた場合、ぜひスターを付けていただき、論文を引用してください。 + +``` +@article{yang2024cogvideox, + title={CogVideoX: Text-to-Video Diffusion Models with An Expert Transformer}, + author={Yang, Zhuoyi and Teng, Jiayan and Zheng, Wendi and Ding, Ming and Huang, Shiyu and Xu, Jiazheng and Yang, Yuanming and Hong, Wenyi and Zhang, Xiaohan and Feng, Guanyu and others}, + journal={arXiv preprint arXiv:2408.06072}, + year={2024} +} +@article{hong2022cogvideo, + title={CogVideo: Large-scale Pretraining for Text-to-Video Generation via Transformers}, + author={Hong, Wenyi and Ding, Ming and Zheng, Wendi and Liu, Xinghan and Tang, Jie}, + journal={arXiv preprint arXiv:2205.15868}, + year={2022} +} +``` + +## ライセンス契約 + +このリポジトリのコードは [Apache 2.0 License](LICENSE) の下で公開されています。 + +CogVideoX-2B モデル (対応するTransformersモジュールやVAEモジュールを含む) は +[Apache 2.0 License](LICENSE) の下で公開されています。 + +CogVideoX-5B モデル（Transformers モジュール、画像生成ビデオとテキスト生成ビデオのバージョンを含む） は +[CogVideoX LICENSE](https://huggingface.co/THUDM/CogVideoX-5b/blob/main/LICENSE) の下で公開されています。 diff --git a/README_zh.md b/README_zh.md new file mode 100644 index 0000000..3eb399e --- /dev/null +++ b/README_zh.md @@ -0,0 +1,412 @@ +# CogVideo & CogVideoX + +[Read this in English](./README.md) + +[日本語で読む](./README_ja.md) + +
+ +
+

+在 🤗 Huggingface Space 或 🤖 ModelScope Space 在线体验 CogVideoX-5B 模型 +

+

+📚 查看 论文 和 使用文档 +

+

+ 👋 加入我们的 微信 和 Discord +

+

+📍 前往 清影 和 API平台 体验更大规模的商业版视频生成模型。 +

+ +## 项目更新 + +- 🔥🔥 **News**: ```2025/03/24```: 我们推出了 [CogKit](https://github.com/THUDM/CogKit) 工具，这是一个微调**CogView4**, **CogVideoX** 系列的微调和推理框架，一个工具包，玩转我们的多模态生成模型。 +- 🔥 **News**: ```2025/02/28```: DDIM Inverse 已经在`CogVideoX-5B` 和 `CogVideoX1.5 -5B` 支持，查看 [here](inference/ddim_inversion.py). +- 🔥 **News**: ```2025/01/08```: 我们更新了基于`diffusers`版本模型的`Lora`微调代码，占用显存更低，详情请见[这里](finetune/README_zh.md)。 +- 🔥 **News**: ```2024/11/15```: 我们发布 `CogVideoX1.5` 模型的diffusers版本，仅需调整部分参数即可沿用之前的代码。 +- 🔥 **News**: ```2024/11/08```: 我们发布 `CogVideoX1.5` 模型。CogVideoX1.5 是 CogVideoX 开源模型的升级版本。 +CogVideoX1.5-5B 系列模型支持 **10秒** 长度的视频和更高的分辨率，其中 `CogVideoX1.5-5B-I2V` 支持 **任意分辨率** 的视频生成，SAT代码已经更新。`diffusers`版本还在适配中。SAT版本代码前往 [这里](https://huggingface.co/THUDM/CogVideoX1.5-5B-SAT) 下载。 +- 🔥**News**: ```2024/10/13```: 成本更低，单卡4090可微调 `CogVideoX-5B` + 的微调框架[cogvideox-factory](https://github.com/a-r-r-o-w/cogvideox-factory)已经推出，多种分辨率微调，欢迎使用。 +- 🔥 **News**: ```2024/10/10```: 我们更新了我们的技术报告,请点击 [这里](https://arxiv.org/pdf/2408.06072) + 查看，附上了更多的训练细节和demo，关于demo，点击[这里](https://yzy-thu.github.io/CogVideoX-demo/) 查看。 +- 🔥 **News**: ```2024/10/09```: 我们在飞书[技术文档](https://zhipu-ai.feishu.cn/wiki/DHCjw1TrJiTyeukfc9RceoSRnCh") + 公开CogVideoX微调指导，以进一步增加分发自由度，公开文档中所有示例可以完全复现 +- 🔥 **News**: ```2024/9/19```: 我们开源 CogVideoX 系列图生视频模型 **CogVideoX-5B-I2V** + 。该模型可以将一张图像作为背景输入，结合提示词一起生成视频，具有更强的可控性。 + 至此，CogVideoX系列模型已经支持文本生成视频，视频续写，图片生成视频三种任务。欢迎前往在线[体验](https://huggingface.co/spaces/THUDM/CogVideoX-5B-Space)。 +- 🔥 **News**: ```2024/9/19```: CogVideoX 训练过程中用于将视频数据转换为文本描述的 Caption + 模型 [CogVLM2-Caption](https://huggingface.co/THUDM/cogvlm2-llama3-caption) + 已经开源。欢迎前往下载并使用。 +- 🔥 ```2024/8/27```: 我们开源 CogVideoX 系列更大的模型 **CogVideoX-5B** + 。我们大幅度优化了模型的推理性能，推理门槛大幅降低，您可以在 `GTX 1080TI` 等早期显卡运行 **CogVideoX-2B**，在 `RTX 3060` + 等桌面端甜品卡运行 **CogVideoX-5B** 模型。 请严格按照[要求](requirements.txt) + 更新安装依赖，推理代码请查看 [cli_demo](inference/cli_demo.py)。同时，**CogVideoX-2B** 模型开源协议已经修改为**Apache 2.0 协议**。 +- 🔥 ```2024/8/6```: 我们开源 **3D Causal VAE**，用于 **CogVideoX-2B**，可以几乎无损地重构视频。 +- 🔥 ```2024/8/6```: 我们开源 CogVideoX 系列视频生成模型的第一个模型, **CogVideoX-2B**。 +- 🌱 **Source**: ```2022/5/19```: 我们开源了 CogVideo 视频生成模型（现在你可以在 `CogVideo` 分支中看到），这是首个开源的基于 + Transformer 的大型文本生成视频模型，您可以访问 [ICLR'23 论文](https://arxiv.org/abs/2205.15868) 查看技术细节。 + +## 目录 + +跳转到指定部分： + +- [快速开始](#快速开始) + - [提示词优化](#提示词优化) + - [SAT](#sat) + - [Diffusers](#diffusers) +- [视频作品](#视频作品) + - [CogVideoX-5B](#cogvideox-5b) + - [CogVideoX-2B](#cogvideox-2b) +- [模型介绍](#模型介绍) +- [友情链接](#友情链接) +- [完整项目代码结构](#完整项目代码结构) + - [Colab 快速使用](#colab-快速使用) + - [inference](#inference) + - [finetune](#finetune) + - [sat](#sat-1) + - [tools](#tools) +- [CogVideo(ICLR'23)](#cogvideoiclr23) +- [引用](#引用) +- [模型协议](#模型协议) + +## 快速开始 + +### 提示词优化 + +在开始运行模型之前，请参考 [这里](inference/convert_demo.py) 查看我们是怎么使用GLM-4(或者同级别的其他产品，例如GPT-4) +大模型对模型进行优化的，这很重要， +由于模型是在长提示词下训练的，一个好的提示词直接影响了视频生成的质量。 + +### SAT + +查看sat文件夹下的 [sat_demo](sat/README.md)：包含了 SAT 权重的推理代码和微调代码，推荐基于此代码进行 CogVideoX +模型结构的改进，研究者使用该代码可以更好的进行快速的迭代和开发。 + +### Diffusers + +``` +pip install -r requirements.txt +``` + +查看[diffusers_demo](inference/cli_demo.py)：包含对推理代码更详细的解释，包括各种关键的参数。 + +欲了解更多关于量化推理的细节，请参考 [diffusers-torchao](https://github.com/sayakpaul/diffusers-torchao/)。使用 Diffusers +和 TorchAO，量化推理也是可能的，这可以实现内存高效的推理，并且在某些情况下编译后速度有所提升。有关在 A100 和 H100 +上使用各种设置的内存和时间基准测试的完整列表，已发布在 [diffusers-torchao](https://github.com/sayakpaul/diffusers-torchao) +上。 + +## 视频作品 + +### CogVideoX-5B + + + + + + + + + + + + + + +
+ + + + + + + +
+ + + + + + + +
+ +### CogVideoX-2B + + + + + + + + +
+ + + + + + + +
+ + +查看画廊的对应提示词，请点击[这里](resources/galary_prompt.md) + +## 模型介绍 + +CogVideoX是 [清影](https://chatglm.cn/video?fr=osm_cogvideox) 同源的开源版本视频生成模型。 +下表展示我们提供的视频生成模型相关基础信息: + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + +
模型名CogVideoX1.5-5B (最新)CogVideoX1.5-5B-I2V (最新)CogVideoX-2BCogVideoX-5BCogVideoX-5B-I2V
发布时间2024年11月8日2024年11月8日2024年8月6日2024年8月27日2024年9月19日
视频分辨率1360 * 768 Min(W, H) = 768
768 ≤ Max(W, H) ≤ 1360
Max(W, H) % 16 = 0 		720 * 480
帧数必须为 16N + 1 其中 N <= 10 (默认 81)必须为 8N + 1 其中 N <= 6 (默认 49)
推理精度BF16(推荐), FP16, FP32，FP8*，INT8，不支持INT4FP16*(推荐), BF16, FP32，FP8*，INT8，不支持INT4BF16(推荐), FP16, FP32，FP8*，INT8，不支持INT4
单GPU显存消耗
SAT BF16: 76GB
diffusers BF16 : 10GB起*
diffusers INT8(torchao): 7G起* 		SAT FP16: 18GB
diffusers FP16: 4GB起*
diffusers INT8(torchao): 3.6G起*		SAT BF16: 26GB
diffusers BF16 : 5GB起*
diffusers INT8(torchao): 4.4G起*
多GPU推理显存消耗BF16: 24GB* using diffusers
FP16: 10GB* using diffusers
BF16: 15GB* using diffusers
推理速度
(Step = 50, FP/BF16)		单卡A100: ~1000秒(5秒视频)
单卡H100: ~550秒(5秒视频)		单卡A100: ~90秒
单卡H100: ~45秒		单卡A100: ~180秒
单卡H100: ~90秒
提示词语言English*
提示词长度上限224 Tokens226 Tokens
视频长度5 秒 或 10 秒6 秒
帧率16 帧 / 秒 8 帧 / 秒
位置编码3d_rope_pos_embed3d_sincos_pos_embed3d_rope_pos_embed3d_rope_pos_embed + learnable_pos_embed
下载链接 (Diffusers)🤗 HuggingFace
🤖 ModelScope
🟣 WiseModel		🤗 HuggingFace
🤖 ModelScope
🟣 WiseModel		🤗 HuggingFace
🤖 ModelScope
🟣 WiseModel		🤗 HuggingFace
🤖 ModelScope
🟣 WiseModel		🤗 HuggingFace
🤖 ModelScope
🟣 WiseModel
下载链接 (SAT)🤗 HuggingFace
🤖 ModelScope
🟣 WiseModel		SAT
+ +**数据解释** + ++ 使用 diffusers 库进行测试时，启用了全部`diffusers`库自带的优化，该方案未测试在非**NVIDIA A100 / H100** + 外的设备上的实际显存 / 内存占用。通常，该方案可以适配于所有 **NVIDIA 安培架构** + 以上的设备。若关闭优化，显存占用会成倍增加，峰值显存约为表格的3倍。但速度提升3-4倍左右。你可以选择性的关闭部分优化，这些优化包括: + +``` +pipe.enable_sequential_cpu_offload() +pipe.vae.enable_slicing() +pipe.vae.enable_tiling() +``` + ++ 多GPU推理时，需要关闭 `enable_sequential_cpu_offload()` 优化。 ++ 使用 INT8 模型会导致推理速度降低，此举是为了满足显存较低的显卡能正常推理并保持较少的视频质量损失，推理速度大幅降低。 ++ CogVideoX-2B 模型采用 `FP16` 精度训练， 搜有 CogVideoX-5B 模型采用 `BF16` 精度训练。我们推荐使用模型训练的精度进行推理。 ++ [PytorchAO](https://github.com/pytorch/ao) 和 [Optimum-quanto](https://github.com/huggingface/optimum-quanto/) + 可以用于量化文本编码器、Transformer 和 VAE 模块，以降低 CogVideoX 的内存需求。这使得在免费的 T4 Colab 或更小显存的 GPU + 上运行模型成为可能！同样值得注意的是，TorchAO 量化完全兼容 `torch.compile`，这可以显著提高推理速度。在 `NVIDIA H100` + 及以上设备上必须使用 `FP8` 精度，这需要源码安装 `torch`、`torchao` Python 包。建议使用 `CUDA 12.4`。 ++ 推理速度测试同样采用了上述显存优化方案，不采用显存优化的情况下，推理速度提升约10%。 只有`diffusers`版本模型支持量化。 ++ 模型仅支持英语输入，其他语言可以通过大模型润色时翻译为英语。 + +## 友情链接 + +我们非常欢迎来自社区的贡献，并积极的贡献开源社区。以下作品已经对CogVideoX进行了适配，欢迎大家使用: + ++ [LeMiCa](https://unicomai.github.io/LeMiCa/): 由中国联通数据科学与人工智能研究院开发的扩散模型推理加速解决方案。它利用基于缓存的技术和全局去噪路径优化，为CogVideoX提供高效的推理支持，在保持视觉一致性和质量的前提下，实现了近2.5倍的无损加速。 ++ [RIFLEx-CogVideoX](https://github.com/thu-ml/RIFLEx): + RIFLEx 是一个视频长度外推的方法，只需一行代码即可将视频生成长度延伸为原先的二倍。RIFLEx 不仅支持 Training-free 的推理，也提供基于 CogVideoX 进行微调的模型，只需在原有长度视频上微调 1000 步即可大大提高长度外推能力。 ++ [CogVideoX-Fun](https://github.com/aigc-apps/CogVideoX-Fun): + CogVideoX-Fun是一个基于CogVideoX结构修改后的的pipeline，支持自由的分辨率，多种启动方式。 ++ [CogStudio](https://github.com/pinokiofactory/cogstudio): CogVideo 的 Gradio Web UI单独实现仓库，支持更多功能的 Web UI。 ++ [Xorbits Inference](https://github.com/xorbitsai/inference): 性能强大且功能全面的分布式推理框架，轻松一键部署你自己的模型或内置的前沿开源模型。 ++ [ComfyUI-CogVideoXWrapper](https://github.com/kijai/ComfyUI-CogVideoXWrapper) 使用ComfyUI框架，将CogVideoX加入到你的工作流中。 ++ [VideoSys](https://github.com/NUS-HPC-AI-Lab/VideoSys): VideoSys 提供了易用且高性能的视频生成基础设施，支持完整的管道，并持续集成最新的模型和技术。 ++ [AutoDL镜像](https://www.codewithgpu.com/i/THUDM/CogVideo/CogVideoX-5b-demo): 由社区成员提供的一键部署Huggingface + Space镜像。 ++ [室内设计微调模型](https://huggingface.co/collections/bertjiazheng/koolcogvideox-66e4762f53287b7f39f8f3ba) 基于 + CogVideoX的微调模型，它专为室内设计而设计 ++ [xDiT](https://github.com/xdit-project/xDiT): xDiT是一个用于在多GPU集群上对DiTs并行推理的引擎。xDiT支持实时图像和视频生成服务。 ++ [CogVideoX-Interpolation](https://github.com/feizc/CogvideX-Interpolation): 基于 CogVideoX 结构修改的管道，旨在为关键帧插值生成提供更大的灵活性。 ++ [DiffSynth-Studio](https://github.com/modelscope/DiffSynth-Studio): DiffSynth 工作室是一款扩散引擎。重构了架构，包括文本编码器、UNet、VAE + 等，在保持与开源社区模型兼容性的同时，提升了计算性能。该框架已经适配 CogVideoX。 ++ [CogVideoX-Controlnet](https://github.com/TheDenk/cogvideox-controlnet): 一个包含 CogvideoX 模型的简单 Controlnet 模块的代码。 ++ [VideoTuna](https://github.com/VideoVerses/VideoTuna)：VideoTuna 是首个集成多种 AI 视频生成模型的仓库，支持文本转视频、图像转视频、文本转图像生成。 ++ [ConsisID](https://github.com/PKU-YuanGroup/ConsisID): 一种身份保持的文本到视频生成模型，基于 CogVideoX-5B，通过频率分解在生成的视频中保持面部一致性。 ++ [教程](https://www.youtube.com/watch?v=5UCkMzP2VLE&ab_channel=SECourses): 一个关于在Windows和云环境中安装和优化CogVideoX1.5-5B-I2V模型的分步指南。特别感谢[FurkanGozukara](https://github.com/FurkanGozukara)的努力和支持！ + + +## 完整项目代码结构 + +本开源仓库将带领开发者快速上手 **CogVideoX** 开源模型的基础调用方式、微调示例。 + +### Colab 快速使用 + +这里提供了三个能直接在免费的 Colab T4上 运行的项目 + ++ [CogVideoX-5B-T2V-Colab.ipynb](https://colab.research.google.com/drive/1pCe5s0bC_xuXbBlpvIH1z0kfdTLQPzCS?usp=sharing): + CogVideoX-5B 文字生成视频 Colab 代码。 ++ [CogVideoX-5B-T2V-Int8-Colab.ipynb](https://colab.research.google.com/drive/1DUffhcjrU-uz7_cpuJO3E_D4BaJT7OPa?usp=sharing): + CogVideoX-5B 文字生成视频量化推理 Colab 代码，运行一次大约需要30分钟。 ++ [CogVideoX-5B-I2V-Colab.ipynb](https://colab.research.google.com/drive/17CqYCqSwz39nZAX2YyonDxosVKUZGzcX?usp=sharing): + CogVideoX-5B 图片生成视频 Colab 代码。 ++ [CogVideoX-5B-V2V-Colab.ipynb](https://colab.research.google.com/drive/1comfGAUJnChl5NwPuO8Ox5_6WCy4kbNN?usp=sharing): + CogVideoX-5B 视频生成视频 Colab 代码。 + +### inference + ++ [cli_demo](inference/cli_demo.py): 更详细的推理代码讲解，常见参数的意义，在这里都会提及。 ++ [cli_demo_quantization](inference/cli_demo_quantization.py): + 量化模型推理代码，可以在显存较低的设备上运行，也可以基于此代码修改，以支持运行FP8等精度的CogVideoX模型。请注意，FP8 + 仅测试通过，且必须将 `torch-nightly`,`torchao`源代码安装，不建议在生产环境中使用。 ++ [diffusers_vae_demo](inference/cli_vae_demo.py): 单独执行VAE的推理代码。 ++ [space demo](inference/gradio_composite_demo): Huggingface Space同款的 GUI 代码，植入了插帧，超分工具。 + +
+ +
+ ++ [convert_demo](inference/convert_demo.py): 如何将用户的输入转换成适合 + CogVideoX的长输入。因为CogVideoX是在长文本上训练的，所以我们需要把输入文本的分布通过LLM转换为和训练一致的长文本。脚本中默认使用GLM-4，也可以替换为GPT、Gemini等任意大语言模型。 ++ [gradio_web_demo](inference/gradio_composite_demo/app.py): 与 Huggingface Space 完全相同的代码实现，快速部署 CogVideoX + GUI体验。 + +### finetune + ++ [train_cogvideox_lora](finetune/README_zh.md): diffusers版本 CogVideoX 模型微调方案和细节。 + +### sat + ++ [sat_demo](sat/README_zh.md): 包含了 SAT 权重的推理代码和微调代码，推荐基于 CogVideoX + 模型结构进行改进，创新的研究者使用改代码以更好的进行快速的堆叠和开发。 + +### tools + +本文件夹包含了一些工具，用于模型的转换 / Caption 等工作。 + ++ [convert_weight_sat2hf](tools/convert_weight_sat2hf.py): 将 SAT 模型权重转换为 Huggingface 模型权重。 ++ [caption_demo](tools/caption/README_zh.md): Caption 工具，对视频理解并用文字输出的模型。 ++ [export_sat_lora_weight](tools/export_sat_lora_weight.py): SAT微调模型导出工具，将 + SAT Lora Adapter 导出为 diffusers 格式。 ++ [load_cogvideox_lora](tools/load_cogvideox_lora.py): 载入diffusers版微调Lora Adapter的工具代码。 ++ [llm_flux_cogvideox](tools/llm_flux_cogvideox/llm_flux_cogvideox.py): 使用开源本地大语言模型 + Flux + + CogVideoX实现自动化生成视频。 ++ [parallel_inference_xdit](tools/parallel_inference/parallel_inference_xdit.py): + 在多个 GPU 上并行化视频生成过程， + 由[xDiT](https://github.com/xdit-project/xDiT)提供支持。 ++ [cogvideox-factory](https://github.com/a-r-r-o-w/cogvideox-factory): CogVideoX低成文微调框架，适配`diffusers` + 版本模型。支持更多分辨率，单卡4090即可微调 CogVideoX-5B 。 + +## CogVideo(ICLR'23) + +[CogVideo: Large-scale Pretraining for Text-to-Video Generation via Transformers](https://arxiv.org/abs/2205.15868) +的官方repo位于[CogVideo branch](https://github.com/THUDM/CogVideo/tree/CogVideo)。 + +**CogVideo可以生成高帧率视频，下面展示了一个32帧的4秒视频。** + +![High-frame-rate sample](https://raw.githubusercontent.com/THUDM/CogVideo/CogVideo/assets/appendix-sample-highframerate.png) + +![Intro images](https://raw.githubusercontent.com/THUDM/CogVideo/CogVideo/assets/intro-image.png) + + +
+ +
+ +CogVideo的demo网站在[https://models.aminer.cn/cogvideo](https://models.aminer.cn/cogvideo/)。您可以在这里体验文本到视频生成。 +*原始输入为中文。* + +## 引用 + +🌟 如果您发现我们的工作有所帮助，欢迎引用我们的文章，留下宝贵的stars + +``` +@article{yang2024cogvideox, + title={CogVideoX: Text-to-Video Diffusion Models with An Expert Transformer}, + author={Yang, Zhuoyi and Teng, Jiayan and Zheng, Wendi and Ding, Ming and Huang, Shiyu and Xu, Jiazheng and Yang, Yuanming and Hong, Wenyi and Zhang, Xiaohan and Feng, Guanyu and others}, + journal={arXiv preprint arXiv:2408.06072}, + year={2024} +} +@article{hong2022cogvideo, + title={CogVideo: Large-scale Pretraining for Text-to-Video Generation via Transformers}, + author={Hong, Wenyi and Ding, Ming and Zheng, Wendi and Liu, Xinghan and Tang, Jie}, + journal={arXiv preprint arXiv:2205.15868}, + year={2022} +} +``` + +## 模型协议 + +本仓库代码使用 [Apache 2.0 协议](LICENSE) 发布。 + +CogVideoX-2B 模型 (包括其对应的Transformers模块，VAE模块) 根据 [Apache 2.0 协议](LICENSE) 许可证发布。 + +CogVideoX-5B 模型 (Transformers 模块，包括图生视频，文生视频版本) +根据 [CogVideoX LICENSE](https://huggingface.co/THUDM/CogVideoX-5b/blob/main/LICENSE) +许可证发布。 diff --git a/finetune/README.md b/finetune/README.md new file mode 100644 index 0000000..aeaf789 --- /dev/null +++ b/finetune/README.md @@ -0,0 +1,148 @@ +# CogVideoX Diffusers Fine-tuning Guide + +[中文阅读](./README_zh.md) + +[日本語で読む](./README_ja.md) + +If you're looking for the fine-tuning instructions for the SAT version, please check [here](../sat/README_zh.md). The +dataset format for this version differs from the one used here. + +🔥🔥 **News**: ```2025/03/24```: We have launched [CogKit](https://github.com/THUDM/CogKit), a fine-tuning and inference framework for the CogView4 and CogVideoX series. We recommend migrate to CogKit, as future fine-tuning work for the CogVideo series will primarily be maintained within CogKit. + +## Hardware Requirements + +| Model | Training Type | Distribution Strategy | Mixed Precision | Training Resolution (FxHxW) | Hardware Requirements | +|----------------------------|----------------|--------------------------------------|-----------------|-----------------------------|-------------------------| +| cogvideox-t2v-2b | lora (rank128) | DDP | fp16 | 49x480x720 | 16GB VRAM (NVIDIA 4080) | +| cogvideox-{t2v, i2v}-5b | lora (rank128) | DDP | bf16 | 49x480x720 | 24GB VRAM (NVIDIA 4090) | +| cogvideox1.5-{t2v, i2v}-5b | lora (rank128) | DDP | bf16 | 81x768x1360 | 35GB VRAM (NVIDIA A100) | +| cogvideox-t2v-2b | sft | DDP | fp16 | 49x480x720 | 36GB VRAM (NVIDIA A100) | +| cogvideox-t2v-2b | sft | 1-GPU zero-2 + opt offload | fp16 | 49x480x720 | 17GB VRAM (NVIDIA 4090) | +| cogvideox-t2v-2b | sft | 8-GPU zero-2 | fp16 | 49x480x720 | 17GB VRAM (NVIDIA 4090) | +| cogvideox-t2v-2b | sft | 8-GPU zero-3 | fp16 | 49x480x720 | 19GB VRAM (NVIDIA 4090) | +| cogvideox-t2v-2b | sft | 8-GPU zero-3 + opt and param offload | bf16 | 49x480x720 | 14GB VRAM (NVIDIA 4080) | +| cogvideox-{t2v, i2v}-5b | sft | 1-GPU zero-2 + opt offload | bf16 | 49x480x720 | 42GB VRAM (NVIDIA A100) | +| cogvideox-{t2v, i2v}-5b | sft | 8-GPU zero-2 | bf16 | 49x480x720 | 42GB VRAM (NVIDIA 4090) | +| cogvideox-{t2v, i2v}-5b | sft | 8-GPU zero-3 | bf16 | 49x480x720 | 43GB VRAM (NVIDIA 4090) | +| cogvideox-{t2v, i2v}-5b | sft | 8-GPU zero-3 + opt and param offload | bf16 | 49x480x720 | 28GB VRAM (NVIDIA 5090) | +| cogvideox1.5-{t2v, i2v}-5b | sft | 1-GPU zero-2 + opt offload | bf16 | 81x768x1360 | 56GB VRAM (NVIDIA A100) | +| cogvideox1.5-{t2v, i2v}-5b | sft | 8-GPU zero-2 | bf16 | 81x768x1360 | 55GB VRAM (NVIDIA A100) | +| cogvideox1.5-{t2v, i2v}-5b | sft | 8-GPU zero-3 | bf16 | 81x768x1360 | 55GB VRAM (NVIDIA A100) | +| cogvideox1.5-{t2v, i2v}-5b | sft | 8-GPU zero-3 + opt and param offload | bf16 | 81x768x1360 | 40GB VRAM (NVIDIA A100) | + +## Install Dependencies + +Since the relevant code has not yet been merged into the official `diffusers` release, you need to fine-tune based on +the diffusers branch. Follow the steps below to install the dependencies: + +```shell +git clone https://github.com/huggingface/diffusers.git +cd diffusers # Now on the Main branch +pip install -e . +``` + +## Prepare the Dataset + +First, you need to prepare your dataset. Depending on your task type (T2V or I2V), the dataset format will vary +slightly: + +``` +. +├── prompts.txt +├── videos +├── videos.txt +├── images # (Optional) For I2V, if not provided, first frame will be extracted from video as reference +└── images.txt # (Optional) For I2V, if not provided, first frame will be extracted from video as reference +``` + +Where: + +- `prompts.txt`: Contains the prompts +- `videos/`: Contains the .mp4 video files +- `videos.txt`: Contains the list of video files in the `videos/` directory +- `images/`: (Optional) Contains the .png reference image files +- `images.txt`: (Optional) Contains the list of reference image files + +You can download a sample dataset ( +T2V) [Disney Steamboat Willie](https://huggingface.co/datasets/Wild-Heart/Disney-VideoGeneration-Dataset). + +If you need to use a validation dataset during training, make sure to provide a validation dataset with the same format +as the training dataset. + +## Running Scripts to Start Fine-tuning + +Before starting training, please note the following resolution requirements: + +1. The number of frames must be a multiple of 8 **plus 1** (i.e., 8N+1), such as 49, 81 ... +2. Recommended video resolutions for each model: + - CogVideoX: 480x720 (height x width) + - CogVideoX1.5: 768x1360 (height x width) +3. For samples (videos or images) that don't match the training resolution, the code will directly resize them. This may + cause aspect ratio distortion and affect training results. It's recommended to preprocess your samples (e.g., using + crop + resize to maintain aspect ratio) before training. + +> **Important Note**: To improve training efficiency, we automatically encode videos and cache the results on disk +> before training. If you modify the data after training, please delete the latent directory under the video directory to +> ensure the latest data is used. + +### LoRA + +```bash +# Modify configuration parameters in train_ddp_t2v.sh +# Main parameters to modify: +# --output_dir: Output directory +# --data_root: Dataset root directory +# --caption_column: Path to prompt file +# --image_column: Optional for I2V, path to reference image file list (remove this parameter to use the first frame of video as image condition) +# --video_column: Path to video file list +# --train_resolution: Training resolution (frames x height x width) +# For other important parameters, please refer to the launch script + +bash train_ddp_t2v.sh # Text-to-Video (T2V) fine-tuning +bash train_ddp_i2v.sh # Image-to-Video (I2V) fine-tuning +``` + +### SFT + +We provide several zero configuration templates in the `configs/` directory. Please choose the appropriate training +configuration based on your needs (configure the `deepspeed_config_file` option in `accelerate_config.yaml`). + +```bash +# Parameters to configure are the same as LoRA training + +bash train_zero_t2v.sh # Text-to-Video (T2V) fine-tuning +bash train_zero_i2v.sh # Image-to-Video (I2V) fine-tuning +``` + +In addition to setting the bash script parameters, you need to set the relevant training options in the zero +configuration file and ensure the zero training configuration matches the parameters in the bash script, such as +batch_size, gradient_accumulation_steps, mixed_precision. For details, please refer to +the [DeepSpeed official documentation](https://www.deepspeed.ai/docs/config-json/) + +When using SFT training, please note: + +1. For SFT training, model offload is not used during validation, so the peak VRAM usage may exceed 24GB. For GPUs with + less than 24GB VRAM, it's recommended to disable validation. + +2. Validation is slow when zero-3 is enabled, so it's recommended to disable validation when using zero-3. + +## Load the Fine-tuned Model + ++ Please refer to [cli_demo.py](../inference/cli_demo.py) for instructions on how to load the fine-tuned model. + ++ For SFT trained models, please first use the `zero_to_fp32.py` script in the `checkpoint-*/` directory to merge the + model weights + +## Best Practices + ++ We included 70 training videos with a resolution of `200 x 480 x 720` (frames x height x width). Through frame + skipping in the data preprocessing, we created two smaller datasets with 49 and 16 frames to speed up experiments. The + maximum frame count recommended by the CogVideoX team is 49 frames. These 70 videos were divided into three groups: + 10, 25, and 50 videos, with similar conceptual nature. ++ Videos with 25 or more frames work best for training new concepts and styles. ++ It's recommended to use an identifier token, which can be specified using `--id_token`, for better training results. + This is similar to Dreambooth training, though regular fine-tuning without using this token will still work. ++ The original repository uses `lora_alpha` set to 1. We found that this value performed poorly in several runs, + possibly due to differences in the model backend and training settings. Our recommendation is to set `lora_alpha` to + be equal to the rank or `rank // 2`. ++ It's advised to use a rank of 64 or higher. diff --git a/finetune/README_ja.md b/finetune/README_ja.md new file mode 100644 index 0000000..e4d4641 --- /dev/null +++ b/finetune/README_ja.md @@ -0,0 +1,124 @@ +# CogVideoX Diffusers ファインチューニングガイド + +[中文阅读](./README_zh.md) + +[Read in English](./README.md) + +SATバージョンのファインチューニング手順については、[こちら](../sat/README_zh.md)をご確認ください。このバージョンのデータセットフォーマットは、こちらのバージョンとは異なります。 + +## ハードウェア要件 + +| モデル | トレーニングタイプ | 分散戦略 | 混合トレーニング精度 | トレーニング解像度（フレーム数×高さ×幅） | ハードウェア要件 | +|----------------------------|-------------------|----------------------------------|-----------------|-----------------------------|----------------------------| +| cogvideox-t2v-2b | lora (rank128) | DDP | fp16 | 49x480x720 | 16GB VRAM (NVIDIA 4080) | +| cogvideox-{t2v, i2v}-5b | lora (rank128) | DDP | bf16 | 49x480x720 | 24GB VRAM (NVIDIA 4090) | +| cogvideox1.5-{t2v, i2v}-5b | lora (rank128) | DDP | bf16 | 81x768x1360 | 35GB VRAM (NVIDIA A100) | +| cogvideox-t2v-2b | sft | DDP | fp16 | 49x480x720 | 36GB VRAM (NVIDIA A100) | +| cogvideox-t2v-2b | sft | 1カード zero-2 + オプティマイゼーションオフロード | fp16 | 49x480x720 | 17GB VRAM (NVIDIA 4090) | +| cogvideox-t2v-2b | sft | 8カード zero-2 | fp16 | 49x480x720 | 17GB VRAM (NVIDIA 4090) | +| cogvideox-t2v-2b | sft | 8カード zero-3 | fp16 | 49x480x720 | 19GB VRAM (NVIDIA 4090) | +| cogvideox-t2v-2b | sft | 8カード zero-3 + オプティマイゼーションとパラメータオフロード | bf16 | 49x480x720 | 14GB VRAM (NVIDIA 4080) | +| cogvideox-{t2v, i2v}-5b | sft | 1カード zero-2 + オプティマイゼーションオフロード | bf16 | 49x480x720 | 42GB VRAM (NVIDIA A100) | +| cogvideox-{t2v, i2v}-5b | sft | 8カード zero-2 | bf16 | 49x480x720 | 42GB VRAM (NVIDIA 4090) | +| cogvideox-{t2v, i2v}-5b | sft | 8カード zero-3 | bf16 | 49x480x720 | 43GB VRAM (NVIDIA 4090) | +| cogvideox-{t2v, i2v}-5b | sft | 8カード zero-3 + オプティマイゼーションとパラメータオフロード | bf16 | 49x480x720 | 28GB VRAM (NVIDIA 5090) | +| cogvideox1.5-{t2v, i2v}-5b | sft | 1カード zero-2 + オプティマイゼーションオフロード | bf16 | 81x768x1360 | 56GB VRAM (NVIDIA A100) | +| cogvideox1.5-{t2v, i2v}-5b | sft | 8カード zero-2 | bf16 | 81x768x1360 | 55GB VRAM (NVIDIA A100) | +| cogvideox1.5-{t2v, i2v}-5b | sft | 8カード zero-3 | bf16 | 81x768x1360 | 55GB VRAM (NVIDIA A100) | +| cogvideox1.5-{t2v, i2v}-5b | sft | 8カード zero-3 + オプティマイゼーションとパラメータオフロード | bf16 | 81x768x1360 | 40GB VRAM (NVIDIA A100) | + + + +## 依存関係のインストール + +関連するコードがまだ `diffusers` の公式リリースに統合されていないため、`diffusers` ブランチを基にファインチューニングを行う必要があります。以下の手順に従って依存関係をインストールしてください： + +```shell +git clone https://github.com/huggingface/diffusers.git +cd diffusers # 現在は Main ブランチ +pip install -e . +``` + +## データセットの準備 + +まず、データセットを準備する必要があります。タスクの種類（T2V または I2V）によって、データセットのフォーマットが少し異なります： + +``` +. +├── prompts.txt +├── videos +├── videos.txt +├── images # (オプション) I2Vの場合。提供されない場合、動画の最初のフレームが参照画像として使用されます +└── images.txt # (オプション) I2Vの場合。提供されない場合、動画の最初のフレームが参照画像として使用されます +``` + +各ファイルの役割は以下の通りです： +- `prompts.txt`: プロンプトを格納 +- `videos/`: .mp4 動画ファイルを格納 +- `videos.txt`: `videos/` フォルダ内の動画ファイルリストを格納 +- `images/`: (オプション) .png 形式の参照画像ファイル +- `images.txt`: (オプション) 参照画像ファイルリスト + +トレーニング中に検証データセットを使用する場合は、トレーニングデータセットと同じフォーマットで検証データセットを提供する必要があります。 + +## スクリプトを実行してファインチューニングを開始 + +トレーニングを開始する前に、以下の解像度設定要件に注意してください： + +1. フレーム数は8の倍数 **+1** (つまり8N+1) でなければなりません。例：49, 81 ... +2. ビデオ解像度はモデルのデフォルトサイズを使用することをお勧めします： + - CogVideoX: 480x720 (高さ×幅) + - CogVideoX1.5: 768x1360 (高さ×幅) +3. トレーニング解像度に合わないサンプル（ビデオや画像）はコード内で自動的にリサイズされます。このため、サンプルのアスペクト比が変形し、トレーニング効果に影響を与える可能性があります。解像度に関しては、事前にサンプルを処理（例えば、アスペクト比を維持するためにクロップ＋リサイズを使用）してからトレーニングを行うことをお勧めします。 + +> **重要な注意**：トレーニング効率を高めるため、トレーニング前にビデオをエンコードし、その結果をディスクにキャッシュします。トレーニング後にデータを変更した場合は、`video`ディレクトリ内の`latent`ディレクトリを削除して、最新のデータを使用するようにしてください。 + +### LoRA + +```bash +# train_ddp_t2v.sh の設定パラメータを変更 +# 主に以下のパラメータを変更する必要があります: +# --output_dir: 出力ディレクトリ +# --data_root: データセットのルートディレクトリ +# --caption_column: テキストプロンプトのファイルパス +# --image_column: I2Vの場合、参照画像のファイルリストのパス（このパラメータを削除すると、デフォルトで動画の最初のフレームが画像条件として使用されます） +# --video_column: 動画ファイルのリストのパス +# --train_resolution: トレーニング解像度（フレーム数×高さ×幅） +# その他の重要なパラメータについては、起動スクリプトを参照してください + +bash train_ddp_t2v.sh # テキストから動画（T2V）微調整 +bash train_ddp_i2v.sh # 画像から動画（I2V）微調整 +``` + +### SFT + +`configs/`ディレクトリにはいくつかのZero構成テンプレートが提供されています。必要に応じて適切なトレーニング設定を選択してください（`accelerate_config.yaml`で`deepspeed_config_file`オプションを設定します）。 + +```bash +# 設定するパラメータはLoRAトレーニングと同様です + +bash train_zero_t2v.sh # テキストから動画（T2V）微調整 +bash train_zero_i2v.sh # 画像から動画（I2V）微調整 +``` + +Bashスクリプトの関連パラメータを設定するだけでなく、Zeroの設定ファイルでトレーニングオプションを設定し、Zeroのトレーニング設定がBashスクリプト内のパラメータと一致していることを確認する必要があります。例えば、`batch_size`、`gradient_accumulation_steps`、`mixed_precision`など、具体的な詳細は[DeepSpeed公式ドキュメント](https://www.deepspeed.ai/docs/config-json/)を参照してください。 + +SFTトレーニングを使用する際に注意すべき点： + +1. SFTトレーニングでは、検証時にモデルオフロードは使用されません。そのため、24GB以下のGPUでは検証時にVRAMのピークが24GBを超える可能性があります。24GB以下のGPUでは、検証を無効にすることをお勧めします。 + +2. Zero-3を有効にすると検証が遅くなるため、Zero-3では検証を無効にすることをお勧めします。 + +## ファインチューニングしたモデルの読み込み + ++ ファインチューニングしたモデルを読み込む方法については、[cli_demo.py](../inference/cli_demo.py)を参照してください。 + ++ SFTトレーニングのモデルについては、まず`checkpoint-*`/ディレクトリ内の`zero_to_fp32.py`スクリプトを使用して、モデルの重みを統合してください。 + +## ベストプラクティス + ++ 解像度が `200 x 480 x 720`（フレーム数 x 高さ x 幅）の70本のトレーニング動画を使用しました。データ前処理でフレームスキップを行い、49フレームおよび16フレームの2つの小さなデータセットを作成して実験速度を向上させました。CogVideoXチームの推奨最大フレーム数制限は49フレームです。これらの70本の動画は、10、25、50本の3つのグループに分け、概念的に類似した性質のものです。 ++ 25本以上の動画を使用することで、新しい概念やスタイルのトレーニングが最適です。 ++ `--id_token` で指定できる識別子トークンを使用すると、トレーニング効果がより良くなります。これはDreamboothトレーニングに似ていますが、このトークンを使用しない通常のファインチューニングでも問題なく動作します。 ++ 元のリポジトリでは `lora_alpha` が1に設定されていますが、この値は多くの実行で効果が悪かったため、モデルのバックエンドやトレーニング設定の違いが影響している可能性があります。私たちの推奨は、`lora_alpha` を rank と同じか、`rank // 2` に設定することです。 ++ rank は64以上に設定することをお勧めします。 diff --git a/finetune/README_zh.md b/finetune/README_zh.md new file mode 100644 index 0000000..ba20114 --- /dev/null +++ b/finetune/README_zh.md @@ -0,0 +1,130 @@ +# CogVideoX diffusers 微调方案 + +[Read this in English](./README.md) + +[日本語で読む](./README_ja.md) + +如果您想查看SAT版本微调，请查看[这里](../sat/README_zh.md)。其数据集格式与本版本不同。 + +🔥🔥 News：我们已正式发布[CogKit](https://github.com/THUDM/CogKit)，这是一个专为CogView4和CogVideoX系列设计的微调与推理框架。我们建议迁移至CogKit，因为未来与CogVideo系列相关的微调工作将主要在CogKit中维护。 + +## 硬件要求 + +| 模型 | 训练类型 | 分布式策略 | 混合训练精度 | 训练分辨率(帧数x高x宽) | 硬件要求 | +|----------------------------|----------------|-----------------------------------|------------|----------------------|-----------------------| +| cogvideox-t2v-2b | lora (rank128) | DDP | fp16 | 49x480x720 | 16G显存 (NVIDIA 4080) | +| cogvideox-{t2v, i2v}-5b | lora (rank128) | DDP | bf16 | 49x480x720 | 24G显存 (NVIDIA 4090) | +| cogvideox1.5-{t2v, i2v}-5b | lora (rank128) | DDP | bf16 | 81x768x1360 | 35G显存 (NVIDIA A100) | +| cogvideox-t2v-2b | sft | DDP | fp16 | 49x480x720 | 36G显存 (NVIDIA A100) | +| cogvideox-t2v-2b | sft | 1卡zero-2 + opt offload | fp16 | 49x480x720 | 17G显存 (NVIDIA 4090) | +| cogvideox-t2v-2b | sft | 8卡zero-2 | fp16 | 49x480x720 | 17G显存 (NVIDIA 4090) | +| cogvideox-t2v-2b | sft | 8卡zero-3 | fp16 | 49x480x720 | 19G显存 (NVIDIA 4090) | +| cogvideox-t2v-2b | sft | 8卡zero-3 + opt and param offload | bf16 | 49x480x720 | 14G显存 (NVIDIA 4080) | +| cogvideox-{t2v, i2v}-5b | sft | 1卡zero-2 + opt offload | bf16 | 49x480x720 | 42G显存 (NVIDIA A100) | +| cogvideox-{t2v, i2v}-5b | sft | 8卡zero-2 | bf16 | 49x480x720 | 42G显存 (NVIDIA 4090) | +| cogvideox-{t2v, i2v}-5b | sft | 8卡zero-3 | bf16 | 49x480x720 | 43G显存 (NVIDIA 4090) | +| cogvideox-{t2v, i2v}-5b | sft | 8卡zero-3 + opt and param offload | bf16 | 49x480x720 | 28G显存 (NVIDIA 5090) | +| cogvideox1.5-{t2v, i2v}-5b | sft | 1卡zero-2 + opt offload | bf16 | 81x768x1360 | 56G显存 (NVIDIA A100) | +| cogvideox1.5-{t2v, i2v}-5b | sft | 8卡zero-2 | bf16 | 81x768x1360 | 55G显存 (NVIDIA A100) | +| cogvideox1.5-{t2v, i2v}-5b | sft | 8卡zero-3 | bf16 | 81x768x1360 | 55G显存 (NVIDIA A100) | +| cogvideox1.5-{t2v, i2v}-5b | sft | 8卡zero-3 + opt and param offload | bf16 | 81x768x1360 | 40G显存 (NVIDIA A100) | + + +## 安装依赖 + +由于相关代码还没有被合并到diffusers发行版，你需要基于diffusers分支进行微调。请按照以下步骤安装依赖： + +```shell +git clone https://github.com/huggingface/diffusers.git +cd diffusers # Now in Main branch +pip install -e . +``` + +## 准备数据集 + +首先，你需要准备数据集。根据你的任务类型（T2V 或 I2V），数据集格式略有不同： + +``` +. +├── prompts.txt +├── videos +├── videos.txt +├── images # (可选) 对于I2V，若不提供，则从视频中提取第一帧作为参考图像 +└── images.txt # (可选) 对于I2V，若不提供，则从视频中提取第一帧作为参考图像 +``` + +其中： +- `prompts.txt`: 存放提示词 +- `videos/`: 存放.mp4视频文件 +- `videos.txt`: 存放 videos 目录中的视频文件列表 +- `images/`: (可选) 存放.png参考图像文件 +- `images.txt`: (可选) 存放参考图像文件列表 + +你可以从这里下载示例数据集(T2V) [迪士尼汽船威利号](https://huggingface.co/datasets/Wild-Heart/Disney-VideoGeneration-Dataset) + +如果需要在训练过程中进行validation，则需要额外提供验证数据集，其中数据格式与训练集相同。 + +## 运行脚本，开始微调 + +在开始训练之前，请注意以下分辨率设置要求： + +1. 帧数必须是8的倍数 **+1** (即8N+1), 例如49, 81 ... +2. 视频分辨率建议使用模型的默认大小： + - CogVideoX: 480x720 (高x宽) + - CogVideoX1.5: 768x1360 (高x宽) +3. 对于不满足训练分辨率的样本（视频或图片）在代码中会直接进行resize。这可能会导致样本的宽高比发生形变从而影响训练效果。建议用户提前对样本在分辨率上进行处理（例如使用crop + resize来维持宽高比）再进行训练。 + +> **重要提示**：为了提高训练效率，我们会在训练前自动对video进行encode并将结果缓存在磁盘。如果在训练后修改了数据，请删除video目录下的latent目录，以确保使用最新的数据。 + +### LoRA + +```bash +# 修改 train_ddp_t2v.sh 中的配置参数 +# 主要需要修改以下参数: +# --output_dir: 输出目录 +# --data_root: 数据集根目录 +# --caption_column: 提示词文件路径 +# --image_column: I2V可选，参考图像文件列表路径 (移除这个参数将默认使用视频第一帧作为image condition) +# --video_column: 视频文件列表路径 +# --train_resolution: 训练分辨率 (帧数x高x宽) +# 其他重要参数请参考启动脚本 + +bash train_ddp_t2v.sh # 文本生成视频 (T2V) 微调 +bash train_ddp_i2v.sh # 图像生成视频 (I2V) 微调 +``` + +### SFT + +我们在`configs/`目录中提供了几个zero配置的模版，请根据你的需求选择合适的训练配置（在`accelerate_config.yaml`中配置`deepspeed_config_file`选项即可）。 + +```bash +# 需要配置的参数与LoRA训练同理 + +bash train_zero_t2v.sh # 文本生成视频 (T2V) 微调 +bash train_zero_i2v.sh # 图像生成视频 (I2V) 微调 +``` + +除了设置bash脚本的相关参数，你还需要在zero的配置文件中设定相关的训练选项，并确保zero的训练配置与bash脚本中的参数一致，例如batch_size，gradient_accumulation_steps，mixed_precision，具体细节请参考[deepspeed官方文档](https://www.deepspeed.ai/docs/config-json/) + +在使用sft训练时，有以下几点需要注意： + +1. 对于sft训练，validation时不会使用model offload，因此显存峰值可能会超出24GB，所以对于24GB以下的显卡，建议关闭validation。 + +2. 开启zero-3时validation会比较慢，建议在zero-3下关闭validation。 + + +## 载入微调的模型 + ++ 请关注[cli_demo.py](../inference/cli_demo.py) 以了解如何加载微调的模型。 + ++ 对于sft训练的模型，请先使用`checkpoint-*/`目录下的`zero_to_fp32.py`脚本合并模型权重 + +## 最佳实践 + ++ 包含70个分辨率为 `200 x 480 x 720`（帧数 x 高 x + 宽）的训练视频。通过数据预处理中的帧跳过，我们创建了两个较小的49帧和16帧数据集，以加快实验速度，因为CogVideoX团队建议的最大帧数限制是49帧。我们将70个视频分成三组，分别为10、25和50个视频。这些视频的概念性质相似。 ++ 25个及以上的视频在训练新概念和风格时效果最佳。 ++ 现使用可以通过 `--id_token` 指定的标识符token进行训练效果更好。这类似于 Dreambooth 训练，但不使用这种token的常规微调也可以工作。 ++ 原始仓库使用 `lora_alpha` 设置为 1。我们发现这个值在多次运行中效果不佳，可能是因为模型后端和训练设置的不同。我们的建议是将 + lora_alpha 设置为与 rank 相同或 rank // 2。 ++ 建议使用 rank 为 64 及以上的设置。 diff --git a/finetune/accelerate_config.yaml b/finetune/accelerate_config.yaml new file mode 100644 index 0000000..7b7750c --- /dev/null +++ b/finetune/accelerate_config.yaml @@ -0,0 +1,21 @@ +compute_environment: LOCAL_MACHINE + +gpu_ids: "0,1,2,3,4,5,6,7" +num_processes: 8 # should be the same as the number of GPUs + +debug: false +deepspeed_config: + deepspeed_config_file: configs/zero2.yaml # e.g. configs/zero2.yaml, need use absolute path + zero3_init_flag: false +distributed_type: DEEPSPEED +downcast_bf16: 'no' +enable_cpu_affinity: false +machine_rank: 0 +main_training_function: main +num_machines: 1 +rdzv_backend: static +same_network: true +tpu_env: [] +tpu_use_cluster: false +tpu_use_sudo: false +use_cpu: false diff --git a/finetune/configs/zero2.yaml b/finetune/configs/zero2.yaml new file mode 100644 index 0000000..b056bd4 --- /dev/null +++ b/finetune/configs/zero2.yaml @@ -0,0 +1,38 @@ +{ + "bf16": { + "enabled": true + }, + "optimizer": { + "type": "AdamW", + "params": { + "lr": "auto", + "weight_decay": "auto", + "torch_adam": true, + "adam_w_mode": true + } + }, + "scheduler": { + "type": "WarmupDecayLR", + "params": { + "warmup_min_lr": "auto", + "warmup_max_lr": "auto", + "warmup_num_steps": "auto", + "total_num_steps": "auto" + } + }, + "zero_optimization": { + "stage": 2, + "allgather_partitions": true, + "allgather_bucket_size": 2e8, + "overlap_comm": true, + "reduce_scatter": true, + "reduce_bucket_size": 5e8, + "contiguous_gradients": true + }, + "gradient_accumulation_steps": 1, + "train_micro_batch_size_per_gpu": 1, + "train_batch_size": "auto", + "gradient_clipping": "auto", + "steps_per_print": 2000, + "wall_clock_breakdown": false +} diff --git a/finetune/configs/zero2_offload.yaml b/finetune/configs/zero2_offload.yaml new file mode 100644 index 0000000..24fdcb4 --- /dev/null +++ b/finetune/configs/zero2_offload.yaml @@ -0,0 +1,42 @@ +{ + "bf16": { + "enabled": true + }, + "optimizer": { + "type": "AdamW", + "params": { + "lr": "auto", + "weight_decay": "auto", + "torch_adam": true, + "adam_w_mode": true + } + }, + "scheduler": { + "type": "WarmupDecayLR", + "params": { + "warmup_min_lr": "auto", + "warmup_max_lr": "auto", + "warmup_num_steps": "auto", + "total_num_steps": "auto" + } + }, + "zero_optimization": { + "stage": 2, + "allgather_partitions": true, + "allgather_bucket_size": 2e8, + "overlap_comm": true, + "reduce_scatter": true, + "reduce_bucket_size": 5e8, + "contiguous_gradients": true, + "offload_optimizer": { + "device": "cpu", + "pin_memory": true + } + }, + "gradient_accumulation_steps": 1, + "train_micro_batch_size_per_gpu": 1, + "train_batch_size": "auto", + "gradient_clipping": "auto", + "steps_per_print": 2000, + "wall_clock_breakdown": false +} diff --git a/finetune/configs/zero3.yaml b/finetune/configs/zero3.yaml new file mode 100644 index 0000000..69c5fd5 --- /dev/null +++ b/finetune/configs/zero3.yaml @@ -0,0 +1,43 @@ +{ + "bf16": { + "enabled": true + }, + "optimizer": { + "type": "AdamW", + "params": { + "lr": "auto", + "weight_decay": "auto", + "torch_adam": true, + "adam_w_mode": true + } + }, + "scheduler": { + "type": "WarmupDecayLR", + "params": { + "warmup_min_lr": "auto", + "warmup_max_lr": "auto", + "warmup_num_steps": "auto", + "total_num_steps": "auto" + } + }, + "zero_optimization": { + "stage": 3, + "overlap_comm": true, + "contiguous_gradients": true, + "reduce_bucket_size": 5e8, + "stage3_prefetch_bucket_size": "auto", + "stage3_param_persistence_threshold": "auto", + "sub_group_size": 1e9, + "stage3_max_live_parameters": 1e9, + "stage3_max_reuse_distance": 1e9, + "stage3_gather_16bit_weights_on_model_save": "auto", + "stage3_prefetch_bucket_size": 5e8, + "stage3_param_persistence_threshold": 1e5 + }, + "gradient_accumulation_steps": 1, + "train_micro_batch_size_per_gpu": 1, + "train_batch_size": "auto", + "gradient_clipping": "auto", + "steps_per_print": 2000, + "wall_clock_breakdown": false +} diff --git a/finetune/configs/zero3_offload.yaml b/finetune/configs/zero3_offload.yaml new file mode 100644 index 0000000..58e6529 --- /dev/null +++ b/finetune/configs/zero3_offload.yaml @@ -0,0 +1,51 @@ +{ + "bf16": { + "enabled": true + }, + "optimizer": { + "type": "AdamW", + "params": { + "lr": "auto", + "weight_decay": "auto", + "torch_adam": true, + "adam_w_mode": true + } + }, + "scheduler": { + "type": "WarmupDecayLR", + "params": { + "warmup_min_lr": "auto", + "warmup_max_lr": "auto", + "warmup_num_steps": "auto", + "total_num_steps": "auto" + } + }, + "zero_optimization": { + "stage": 3, + "offload_optimizer": { + "device": "cpu", + "pin_memory": true + }, + "offload_param": { + "device": "cpu", + "pin_memory": true + }, + "overlap_comm": true, + "contiguous_gradients": true, + "reduce_bucket_size": 5e8, + "stage3_prefetch_bucket_size": "auto", + "stage3_param_persistence_threshold": "auto", + "sub_group_size": 1e9, + "stage3_max_live_parameters": 1e9, + "stage3_max_reuse_distance": 1e9, + "stage3_gather_16bit_weights_on_model_save": "auto", + "stage3_prefetch_bucket_size": 5e8, + "stage3_param_persistence_threshold": 1e6 + }, + "gradient_accumulation_steps": 1, + "train_micro_batch_size_per_gpu": 1, + "train_batch_size": "auto", + "gradient_clipping": "auto", + "steps_per_print": 2000, + "wall_clock_breakdown": false +} diff --git a/finetune/constants.py b/finetune/constants.py new file mode 100644 index 0000000..89dde58 --- /dev/null +++ b/finetune/constants.py @@ -0,0 +1,2 @@ +LOG_NAME = "trainer" +LOG_LEVEL = "INFO" diff --git a/finetune/datasets/__init__.py b/finetune/datasets/__init__.py new file mode 100644 index 0000000..25f24e1 --- /dev/null +++ b/finetune/datasets/__init__.py @@ -0,0 +1,12 @@ +from .bucket_sampler import BucketSampler +from .i2v_dataset import I2VDatasetWithBuckets, I2VDatasetWithResize +from .t2v_dataset import T2VDatasetWithBuckets, T2VDatasetWithResize + + +__all__ = [ + "I2VDatasetWithResize", + "I2VDatasetWithBuckets", + "T2VDatasetWithResize", + "T2VDatasetWithBuckets", + "BucketSampler", +] diff --git a/finetune/datasets/bucket_sampler.py b/finetune/datasets/bucket_sampler.py new file mode 100644 index 0000000..3e8c5c9 --- /dev/null +++ b/finetune/datasets/bucket_sampler.py @@ -0,0 +1,79 @@ +import logging +import random + +from torch.utils.data import Dataset, Sampler + + +logger = logging.getLogger(__name__) + + +class BucketSampler(Sampler): + r""" + PyTorch Sampler that groups 3D data by height, width and frames. + + Args: + data_source (`VideoDataset`): + A PyTorch dataset object that is an instance of `VideoDataset`. + batch_size (`int`, defaults to `8`): + The batch size to use for training. + shuffle (`bool`, defaults to `True`): + Whether or not to shuffle the data in each batch before dispatching to dataloader. + drop_last (`bool`, defaults to `False`): + Whether or not to drop incomplete buckets of data after completely iterating over all data + in the dataset. If set to True, only batches that have `batch_size` number of entries will + be yielded. If set to False, it is guaranteed that all data in the dataset will be processed + and batches that do not have `batch_size` number of entries will also be yielded. + """ + + def __init__( + self, + data_source: Dataset, + batch_size: int = 8, + shuffle: bool = True, + drop_last: bool = False, + ) -> None: + self.data_source = data_source + self.batch_size = batch_size + self.shuffle = shuffle + self.drop_last = drop_last + + self.buckets = {resolution: [] for resolution in data_source.video_resolution_buckets} + + self._raised_warning_for_drop_last = False + + def __len__(self): + if self.drop_last and not self._raised_warning_for_drop_last: + self._raised_warning_for_drop_last = True + logger.warning( + "Calculating the length for bucket sampler is not possible when `drop_last` is set to True. This may cause problems when setting the number of epochs used for training." + ) + return (len(self.data_source) + self.batch_size - 1) // self.batch_size + + def __iter__(self): + for index, data in enumerate(self.data_source): + video_metadata = data["video_metadata"] + f, h, w = ( + video_metadata["num_frames"], + video_metadata["height"], + video_metadata["width"], + ) + + self.buckets[(f, h, w)].append(data) + if len(self.buckets[(f, h, w)]) == self.batch_size: + if self.shuffle: + random.shuffle(self.buckets[(f, h, w)]) + yield self.buckets[(f, h, w)] + del self.buckets[(f, h, w)] + self.buckets[(f, h, w)] = [] + + if self.drop_last: + return + + for fhw, bucket in list(self.buckets.items()): + if len(bucket) == 0: + continue + if self.shuffle: + random.shuffle(bucket) + yield bucket + del self.buckets[fhw] + self.buckets[fhw] = [] diff --git a/finetune/datasets/i2v_dataset.py b/finetune/datasets/i2v_dataset.py new file mode 100644 index 0000000..f1f1d48 --- /dev/null +++ b/finetune/datasets/i2v_dataset.py @@ -0,0 +1,325 @@ +import hashlib +from pathlib import Path +from typing import TYPE_CHECKING, Any, Dict, List, Tuple + +import torch +from accelerate.logging import get_logger +from safetensors.torch import load_file, save_file +from torch.utils.data import Dataset +from torchvision import transforms +from typing_extensions import override + +from finetune.constants import LOG_LEVEL, LOG_NAME + +from .utils import ( + load_images, + load_images_from_videos, + load_prompts, + load_videos, + preprocess_image_with_resize, + preprocess_video_with_buckets, + preprocess_video_with_resize, +) + + +if TYPE_CHECKING: + from finetune.trainer import Trainer + +# Must import after torch because this can sometimes lead to a nasty segmentation fault, or stack smashing error +# Very few bug reports but it happens. Look in decord Github issues for more relevant information. +import decord # isort:skip + +decord.bridge.set_bridge("torch") + +logger = get_logger(LOG_NAME, LOG_LEVEL) + + +class BaseI2VDataset(Dataset): + """ + Base dataset class for Image-to-Video (I2V) training. + + This dataset loads prompts, videos and corresponding conditioning images for I2V training. + + Args: + data_root (str): Root directory containing the dataset files + caption_column (str): Path to file containing text prompts/captions + video_column (str): Path to file containing video paths + image_column (str): Path to file containing image paths + device (torch.device): Device to load the data on + encode_video_fn (Callable[[torch.Tensor], torch.Tensor], optional): Function to encode videos + """ + + def __init__( + self, + data_root: str, + caption_column: str, + video_column: str, + image_column: str | None, + device: torch.device, + trainer: "Trainer" = None, + *args, + **kwargs, + ) -> None: + super().__init__() + + data_root = Path(data_root) + self.prompts = load_prompts(data_root / caption_column) + self.videos = load_videos(data_root / video_column) + if image_column is not None: + self.images = load_images(data_root / image_column) + else: + self.images = load_images_from_videos(self.videos) + self.trainer = trainer + + self.device = device + self.encode_video = trainer.encode_video + self.encode_text = trainer.encode_text + + # Check if number of prompts matches number of videos and images + if not (len(self.videos) == len(self.prompts) == len(self.images)): + raise ValueError( + f"Expected length of prompts, videos and images to be the same but found {len(self.prompts)=}, {len(self.videos)=} and {len(self.images)=}. Please ensure that the number of caption prompts, videos and images match in your dataset." + ) + + # Check if all video files exist + if any(not path.is_file() for path in self.videos): + raise ValueError( + f"Some video files were not found. Please ensure that all video files exist in the dataset directory. Missing file: {next(path for path in self.videos if not path.is_file())}" + ) + + # Check if all image files exist + if any(not path.is_file() for path in self.images): + raise ValueError( + f"Some image files were not found. Please ensure that all image files exist in the dataset directory. Missing file: {next(path for path in self.images if not path.is_file())}" + ) + + def __len__(self) -> int: + return len(self.videos) + + def __getitem__(self, index: int) -> Dict[str, Any]: + if isinstance(index, list): + # Here, index is actually a list of data objects that we need to return. + # The BucketSampler should ideally return indices. But, in the sampler, we'd like + # to have information about num_frames, height and width. Since this is not stored + # as metadata, we need to read the video to get this information. You could read this + # information without loading the full video in memory, but we do it anyway. In order + # to not load the video twice (once to get the metadata, and once to return the loaded video + # based on sampled indices), we cache it in the BucketSampler. When the sampler is + # to yield, we yield the cache data instead of indices. So, this special check ensures + # that data is not loaded a second time. PRs are welcome for improvements. + return index + + prompt = self.prompts[index] + video = self.videos[index] + image = self.images[index] + train_resolution_str = "x".join(str(x) for x in self.trainer.args.train_resolution) + + cache_dir = self.trainer.args.data_root / "cache" + video_latent_dir = ( + cache_dir / "video_latent" / self.trainer.args.model_name / train_resolution_str + ) + prompt_embeddings_dir = cache_dir / "prompt_embeddings" + video_latent_dir.mkdir(parents=True, exist_ok=True) + prompt_embeddings_dir.mkdir(parents=True, exist_ok=True) + + prompt_hash = str(hashlib.sha256(prompt.encode()).hexdigest()) + prompt_embedding_path = prompt_embeddings_dir / (prompt_hash + ".safetensors") + encoded_video_path = video_latent_dir / (video.stem + ".safetensors") + + if prompt_embedding_path.exists(): + prompt_embedding = load_file(prompt_embedding_path)["prompt_embedding"] + logger.debug( + f"process {self.trainer.accelerator.process_index}: Loaded prompt embedding from {prompt_embedding_path}", + main_process_only=False, + ) + else: + prompt_embedding = self.encode_text(prompt) + prompt_embedding = prompt_embedding.to("cpu") + # [1, seq_len, hidden_size] -> [seq_len, hidden_size] + prompt_embedding = prompt_embedding[0] + save_file({"prompt_embedding": prompt_embedding}, prompt_embedding_path) + logger.info( + f"Saved prompt embedding to {prompt_embedding_path}", main_process_only=False + ) + + if encoded_video_path.exists(): + encoded_video = load_file(encoded_video_path)["encoded_video"] + logger.debug(f"Loaded encoded video from {encoded_video_path}", main_process_only=False) + # shape of image: [C, H, W] + _, image = self.preprocess(None, self.images[index]) + image = self.image_transform(image) + else: + frames, image = self.preprocess(video, image) + frames = frames.to(self.device) + image = image.to(self.device) + image = self.image_transform(image) + # Current shape of frames: [F, C, H, W] + frames = self.video_transform(frames) + + # Convert to [B, C, F, H, W] + frames = frames.unsqueeze(0) + frames = frames.permute(0, 2, 1, 3, 4).contiguous() + encoded_video = self.encode_video(frames) + + # [1, C, F, H, W] -> [C, F, H, W] + encoded_video = encoded_video[0] + encoded_video = encoded_video.to("cpu") + image = image.to("cpu") + save_file({"encoded_video": encoded_video}, encoded_video_path) + logger.info(f"Saved encoded video to {encoded_video_path}", main_process_only=False) + + # shape of encoded_video: [C, F, H, W] + # shape of image: [C, H, W] + return { + "image": image, + "prompt_embedding": prompt_embedding, + "encoded_video": encoded_video, + "video_metadata": { + "num_frames": encoded_video.shape[1], + "height": encoded_video.shape[2], + "width": encoded_video.shape[3], + }, + } + + def preprocess( + self, video_path: Path | None, image_path: Path | None + ) -> Tuple[torch.Tensor, torch.Tensor]: + """ + Loads and preprocesses a video and an image. + If either path is None, no preprocessing will be done for that input. + + Args: + video_path: Path to the video file to load + image_path: Path to the image file to load + + Returns: + A tuple containing: + - video(torch.Tensor) of shape [F, C, H, W] where F is number of frames, + C is number of channels, H is height and W is width + - image(torch.Tensor) of shape [C, H, W] + """ + raise NotImplementedError("Subclass must implement this method") + + def video_transform(self, frames: torch.Tensor) -> torch.Tensor: + """ + Applies transformations to a video. + + Args: + frames (torch.Tensor): A 4D tensor representing a video + with shape [F, C, H, W] where: + - F is number of frames + - C is number of channels (3 for RGB) + - H is height + - W is width + + Returns: + torch.Tensor: The transformed video tensor + """ + raise NotImplementedError("Subclass must implement this method") + + def image_transform(self, image: torch.Tensor) -> torch.Tensor: + """ + Applies transformations to an image. + + Args: + image (torch.Tensor): A 3D tensor representing an image + with shape [C, H, W] where: + - C is number of channels (3 for RGB) + - H is height + - W is width + + Returns: + torch.Tensor: The transformed image tensor + """ + raise NotImplementedError("Subclass must implement this method") + + +class I2VDatasetWithResize(BaseI2VDataset): + """ + A dataset class for image-to-video generation that resizes inputs to fixed dimensions. + + This class preprocesses videos and images by resizing them to specified dimensions: + - Videos are resized to max_num_frames x height x width + - Images are resized to height x width + + Args: + max_num_frames (int): Maximum number of frames to extract from videos + height (int): Target height for resizing videos and images + width (int): Target width for resizing videos and images + """ + + def __init__(self, max_num_frames: int, height: int, width: int, *args, **kwargs) -> None: + super().__init__(*args, **kwargs) + + self.max_num_frames = max_num_frames + self.height = height + self.width = width + + self.__frame_transforms = transforms.Compose( + [transforms.Lambda(lambda x: x / 255.0 * 2.0 - 1.0)] + ) + self.__image_transforms = self.__frame_transforms + + @override + def preprocess( + self, video_path: Path | None, image_path: Path | None + ) -> Tuple[torch.Tensor, torch.Tensor]: + if video_path is not None: + video = preprocess_video_with_resize( + video_path, self.max_num_frames, self.height, self.width + ) + else: + video = None + if image_path is not None: + image = preprocess_image_with_resize(image_path, self.height, self.width) + else: + image = None + return video, image + + @override + def video_transform(self, frames: torch.Tensor) -> torch.Tensor: + return torch.stack([self.__frame_transforms(f) for f in frames], dim=0) + + @override + def image_transform(self, image: torch.Tensor) -> torch.Tensor: + return self.__image_transforms(image) + + +class I2VDatasetWithBuckets(BaseI2VDataset): + def __init__( + self, + video_resolution_buckets: List[Tuple[int, int, int]], + vae_temporal_compression_ratio: int, + vae_height_compression_ratio: int, + vae_width_compression_ratio: int, + *args, + **kwargs, + ) -> None: + super().__init__(*args, **kwargs) + + self.video_resolution_buckets = [ + ( + int(b[0] / vae_temporal_compression_ratio), + int(b[1] / vae_height_compression_ratio), + int(b[2] / vae_width_compression_ratio), + ) + for b in video_resolution_buckets + ] + self.__frame_transforms = transforms.Compose( + [transforms.Lambda(lambda x: x / 255.0 * 2.0 - 1.0)] + ) + self.__image_transforms = self.__frame_transforms + + @override + def preprocess(self, video_path: Path, image_path: Path) -> Tuple[torch.Tensor, torch.Tensor]: + video = preprocess_video_with_buckets(video_path, self.video_resolution_buckets) + image = preprocess_image_with_resize(image_path, video.shape[2], video.shape[3]) + return video, image + + @override + def video_transform(self, frames: torch.Tensor) -> torch.Tensor: + return torch.stack([self.__frame_transforms(f) for f in frames], dim=0) + + @override + def image_transform(self, image: torch.Tensor) -> torch.Tensor: + return self.__image_transforms(image) diff --git a/finetune/datasets/t2v_dataset.py b/finetune/datasets/t2v_dataset.py new file mode 100644 index 0000000..9967a5c --- /dev/null +++ b/finetune/datasets/t2v_dataset.py @@ -0,0 +1,264 @@ +import hashlib +from pathlib import Path +from typing import TYPE_CHECKING, Any, Dict, List, Tuple + +import torch +from accelerate.logging import get_logger +from safetensors.torch import load_file, save_file +from torch.utils.data import Dataset +from torchvision import transforms +from typing_extensions import override + +from finetune.constants import LOG_LEVEL, LOG_NAME + +from .utils import ( + load_prompts, + load_videos, + preprocess_video_with_buckets, + preprocess_video_with_resize, +) + + +if TYPE_CHECKING: + from finetune.trainer import Trainer + +# Must import after torch because this can sometimes lead to a nasty segmentation fault, or stack smashing error +# Very few bug reports but it happens. Look in decord Github issues for more relevant information. +import decord # isort:skip + +decord.bridge.set_bridge("torch") + +logger = get_logger(LOG_NAME, LOG_LEVEL) + + +class BaseT2VDataset(Dataset): + """ + Base dataset class for Text-to-Video (T2V) training. + + This dataset loads prompts and videos for T2V training. + + Args: + data_root (str): Root directory containing the dataset files + caption_column (str): Path to file containing text prompts/captions + video_column (str): Path to file containing video paths + device (torch.device): Device to load the data on + encode_video_fn (Callable[[torch.Tensor], torch.Tensor], optional): Function to encode videos + """ + + def __init__( + self, + data_root: str, + caption_column: str, + video_column: str, + device: torch.device = None, + trainer: "Trainer" = None, + *args, + **kwargs, + ) -> None: + super().__init__() + + data_root = Path(data_root) + self.prompts = load_prompts(data_root / caption_column) + self.videos = load_videos(data_root / video_column) + self.device = device + self.encode_video = trainer.encode_video + self.encode_text = trainer.encode_text + self.trainer = trainer + + # Check if all video files exist + if any(not path.is_file() for path in self.videos): + raise ValueError( + f"Some video files were not found. Please ensure that all video files exist in the dataset directory. Missing file: {next(path for path in self.videos if not path.is_file())}" + ) + + # Check if number of prompts matches number of videos + if len(self.videos) != len(self.prompts): + raise ValueError( + f"Expected length of prompts and videos to be the same but found {len(self.prompts)=} and {len(self.videos)=}. Please ensure that the number of caption prompts and videos match in your dataset." + ) + + def __len__(self) -> int: + return len(self.videos) + + def __getitem__(self, index: int) -> Dict[str, Any]: + if isinstance(index, list): + # Here, index is actually a list of data objects that we need to return. + # The BucketSampler should ideally return indices. But, in the sampler, we'd like + # to have information about num_frames, height and width. Since this is not stored + # as metadata, we need to read the video to get this information. You could read this + # information without loading the full video in memory, but we do it anyway. In order + # to not load the video twice (once to get the metadata, and once to return the loaded video + # based on sampled indices), we cache it in the BucketSampler. When the sampler is + # to yield, we yield the cache data instead of indices. So, this special check ensures + # that data is not loaded a second time. PRs are welcome for improvements. + return index + + prompt = self.prompts[index] + video = self.videos[index] + train_resolution_str = "x".join(str(x) for x in self.trainer.args.train_resolution) + + cache_dir = self.trainer.args.data_root / "cache" + video_latent_dir = ( + cache_dir / "video_latent" / self.trainer.args.model_name / train_resolution_str + ) + prompt_embeddings_dir = cache_dir / "prompt_embeddings" + video_latent_dir.mkdir(parents=True, exist_ok=True) + prompt_embeddings_dir.mkdir(parents=True, exist_ok=True) + + prompt_hash = str(hashlib.sha256(prompt.encode()).hexdigest()) + prompt_embedding_path = prompt_embeddings_dir / (prompt_hash + ".safetensors") + encoded_video_path = video_latent_dir / (video.stem + ".safetensors") + + if prompt_embedding_path.exists(): + prompt_embedding = load_file(prompt_embedding_path)["prompt_embedding"] + logger.debug( + f"process {self.trainer.accelerator.process_index}: Loaded prompt embedding from {prompt_embedding_path}", + main_process_only=False, + ) + else: + prompt_embedding = self.encode_text(prompt) + prompt_embedding = prompt_embedding.to("cpu") + # [1, seq_len, hidden_size] -> [seq_len, hidden_size] + prompt_embedding = prompt_embedding[0] + save_file({"prompt_embedding": prompt_embedding}, prompt_embedding_path) + logger.info( + f"Saved prompt embedding to {prompt_embedding_path}", main_process_only=False + ) + + if encoded_video_path.exists(): + # encoded_video = torch.load(encoded_video_path, weights_only=True) + encoded_video = load_file(encoded_video_path)["encoded_video"] + logger.debug(f"Loaded encoded video from {encoded_video_path}", main_process_only=False) + # shape of image: [C, H, W] + else: + frames = self.preprocess(video) + frames = frames.to(self.device) + # Current shape of frames: [F, C, H, W] + frames = self.video_transform(frames) + # Convert to [B, C, F, H, W] + frames = frames.unsqueeze(0) + frames = frames.permute(0, 2, 1, 3, 4).contiguous() + encoded_video = self.encode_video(frames) + + # [1, C, F, H, W] -> [C, F, H, W] + encoded_video = encoded_video[0] + encoded_video = encoded_video.to("cpu") + save_file({"encoded_video": encoded_video}, encoded_video_path) + logger.info(f"Saved encoded video to {encoded_video_path}", main_process_only=False) + + # shape of encoded_video: [C, F, H, W] + return { + "prompt_embedding": prompt_embedding, + "encoded_video": encoded_video, + "video_metadata": { + "num_frames": encoded_video.shape[1], + "height": encoded_video.shape[2], + "width": encoded_video.shape[3], + }, + } + + def preprocess(self, video_path: Path) -> torch.Tensor: + """ + Loads and preprocesses a video. + + Args: + video_path: Path to the video file to load. + + Returns: + torch.Tensor: Video tensor of shape [F, C, H, W] where: + - F is number of frames + - C is number of channels (3 for RGB) + - H is height + - W is width + """ + raise NotImplementedError("Subclass must implement this method") + + def video_transform(self, frames: torch.Tensor) -> torch.Tensor: + """ + Applies transformations to a video. + + Args: + frames (torch.Tensor): A 4D tensor representing a video + with shape [F, C, H, W] where: + - F is number of frames + - C is number of channels (3 for RGB) + - H is height + - W is width + + Returns: + torch.Tensor: The transformed video tensor with the same shape as the input + """ + raise NotImplementedError("Subclass must implement this method") + + +class T2VDatasetWithResize(BaseT2VDataset): + """ + A dataset class for text-to-video generation that resizes inputs to fixed dimensions. + + This class preprocesses videos by resizing them to specified dimensions: + - Videos are resized to max_num_frames x height x width + + Args: + max_num_frames (int): Maximum number of frames to extract from videos + height (int): Target height for resizing videos + width (int): Target width for resizing videos + """ + + def __init__(self, max_num_frames: int, height: int, width: int, *args, **kwargs) -> None: + super().__init__(*args, **kwargs) + + self.max_num_frames = max_num_frames + self.height = height + self.width = width + + self.__frame_transform = transforms.Compose( + [transforms.Lambda(lambda x: x / 255.0 * 2.0 - 1.0)] + ) + + @override + def preprocess(self, video_path: Path) -> torch.Tensor: + return preprocess_video_with_resize( + video_path, + self.max_num_frames, + self.height, + self.width, + ) + + @override + def video_transform(self, frames: torch.Tensor) -> torch.Tensor: + return torch.stack([self.__frame_transform(f) for f in frames], dim=0) + + +class T2VDatasetWithBuckets(BaseT2VDataset): + def __init__( + self, + video_resolution_buckets: List[Tuple[int, int, int]], + vae_temporal_compression_ratio: int, + vae_height_compression_ratio: int, + vae_width_compression_ratio: int, + *args, + **kwargs, + ) -> None: + """ """ + super().__init__(*args, **kwargs) + + self.video_resolution_buckets = [ + ( + int(b[0] / vae_temporal_compression_ratio), + int(b[1] / vae_height_compression_ratio), + int(b[2] / vae_width_compression_ratio), + ) + for b in video_resolution_buckets + ] + + self.__frame_transform = transforms.Compose( + [transforms.Lambda(lambda x: x / 255.0 * 2.0 - 1.0)] + ) + + @override + def preprocess(self, video_path: Path) -> torch.Tensor: + return preprocess_video_with_buckets(video_path, self.video_resolution_buckets) + + @override + def video_transform(self, frames: torch.Tensor) -> torch.Tensor: + return torch.stack([self.__frame_transform(f) for f in frames], dim=0) diff --git a/finetune/datasets/utils.py b/finetune/datasets/utils.py new file mode 100644 index 0000000..e2667d6 --- /dev/null +++ b/finetune/datasets/utils.py @@ -0,0 +1,196 @@ +import logging +from pathlib import Path +from typing import List, Tuple + +import cv2 +import torch +from torchvision.transforms.functional import resize + + +# Must import after torch because this can sometimes lead to a nasty segmentation fault, or stack smashing error +# Very few bug reports but it happens. Look in decord Github issues for more relevant information. +import decord # isort:skip + +decord.bridge.set_bridge("torch") + + +########## loaders ########## + + +def load_prompts(prompt_path: Path) -> List[str]: + with open(prompt_path, "r", encoding="utf-8") as file: + return [line.strip() for line in file.readlines() if len(line.strip()) > 0] + + +def load_videos(video_path: Path) -> List[Path]: + with open(video_path, "r", encoding="utf-8") as file: + return [ + video_path.parent / line.strip() for line in file.readlines() if len(line.strip()) > 0 + ] + + +def load_images(image_path: Path) -> List[Path]: + with open(image_path, "r", encoding="utf-8") as file: + return [ + image_path.parent / line.strip() for line in file.readlines() if len(line.strip()) > 0 + ] + + +def load_images_from_videos(videos_path: List[Path]) -> List[Path]: + first_frames_dir = videos_path[0].parent.parent / "first_frames" + first_frames_dir.mkdir(exist_ok=True) + + first_frame_paths = [] + for video_path in videos_path: + frame_path = first_frames_dir / f"{video_path.stem}.png" + if frame_path.exists(): + first_frame_paths.append(frame_path) + continue + + # Open video + cap = cv2.VideoCapture(str(video_path)) + + # Read first frame + ret, frame = cap.read() + if not ret: + raise RuntimeError(f"Failed to read video: {video_path}") + + # Save frame as PNG with same name as video + cv2.imwrite(str(frame_path), frame) + logging.info(f"Saved first frame to {frame_path}") + + # Release video capture + cap.release() + + first_frame_paths.append(frame_path) + + return first_frame_paths + + +########## preprocessors ########## + + +def preprocess_image_with_resize( + image_path: Path | str, + height: int, + width: int, +) -> torch.Tensor: + """ + Loads and resizes a single image. + + Args: + image_path: Path to the image file. + height: Target height for resizing. + width: Target width for resizing. + + Returns: + torch.Tensor: Image tensor with shape [C, H, W] where: + C = number of channels (3 for RGB) + H = height + W = width + """ + if isinstance(image_path, str): + image_path = Path(image_path) + image = cv2.imread(image_path.as_posix()) + image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB) + image = cv2.resize(image, (width, height)) + image = torch.from_numpy(image).float() + image = image.permute(2, 0, 1).contiguous() + return image + + +def preprocess_video_with_resize( + video_path: Path | str, + max_num_frames: int, + height: int, + width: int, +) -> torch.Tensor: + """ + Loads and resizes a single video. + + The function processes the video through these steps: + 1. If video frame count > max_num_frames, downsample frames evenly + 2. If video dimensions don't match (height, width), resize frames + + Args: + video_path: Path to the video file. + max_num_frames: Maximum number of frames to keep. + height: Target height for resizing. + width: Target width for resizing. + + Returns: + A torch.Tensor with shape [F, C, H, W] where: + F = number of frames + C = number of channels (3 for RGB) + H = height + W = width + """ + if isinstance(video_path, str): + video_path = Path(video_path) + video_reader = decord.VideoReader(uri=video_path.as_posix(), width=width, height=height) + video_num_frames = len(video_reader) + if video_num_frames < max_num_frames: + # Get all frames first + frames = video_reader.get_batch(list(range(video_num_frames))) + # Repeat the last frame until we reach max_num_frames + last_frame = frames[-1:] + num_repeats = max_num_frames - video_num_frames + repeated_frames = last_frame.repeat(num_repeats, 1, 1, 1) + frames = torch.cat([frames, repeated_frames], dim=0) + return frames.float().permute(0, 3, 1, 2).contiguous() + else: + indices = list(range(0, video_num_frames, video_num_frames // max_num_frames)) + frames = video_reader.get_batch(indices) + frames = frames[:max_num_frames].float() + frames = frames.permute(0, 3, 1, 2).contiguous() + return frames + + +def preprocess_video_with_buckets( + video_path: Path, + resolution_buckets: List[Tuple[int, int, int]], +) -> torch.Tensor: + """ + Args: + video_path: Path to the video file. + resolution_buckets: List of tuples (num_frames, height, width) representing + available resolution buckets. + + Returns: + torch.Tensor: Video tensor with shape [F, C, H, W] where: + F = number of frames + C = number of channels (3 for RGB) + H = height + W = width + + The function processes the video through these steps: + 1. Finds nearest frame bucket <= video frame count + 2. Downsamples frames evenly to match bucket size + 3. Finds nearest resolution bucket based on dimensions + 4. Resizes frames to match bucket resolution + """ + video_reader = decord.VideoReader(uri=video_path.as_posix()) + video_num_frames = len(video_reader) + resolution_buckets = [bucket for bucket in resolution_buckets if bucket[0] <= video_num_frames] + if len(resolution_buckets) == 0: + raise ValueError( + f"video frame count in {video_path} is less than all frame buckets {resolution_buckets}" + ) + + nearest_frame_bucket = min( + resolution_buckets, + key=lambda bucket: video_num_frames - bucket[0], + default=1, + )[0] + frame_indices = list(range(0, video_num_frames, video_num_frames // nearest_frame_bucket)) + frames = video_reader.get_batch(frame_indices) + frames = frames[:nearest_frame_bucket].float() + frames = frames.permute(0, 3, 1, 2).contiguous() + + nearest_res = min( + resolution_buckets, key=lambda x: abs(x[1] - frames.shape[2]) + abs(x[2] - frames.shape[3]) + ) + nearest_res = (nearest_res[1], nearest_res[2]) + frames = torch.stack([resize(f, nearest_res) for f in frames], dim=0) + + return frames diff --git a/finetune/models/__init__.py b/finetune/models/__init__.py new file mode 100644 index 0000000..fa8142e --- /dev/null +++ b/finetune/models/__init__.py @@ -0,0 +1,12 @@ +import importlib +from pathlib import Path + + +package_dir = Path(__file__).parent + +for subdir in package_dir.iterdir(): + if subdir.is_dir() and not subdir.name.startswith("_"): + for module_path in subdir.glob("*.py"): + module_name = module_path.stem + full_module_name = f".{subdir.name}.{module_name}" + importlib.import_module(full_module_name, package=__name__) diff --git a/finetune/models/cogvideox1_5_i2v/lora_trainer.py b/finetune/models/cogvideox1_5_i2v/lora_trainer.py new file mode 100644 index 0000000..d72c89b --- /dev/null +++ b/finetune/models/cogvideox1_5_i2v/lora_trainer.py @@ -0,0 +1,9 @@ +from ..cogvideox_i2v.lora_trainer import CogVideoXI2VLoraTrainer +from ..utils import register + + +class CogVideoX1_5I2VLoraTrainer(CogVideoXI2VLoraTrainer): + pass + + +register("cogvideox1.5-i2v", "lora", CogVideoX1_5I2VLoraTrainer) diff --git a/finetune/models/cogvideox1_5_i2v/sft_trainer.py b/finetune/models/cogvideox1_5_i2v/sft_trainer.py new file mode 100644 index 0000000..63e4709 --- /dev/null +++ b/finetune/models/cogvideox1_5_i2v/sft_trainer.py @@ -0,0 +1,9 @@ +from ..cogvideox_i2v.sft_trainer import CogVideoXI2VSftTrainer +from ..utils import register + + +class CogVideoX1_5I2VSftTrainer(CogVideoXI2VSftTrainer): + pass + + +register("cogvideox1.5-i2v", "sft", CogVideoX1_5I2VSftTrainer) diff --git a/finetune/models/cogvideox1_5_t2v/lora_trainer.py b/finetune/models/cogvideox1_5_t2v/lora_trainer.py new file mode 100644 index 0000000..b921660 --- /dev/null +++ b/finetune/models/cogvideox1_5_t2v/lora_trainer.py @@ -0,0 +1,9 @@ +from ..cogvideox_t2v.lora_trainer import CogVideoXT2VLoraTrainer +from ..utils import register + + +class CogVideoX1_5T2VLoraTrainer(CogVideoXT2VLoraTrainer): + pass + + +register("cogvideox1.5-t2v", "lora", CogVideoX1_5T2VLoraTrainer) diff --git a/finetune/models/cogvideox1_5_t2v/sft_trainer.py b/finetune/models/cogvideox1_5_t2v/sft_trainer.py new file mode 100644 index 0000000..6220bb3 --- /dev/null +++ b/finetune/models/cogvideox1_5_t2v/sft_trainer.py @@ -0,0 +1,9 @@ +from ..cogvideox_t2v.sft_trainer import CogVideoXT2VSftTrainer +from ..utils import register + + +class CogVideoX1_5T2VSftTrainer(CogVideoXT2VSftTrainer): + pass + + +register("cogvideox1.5-t2v", "sft", CogVideoX1_5T2VSftTrainer) diff --git a/finetune/models/cogvideox_i2v/lora_trainer.py b/finetune/models/cogvideox_i2v/lora_trainer.py new file mode 100644 index 0000000..793cf76 --- /dev/null +++ b/finetune/models/cogvideox_i2v/lora_trainer.py @@ -0,0 +1,272 @@ +from typing import Any, Dict, List, Tuple + +import torch +from diffusers import ( + AutoencoderKLCogVideoX, + CogVideoXDPMScheduler, + CogVideoXImageToVideoPipeline, + CogVideoXTransformer3DModel, +) +from diffusers.models.embeddings import get_3d_rotary_pos_embed +from PIL import Image +from numpy import dtype +from transformers import AutoTokenizer, T5EncoderModel +from typing_extensions import override + +from finetune.schemas import Components +from finetune.trainer import Trainer +from finetune.utils import unwrap_model + +from ..utils import register + + +class CogVideoXI2VLoraTrainer(Trainer): + UNLOAD_LIST = ["text_encoder"] + + @override + def load_components(self) -> Dict[str, Any]: + components = Components() + model_path = str(self.args.model_path) + + components.pipeline_cls = CogVideoXImageToVideoPipeline + + components.tokenizer = AutoTokenizer.from_pretrained(model_path, subfolder="tokenizer") + + components.text_encoder = T5EncoderModel.from_pretrained( + model_path, subfolder="text_encoder" + ) + + components.transformer = CogVideoXTransformer3DModel.from_pretrained( + model_path, subfolder="transformer" + ) + + components.vae = AutoencoderKLCogVideoX.from_pretrained(model_path, subfolder="vae") + + components.scheduler = CogVideoXDPMScheduler.from_pretrained( + model_path, subfolder="scheduler" + ) + + return components + + @override + def initialize_pipeline(self) -> CogVideoXImageToVideoPipeline: + pipe = CogVideoXImageToVideoPipeline( + tokenizer=self.components.tokenizer, + text_encoder=self.components.text_encoder, + vae=self.components.vae, + transformer=unwrap_model(self.accelerator, self.components.transformer), + scheduler=self.components.scheduler, + ) + return pipe + + @override + def encode_video(self, video: torch.Tensor) -> torch.Tensor: + # shape of input video: [B, C, F, H, W] + vae = self.components.vae + video = video.to(vae.device, dtype=vae.dtype) + latent_dist = vae.encode(video).latent_dist + latent = latent_dist.sample() * vae.config.scaling_factor + return latent + + @override + def encode_text(self, prompt: str) -> torch.Tensor: + prompt_token_ids = self.components.tokenizer( + prompt, + padding="max_length", + max_length=self.state.transformer_config.max_text_seq_length, + truncation=True, + add_special_tokens=True, + return_tensors="pt", + ) + prompt_token_ids = prompt_token_ids.input_ids + prompt_embedding = self.components.text_encoder( + prompt_token_ids.to(self.accelerator.device) + )[0] + return prompt_embedding + + @override + def collate_fn(self, samples: List[Dict[str, Any]]) -> Dict[str, Any]: + ret = {"encoded_videos": [], "prompt_embedding": [], "images": []} + + for sample in samples: + encoded_video = sample["encoded_video"] + prompt_embedding = sample["prompt_embedding"] + image = sample["image"] + + ret["encoded_videos"].append(encoded_video) + ret["prompt_embedding"].append(prompt_embedding) + ret["images"].append(image) + + ret["encoded_videos"] = torch.stack(ret["encoded_videos"]) + ret["prompt_embedding"] = torch.stack(ret["prompt_embedding"]) + ret["images"] = torch.stack(ret["images"]) + + return ret + + @override + def compute_loss(self, batch) -> torch.Tensor: + prompt_embedding = batch["prompt_embedding"] + latent = batch["encoded_videos"] + images = batch["images"] + + # Shape of prompt_embedding: [B, seq_len, hidden_size] + # Shape of latent: [B, C, F, H, W] + # Shape of images: [B, C, H, W] + + patch_size_t = self.state.transformer_config.patch_size_t + if patch_size_t is not None: + ncopy = latent.shape[2] % patch_size_t + # Copy the first frame ncopy times to match patch_size_t + first_frame = latent[:, :, :1, :, :] # Get first frame [B, C, 1, H, W] + latent = torch.cat([first_frame.repeat(1, 1, ncopy, 1, 1), latent], dim=2) + assert latent.shape[2] % patch_size_t == 0 + + batch_size, num_channels, num_frames, height, width = latent.shape + + # Get prompt embeddings + _, seq_len, _ = prompt_embedding.shape + prompt_embedding = prompt_embedding.view(batch_size, seq_len, -1).to(dtype=latent.dtype) + + # Add frame dimension to images [B,C,H,W] -> [B,C,F,H,W] + images = images.unsqueeze(2) + # Add noise to images + image_noise_sigma = torch.normal( + mean=-3.0, std=0.5, size=(1,), device=self.accelerator.device + ) + image_noise_sigma = torch.exp(image_noise_sigma).to(dtype=images.dtype) + noisy_images = ( + images + torch.randn_like(images) * image_noise_sigma[:, None, None, None, None] + ) + image_latent_dist = self.components.vae.encode( + noisy_images.to(dtype=self.components.vae.dtype) + ).latent_dist + image_latents = image_latent_dist.sample() * self.components.vae.config.scaling_factor + + # Sample a random timestep for each sample + timesteps = torch.randint( + 0, + self.components.scheduler.config.num_train_timesteps, + (batch_size,), + device=self.accelerator.device, + ) + timesteps = timesteps.long() + + # from [B, C, F, H, W] to [B, F, C, H, W] + latent = latent.permute(0, 2, 1, 3, 4) + image_latents = image_latents.permute(0, 2, 1, 3, 4) + assert (latent.shape[0], *latent.shape[2:]) == ( + image_latents.shape[0], + *image_latents.shape[2:], + ) + + # Padding image_latents to the same frame number as latent + padding_shape = (latent.shape[0], latent.shape[1] - 1, *latent.shape[2:]) + latent_padding = image_latents.new_zeros(padding_shape) + image_latents = torch.cat([image_latents, latent_padding], dim=1) + + # Add noise to latent + noise = torch.randn_like(latent) + latent_noisy = self.components.scheduler.add_noise(latent, noise, timesteps) + + # Concatenate latent and image_latents in the channel dimension + latent_img_noisy = torch.cat([latent_noisy, image_latents], dim=2) + + # Prepare rotary embeds + vae_scale_factor_spatial = 2 ** (len(self.components.vae.config.block_out_channels) - 1) + transformer_config = self.state.transformer_config + rotary_emb = ( + self.prepare_rotary_positional_embeddings( + height=height * vae_scale_factor_spatial, + width=width * vae_scale_factor_spatial, + num_frames=num_frames, + transformer_config=transformer_config, + vae_scale_factor_spatial=vae_scale_factor_spatial, + device=self.accelerator.device, + ) + if transformer_config.use_rotary_positional_embeddings + else None + ) + + # Predict noise, For CogVideoX1.5 Only. + ofs_emb = ( + None + if self.state.transformer_config.ofs_embed_dim is None + else latent.new_full((1,), fill_value=2.0) + ) + predicted_noise = self.components.transformer( + hidden_states=latent_img_noisy, + encoder_hidden_states=prompt_embedding, + timestep=timesteps, + ofs=ofs_emb, + image_rotary_emb=rotary_emb, + return_dict=False, + )[0] + + # Denoise + latent_pred = self.components.scheduler.get_velocity( + predicted_noise, latent_noisy, timesteps + ) + + alphas_cumprod = self.components.scheduler.alphas_cumprod[timesteps] + weights = 1 / (1 - alphas_cumprod) + while len(weights.shape) < len(latent_pred.shape): + weights = weights.unsqueeze(-1) + + loss = torch.mean((weights * (latent_pred - latent) ** 2).reshape(batch_size, -1), dim=1) + loss = loss.mean() + + return loss + + @override + def validation_step( + self, eval_data: Dict[str, Any], pipe: CogVideoXImageToVideoPipeline + ) -> List[Tuple[str, Image.Image | List[Image.Image]]]: + """ + Return the data that needs to be saved. For videos, the data format is List[PIL], + and for images, the data format is PIL + """ + prompt, image, video = eval_data["prompt"], eval_data["image"], eval_data["video"] + + video_generate = pipe( + num_frames=self.state.train_frames, + height=self.state.train_height, + width=self.state.train_width, + prompt=prompt, + image=image, + generator=self.state.generator, + ).frames[0] + return [("video", video_generate)] + + def prepare_rotary_positional_embeddings( + self, + height: int, + width: int, + num_frames: int, + transformer_config: Dict, + vae_scale_factor_spatial: int, + device: torch.device, + ) -> Tuple[torch.Tensor, torch.Tensor]: + grid_height = height // (vae_scale_factor_spatial * transformer_config.patch_size) + grid_width = width // (vae_scale_factor_spatial * transformer_config.patch_size) + + if transformer_config.patch_size_t is None: + base_num_frames = num_frames + else: + base_num_frames = ( + num_frames + transformer_config.patch_size_t - 1 + ) // transformer_config.patch_size_t + + freqs_cos, freqs_sin = get_3d_rotary_pos_embed( + embed_dim=transformer_config.attention_head_dim, + crops_coords=None, + grid_size=(grid_height, grid_width), + temporal_size=base_num_frames, + grid_type="slice", + max_size=(grid_height, grid_width), + device=device, + ) + + return freqs_cos, freqs_sin + + +register("cogvideox-i2v", "lora", CogVideoXI2VLoraTrainer) diff --git a/finetune/models/cogvideox_i2v/sft_trainer.py b/finetune/models/cogvideox_i2v/sft_trainer.py new file mode 100644 index 0000000..b55bee8 --- /dev/null +++ b/finetune/models/cogvideox_i2v/sft_trainer.py @@ -0,0 +1,9 @@ +from ..cogvideox_i2v.lora_trainer import CogVideoXI2VLoraTrainer +from ..utils import register + + +class CogVideoXI2VSftTrainer(CogVideoXI2VLoraTrainer): + pass + + +register("cogvideox-i2v", "sft", CogVideoXI2VSftTrainer) diff --git a/finetune/models/cogvideox_t2v/lora_trainer.py b/finetune/models/cogvideox_t2v/lora_trainer.py new file mode 100644 index 0000000..5f0ec1c --- /dev/null +++ b/finetune/models/cogvideox_t2v/lora_trainer.py @@ -0,0 +1,228 @@ +from typing import Any, Dict, List, Tuple + +import torch +from diffusers import ( + AutoencoderKLCogVideoX, + CogVideoXDPMScheduler, + CogVideoXPipeline, + CogVideoXTransformer3DModel, +) +from diffusers.models.embeddings import get_3d_rotary_pos_embed +from PIL import Image +from transformers import AutoTokenizer, T5EncoderModel +from typing_extensions import override + +from finetune.schemas import Components +from finetune.trainer import Trainer +from finetune.utils import unwrap_model + +from ..utils import register + + +class CogVideoXT2VLoraTrainer(Trainer): + UNLOAD_LIST = ["text_encoder", "vae"] + + @override + def load_components(self) -> Components: + components = Components() + model_path = str(self.args.model_path) + + components.pipeline_cls = CogVideoXPipeline + + components.tokenizer = AutoTokenizer.from_pretrained(model_path, subfolder="tokenizer") + + components.text_encoder = T5EncoderModel.from_pretrained( + model_path, subfolder="text_encoder" + ) + + components.transformer = CogVideoXTransformer3DModel.from_pretrained( + model_path, subfolder="transformer" + ) + + components.vae = AutoencoderKLCogVideoX.from_pretrained(model_path, subfolder="vae") + + components.scheduler = CogVideoXDPMScheduler.from_pretrained( + model_path, subfolder="scheduler" + ) + + return components + + @override + def initialize_pipeline(self) -> CogVideoXPipeline: + pipe = CogVideoXPipeline( + tokenizer=self.components.tokenizer, + text_encoder=self.components.text_encoder, + vae=self.components.vae, + transformer=unwrap_model(self.accelerator, self.components.transformer), + scheduler=self.components.scheduler, + ) + return pipe + + @override + def encode_video(self, video: torch.Tensor) -> torch.Tensor: + # shape of input video: [B, C, F, H, W] + vae = self.components.vae + video = video.to(vae.device, dtype=vae.dtype) + latent_dist = vae.encode(video).latent_dist + latent = latent_dist.sample() * vae.config.scaling_factor + return latent + + @override + def encode_text(self, prompt: str) -> torch.Tensor: + prompt_token_ids = self.components.tokenizer( + prompt, + padding="max_length", + max_length=self.state.transformer_config.max_text_seq_length, + truncation=True, + add_special_tokens=True, + return_tensors="pt", + ) + prompt_token_ids = prompt_token_ids.input_ids + prompt_embedding = self.components.text_encoder( + prompt_token_ids.to(self.accelerator.device) + )[0] + return prompt_embedding + + @override + def collate_fn(self, samples: List[Dict[str, Any]]) -> Dict[str, Any]: + ret = {"encoded_videos": [], "prompt_embedding": []} + + for sample in samples: + encoded_video = sample["encoded_video"] + prompt_embedding = sample["prompt_embedding"] + + ret["encoded_videos"].append(encoded_video) + ret["prompt_embedding"].append(prompt_embedding) + + ret["encoded_videos"] = torch.stack(ret["encoded_videos"]) + ret["prompt_embedding"] = torch.stack(ret["prompt_embedding"]) + + return ret + + @override + def compute_loss(self, batch) -> torch.Tensor: + prompt_embedding = batch["prompt_embedding"] + latent = batch["encoded_videos"] + + # Shape of prompt_embedding: [B, seq_len, hidden_size] + # Shape of latent: [B, C, F, H, W] + + patch_size_t = self.state.transformer_config.patch_size_t + if patch_size_t is not None: + ncopy = latent.shape[2] % patch_size_t + # Copy the first frame ncopy times to match patch_size_t + first_frame = latent[:, :, :1, :, :] # Get first frame [B, C, 1, H, W] + latent = torch.cat([first_frame.repeat(1, 1, ncopy, 1, 1), latent], dim=2) + assert latent.shape[2] % patch_size_t == 0 + + batch_size, num_channels, num_frames, height, width = latent.shape + + # Get prompt embeddings + _, seq_len, _ = prompt_embedding.shape + prompt_embedding = prompt_embedding.view(batch_size, seq_len, -1).to(dtype=latent.dtype) + + # Sample a random timestep for each sample + timesteps = torch.randint( + 0, + self.components.scheduler.config.num_train_timesteps, + (batch_size,), + device=self.accelerator.device, + ) + timesteps = timesteps.long() + + # Add noise to latent + latent = latent.permute(0, 2, 1, 3, 4) # from [B, C, F, H, W] to [B, F, C, H, W] + noise = torch.randn_like(latent) + latent_added_noise = self.components.scheduler.add_noise(latent, noise, timesteps) + + # Prepare rotary embeds + vae_scale_factor_spatial = 2 ** (len(self.components.vae.config.block_out_channels) - 1) + transformer_config = self.state.transformer_config + rotary_emb = ( + self.prepare_rotary_positional_embeddings( + height=height * vae_scale_factor_spatial, + width=width * vae_scale_factor_spatial, + num_frames=num_frames, + transformer_config=transformer_config, + vae_scale_factor_spatial=vae_scale_factor_spatial, + device=self.accelerator.device, + ) + if transformer_config.use_rotary_positional_embeddings + else None + ) + + # Predict noise + predicted_noise = self.components.transformer( + hidden_states=latent_added_noise, + encoder_hidden_states=prompt_embedding, + timestep=timesteps, + image_rotary_emb=rotary_emb, + return_dict=False, + )[0] + + # Denoise + latent_pred = self.components.scheduler.get_velocity( + predicted_noise, latent_added_noise, timesteps + ) + + alphas_cumprod = self.components.scheduler.alphas_cumprod[timesteps] + weights = 1 / (1 - alphas_cumprod) + while len(weights.shape) < len(latent_pred.shape): + weights = weights.unsqueeze(-1) + + loss = torch.mean((weights * (latent_pred - latent) ** 2).reshape(batch_size, -1), dim=1) + loss = loss.mean() + + return loss + + @override + def validation_step( + self, eval_data: Dict[str, Any], pipe: CogVideoXPipeline + ) -> List[Tuple[str, Image.Image | List[Image.Image]]]: + """ + Return the data that needs to be saved. For videos, the data format is List[PIL], + and for images, the data format is PIL + """ + prompt, image, video = eval_data["prompt"], eval_data["image"], eval_data["video"] + + video_generate = pipe( + num_frames=self.state.train_frames, + height=self.state.train_height, + width=self.state.train_width, + prompt=prompt, + generator=self.state.generator, + ).frames[0] + return [("video", video_generate)] + + def prepare_rotary_positional_embeddings( + self, + height: int, + width: int, + num_frames: int, + transformer_config: Dict, + vae_scale_factor_spatial: int, + device: torch.device, + ) -> Tuple[torch.Tensor, torch.Tensor]: + grid_height = height // (vae_scale_factor_spatial * transformer_config.patch_size) + grid_width = width // (vae_scale_factor_spatial * transformer_config.patch_size) + + if transformer_config.patch_size_t is None: + base_num_frames = num_frames + else: + base_num_frames = ( + num_frames + transformer_config.patch_size_t - 1 + ) // transformer_config.patch_size_t + freqs_cos, freqs_sin = get_3d_rotary_pos_embed( + embed_dim=transformer_config.attention_head_dim, + crops_coords=None, + grid_size=(grid_height, grid_width), + temporal_size=base_num_frames, + grid_type="slice", + max_size=(grid_height, grid_width), + device=device, + ) + + return freqs_cos, freqs_sin + + +register("cogvideox-t2v", "lora", CogVideoXT2VLoraTrainer) diff --git a/finetune/models/cogvideox_t2v/sft_trainer.py b/finetune/models/cogvideox_t2v/sft_trainer.py new file mode 100644 index 0000000..12b239e --- /dev/null +++ b/finetune/models/cogvideox_t2v/sft_trainer.py @@ -0,0 +1,9 @@ +from ..cogvideox_t2v.lora_trainer import CogVideoXT2VLoraTrainer +from ..utils import register + + +class CogVideoXT2VSftTrainer(CogVideoXT2VLoraTrainer): + pass + + +register("cogvideox-t2v", "sft", CogVideoXT2VSftTrainer) diff --git a/finetune/models/utils.py b/finetune/models/utils.py new file mode 100644 index 0000000..2028672 --- /dev/null +++ b/finetune/models/utils.py @@ -0,0 +1,59 @@ +from typing import Dict, Literal + +from finetune.trainer import Trainer + + +SUPPORTED_MODELS: Dict[str, Dict[str, Trainer]] = {} + + +def register(model_name: str, training_type: Literal["lora", "sft"], trainer_cls: Trainer): + """Register a model and its associated functions for a specific training type. + + Args: + model_name (str): Name of the model to register (e.g. "cogvideox-5b") + training_type (Literal["lora", "sft"]): Type of training - either "lora" or "sft" + trainer_cls (Trainer): Trainer class to register. + """ + + # Check if model_name and training_type exists in SUPPORTED_MODELS + if model_name not in SUPPORTED_MODELS: + SUPPORTED_MODELS[model_name] = {} + else: + if training_type in SUPPORTED_MODELS[model_name]: + raise ValueError(f"Training type {training_type} already exists for model {model_name}") + + SUPPORTED_MODELS[model_name][training_type] = trainer_cls + + +def show_supported_models(): + """Print all currently supported models and their training types.""" + + print("\nSupported Models:") + print("================") + + for model_name, training_types in SUPPORTED_MODELS.items(): + print(f"\n{model_name}") + print("-" * len(model_name)) + for training_type in training_types: + print(f" • {training_type}") + + +def get_model_cls(model_type: str, training_type: Literal["lora", "sft"]) -> Trainer: + """Get the trainer class for a specific model and training type.""" + if model_type not in SUPPORTED_MODELS: + print(f"\nModel '{model_type}' is not supported.") + print("\nSupported models are:") + for supported_model in SUPPORTED_MODELS: + print(f" • {supported_model}") + raise ValueError(f"Model '{model_type}' is not supported") + + if training_type not in SUPPORTED_MODELS[model_type]: + print(f"\nTraining type '{training_type}' is not supported for model '{model_type}'.") + print(f"\nSupported training types for '{model_type}' are:") + for supported_type in SUPPORTED_MODELS[model_type]: + print(f" • {supported_type}") + raise ValueError( + f"Training type '{training_type}' is not supported for model '{model_type}'" + ) + + return SUPPORTED_MODELS[model_type][training_type] diff --git a/finetune/schemas/__init__.py b/finetune/schemas/__init__.py new file mode 100644 index 0000000..73f547b --- /dev/null +++ b/finetune/schemas/__init__.py @@ -0,0 +1,6 @@ +from .args import Args +from .components import Components +from .state import State + + +__all__ = ["Args", "State", "Components"] diff --git a/finetune/schemas/args.py b/finetune/schemas/args.py new file mode 100644 index 0000000..bba7d01 --- /dev/null +++ b/finetune/schemas/args.py @@ -0,0 +1,254 @@ +import argparse +import datetime +import logging +from pathlib import Path +from typing import Any, List, Literal, Tuple + +from pydantic import BaseModel, ValidationInfo, field_validator + + +class Args(BaseModel): + ########## Model ########## + model_path: Path + model_name: str + model_type: Literal["i2v", "t2v"] + training_type: Literal["lora", "sft"] = "lora" + + ########## Output ########## + output_dir: Path = Path("train_results/{:%Y-%m-%d-%H-%M-%S}".format(datetime.datetime.now())) + report_to: Literal["tensorboard", "wandb", "all"] | None = None + tracker_name: str = "finetrainer-cogvideo" + + ########## Data ########### + data_root: Path + caption_column: Path + image_column: Path | None = None + video_column: Path + + ########## Training ######### + resume_from_checkpoint: Path | None = None + + seed: int | None = None + train_epochs: int + train_steps: int | None = None + checkpointing_steps: int = 200 + checkpointing_limit: int = 10 + + batch_size: int + gradient_accumulation_steps: int = 1 + + train_resolution: Tuple[int, int, int] # shape: (frames, height, width) + + #### deprecated args: video_resolution_buckets + # if use bucket for training, should not be None + # Note1: At least one frame rate in the bucket must be less than or equal to the frame rate of any video in the dataset + # Note2: For cogvideox, cogvideox1.5 + # The frame rate set in the bucket must be an integer multiple of 8 (spatial_compression_rate[4] * path_t[2] = 8) + # The height and width set in the bucket must be an integer multiple of 8 (temporal_compression_rate[8]) + # video_resolution_buckets: List[Tuple[int, int, int]] | None = None + + mixed_precision: Literal["no", "fp16", "bf16"] + + learning_rate: float = 2e-5 + optimizer: str = "adamw" + beta1: float = 0.9 + beta2: float = 0.95 + beta3: float = 0.98 + epsilon: float = 1e-8 + weight_decay: float = 1e-4 + max_grad_norm: float = 1.0 + + lr_scheduler: str = "constant_with_warmup" + lr_warmup_steps: int = 100 + lr_num_cycles: int = 1 + lr_power: float = 1.0 + + num_workers: int = 8 + pin_memory: bool = True + + gradient_checkpointing: bool = True + enable_slicing: bool = True + enable_tiling: bool = True + nccl_timeout: int = 1800 + + ########## Lora ########## + rank: int = 128 + lora_alpha: int = 64 + target_modules: List[str] = ["to_q", "to_k", "to_v", "to_out.0"] + + ########## Validation ########## + do_validation: bool = False + validation_steps: int | None # if set, should be a multiple of checkpointing_steps + validation_dir: Path | None # if set do_validation, should not be None + validation_prompts: str | None # if set do_validation, should not be None + validation_images: str | None # if set do_validation and model_type == i2v, should not be None + validation_videos: str | None # if set do_validation and model_type == v2v, should not be None + gen_fps: int = 15 + + #### deprecated args: gen_video_resolution + # 1. If set do_validation, should not be None + # 2. Suggest selecting the bucket from `video_resolution_buckets` that is closest to the resolution you have chosen for fine-tuning + # or the resolution recommended by the model + # 3. Note: For cogvideox, cogvideox1.5 + # The frame rate set in the bucket must be an integer multiple of 8 (spatial_compression_rate[4] * path_t[2] = 8) + # The height and width set in the bucket must be an integer multiple of 8 (temporal_compression_rate[8]) + # gen_video_resolution: Tuple[int, int, int] | None # shape: (frames, height, width) + + @field_validator("image_column") + def validate_image_column(cls, v: str | None, info: ValidationInfo) -> str | None: + values = info.data + if values.get("model_type") == "i2v" and not v: + logging.warning( + "No `image_column` specified for i2v model. Will automatically extract first frames from videos as conditioning images." + ) + return v + + @field_validator("validation_dir", "validation_prompts") + def validate_validation_required_fields(cls, v: Any, info: ValidationInfo) -> Any: + values = info.data + if values.get("do_validation") and not v: + field_name = info.field_name + raise ValueError(f"{field_name} must be specified when do_validation is True") + return v + + @field_validator("validation_images") + def validate_validation_images(cls, v: str | None, info: ValidationInfo) -> str | None: + values = info.data + if values.get("do_validation") and values.get("model_type") == "i2v" and not v: + raise ValueError( + "validation_images must be specified when do_validation is True and model_type is i2v" + ) + return v + + @field_validator("validation_videos") + def validate_validation_videos(cls, v: str | None, info: ValidationInfo) -> str | None: + values = info.data + if values.get("do_validation") and values.get("model_type") == "v2v" and not v: + raise ValueError( + "validation_videos must be specified when do_validation is True and model_type is v2v" + ) + return v + + @field_validator("validation_steps") + def validate_validation_steps(cls, v: int | None, info: ValidationInfo) -> int | None: + values = info.data + if values.get("do_validation"): + if v is None: + raise ValueError("validation_steps must be specified when do_validation is True") + if values.get("checkpointing_steps") and v % values["checkpointing_steps"] != 0: + raise ValueError("validation_steps must be a multiple of checkpointing_steps") + return v + + @field_validator("train_resolution") + def validate_train_resolution(cls, v: Tuple[int, int, int], info: ValidationInfo) -> str: + try: + frames, height, width = v + + # Check if (frames - 1) is multiple of 8 + if (frames - 1) % 8 != 0: + raise ValueError("Number of frames - 1 must be a multiple of 8") + + # Check resolution for cogvideox-5b models + model_name = info.data.get("model_name", "") + if model_name in ["cogvideox-5b-i2v", "cogvideox-5b-t2v"]: + if (height, width) != (480, 720): + raise ValueError( + "For cogvideox-5b models, height must be 480 and width must be 720" + ) + + return v + + except ValueError as e: + if ( + str(e) == "not enough values to unpack (expected 3, got 0)" + or str(e) == "invalid literal for int() with base 10" + ): + raise ValueError("train_resolution must be in format 'frames x height x width'") + raise e + + @field_validator("mixed_precision") + def validate_mixed_precision(cls, v: str, info: ValidationInfo) -> str: + if v == "fp16" and "cogvideox-2b" not in str(info.data.get("model_path", "")).lower(): + logging.warning( + "All CogVideoX models except cogvideox-2b were trained with bfloat16. " + "Using fp16 precision may lead to training instability." + ) + return v + + @classmethod + def parse_args(cls): + """Parse command line arguments and return Args instance""" + parser = argparse.ArgumentParser() + # Required arguments + parser.add_argument("--model_path", type=str, required=True) + parser.add_argument("--model_name", type=str, required=True) + parser.add_argument("--model_type", type=str, required=True) + parser.add_argument("--training_type", type=str, required=True) + parser.add_argument("--output_dir", type=str, required=True) + parser.add_argument("--data_root", type=str, required=True) + parser.add_argument("--caption_column", type=str, required=True) + parser.add_argument("--video_column", type=str, required=True) + parser.add_argument("--train_resolution", type=str, required=True) + parser.add_argument("--report_to", type=str, required=True) + + # Training hyperparameters + parser.add_argument("--seed", type=int, default=42) + parser.add_argument("--train_epochs", type=int, default=10) + parser.add_argument("--train_steps", type=int, default=None) + parser.add_argument("--gradient_accumulation_steps", type=int, default=1) + parser.add_argument("--batch_size", type=int, default=1) + parser.add_argument("--learning_rate", type=float, default=2e-5) + parser.add_argument("--optimizer", type=str, default="adamw") + parser.add_argument("--beta1", type=float, default=0.9) + parser.add_argument("--beta2", type=float, default=0.95) + parser.add_argument("--beta3", type=float, default=0.98) + parser.add_argument("--epsilon", type=float, default=1e-8) + parser.add_argument("--weight_decay", type=float, default=1e-4) + parser.add_argument("--max_grad_norm", type=float, default=1.0) + + # Learning rate scheduler + parser.add_argument("--lr_scheduler", type=str, default="constant_with_warmup") + parser.add_argument("--lr_warmup_steps", type=int, default=100) + parser.add_argument("--lr_num_cycles", type=int, default=1) + parser.add_argument("--lr_power", type=float, default=1.0) + + # Data loading + parser.add_argument("--num_workers", type=int, default=8) + parser.add_argument("--pin_memory", type=bool, default=True) + parser.add_argument("--image_column", type=str, default=None) + + # Model configuration + parser.add_argument("--mixed_precision", type=str, default="no") + parser.add_argument("--gradient_checkpointing", type=bool, default=True) + parser.add_argument("--enable_slicing", type=bool, default=True) + parser.add_argument("--enable_tiling", type=bool, default=True) + parser.add_argument("--nccl_timeout", type=int, default=1800) + + # LoRA parameters + parser.add_argument("--rank", type=int, default=128) + parser.add_argument("--lora_alpha", type=int, default=64) + parser.add_argument( + "--target_modules", type=str, nargs="+", default=["to_q", "to_k", "to_v", "to_out.0"] + ) + + # Checkpointing + parser.add_argument("--checkpointing_steps", type=int, default=200) + parser.add_argument("--checkpointing_limit", type=int, default=10) + parser.add_argument("--resume_from_checkpoint", type=str, default=None) + + # Validation + parser.add_argument("--do_validation", type=lambda x: x.lower() == 'true', default=False) + parser.add_argument("--validation_steps", type=int, default=None) + parser.add_argument("--validation_dir", type=str, default=None) + parser.add_argument("--validation_prompts", type=str, default=None) + parser.add_argument("--validation_images", type=str, default=None) + parser.add_argument("--validation_videos", type=str, default=None) + parser.add_argument("--gen_fps", type=int, default=15) + + args = parser.parse_args() + + # Convert video_resolution_buckets string to list of tuples + frames, height, width = args.train_resolution.split("x") + args.train_resolution = (int(frames), int(height), int(width)) + + return cls(**vars(args)) diff --git a/finetune/schemas/components.py b/finetune/schemas/components.py new file mode 100644 index 0000000..5edd38e --- /dev/null +++ b/finetune/schemas/components.py @@ -0,0 +1,28 @@ +from typing import Any + +from pydantic import BaseModel + + +class Components(BaseModel): + # pipeline cls + pipeline_cls: Any = None + + # Tokenizers + tokenizer: Any = None + tokenizer_2: Any = None + tokenizer_3: Any = None + + # Text encoders + text_encoder: Any = None + text_encoder_2: Any = None + text_encoder_3: Any = None + + # Autoencoder + vae: Any = None + + # Denoiser + transformer: Any = None + unet: Any = None + + # Scheduler + scheduler: Any = None diff --git a/finetune/schemas/state.py b/finetune/schemas/state.py new file mode 100644 index 0000000..36fa0fa --- /dev/null +++ b/finetune/schemas/state.py @@ -0,0 +1,29 @@ +from pathlib import Path +from typing import Any, Dict, List + +import torch +from pydantic import BaseModel + + +class State(BaseModel): + model_config = {"arbitrary_types_allowed": True} + + train_frames: int + train_height: int + train_width: int + + transformer_config: Dict[str, Any] = None + + weight_dtype: torch.dtype = torch.float32 # dtype for mixed precision training + num_trainable_parameters: int = 0 + overwrote_max_train_steps: bool = False + num_update_steps_per_epoch: int = 0 + total_batch_size_count: int = 0 + + generator: torch.Generator | None = None + + validation_prompts: List[str] = [] + validation_images: List[Path | None] = [] + validation_videos: List[Path | None] = [] + + using_deepspeed: bool = False diff --git a/finetune/scripts/extract_images.py b/finetune/scripts/extract_images.py new file mode 100644 index 0000000..eca6c21 --- /dev/null +++ b/finetune/scripts/extract_images.py @@ -0,0 +1,60 @@ +import argparse +import os +from pathlib import Path + +import cv2 + + +def parse_args(): + parser = argparse.ArgumentParser() + parser.add_argument( + "--datadir", + type=str, + required=True, + help="Root directory containing videos.txt and video subdirectory", + ) + return parser.parse_args() + + +args = parse_args() + +# Create data/images directory if it doesn't exist +data_dir = Path(args.datadir) +image_dir = data_dir / "images" +image_dir.mkdir(exist_ok=True) + +# Read videos.txt +videos_file = data_dir / "videos.txt" +with open(videos_file, "r") as f: + video_paths = [line.strip() for line in f.readlines() if line.strip()] + +# Process each video file and collect image paths +image_paths = [] +for video_rel_path in video_paths: + video_path = data_dir / video_rel_path + + # Open video + cap = cv2.VideoCapture(str(video_path)) + + # Read first frame + ret, frame = cap.read() + if not ret: + print(f"Failed to read video: {video_path}") + continue + + # Save frame as PNG with same name as video + image_name = f"images/{video_path.stem}.png" + image_path = data_dir / image_name + cv2.imwrite(str(image_path), frame) + + # Release video capture + cap.release() + + print(f"Extracted first frame from {video_path} to {image_path}") + image_paths.append(image_name) + +# Write images.txt +images_file = data_dir / "images.txt" +with open(images_file, "w") as f: + for path in image_paths: + f.write(f"{path}\n") diff --git a/finetune/train.py b/finetune/train.py new file mode 100644 index 0000000..660a744 --- /dev/null +++ b/finetune/train.py @@ -0,0 +1,19 @@ +import sys +from pathlib import Path + + +sys.path.append(str(Path(__file__).parent.parent)) + +from finetune.models.utils import get_model_cls +from finetune.schemas import Args + + +def main(): + args = Args.parse_args() + trainer_cls = get_model_cls(args.model_name, args.training_type) + trainer = trainer_cls(args) + trainer.fit() + + +if __name__ == "__main__": + main() diff --git a/finetune/train_ddp_i2v.sh b/finetune/train_ddp_i2v.sh new file mode 100644 index 0000000..d05661f --- /dev/null +++ b/finetune/train_ddp_i2v.sh @@ -0,0 +1,70 @@ +#!/usr/bin/env bash + +# Prevent tokenizer parallelism issues +export TOKENIZERS_PARALLELISM=false + +# Model Configuration +MODEL_ARGS=( + --model_path "THUDM/CogVideoX1.5-5B-I2V" + --model_name "cogvideox1.5-i2v" # ["cogvideox-i2v"] + --model_type "i2v" + --training_type "lora" +) + +# Output Configuration +OUTPUT_ARGS=( + --output_dir "/path/to/your/output_dir" + --report_to "tensorboard" +) + +# Data Configuration +DATA_ARGS=( + --data_root "/absolute/path/to/your/data_root" + --caption_column "prompt.txt" + --video_column "videos.txt" + # --image_column "images.txt" # comment this line will use first frame of video as image conditioning + --train_resolution "81x768x1360" # (frames x height x width), frames should be 8N+1 +) + +# Training Configuration +TRAIN_ARGS=( + --train_epochs 10 # number of training epochs + --seed 42 # random seed + --batch_size 1 + --gradient_accumulation_steps 1 + --mixed_precision "bf16" # ["no", "fp16"] # Only CogVideoX-2B supports fp16 training +) + +# System Configuration +SYSTEM_ARGS=( + --num_workers 8 + --pin_memory True + --nccl_timeout 1800 +) + +# Checkpointing Configuration +CHECKPOINT_ARGS=( + --checkpointing_steps 10 # save checkpoint every x steps + --checkpointing_limit 2 # maximum number of checkpoints to keep, after which the oldest one is deleted + --resume_from_checkpoint "/absolute/path/to/checkpoint_dir" # if you want to resume from a checkpoint, otherwise, comment this line +) + +# Validation Configuration +VALIDATION_ARGS=( + --do_validation false # ["true", "false"] + --validation_dir "/absolute/path/to/your/validation_set" + --validation_steps 20 # should be multiple of checkpointing_steps + --validation_prompts "prompts.txt" + --validation_images "images.txt" + --gen_fps 16 +) + +# Combine all arguments and launch training +accelerate launch train.py \ + "${MODEL_ARGS[@]}" \ + "${OUTPUT_ARGS[@]}" \ + "${DATA_ARGS[@]}" \ + "${TRAIN_ARGS[@]}" \ + "${SYSTEM_ARGS[@]}" \ + "${CHECKPOINT_ARGS[@]}" \ + "${VALIDATION_ARGS[@]}" diff --git a/finetune/train_ddp_t2v.sh b/finetune/train_ddp_t2v.sh new file mode 100644 index 0000000..994a440 --- /dev/null +++ b/finetune/train_ddp_t2v.sh @@ -0,0 +1,68 @@ +#!/usr/bin/env bash + +# Prevent tokenizer parallelism issues +export TOKENIZERS_PARALLELISM=false + +# Model Configuration +MODEL_ARGS=( + --model_path "THUDM/CogVideoX1.5-5B" + --model_name "cogvideox1.5-t2v" # ["cogvideox-t2v"] + --model_type "t2v" + --training_type "lora" +) + +# Output Configuration +OUTPUT_ARGS=( + --output_dir "/absolute/path/to/your/output_dir" + --report_to "tensorboard" +) + +# Data Configuration +DATA_ARGS=( + --data_root "/absolute/path/to/your/data_root" + --caption_column "prompt.txt" + --video_column "videos.txt" + --train_resolution "81x768x1360" # (frames x height x width), frames should be 8N+1 +) + +# Training Configuration +TRAIN_ARGS=( + --train_epochs 10 # number of training epochs + --seed 42 # random seed + --batch_size 1 + --gradient_accumulation_steps 1 + --mixed_precision "bf16" # ["no", "fp16"] # Only CogVideoX-2B supports fp16 training +) + +# System Configuration +SYSTEM_ARGS=( + --num_workers 8 + --pin_memory True + --nccl_timeout 1800 +) + +# Checkpointing Configuration +CHECKPOINT_ARGS=( + --checkpointing_steps 10 # save checkpoint every x steps + --checkpointing_limit 2 # maximum number of checkpoints to keep, after which the oldest one is deleted + --resume_from_checkpoint "/absolute/path/to/checkpoint_dir" # if you want to resume from a checkpoint, otherwise, comment this line +) + +# Validation Configuration +VALIDATION_ARGS=( + --do_validation false # ["true", "false"] + --validation_dir "/absolute/path/to/your/validation_set" + --validation_steps 20 # should be multiple of checkpointing_steps + --validation_prompts "prompts.txt" + --gen_fps 16 +) + +# Combine all arguments and launch training +accelerate launch train.py \ + "${MODEL_ARGS[@]}" \ + "${OUTPUT_ARGS[@]}" \ + "${DATA_ARGS[@]}" \ + "${TRAIN_ARGS[@]}" \ + "${SYSTEM_ARGS[@]}" \ + "${CHECKPOINT_ARGS[@]}" \ + "${VALIDATION_ARGS[@]}" diff --git a/finetune/train_zero_i2v.sh b/finetune/train_zero_i2v.sh new file mode 100644 index 0000000..03f2c93 --- /dev/null +++ b/finetune/train_zero_i2v.sh @@ -0,0 +1,73 @@ +#!/usr/bin/env bash + +# Prevent tokenizer parallelism issues +export TOKENIZERS_PARALLELISM=false + +# Model Configuration +MODEL_ARGS=( + --model_path "THUDM/CogVideoX1.5-5B-I2V" + --model_name "cogvideox1.5-i2v" # ["cogvideox-i2v"] + --model_type "i2v" + --training_type "sft" +) + +# Output Configuration +OUTPUT_ARGS=( + --output_dir "/absolute/path/to/your/output_dir" + --report_to "tensorboard" +) + +# Data Configuration +DATA_ARGS=( + --data_root "/absolute/path/to/your/data_root" + --caption_column "prompt.txt" + --video_column "videos.txt" + # --image_column "images.txt" # comment this line will use first frame of video as image conditioning + --train_resolution "81x768x1360" # (frames x height x width), frames should be 8N+1 and height, width should be multiples of 16 +) + +# Training Configuration +TRAIN_ARGS=( + --train_epochs 10 # number of training epochs + --seed 42 # random seed + + ######### Please keep consistent with deepspeed config file ########## + --batch_size 1 + --gradient_accumulation_steps 1 + --mixed_precision "bf16" # ["no", "fp16"] Only CogVideoX-2B supports fp16 training + ######################################################################## +) + +# System Configuration +SYSTEM_ARGS=( + --num_workers 8 + --pin_memory True + --nccl_timeout 1800 +) + +# Checkpointing Configuration +CHECKPOINT_ARGS=( + --checkpointing_steps 10 # save checkpoint every x steps + --checkpointing_limit 2 # maximum number of checkpoints to keep, after which the oldest one is deleted + # --resume_from_checkpoint "/absolute/path/to/checkpoint_dir" # if you want to resume from a checkpoint, otherwise, comment this line +) + +# Validation Configuration +VALIDATION_ARGS=( + --do_validation false # ["true", "false"] + --validation_dir "/absolute/path/to/validation_set" + --validation_steps 20 # should be multiple of checkpointing_steps + --validation_prompts "prompts.txt" + --validation_images "images.txt" + --gen_fps 16 +) + +# Combine all arguments and launch training +accelerate launch --config_file accelerate_config.yaml train.py \ + "${MODEL_ARGS[@]}" \ + "${OUTPUT_ARGS[@]}" \ + "${DATA_ARGS[@]}" \ + "${TRAIN_ARGS[@]}" \ + "${SYSTEM_ARGS[@]}" \ + "${CHECKPOINT_ARGS[@]}" \ + "${VALIDATION_ARGS[@]}" diff --git a/finetune/train_zero_t2v.sh b/finetune/train_zero_t2v.sh new file mode 100644 index 0000000..265b08d --- /dev/null +++ b/finetune/train_zero_t2v.sh @@ -0,0 +1,71 @@ +#!/usr/bin/env bash + +# Prevent tokenizer parallelism issues +export TOKENIZERS_PARALLELISM=false + +# Model Configuration +MODEL_ARGS=( + --model_path "THUDM/CogVideoX1.5-5B" + --model_name "cogvideox1.5-t2v" # ["cogvideox-t2v"] + --model_type "t2v" + --training_type "sft" +) + +# Output Configuration +OUTPUT_ARGS=( + --output_dir "/absolute/path/to/your/output_dir" + --report_to "tensorboard" +) + +# Data Configuration +DATA_ARGS=( + --data_root "/absolute/path/to/your/data_root" + --caption_column "prompt.txt" + --video_column "videos.txt" + --train_resolution "81x768x1360" # (frames x height x width), frames should be 8N+1 and height, width should be multiples of 16 +) + +# Training Configuration +TRAIN_ARGS=( + --train_epochs 10 # number of training epochs + --seed 42 # random seed + + ######### Please keep consistent with deepspeed config file ########## + --batch_size 1 + --gradient_accumulation_steps 1 + --mixed_precision "bf16" # ["no", "fp16"] Only CogVideoX-2B supports fp16 training + ######################################################################## +) + +# System Configuration +SYSTEM_ARGS=( + --num_workers 8 + --pin_memory True + --nccl_timeout 1800 +) + +# Checkpointing Configuration +CHECKPOINT_ARGS=( + --checkpointing_steps 10 # save checkpoint every x steps + --checkpointing_limit 2 # maximum number of checkpoints to keep, after which the oldest one is deleted + # --resume_from_checkpoint "/absolute/path/to/checkpoint_dir" # if you want to resume from a checkpoint, otherwise, comment this line +) + +# Validation Configuration +VALIDATION_ARGS=( + --do_validation false # ["true", "false"] + --validation_dir "/absolute/path/to/validation_set" + --validation_steps 20 # should be multiple of checkpointing_steps + --validation_prompts "prompts.txt" + --gen_fps 16 +) + +# Combine all arguments and launch training +accelerate launch --config_file accelerate_config.yaml train.py \ + "${MODEL_ARGS[@]}" \ + "${OUTPUT_ARGS[@]}" \ + "${DATA_ARGS[@]}" \ + "${TRAIN_ARGS[@]}" \ + "${SYSTEM_ARGS[@]}" \ + "${CHECKPOINT_ARGS[@]}" \ + "${VALIDATION_ARGS[@]}" diff --git a/finetune/trainer.py b/finetune/trainer.py new file mode 100644 index 0000000..5746fee --- /dev/null +++ b/finetune/trainer.py @@ -0,0 +1,811 @@ +import hashlib +import json +import logging +import math +from datetime import timedelta +from pathlib import Path +from typing import Any, Dict, List, Tuple + +import diffusers +import torch +import transformers +import wandb +from accelerate.accelerator import Accelerator, DistributedType +from accelerate.logging import get_logger +from accelerate.utils import ( + DistributedDataParallelKwargs, + InitProcessGroupKwargs, + ProjectConfiguration, + gather_object, + set_seed, +) +from diffusers.optimization import get_scheduler +from diffusers.pipelines import DiffusionPipeline +from diffusers.utils.export_utils import export_to_video +from peft import LoraConfig, get_peft_model_state_dict, set_peft_model_state_dict +from PIL import Image +from torch.utils.data import DataLoader, Dataset +from tqdm import tqdm + +from finetune.constants import LOG_LEVEL, LOG_NAME +from finetune.datasets import I2VDatasetWithResize, T2VDatasetWithResize +from finetune.datasets.utils import ( + load_images, + load_prompts, + load_videos, + preprocess_image_with_resize, + preprocess_video_with_resize, +) +from finetune.schemas import Args, Components, State +from finetune.utils import ( + cast_training_params, + free_memory, + get_intermediate_ckpt_path, + get_latest_ckpt_path_to_resume_from, + get_memory_statistics, + get_optimizer, + string_to_filename, + unload_model, + unwrap_model, +) + + +logger = get_logger(LOG_NAME, LOG_LEVEL) + +_DTYPE_MAP = { + "fp32": torch.float32, + "fp16": torch.float16, # FP16 is Only Support for CogVideoX-2B + "bf16": torch.bfloat16, +} + + +class Trainer: + # If set, should be a list of components to unload (refer to `Components``) + UNLOAD_LIST: List[str] = None + + def __init__(self, args: Args) -> None: + self.args = args + self.state = State( + weight_dtype=self.__get_training_dtype(), + train_frames=self.args.train_resolution[0], + train_height=self.args.train_resolution[1], + train_width=self.args.train_resolution[2], + ) + + self.components: Components = self.load_components() + self.accelerator: Accelerator = None + self.dataset: Dataset = None + self.data_loader: DataLoader = None + + self.optimizer = None + self.lr_scheduler = None + + self._init_distributed() + self._init_logging() + self._init_directories() + + self.state.using_deepspeed = self.accelerator.state.deepspeed_plugin is not None + + def _init_distributed(self): + logging_dir = Path(self.args.output_dir, "logs") + project_config = ProjectConfiguration( + project_dir=self.args.output_dir, logging_dir=logging_dir + ) + ddp_kwargs = DistributedDataParallelKwargs(find_unused_parameters=True) + init_process_group_kwargs = InitProcessGroupKwargs( + backend="nccl", timeout=timedelta(seconds=self.args.nccl_timeout) + ) + mixed_precision = "no" if torch.backends.mps.is_available() else self.args.mixed_precision + report_to = None if self.args.report_to.lower() == "none" else self.args.report_to + + accelerator = Accelerator( + project_config=project_config, + gradient_accumulation_steps=self.args.gradient_accumulation_steps, + mixed_precision=mixed_precision, + log_with=report_to, + kwargs_handlers=[ddp_kwargs, init_process_group_kwargs], + ) + + # Disable AMP for MPS. + if torch.backends.mps.is_available(): + accelerator.native_amp = False + + self.accelerator = accelerator + + if self.args.seed is not None: + set_seed(self.args.seed) + + def _init_logging(self) -> None: + logging.basicConfig( + format="%(asctime)s - %(levelname)s - %(name)s - %(message)s", + datefmt="%m/%d/%Y %H:%M:%S", + level=LOG_LEVEL, + ) + if self.accelerator.is_local_main_process: + transformers.utils.logging.set_verbosity_warning() + diffusers.utils.logging.set_verbosity_info() + else: + transformers.utils.logging.set_verbosity_error() + diffusers.utils.logging.set_verbosity_error() + + logger.info("Initialized Trainer") + logger.info(f"Accelerator state: \n{self.accelerator.state}", main_process_only=False) + + def _init_directories(self) -> None: + if self.accelerator.is_main_process: + self.args.output_dir = Path(self.args.output_dir) + self.args.output_dir.mkdir(parents=True, exist_ok=True) + + def check_setting(self) -> None: + # Check for unload_list + if self.UNLOAD_LIST is None: + logger.warning( + "\033[91mNo unload_list specified for this Trainer. All components will be loaded to GPU during training.\033[0m" + ) + else: + for name in self.UNLOAD_LIST: + if name not in self.components.model_fields: + raise ValueError(f"Invalid component name in unload_list: {name}") + + def prepare_models(self) -> None: + logger.info("Initializing models") + + if self.components.vae is not None: + if self.args.enable_slicing: + self.components.vae.enable_slicing() + if self.args.enable_tiling: + self.components.vae.enable_tiling() + + self.state.transformer_config = self.components.transformer.config + + def prepare_dataset(self) -> None: + logger.info("Initializing dataset and dataloader") + + if self.args.model_type == "i2v": + self.dataset = I2VDatasetWithResize( + **(self.args.model_dump()), + device=self.accelerator.device, + max_num_frames=self.state.train_frames, + height=self.state.train_height, + width=self.state.train_width, + trainer=self, + ) + elif self.args.model_type == "t2v": + self.dataset = T2VDatasetWithResize( + **(self.args.model_dump()), + device=self.accelerator.device, + max_num_frames=self.state.train_frames, + height=self.state.train_height, + width=self.state.train_width, + trainer=self, + ) + else: + raise ValueError(f"Invalid model type: {self.args.model_type}") + + # Prepare VAE and text encoder for encoding + self.components.vae.requires_grad_(False) + self.components.text_encoder.requires_grad_(False) + self.components.vae = self.components.vae.to( + self.accelerator.device, dtype=self.state.weight_dtype + ) + self.components.text_encoder = self.components.text_encoder.to( + self.accelerator.device, dtype=self.state.weight_dtype + ) + + # Precompute latent for video and prompt embedding + logger.info("Precomputing latent for video and prompt embedding ...") + tmp_data_loader = torch.utils.data.DataLoader( + self.dataset, + collate_fn=self.collate_fn, + batch_size=1, + num_workers=0, + pin_memory=self.args.pin_memory, + ) + tmp_data_loader = self.accelerator.prepare_data_loader(tmp_data_loader) + for _ in tmp_data_loader: + ... + self.accelerator.wait_for_everyone() + logger.info("Precomputing latent for video and prompt embedding ... Done") + + unload_model(self.components.vae) + unload_model(self.components.text_encoder) + free_memory() + + self.data_loader = torch.utils.data.DataLoader( + self.dataset, + collate_fn=self.collate_fn, + batch_size=self.args.batch_size, + num_workers=self.args.num_workers, + pin_memory=self.args.pin_memory, + shuffle=True, + ) + + def prepare_trainable_parameters(self): + logger.info("Initializing trainable parameters") + + # For mixed precision training we cast all non-trainable weights to half-precision + # as these weights are only used for inference, keeping weights in full precision is not required. + weight_dtype = self.state.weight_dtype + + if torch.backends.mps.is_available() and weight_dtype == torch.bfloat16: + # due to pytorch#99272, MPS does not yet support bfloat16. + raise ValueError( + "Mixed precision training with bfloat16 is not supported on MPS. Please use fp16 (recommended) or fp32 instead." + ) + + # For LoRA, we freeze all the parameters + # For SFT, we train all the parameters in transformer model + for attr_name, component in vars(self.components).items(): + if hasattr(component, "requires_grad_"): + if self.args.training_type == "sft" and attr_name == "transformer": + component.requires_grad_(True) + else: + component.requires_grad_(False) + + if self.args.training_type == "lora": + transformer_lora_config = LoraConfig( + r=self.args.rank, + lora_alpha=self.args.lora_alpha, + init_lora_weights=True, + target_modules=self.args.target_modules, + ) + self.components.transformer.add_adapter(transformer_lora_config) + self.__prepare_saving_loading_hooks(transformer_lora_config) + + # Load components needed for training to GPU (except transformer), and cast them to the specified data type + ignore_list = ["transformer"] + self.UNLOAD_LIST + self.__move_components_to_device(dtype=weight_dtype, ignore_list=ignore_list) + + if self.args.gradient_checkpointing: + self.components.transformer.enable_gradient_checkpointing() + + def prepare_optimizer(self) -> None: + logger.info("Initializing optimizer and lr scheduler") + + # Make sure the trainable params are in float32 + cast_training_params([self.components.transformer], dtype=torch.float32) + + # For LoRA, we only want to train the LoRA weights + # For SFT, we want to train all the parameters + trainable_parameters = list( + filter(lambda p: p.requires_grad, self.components.transformer.parameters()) + ) + transformer_parameters_with_lr = { + "params": trainable_parameters, + "lr": self.args.learning_rate, + } + params_to_optimize = [transformer_parameters_with_lr] + self.state.num_trainable_parameters = sum(p.numel() for p in trainable_parameters) + + use_deepspeed_opt = ( + self.accelerator.state.deepspeed_plugin is not None + and "optimizer" in self.accelerator.state.deepspeed_plugin.deepspeed_config + ) + optimizer = get_optimizer( + params_to_optimize=params_to_optimize, + optimizer_name=self.args.optimizer, + learning_rate=self.args.learning_rate, + beta1=self.args.beta1, + beta2=self.args.beta2, + beta3=self.args.beta3, + epsilon=self.args.epsilon, + weight_decay=self.args.weight_decay, + use_deepspeed=use_deepspeed_opt, + ) + + num_update_steps_per_epoch = math.ceil( + len(self.data_loader) / self.args.gradient_accumulation_steps + ) + if self.args.train_steps is None: + self.args.train_steps = self.args.train_epochs * num_update_steps_per_epoch + self.state.overwrote_max_train_steps = True + + use_deepspeed_lr_scheduler = ( + self.accelerator.state.deepspeed_plugin is not None + and "scheduler" in self.accelerator.state.deepspeed_plugin.deepspeed_config + ) + total_training_steps = self.args.train_steps * self.accelerator.num_processes + num_warmup_steps = self.args.lr_warmup_steps * self.accelerator.num_processes + + if use_deepspeed_lr_scheduler: + from accelerate.utils import DummyScheduler + + lr_scheduler = DummyScheduler( + name=self.args.lr_scheduler, + optimizer=optimizer, + total_num_steps=total_training_steps, + num_warmup_steps=num_warmup_steps, + ) + else: + lr_scheduler = get_scheduler( + name=self.args.lr_scheduler, + optimizer=optimizer, + num_warmup_steps=num_warmup_steps, + num_training_steps=total_training_steps, + num_cycles=self.args.lr_num_cycles, + power=self.args.lr_power, + ) + + self.optimizer = optimizer + self.lr_scheduler = lr_scheduler + + def prepare_for_training(self) -> None: + self.components.transformer, self.optimizer, self.data_loader, self.lr_scheduler = ( + self.accelerator.prepare( + self.components.transformer, self.optimizer, self.data_loader, self.lr_scheduler + ) + ) + + # We need to recalculate our total training steps as the size of the training dataloader may have changed. + num_update_steps_per_epoch = math.ceil( + len(self.data_loader) / self.args.gradient_accumulation_steps + ) + if self.state.overwrote_max_train_steps: + self.args.train_steps = self.args.train_epochs * num_update_steps_per_epoch + # Afterwards we recalculate our number of training epochs + self.args.train_epochs = math.ceil(self.args.train_steps / num_update_steps_per_epoch) + self.state.num_update_steps_per_epoch = num_update_steps_per_epoch + + def prepare_for_validation(self): + validation_prompts = load_prompts(self.args.validation_dir / self.args.validation_prompts) + + if self.args.validation_images is not None: + validation_images = load_images(self.args.validation_dir / self.args.validation_images) + else: + validation_images = [None] * len(validation_prompts) + + if self.args.validation_videos is not None: + validation_videos = load_videos(self.args.validation_dir / self.args.validation_videos) + else: + validation_videos = [None] * len(validation_prompts) + + self.state.validation_prompts = validation_prompts + self.state.validation_images = validation_images + self.state.validation_videos = validation_videos + + def prepare_trackers(self) -> None: + logger.info("Initializing trackers") + + tracker_name = self.args.tracker_name or "finetrainers-experiment" + self.accelerator.init_trackers(tracker_name, config=self.args.model_dump()) + + def train(self) -> None: + logger.info("Starting training") + + memory_statistics = get_memory_statistics() + logger.info(f"Memory before training start: {json.dumps(memory_statistics, indent=4)}") + + self.state.total_batch_size_count = ( + self.args.batch_size + * self.accelerator.num_processes + * self.args.gradient_accumulation_steps + ) + info = { + "trainable parameters": self.state.num_trainable_parameters, + "total samples": len(self.dataset), + "train epochs": self.args.train_epochs, + "train steps": self.args.train_steps, + "batches per device": self.args.batch_size, + "total batches observed per epoch": len(self.data_loader), + "train batch size total count": self.state.total_batch_size_count, + "gradient accumulation steps": self.args.gradient_accumulation_steps, + } + logger.info(f"Training configuration: {json.dumps(info, indent=4)}") + + global_step = 0 + first_epoch = 0 + initial_global_step = 0 + + # Potentially load in the weights and states from a previous save + ( + resume_from_checkpoint_path, + initial_global_step, + global_step, + first_epoch, + ) = get_latest_ckpt_path_to_resume_from( + resume_from_checkpoint=self.args.resume_from_checkpoint, + num_update_steps_per_epoch=self.state.num_update_steps_per_epoch, + ) + if resume_from_checkpoint_path is not None: + self.accelerator.load_state(resume_from_checkpoint_path) + + progress_bar = tqdm( + range(0, self.args.train_steps), + initial=initial_global_step, + desc="Training steps", + disable=not self.accelerator.is_local_main_process, + ) + + accelerator = self.accelerator + generator = torch.Generator(device=accelerator.device) + if self.args.seed is not None: + generator = generator.manual_seed(self.args.seed) + self.state.generator = generator + + free_memory() + for epoch in range(first_epoch, self.args.train_epochs): + logger.debug(f"Starting epoch ({epoch + 1}/{self.args.train_epochs})") + + self.components.transformer.train() + models_to_accumulate = [self.components.transformer] + + for step, batch in enumerate(self.data_loader): + logger.debug(f"Starting step {step + 1}") + logs = {} + + with accelerator.accumulate(models_to_accumulate): + # These weighting schemes use a uniform timestep sampling and instead post-weight the loss + loss = self.compute_loss(batch) + accelerator.backward(loss) + + if accelerator.sync_gradients: + if accelerator.distributed_type == DistributedType.DEEPSPEED: + grad_norm = self.components.transformer.get_global_grad_norm() + # In some cases the grad norm may not return a float + if torch.is_tensor(grad_norm): + grad_norm = grad_norm.item() + else: + grad_norm = accelerator.clip_grad_norm_( + self.components.transformer.parameters(), self.args.max_grad_norm + ) + if torch.is_tensor(grad_norm): + grad_norm = grad_norm.item() + + logs["grad_norm"] = grad_norm + + self.optimizer.step() + self.lr_scheduler.step() + self.optimizer.zero_grad() + + # Checks if the accelerator has performed an optimization step behind the scenes + if accelerator.sync_gradients: + progress_bar.update(1) + global_step += 1 + self.__maybe_save_checkpoint(global_step) + + logs["loss"] = loss.detach().item() + logs["lr"] = self.lr_scheduler.get_last_lr()[0] + progress_bar.set_postfix(logs) + + # Maybe run validation + should_run_validation = ( + self.args.do_validation and global_step % self.args.validation_steps == 0 + ) + if should_run_validation: + del loss + free_memory() + self.validate(global_step) + + accelerator.log(logs, step=global_step) + + if global_step >= self.args.train_steps: + break + + memory_statistics = get_memory_statistics() + logger.info( + f"Memory after epoch {epoch + 1}: {json.dumps(memory_statistics, indent=4)}" + ) + + accelerator.wait_for_everyone() + self.__maybe_save_checkpoint(global_step, must_save=True) + if self.args.do_validation: + free_memory() + self.validate(global_step) + + del self.components + free_memory() + memory_statistics = get_memory_statistics() + logger.info(f"Memory after training end: {json.dumps(memory_statistics, indent=4)}") + + accelerator.end_training() + + def validate(self, step: int) -> None: + logger.info("Starting validation") + + accelerator = self.accelerator + num_validation_samples = len(self.state.validation_prompts) + + if num_validation_samples == 0: + logger.warning("No validation samples found. Skipping validation.") + return + + self.components.transformer.eval() + torch.set_grad_enabled(False) + + memory_statistics = get_memory_statistics() + logger.info(f"Memory before validation start: {json.dumps(memory_statistics, indent=4)}") + + ##### Initialize pipeline ##### + pipe = self.initialize_pipeline() + + if self.state.using_deepspeed: + # Can't using model_cpu_offload in deepspeed, + # so we need to move all components in pipe to device + # pipe.to(self.accelerator.device, dtype=self.state.weight_dtype) + self.__move_components_to_device( + dtype=self.state.weight_dtype, ignore_list=["transformer"] + ) + else: + # if not using deepspeed, use model_cpu_offload to further reduce memory usage + # Or use pipe.enable_sequential_cpu_offload() to further reduce memory usage + pipe.enable_model_cpu_offload(device=self.accelerator.device) + + # Convert all model weights to training dtype + # Note, this will change LoRA weights in self.components.transformer to training dtype, rather than keep them in fp32 + pipe = pipe.to(dtype=self.state.weight_dtype) + + ################################# + + all_processes_artifacts = [] + for i in range(num_validation_samples): + if self.state.using_deepspeed and self.accelerator.deepspeed_plugin.zero_stage != 3: + # Skip current validation on all processes but one + if i % accelerator.num_processes != accelerator.process_index: + continue + + prompt = self.state.validation_prompts[i] + image = self.state.validation_images[i] + video = self.state.validation_videos[i] + + if image is not None: + image = preprocess_image_with_resize( + image, self.state.train_height, self.state.train_width + ) + # Convert image tensor (C, H, W) to PIL images + image = image.to(torch.uint8) + image = image.permute(1, 2, 0).cpu().numpy() + image = Image.fromarray(image) + + if video is not None: + video = preprocess_video_with_resize( + video, self.state.train_frames, self.state.train_height, self.state.train_width + ) + # Convert video tensor (F, C, H, W) to list of PIL images + video = video.round().clamp(0, 255).to(torch.uint8) + video = [Image.fromarray(frame.permute(1, 2, 0).cpu().numpy()) for frame in video] + + logger.debug( + f"Validating sample {i + 1}/{num_validation_samples} on process {accelerator.process_index}. Prompt: {prompt}", + main_process_only=False, + ) + validation_artifacts = self.validation_step( + {"prompt": prompt, "image": image, "video": video}, pipe + ) + + if ( + self.state.using_deepspeed + and self.accelerator.deepspeed_plugin.zero_stage == 3 + and not accelerator.is_main_process + ): + continue + + prompt_filename = string_to_filename(prompt)[:25] + # Calculate hash of reversed prompt as a unique identifier + reversed_prompt = prompt[::-1] + hash_suffix = hashlib.md5(reversed_prompt.encode()).hexdigest()[:5] + + artifacts = { + "image": {"type": "image", "value": image}, + "video": {"type": "video", "value": video}, + } + for i, (artifact_type, artifact_value) in enumerate(validation_artifacts): + artifacts.update( + {f"artifact_{i}": {"type": artifact_type, "value": artifact_value}} + ) + logger.debug( + f"Validation artifacts on process {accelerator.process_index}: {list(artifacts.keys())}", + main_process_only=False, + ) + + for key, value in list(artifacts.items()): + artifact_type = value["type"] + artifact_value = value["value"] + if artifact_type not in ["image", "video"] or artifact_value is None: + continue + + extension = "png" if artifact_type == "image" else "mp4" + filename = f"validation-{step}-{accelerator.process_index}-{prompt_filename}-{hash_suffix}.{extension}" + validation_path = self.args.output_dir / "validation_res" + validation_path.mkdir(parents=True, exist_ok=True) + filename = str(validation_path / filename) + + if artifact_type == "image": + logger.debug(f"Saving image to {filename}") + artifact_value.save(filename) + artifact_value = wandb.Image(filename) + elif artifact_type == "video": + logger.debug(f"Saving video to {filename}") + export_to_video(artifact_value, filename, fps=self.args.gen_fps) + artifact_value = wandb.Video(filename, caption=prompt) + + all_processes_artifacts.append(artifact_value) + + all_artifacts = gather_object(all_processes_artifacts) + + if accelerator.is_main_process: + tracker_key = "validation" + for tracker in accelerator.trackers: + if tracker.name == "wandb": + image_artifacts = [ + artifact for artifact in all_artifacts if isinstance(artifact, wandb.Image) + ] + video_artifacts = [ + artifact for artifact in all_artifacts if isinstance(artifact, wandb.Video) + ] + tracker.log( + { + tracker_key: {"images": image_artifacts, "videos": video_artifacts}, + }, + step=step, + ) + + ########## Clean up ########## + if self.state.using_deepspeed: + del pipe + # Unload models except those needed for training + self.__move_components_to_cpu(unload_list=self.UNLOAD_LIST) + else: + pipe.remove_all_hooks() + del pipe + # Load models except those not needed for training + self.__move_components_to_device( + dtype=self.state.weight_dtype, ignore_list=self.UNLOAD_LIST + ) + self.components.transformer.to(self.accelerator.device, dtype=self.state.weight_dtype) + + # Change trainable weights back to fp32 to keep with dtype after prepare the model + cast_training_params([self.components.transformer], dtype=torch.float32) + + free_memory() + accelerator.wait_for_everyone() + ################################ + + memory_statistics = get_memory_statistics() + logger.info(f"Memory after validation end: {json.dumps(memory_statistics, indent=4)}") + torch.cuda.reset_peak_memory_stats(accelerator.device) + + torch.set_grad_enabled(True) + self.components.transformer.train() + + def fit(self): + self.check_setting() + self.prepare_models() + self.prepare_dataset() + self.prepare_trainable_parameters() + self.prepare_optimizer() + self.prepare_for_training() + if self.args.do_validation: + self.prepare_for_validation() + self.prepare_trackers() + self.train() + + def collate_fn(self, examples: List[Dict[str, Any]]): + raise NotImplementedError + + def load_components(self) -> Components: + raise NotImplementedError + + def initialize_pipeline(self) -> DiffusionPipeline: + raise NotImplementedError + + def encode_video(self, video: torch.Tensor) -> torch.Tensor: + # shape of input video: [B, C, F, H, W], where B = 1 + # shape of output video: [B, C', F', H', W'], where B = 1 + raise NotImplementedError + + def encode_text(self, text: str) -> torch.Tensor: + # shape of output text: [batch size, sequence length, embedding dimension] + raise NotImplementedError + + def compute_loss(self, batch) -> torch.Tensor: + raise NotImplementedError + + def validation_step(self) -> List[Tuple[str, Image.Image | List[Image.Image]]]: + raise NotImplementedError + + def __get_training_dtype(self) -> torch.dtype: + if self.args.mixed_precision == "no": + return _DTYPE_MAP["fp32"] + elif self.args.mixed_precision == "fp16": + return _DTYPE_MAP["fp16"] + elif self.args.mixed_precision == "bf16": + return _DTYPE_MAP["bf16"] + else: + raise ValueError(f"Invalid mixed precision: {self.args.mixed_precision}") + + def __move_components_to_device(self, dtype, ignore_list: List[str] = []): + ignore_list = set(ignore_list) + components = self.components.model_dump() + for name, component in components.items(): + if not isinstance(component, type) and hasattr(component, "to"): + if name not in ignore_list: + setattr( + self.components, name, component.to(self.accelerator.device, dtype=dtype) + ) + + def __move_components_to_cpu(self, unload_list: List[str] = []): + unload_list = set(unload_list) + components = self.components.model_dump() + for name, component in components.items(): + if not isinstance(component, type) and hasattr(component, "to"): + if name in unload_list: + setattr(self.components, name, component.to("cpu")) + + def __prepare_saving_loading_hooks(self, transformer_lora_config): + # create custom saving & loading hooks so that `accelerator.save_state(...)` serializes in a nice format + def save_model_hook(models, weights, output_dir): + if self.accelerator.is_main_process: + transformer_lora_layers_to_save = None + + for model in models: + if isinstance( + unwrap_model(self.accelerator, model), + type(unwrap_model(self.accelerator, self.components.transformer)), + ): + model = unwrap_model(self.accelerator, model) + transformer_lora_layers_to_save = get_peft_model_state_dict(model) + else: + raise ValueError(f"Unexpected save model: {model.__class__}") + + # make sure to pop weight so that corresponding model is not saved again + if weights: + weights.pop() + + self.components.pipeline_cls.save_lora_weights( + output_dir, + transformer_lora_layers=transformer_lora_layers_to_save, + ) + + def load_model_hook(models, input_dir): + if not self.accelerator.distributed_type == DistributedType.DEEPSPEED: + while len(models) > 0: + model = models.pop() + if isinstance( + unwrap_model(self.accelerator, model), + type(unwrap_model(self.accelerator, self.components.transformer)), + ): + transformer_ = unwrap_model(self.accelerator, model) + else: + raise ValueError( + f"Unexpected save model: {unwrap_model(self.accelerator, model).__class__}" + ) + else: + transformer_ = unwrap_model( + self.accelerator, self.components.transformer + ).__class__.from_pretrained(self.args.model_path, subfolder="transformer") + transformer_.add_adapter(transformer_lora_config) + + lora_state_dict = self.components.pipeline_cls.lora_state_dict(input_dir) + transformer_state_dict = { + f'{k.replace("transformer.", "")}': v + for k, v in lora_state_dict.items() + if k.startswith("transformer.") + } + incompatible_keys = set_peft_model_state_dict( + transformer_, transformer_state_dict, adapter_name="default" + ) + if incompatible_keys is not None: + # check only for unexpected keys + unexpected_keys = getattr(incompatible_keys, "unexpected_keys", None) + if unexpected_keys: + logger.warning( + f"Loading adapter weights from state_dict led to unexpected keys not found in the model: " + f" {unexpected_keys}. " + ) + + self.accelerator.register_save_state_pre_hook(save_model_hook) + self.accelerator.register_load_state_pre_hook(load_model_hook) + + def __maybe_save_checkpoint(self, global_step: int, must_save: bool = False): + if ( + self.accelerator.distributed_type == DistributedType.DEEPSPEED + or self.accelerator.is_main_process + ): + if must_save or global_step % self.args.checkpointing_steps == 0: + # for training + save_path = get_intermediate_ckpt_path( + checkpointing_limit=self.args.checkpointing_limit, + step=global_step, + output_dir=self.args.output_dir, + ) + self.accelerator.save_state(save_path, safe_serialization=True) diff --git a/finetune/utils/__init__.py b/finetune/utils/__init__.py new file mode 100644 index 0000000..9ff4912 --- /dev/null +++ b/finetune/utils/__init__.py @@ -0,0 +1,5 @@ +from .checkpointing import * +from .file_utils import * +from .memory_utils import * +from .optimizer_utils import * +from .torch_utils import * diff --git a/finetune/utils/checkpointing.py b/finetune/utils/checkpointing.py new file mode 100644 index 0000000..9c1ccb5 --- /dev/null +++ b/finetune/utils/checkpointing.py @@ -0,0 +1,57 @@ +import os +from pathlib import Path +from typing import Tuple + +from accelerate.logging import get_logger + +from finetune.constants import LOG_LEVEL, LOG_NAME + +from ..utils.file_utils import delete_files, find_files + + +logger = get_logger(LOG_NAME, LOG_LEVEL) + + +def get_latest_ckpt_path_to_resume_from( + resume_from_checkpoint: str | None, num_update_steps_per_epoch: int +) -> Tuple[str | None, int, int, int]: + if resume_from_checkpoint is None: + initial_global_step = 0 + global_step = 0 + first_epoch = 0 + resume_from_checkpoint_path = None + else: + resume_from_checkpoint_path = Path(resume_from_checkpoint) + if not resume_from_checkpoint_path.exists(): + logger.info( + f"Checkpoint '{resume_from_checkpoint}' does not exist. Starting a new training run." + ) + initial_global_step = 0 + global_step = 0 + first_epoch = 0 + resume_from_checkpoint_path = None + else: + logger.info(f"Resuming from checkpoint {resume_from_checkpoint}") + global_step = int(resume_from_checkpoint_path.name.split("-")[1]) + + initial_global_step = global_step + first_epoch = global_step // num_update_steps_per_epoch + + return resume_from_checkpoint_path, initial_global_step, global_step, first_epoch + + +def get_intermediate_ckpt_path(checkpointing_limit: int, step: int, output_dir: str) -> str: + # before saving state, check if this save would set us over the `checkpointing_limit` + if checkpointing_limit is not None: + checkpoints = find_files(output_dir, prefix="checkpoint") + + # before we save the new checkpoint, we need to have at_most `checkpoints_total_limit - 1` checkpoints + if len(checkpoints) >= checkpointing_limit: + num_to_remove = len(checkpoints) - checkpointing_limit + 1 + checkpoints_to_remove = checkpoints[0:num_to_remove] + delete_files(checkpoints_to_remove) + + logger.info(f"Checkpointing at step {step}") + save_path = os.path.join(output_dir, f"checkpoint-{step}") + logger.info(f"Saving state to {save_path}") + return save_path diff --git a/finetune/utils/file_utils.py b/finetune/utils/file_utils.py new file mode 100644 index 0000000..38b1105 --- /dev/null +++ b/finetune/utils/file_utils.py @@ -0,0 +1,48 @@ +import logging +import os +import shutil +from pathlib import Path +from typing import Any, Dict, List, Union + +from accelerate.logging import get_logger + +from finetune.constants import LOG_LEVEL, LOG_NAME + + +logger = get_logger(LOG_NAME, LOG_LEVEL) + + +def find_files(dir: Union[str, Path], prefix: str = "checkpoint") -> List[str]: + if not isinstance(dir, Path): + dir = Path(dir) + if not dir.exists(): + return [] + checkpoints = os.listdir(dir.as_posix()) + checkpoints = [c for c in checkpoints if c.startswith(prefix)] + checkpoints = sorted(checkpoints, key=lambda x: int(x.split("-")[1])) + checkpoints = [dir / c for c in checkpoints] + return checkpoints + + +def delete_files(dirs: Union[str, List[str], Path, List[Path]]) -> None: + if not isinstance(dirs, list): + dirs = [dirs] + dirs = [Path(d) if isinstance(d, str) else d for d in dirs] + logger.info(f"Deleting files: {dirs}") + for dir in dirs: + if not dir.exists(): + continue + shutil.rmtree(dir, ignore_errors=True) + + +def string_to_filename(s: str) -> str: + return ( + s.replace(" ", "-") + .replace("/", "-") + .replace(":", "-") + .replace(".", "-") + .replace(",", "-") + .replace(";", "-") + .replace("!", "-") + .replace("?", "-") + ) diff --git a/finetune/utils/memory_utils.py b/finetune/utils/memory_utils.py new file mode 100644 index 0000000..f247e7d --- /dev/null +++ b/finetune/utils/memory_utils.py @@ -0,0 +1,66 @@ +import gc +from typing import Any, Dict, Union + +import torch +from accelerate.logging import get_logger + +from finetune.constants import LOG_LEVEL, LOG_NAME + + +logger = get_logger(LOG_NAME, LOG_LEVEL) + + +def get_memory_statistics(precision: int = 3) -> Dict[str, Any]: + memory_allocated = None + memory_reserved = None + max_memory_allocated = None + max_memory_reserved = None + + if torch.cuda.is_available(): + device = torch.cuda.current_device() + memory_allocated = torch.cuda.memory_allocated(device) + memory_reserved = torch.cuda.memory_reserved(device) + max_memory_allocated = torch.cuda.max_memory_allocated(device) + max_memory_reserved = torch.cuda.max_memory_reserved(device) + + elif torch.mps.is_available(): + memory_allocated = torch.mps.current_allocated_memory() + + else: + logger.warning("No CUDA, MPS, or ROCm device found. Memory statistics are not available.") + + return { + "memory_allocated": round(bytes_to_gigabytes(memory_allocated), ndigits=precision), + "memory_reserved": round(bytes_to_gigabytes(memory_reserved), ndigits=precision), + "max_memory_allocated": round(bytes_to_gigabytes(max_memory_allocated), ndigits=precision), + "max_memory_reserved": round(bytes_to_gigabytes(max_memory_reserved), ndigits=precision), + } + + +def bytes_to_gigabytes(x: int) -> float: + if x is not None: + return x / 1024**3 + + +def free_memory() -> None: + if torch.cuda.is_available(): + gc.collect() + torch.cuda.empty_cache() + torch.cuda.ipc_collect() + + # TODO(aryan): handle non-cuda devices + + +def unload_model(model): + model.to("cpu") + + +def make_contiguous( + x: Union[torch.Tensor, Dict[str, torch.Tensor]], +) -> Union[torch.Tensor, Dict[str, torch.Tensor]]: + if isinstance(x, torch.Tensor): + return x.contiguous() + elif isinstance(x, dict): + return {k: make_contiguous(v) for k, v in x.items()} + else: + return x diff --git a/finetune/utils/optimizer_utils.py b/finetune/utils/optimizer_utils.py new file mode 100644 index 0000000..94aeafd --- /dev/null +++ b/finetune/utils/optimizer_utils.py @@ -0,0 +1,191 @@ +import inspect + +import torch +from accelerate.logging import get_logger + +from finetune.constants import LOG_LEVEL, LOG_NAME + + +logger = get_logger(LOG_NAME, LOG_LEVEL) + + +def get_optimizer( + params_to_optimize, + optimizer_name: str = "adam", + learning_rate: float = 1e-3, + beta1: float = 0.9, + beta2: float = 0.95, + beta3: float = 0.98, + epsilon: float = 1e-8, + weight_decay: float = 1e-4, + prodigy_decouple: bool = False, + prodigy_use_bias_correction: bool = False, + prodigy_safeguard_warmup: bool = False, + use_8bit: bool = False, + use_4bit: bool = False, + use_torchao: bool = False, + use_deepspeed: bool = False, + use_cpu_offload_optimizer: bool = False, + offload_gradients: bool = False, +) -> torch.optim.Optimizer: + optimizer_name = optimizer_name.lower() + + # Use DeepSpeed optimzer + if use_deepspeed: + from accelerate.utils import DummyOptim + + return DummyOptim( + params_to_optimize, + lr=learning_rate, + betas=(beta1, beta2), + eps=epsilon, + weight_decay=weight_decay, + ) + + if use_8bit and use_4bit: + raise ValueError("Cannot set both `use_8bit` and `use_4bit` to True.") + + if (use_torchao and (use_8bit or use_4bit)) or use_cpu_offload_optimizer: + try: + import torchao + + torchao.__version__ + except ImportError: + raise ImportError( + "To use optimizers from torchao, please install the torchao library: `USE_CPP=0 pip install torchao`." + ) + + if not use_torchao and use_4bit: + raise ValueError("4-bit Optimizers are only supported with torchao.") + + # Optimizer creation + supported_optimizers = ["adam", "adamw", "prodigy", "came"] + if optimizer_name not in supported_optimizers: + logger.warning( + f"Unsupported choice of optimizer: {optimizer_name}. Supported optimizers include {supported_optimizers}. Defaulting to `AdamW`." + ) + optimizer_name = "adamw" + + if (use_8bit or use_4bit) and optimizer_name not in ["adam", "adamw"]: + raise ValueError( + "`use_8bit` and `use_4bit` can only be used with the Adam and AdamW optimizers." + ) + + if use_8bit: + try: + import bitsandbytes as bnb + except ImportError: + raise ImportError( + "To use 8-bit Adam, please install the bitsandbytes library: `pip install bitsandbytes`." + ) + + if optimizer_name == "adamw": + if use_torchao: + from torchao.prototype.low_bit_optim import AdamW4bit, AdamW8bit + + optimizer_class = ( + AdamW8bit if use_8bit else AdamW4bit if use_4bit else torch.optim.AdamW + ) + else: + optimizer_class = bnb.optim.AdamW8bit if use_8bit else torch.optim.AdamW + + init_kwargs = { + "betas": (beta1, beta2), + "eps": epsilon, + "weight_decay": weight_decay, + } + + elif optimizer_name == "adam": + if use_torchao: + from torchao.prototype.low_bit_optim import Adam4bit, Adam8bit + + optimizer_class = Adam8bit if use_8bit else Adam4bit if use_4bit else torch.optim.Adam + else: + optimizer_class = bnb.optim.Adam8bit if use_8bit else torch.optim.Adam + + init_kwargs = { + "betas": (beta1, beta2), + "eps": epsilon, + "weight_decay": weight_decay, + } + + elif optimizer_name == "prodigy": + try: + import prodigyopt + except ImportError: + raise ImportError( + "To use Prodigy, please install the prodigyopt library: `pip install prodigyopt`" + ) + + optimizer_class = prodigyopt.Prodigy + + if learning_rate <= 0.1: + logger.warning( + "Learning rate is too low. When using prodigy, it's generally better to set learning rate around 1.0" + ) + + init_kwargs = { + "lr": learning_rate, + "betas": (beta1, beta2), + "beta3": beta3, + "eps": epsilon, + "weight_decay": weight_decay, + "decouple": prodigy_decouple, + "use_bias_correction": prodigy_use_bias_correction, + "safeguard_warmup": prodigy_safeguard_warmup, + } + + elif optimizer_name == "came": + try: + import came_pytorch + except ImportError: + raise ImportError( + "To use CAME, please install the came-pytorch library: `pip install came-pytorch`" + ) + + optimizer_class = came_pytorch.CAME + + init_kwargs = { + "lr": learning_rate, + "eps": (1e-30, 1e-16), + "betas": (beta1, beta2, beta3), + "weight_decay": weight_decay, + } + + if use_cpu_offload_optimizer: + from torchao.prototype.low_bit_optim import CPUOffloadOptimizer + + if "fused" in inspect.signature(optimizer_class.__init__).parameters: + init_kwargs.update({"fused": True}) + + optimizer = CPUOffloadOptimizer( + params_to_optimize, + optimizer_class=optimizer_class, + offload_gradients=offload_gradients, + **init_kwargs, + ) + else: + optimizer = optimizer_class(params_to_optimize, **init_kwargs) + + return optimizer + + +def gradient_norm(parameters): + norm = 0 + for param in parameters: + if param.grad is None: + continue + local_norm = param.grad.detach().data.norm(2) + norm += local_norm.item() ** 2 + norm = norm**0.5 + return norm + + +def max_gradient(parameters): + max_grad_value = float("-inf") + for param in parameters: + if param.grad is None: + continue + local_max_grad = param.grad.detach().data.abs().max() + max_grad_value = max(max_grad_value, local_max_grad.item()) + return max_grad_value diff --git a/finetune/utils/torch_utils.py b/finetune/utils/torch_utils.py new file mode 100644 index 0000000..867a8bf --- /dev/null +++ b/finetune/utils/torch_utils.py @@ -0,0 +1,52 @@ +from typing import Dict, List, Optional, Union + +import torch +from accelerate import Accelerator +from diffusers.utils.torch_utils import is_compiled_module + + +def unwrap_model(accelerator: Accelerator, model): + model = accelerator.unwrap_model(model) + model = model._orig_mod if is_compiled_module(model) else model + return model + + +def align_device_and_dtype( + x: Union[torch.Tensor, Dict[str, torch.Tensor]], + device: Optional[torch.device] = None, + dtype: Optional[torch.dtype] = None, +): + if isinstance(x, torch.Tensor): + if device is not None: + x = x.to(device) + if dtype is not None: + x = x.to(dtype) + elif isinstance(x, dict): + if device is not None: + x = {k: align_device_and_dtype(v, device, dtype) for k, v in x.items()} + if dtype is not None: + x = {k: align_device_and_dtype(v, device, dtype) for k, v in x.items()} + return x + + +def expand_tensor_to_dims(tensor, ndim): + while len(tensor.shape) < ndim: + tensor = tensor.unsqueeze(-1) + return tensor + + +def cast_training_params(model: Union[torch.nn.Module, List[torch.nn.Module]], dtype=torch.float32): + """ + Casts the training parameters of the model to the specified data type. + + Args: + model: The PyTorch model whose parameters will be cast. + dtype: The data type to which the model parameters will be cast. + """ + if not isinstance(model, list): + model = [model] + for m in model: + for param in m.parameters(): + # only upcast trainable parameters into fp32 + if param.requires_grad: + param.data = param.to(dtype) diff --git a/inference/cli_demo.py b/inference/cli_demo.py new file mode 100644 index 0000000..2e28165 --- /dev/null +++ b/inference/cli_demo.py @@ -0,0 +1,271 @@ +""" +This script demonstrates how to generate a video using the CogVideoX model with the Hugging Face `diffusers` pipeline. +The script supports different types of video generation, including text-to-video (t2v), image-to-video (i2v), +and video-to-video (v2v), depending on the input data and different weight. + +- text-to-video: THUDM/CogVideoX-5b, THUDM/CogVideoX-2b or THUDM/CogVideoX1.5-5b +- video-to-video: THUDM/CogVideoX-5b, THUDM/CogVideoX-2b or THUDM/CogVideoX1.5-5b +- image-to-video: THUDM/CogVideoX-5b-I2V or THUDM/CogVideoX1.5-5b-I2V + +Running the Script: +To run the script, use the following command with appropriate arguments: + +```bash +$ python cli_demo.py --prompt "A girl riding a bike." --model_path THUDM/CogVideoX1.5-5b --generate_type "t2v" +``` + +You can change `pipe.enable_sequential_cpu_offload()` to `pipe.enable_model_cpu_offload()` to speed up inference, but this will use more GPU memory + +Additional options are available to specify the model path, guidance scale, number of inference steps, video generation type, and output paths. + +""" + +import argparse +import logging +from typing import Literal, Optional + +import torch + +from diffusers import ( + CogVideoXDPMScheduler, + CogVideoXImageToVideoPipeline, + CogVideoXPipeline, + CogVideoXVideoToVideoPipeline, +) +from diffusers.utils import export_to_video, load_image, load_video + + +logging.basicConfig(level=logging.INFO) + +# Recommended resolution for each model (width, height) +RESOLUTION_MAP = { + # cogvideox1.5-* + "cogvideox1.5-5b-i2v": (768, 1360), + "cogvideox1.5-5b": (768, 1360), + # cogvideox-* + "cogvideox-5b-i2v": (480, 720), + "cogvideox-5b": (480, 720), + "cogvideox-2b": (480, 720), +} + + +def generate_video( + prompt: str, + model_path: str, + lora_path: str = None, + lora_rank: int = 128, + num_frames: int = 81, + width: Optional[int] = None, + height: Optional[int] = None, + output_path: str = "./output.mp4", + image_or_video_path: str = "", + num_inference_steps: int = 50, + guidance_scale: float = 6.0, + num_videos_per_prompt: int = 1, + dtype: torch.dtype = torch.bfloat16, + generate_type: str = Literal["t2v", "i2v", "v2v"], # i2v: image to video, v2v: video to video + seed: int = 42, + fps: int = 16, +): + """ + Generates a video based on the given prompt and saves it to the specified path. + + Parameters: + - prompt (str): The description of the video to be generated. + - model_path (str): The path of the pre-trained model to be used. + - lora_path (str): The path of the LoRA weights to be used. + - lora_rank (int): The rank of the LoRA weights. + - output_path (str): The path where the generated video will be saved. + - num_inference_steps (int): Number of steps for the inference process. More steps can result in better quality. + - num_frames (int): Number of frames to generate. CogVideoX1.0 generates 49 frames for 6 seconds at 8 fps, while CogVideoX1.5 produces either 81 or 161 frames, corresponding to 5 seconds or 10 seconds at 16 fps. + - width (int): The width of the generated video, applicable only for CogVideoX1.5-5B-I2V + - height (int): The height of the generated video, applicable only for CogVideoX1.5-5B-I2V + - guidance_scale (float): The scale for classifier-free guidance. Higher values can lead to better alignment with the prompt. + - num_videos_per_prompt (int): Number of videos to generate per prompt. + - dtype (torch.dtype): The data type for computation (default is torch.bfloat16). + - generate_type (str): The type of video generation (e.g., 't2v', 'i2v', 'v2v').· + - seed (int): The seed for reproducibility. + - fps (int): The frames per second for the generated video. + """ + + # 1. Load the pre-trained CogVideoX pipeline with the specified precision (bfloat16). + # add device_map="balanced" in the from_pretrained function and remove the enable_model_cpu_offload() + # function to use Multi GPUs. + + image = None + video = None + + model_name = model_path.split("/")[-1].lower() + desired_resolution = RESOLUTION_MAP[model_name] + if width is None or height is None: + height, width = desired_resolution + logging.info( + f"\033[1mUsing default resolution {desired_resolution} for {model_name}\033[0m" + ) + elif (height, width) != desired_resolution: + if generate_type == "i2v": + # For i2v models, use user-defined width and height + logging.warning( + f"\033[1;31mThe width({width}) and height({height}) are not recommended for {model_name}. The best resolution is {desired_resolution}.\033[0m" + ) + else: + # Otherwise, use the recommended width and height + logging.warning( + f"\033[1;31m{model_name} is not supported for custom resolution. Setting back to default resolution {desired_resolution}.\033[0m" + ) + height, width = desired_resolution + + if generate_type == "i2v": + pipe = CogVideoXImageToVideoPipeline.from_pretrained(model_path, torch_dtype=dtype) + image = load_image(image=image_or_video_path) + elif generate_type == "t2v": + pipe = CogVideoXPipeline.from_pretrained(model_path, torch_dtype=dtype) + else: + pipe = CogVideoXVideoToVideoPipeline.from_pretrained(model_path, torch_dtype=dtype) + video = load_video(image_or_video_path) + + # If you're using with lora, add this code + if lora_path: + pipe.load_lora_weights( + lora_path, weight_name="pytorch_lora_weights.safetensors", adapter_name="test_1" + ) + pipe.fuse_lora(components=["transformer"], lora_scale=1.0) + + # 2. Set Scheduler. + # Can be changed to `CogVideoXDPMScheduler` or `CogVideoXDDIMScheduler`. + # We recommend using `CogVideoXDDIMScheduler` for CogVideoX-2B. + # using `CogVideoXDPMScheduler` for CogVideoX-5B / CogVideoX-5B-I2V. + + # pipe.scheduler = CogVideoXDDIMScheduler.from_config(pipe.scheduler.config, timestep_spacing="trailing") + pipe.scheduler = CogVideoXDPMScheduler.from_config( + pipe.scheduler.config, timestep_spacing="trailing" + ) + + # 3. Enable CPU offload for the model. + # turn off if you have multiple GPUs or enough GPU memory(such as H100) and it will cost less time in inference + # and enable to("cuda") + # pipe.to("cuda") + + # pipe.enable_model_cpu_offload() + pipe.enable_sequential_cpu_offload() + pipe.vae.enable_slicing() + pipe.vae.enable_tiling() + + # 4. Generate the video frames based on the prompt. + # `num_frames` is the Number of frames to generate. + if generate_type == "i2v": + video_generate = pipe( + height=height, + width=width, + prompt=prompt, + image=image, + # The path of the image, the resolution of video will be the same as the image for CogVideoX1.5-5B-I2V, otherwise it will be 720 * 480 + num_videos_per_prompt=num_videos_per_prompt, # Number of videos to generate per prompt + num_inference_steps=num_inference_steps, # Number of inference steps + num_frames=num_frames, # Number of frames to generate + use_dynamic_cfg=True, # This id used for DPM scheduler, for DDIM scheduler, it should be False + guidance_scale=guidance_scale, + generator=torch.Generator().manual_seed(seed), # Set the seed for reproducibility + ).frames[0] + elif generate_type == "t2v": + video_generate = pipe( + height=height, + width=width, + prompt=prompt, + num_videos_per_prompt=num_videos_per_prompt, + num_inference_steps=num_inference_steps, + num_frames=num_frames, + use_dynamic_cfg=True, + guidance_scale=guidance_scale, + generator=torch.Generator().manual_seed(seed), + ).frames[0] + else: + video_generate = pipe( + height=height, + width=width, + prompt=prompt, + video=video, # The path of the video to be used as the background of the video + num_videos_per_prompt=num_videos_per_prompt, + num_inference_steps=num_inference_steps, + num_frames=num_frames, + use_dynamic_cfg=True, + guidance_scale=guidance_scale, + generator=torch.Generator().manual_seed(seed), # Set the seed for reproducibility + ).frames[0] + export_to_video(video_generate, output_path, fps=fps) + + +if __name__ == "__main__": + parser = argparse.ArgumentParser( + description="Generate a video from a text prompt using CogVideoX" + ) + parser.add_argument( + "--prompt", type=str, required=True, help="The description of the video to be generated" + ) + parser.add_argument( + "--image_or_video_path", + type=str, + default=None, + help="The path of the image to be used as the background of the video", + ) + parser.add_argument( + "--model_path", + type=str, + default="THUDM/CogVideoX1.5-5B", + help="Path of the pre-trained model use", + ) + parser.add_argument( + "--lora_path", type=str, default=None, help="The path of the LoRA weights to be used" + ) + parser.add_argument("--lora_rank", type=int, default=128, help="The rank of the LoRA weights") + parser.add_argument( + "--output_path", type=str, default="./output.mp4", help="The path save generated video" + ) + parser.add_argument( + "--guidance_scale", type=float, default=6.0, help="The scale for classifier-free guidance" + ) + parser.add_argument("--num_inference_steps", type=int, default=50, help="Inference steps") + parser.add_argument( + "--num_frames", type=int, default=81, help="Number of steps for the inference process" + ) + parser.add_argument("--width", type=int, default=None, help="The width of the generated video") + parser.add_argument( + "--height", type=int, default=None, help="The height of the generated video" + ) + parser.add_argument( + "--fps", type=int, default=16, help="The frames per second for the generated video" + ) + parser.add_argument( + "--num_videos_per_prompt", + type=int, + default=1, + help="Number of videos to generate per prompt", + ) + parser.add_argument( + "--generate_type", type=str, default="t2v", help="The type of video generation" + ) + parser.add_argument( + "--dtype", type=str, default="bfloat16", help="The data type for computation" + ) + parser.add_argument("--seed", type=int, default=42, help="The seed for reproducibility") + + args = parser.parse_args() + dtype = torch.float16 if args.dtype == "float16" else torch.bfloat16 + generate_video( + prompt=args.prompt, + model_path=args.model_path, + lora_path=args.lora_path, + lora_rank=args.lora_rank, + output_path=args.output_path, + num_frames=args.num_frames, + width=args.width, + height=args.height, + image_or_video_path=args.image_or_video_path, + num_inference_steps=args.num_inference_steps, + guidance_scale=args.guidance_scale, + num_videos_per_prompt=args.num_videos_per_prompt, + dtype=dtype, + generate_type=args.generate_type, + seed=args.seed, + fps=args.fps, + ) diff --git a/inference/cli_demo_quantization.py b/inference/cli_demo_quantization.py new file mode 100644 index 0000000..a91c04c --- /dev/null +++ b/inference/cli_demo_quantization.py @@ -0,0 +1,161 @@ +""" +This script demonstrates how to generate a video from a text prompt using CogVideoX with quantization. + +Note: + +Must install the `torchao`，`torch` library FROM SOURCE to use the quantization feature. +Only NVIDIA GPUs like H100 or higher are supported om FP-8 quantization. + +ALL quantization schemes must use with NVIDIA GPUs. + +# Run the script: + +python cli_demo_quantization.py --prompt "A girl riding a bike." --model_path THUDM/CogVideoX-2b --quantization_scheme fp8 --dtype float16 +python cli_demo_quantization.py --prompt "A girl riding a bike." --model_path THUDM/CogVideoX-5b --quantization_scheme fp8 --dtype bfloat16 + +""" + +import argparse +import os +import torch +import torch._dynamo +from diffusers import ( + AutoencoderKLCogVideoX, + CogVideoXTransformer3DModel, + CogVideoXPipeline, + CogVideoXDPMScheduler, +) +from diffusers.utils import export_to_video +from transformers import T5EncoderModel +from torchao.quantization import quantize_, int8_weight_only +from torchao.float8.inference import ActivationCasting, QuantConfig, quantize_to_float8 + +os.environ["TORCH_LOGS"] = "+dynamo,output_code,graph_breaks,recompiles" +torch._dynamo.config.suppress_errors = True +torch.set_float32_matmul_precision("high") +torch._inductor.config.conv_1x1_as_mm = True +torch._inductor.config.coordinate_descent_tuning = True +torch._inductor.config.epilogue_fusion = False +torch._inductor.config.coordinate_descent_check_all_directions = True + + +def quantize_model(part, quantization_scheme): + if quantization_scheme == "int8": + quantize_(part, int8_weight_only()) + elif quantization_scheme == "fp8": + quantize_to_float8(part, QuantConfig(ActivationCasting.DYNAMIC)) + return part + + +def generate_video( + prompt: str, + model_path: str, + output_path: str = "./output.mp4", + num_inference_steps: int = 50, + guidance_scale: float = 6.0, + num_videos_per_prompt: int = 1, + quantization_scheme: str = "fp8", + dtype: torch.dtype = torch.bfloat16, + num_frames: int = 81, + fps: int = 8, + seed: int = 42, +): + """ + Generates a video based on the given prompt and saves it to the specified path. + + Parameters: + - prompt (str): The description of the video to be generated. + - model_path (str): The path of the pre-trained model to be used. + - output_path (str): The path where the generated video will be saved. + - num_inference_steps (int): Number of steps for the inference process. More steps can result in better quality. + - guidance_scale (float): The scale for classifier-free guidance. Higher values can lead to better alignment with the prompt. + - num_videos_per_prompt (int): Number of videos to generate per prompt. + - quantization_scheme (str): The quantization scheme to use ('int8', 'fp8'). + - dtype (torch.dtype): The data type for computation (default is torch.bfloat16). + """ + text_encoder = T5EncoderModel.from_pretrained( + model_path, subfolder="text_encoder", torch_dtype=dtype + ) + text_encoder = quantize_model(part=text_encoder, quantization_scheme=quantization_scheme) + transformer = CogVideoXTransformer3DModel.from_pretrained( + model_path, subfolder="transformer", torch_dtype=dtype + ) + transformer = quantize_model(part=transformer, quantization_scheme=quantization_scheme) + vae = AutoencoderKLCogVideoX.from_pretrained(model_path, subfolder="vae", torch_dtype=dtype) + vae = quantize_model(part=vae, quantization_scheme=quantization_scheme) + pipe = CogVideoXPipeline.from_pretrained( + model_path, + text_encoder=text_encoder, + transformer=transformer, + vae=vae, + torch_dtype=dtype, + ) + pipe.scheduler = CogVideoXDPMScheduler.from_config( + pipe.scheduler.config, timestep_spacing="trailing" + ) + pipe.enable_model_cpu_offload() + pipe.vae.enable_slicing() + pipe.vae.enable_tiling() + + video = pipe( + prompt=prompt, + num_videos_per_prompt=num_videos_per_prompt, + num_inference_steps=num_inference_steps, + num_frames=num_frames, + use_dynamic_cfg=True, + guidance_scale=guidance_scale, + generator=torch.Generator(device="cuda").manual_seed(seed), + ).frames[0] + + export_to_video(video, output_path, fps=fps) + + +if __name__ == "__main__": + parser = argparse.ArgumentParser( + description="Generate a video from a text prompt using CogVideoX" + ) + parser.add_argument( + "--prompt", type=str, required=True, help="The description of the video to be generated" + ) + parser.add_argument( + "--model_path", type=str, default="THUDM/CogVideoX-5b", help="Path of the pre-trained model" + ) + parser.add_argument( + "--output_path", type=str, default="./output.mp4", help="Path to save generated video" + ) + parser.add_argument("--num_inference_steps", type=int, default=50, help="Inference steps") + parser.add_argument( + "--guidance_scale", type=float, default=6.0, help="Classifier-free guidance scale" + ) + parser.add_argument( + "--num_videos_per_prompt", type=int, default=1, help="Videos to generate per prompt" + ) + parser.add_argument( + "--dtype", type=str, default="bfloat16", help="Data type (e.g., 'float16', 'bfloat16')" + ) + parser.add_argument( + "--quantization_scheme", + type=str, + default="fp8", + choices=["int8", "fp8"], + help="Quantization scheme", + ) + parser.add_argument("--num_frames", type=int, default=81, help="Number of frames in the video") + parser.add_argument("--fps", type=int, default=16, help="Frames per second for output video") + parser.add_argument("--seed", type=int, default=42, help="Random seed for reproducibility") + + args = parser.parse_args() + dtype = torch.float16 if args.dtype == "float16" else torch.bfloat16 + generate_video( + prompt=args.prompt, + model_path=args.model_path, + output_path=args.output_path, + num_inference_steps=args.num_inference_steps, + guidance_scale=args.guidance_scale, + num_videos_per_prompt=args.num_videos_per_prompt, + quantization_scheme=args.quantization_scheme, + dtype=dtype, + num_frames=args.num_frames, + fps=args.fps, + seed=args.seed, + ) diff --git a/inference/cli_vae_demo.py b/inference/cli_vae_demo.py new file mode 100644 index 0000000..508aedb --- /dev/null +++ b/inference/cli_vae_demo.py @@ -0,0 +1,162 @@ +""" +This script is designed to demonstrate how to use the CogVideoX-2b VAE model for video encoding and decoding. +It allows you to encode a video into a latent representation, decode it back into a video, or perform both operations sequentially. +Before running the script, make sure to clone the CogVideoX Hugging Face model repository and set the +`{your local diffusers path}` argument to the path of the cloned repository. + +Command 1: Encoding Video +Encodes the video located at ../resources/videos/1.mp4 using the CogVideoX-5b VAE model. +Memory Usage: ~18GB of GPU memory for encoding. + +If you do not have enough GPU memory, we provide a pre-encoded tensor file (encoded.pt) in the resources folder, +and you can still run the decoding command. + +$ python cli_vae_demo.py --model_path {your local diffusers path}/CogVideoX-2b/vae/ --video_path ../resources/videos/1.mp4 --mode encode + +Command 2: Decoding Video + +Decodes the latent representation stored in encoded.pt back into a video. +Memory Usage: ~4GB of GPU memory for decoding. +$ python cli_vae_demo.py --model_path {your local diffusers path}/CogVideoX-2b/vae/ --encoded_path ./encoded.pt --mode decode + +Command 3: Encoding and Decoding Video +Encodes the video located at ../resources/videos/1.mp4 and then immediately decodes it. +Memory Usage: 34GB for encoding + 19GB for decoding (sequentially). +$ python cli_vae_demo.py --model_path {your local diffusers path}/CogVideoX-2b/vae/ --video_path ../resources/videos/1.mp4 --mode both +""" + +import argparse +import torch +import imageio +from diffusers import AutoencoderKLCogVideoX +from torchvision import transforms +import numpy as np + + +def encode_video(model_path, video_path, dtype, device): + """ + Loads a pre-trained AutoencoderKLCogVideoX model and encodes the video frames. + + Parameters: + - model_path (str): The path to the pre-trained model. + - video_path (str): The path to the video file. + - dtype (torch.dtype): The data type for computation. + - device (str): The device to use for computation (e.g., "cuda" or "cpu"). + + Returns: + - torch.Tensor: The encoded video frames. + """ + + model = AutoencoderKLCogVideoX.from_pretrained(model_path, torch_dtype=dtype).to(device) + + model.enable_slicing() + model.enable_tiling() + + video_reader = imageio.get_reader(video_path, "ffmpeg") + + frames = [transforms.ToTensor()(frame) for frame in video_reader] + video_reader.close() + + frames_tensor = torch.stack(frames).to(device).permute(1, 0, 2, 3).unsqueeze(0).to(dtype) + + with torch.no_grad(): + encoded_frames = model.encode(frames_tensor)[0].sample() + return encoded_frames + + +def decode_video(model_path, encoded_tensor_path, dtype, device): + """ + Loads a pre-trained AutoencoderKLCogVideoX model and decodes the encoded video frames. + + Parameters: + - model_path (str): The path to the pre-trained model. + - encoded_tensor_path (str): The path to the encoded tensor file. + - dtype (torch.dtype): The data type for computation. + - device (str): The device to use for computation (e.g., "cuda" or "cpu"). + + Returns: + - torch.Tensor: The decoded video frames. + """ + model = AutoencoderKLCogVideoX.from_pretrained(model_path, torch_dtype=dtype).to(device) + encoded_frames = torch.load(encoded_tensor_path, weights_only=True).to(device).to(dtype) + with torch.no_grad(): + decoded_frames = model.decode(encoded_frames).sample + return decoded_frames + + +def save_video(tensor, output_path): + """ + Saves the video frames to a video file. + + Parameters: + - tensor (torch.Tensor): The video frames' tensor. + - output_path (str): The path to save the output video. + """ + tensor = tensor.to(dtype=torch.float32) + frames = tensor[0].squeeze(0).permute(1, 2, 3, 0).cpu().numpy() + frames = np.clip(frames, 0, 1) * 255 + frames = frames.astype(np.uint8) + writer = imageio.get_writer(output_path + "/output.mp4", fps=8) + for frame in frames: + writer.append_data(frame) + writer.close() + + +if __name__ == "__main__": + parser = argparse.ArgumentParser(description="CogVideoX encode/decode demo") + parser.add_argument( + "--model_path", type=str, required=True, help="The path to the CogVideoX model" + ) + parser.add_argument("--video_path", type=str, help="The path to the video file (for encoding)") + parser.add_argument( + "--encoded_path", type=str, help="The path to the encoded tensor file (for decoding)" + ) + parser.add_argument( + "--output_path", type=str, default=".", help="The path to save the output file" + ) + parser.add_argument( + "--mode", + type=str, + choices=["encode", "decode", "both"], + required=True, + help="Mode: encode, decode, or both", + ) + parser.add_argument( + "--dtype", + type=str, + default="bfloat16", + help="The data type for computation (e.g., 'float16' or 'bfloat16')", + ) + parser.add_argument( + "--device", + type=str, + default="cuda", + help="The device to use for computation (e.g., 'cuda' or 'cpu')", + ) + args = parser.parse_args() + + device = torch.device(args.device) + dtype = torch.float16 if args.dtype == "float16" else torch.bfloat16 + + if args.mode == "encode": + assert args.video_path, "Video path must be provided for encoding." + encoded_output = encode_video(args.model_path, args.video_path, dtype, device) + torch.save(encoded_output, args.output_path + "/encoded.pt") + print( + f"Finished encoding the video to a tensor, save it to a file at {encoded_output}/encoded.pt" + ) + elif args.mode == "decode": + assert args.encoded_path, "Encoded tensor path must be provided for decoding." + decoded_output = decode_video(args.model_path, args.encoded_path, dtype, device) + save_video(decoded_output, args.output_path) + print( + f"Finished decoding the video and saved it to a file at {args.output_path}/output.mp4" + ) + elif args.mode == "both": + assert args.video_path, "Video path must be provided for encoding." + encoded_output = encode_video(args.model_path, args.video_path, dtype, device) + torch.save(encoded_output, args.output_path + "/encoded.pt") + decoded_output = decode_video( + args.model_path, args.output_path + "/encoded.pt", dtype, device + ) + save_video(decoded_output, args.output_path) diff --git a/inference/convert_demo.py b/inference/convert_demo.py new file mode 100644 index 0000000..2c423fc --- /dev/null +++ b/inference/convert_demo.py @@ -0,0 +1,155 @@ +""" +The CogVideoX model is designed to generate high-quality videos based on detailed and highly descriptive prompts. +The model performs best when provided with refined, granular prompts, which enhance the quality of video generation. +This script is designed to assist with transforming simple user inputs into detailed prompts suitable for CogVideoX. +It can handle both text-to-video (t2v) and image-to-video (i2v) conversions. + +- For text-to-video, simply provide the prompt. +- For image-to-video, provide the path to the image file and an optional user input. +The image will be encoded and sent as part of the request to Azure OpenAI. + +### How to run: +Run the script for **text-to-video**: + $ python convert_demo.py --prompt "A girl riding a bike." --type "t2v" + +Run the script for **image-to-video**: + $ python convert_demo.py --prompt "the cat is running" --type "i2v" --image_path "/path/to/your/image.jpg" +""" + +import argparse +from openai import OpenAI, AzureOpenAI +import base64 +from mimetypes import guess_type + +sys_prompt_t2v = """You are part of a team of bots that creates videos. You work with an assistant bot that will draw anything you say in square brackets. + +For example , outputting " a beautiful morning in the woods with the sun peaking through the trees " will trigger your partner bot to output an video of a forest morning , as described. You will be prompted by people looking to create detailed , amazing videos. The way to accomplish this is to take their short prompts and make them extremely detailed and descriptive. +There are a few rules to follow: + +You will only ever output a single video description per user request. + +When modifications are requested , you should not simply make the description longer . You should refactor the entire description to integrate the suggestions. +Other times the user will not want modifications , but instead want a new image . In this case , you should ignore your previous conversation with the user. + +Video descriptions must have the same num of words as examples below. Extra words will be ignored. +""" + +sys_prompt_i2v = """ +**Objective**: **Give a highly descriptive video caption based on input image and user input. **. As an expert, delve deep into the image with a discerning eye, leveraging rich creativity, meticulous thought. When describing the details of an image, include appropriate dynamic information to ensure that the video caption contains reasonable actions and plots. If user input is not empty, then the caption should be expanded according to the user's input. + +**Note**: The input image is the first frame of the video, and the output video caption should describe the motion starting from the current image. User input is optional and can be empty. + +**Note**: Don't contain camera transitions!!! Don't contain screen switching!!! Don't contain perspective shifts !!! + +**Answering Style**: +Answers should be comprehensive, conversational, and use complete sentences. The answer should be in English no matter what the user's input is. Provide context where necessary and maintain a certain tone. Begin directly without introductory phrases like "The image/video showcases" "The photo captures" and more. For example, say "A woman is on a beach", instead of "A woman is depicted in the image". + +**Output Format**: "[highly descriptive image caption here]" + +user input: +""" + + +def image_to_url(image_path): + mime_type, _ = guess_type(image_path) + if mime_type is None: + mime_type = "application/octet-stream" + with open(image_path, "rb") as image_file: + base64_encoded_data = base64.b64encode(image_file.read()).decode("utf-8") + return f"data:{mime_type};base64,{base64_encoded_data}" + + +def convert_prompt(prompt: str, retry_times: int = 3, type: str = "t2v", image_path: str = None): + """ + Convert a prompt to a format that can be used by the model for inference + """ + + client = OpenAI() + ## If you using with Azure OpenAI, please uncomment the below line and comment the above line + # client = AzureOpenAI( + # api_key="", + # api_version="", + # azure_endpoint="" + # ) + + text = prompt.strip() + for i in range(retry_times): + if type == "t2v": + response = client.chat.completions.create( + messages=[ + {"role": "system", "content": f"{sys_prompt_t2v}"}, + { + "role": "user", + "content": 'Create an imaginative video descriptive caption or modify an earlier caption for the user input : " a girl is on the beach"', + }, + { + "role": "assistant", + "content": "A radiant woman stands on a deserted beach, arms outstretched, wearing a beige trench coat, white blouse, light blue jeans, and chic boots, against a backdrop of soft sky and sea. Moments later, she is seen mid-twirl, arms exuberant, with the lighting suggesting dawn or dusk. Then, she runs along the beach, her attire complemented by an off-white scarf and black ankle boots, the tranquil sea behind her. Finally, she holds a paper airplane, her pose reflecting joy and freedom, with the ocean's gentle waves and the sky's soft pastel hues enhancing the serene ambiance.", + }, + { + "role": "user", + "content": 'Create an imaginative video descriptive caption or modify an earlier caption for the user input : " A man jogging on a football field"', + }, + { + "role": "assistant", + "content": "A determined man in athletic attire, including a blue long-sleeve shirt, black shorts, and blue socks, jogs around a snow-covered soccer field, showcasing his solitary exercise in a quiet, overcast setting. His long dreadlocks, focused expression, and the serene winter backdrop highlight his dedication to fitness. As he moves, his attire, consisting of a blue sports sweatshirt, black athletic pants, gloves, and sneakers, grips the snowy ground. He is seen running past a chain-link fence enclosing the playground area, with a basketball hoop and children's slide, suggesting a moment of solitary exercise amidst the empty field.", + }, + { + "role": "user", + "content": 'Create an imaginative video descriptive caption or modify an earlier caption for the user input : " A woman is dancing, HD footage, close-up"', + }, + { + "role": "assistant", + "content": "A young woman with her hair in an updo and wearing a teal hoodie stands against a light backdrop, initially looking over her shoulder with a contemplative expression. She then confidently makes a subtle dance move, suggesting rhythm and movement. Next, she appears poised and focused, looking directly at the camera. Her expression shifts to one of introspection as she gazes downward slightly. Finally, she dances with confidence, her left hand over her heart, symbolizing a poignant moment, all while dressed in the same teal hoodie against a plain, light-colored background.", + }, + { + "role": "user", + "content": f'Create an imaginative video descriptive caption or modify an earlier caption in ENGLISH for the user input: " {text} "', + }, + ], + model="glm-4-plus", # glm-4-plus and gpt-4o have be tested + temperature=0.01, + top_p=0.7, + stream=False, + max_tokens=250, + ) + else: + response = client.chat.completions.create( + model="gpt-4o", + messages=[ + {"role": "system", "content": f"{sys_prompt_i2v}"}, + { + "role": "user", + "content": [ + {"type": "text", "text": prompt}, + { + "type": "image_url", + "image_url": { + "url": image_to_url(image_path), + }, + }, + ], + }, + ], + temperature=0.01, + top_p=0.7, + stream=False, + max_tokens=250, + ) + if response.choices: + return response.choices[0].message.content + return prompt + + +if __name__ == "__main__": + parser = argparse.ArgumentParser() + parser.add_argument("--prompt", type=str, required=True, help="Prompt to convert") + parser.add_argument( + "--retry_times", type=int, default=3, help="Number of times to retry the conversion" + ) + parser.add_argument("--type", type=str, default="t2v", help="Type of conversion (t2v or i2v)") + parser.add_argument("--image_path", type=str, default=None, help="Path to the image file") + args = parser.parse_args() + + converted_prompt = convert_prompt(args.prompt, args.retry_times, args.type, args.image_path) + print(converted_prompt) diff --git a/inference/ddim_inversion.py b/inference/ddim_inversion.py new file mode 100644 index 0000000..1ca682e --- /dev/null +++ b/inference/ddim_inversion.py @@ -0,0 +1,519 @@ +""" +This script performs DDIM inversion for video frames using a pre-trained model and generates +a video reconstruction based on a provided prompt. It utilizes the CogVideoX pipeline to +process video frames, apply the DDIM inverse scheduler, and produce an output video. + +**Please notice that this script is based on the CogVideoX 5B model, and would not generate +a good result for 2B variants.** + +Usage: + python ddim_inversion.py + --model-path /path/to/model + --prompt "a prompt" + --video-path /path/to/video.mp4 + --output-path /path/to/output + +For more details about the cli arguments, please run `python ddim_inversion.py --help`. + +Author: + LittleNyima +""" + +import argparse +import math +import os +from typing import Any, Dict, List, Optional, Tuple, TypedDict, Union, cast + +import torch +import torch.nn.functional as F +import torchvision.transforms as T +from diffusers.models.attention_processor import Attention, CogVideoXAttnProcessor2_0 +from diffusers.models.autoencoders import AutoencoderKLCogVideoX +from diffusers.models.embeddings import apply_rotary_emb +from diffusers.models.transformers.cogvideox_transformer_3d import ( + CogVideoXBlock, + CogVideoXTransformer3DModel, +) +from diffusers.pipelines.cogvideo.pipeline_cogvideox import CogVideoXPipeline, retrieve_timesteps +from diffusers.schedulers import CogVideoXDDIMScheduler, DDIMInverseScheduler +from diffusers.utils import export_to_video + +# Must import after torch because this can sometimes lead to a nasty segmentation fault, or stack smashing error. +# Very few bug reports but it happens. Look in decord Github issues for more relevant information. +import decord # isort: skip + + +class DDIMInversionArguments(TypedDict): + model_path: str + prompt: str + video_path: str + output_path: str + guidance_scale: float + num_inference_steps: int + skip_frames_start: int + skip_frames_end: int + frame_sample_step: Optional[int] + max_num_frames: int + width: int + height: int + fps: int + dtype: torch.dtype + seed: int + device: torch.device + + +def get_args() -> DDIMInversionArguments: + parser = argparse.ArgumentParser() + + parser.add_argument( + "--model_path", type=str, required=True, help="Path of the pretrained model" + ) + parser.add_argument( + "--prompt", type=str, required=True, help="Prompt for the direct sample procedure" + ) + parser.add_argument( + "--video_path", type=str, required=True, help="Path of the video for inversion" + ) + parser.add_argument( + "--output_path", type=str, default="output", help="Path of the output videos" + ) + parser.add_argument( + "--guidance_scale", type=float, default=6.0, help="Classifier-free guidance scale" + ) + parser.add_argument( + "--num_inference_steps", type=int, default=50, help="Number of inference steps" + ) + parser.add_argument( + "--skip_frames_start", type=int, default=0, help="Number of skipped frames from the start" + ) + parser.add_argument( + "--skip_frames_end", type=int, default=0, help="Number of skipped frames from the end" + ) + parser.add_argument( + "--frame_sample_step", type=int, default=None, help="Temporal stride of the sampled frames" + ) + parser.add_argument( + "--max_num_frames", type=int, default=81, help="Max number of sampled frames" + ) + parser.add_argument("--width", type=int, default=720, help="Resized width of the video frames") + parser.add_argument( + "--height", type=int, default=480, help="Resized height of the video frames" + ) + parser.add_argument("--fps", type=int, default=8, help="Frame rate of the output videos") + parser.add_argument( + "--dtype", type=str, default="bf16", choices=["bf16", "fp16"], help="Dtype of the model" + ) + parser.add_argument("--seed", type=int, default=42, help="Seed for the random number generator") + parser.add_argument( + "--device", type=str, default="cuda", choices=["cuda", "cpu"], help="Device for inference" + ) + + args = parser.parse_args() + args.dtype = torch.bfloat16 if args.dtype == "bf16" else torch.float16 + args.device = torch.device(args.device) + + return DDIMInversionArguments(**vars(args)) + + +class CogVideoXAttnProcessor2_0ForDDIMInversion(CogVideoXAttnProcessor2_0): + def __init__(self): + super().__init__() + + def calculate_attention( + self, + query: torch.Tensor, + key: torch.Tensor, + value: torch.Tensor, + attn: Attention, + batch_size: int, + image_seq_length: int, + text_seq_length: int, + attention_mask: Optional[torch.Tensor], + image_rotary_emb: Optional[torch.Tensor], + ) -> Tuple[torch.Tensor, torch.Tensor]: + inner_dim = key.shape[-1] + head_dim = inner_dim // attn.heads + + query = query.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2) + key = key.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2) + value = value.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2) + + if attn.norm_q is not None: + query = attn.norm_q(query) + if attn.norm_k is not None: + key = attn.norm_k(key) + + # Apply RoPE if needed + if image_rotary_emb is not None: + query[:, :, text_seq_length:] = apply_rotary_emb( + query[:, :, text_seq_length:], image_rotary_emb + ) + if not attn.is_cross_attention: + if key.size(2) == query.size(2): # Attention for reference hidden states + key[:, :, text_seq_length:] = apply_rotary_emb( + key[:, :, text_seq_length:], image_rotary_emb + ) + else: # RoPE should be applied to each group of image tokens + key[:, :, text_seq_length : text_seq_length + image_seq_length] = ( + apply_rotary_emb( + key[:, :, text_seq_length : text_seq_length + image_seq_length], + image_rotary_emb, + ) + ) + key[:, :, text_seq_length * 2 + image_seq_length :] = apply_rotary_emb( + key[:, :, text_seq_length * 2 + image_seq_length :], image_rotary_emb + ) + + hidden_states = F.scaled_dot_product_attention( + query, key, value, attn_mask=attention_mask, dropout_p=0.0, is_causal=False + ) + + hidden_states = hidden_states.transpose(1, 2).reshape(batch_size, -1, attn.heads * head_dim) + + # linear proj + hidden_states = attn.to_out[0](hidden_states) + # dropout + hidden_states = attn.to_out[1](hidden_states) + + encoder_hidden_states, hidden_states = hidden_states.split( + [text_seq_length, hidden_states.size(1) - text_seq_length], dim=1 + ) + return hidden_states, encoder_hidden_states + + def __call__( + self, + attn: Attention, + hidden_states: torch.Tensor, + encoder_hidden_states: torch.Tensor, + attention_mask: Optional[torch.Tensor] = None, + image_rotary_emb: Optional[torch.Tensor] = None, + ) -> Tuple[torch.Tensor, torch.Tensor]: + image_seq_length = hidden_states.size(1) + text_seq_length = encoder_hidden_states.size(1) + + hidden_states = torch.cat([encoder_hidden_states, hidden_states], dim=1) + + batch_size, sequence_length, _ = ( + hidden_states.shape if encoder_hidden_states is None else encoder_hidden_states.shape + ) + + if attention_mask is not None: + attention_mask = attn.prepare_attention_mask( + attention_mask, sequence_length, batch_size + ) + attention_mask = attention_mask.view( + batch_size, attn.heads, -1, attention_mask.shape[-1] + ) + + query = attn.to_q(hidden_states) + key = attn.to_k(hidden_states) + value = attn.to_v(hidden_states) + + query, query_reference = query.chunk(2) + key, key_reference = key.chunk(2) + value, value_reference = value.chunk(2) + batch_size = batch_size // 2 + + hidden_states, encoder_hidden_states = self.calculate_attention( + query=query, + key=torch.cat((key, key_reference), dim=1), + value=torch.cat((value, value_reference), dim=1), + attn=attn, + batch_size=batch_size, + image_seq_length=image_seq_length, + text_seq_length=text_seq_length, + attention_mask=attention_mask, + image_rotary_emb=image_rotary_emb, + ) + hidden_states_reference, encoder_hidden_states_reference = self.calculate_attention( + query=query_reference, + key=key_reference, + value=value_reference, + attn=attn, + batch_size=batch_size, + image_seq_length=image_seq_length, + text_seq_length=text_seq_length, + attention_mask=attention_mask, + image_rotary_emb=image_rotary_emb, + ) + + return ( + torch.cat((hidden_states, hidden_states_reference)), + torch.cat((encoder_hidden_states, encoder_hidden_states_reference)), + ) + + +class OverrideAttnProcessors: + def __init__(self, transformer: CogVideoXTransformer3DModel): + self.transformer = transformer + self.original_processors = {} + + def __enter__(self): + for block in self.transformer.transformer_blocks: + block = cast(CogVideoXBlock, block) + self.original_processors[id(block)] = block.attn1.get_processor() + block.attn1.set_processor(CogVideoXAttnProcessor2_0ForDDIMInversion()) + + def __exit__(self, _0, _1, _2): + for block in self.transformer.transformer_blocks: + block = cast(CogVideoXBlock, block) + block.attn1.set_processor(self.original_processors[id(block)]) + + +def get_video_frames( + video_path: str, + width: int, + height: int, + skip_frames_start: int, + skip_frames_end: int, + max_num_frames: int, + frame_sample_step: Optional[int], +) -> torch.FloatTensor: + with decord.bridge.use_torch(): + video_reader = decord.VideoReader(uri=video_path, width=width, height=height) + video_num_frames = len(video_reader) + start_frame = min(skip_frames_start, video_num_frames) + end_frame = max(0, video_num_frames - skip_frames_end) + + if end_frame <= start_frame: + indices = [start_frame] + elif end_frame - start_frame <= max_num_frames: + indices = list(range(start_frame, end_frame)) + else: + step = frame_sample_step or (end_frame - start_frame) // max_num_frames + indices = list(range(start_frame, end_frame, step)) + + frames = video_reader.get_batch(indices=indices) + frames = frames[:max_num_frames].float() # ensure that we don't go over the limit + + # Choose first (4k + 1) frames as this is how many is required by the VAE + selected_num_frames = frames.size(0) + remainder = (3 + selected_num_frames) % 4 + if remainder != 0: + frames = frames[:-remainder] + assert frames.size(0) % 4 == 1 + + # Normalize the frames + transform = T.Lambda(lambda x: x / 255.0 * 2.0 - 1.0) + frames = torch.stack(tuple(map(transform, frames)), dim=0) + + return frames.permute(0, 3, 1, 2).contiguous() # [F, C, H, W] + + +def encode_video_frames( + vae: AutoencoderKLCogVideoX, video_frames: torch.FloatTensor +) -> torch.FloatTensor: + video_frames = video_frames.to(device=vae.device, dtype=vae.dtype) + video_frames = video_frames.unsqueeze(0).permute(0, 2, 1, 3, 4) # [B, C, F, H, W] + latent_dist = vae.encode(x=video_frames).latent_dist.sample().transpose(1, 2) + return latent_dist * vae.config.scaling_factor + + +def export_latents_to_video( + pipeline: CogVideoXPipeline, latents: torch.FloatTensor, video_path: str, fps: int +): + video = pipeline.decode_latents(latents) + frames = pipeline.video_processor.postprocess_video(video=video, output_type="pil") + export_to_video(video_frames=frames[0], output_video_path=video_path, fps=fps) + + +# Modified from CogVideoXPipeline.__call__ +def sample( + pipeline: CogVideoXPipeline, + latents: torch.FloatTensor, + scheduler: Union[DDIMInverseScheduler, CogVideoXDDIMScheduler], + prompt: Optional[Union[str, List[str]]] = None, + negative_prompt: Optional[Union[str, List[str]]] = None, + num_inference_steps: int = 50, + guidance_scale: float = 6, + use_dynamic_cfg: bool = False, + eta: float = 0.0, + generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None, + attention_kwargs: Optional[Dict[str, Any]] = None, + reference_latents: torch.FloatTensor = None, +) -> torch.FloatTensor: + pipeline._guidance_scale = guidance_scale + pipeline._attention_kwargs = attention_kwargs + pipeline._interrupt = False + + device = pipeline._execution_device + + # here `guidance_scale` is defined analog to the guidance weight `w` of equation (2) + # of the Imagen paper: https://arxiv.org/pdf/2205.11487.pdf . `guidance_scale = 1` + # corresponds to doing no classifier free guidance. + do_classifier_free_guidance = guidance_scale > 1.0 + + # 3. Encode input prompt + prompt_embeds, negative_prompt_embeds = pipeline.encode_prompt( + prompt, + negative_prompt, + do_classifier_free_guidance, + device=device, + ) + if do_classifier_free_guidance: + prompt_embeds = torch.cat([negative_prompt_embeds, prompt_embeds], dim=0) + if reference_latents is not None: + prompt_embeds = torch.cat([prompt_embeds] * 2, dim=0) + + # 4. Prepare timesteps + timesteps, num_inference_steps = retrieve_timesteps(scheduler, num_inference_steps, device) + pipeline._num_timesteps = len(timesteps) + + # 5. Prepare latents. + latents = latents.to(device=device) * scheduler.init_noise_sigma + + # 6. Prepare extra step kwargs. TODO: Logic should ideally just be moved out of the pipeline + extra_step_kwargs = pipeline.prepare_extra_step_kwargs(generator, eta) + if isinstance( + scheduler, DDIMInverseScheduler + ): # Inverse scheduler does not accept extra kwargs + extra_step_kwargs = {} + + # 7. Create rotary embeds if required + image_rotary_emb = ( + pipeline._prepare_rotary_positional_embeddings( + height=latents.size(3) * pipeline.vae_scale_factor_spatial, + width=latents.size(4) * pipeline.vae_scale_factor_spatial, + num_frames=latents.size(1), + device=device, + ) + if pipeline.transformer.config.use_rotary_positional_embeddings + else None + ) + + # 8. Denoising loop + num_warmup_steps = max(len(timesteps) - num_inference_steps * scheduler.order, 0) + + trajectory = torch.zeros_like(latents).unsqueeze(0).repeat(len(timesteps), 1, 1, 1, 1, 1) + with pipeline.progress_bar(total=num_inference_steps) as progress_bar: + for i, t in enumerate(timesteps): + if pipeline.interrupt: + continue + + latent_model_input = ( + torch.cat([latents] * 2) if do_classifier_free_guidance else latents + ) + if reference_latents is not None: + reference = reference_latents[i] + reference = torch.cat([reference] * 2) if do_classifier_free_guidance else reference + latent_model_input = torch.cat([latent_model_input, reference], dim=0) + latent_model_input = scheduler.scale_model_input(latent_model_input, t) + + # broadcast to batch dimension in a way that's compatible with ONNX/Core ML + timestep = t.expand(latent_model_input.shape[0]) + + # predict noise model_output + noise_pred = pipeline.transformer( + hidden_states=latent_model_input, + encoder_hidden_states=prompt_embeds, + timestep=timestep, + image_rotary_emb=image_rotary_emb, + attention_kwargs=attention_kwargs, + return_dict=False, + )[0] + noise_pred = noise_pred.float() + + if reference_latents is not None: # Recover the original batch size + noise_pred, _ = noise_pred.chunk(2) + + # perform guidance + if use_dynamic_cfg: + pipeline._guidance_scale = 1 + guidance_scale * ( + ( + 1 + - math.cos( + math.pi + * ((num_inference_steps - t.item()) / num_inference_steps) ** 5.0 + ) + ) + / 2 + ) + if do_classifier_free_guidance: + noise_pred_uncond, noise_pred_text = noise_pred.chunk(2) + noise_pred = noise_pred_uncond + pipeline.guidance_scale * ( + noise_pred_text - noise_pred_uncond + ) + + # compute the noisy sample x_t-1 -> x_t + latents = scheduler.step( + noise_pred, t, latents, **extra_step_kwargs, return_dict=False + )[0] + latents = latents.to(prompt_embeds.dtype) + trajectory[i] = latents + + if i == len(timesteps) - 1 or ( + (i + 1) > num_warmup_steps and (i + 1) % scheduler.order == 0 + ): + progress_bar.update() + + # Offload all models + pipeline.maybe_free_model_hooks() + + return trajectory + + +@torch.no_grad() +def ddim_inversion( + model_path: str, + prompt: str, + video_path: str, + output_path: str, + guidance_scale: float, + num_inference_steps: int, + skip_frames_start: int, + skip_frames_end: int, + frame_sample_step: Optional[int], + max_num_frames: int, + width: int, + height: int, + fps: int, + dtype: torch.dtype, + seed: int, + device: torch.device, +): + pipeline: CogVideoXPipeline = CogVideoXPipeline.from_pretrained( + model_path, torch_dtype=dtype + ).to(device=device) + if not pipeline.transformer.config.use_rotary_positional_embeddings: + raise NotImplementedError("This script supports CogVideoX 5B model only.") + video_frames = get_video_frames( + video_path=video_path, + width=width, + height=height, + skip_frames_start=skip_frames_start, + skip_frames_end=skip_frames_end, + max_num_frames=max_num_frames, + frame_sample_step=frame_sample_step, + ).to(device=device) + video_latents = encode_video_frames(vae=pipeline.vae, video_frames=video_frames) + inverse_scheduler = DDIMInverseScheduler(**pipeline.scheduler.config) + inverse_latents = sample( + pipeline=pipeline, + latents=video_latents, + scheduler=inverse_scheduler, + prompt="", + num_inference_steps=num_inference_steps, + guidance_scale=guidance_scale, + generator=torch.Generator(device=device).manual_seed(seed), + ) + with OverrideAttnProcessors(transformer=pipeline.transformer): + recon_latents = sample( + pipeline=pipeline, + latents=torch.randn_like(video_latents), + scheduler=pipeline.scheduler, + prompt=prompt, + num_inference_steps=num_inference_steps, + guidance_scale=guidance_scale, + generator=torch.Generator(device=device).manual_seed(seed), + reference_latents=reversed(inverse_latents), + ) + filename, _ = os.path.splitext(os.path.basename(video_path)) + inverse_video_path = os.path.join(output_path, f"{filename}_inversion.mp4") + recon_video_path = os.path.join(output_path, f"{filename}_reconstruction.mp4") + export_latents_to_video(pipeline, inverse_latents[-1], inverse_video_path, fps) + export_latents_to_video(pipeline, recon_latents[-1], recon_video_path, fps) + + +if __name__ == "__main__": + arguments = get_args() + ddim_inversion(**arguments) diff --git a/inference/encoded.pt b/inference/encoded.pt new file mode 100644 index 0000000..4367fec Binary files /dev/null and b/inference/encoded.pt differ diff --git a/inference/gradio_composite_demo/README.md b/inference/gradio_composite_demo/README.md new file mode 100644 index 0000000..50dfde7 --- /dev/null +++ b/inference/gradio_composite_demo/README.md @@ -0,0 +1,45 @@ +--- +title: CogVideoX-5B +emoji: 🎥 +colorFrom: yellow +colorTo: blue +sdk: gradio +sdk_version: 4.42.0 +suggested_hardware: a10g-large +suggested_storage: large +app_port: 7860 +app_file: app.py +models: + - THUDM/CogVideoX-5b +tags: + - cogvideox + - video-generation + - thudm +short_description: Text-to-Video +disable_embedding: false +--- + +# Gradio Composite Demo + +This Gradio demo integrates the CogVideoX-5B model, allowing you to perform video inference directly in your browser. It +supports features like UpScale, RIFE, and other functionalities. + +## Environment Setup + +Set the following environment variables in your system: + ++ OPENAI_API_KEY = your_api_key ++ OPENAI_BASE_URL= your_base_url ++ GRADIO_TEMP_DIR= gradio_tmp + +## Installation + +```bash +pip install -r requirements.txt +``` + +## Running the code + +```bash +python app.py +``` diff --git a/inference/gradio_composite_demo/app.py b/inference/gradio_composite_demo/app.py new file mode 100644 index 0000000..085371f --- /dev/null +++ b/inference/gradio_composite_demo/app.py @@ -0,0 +1,513 @@ +""" +THis is the main file for the gradio web demo. It uses the CogVideoX-5B model to generate videos gradio web demo. +set environment variable OPENAI_API_KEY to use the OpenAI API to enhance the prompt. + +Usage: + OpenAI_API_KEY=your_openai_api_key OPENAI_BASE_URL=https://api.openai.com/v1 python inference/gradio_web_demo.py +""" + +import math +import os +import random +import threading +import time + +import cv2 +import tempfile +import imageio_ffmpeg +import gradio as gr +import torch +from PIL import Image +from diffusers import ( + CogVideoXPipeline, + CogVideoXDPMScheduler, + CogVideoXVideoToVideoPipeline, + CogVideoXImageToVideoPipeline, + CogVideoXTransformer3DModel, +) +from diffusers.utils import load_video, load_image +from datetime import datetime, timedelta + +from diffusers.image_processor import VaeImageProcessor +from openai import OpenAI +from moviepy import VideoFileClip +import utils +from rife_model import load_rife_model, rife_inference_with_latents +from huggingface_hub import hf_hub_download, snapshot_download + +device = "cuda" if torch.cuda.is_available() else "cpu" + +MODEL = "THUDM/CogVideoX-5b" + +hf_hub_download( + repo_id="ai-forever/Real-ESRGAN", filename="RealESRGAN_x4.pth", local_dir="model_real_esran" +) +snapshot_download(repo_id="AlexWortega/RIFE", local_dir="model_rife") + +pipe = CogVideoXPipeline.from_pretrained(MODEL, torch_dtype=torch.bfloat16).to(device) +pipe.scheduler = CogVideoXDPMScheduler.from_config( + pipe.scheduler.config, timestep_spacing="trailing" +) +pipe_video = CogVideoXVideoToVideoPipeline.from_pretrained( + MODEL, + transformer=pipe.transformer, + vae=pipe.vae, + scheduler=pipe.scheduler, + tokenizer=pipe.tokenizer, + text_encoder=pipe.text_encoder, + torch_dtype=torch.bfloat16, +).to(device) + +pipe_image = CogVideoXImageToVideoPipeline.from_pretrained( + MODEL, + transformer=CogVideoXTransformer3DModel.from_pretrained( + MODEL, subfolder="transformer", torch_dtype=torch.bfloat16 + ), + vae=pipe.vae, + scheduler=pipe.scheduler, + tokenizer=pipe.tokenizer, + text_encoder=pipe.text_encoder, + torch_dtype=torch.bfloat16, +).to(device) + + +# pipe.transformer.to(memory_format=torch.channels_last) +# pipe.transformer = torch.compile(pipe.transformer, mode="max-autotune", fullgraph=True) +# pipe_image.transformer.to(memory_format=torch.channels_last) +# pipe_image.transformer = torch.compile(pipe_image.transformer, mode="max-autotune", fullgraph=True) + +os.makedirs("./output", exist_ok=True) +os.makedirs("./gradio_tmp", exist_ok=True) + +upscale_model = utils.load_sd_upscale("model_real_esran/RealESRGAN_x4.pth", device) +frame_interpolation_model = load_rife_model("model_rife") + +sys_prompt = """You are part of a team of bots that creates videos. You work with an assistant bot that will draw anything you say in square brackets. + +For example , outputting " a beautiful morning in the woods with the sun peaking through the trees " will trigger your partner bot to output an video of a forest morning , as described. You will be prompted by people looking to create detailed , amazing videos. The way to accomplish this is to take their short prompts and make them extremely detailed and descriptive. +There are a few rules to follow: + +You will only ever output a single video description per user request. + +When modifications are requested , you should not simply make the description longer . You should refactor the entire description to integrate the suggestions. +Other times the user will not want modifications , but instead want a new image . In this case , you should ignore your previous conversation with the user. + +Video descriptions must have the same num of words as examples below. Extra words will be ignored. +""" + + +def resize_if_unfit(input_video, progress=gr.Progress(track_tqdm=True)): + width, height = get_video_dimensions(input_video) + + if width == 720 and height == 480: + processed_video = input_video + else: + processed_video = center_crop_resize(input_video) + return processed_video + + +def get_video_dimensions(input_video_path): + reader = imageio_ffmpeg.read_frames(input_video_path) + metadata = next(reader) + return metadata["size"] + + +def center_crop_resize(input_video_path, target_width=720, target_height=480): + cap = cv2.VideoCapture(input_video_path) + + orig_width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH)) + orig_height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT)) + orig_fps = cap.get(cv2.CAP_PROP_FPS) + total_frames = int(cap.get(cv2.CAP_PROP_FRAME_COUNT)) + + width_factor = target_width / orig_width + height_factor = target_height / orig_height + resize_factor = max(width_factor, height_factor) + + inter_width = int(orig_width * resize_factor) + inter_height = int(orig_height * resize_factor) + + target_fps = 8 + ideal_skip = max(0, math.ceil(orig_fps / target_fps) - 1) + skip = min(5, ideal_skip) # Cap at 5 + + while (total_frames / (skip + 1)) < 49 and skip > 0: + skip -= 1 + + processed_frames = [] + frame_count = 0 + total_read = 0 + + while frame_count < 49 and total_read < total_frames: + ret, frame = cap.read() + if not ret: + break + + if total_read % (skip + 1) == 0: + resized = cv2.resize(frame, (inter_width, inter_height), interpolation=cv2.INTER_AREA) + + start_x = (inter_width - target_width) // 2 + start_y = (inter_height - target_height) // 2 + cropped = resized[start_y : start_y + target_height, start_x : start_x + target_width] + + processed_frames.append(cropped) + frame_count += 1 + + total_read += 1 + + cap.release() + + with tempfile.NamedTemporaryFile(suffix=".mp4", delete=False) as temp_file: + temp_video_path = temp_file.name + fourcc = cv2.VideoWriter_fourcc(*"mp4v") + out = cv2.VideoWriter(temp_video_path, fourcc, target_fps, (target_width, target_height)) + + for frame in processed_frames: + out.write(frame) + + out.release() + + return temp_video_path + + +def convert_prompt(prompt: str, retry_times: int = 3) -> str: + if not os.environ.get("OPENAI_API_KEY"): + return prompt + client = OpenAI() + text = prompt.strip() + + for i in range(retry_times): + response = client.chat.completions.create( + messages=[ + {"role": "system", "content": sys_prompt}, + { + "role": "user", + "content": 'Create an imaginative video descriptive caption or modify an earlier caption for the user input : "a girl is on the beach"', + }, + { + "role": "assistant", + "content": "A radiant woman stands on a deserted beach, arms outstretched, wearing a beige trench coat, white blouse, light blue jeans, and chic boots, against a backdrop of soft sky and sea. Moments later, she is seen mid-twirl, arms exuberant, with the lighting suggesting dawn or dusk. Then, she runs along the beach, her attire complemented by an off-white scarf and black ankle boots, the tranquil sea behind her. Finally, she holds a paper airplane, her pose reflecting joy and freedom, with the ocean's gentle waves and the sky's soft pastel hues enhancing the serene ambiance.", + }, + { + "role": "user", + "content": 'Create an imaginative video descriptive caption or modify an earlier caption for the user input : "A man jogging on a football field"', + }, + { + "role": "assistant", + "content": "A determined man in athletic attire, including a blue long-sleeve shirt, black shorts, and blue socks, jogs around a snow-covered soccer field, showcasing his solitary exercise in a quiet, overcast setting. His long dreadlocks, focused expression, and the serene winter backdrop highlight his dedication to fitness. As he moves, his attire, consisting of a blue sports sweatshirt, black athletic pants, gloves, and sneakers, grips the snowy ground. He is seen running past a chain-link fence enclosing the playground area, with a basketball hoop and children's slide, suggesting a moment of solitary exercise amidst the empty field.", + }, + { + "role": "user", + "content": 'Create an imaginative video descriptive caption or modify an earlier caption for the user input : " A woman is dancing, HD footage, close-up"', + }, + { + "role": "assistant", + "content": "A young woman with her hair in an updo and wearing a teal hoodie stands against a light backdrop, initially looking over her shoulder with a contemplative expression. She then confidently makes a subtle dance move, suggesting rhythm and movement. Next, she appears poised and focused, looking directly at the camera. Her expression shifts to one of introspection as she gazes downward slightly. Finally, she dances with confidence, her left hand over her heart, symbolizing a poignant moment, all while dressed in the same teal hoodie against a plain, light-colored background.", + }, + { + "role": "user", + "content": f'Create an imaginative video descriptive caption or modify an earlier caption in ENGLISH for the user input: "{text}"', + }, + ], + model="glm-4-plus", + temperature=0.01, + top_p=0.7, + stream=False, + max_tokens=200, + ) + if response.choices: + return response.choices[0].message.content + return prompt + + +def infer( + prompt: str, + image_input: str, + video_input: str, + video_strenght: float, + num_inference_steps: int, + guidance_scale: float, + seed: int = -1, + progress=gr.Progress(track_tqdm=True), +): + if seed == -1: + seed = random.randint(0, 2**8 - 1) + + if video_input is not None: + video = load_video(video_input)[:49] # Limit to 49 frames + video_pt = pipe_video( + video=video, + prompt=prompt, + num_inference_steps=num_inference_steps, + num_videos_per_prompt=1, + strength=video_strenght, + use_dynamic_cfg=True, + output_type="pt", + guidance_scale=guidance_scale, + generator=torch.Generator(device="cpu").manual_seed(seed), + ).frames + elif image_input is not None: + image_input = Image.fromarray(image_input).resize(size=(720, 480)) # Convert to PIL + image = load_image(image_input) + video_pt = pipe_image( + image=image, + prompt=prompt, + num_inference_steps=num_inference_steps, + num_videos_per_prompt=1, + use_dynamic_cfg=True, + output_type="pt", + guidance_scale=guidance_scale, + generator=torch.Generator(device="cpu").manual_seed(seed), + ).frames + else: + video_pt = pipe( + prompt=prompt, + num_videos_per_prompt=1, + num_inference_steps=num_inference_steps, + num_frames=49, + use_dynamic_cfg=True, + output_type="pt", + guidance_scale=guidance_scale, + generator=torch.Generator(device="cpu").manual_seed(seed), + ).frames + + return (video_pt, seed) + + +def convert_to_gif(video_path): + clip = VideoFileClip(video_path) + clip = clip.with_fps(8) + clip = clip.resized(height=240) + gif_path = video_path.replace(".mp4", ".gif") + clip.write_gif(gif_path, fps=8) + return gif_path + + +def delete_old_files(): + while True: + now = datetime.now() + cutoff = now - timedelta(minutes=10) + directories = ["./output", "./gradio_tmp"] + + for directory in directories: + for filename in os.listdir(directory): + file_path = os.path.join(directory, filename) + if os.path.isfile(file_path): + file_mtime = datetime.fromtimestamp(os.path.getmtime(file_path)) + if file_mtime < cutoff: + os.remove(file_path) + time.sleep(600) + + +threading.Thread(target=delete_old_files, daemon=True).start() +examples_videos = [ + ["example_videos/horse.mp4"], + ["example_videos/kitten.mp4"], + ["example_videos/train_running.mp4"], +] +examples_images = [ + ["example_images/beach.png"], + ["example_images/street.png"], + ["example_images/camping.png"], +] + +with gr.Blocks() as demo: + gr.Markdown(""" +
+ CogVideoX-5B Huggingface Space🤗 +
+
+ 🤗 5B(T2V) Model Hub | + 🤗 5B(I2V) Model Hub | + 🌐 Github | + 📜 arxiv +
+
+ If the Space is too busy, duplicate it to use privately + +
+
+ ⚠️ This demo is for academic research and experimental use only. +
+ """) + with gr.Row(): + with gr.Column(): + with gr.Accordion( + "I2V: Image Input (cannot be used simultaneously with video input)", open=False + ): + image_input = gr.Image(label="Input Image (will be cropped to 720 * 480)") + examples_component_images = gr.Examples( + examples_images, inputs=[image_input], cache_examples=False + ) + with gr.Accordion( + "V2V: Video Input (cannot be used simultaneously with image input)", open=False + ): + video_input = gr.Video( + label="Input Video (will be cropped to 49 frames, 6 seconds at 8fps)" + ) + strength = gr.Slider(0.1, 1.0, value=0.8, step=0.01, label="Strength") + examples_component_videos = gr.Examples( + examples_videos, inputs=[video_input], cache_examples=False + ) + prompt = gr.Textbox( + label="Prompt (Less than 200 Words)", placeholder="Enter your prompt here", lines=5 + ) + + with gr.Row(): + gr.Markdown( + "✨Upon pressing the enhanced prompt button, we will use [GLM-4 Model](https://github.com/THUDM/GLM-4) to polish the prompt and overwrite the original one." + ) + enhance_button = gr.Button("✨ Enhance Prompt(Optional)") + with gr.Group(): + with gr.Column(): + with gr.Row(): + seed_param = gr.Number( + label="Inference Seed (Enter a positive number, -1 for random)", + value=-1, + ) + with gr.Row(): + enable_scale = gr.Checkbox( + label="Super-Resolution (720 × 480 -> 2880 × 1920)", value=False + ) + enable_rife = gr.Checkbox( + label="Frame Interpolation (8fps -> 16fps)", value=False + ) + gr.Markdown( + "✨In this demo, we use [RIFE](https://github.com/hzwer/ECCV2022-RIFE) for frame interpolation and [Real-ESRGAN](https://github.com/xinntao/Real-ESRGAN) for upscaling(Super-Resolution).
    The entire process is based on open-source solutions." + ) + + generate_button = gr.Button("🎬 Generate Video") + + with gr.Column(): + video_output = gr.Video(label="CogVideoX Generate Video", width=720, height=480) + with gr.Row(): + download_video_button = gr.File(label="📥 Download Video", visible=False) + download_gif_button = gr.File(label="📥 Download GIF", visible=False) + seed_text = gr.Number(label="Seed Used for Video Generation", visible=False) + + gr.Markdown(""" + +
+ 🎥 Video Gallery(For 5B) +
+ + + + + + + + + + + + + + + + + + + + + + + + +
+

A garden comes to life as a kaleidoscope of butterflies flutters amidst the blossoms, their delicate wings casting shadows on the petals below. In the background, a grand fountain cascades water with a gentle splendor, its rhythmic sound providing a soothing backdrop. Beneath the cool shade of a mature tree, a solitary wooden chair invites solitude and reflection, its smooth surface worn by the touch of countless visitors seeking a moment of tranquility in nature's embrace.

+ 		+ + +

A small boy, head bowed and determination etched on his face, sprints through the torrential downpour as lightning crackles and thunder rumbles in the distance. The relentless rain pounds the ground, creating a chaotic dance of water droplets that mirror the dramatic sky's anger. In the far background, the silhouette of a cozy home beckons, a faint beacon of safety and warmth amidst the fierce weather. The scene is one of perseverance and the unyielding spirit of a child braving the elements.

+ 		+ +
+

A suited astronaut, with the red dust of Mars clinging to their boots, reaches out to shake hands with an alien being, their skin a shimmering blue, under the pink-tinged sky of the fourth planet. In the background, a sleek silver rocket, a beacon of human ingenuity, stands tall, its engines powered down, as the two representatives of different worlds exchange a historic greeting amidst the desolate beauty of the Martian landscape.

+ 		+ + +

An elderly gentleman, with a serene expression, sits at the water's edge, a steaming cup of tea by his side. He is engrossed in his artwork, brush in hand, as he renders an oil painting on a canvas that's propped up against a small, weathered table. The sea breeze whispers through his silver hair, gently billowing his loose-fitting white shirt, while the salty air adds an intangible element to his masterpiece in progress. The scene is one of tranquility and inspiration, with the artist's canvas capturing the vibrant hues of the setting sun reflecting off the tranquil sea.

+ 		+ +
+

In a dimly lit bar, purplish light bathes the face of a mature man, his eyes blinking thoughtfully as he ponders in close-up, the background artfully blurred to focus on his introspective expression, the ambiance of the bar a mere suggestion of shadows and soft lighting.

+ 		+ + +

A golden retriever, sporting sleek black sunglasses, with its lengthy fur flowing in the breeze, sprints playfully across a rooftop terrace, recently refreshed by a light rain. The scene unfolds from a distance, the dog's energetic bounds growing larger as it approaches the camera, its tail wagging with unrestrained joy, while droplets of water glisten on the concrete behind it. The overcast sky provides a dramatic backdrop, emphasizing the vibrant golden coat of the canine as it dashes towards the viewer.

+ 		+ +
+

On a brilliant sunny day, the lakeshore is lined with an array of willow trees, their slender branches swaying gently in the soft breeze. The tranquil surface of the lake reflects the clear blue sky, while several elegant swans glide gracefully through the still water, leaving behind delicate ripples that disturb the mirror-like quality of the lake. The scene is one of serene beauty, with the willows' greenery providing a picturesque frame for the peaceful avian visitors.

+ 		+ + +

A Chinese mother, draped in a soft, pastel-colored robe, gently rocks back and forth in a cozy rocking chair positioned in the tranquil setting of a nursery. The dimly lit bedroom is adorned with whimsical mobiles dangling from the ceiling, casting shadows that dance on the walls. Her baby, swaddled in a delicate, patterned blanket, rests against her chest, the child's earlier cries now replaced by contented coos as the mother's soothing voice lulls the little one to sleep. The scent of lavender fills the air, adding to the serene atmosphere, while a warm, orange glow from a nearby nightlight illuminates the scene with a gentle hue, capturing a moment of tender love and comfort.

+ 		+ +
+ """) + + def generate( + prompt, + image_input, + video_input, + video_strength, + seed_value, + scale_status, + rife_status, + progress=gr.Progress(track_tqdm=True), + ): + latents, seed = infer( + prompt, + image_input, + video_input, + video_strength, + num_inference_steps=50, # NOT Changed + guidance_scale=7.0, # NOT Changed + seed=seed_value, + progress=progress, + ) + if scale_status: + latents = utils.upscale_batch_and_concatenate(upscale_model, latents, device) + if rife_status: + latents = rife_inference_with_latents(frame_interpolation_model, latents) + + batch_size = latents.shape[0] + batch_video_frames = [] + for batch_idx in range(batch_size): + pt_image = latents[batch_idx] + pt_image = torch.stack([pt_image[i] for i in range(pt_image.shape[0])]) + + image_np = VaeImageProcessor.pt_to_numpy(pt_image) + image_pil = VaeImageProcessor.numpy_to_pil(image_np) + batch_video_frames.append(image_pil) + + video_path = utils.save_video( + batch_video_frames[0], fps=math.ceil((len(batch_video_frames[0]) - 1) / 6) + ) + video_update = gr.update(visible=True, value=video_path) + gif_path = convert_to_gif(video_path) + gif_update = gr.update(visible=True, value=gif_path) + seed_update = gr.update(visible=True, value=seed) + + return video_path, video_update, gif_update, seed_update + + def enhance_prompt_func(prompt): + return convert_prompt(prompt, retry_times=1) + + generate_button.click( + generate, + inputs=[prompt, image_input, video_input, strength, seed_param, enable_scale, enable_rife], + outputs=[video_output, download_video_button, download_gif_button, seed_text], + ) + + enhance_button.click(enhance_prompt_func, inputs=[prompt], outputs=[prompt]) + video_input.upload(resize_if_unfit, inputs=[video_input], outputs=[video_input]) + +if __name__ == "__main__": + demo.queue(max_size=15) + demo.launch() diff --git a/inference/gradio_composite_demo/example_images/beach.png b/inference/gradio_composite_demo/example_images/beach.png new file mode 100644 index 0000000..8ea6b60 Binary files /dev/null and b/inference/gradio_composite_demo/example_images/beach.png differ diff --git a/inference/gradio_composite_demo/example_images/camping.png b/inference/gradio_composite_demo/example_images/camping.png new file mode 100644 index 0000000..a0b809c Binary files /dev/null and b/inference/gradio_composite_demo/example_images/camping.png differ diff --git a/inference/gradio_composite_demo/example_images/street.png b/inference/gradio_composite_demo/example_images/street.png new file mode 100644 index 0000000..0b95e34 Binary files /dev/null and b/inference/gradio_composite_demo/example_images/street.png differ diff --git a/inference/gradio_composite_demo/requirements.txt b/inference/gradio_composite_demo/requirements.txt new file mode 100644 index 0000000..50ce31c --- /dev/null +++ b/inference/gradio_composite_demo/requirements.txt @@ -0,0 +1,19 @@ +spaces>=0.29.3 +safetensors>=0.4.5 +spandrel>=0.4.0 +tqdm>=4.66.5 +scikit-video>=1.1.11 +diffusers>=0.31.0 +transformers>=4.44.0 +accelerate>=0.34.2 +opencv-python>=4.10.0.84 +sentencepiece>=0.2.0 +numpy==1.26.0 +torch>=2.5.0 +torchvision>=0.20.0 +gradio>=5.4.0 +imageio>=2.34.2 +imageio-ffmpeg>=0.5.1 +openai>=1.45.0 +moviepy>=2.0.0 +pillow==9.5.0 diff --git a/inference/gradio_composite_demo/rife/IFNet.py b/inference/gradio_composite_demo/rife/IFNet.py new file mode 100644 index 0000000..c395783 --- /dev/null +++ b/inference/gradio_composite_demo/rife/IFNet.py @@ -0,0 +1,136 @@ +from .refine import * + + +def deconv(in_planes, out_planes, kernel_size=4, stride=2, padding=1): + return nn.Sequential( + torch.nn.ConvTranspose2d( + in_channels=in_planes, out_channels=out_planes, kernel_size=4, stride=2, padding=1 + ), + nn.PReLU(out_planes), + ) + + +def conv(in_planes, out_planes, kernel_size=3, stride=1, padding=1, dilation=1): + return nn.Sequential( + nn.Conv2d( + in_planes, + out_planes, + kernel_size=kernel_size, + stride=stride, + padding=padding, + dilation=dilation, + bias=True, + ), + nn.PReLU(out_planes), + ) + + +class IFBlock(nn.Module): + def __init__(self, in_planes, c=64): + super(IFBlock, self).__init__() + self.conv0 = nn.Sequential( + conv(in_planes, c // 2, 3, 2, 1), + conv(c // 2, c, 3, 2, 1), + ) + self.convblock = nn.Sequential( + conv(c, c), + conv(c, c), + conv(c, c), + conv(c, c), + conv(c, c), + conv(c, c), + conv(c, c), + conv(c, c), + ) + self.lastconv = nn.ConvTranspose2d(c, 5, 4, 2, 1) + + def forward(self, x, flow, scale): + if scale != 1: + x = F.interpolate(x, scale_factor=1.0 / scale, mode="bilinear", align_corners=False) + if flow != None: + flow = ( + F.interpolate(flow, scale_factor=1.0 / scale, mode="bilinear", align_corners=False) + * 1.0 + / scale + ) + x = torch.cat((x, flow), 1) + x = self.conv0(x) + x = self.convblock(x) + x + tmp = self.lastconv(x) + tmp = F.interpolate(tmp, scale_factor=scale * 2, mode="bilinear", align_corners=False) + flow = tmp[:, :4] * scale * 2 + mask = tmp[:, 4:5] + return flow, mask + + +class IFNet(nn.Module): + def __init__(self): + super(IFNet, self).__init__() + self.block0 = IFBlock(6, c=240) + self.block1 = IFBlock(13 + 4, c=150) + self.block2 = IFBlock(13 + 4, c=90) + self.block_tea = IFBlock(16 + 4, c=90) + self.contextnet = Contextnet() + self.unet = Unet() + + def forward(self, x, scale=[4, 2, 1], timestep=0.5): + img0 = x[:, :3] + img1 = x[:, 3:6] + gt = x[:, 6:] # In inference time, gt is None + flow_list = [] + merged = [] + mask_list = [] + warped_img0 = img0 + warped_img1 = img1 + flow = None + loss_distill = 0 + stu = [self.block0, self.block1, self.block2] + for i in range(3): + if flow != None: + flow_d, mask_d = stu[i]( + torch.cat((img0, img1, warped_img0, warped_img1, mask), 1), flow, scale=scale[i] + ) + flow = flow + flow_d + mask = mask + mask_d + else: + flow, mask = stu[i](torch.cat((img0, img1), 1), None, scale=scale[i]) + mask_list.append(torch.sigmoid(mask)) + flow_list.append(flow) + warped_img0 = warp(img0, flow[:, :2]) + warped_img1 = warp(img1, flow[:, 2:4]) + merged_student = (warped_img0, warped_img1) + merged.append(merged_student) + if gt.shape[1] == 3: + flow_d, mask_d = self.block_tea( + torch.cat((img0, img1, warped_img0, warped_img1, mask, gt), 1), flow, scale=1 + ) + flow_teacher = flow + flow_d + warped_img0_teacher = warp(img0, flow_teacher[:, :2]) + warped_img1_teacher = warp(img1, flow_teacher[:, 2:4]) + mask_teacher = torch.sigmoid(mask + mask_d) + merged_teacher = warped_img0_teacher * mask_teacher + warped_img1_teacher * ( + 1 - mask_teacher + ) + else: + flow_teacher = None + merged_teacher = None + for i in range(3): + merged[i] = merged[i][0] * mask_list[i] + merged[i][1] * (1 - mask_list[i]) + if gt.shape[1] == 3: + loss_mask = ( + ( + (merged[i] - gt).abs().mean(1, True) + > (merged_teacher - gt).abs().mean(1, True) + 0.01 + ) + .float() + .detach() + ) + loss_distill += ( + ((flow_teacher.detach() - flow_list[i]) ** 2).mean(1, True) ** 0.5 * loss_mask + ).mean() + c0 = self.contextnet(img0, flow[:, :2]) + c1 = self.contextnet(img1, flow[:, 2:4]) + tmp = self.unet(img0, img1, warped_img0, warped_img1, mask, flow, c0, c1) + res = tmp[:, :3] * 2 - 1 + merged[2] = torch.clamp(merged[2] + res, 0, 1) + return flow_list, mask_list[2], merged, flow_teacher, merged_teacher, loss_distill diff --git a/inference/gradio_composite_demo/rife/IFNet_2R.py b/inference/gradio_composite_demo/rife/IFNet_2R.py new file mode 100644 index 0000000..8ec9eaa --- /dev/null +++ b/inference/gradio_composite_demo/rife/IFNet_2R.py @@ -0,0 +1,136 @@ +from .refine_2R import * + + +def deconv(in_planes, out_planes, kernel_size=4, stride=2, padding=1): + return nn.Sequential( + torch.nn.ConvTranspose2d( + in_channels=in_planes, out_channels=out_planes, kernel_size=4, stride=2, padding=1 + ), + nn.PReLU(out_planes), + ) + + +def conv(in_planes, out_planes, kernel_size=3, stride=1, padding=1, dilation=1): + return nn.Sequential( + nn.Conv2d( + in_planes, + out_planes, + kernel_size=kernel_size, + stride=stride, + padding=padding, + dilation=dilation, + bias=True, + ), + nn.PReLU(out_planes), + ) + + +class IFBlock(nn.Module): + def __init__(self, in_planes, c=64): + super(IFBlock, self).__init__() + self.conv0 = nn.Sequential( + conv(in_planes, c // 2, 3, 1, 1), + conv(c // 2, c, 3, 2, 1), + ) + self.convblock = nn.Sequential( + conv(c, c), + conv(c, c), + conv(c, c), + conv(c, c), + conv(c, c), + conv(c, c), + conv(c, c), + conv(c, c), + ) + self.lastconv = nn.ConvTranspose2d(c, 5, 4, 2, 1) + + def forward(self, x, flow, scale): + if scale != 1: + x = F.interpolate(x, scale_factor=1.0 / scale, mode="bilinear", align_corners=False) + if flow != None: + flow = ( + F.interpolate(flow, scale_factor=1.0 / scale, mode="bilinear", align_corners=False) + * 1.0 + / scale + ) + x = torch.cat((x, flow), 1) + x = self.conv0(x) + x = self.convblock(x) + x + tmp = self.lastconv(x) + tmp = F.interpolate(tmp, scale_factor=scale, mode="bilinear", align_corners=False) + flow = tmp[:, :4] * scale + mask = tmp[:, 4:5] + return flow, mask + + +class IFNet(nn.Module): + def __init__(self): + super(IFNet, self).__init__() + self.block0 = IFBlock(6, c=240) + self.block1 = IFBlock(13 + 4, c=150) + self.block2 = IFBlock(13 + 4, c=90) + self.block_tea = IFBlock(16 + 4, c=90) + self.contextnet = Contextnet() + self.unet = Unet() + + def forward(self, x, scale=[4, 2, 1], timestep=0.5): + img0 = x[:, :3] + img1 = x[:, 3:6] + gt = x[:, 6:] # In inference time, gt is None + flow_list = [] + merged = [] + mask_list = [] + warped_img0 = img0 + warped_img1 = img1 + flow = None + loss_distill = 0 + stu = [self.block0, self.block1, self.block2] + for i in range(3): + if flow != None: + flow_d, mask_d = stu[i]( + torch.cat((img0, img1, warped_img0, warped_img1, mask), 1), flow, scale=scale[i] + ) + flow = flow + flow_d + mask = mask + mask_d + else: + flow, mask = stu[i](torch.cat((img0, img1), 1), None, scale=scale[i]) + mask_list.append(torch.sigmoid(mask)) + flow_list.append(flow) + warped_img0 = warp(img0, flow[:, :2]) + warped_img1 = warp(img1, flow[:, 2:4]) + merged_student = (warped_img0, warped_img1) + merged.append(merged_student) + if gt.shape[1] == 3: + flow_d, mask_d = self.block_tea( + torch.cat((img0, img1, warped_img0, warped_img1, mask, gt), 1), flow, scale=1 + ) + flow_teacher = flow + flow_d + warped_img0_teacher = warp(img0, flow_teacher[:, :2]) + warped_img1_teacher = warp(img1, flow_teacher[:, 2:4]) + mask_teacher = torch.sigmoid(mask + mask_d) + merged_teacher = warped_img0_teacher * mask_teacher + warped_img1_teacher * ( + 1 - mask_teacher + ) + else: + flow_teacher = None + merged_teacher = None + for i in range(3): + merged[i] = merged[i][0] * mask_list[i] + merged[i][1] * (1 - mask_list[i]) + if gt.shape[1] == 3: + loss_mask = ( + ( + (merged[i] - gt).abs().mean(1, True) + > (merged_teacher - gt).abs().mean(1, True) + 0.01 + ) + .float() + .detach() + ) + loss_distill += ( + ((flow_teacher.detach() - flow_list[i]) ** 2).mean(1, True) ** 0.5 * loss_mask + ).mean() + c0 = self.contextnet(img0, flow[:, :2]) + c1 = self.contextnet(img1, flow[:, 2:4]) + tmp = self.unet(img0, img1, warped_img0, warped_img1, mask, flow, c0, c1) + res = tmp[:, :3] * 2 - 1 + merged[2] = torch.clamp(merged[2] + res, 0, 1) + return flow_list, mask_list[2], merged, flow_teacher, merged_teacher, loss_distill diff --git a/inference/gradio_composite_demo/rife/IFNet_HDv3.py b/inference/gradio_composite_demo/rife/IFNet_HDv3.py new file mode 100644 index 0000000..ad4a727 --- /dev/null +++ b/inference/gradio_composite_demo/rife/IFNet_HDv3.py @@ -0,0 +1,160 @@ +import torch +import torch.nn as nn +import torch.nn.functional as F +from .warplayer import warp + +device = torch.device("cuda" if torch.cuda.is_available() else "cpu") + + +def conv(in_planes, out_planes, kernel_size=3, stride=1, padding=1, dilation=1): + return nn.Sequential( + nn.Conv2d( + in_planes, + out_planes, + kernel_size=kernel_size, + stride=stride, + padding=padding, + dilation=dilation, + bias=True, + ), + nn.PReLU(out_planes), + ) + + +def conv_bn(in_planes, out_planes, kernel_size=3, stride=1, padding=1, dilation=1): + return nn.Sequential( + nn.Conv2d( + in_planes, + out_planes, + kernel_size=kernel_size, + stride=stride, + padding=padding, + dilation=dilation, + bias=False, + ), + nn.BatchNorm2d(out_planes), + nn.PReLU(out_planes), + ) + + +class IFBlock(nn.Module): + def __init__(self, in_planes, c=64): + super(IFBlock, self).__init__() + self.conv0 = nn.Sequential( + conv(in_planes, c // 2, 3, 2, 1), + conv(c // 2, c, 3, 2, 1), + ) + self.convblock0 = nn.Sequential(conv(c, c), conv(c, c)) + self.convblock1 = nn.Sequential(conv(c, c), conv(c, c)) + self.convblock2 = nn.Sequential(conv(c, c), conv(c, c)) + self.convblock3 = nn.Sequential(conv(c, c), conv(c, c)) + self.conv1 = nn.Sequential( + nn.ConvTranspose2d(c, c // 2, 4, 2, 1), + nn.PReLU(c // 2), + nn.ConvTranspose2d(c // 2, 4, 4, 2, 1), + ) + self.conv2 = nn.Sequential( + nn.ConvTranspose2d(c, c // 2, 4, 2, 1), + nn.PReLU(c // 2), + nn.ConvTranspose2d(c // 2, 1, 4, 2, 1), + ) + + def forward(self, x, flow, scale=1): + x = F.interpolate( + x, + scale_factor=1.0 / scale, + mode="bilinear", + align_corners=False, + recompute_scale_factor=False, + ) + flow = ( + F.interpolate( + flow, + scale_factor=1.0 / scale, + mode="bilinear", + align_corners=False, + recompute_scale_factor=False, + ) + * 1.0 + / scale + ) + feat = self.conv0(torch.cat((x, flow), 1)) + feat = self.convblock0(feat) + feat + feat = self.convblock1(feat) + feat + feat = self.convblock2(feat) + feat + feat = self.convblock3(feat) + feat + flow = self.conv1(feat) + mask = self.conv2(feat) + flow = ( + F.interpolate( + flow, + scale_factor=scale, + mode="bilinear", + align_corners=False, + recompute_scale_factor=False, + ) + * scale + ) + mask = F.interpolate( + mask, + scale_factor=scale, + mode="bilinear", + align_corners=False, + recompute_scale_factor=False, + ) + return flow, mask + + +class IFNet(nn.Module): + def __init__(self): + super(IFNet, self).__init__() + self.block0 = IFBlock(7 + 4, c=90) + self.block1 = IFBlock(7 + 4, c=90) + self.block2 = IFBlock(7 + 4, c=90) + self.block_tea = IFBlock(10 + 4, c=90) + # self.contextnet = Contextnet() + # self.unet = Unet() + + def forward(self, x, scale_list=[4, 2, 1], training=False): + if training == False: + channel = x.shape[1] // 2 + img0 = x[:, :channel] + img1 = x[:, channel:] + flow_list = [] + merged = [] + mask_list = [] + warped_img0 = img0 + warped_img1 = img1 + flow = (x[:, :4]).detach() * 0 + mask = (x[:, :1]).detach() * 0 + loss_cons = 0 + block = [self.block0, self.block1, self.block2] + for i in range(3): + f0, m0 = block[i]( + torch.cat((warped_img0[:, :3], warped_img1[:, :3], mask), 1), + flow, + scale=scale_list[i], + ) + f1, m1 = block[i]( + torch.cat((warped_img1[:, :3], warped_img0[:, :3], -mask), 1), + torch.cat((flow[:, 2:4], flow[:, :2]), 1), + scale=scale_list[i], + ) + flow = flow + (f0 + torch.cat((f1[:, 2:4], f1[:, :2]), 1)) / 2 + mask = mask + (m0 + (-m1)) / 2 + mask_list.append(mask) + flow_list.append(flow) + warped_img0 = warp(img0, flow[:, :2]) + warped_img1 = warp(img1, flow[:, 2:4]) + merged.append((warped_img0, warped_img1)) + """ + c0 = self.contextnet(img0, flow[:, :2]) + c1 = self.contextnet(img1, flow[:, 2:4]) + tmp = self.unet(img0, img1, warped_img0, warped_img1, mask, flow, c0, c1) + res = tmp[:, 1:4] * 2 - 1 + """ + for i in range(3): + mask_list[i] = torch.sigmoid(mask_list[i]) + merged[i] = merged[i][0] * mask_list[i] + merged[i][1] * (1 - mask_list[i]) + # merged[i] = torch.clamp(merged[i] + res, 0, 1) + return flow_list, mask_list[2], merged diff --git a/inference/gradio_composite_demo/rife/IFNet_m.py b/inference/gradio_composite_demo/rife/IFNet_m.py new file mode 100644 index 0000000..8547ef8 --- /dev/null +++ b/inference/gradio_composite_demo/rife/IFNet_m.py @@ -0,0 +1,144 @@ +from .refine import * + + +def deconv(in_planes, out_planes, kernel_size=4, stride=2, padding=1): + return nn.Sequential( + torch.nn.ConvTranspose2d( + in_channels=in_planes, out_channels=out_planes, kernel_size=4, stride=2, padding=1 + ), + nn.PReLU(out_planes), + ) + + +def conv(in_planes, out_planes, kernel_size=3, stride=1, padding=1, dilation=1): + return nn.Sequential( + nn.Conv2d( + in_planes, + out_planes, + kernel_size=kernel_size, + stride=stride, + padding=padding, + dilation=dilation, + bias=True, + ), + nn.PReLU(out_planes), + ) + + +class IFBlock(nn.Module): + def __init__(self, in_planes, c=64): + super(IFBlock, self).__init__() + self.conv0 = nn.Sequential( + conv(in_planes, c // 2, 3, 2, 1), + conv(c // 2, c, 3, 2, 1), + ) + self.convblock = nn.Sequential( + conv(c, c), + conv(c, c), + conv(c, c), + conv(c, c), + conv(c, c), + conv(c, c), + conv(c, c), + conv(c, c), + ) + self.lastconv = nn.ConvTranspose2d(c, 5, 4, 2, 1) + + def forward(self, x, flow, scale): + if scale != 1: + x = F.interpolate(x, scale_factor=1.0 / scale, mode="bilinear", align_corners=False) + if flow != None: + flow = ( + F.interpolate(flow, scale_factor=1.0 / scale, mode="bilinear", align_corners=False) + * 1.0 + / scale + ) + x = torch.cat((x, flow), 1) + x = self.conv0(x) + x = self.convblock(x) + x + tmp = self.lastconv(x) + tmp = F.interpolate(tmp, scale_factor=scale * 2, mode="bilinear", align_corners=False) + flow = tmp[:, :4] * scale * 2 + mask = tmp[:, 4:5] + return flow, mask + + +class IFNet_m(nn.Module): + def __init__(self): + super(IFNet_m, self).__init__() + self.block0 = IFBlock(6 + 1, c=240) + self.block1 = IFBlock(13 + 4 + 1, c=150) + self.block2 = IFBlock(13 + 4 + 1, c=90) + self.block_tea = IFBlock(16 + 4 + 1, c=90) + self.contextnet = Contextnet() + self.unet = Unet() + + def forward(self, x, scale=[4, 2, 1], timestep=0.5, returnflow=False): + timestep = (x[:, :1].clone() * 0 + 1) * timestep + img0 = x[:, :3] + img1 = x[:, 3:6] + gt = x[:, 6:] # In inference time, gt is None + flow_list = [] + merged = [] + mask_list = [] + warped_img0 = img0 + warped_img1 = img1 + flow = None + loss_distill = 0 + stu = [self.block0, self.block1, self.block2] + for i in range(3): + if flow != None: + flow_d, mask_d = stu[i]( + torch.cat((img0, img1, timestep, warped_img0, warped_img1, mask), 1), + flow, + scale=scale[i], + ) + flow = flow + flow_d + mask = mask + mask_d + else: + flow, mask = stu[i](torch.cat((img0, img1, timestep), 1), None, scale=scale[i]) + mask_list.append(torch.sigmoid(mask)) + flow_list.append(flow) + warped_img0 = warp(img0, flow[:, :2]) + warped_img1 = warp(img1, flow[:, 2:4]) + merged_student = (warped_img0, warped_img1) + merged.append(merged_student) + if gt.shape[1] == 3: + flow_d, mask_d = self.block_tea( + torch.cat((img0, img1, timestep, warped_img0, warped_img1, mask, gt), 1), + flow, + scale=1, + ) + flow_teacher = flow + flow_d + warped_img0_teacher = warp(img0, flow_teacher[:, :2]) + warped_img1_teacher = warp(img1, flow_teacher[:, 2:4]) + mask_teacher = torch.sigmoid(mask + mask_d) + merged_teacher = warped_img0_teacher * mask_teacher + warped_img1_teacher * ( + 1 - mask_teacher + ) + else: + flow_teacher = None + merged_teacher = None + for i in range(3): + merged[i] = merged[i][0] * mask_list[i] + merged[i][1] * (1 - mask_list[i]) + if gt.shape[1] == 3: + loss_mask = ( + ( + (merged[i] - gt).abs().mean(1, True) + > (merged_teacher - gt).abs().mean(1, True) + 0.01 + ) + .float() + .detach() + ) + loss_distill += ( + ((flow_teacher.detach() - flow_list[i]) ** 2).mean(1, True) ** 0.5 * loss_mask + ).mean() + if returnflow: + return flow + else: + c0 = self.contextnet(img0, flow[:, :2]) + c1 = self.contextnet(img1, flow[:, 2:4]) + tmp = self.unet(img0, img1, warped_img0, warped_img1, mask, flow, c0, c1) + res = tmp[:, :3] * 2 - 1 + merged[2] = torch.clamp(merged[2] + res, 0, 1) + return flow_list, mask_list[2], merged, flow_teacher, merged_teacher, loss_distill diff --git a/inference/gradio_composite_demo/rife/RIFE.py b/inference/gradio_composite_demo/rife/RIFE.py new file mode 100644 index 0000000..a7039ef --- /dev/null +++ b/inference/gradio_composite_demo/rife/RIFE.py @@ -0,0 +1,95 @@ +from torch.optim import AdamW +from torch.nn.parallel import DistributedDataParallel as DDP +from .IFNet import * +from .IFNet_m import * +from .loss import * +from .laplacian import * +from .refine import * + +device = torch.device("cuda" if torch.cuda.is_available() else "cpu") + + +class Model: + def __init__(self, local_rank=-1, arbitrary=False): + if arbitrary == True: + self.flownet = IFNet_m() + else: + self.flownet = IFNet() + self.device() + self.optimG = AdamW( + self.flownet.parameters(), lr=1e-6, weight_decay=1e-3 + ) # use large weight decay may avoid NaN loss + self.epe = EPE() + self.lap = LapLoss() + self.sobel = SOBEL() + if local_rank != -1: + self.flownet = DDP(self.flownet, device_ids=[local_rank], output_device=local_rank) + + def train(self): + self.flownet.train() + + def eval(self): + self.flownet.eval() + + def device(self): + self.flownet.to(device) + + def load_model(self, path, rank=0): + def convert(param): + return {k.replace("module.", ""): v for k, v in param.items() if "module." in k} + + if rank <= 0: + self.flownet.load_state_dict(convert(torch.load("{}/flownet.pkl".format(path)))) + + def save_model(self, path, rank=0): + if rank == 0: + torch.save(self.flownet.state_dict(), "{}/flownet.pkl".format(path)) + + def inference(self, img0, img1, scale=1, scale_list=[4, 2, 1], TTA=False, timestep=0.5): + for i in range(3): + scale_list[i] = scale_list[i] * 1.0 / scale + imgs = torch.cat((img0, img1), 1) + flow, mask, merged, flow_teacher, merged_teacher, loss_distill = self.flownet( + imgs, scale_list, timestep=timestep + ) + if TTA == False: + return merged[2] + else: + flow2, mask2, merged2, flow_teacher2, merged_teacher2, loss_distill2 = self.flownet( + imgs.flip(2).flip(3), scale_list, timestep=timestep + ) + return (merged[2] + merged2[2].flip(2).flip(3)) / 2 + + def update(self, imgs, gt, learning_rate=0, mul=1, training=True, flow_gt=None): + for param_group in self.optimG.param_groups: + param_group["lr"] = learning_rate + img0 = imgs[:, :3] + img1 = imgs[:, 3:] + if training: + self.train() + else: + self.eval() + flow, mask, merged, flow_teacher, merged_teacher, loss_distill = self.flownet( + torch.cat((imgs, gt), 1), scale=[4, 2, 1] + ) + loss_l1 = (self.lap(merged[2], gt)).mean() + loss_tea = (self.lap(merged_teacher, gt)).mean() + if training: + self.optimG.zero_grad() + loss_G = ( + loss_l1 + loss_tea + loss_distill * 0.01 + ) # when training RIFEm, the weight of loss_distill should be 0.005 or 0.002 + loss_G.backward() + self.optimG.step() + else: + flow_teacher = flow[2] + return merged[2], { + "merged_tea": merged_teacher, + "mask": mask, + "mask_tea": mask, + "flow": flow[2][:, :2], + "flow_tea": flow_teacher, + "loss_l1": loss_l1, + "loss_tea": loss_tea, + "loss_distill": loss_distill, + } diff --git a/inference/gradio_composite_demo/rife/RIFE_HDv3.py b/inference/gradio_composite_demo/rife/RIFE_HDv3.py new file mode 100644 index 0000000..182c78e --- /dev/null +++ b/inference/gradio_composite_demo/rife/RIFE_HDv3.py @@ -0,0 +1,88 @@ +import torch +import torch.nn as nn +import numpy as np +from torch.optim import AdamW +import torch.optim as optim +import itertools +from .warplayer import warp +from torch.nn.parallel import DistributedDataParallel as DDP +from .IFNet_HDv3 import * +import torch.nn.functional as F +from .loss import * + +device = torch.device("cuda" if torch.cuda.is_available() else "cpu") + + +class Model: + def __init__(self, local_rank=-1): + self.flownet = IFNet() + self.device() + self.optimG = AdamW(self.flownet.parameters(), lr=1e-6, weight_decay=1e-4) + self.epe = EPE() + # self.vgg = VGGPerceptualLoss().to(device) + self.sobel = SOBEL() + if local_rank != -1: + self.flownet = DDP(self.flownet, device_ids=[local_rank], output_device=local_rank) + + def train(self): + self.flownet.train() + + def eval(self): + self.flownet.eval() + + def device(self): + self.flownet.to(device) + + def load_model(self, path, rank=0): + def convert(param): + if rank == -1: + return {k.replace("module.", ""): v for k, v in param.items() if "module." in k} + else: + return param + + if rank <= 0: + if torch.cuda.is_available(): + self.flownet.load_state_dict(convert(torch.load("{}/flownet.pkl".format(path)))) + else: + self.flownet.load_state_dict( + convert(torch.load("{}/flownet.pkl".format(path), map_location="cpu")) + ) + + def save_model(self, path, rank=0): + if rank == 0: + torch.save(self.flownet.state_dict(), "{}/flownet.pkl".format(path)) + + def inference(self, img0, img1, scale=1.0): + imgs = torch.cat((img0, img1), 1) + scale_list = [4 / scale, 2 / scale, 1 / scale] + flow, mask, merged = self.flownet(imgs, scale_list) + return merged[2] + + def update(self, imgs, gt, learning_rate=0, mul=1, training=True, flow_gt=None): + for param_group in self.optimG.param_groups: + param_group["lr"] = learning_rate + img0 = imgs[:, :3] + img1 = imgs[:, 3:] + if training: + self.train() + else: + self.eval() + scale = [4, 2, 1] + flow, mask, merged = self.flownet(torch.cat((imgs, gt), 1), scale=scale, training=training) + loss_l1 = (merged[2] - gt).abs().mean() + loss_smooth = self.sobel(flow[2], flow[2] * 0).mean() + # loss_vgg = self.vgg(merged[2], gt) + if training: + self.optimG.zero_grad() + loss_G = loss_cons + loss_smooth * 0.1 + loss_G.backward() + self.optimG.step() + else: + flow_teacher = flow[2] + return merged[2], { + "mask": mask, + "flow": flow[2][:, :2], + "loss_l1": loss_l1, + "loss_cons": loss_cons, + "loss_smooth": loss_smooth, + } diff --git a/inference/gradio_composite_demo/rife/__init__.py b/inference/gradio_composite_demo/rife/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/inference/gradio_composite_demo/rife/laplacian.py b/inference/gradio_composite_demo/rife/laplacian.py new file mode 100644 index 0000000..dac3994 --- /dev/null +++ b/inference/gradio_composite_demo/rife/laplacian.py @@ -0,0 +1,77 @@ +import torch +import numpy as np +import torch.nn as nn +import torch.nn.functional as F + +device = torch.device("cuda" if torch.cuda.is_available() else "cpu") + +import torch + + +def gauss_kernel(size=5, channels=3): + kernel = torch.tensor( + [ + [1.0, 4.0, 6.0, 4.0, 1], + [4.0, 16.0, 24.0, 16.0, 4.0], + [6.0, 24.0, 36.0, 24.0, 6.0], + [4.0, 16.0, 24.0, 16.0, 4.0], + [1.0, 4.0, 6.0, 4.0, 1.0], + ] + ) + kernel /= 256.0 + kernel = kernel.repeat(channels, 1, 1, 1) + kernel = kernel.to(device) + return kernel + + +def downsample(x): + return x[:, :, ::2, ::2] + + +def upsample(x): + cc = torch.cat( + [x, torch.zeros(x.shape[0], x.shape[1], x.shape[2], x.shape[3]).to(device)], dim=3 + ) + cc = cc.view(x.shape[0], x.shape[1], x.shape[2] * 2, x.shape[3]) + cc = cc.permute(0, 1, 3, 2) + cc = torch.cat( + [cc, torch.zeros(x.shape[0], x.shape[1], x.shape[3], x.shape[2] * 2).to(device)], dim=3 + ) + cc = cc.view(x.shape[0], x.shape[1], x.shape[3] * 2, x.shape[2] * 2) + x_up = cc.permute(0, 1, 3, 2) + return conv_gauss(x_up, 4 * gauss_kernel(channels=x.shape[1])) + + +def conv_gauss(img, kernel): + img = torch.nn.functional.pad(img, (2, 2, 2, 2), mode="reflect") + out = torch.nn.functional.conv2d(img, kernel, groups=img.shape[1]) + return out + + +def laplacian_pyramid(img, kernel, max_levels=3): + current = img + pyr = [] + for level in range(max_levels): + filtered = conv_gauss(current, kernel) + down = downsample(filtered) + up = upsample(down) + diff = current - up + pyr.append(diff) + current = down + return pyr + + +class LapLoss(torch.nn.Module): + def __init__(self, max_levels=5, channels=3): + super(LapLoss, self).__init__() + self.max_levels = max_levels + self.gauss_kernel = gauss_kernel(channels=channels) + + def forward(self, input, target): + pyr_input = laplacian_pyramid( + img=input, kernel=self.gauss_kernel, max_levels=self.max_levels + ) + pyr_target = laplacian_pyramid( + img=target, kernel=self.gauss_kernel, max_levels=self.max_levels + ) + return sum(torch.nn.functional.l1_loss(a, b) for a, b in zip(pyr_input, pyr_target)) diff --git a/inference/gradio_composite_demo/rife/loss.py b/inference/gradio_composite_demo/rife/loss.py new file mode 100644 index 0000000..8ed7564 --- /dev/null +++ b/inference/gradio_composite_demo/rife/loss.py @@ -0,0 +1,130 @@ +import torch +import numpy as np +import torch.nn as nn +import torch.nn.functional as F +import torchvision.models as models + +device = torch.device("cuda" if torch.cuda.is_available() else "cpu") + + +class EPE(nn.Module): + def __init__(self): + super(EPE, self).__init__() + + def forward(self, flow, gt, loss_mask): + loss_map = (flow - gt.detach()) ** 2 + loss_map = (loss_map.sum(1, True) + 1e-6) ** 0.5 + return loss_map * loss_mask + + +class Ternary(nn.Module): + def __init__(self): + super(Ternary, self).__init__() + patch_size = 7 + out_channels = patch_size * patch_size + self.w = np.eye(out_channels).reshape((patch_size, patch_size, 1, out_channels)) + self.w = np.transpose(self.w, (3, 2, 0, 1)) + self.w = torch.tensor(self.w).float().to(device) + + def transform(self, img): + patches = F.conv2d(img, self.w, padding=3, bias=None) + transf = patches - img + transf_norm = transf / torch.sqrt(0.81 + transf**2) + return transf_norm + + def rgb2gray(self, rgb): + r, g, b = rgb[:, 0:1, :, :], rgb[:, 1:2, :, :], rgb[:, 2:3, :, :] + gray = 0.2989 * r + 0.5870 * g + 0.1140 * b + return gray + + def hamming(self, t1, t2): + dist = (t1 - t2) ** 2 + dist_norm = torch.mean(dist / (0.1 + dist), 1, True) + return dist_norm + + def valid_mask(self, t, padding): + n, _, h, w = t.size() + inner = torch.ones(n, 1, h - 2 * padding, w - 2 * padding).type_as(t) + mask = F.pad(inner, [padding] * 4) + return mask + + def forward(self, img0, img1): + img0 = self.transform(self.rgb2gray(img0)) + img1 = self.transform(self.rgb2gray(img1)) + return self.hamming(img0, img1) * self.valid_mask(img0, 1) + + +class SOBEL(nn.Module): + def __init__(self): + super(SOBEL, self).__init__() + self.kernelX = torch.tensor( + [ + [1, 0, -1], + [2, 0, -2], + [1, 0, -1], + ] + ).float() + self.kernelY = self.kernelX.clone().T + self.kernelX = self.kernelX.unsqueeze(0).unsqueeze(0).to(device) + self.kernelY = self.kernelY.unsqueeze(0).unsqueeze(0).to(device) + + def forward(self, pred, gt): + N, C, H, W = pred.shape[0], pred.shape[1], pred.shape[2], pred.shape[3] + img_stack = torch.cat([pred.reshape(N * C, 1, H, W), gt.reshape(N * C, 1, H, W)], 0) + sobel_stack_x = F.conv2d(img_stack, self.kernelX, padding=1) + sobel_stack_y = F.conv2d(img_stack, self.kernelY, padding=1) + pred_X, gt_X = sobel_stack_x[: N * C], sobel_stack_x[N * C :] + pred_Y, gt_Y = sobel_stack_y[: N * C], sobel_stack_y[N * C :] + + L1X, L1Y = torch.abs(pred_X - gt_X), torch.abs(pred_Y - gt_Y) + loss = L1X + L1Y + return loss + + +class MeanShift(nn.Conv2d): + def __init__(self, data_mean, data_std, data_range=1, norm=True): + c = len(data_mean) + super(MeanShift, self).__init__(c, c, kernel_size=1) + std = torch.Tensor(data_std) + self.weight.data = torch.eye(c).view(c, c, 1, 1) + if norm: + self.weight.data.div_(std.view(c, 1, 1, 1)) + self.bias.data = -1 * data_range * torch.Tensor(data_mean) + self.bias.data.div_(std) + else: + self.weight.data.mul_(std.view(c, 1, 1, 1)) + self.bias.data = data_range * torch.Tensor(data_mean) + self.requires_grad = False + + +class VGGPerceptualLoss(torch.nn.Module): + def __init__(self, rank=0): + super(VGGPerceptualLoss, self).__init__() + blocks = [] + pretrained = True + self.vgg_pretrained_features = models.vgg19(pretrained=pretrained).features + self.normalize = MeanShift([0.485, 0.456, 0.406], [0.229, 0.224, 0.225], norm=True).cuda() + for param in self.parameters(): + param.requires_grad = False + + def forward(self, X, Y, indices=None): + X = self.normalize(X) + Y = self.normalize(Y) + indices = [2, 7, 12, 21, 30] + weights = [1.0 / 2.6, 1.0 / 4.8, 1.0 / 3.7, 1.0 / 5.6, 10 / 1.5] + k = 0 + loss = 0 + for i in range(indices[-1]): + X = self.vgg_pretrained_features[i](X) + Y = self.vgg_pretrained_features[i](Y) + if (i + 1) in indices: + loss += weights[k] * (X - Y.detach()).abs().mean() * 0.1 + k += 1 + return loss + + +if __name__ == "__main__": + img0 = torch.zeros(3, 3, 256, 256).float().to(device) + img1 = torch.tensor(np.random.normal(0, 1, (3, 3, 256, 256))).float().to(device) + ternary_loss = Ternary() + print(ternary_loss(img0, img1).shape) diff --git a/inference/gradio_composite_demo/rife/pytorch_msssim/__init__.py b/inference/gradio_composite_demo/rife/pytorch_msssim/__init__.py new file mode 100644 index 0000000..f1e4ec4 --- /dev/null +++ b/inference/gradio_composite_demo/rife/pytorch_msssim/__init__.py @@ -0,0 +1,256 @@ +import torch +import torch.nn.functional as F +from math import exp +import numpy as np + +device = torch.device("cuda" if torch.cuda.is_available() else "cpu") + + +def gaussian(window_size, sigma): + gauss = torch.Tensor( + [exp(-((x - window_size // 2) ** 2) / float(2 * sigma**2)) for x in range(window_size)] + ) + return gauss / gauss.sum() + + +def create_window(window_size, channel=1): + _1D_window = gaussian(window_size, 1.5).unsqueeze(1) + _2D_window = _1D_window.mm(_1D_window.t()).float().unsqueeze(0).unsqueeze(0).to(device) + window = _2D_window.expand(channel, 1, window_size, window_size).contiguous() + return window + + +def create_window_3d(window_size, channel=1): + _1D_window = gaussian(window_size, 1.5).unsqueeze(1) + _2D_window = _1D_window.mm(_1D_window.t()) + _3D_window = _2D_window.unsqueeze(2) @ (_1D_window.t()) + window = ( + _3D_window.expand(1, channel, window_size, window_size, window_size).contiguous().to(device) + ) + return window + + +def ssim(img1, img2, window_size=11, window=None, size_average=True, full=False, val_range=None): + # Value range can be different from 255. Other common ranges are 1 (sigmoid) and 2 (tanh). + if val_range is None: + if torch.max(img1) > 128: + max_val = 255 + else: + max_val = 1 + + if torch.min(img1) < -0.5: + min_val = -1 + else: + min_val = 0 + L = max_val - min_val + else: + L = val_range + + padd = 0 + (_, channel, height, width) = img1.size() + if window is None: + real_size = min(window_size, height, width) + window = create_window(real_size, channel=channel).to(img1.device) + + # mu1 = F.conv2d(img1, window, padding=padd, groups=channel) + # mu2 = F.conv2d(img2, window, padding=padd, groups=channel) + mu1 = F.conv2d( + F.pad(img1, (5, 5, 5, 5), mode="replicate"), window, padding=padd, groups=channel + ) + mu2 = F.conv2d( + F.pad(img2, (5, 5, 5, 5), mode="replicate"), window, padding=padd, groups=channel + ) + + mu1_sq = mu1.pow(2) + mu2_sq = mu2.pow(2) + mu1_mu2 = mu1 * mu2 + + sigma1_sq = ( + F.conv2d( + F.pad(img1 * img1, (5, 5, 5, 5), "replicate"), window, padding=padd, groups=channel + ) + - mu1_sq + ) + sigma2_sq = ( + F.conv2d( + F.pad(img2 * img2, (5, 5, 5, 5), "replicate"), window, padding=padd, groups=channel + ) + - mu2_sq + ) + sigma12 = ( + F.conv2d( + F.pad(img1 * img2, (5, 5, 5, 5), "replicate"), window, padding=padd, groups=channel + ) + - mu1_mu2 + ) + + C1 = (0.01 * L) ** 2 + C2 = (0.03 * L) ** 2 + + v1 = 2.0 * sigma12 + C2 + v2 = sigma1_sq + sigma2_sq + C2 + cs = torch.mean(v1 / v2) # contrast sensitivity + + ssim_map = ((2 * mu1_mu2 + C1) * v1) / ((mu1_sq + mu2_sq + C1) * v2) + + if size_average: + ret = ssim_map.mean() + else: + ret = ssim_map.mean(1).mean(1).mean(1) + + if full: + return ret, cs + return ret + + +def ssim_matlab( + img1, img2, window_size=11, window=None, size_average=True, full=False, val_range=None +): + # Value range can be different from 255. Other common ranges are 1 (sigmoid) and 2 (tanh). + if val_range is None: + if torch.max(img1) > 128: + max_val = 255 + else: + max_val = 1 + + if torch.min(img1) < -0.5: + min_val = -1 + else: + min_val = 0 + L = max_val - min_val + else: + L = val_range + + padd = 0 + (_, _, height, width) = img1.size() + if window is None: + real_size = min(window_size, height, width) + window = create_window_3d(real_size, channel=1).to(img1.device, dtype=img1.dtype) + # Channel is set to 1 since we consider color images as volumetric images + + img1 = img1.unsqueeze(1) + img2 = img2.unsqueeze(1) + + mu1 = F.conv3d( + F.pad(img1, (5, 5, 5, 5, 5, 5), mode="replicate"), window, padding=padd, groups=1 + ) + mu2 = F.conv3d( + F.pad(img2, (5, 5, 5, 5, 5, 5), mode="replicate"), window, padding=padd, groups=1 + ) + + mu1_sq = mu1.pow(2) + mu2_sq = mu2.pow(2) + mu1_mu2 = mu1 * mu2 + + sigma1_sq = ( + F.conv3d( + F.pad(img1 * img1, (5, 5, 5, 5, 5, 5), "replicate"), window, padding=padd, groups=1 + ) + - mu1_sq + ) + sigma2_sq = ( + F.conv3d( + F.pad(img2 * img2, (5, 5, 5, 5, 5, 5), "replicate"), window, padding=padd, groups=1 + ) + - mu2_sq + ) + sigma12 = ( + F.conv3d( + F.pad(img1 * img2, (5, 5, 5, 5, 5, 5), "replicate"), window, padding=padd, groups=1 + ) + - mu1_mu2 + ) + + C1 = (0.01 * L) ** 2 + C2 = (0.03 * L) ** 2 + + v1 = 2.0 * sigma12 + C2 + v2 = sigma1_sq + sigma2_sq + C2 + cs = torch.mean(v1 / v2) # contrast sensitivity + + ssim_map = ((2 * mu1_mu2 + C1) * v1) / ((mu1_sq + mu2_sq + C1) * v2) + + if size_average: + ret = ssim_map.mean() + else: + ret = ssim_map.mean(1).mean(1).mean(1) + + if full: + return ret, cs + return ret + + +def msssim(img1, img2, window_size=11, size_average=True, val_range=None, normalize=False): + device = img1.device + weights = torch.FloatTensor([0.0448, 0.2856, 0.3001, 0.2363, 0.1333]).to(device) + levels = weights.size()[0] + mssim = [] + mcs = [] + for _ in range(levels): + sim, cs = ssim( + img1, + img2, + window_size=window_size, + size_average=size_average, + full=True, + val_range=val_range, + ) + mssim.append(sim) + mcs.append(cs) + + img1 = F.avg_pool2d(img1, (2, 2)) + img2 = F.avg_pool2d(img2, (2, 2)) + + mssim = torch.stack(mssim) + mcs = torch.stack(mcs) + + # Normalize (to avoid NaNs during training unstable models, not compliant with original definition) + if normalize: + mssim = (mssim + 1) / 2 + mcs = (mcs + 1) / 2 + + pow1 = mcs**weights + pow2 = mssim**weights + # From Matlab implementation https://ece.uwaterloo.ca/~z70wang/research/iwssim/ + output = torch.prod(pow1[:-1] * pow2[-1]) + return output + + +# Classes to re-use window +class SSIM(torch.nn.Module): + def __init__(self, window_size=11, size_average=True, val_range=None): + super(SSIM, self).__init__() + self.window_size = window_size + self.size_average = size_average + self.val_range = val_range + + # Assume 3 channel for SSIM + self.channel = 3 + self.window = create_window(window_size, channel=self.channel) + + def forward(self, img1, img2): + (_, channel, _, _) = img1.size() + + if channel == self.channel and self.window.dtype == img1.dtype: + window = self.window + else: + window = create_window(self.window_size, channel).to(img1.device).type(img1.dtype) + self.window = window + self.channel = channel + + _ssim = ssim( + img1, img2, window=window, window_size=self.window_size, size_average=self.size_average + ) + dssim = (1 - _ssim) / 2 + return dssim + + +class MSSSIM(torch.nn.Module): + def __init__(self, window_size=11, size_average=True, channel=3): + super(MSSSIM, self).__init__() + self.window_size = window_size + self.size_average = size_average + self.channel = channel + + def forward(self, img1, img2): + return msssim(img1, img2, window_size=self.window_size, size_average=self.size_average) diff --git a/inference/gradio_composite_demo/rife/refine.py b/inference/gradio_composite_demo/rife/refine.py new file mode 100644 index 0000000..e956318 --- /dev/null +++ b/inference/gradio_composite_demo/rife/refine.py @@ -0,0 +1,136 @@ +import torch +import torch.nn as nn +from .warplayer import warp +import torch.nn.functional as F + +device = torch.device("cuda" if torch.cuda.is_available() else "cpu") + + +def conv(in_planes, out_planes, kernel_size=3, stride=1, padding=1, dilation=1): + return nn.Sequential( + nn.Conv2d( + in_planes, + out_planes, + kernel_size=kernel_size, + stride=stride, + padding=padding, + dilation=dilation, + bias=True, + ), + nn.PReLU(out_planes), + ) + + +def deconv(in_planes, out_planes, kernel_size=4, stride=2, padding=1): + return nn.Sequential( + torch.nn.ConvTranspose2d( + in_channels=in_planes, + out_channels=out_planes, + kernel_size=4, + stride=2, + padding=1, + bias=True, + ), + nn.PReLU(out_planes), + ) + + +class Conv2(nn.Module): + def __init__(self, in_planes, out_planes, stride=2): + super(Conv2, self).__init__() + self.conv1 = conv(in_planes, out_planes, 3, stride, 1) + self.conv2 = conv(out_planes, out_planes, 3, 1, 1) + + def forward(self, x): + x = self.conv1(x) + x = self.conv2(x) + return x + + +c = 16 + + +class Contextnet(nn.Module): + def __init__(self): + super(Contextnet, self).__init__() + self.conv1 = Conv2(3, c) + self.conv2 = Conv2(c, 2 * c) + self.conv3 = Conv2(2 * c, 4 * c) + self.conv4 = Conv2(4 * c, 8 * c) + + def forward(self, x, flow): + x = self.conv1(x) + flow = ( + F.interpolate( + flow, + scale_factor=0.5, + mode="bilinear", + align_corners=False, + recompute_scale_factor=False, + ) + * 0.5 + ) + f1 = warp(x, flow) + x = self.conv2(x) + flow = ( + F.interpolate( + flow, + scale_factor=0.5, + mode="bilinear", + align_corners=False, + recompute_scale_factor=False, + ) + * 0.5 + ) + f2 = warp(x, flow) + x = self.conv3(x) + flow = ( + F.interpolate( + flow, + scale_factor=0.5, + mode="bilinear", + align_corners=False, + recompute_scale_factor=False, + ) + * 0.5 + ) + f3 = warp(x, flow) + x = self.conv4(x) + flow = ( + F.interpolate( + flow, + scale_factor=0.5, + mode="bilinear", + align_corners=False, + recompute_scale_factor=False, + ) + * 0.5 + ) + f4 = warp(x, flow) + return [f1, f2, f3, f4] + + +class Unet(nn.Module): + def __init__(self): + super(Unet, self).__init__() + self.down0 = Conv2(17, 2 * c) + self.down1 = Conv2(4 * c, 4 * c) + self.down2 = Conv2(8 * c, 8 * c) + self.down3 = Conv2(16 * c, 16 * c) + self.up0 = deconv(32 * c, 8 * c) + self.up1 = deconv(16 * c, 4 * c) + self.up2 = deconv(8 * c, 2 * c) + self.up3 = deconv(4 * c, c) + self.conv = nn.Conv2d(c, 3, 3, 1, 1) + + def forward(self, img0, img1, warped_img0, warped_img1, mask, flow, c0, c1): + s0 = self.down0(torch.cat((img0, img1, warped_img0, warped_img1, mask, flow), 1)) + s1 = self.down1(torch.cat((s0, c0[0], c1[0]), 1)) + s2 = self.down2(torch.cat((s1, c0[1], c1[1]), 1)) + s3 = self.down3(torch.cat((s2, c0[2], c1[2]), 1)) + x = self.up0(torch.cat((s3, c0[3], c1[3]), 1)) + x = self.up1(torch.cat((x, s2), 1)) + x = self.up2(torch.cat((x, s1), 1)) + x = self.up3(torch.cat((x, s0), 1)) + x = self.conv(x) + return torch.sigmoid(x) diff --git a/inference/gradio_composite_demo/rife/refine_2R.py b/inference/gradio_composite_demo/rife/refine_2R.py new file mode 100644 index 0000000..6bc1515 --- /dev/null +++ b/inference/gradio_composite_demo/rife/refine_2R.py @@ -0,0 +1,127 @@ +import torch +import torch.nn as nn +from .warplayer import warp +import torch.nn.functional as F + +device = torch.device("cuda" if torch.cuda.is_available() else "cpu") + + +def conv(in_planes, out_planes, kernel_size=3, stride=1, padding=1, dilation=1): + return nn.Sequential( + nn.Conv2d( + in_planes, + out_planes, + kernel_size=kernel_size, + stride=stride, + padding=padding, + dilation=dilation, + bias=True, + ), + nn.PReLU(out_planes), + ) + + +def deconv(in_planes, out_planes, kernel_size=4, stride=2, padding=1): + return nn.Sequential( + torch.nn.ConvTranspose2d( + in_channels=in_planes, + out_channels=out_planes, + kernel_size=4, + stride=2, + padding=1, + bias=True, + ), + nn.PReLU(out_planes), + ) + + +class Conv2(nn.Module): + def __init__(self, in_planes, out_planes, stride=2): + super(Conv2, self).__init__() + self.conv1 = conv(in_planes, out_planes, 3, stride, 1) + self.conv2 = conv(out_planes, out_planes, 3, 1, 1) + + def forward(self, x): + x = self.conv1(x) + x = self.conv2(x) + return x + + +c = 16 + + +class Contextnet(nn.Module): + def __init__(self): + super(Contextnet, self).__init__() + self.conv1 = Conv2(3, c, 1) + self.conv2 = Conv2(c, 2 * c) + self.conv3 = Conv2(2 * c, 4 * c) + self.conv4 = Conv2(4 * c, 8 * c) + + def forward(self, x, flow): + x = self.conv1(x) + # flow = F.interpolate(flow, scale_factor=0.5, mode="bilinear", align_corners=False, recompute_scale_factor=False) * 0.5 + f1 = warp(x, flow) + x = self.conv2(x) + flow = ( + F.interpolate( + flow, + scale_factor=0.5, + mode="bilinear", + align_corners=False, + recompute_scale_factor=False, + ) + * 0.5 + ) + f2 = warp(x, flow) + x = self.conv3(x) + flow = ( + F.interpolate( + flow, + scale_factor=0.5, + mode="bilinear", + align_corners=False, + recompute_scale_factor=False, + ) + * 0.5 + ) + f3 = warp(x, flow) + x = self.conv4(x) + flow = ( + F.interpolate( + flow, + scale_factor=0.5, + mode="bilinear", + align_corners=False, + recompute_scale_factor=False, + ) + * 0.5 + ) + f4 = warp(x, flow) + return [f1, f2, f3, f4] + + +class Unet(nn.Module): + def __init__(self): + super(Unet, self).__init__() + self.down0 = Conv2(17, 2 * c, 1) + self.down1 = Conv2(4 * c, 4 * c) + self.down2 = Conv2(8 * c, 8 * c) + self.down3 = Conv2(16 * c, 16 * c) + self.up0 = deconv(32 * c, 8 * c) + self.up1 = deconv(16 * c, 4 * c) + self.up2 = deconv(8 * c, 2 * c) + self.up3 = deconv(4 * c, c) + self.conv = nn.Conv2d(c, 3, 3, 2, 1) + + def forward(self, img0, img1, warped_img0, warped_img1, mask, flow, c0, c1): + s0 = self.down0(torch.cat((img0, img1, warped_img0, warped_img1, mask, flow), 1)) + s1 = self.down1(torch.cat((s0, c0[0], c1[0]), 1)) + s2 = self.down2(torch.cat((s1, c0[1], c1[1]), 1)) + s3 = self.down3(torch.cat((s2, c0[2], c1[2]), 1)) + x = self.up0(torch.cat((s3, c0[3], c1[3]), 1)) + x = self.up1(torch.cat((x, s2), 1)) + x = self.up2(torch.cat((x, s1), 1)) + x = self.up3(torch.cat((x, s0), 1)) + x = self.conv(x) + return torch.sigmoid(x) diff --git a/inference/gradio_composite_demo/rife/warplayer.py b/inference/gradio_composite_demo/rife/warplayer.py new file mode 100644 index 0000000..ff796e8 --- /dev/null +++ b/inference/gradio_composite_demo/rife/warplayer.py @@ -0,0 +1,34 @@ +import torch +import torch.nn as nn + +device = torch.device("cuda" if torch.cuda.is_available() else "cpu") +backwarp_tenGrid = {} + + +def warp(tenInput, tenFlow): + k = (str(tenFlow.device), str(tenFlow.size())) + if k not in backwarp_tenGrid: + tenHorizontal = ( + torch.linspace(-1.0, 1.0, tenFlow.shape[3], device=device) + .view(1, 1, 1, tenFlow.shape[3]) + .expand(tenFlow.shape[0], -1, tenFlow.shape[2], -1) + ) + tenVertical = ( + torch.linspace(-1.0, 1.0, tenFlow.shape[2], device=device) + .view(1, 1, tenFlow.shape[2], 1) + .expand(tenFlow.shape[0], -1, -1, tenFlow.shape[3]) + ) + backwarp_tenGrid[k] = torch.cat([tenHorizontal, tenVertical], 1).to(device) + + tenFlow = torch.cat( + [ + tenFlow[:, 0:1, :, :] / ((tenInput.shape[3] - 1.0) / 2.0), + tenFlow[:, 1:2, :, :] / ((tenInput.shape[2] - 1.0) / 2.0), + ], + 1, + ) + + g = (backwarp_tenGrid[k] + tenFlow).permute(0, 2, 3, 1) + return torch.nn.functional.grid_sample( + input=tenInput, grid=g, mode="bilinear", padding_mode="border", align_corners=True + ) diff --git a/inference/gradio_composite_demo/rife_model.py b/inference/gradio_composite_demo/rife_model.py new file mode 100644 index 0000000..c21a713 --- /dev/null +++ b/inference/gradio_composite_demo/rife_model.py @@ -0,0 +1,183 @@ +import torch +from diffusers.image_processor import VaeImageProcessor +from torch.nn import functional as F +import cv2 +import utils +from rife.pytorch_msssim import ssim_matlab +import numpy as np +import logging +import skvideo.io +from rife.RIFE_HDv3 import Model +from huggingface_hub import hf_hub_download, snapshot_download + +logger = logging.getLogger(__name__) + +device = "cuda" if torch.cuda.is_available() else "cpu" + + +def pad_image(img, scale): + _, _, h, w = img.shape + tmp = max(32, int(32 / scale)) + ph = ((h - 1) // tmp + 1) * tmp + pw = ((w - 1) // tmp + 1) * tmp + padding = (0, pw - w, 0, ph - h) + return F.pad(img, padding), padding + + +def make_inference(model, I0, I1, upscale_amount, n): + middle = model.inference(I0, I1, upscale_amount) + if n == 1: + return [middle] + first_half = make_inference(model, I0, middle, upscale_amount, n=n // 2) + second_half = make_inference(model, middle, I1, upscale_amount, n=n // 2) + if n % 2: + return [*first_half, middle, *second_half] + else: + return [*first_half, *second_half] + + +@torch.inference_mode() +def ssim_interpolation_rife(model, samples, exp=1, upscale_amount=1, output_device="cpu"): + print(f"samples dtype:{samples.dtype}") + print(f"samples shape:{samples.shape}") + output = [] + pbar = utils.ProgressBar(samples.shape[0], desc="RIFE inference") + # [f, c, h, w] + for b in range(samples.shape[0]): + frame = samples[b : b + 1] + _, _, h, w = frame.shape + + I0 = samples[b : b + 1] + I1 = samples[b + 1 : b + 2] if b + 2 < samples.shape[0] else samples[-1:] + + I0, padding = pad_image(I0, upscale_amount) + I0 = I0.to(torch.float) + I1, _ = pad_image(I1, upscale_amount) + I1 = I1.to(torch.float) + + # [c, h, w] + I0_small = F.interpolate(I0, (32, 32), mode="bilinear", align_corners=False) + I1_small = F.interpolate(I1, (32, 32), mode="bilinear", align_corners=False) + + ssim = ssim_matlab(I0_small[:, :3], I1_small[:, :3]) + + if ssim > 0.996: + I1 = samples[b : b + 1] + # print(f'upscale_amount:{upscale_amount}') + # print(f'ssim:{upscale_amount}') + # print(f'I0 shape:{I0.shape}') + # print(f'I1 shape:{I1.shape}') + I1, padding = pad_image(I1, upscale_amount) + # print(f'I0 shape:{I0.shape}') + # print(f'I1 shape:{I1.shape}') + I1 = make_inference(model, I0, I1, upscale_amount, 1) + + # print(f'I0 shape:{I0.shape}') + # print(f'I1[0] shape:{I1[0].shape}') + I1 = I1[0] + + # print(f'I1[0] unpadded shape:{I1.shape}') + I1_small = F.interpolate(I1, (32, 32), mode="bilinear", align_corners=False) + ssim = ssim_matlab(I0_small[:, :3], I1_small[:, :3]) + if padding[3] > 0 and padding[1] > 0: + frame = I1[:, :, : -padding[3], : -padding[1]] + elif padding[3] > 0: + frame = I1[:, :, : -padding[3], :] + elif padding[1] > 0: + frame = I1[:, :, :, : -padding[1]] + else: + frame = I1 + + tmp_output = [] + if ssim < 0.2: + for i in range((2**exp) - 1): + tmp_output.append(I0) + + else: + tmp_output = make_inference(model, I0, I1, upscale_amount, 2**exp - 1) if exp else [] + + frame, _ = pad_image(frame, upscale_amount) + # print(f'frame shape:{frame.shape}') + + frame = F.interpolate(frame, size=(h, w)) + output.append(frame.to(output_device)) + for i, tmp_frame in enumerate(tmp_output): + # tmp_frame, _ = pad_image(tmp_frame, upscale_amount) + tmp_frame = F.interpolate(tmp_frame, size=(h, w)) + output.append(tmp_frame.to(output_device)) + pbar.update(1) + return output + + +def load_rife_model(model_path): + model = Model() + model.load_model(model_path, -1) + model.eval() + return model + + +# Create a generator that yields each frame, similar to cv2.VideoCapture +def frame_generator(video_capture): + while True: + ret, frame = video_capture.read() + if not ret: + break + yield frame + video_capture.release() + + +def rife_inference_with_path(model, video_path): + # Open the video file + video_capture = cv2.VideoCapture(video_path) + fps = video_capture.get(cv2.CAP_PROP_FPS) # Get the frames per second + tot_frame = int(video_capture.get(cv2.CAP_PROP_FRAME_COUNT)) # Total frames in the video + pt_frame_data = [] + pt_frame = skvideo.io.vreader(video_path) + # Cyclic reading of the video frames + while video_capture.isOpened(): + ret, frame = video_capture.read() + + if not ret: + break + + # BGR to RGB + frame_rgb = frame[..., ::-1] + frame_rgb = frame_rgb.copy() + tensor = torch.from_numpy(frame_rgb).float().to("cpu", non_blocking=True).float() / 255.0 + pt_frame_data.append(tensor.permute(2, 0, 1)) # to [c, h, w,] + + pt_frame = torch.from_numpy(np.stack(pt_frame_data)) + pt_frame = pt_frame.to(device) + pbar = utils.ProgressBar(tot_frame, desc="RIFE inference") + frames = ssim_interpolation_rife(model, pt_frame) + pt_image = torch.stack([frames[i].squeeze(0) for i in range(len(frames))]) + image_np = VaeImageProcessor.pt_to_numpy(pt_image) # (to [49, 512, 480, 3]) + image_pil = VaeImageProcessor.numpy_to_pil(image_np) + video_path = utils.save_video(image_pil, fps=16) + if pbar: + pbar.update(1) + return video_path + + +def rife_inference_with_latents(model, latents): + rife_results = [] + latents = latents.to(device) + for i in range(latents.size(0)): + # [f, c, w, h] + latent = latents[i] + + frames = ssim_interpolation_rife(model, latent) + pt_image = torch.stack( + [frames[i].squeeze(0) for i in range(len(frames))] + ) # (to [f, c, w, h]) + rife_results.append(pt_image) + + return torch.stack(rife_results) + + +# if __name__ == "__main__": +# snapshot_download(repo_id="AlexWortega/RIFE", local_dir="model_rife") +# model = load_rife_model("model_rife") + +# video_path = rife_inference_with_path(model, "/mnt/ceph/develop/jiawei/CogVideo/output/20241003_130720.mp4") +# print(video_path) diff --git a/inference/gradio_composite_demo/utils.py b/inference/gradio_composite_demo/utils.py new file mode 100644 index 0000000..d39f227 --- /dev/null +++ b/inference/gradio_composite_demo/utils.py @@ -0,0 +1,244 @@ +import math +from typing import Union, List + +import torch +import os +from datetime import datetime +import numpy as np +import itertools +import PIL.Image +import safetensors.torch +import tqdm +import logging +from diffusers.utils import export_to_video +from spandrel import ModelLoader + +logger = logging.getLogger(__file__) + + +def load_torch_file(ckpt, device=None, dtype=torch.float16): + if device is None: + device = torch.device("cpu") + if ckpt.lower().endswith(".safetensors") or ckpt.lower().endswith(".sft"): + sd = safetensors.torch.load_file(ckpt, device=device.type) + else: + if "weights_only" not in torch.load.__code__.co_varnames: + logger.warning( + "Warning torch.load doesn't support weights_only on this pytorch version, loading unsafely." + ) + + pl_sd = torch.load(ckpt, map_location=device, weights_only=True) + if "global_step" in pl_sd: + logger.debug(f"Global Step: {pl_sd['global_step']}") + if "state_dict" in pl_sd: + sd = pl_sd["state_dict"] + elif "params_ema" in pl_sd: + sd = pl_sd["params_ema"] + else: + sd = pl_sd + + sd = {k: v.to(dtype) for k, v in sd.items()} + return sd + + +def state_dict_prefix_replace(state_dict, replace_prefix, filter_keys=False): + if filter_keys: + out = {} + else: + out = state_dict + for rp in replace_prefix: + replace = list( + map( + lambda a: (a, "{}{}".format(replace_prefix[rp], a[len(rp) :])), + filter(lambda a: a.startswith(rp), state_dict.keys()), + ) + ) + for x in replace: + w = state_dict.pop(x[0]) + out[x[1]] = w + return out + + +def module_size(module): + module_mem = 0 + sd = module.state_dict() + for k in sd: + t = sd[k] + module_mem += t.nelement() * t.element_size() + return module_mem + + +def get_tiled_scale_steps(width, height, tile_x, tile_y, overlap): + return math.ceil((height / (tile_y - overlap))) * math.ceil((width / (tile_x - overlap))) + + +@torch.inference_mode() +def tiled_scale_multidim( + samples, + function, + tile=(64, 64), + overlap=8, + upscale_amount=4, + out_channels=3, + output_device="cpu", + pbar=None, +): + dims = len(tile) + print(f"samples dtype:{samples.dtype}") + output = torch.empty( + [samples.shape[0], out_channels] + + list(map(lambda a: round(a * upscale_amount), samples.shape[2:])), + device=output_device, + ) + + for b in range(samples.shape[0]): + s = samples[b : b + 1] + out = torch.zeros( + [s.shape[0], out_channels] + + list(map(lambda a: round(a * upscale_amount), s.shape[2:])), + device=output_device, + ) + out_div = torch.zeros( + [s.shape[0], out_channels] + + list(map(lambda a: round(a * upscale_amount), s.shape[2:])), + device=output_device, + ) + + for it in itertools.product( + *map(lambda a: range(0, a[0], a[1] - overlap), zip(s.shape[2:], tile)) + ): + s_in = s + upscaled = [] + + for d in range(dims): + pos = max(0, min(s.shape[d + 2] - overlap, it[d])) + l = min(tile[d], s.shape[d + 2] - pos) + s_in = s_in.narrow(d + 2, pos, l) + upscaled.append(round(pos * upscale_amount)) + + ps = function(s_in).to(output_device) + mask = torch.ones_like(ps) + feather = round(overlap * upscale_amount) + for t in range(feather): + for d in range(2, dims + 2): + m = mask.narrow(d, t, 1) + m *= (1.0 / feather) * (t + 1) + m = mask.narrow(d, mask.shape[d] - 1 - t, 1) + m *= (1.0 / feather) * (t + 1) + + o = out + o_d = out_div + for d in range(dims): + o = o.narrow(d + 2, upscaled[d], mask.shape[d + 2]) + o_d = o_d.narrow(d + 2, upscaled[d], mask.shape[d + 2]) + + o += ps * mask + o_d += mask + + if pbar is not None: + pbar.update(1) + + output[b : b + 1] = out / out_div + return output + + +def tiled_scale( + samples, + function, + tile_x=64, + tile_y=64, + overlap=8, + upscale_amount=4, + out_channels=3, + output_device="cpu", + pbar=None, +): + return tiled_scale_multidim( + samples, + function, + (tile_y, tile_x), + overlap, + upscale_amount, + out_channels, + output_device, + pbar, + ) + + +def load_sd_upscale(ckpt, inf_device): + sd = load_torch_file(ckpt, device=inf_device) + if "module.layers.0.residual_group.blocks.0.norm1.weight" in sd: + sd = state_dict_prefix_replace(sd, {"module.": ""}) + out = ModelLoader().load_from_state_dict(sd).half() + return out + + +def upscale(upscale_model, tensor: torch.Tensor, inf_device, output_device="cpu") -> torch.Tensor: + memory_required = module_size(upscale_model.model) + memory_required += ( + (512 * 512 * 3) * tensor.element_size() * max(upscale_model.scale, 1.0) * 384.0 + ) # The 384.0 is an estimate of how much some of these models take, TODO: make it more accurate + memory_required += tensor.nelement() * tensor.element_size() + print(f"UPScaleMemory required: {memory_required / 1024 / 1024 / 1024} GB") + + upscale_model.to(inf_device) + tile = 512 + overlap = 32 + + steps = tensor.shape[0] * get_tiled_scale_steps( + tensor.shape[3], tensor.shape[2], tile_x=tile, tile_y=tile, overlap=overlap + ) + + pbar = ProgressBar(steps, desc="Tiling and Upscaling") + + s = tiled_scale( + samples=tensor.to(torch.float16), + function=lambda a: upscale_model(a), + tile_x=tile, + tile_y=tile, + overlap=overlap, + upscale_amount=upscale_model.scale, + pbar=pbar, + ) + + upscale_model.to(output_device) + return s + + +def upscale_batch_and_concatenate( + upscale_model, latents, inf_device, output_device="cpu" +) -> torch.Tensor: + upscaled_latents = [] + for i in range(latents.size(0)): + latent = latents[i] + upscaled_latent = upscale(upscale_model, latent, inf_device, output_device) + upscaled_latents.append(upscaled_latent) + return torch.stack(upscaled_latents) + + +def save_video(tensor: Union[List[np.ndarray], List[PIL.Image.Image]], fps: int = 8): + timestamp = datetime.now().strftime("%Y%m%d_%H%M%S") + video_path = f"./output/{timestamp}.mp4" + os.makedirs(os.path.dirname(video_path), exist_ok=True) + export_to_video(tensor, video_path, fps=fps) + return video_path + + +class ProgressBar: + def __init__(self, total, desc=None): + self.total = total + self.current = 0 + self.b_unit = tqdm.tqdm( + total=total, desc="ProgressBar context index: 0" if desc is None else desc + ) + + def update(self, value): + if value > self.total: + value = self.total + self.current = value + if self.b_unit is not None: + self.b_unit.set_description("ProgressBar context index: {}".format(self.current)) + self.b_unit.refresh() + + # 更新进度 + self.b_unit.update(self.current) diff --git a/inference/gradio_web_demo.py b/inference/gradio_web_demo.py new file mode 100644 index 0000000..955c3b7 --- /dev/null +++ b/inference/gradio_web_demo.py @@ -0,0 +1,215 @@ +""" +THis is the main file for the gradio web demo. It uses the CogVideoX-2B model to generate videos gradio web demo. +set environment variable OPENAI_API_KEY to use the OpenAI API to enhance the prompt. + +This demo only supports the text-to-video generation model. +If you wish to use the image-to-video or video-to-video generation models, +please use the gradio_composite_demo to implement the full GUI functionality. + +Usage: + OpenAI_API_KEY=your_openai_api_key OpenAI_BASE_URL=https://api.openai.com/v1 python inference/gradio_web_demo.py +""" + +import os +import threading +import time + +import gradio as gr +import torch +from diffusers import CogVideoXPipeline +from diffusers.utils import export_to_video +from datetime import datetime, timedelta +from openai import OpenAI +from moviepy import VideoFileClip + +pipe = CogVideoXPipeline.from_pretrained("THUDM/CogVideoX-5b", torch_dtype=torch.bfloat16).to( + "cuda" +) + +pipe.vae.enable_slicing() +pipe.vae.enable_tiling() + +os.makedirs("./output", exist_ok=True) +os.makedirs("./gradio_tmp", exist_ok=True) + +sys_prompt = """You are part of a team of bots that creates videos. You work with an assistant bot that will draw anything you say in square brackets. + +For example , outputting " a beautiful morning in the woods with the sun peaking through the trees " will trigger your partner bot to output an video of a forest morning , as described. You will be prompted by people looking to create detailed , amazing videos. The way to accomplish this is to take their short prompts and make them extremely detailed and descriptive. +There are a few rules to follow: + +You will only ever output a single video description per user request. + +When modifications are requested , you should not simply make the description longer . You should refactor the entire description to integrate the suggestions. +Other times the user will not want modifications , but instead want a new image . In this case , you should ignore your previous conversation with the user. + +Video descriptions must have the same num of words as examples below. Extra words will be ignored. +""" + + +def convert_prompt(prompt: str, retry_times: int = 3) -> str: + if not os.environ.get("OPENAI_API_KEY"): + return prompt + + client = OpenAI() + text = prompt.strip() + + for i in range(retry_times): + response = client.chat.completions.create( + messages=[ + {"role": "system", "content": sys_prompt}, + { + "role": "user", + "content": 'Create an imaginative video descriptive caption or modify an earlier caption for the user input : "a girl is on the beach"', + }, + { + "role": "assistant", + "content": "A radiant woman stands on a deserted beach, arms outstretched, wearing a beige trench coat, white blouse, light blue jeans, and chic boots, against a backdrop of soft sky and sea. Moments later, she is seen mid-twirl, arms exuberant, with the lighting suggesting dawn or dusk. Then, she runs along the beach, her attire complemented by an off-white scarf and black ankle boots, the tranquil sea behind her. Finally, she holds a paper airplane, her pose reflecting joy and freedom, with the ocean's gentle waves and the sky's soft pastel hues enhancing the serene ambiance.", + }, + { + "role": "user", + "content": 'Create an imaginative video descriptive caption or modify an earlier caption for the user input : "A man jogging on a football field"', + }, + { + "role": "assistant", + "content": "A determined man in athletic attire, including a blue long-sleeve shirt, black shorts, and blue socks, jogs around a snow-covered soccer field, showcasing his solitary exercise in a quiet, overcast setting. His long dreadlocks, focused expression, and the serene winter backdrop highlight his dedication to fitness. As he moves, his attire, consisting of a blue sports sweatshirt, black athletic pants, gloves, and sneakers, grips the snowy ground. He is seen running past a chain-link fence enclosing the playground area, with a basketball hoop and children's slide, suggesting a moment of solitary exercise amidst the empty field.", + }, + { + "role": "user", + "content": 'Create an imaginative video descriptive caption or modify an earlier caption for the user input : " A woman is dancing, HD footage, close-up"', + }, + { + "role": "assistant", + "content": "A young woman with her hair in an updo and wearing a teal hoodie stands against a light backdrop, initially looking over her shoulder with a contemplative expression. She then confidently makes a subtle dance move, suggesting rhythm and movement. Next, she appears poised and focused, looking directly at the camera. Her expression shifts to one of introspection as she gazes downward slightly. Finally, she dances with confidence, her left hand over her heart, symbolizing a poignant moment, all while dressed in the same teal hoodie against a plain, light-colored background.", + }, + { + "role": "user", + "content": f'Create an imaginative video descriptive caption or modify an earlier caption in ENGLISH for the user input: "{text}"', + }, + ], + model="glm-4-plus", + temperature=0.01, + top_p=0.7, + stream=False, + max_tokens=250, + ) + if response.choices: + return response.choices[0].message.content + return prompt + + +def infer( + prompt: str, + num_inference_steps: int, + guidance_scale: float, + progress=gr.Progress(track_tqdm=True), +): + torch.cuda.empty_cache() + video = pipe( + prompt=prompt, + num_videos_per_prompt=1, + num_inference_steps=num_inference_steps, + num_frames=49, + guidance_scale=guidance_scale, + ).frames[0] + + return video + + +def save_video(tensor): + timestamp = datetime.now().strftime("%Y%m%d_%H%M%S") + video_path = f"./output/{timestamp}.mp4" + os.makedirs(os.path.dirname(video_path), exist_ok=True) + export_to_video(tensor, video_path) + return video_path + + +def convert_to_gif(video_path): + clip = VideoFileClip(video_path) + clip = clip.with_fps(8) + clip = clip.resized(height=240) + gif_path = video_path.replace(".mp4", ".gif") + clip.write_gif(gif_path, fps=8) + return gif_path + + +def delete_old_files(): + while True: + now = datetime.now() + cutoff = now - timedelta(minutes=10) + directories = ["./output", "./gradio_tmp"] + + for directory in directories: + for filename in os.listdir(directory): + file_path = os.path.join(directory, filename) + if os.path.isfile(file_path): + file_mtime = datetime.fromtimestamp(os.path.getmtime(file_path)) + if file_mtime < cutoff: + os.remove(file_path) + time.sleep(600) + + +threading.Thread(target=delete_old_files, daemon=True).start() + +with gr.Blocks() as demo: + gr.Markdown(""" +
+ CogVideoX Gradio Simple Space🤗 + """) + + with gr.Row(): + with gr.Column(): + prompt = gr.Textbox( + label="Prompt (Less than 200 Words)", placeholder="Enter your prompt here", lines=5 + ) + + with gr.Row(): + gr.Markdown( + "✨Upon pressing the enhanced prompt button, we will use [GLM-4 Model](https://github.com/THUDM/GLM-4) to polish the prompt and overwrite the original one." + ) + enhance_button = gr.Button("✨ Enhance Prompt(Optional)") + + with gr.Column(): + gr.Markdown( + "**Optional Parameters** (default values are recommended)
" + "Increasing the number of inference steps will produce more detailed videos, but it will slow down the process.
" + "50 steps are recommended for most cases.
" + ) + with gr.Row(): + num_inference_steps = gr.Number(label="Inference Steps", value=50) + guidance_scale = gr.Number(label="Guidance Scale", value=6.0) + generate_button = gr.Button("🎬 Generate Video") + + with gr.Column(): + video_output = gr.Video(label="CogVideoX Generate Video", width=720, height=480) + with gr.Row(): + download_video_button = gr.File(label="📥 Download Video", visible=False) + download_gif_button = gr.File(label="📥 Download GIF", visible=False) + + def generate( + prompt, + num_inference_steps, + guidance_scale, + model_choice, + progress=gr.Progress(track_tqdm=True), + ): + tensor = infer(prompt, num_inference_steps, guidance_scale, progress=progress) + video_path = save_video(tensor) + video_update = gr.update(visible=True, value=video_path) + gif_path = convert_to_gif(video_path) + gif_update = gr.update(visible=True, value=gif_path) + + return video_path, video_update, gif_update + + def enhance_prompt_func(prompt): + return convert_prompt(prompt, retry_times=1) + + generate_button.click( + generate, + inputs=[prompt, num_inference_steps, guidance_scale], + outputs=[video_output, download_video_button, download_gif_button], + ) + + enhance_button.click(enhance_prompt_func, inputs=[prompt], outputs=[prompt]) + +if __name__ == "__main__": + demo.launch() diff --git a/pyproject.toml b/pyproject.toml new file mode 100644 index 0000000..8c9fa17 --- /dev/null +++ b/pyproject.toml @@ -0,0 +1,43 @@ +[project] +name = "cogvideo" +version = "0.1.0" +readme = "README.md" +requires-python = ">=3.8" +dependencies = [ + "datasets>=2.14.4", + "decord>=0.6.0", + "deepspeed>=0.16.4", + "diffusers>=0.32.2", + "opencv-python>=4.11.0.86", + "peft>=0.13.2", + "pydantic>=2.10.6", + "sentencepiece>=0.2.0", + "torch>=2.5.1", + "torchvision>=0.20.1", + "transformers>=4.46.3", + "wandb>=0.19.7", +] + + +[tool.ruff] +exclude = ['.git', '.mypy_cache', '.ruff_cache', '.venv', 'dist'] +target-version = 'py310' +line-length = 100 + +[tool.ruff.format] +line-ending = 'lf' +quote-style = 'preserve' + +[tool.ruff.lint] +select = ["E", "W"] +ignore = [ + "E999", + "EXE001", + "UP009", + "F401", + "TID252", + "F403", + "F841", + "E501", + "W293", +] diff --git a/requirements.txt b/requirements.txt new file mode 100644 index 0000000..97d4365 --- /dev/null +++ b/requirements.txt @@ -0,0 +1,15 @@ +diffusers>=0.35.2 +accelerate>=1.11.0 +transformers>=4.57.1 +numpy==1.26.0 +torch>=2.8.0 +torchvision>=0.23.0 +sentencepiece>=0.2.1 +SwissArmyTransformer>=0.4.12 +gradio>=5.49.1 +imageio>=2.37.0 +imageio-ffmpeg>=0.6.0 +openai>=2.2.0 +moviepy>=2.2.1 +scikit-video>=1.1.11 +pydantic>=2.12.0 diff --git a/resources/WECHAT.md b/resources/WECHAT.md new file mode 100644 index 0000000..8c56211 --- /dev/null +++ b/resources/WECHAT.md @@ -0,0 +1,6 @@ +
+ + +

扫码关注公众号，加入「 CogVideoX 交流群」

+

Scan the QR code to follow the official account and join the "CogVLM Discussion Group"

+
diff --git a/resources/galary_prompt.md b/resources/galary_prompt.md new file mode 100644 index 0000000..b848211 --- /dev/null +++ b/resources/galary_prompt.md @@ -0,0 +1,31 @@ +## CogVideoX-5B + +Videos 1-8: + +1. A garden comes to life as a kaleidoscope of butterflies flutters amidst the blossoms, their delicate wings casting shadows on the petals below. In the background, a grand fountain cascades water with a gentle splendor, its rhythmic sound providing a soothing backdrop. Beneath the cool shade of a mature tree, a solitary wooden chair invites solitude and reflection, its smooth surface worn by the touch of countless visitors seeking a moment of tranquility in nature's embrace. + +2. A small boy, head bowed and determination etched on his face, sprints through the torrential downpour as闪电 crackles and 雷鸣 rumbles in the distance. The relentless rain pounds the ground, creating a chaotic dance of water droplets that mirror the Dramatic sky's anger. In the far background, the silhouette of a cozy home beckons, a faint beacon of safety and warmth amidst the fierce weather. The scene is one of perseverance and the unyielding spirit of a child braving the elements. + +3. A suited astronaut, with the red dust of Mars clinging to their boots, reaches out to shake hands with an alien being, their skin a shimmering blue, under the pink-tinged sky of the fourth planet. In the background, a sleek silver rocket, a beacon of human ingenuity, stands tall, its engines powered down, as the two representatives of different worlds exchange a historic greeting amidst the desolate beauty of the Martian landscape. + +4. An elderly gentleman, with a serene expression, sits at the water's edge, a steaming cup of tea by his side. He is engrossed in his artwork, brush in hand, as he renders an oil painting on a canvas that's propped up against a small, weathered table. The sea breeze whispers through his silver hair, gently billowing his loose-fitting white shirt, while the salty air adds an intangible element to his masterpiece in progress. The scene is one of tranquility and inspiration, with the artist's canvas capturing the vibrant hues of the setting sun reflecting off the tranquil sea. + +5. In a dimly lit bar, purplish light bathes the face of a mature man, his eyes blinking thoughtfully as he ponders in close-up, the background artfully blurred to focus on his introspective expression, the ambiance of the bar a mere suggestion of shadows and soft lighting. + +6. A golden retriever, sporting sleek black sunglasses, with its lengthy fur flowing in the breeze, sprints playfully across a rooftop terrace, recently refreshed by a light rain. The scene unfolds from a distance, the dog's energetic bounds growing larger as it approaches the camera, its tail wagging with unrestrained joy, while droplets of water glisten on the concrete behind it. The overcast sky provides a dramatic backdrop, emphasizing the vibrant golden coat of the canine as it dashes towards the viewer. + +7. On a brilliant sunny day, the lakeshore is lined with an array of willow trees, their slender branches swaying gently in the soft breeze. The tranquil surface of the lake reflects the clear blue sky, while several elegant swans glide gracefully through the still water, leaving behind delicate ripples that disturb the mirror-like quality of the lake. The scene is one of serene beauty, with the willows' greenery providing a picturesque frame for the peaceful avian visitors. + +8. A Chinese mother, draped in a soft, pastel-colored robe, gently rocks back and forth in a cozy rocking chair positioned in the tranquil setting of a nursery. The dimly lit bedroom is adorned with whimsical mobiles dangling from the ceiling, casting shadows that dance on the walls. Her baby, swaddled in a delicate, patterned blanket, rests against her chest, the child's earlier cries now replaced by contented coos as the mother's soothing voice lulls the little one to sleep. The scent of lavender fills the air, adding to the serene atmosphere, while a warm, orange glow from a nearby nightlight illuminates the scene with a gentle hue, capturing a moment of tender love and comfort. + +## CogVideoX-2B + +Videos 1-4: + +1. A detailed wooden toy ship with intricately carved masts and sails is seen gliding smoothly over a plush, blue carpet that mimics the waves of the sea. The ship's hull is painted a rich brown, with tiny windows. The carpet, soft and textured, provides a perfect backdrop, resembling an oceanic expanse. Surrounding the ship are various other toys and children's items, hinting at a playful environment. The scene captures the innocence and imagination of childhood, with the toy ship's journey symbolizing endless adventures in a whimsical, indoor setting. + +2. The camera follows behind a white vintage SUV with a black roof rack as it speeds up a steep dirt road surrounded by pine trees on a steep mountain slope, dust kicks up from its tires, the sunlight shines on the SUV as it speeds along the dirt road, casting a warm glow over the scene. The dirt road curves gently into the distance, with no other cars or vehicles in sight. The trees on either side of the road are redwoods, with patches of greenery scattered throughout. The car is seen from the rear following the curve with ease, making it seem as if it is on a rugged drive through the rugged terrain. The dirt road itself is surrounded by steep hills and mountains, with a clear blue sky above with wispy clouds. + +3. A street artist, clad in a worn-out denim jacket and a colorful bandana, stands before a vast concrete wall in the heart, holding a can of spray paint, spray-painting a colorful bird on a mottled wall. + +4. In the haunting backdrop of a war-torn city, where ruins and crumbled walls tell a story of devastation, a poignant close-up frames a young girl. Her face is smudged with ash, a silent testament to the chaos around her. Her eyes glistening with a mix of sorrow and resilience, capturing the raw emotion of a world that has lost its innocence to the ravages of conflict. diff --git a/resources/hf_lora_weights.png b/resources/hf_lora_weights.png new file mode 100644 index 0000000..cdf13cc Binary files /dev/null and b/resources/hf_lora_weights.png differ diff --git a/resources/logo.svg b/resources/logo.svg new file mode 100644 index 0000000..d0b8d44 --- /dev/null +++ b/resources/logo.svg @@ -0,0 +1,142 @@ + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + diff --git a/resources/web_demo.png b/resources/web_demo.png new file mode 100644 index 0000000..53cf691 Binary files /dev/null and b/resources/web_demo.png differ diff --git a/resources/wechat.jpg b/resources/wechat.jpg new file mode 100644 index 0000000..8c67652 Binary files /dev/null and b/resources/wechat.jpg differ diff --git a/tools/caption/README.md b/tools/caption/README.md new file mode 100644 index 0000000..2cb8575 --- /dev/null +++ b/tools/caption/README.md @@ -0,0 +1,67 @@ +# Video Caption + +Typically, most video data does not come with corresponding descriptive text, so it is necessary to convert the video +data into textual descriptions to provide the essential training data for text-to-video models. + +## Update and News +- 🔥🔥 **News**: ```2024/9/19```: The caption model used in the CogVideoX training process to convert video data into text + descriptions, [CogVLM2-Caption](https://huggingface.co/THUDM/cogvlm2-llama3-caption), is now open-source. Feel + free to download and use it. + + +## Video Caption via CogVLM2-Caption + +🤗 [Hugging Face](https://huggingface.co/THUDM/cogvlm2-llama3-caption) | 🤖 [ModelScope](https://modelscope.cn/models/ZhipuAI/cogvlm2-llama3-caption/) + +CogVLM2-Caption is a video captioning model used to generate training data for the CogVideoX model. + +### Install +```shell +pip install -r requirements.txt +``` + +### Usage + +```shell +python video_caption.py +``` + +Example: +
+ +
+ +## Video Caption via CogVLM2-Video + +[Code](https://github.com/THUDM/CogVLM2/tree/main/video_demo) | 🤗 [Hugging Face](https://huggingface.co/THUDM/cogvlm2-video-llama3-chat) | 🤖 [ModelScope](https://modelscope.cn/models/ZhipuAI/cogvlm2-video-llama3-chat) | 📑 [Blog](https://cogvlm2-video.github.io/) ｜ [💬 Online Demo](http://cogvlm2-online.cogviewai.cn:7868/) + +CogVLM2-Video is a versatile video understanding model equipped with timestamp-based question answering capabilities. +Users can input prompts such as `Please describe this video in detail.` to the model to obtain a detailed video caption: +
+ +
+ +Users can use the provided [code](https://github.com/THUDM/CogVLM2/tree/main/video_demo) to load the model or configure a RESTful API to generate video captions. + +## Citation + +🌟 If you find our work helpful, please leave us a star and cite our paper. + +CogVLM2-Caption: +``` +@article{yang2024cogvideox, + title={CogVideoX: Text-to-Video Diffusion Models with An Expert Transformer}, + author={Yang, Zhuoyi and Teng, Jiayan and Zheng, Wendi and Ding, Ming and Huang, Shiyu and Xu, Jiazheng and Yang, Yuanming and Hong, Wenyi and Zhang, Xiaohan and Feng, Guanyu and others}, + journal={arXiv preprint arXiv:2408.06072}, + year={2024} +} +``` +CogVLM2-Video: +``` +@article{hong2024cogvlm2, + title={CogVLM2: Visual Language Models for Image and Video Understanding}, + author={Hong, Wenyi and Wang, Weihan and Ding, Ming and Yu, Wenmeng and Lv, Qingsong and Wang, Yan and Cheng, Yean and Huang, Shiyu and Ji, Junhui and Xue, Zhao and others}, + journal={arXiv preprint arXiv:2408.16500}, + year={2024} +} +``` diff --git a/tools/caption/README_ja.md b/tools/caption/README_ja.md new file mode 100644 index 0000000..e328a9a --- /dev/null +++ b/tools/caption/README_ja.md @@ -0,0 +1,65 @@ +# ビデオキャプション + +通常、ほとんどのビデオデータには対応する説明文が付いていないため、ビデオデータをテキストの説明に変換して、テキストからビデオへのモデルに必要なトレーニングデータを提供する必要があります。 + +## 更新とニュース +- 🔥🔥 **ニュース**: ```2024/9/19```：CogVideoX + のトレーニングプロセスで、ビデオデータをテキストに変換するためのキャプションモデル [CogVLM2-Caption](https://huggingface.co/THUDM/cogvlm2-llama3-caption) + がオープンソース化されました。ぜひダウンロードしてご利用ください。 +## CogVLM2-Captionによるビデオキャプション + +🤗 [Hugging Face](https://huggingface.co/THUDM/cogvlm2-llama3-caption) | 🤖 [ModelScope](https://modelscope.cn/models/ZhipuAI/cogvlm2-llama3-caption/) + +CogVLM2-Captionは、CogVideoXモデルのトレーニングデータを生成するために使用されるビデオキャプションモデルです。 + +### インストール +```shell +pip install -r requirements.txt +``` + +### 使用方法 +```shell +python video_caption.py +``` + +例: +
+ +
+ + + +## CogVLM2-Video を使用したビデオキャプション + +[Code](https://github.com/THUDM/CogVLM2/tree/main/video_demo) | 🤗 [Hugging Face](https://huggingface.co/THUDM/cogvlm2-video-llama3-chat) | 🤖 [ModelScope](https://modelscope.cn/models/ZhipuAI/cogvlm2-video-llama3-chat) | 📑 [Blog](https://cogvlm2-video.github.io/) ｜ [💬 Online Demo](http://cogvlm2-online.cogviewai.cn:7868/) + + +CogVLM2-Video は、タイムスタンプベースの質問応答機能を備えた多機能なビデオ理解モデルです。ユーザーは `このビデオを詳細に説明してください。` などのプロンプトをモデルに入力して、詳細なビデオキャプションを取得できます： +
+ +
+ +ユーザーは提供された[コード](https://github.com/THUDM/CogVLM2/tree/main/video_demo)を使用してモデルをロードするか、RESTful API を構成してビデオキャプションを生成できます。 + +## Citation + +🌟 If you find our work helpful, please leave us a star and cite our paper. + +CogVLM2-Caption: +``` +@article{yang2024cogvideox, + title={CogVideoX: Text-to-Video Diffusion Models with An Expert Transformer}, + author={Yang, Zhuoyi and Teng, Jiayan and Zheng, Wendi and Ding, Ming and Huang, Shiyu and Xu, Jiazheng and Yang, Yuanming and Hong, Wenyi and Zhang, Xiaohan and Feng, Guanyu and others}, + journal={arXiv preprint arXiv:2408.06072}, + year={2024} +} +``` +CogVLM2-Video: +``` +@article{hong2024cogvlm2, + title={CogVLM2: Visual Language Models for Image and Video Understanding}, + author={Hong, Wenyi and Wang, Weihan and Ding, Ming and Yu, Wenmeng and Lv, Qingsong and Wang, Yan and Cheng, Yean and Huang, Shiyu and Ji, Junhui and Xue, Zhao and others}, + journal={arXiv preprint arXiv:2408.16500}, + year={2024} +} +``` diff --git a/tools/caption/README_zh.md b/tools/caption/README_zh.md new file mode 100644 index 0000000..8cc79e1 --- /dev/null +++ b/tools/caption/README_zh.md @@ -0,0 +1,67 @@ +# 视频Caption + +通常，大多数视频数据不带有相应的描述性文本，因此需要将视频数据转换为文本描述，以提供必要的训练数据用于文本到视频模型。 + +## 项目更新 +- 🔥🔥 **News**: ```2024/9/19```: CogVideoX 训练过程中用于将视频数据转换为文本描述的 Caption + 模型 [CogVLM2-Caption](https://huggingface.co/THUDM/cogvlm2-llama3-caption) + 已经开源。欢迎前往下载并使用。 + +## 通过 CogVLM2-Caption 模型生成视频Caption + +🤗 [Hugging Face](https://huggingface.co/THUDM/cogvlm2-llama3-caption) | 🤖 [ModelScope](https://modelscope.cn/models/ZhipuAI/cogvlm2-llama3-caption/) + +CogVLM2-Caption是用于生成CogVideoX模型训练数据的视频caption模型。 + +### 安装依赖 +```shell +pip install -r requirements.txt +``` + +### 运行caption模型 + +```shell +python video_caption.py +``` + +示例： +
+ +
+ +## 通过 CogVLM2-Video 模型生成视频Caption + +[Code](https://github.com/THUDM/CogVLM2/tree/main/video_demo) | 🤗 [Hugging Face](https://huggingface.co/THUDM/cogvlm2-video-llama3-chat) | 🤖 [ModelScope](https://modelscope.cn/models/ZhipuAI/cogvlm2-video-llama3-chat) | 📑 [Blog](https://cogvlm2-video.github.io/) ｜ [💬 Online Demo](http://cogvlm2-online.cogviewai.cn:7868/) + +CogVLM2-Video 是一个多功能的视频理解模型，具备基于时间戳的问题回答能力。用户可以输入诸如 `Describe this video in detail.` 的提示语给模型，以获得详细的视频Caption： + + +
+ +
+ +用户可以使用提供的[代码](https://github.com/THUDM/CogVLM2/tree/main/video_demo)加载模型或配置 RESTful API 来生成视频Caption。 + + +## Citation + +🌟 If you find our work helpful, please leave us a star and cite our paper. + +CogVLM2-Caption: +``` +@article{yang2024cogvideox, + title={CogVideoX: Text-to-Video Diffusion Models with An Expert Transformer}, + author={Yang, Zhuoyi and Teng, Jiayan and Zheng, Wendi and Ding, Ming and Huang, Shiyu and Xu, Jiazheng and Yang, Yuanming and Hong, Wenyi and Zhang, Xiaohan and Feng, Guanyu and others}, + journal={arXiv preprint arXiv:2408.06072}, + year={2024} +} +``` +CogVLM2-Video: +``` +@article{hong2024cogvlm2, + title={CogVLM2: Visual Language Models for Image and Video Understanding}, + author={Hong, Wenyi and Wang, Weihan and Ding, Ming and Yu, Wenmeng and Lv, Qingsong and Wang, Yan and Cheng, Yean and Huang, Shiyu and Ji, Junhui and Xue, Zhao and others}, + journal={arXiv preprint arXiv:2408.16500}, + year={2024} +} +``` diff --git a/tools/caption/assests/CogVLM2-Caption-example.png b/tools/caption/assests/CogVLM2-Caption-example.png new file mode 100644 index 0000000..8612bfe Binary files /dev/null and b/tools/caption/assests/CogVLM2-Caption-example.png differ diff --git a/tools/caption/assests/cogvlm2-video-example.png b/tools/caption/assests/cogvlm2-video-example.png new file mode 100644 index 0000000..b6c3c7d Binary files /dev/null and b/tools/caption/assests/cogvlm2-video-example.png differ diff --git a/tools/caption/requirements.txt b/tools/caption/requirements.txt new file mode 100644 index 0000000..1c909e4 --- /dev/null +++ b/tools/caption/requirements.txt @@ -0,0 +1,23 @@ +decord>=0.6.0 +#根据https://download.pytorch.org/whl/torch/，python版本为[3.8,3.11] +torch==2.1.0 +torchvision== 0.16.0 +pytorchvideo==0.1.5 +xformers +transformers==4.42.4 +#git+https://github.com/huggingface/transformers.git +huggingface-hub>=0.23.0 +pillow +chainlit>=1.0 +pydantic>=2.7.1 +timm>=0.9.16 +openai>=1.30.1 +loguru>=0.7.2 +pydantic>=2.7.1 +einops +sse-starlette>=2.1.0 +flask +gunicorn +gevent +requests +gradio diff --git a/tools/caption/video_caption.py b/tools/caption/video_caption.py new file mode 100644 index 0000000..7b2bd71 --- /dev/null +++ b/tools/caption/video_caption.py @@ -0,0 +1,112 @@ +import io + +import argparse +import numpy as np +import torch +from decord import cpu, VideoReader, bridge +from transformers import AutoModelForCausalLM, AutoTokenizer + +MODEL_PATH = "THUDM/cogvlm2-llama3-caption" + +DEVICE = 'cuda' if torch.cuda.is_available() else 'cpu' +TORCH_TYPE = ( + torch.bfloat16 + if torch.cuda.is_available() and torch.cuda.get_device_capability()[0] >= 8 + else torch.float16 +) + +parser = argparse.ArgumentParser(description="CogVLM2-Video CLI Demo") +parser.add_argument( + '--quant', type=int, choices=[4, 8], help='Enable 4-bit or 8-bit precision loading', default=0 +) +args = parser.parse_args([]) + + +def load_video(video_data, strategy='chat'): + bridge.set_bridge('torch') + mp4_stream = video_data + num_frames = 24 + decord_vr = VideoReader(io.BytesIO(mp4_stream), ctx=cpu(0)) + + frame_id_list = None + total_frames = len(decord_vr) + if strategy == 'base': + clip_end_sec = 60 + clip_start_sec = 0 + start_frame = int(clip_start_sec * decord_vr.get_avg_fps()) + end_frame = ( + min(total_frames, int(clip_end_sec * decord_vr.get_avg_fps())) + if clip_end_sec is not None + else total_frames + ) + frame_id_list = np.linspace(start_frame, end_frame - 1, num_frames, dtype=int) + elif strategy == 'chat': + timestamps = decord_vr.get_frame_timestamp(np.arange(total_frames)) + timestamps = [i[0] for i in timestamps] + max_second = round(max(timestamps)) + 1 + frame_id_list = [] + for second in range(max_second): + closest_num = min(timestamps, key=lambda x: abs(x - second)) + index = timestamps.index(closest_num) + frame_id_list.append(index) + if len(frame_id_list) >= num_frames: + break + + video_data = decord_vr.get_batch(frame_id_list) + video_data = video_data.permute(3, 0, 1, 2) + return video_data + + +tokenizer = AutoTokenizer.from_pretrained( + MODEL_PATH, + trust_remote_code=True, +) + +model = ( + AutoModelForCausalLM.from_pretrained(MODEL_PATH, torch_dtype=TORCH_TYPE, trust_remote_code=True) + .eval() + .to(DEVICE) +) + + +def predict(prompt, video_data, temperature): + strategy = 'chat' + + video = load_video(video_data, strategy=strategy) + + history = [] + query = prompt + inputs = model.build_conversation_input_ids( + tokenizer=tokenizer, query=query, images=[video], history=history, template_version=strategy + ) + inputs = { + 'input_ids': inputs['input_ids'].unsqueeze(0).to('cuda'), + 'token_type_ids': inputs['token_type_ids'].unsqueeze(0).to('cuda'), + 'attention_mask': inputs['attention_mask'].unsqueeze(0).to('cuda'), + 'images': [[inputs['images'][0].to('cuda').to(TORCH_TYPE)]], + } + gen_kwargs = { + "max_new_tokens": 2048, + "pad_token_id": 128002, + "top_k": 1, + "do_sample": False, + "top_p": 0.1, + "temperature": temperature, + } + with torch.no_grad(): + outputs = model.generate(**inputs, **gen_kwargs) + outputs = outputs[:, inputs['input_ids'].shape[1] :] + response = tokenizer.decode(outputs[0], skip_special_tokens=True) + return response + + +def test(): + prompt = "Please describe this video in detail." + temperature = 0.1 + video_data = open('test.mp4', 'rb').read() + response = predict(prompt, video_data, temperature) + print(response) + + +if __name__ == '__main__': + test() diff --git a/tools/convert_weight_deepspeed2hf.py b/tools/convert_weight_deepspeed2hf.py new file mode 100644 index 0000000..7d228e1 --- /dev/null +++ b/tools/convert_weight_deepspeed2hf.py @@ -0,0 +1,928 @@ +#!/usr/bin/env python + +# Copyright (c) Microsoft Corporation. +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team + +# This script extracts fp32 consolidated weights from a zero 1, 2 and 3 DeepSpeed checkpoints. It gets +# copied into the top level checkpoint dir, so the user can easily do the conversion at any point in +# the future. Once extracted, the weights don't require DeepSpeed and can be used in any +# application. +# +# example: +# python zero_to_fp32.py . output_dir/ +# or +# python zero_to_fp32.py . output_dir/ --safe_serialization + +import argparse +import torch +import glob +import math +import os +import re +import gc +import json +import numpy as np +from tqdm import tqdm +from collections import OrderedDict +from dataclasses import dataclass + +# while this script doesn't use deepspeed to recover data, since the checkpoints are pickled with +# DeepSpeed data structures it has to be available in the current python environment. +from deepspeed.utils import logger +from deepspeed.checkpoint.constants import ( + DS_VERSION, + OPTIMIZER_STATE_DICT, + SINGLE_PARTITION_OF_FP32_GROUPS, + FP32_FLAT_GROUPS, + ZERO_STAGE, + PARTITION_COUNT, + PARAM_SHAPES, + BUFFER_NAMES, + FROZEN_PARAM_SHAPES, + FROZEN_PARAM_FRAGMENTS, +) + + +@dataclass +class zero_model_state: + buffers: dict() + param_shapes: dict() + shared_params: list + ds_version: int + frozen_param_shapes: dict() + frozen_param_fragments: dict() + + +debug = 0 + +# load to cpu +device = torch.device('cpu') + + +def atoi(text): + return int(text) if text.isdigit() else text + + +def natural_keys(text): + ''' + alist.sort(key=natural_keys) sorts in human order + http://nedbatchelder.com/blog/200712/human_sorting.html + (See Toothy's implementation in the comments) + ''' + return [atoi(c) for c in re.split(r'(\d+)', text)] + + +def get_model_state_file(checkpoint_dir, zero_stage): + if not os.path.isdir(checkpoint_dir): + raise FileNotFoundError(f"Directory '{checkpoint_dir}' doesn't exist") + + # there should be only one file + if zero_stage <= 2: + file = os.path.join(checkpoint_dir, "mp_rank_00_model_states.pt") + elif zero_stage == 3: + file = os.path.join(checkpoint_dir, "zero_pp_rank_0_mp_rank_00_model_states.pt") + + if not os.path.exists(file): + raise FileNotFoundError(f"can't find model states file at '{file}'") + + return file + + +def get_checkpoint_files(checkpoint_dir, glob_pattern): + # XXX: need to test that this simple glob rule works for multi-node setup too + ckpt_files = sorted(glob.glob(os.path.join(checkpoint_dir, glob_pattern)), key=natural_keys) + + if len(ckpt_files) == 0: + raise FileNotFoundError(f"can't find {glob_pattern} files in directory '{checkpoint_dir}'") + + return ckpt_files + + +def get_optim_files(checkpoint_dir): + return get_checkpoint_files(checkpoint_dir, "*_optim_states.pt") + + +def get_model_state_files(checkpoint_dir): + return get_checkpoint_files(checkpoint_dir, "*_model_states.pt") + + +def parse_model_states(files): + zero_model_states = [] + for file in files: + state_dict = torch.load(file, map_location=device, weights_only=False) + + if BUFFER_NAMES not in state_dict: + raise ValueError(f"{file} is not a model state checkpoint") + buffer_names = state_dict[BUFFER_NAMES] + if debug: + print("Found buffers:", buffer_names) + + # recover just the buffers while restoring them to fp32 if they were saved in fp16 + buffers = {k: v.float() for k, v in state_dict["module"].items() if k in buffer_names} + param_shapes = state_dict[PARAM_SHAPES] + + # collect parameters that are included in param_shapes + param_names = [] + for s in param_shapes: + for name in s.keys(): + param_names.append(name) + + # update with frozen parameters + frozen_param_shapes = state_dict.get(FROZEN_PARAM_SHAPES, None) + if frozen_param_shapes is not None: + if debug: + print(f"Found frozen_param_shapes: {frozen_param_shapes}") + param_names += list(frozen_param_shapes.keys()) + + # handle shared params + shared_params = [[k, v] for k, v in state_dict["shared_params"].items()] + + ds_version = state_dict.get(DS_VERSION, None) + + frozen_param_fragments = state_dict.get(FROZEN_PARAM_FRAGMENTS, None) + + z_model_state = zero_model_state( + buffers=buffers, + param_shapes=param_shapes, + shared_params=shared_params, + ds_version=ds_version, + frozen_param_shapes=frozen_param_shapes, + frozen_param_fragments=frozen_param_fragments, + ) + zero_model_states.append(z_model_state) + + return zero_model_states + + +def parse_optim_states(files, ds_checkpoint_dir): + total_files = len(files) + state_dicts = [] + for f in tqdm(files, desc='Loading checkpoint shards'): + state_dict = torch.load(f, map_location=device, mmap=True, weights_only=False) + # immediately discard the potentially huge 2 optimizer states as we only care for fp32 master weights + # and also handle the case where it was already removed by another helper script + state_dict["optimizer_state_dict"].pop("optimizer_state_dict", None) + state_dicts.append(state_dict) + + if ZERO_STAGE not in state_dicts[0][OPTIMIZER_STATE_DICT]: + raise ValueError(f"{files[0]} is not a zero checkpoint") + zero_stage = state_dicts[0][OPTIMIZER_STATE_DICT][ZERO_STAGE] + world_size = state_dicts[0][OPTIMIZER_STATE_DICT][PARTITION_COUNT] + + # For ZeRO-2 each param group can have different partition_count as data parallelism for expert + # parameters can be different from data parallelism for non-expert parameters. So we can just + # use the max of the partition_count to get the dp world_size. + + if type(world_size) is list: + world_size = max(world_size) + + if world_size != total_files: + raise ValueError( + f"Expected {world_size} of '*_optim_states.pt' under '{ds_checkpoint_dir}' but found {total_files} files. " + "Possibly due to an overwrite of an old checkpoint, or a checkpoint didn't get saved by one or more processes." + ) + + # the groups are named differently in each stage + if zero_stage <= 2: + fp32_groups_key = SINGLE_PARTITION_OF_FP32_GROUPS + elif zero_stage == 3: + fp32_groups_key = FP32_FLAT_GROUPS + else: + raise ValueError(f"unknown zero stage {zero_stage}") + + fp32_flat_groups = [ + state_dicts[i][OPTIMIZER_STATE_DICT][fp32_groups_key] for i in range(len(state_dicts)) + ] + return zero_stage, world_size, fp32_flat_groups + + +def _get_fp32_state_dict_from_zero_checkpoint(ds_checkpoint_dir, exclude_frozen_parameters): + """ + Returns fp32 state_dict reconstructed from ds checkpoint + + Args: + - ``ds_checkpoint_dir``: path to the deepspeed checkpoint folder (where the optimizer files are) + + """ + print(f"Processing zero checkpoint '{ds_checkpoint_dir}'") + + optim_files = get_optim_files(ds_checkpoint_dir) + zero_stage, world_size, fp32_flat_groups = parse_optim_states(optim_files, ds_checkpoint_dir) + print(f"Detected checkpoint of type zero stage {zero_stage}, world_size: {world_size}") + + model_files = get_model_state_files(ds_checkpoint_dir) + + zero_model_states = parse_model_states(model_files) + print(f'Parsing checkpoint created by deepspeed=={zero_model_states[0].ds_version}') + + if zero_stage <= 2: + return _get_fp32_state_dict_from_zero2_checkpoint( + world_size, fp32_flat_groups, zero_model_states, exclude_frozen_parameters + ) + elif zero_stage == 3: + return _get_fp32_state_dict_from_zero3_checkpoint( + world_size, fp32_flat_groups, zero_model_states, exclude_frozen_parameters + ) + + +def _zero2_merge_frozen_params(state_dict, zero_model_states): + if ( + zero_model_states[0].frozen_param_shapes is None + or len(zero_model_states[0].frozen_param_shapes) == 0 + ): + return + + frozen_param_shapes = zero_model_states[0].frozen_param_shapes + frozen_param_fragments = zero_model_states[0].frozen_param_fragments + + if debug: + num_elem = sum(s.numel() for s in frozen_param_shapes.values()) + print(f'rank 0: {FROZEN_PARAM_SHAPES}.numel = {num_elem}') + + wanted_params = len(frozen_param_shapes) + wanted_numel = sum(s.numel() for s in frozen_param_shapes.values()) + avail_numel = sum([p.numel() for p in frozen_param_fragments.values()]) + print(f'Frozen params: Have {avail_numel} numels to process.') + print(f'Frozen params: Need {wanted_numel} numels in {wanted_params} params') + + total_params = 0 + total_numel = 0 + for name, shape in frozen_param_shapes.items(): + total_params += 1 + unpartitioned_numel = shape.numel() + total_numel += unpartitioned_numel + + state_dict[name] = frozen_param_fragments[name] + + if debug: + print(f"{name} full shape: {shape} unpartitioned numel {unpartitioned_numel} ") + + print(f"Reconstructed Frozen fp32 state dict with {total_params} params {total_numel} elements") + + +def _has_callable(obj, fn): + attr = getattr(obj, fn, None) + return callable(attr) + + +def _zero2_merge_trainable_params(state_dict, world_size, fp32_flat_groups, zero_model_states): + param_shapes = zero_model_states[0].param_shapes + + # Reconstruction protocol: + # + # XXX: document this + + if debug: + for i in range(world_size): + for j in range(len(fp32_flat_groups[0])): + print(f"{FP32_FLAT_GROUPS}[{i}][{j}].shape={fp32_flat_groups[i][j].shape}") + + # XXX: memory usage doubles here (zero2) + num_param_groups = len(fp32_flat_groups[0]) + merged_single_partition_of_fp32_groups = [] + for i in range(num_param_groups): + merged_partitions = [sd[i] for sd in fp32_flat_groups] + full_single_fp32_vector = torch.cat(merged_partitions, 0) + merged_single_partition_of_fp32_groups.append(full_single_fp32_vector) + avail_numel = sum( + [ + full_single_fp32_vector.numel() + for full_single_fp32_vector in merged_single_partition_of_fp32_groups + ] + ) + + if debug: + wanted_params = sum([len(shapes) for shapes in param_shapes]) + wanted_numel = sum( + [sum(shape.numel() for shape in shapes.values()) for shapes in param_shapes] + ) + # not asserting if there is a mismatch due to possible padding + print(f"Have {avail_numel} numels to process.") + print(f"Need {wanted_numel} numels in {wanted_params} params.") + + # params + # XXX: for huge models that can't fit into the host's RAM we will have to recode this to support + # out-of-core computing solution + total_numel = 0 + total_params = 0 + for shapes, full_single_fp32_vector in zip( + param_shapes, merged_single_partition_of_fp32_groups + ): + offset = 0 + avail_numel = full_single_fp32_vector.numel() + for name, shape in shapes.items(): + unpartitioned_numel = ( + shape.numel() if _has_callable(shape, 'numel') else math.prod(shape) + ) + total_numel += unpartitioned_numel + total_params += 1 + + if debug: + print(f"{name} full shape: {shape} unpartitioned numel {unpartitioned_numel} ") + state_dict[name] = full_single_fp32_vector.narrow(0, offset, unpartitioned_numel).view( + shape + ) + offset += unpartitioned_numel + + # Z2 started to align to 2*world_size to improve nccl performance. Therefore both offset and + # avail_numel can differ by anywhere between 0..2*world_size. Due to two unrelated complex + # paddings performed in the code it's almost impossible to predict the exact numbers w/o the + # live optimizer object, so we are checking that the numbers are within the right range + align_to = 2 * world_size + + def zero2_align(x): + return align_to * math.ceil(x / align_to) + + if debug: + print(f"original offset={offset}, avail_numel={avail_numel}") + + offset = zero2_align(offset) + avail_numel = zero2_align(avail_numel) + + if debug: + print(f"aligned offset={offset}, avail_numel={avail_numel}") + + # Sanity check + if offset != avail_numel: + raise ValueError(f"consumed {offset} numels out of {avail_numel} - something is wrong") + + print(f"Reconstructed fp32 state dict with {total_params} params {total_numel} elements") + + +def _get_fp32_state_dict_from_zero2_checkpoint( + world_size, fp32_flat_groups, zero_model_states, exclude_frozen_parameters +): + state_dict = OrderedDict() + + # buffers + buffers = zero_model_states[0].buffers + state_dict.update(buffers) + if debug: + print(f"added {len(buffers)} buffers") + + if not exclude_frozen_parameters: + _zero2_merge_frozen_params(state_dict, zero_model_states) + + _zero2_merge_trainable_params(state_dict, world_size, fp32_flat_groups, zero_model_states) + + # recover shared parameters + for pair in zero_model_states[0].shared_params: + if pair[1] in state_dict: + state_dict[pair[0]] = state_dict[pair[1]] + + return state_dict + + +def zero3_partitioned_param_info(unpartitioned_numel, world_size): + remainder = unpartitioned_numel % world_size + padding_numel = (world_size - remainder) if remainder else 0 + partitioned_numel = math.ceil(unpartitioned_numel / world_size) + return partitioned_numel, padding_numel + + +def _zero3_merge_frozen_params(state_dict, world_size, zero_model_states): + if ( + zero_model_states[0].frozen_param_shapes is None + or len(zero_model_states[0].frozen_param_shapes) == 0 + ): + return + + if debug: + for i in range(world_size): + num_elem = sum(s.numel() for s in zero_model_states[i].frozen_param_fragments.values()) + print(f'rank {i}: {FROZEN_PARAM_SHAPES}.numel = {num_elem}') + + frozen_param_shapes = zero_model_states[0].frozen_param_shapes + wanted_params = len(frozen_param_shapes) + wanted_numel = sum(s.numel() for s in frozen_param_shapes.values()) + avail_numel = ( + sum([p.numel() for p in zero_model_states[0].frozen_param_fragments.values()]) + * world_size + ) + print(f'Frozen params: Have {avail_numel} numels to process.') + print(f'Frozen params: Need {wanted_numel} numels in {wanted_params} params') + + total_params = 0 + total_numel = 0 + for name, shape in zero_model_states[0].frozen_param_shapes.items(): + total_params += 1 + unpartitioned_numel = shape.numel() + total_numel += unpartitioned_numel + + param_frags = tuple( + model_state.frozen_param_fragments[name] for model_state in zero_model_states + ) + state_dict[name] = torch.cat(param_frags, 0).narrow(0, 0, unpartitioned_numel).view(shape) + + partitioned_numel, partitioned_padding_numel = zero3_partitioned_param_info( + unpartitioned_numel, world_size + ) + + if debug: + print( + f"Frozen params: {total_params} {name} full shape: {shape} partition0 numel={partitioned_numel} partitioned_padding_numel={partitioned_padding_numel}" + ) + + print(f"Reconstructed Frozen fp32 state dict with {total_params} params {total_numel} elements") + + +class GatheredTensor: + """ + A pseudo tensor that collects partitioned weights. + It is more memory efficient when there are multiple groups. + """ + + def __init__(self, flat_groups, flat_groups_offset, offset, partitioned_numel, shape): + self.flat_groups = flat_groups + self.flat_groups_offset = flat_groups_offset + self.offset = offset + self.partitioned_numel = partitioned_numel + self.shape = shape + self.dtype = self.flat_groups[0][0].dtype + + def contiguous(self): + """ + Merge partitioned weights from flat_groups into a single tensor. + """ + end_idx = self.offset + self.partitioned_numel + world_size = len(self.flat_groups) + pad_flat_param_chunks = [] + + for rank_i in range(world_size): + # for each rank, we need to collect weights from related group/groups + flat_groups_at_rank_i = self.flat_groups[rank_i] + start_group_id = None + end_group_id = None + for group_id in range(len(self.flat_groups_offset)): + if ( + self.flat_groups_offset[group_id] + <= self.offset + < self.flat_groups_offset[group_id + 1] + ): + start_group_id = group_id + if ( + self.flat_groups_offset[group_id] + < end_idx + <= self.flat_groups_offset[group_id + 1] + ): + end_group_id = group_id + break + # collect weights from related group/groups + for group_id in range(start_group_id, end_group_id + 1): + flat_tensor = flat_groups_at_rank_i[group_id] + start_offset = self.offset - self.flat_groups_offset[group_id] + end_offset = ( + min(end_idx, self.flat_groups_offset[group_id + 1]) + - self.flat_groups_offset[group_id] + ) + pad_flat_param_chunks.append(flat_tensor[start_offset:end_offset]) + + # collect weights from all ranks + pad_flat_param = torch.cat(pad_flat_param_chunks, dim=0) + param = pad_flat_param[: self.shape.numel()].view(self.shape).contiguous() + return param + + +def _zero3_merge_trainable_params(state_dict, world_size, fp32_flat_groups, zero_model_states): + param_shapes = zero_model_states[0].param_shapes + avail_numel = sum([flat_group.numel() for flat_group in fp32_flat_groups[0]]) * world_size + + # Reconstruction protocol: For zero3 we need to zip the partitions together at boundary of each + # param, re-consolidating each param, while dealing with padding if any + + # merge list of dicts, preserving order + param_shapes = {k: v for d in param_shapes for k, v in d.items()} + + if debug: + for i in range(world_size): + print(f"{FP32_FLAT_GROUPS}[{i}].shape={fp32_flat_groups[i].shape}") + + wanted_params = len(param_shapes) + wanted_numel = sum(shape.numel() for shape in param_shapes.values()) + # not asserting if there is a mismatch due to possible padding + avail_numel = fp32_flat_groups[0].numel() * world_size + print(f"Trainable params: Have {avail_numel} numels to process.") + print(f"Trainable params: Need {wanted_numel} numels in {wanted_params} params.") + + # params + # XXX: for huge models that can't fit into the host's RAM we will have to recode this to support + # out-of-core computing solution + offset = 0 + total_numel = 0 + total_params = 0 + flat_groups_offset = [0] + list( + np.cumsum([flat_tensor.numel() for flat_tensor in fp32_flat_groups[0]]) + ) + for name, shape in tqdm(param_shapes.items(), desc='Gathering sharded weights'): + unpartitioned_numel = shape.numel() + total_numel += unpartitioned_numel + total_params += 1 + partitioned_numel, partitioned_padding_numel = zero3_partitioned_param_info( + unpartitioned_numel, world_size + ) + + if debug: + print( + f"Trainable params: {total_params} {name} full shape: {shape} partition0 numel={partitioned_numel} partitioned_padding_numel={partitioned_padding_numel}" + ) + + # memory efficient tensor + tensor = GatheredTensor( + fp32_flat_groups, flat_groups_offset, offset, partitioned_numel, shape + ) + state_dict[name] = tensor + offset += partitioned_numel + + offset *= world_size + + # Sanity check + if offset != avail_numel: + raise ValueError(f"consumed {offset} numels out of {avail_numel} - something is wrong") + + print( + f"Reconstructed Trainable fp32 state dict with {total_params} params {total_numel} elements" + ) + + +def _get_fp32_state_dict_from_zero3_checkpoint( + world_size, fp32_flat_groups, zero_model_states, exclude_frozen_parameters +): + state_dict = OrderedDict() + + # buffers + buffers = zero_model_states[0].buffers + state_dict.update(buffers) + if debug: + print(f"added {len(buffers)} buffers") + + if not exclude_frozen_parameters: + _zero3_merge_frozen_params(state_dict, world_size, zero_model_states) + + _zero3_merge_trainable_params(state_dict, world_size, fp32_flat_groups, zero_model_states) + + # recover shared parameters + for pair in zero_model_states[0].shared_params: + if pair[1] in state_dict: + state_dict[pair[0]] = state_dict[pair[1]] + + return state_dict + + +def to_torch_tensor(state_dict, return_empty_tensor=False): + """ + Convert state_dict of GatheredTensor to torch tensor + """ + torch_state_dict = {} + converted_tensors = {} + for name, tensor in state_dict.items(): + tensor_id = id(tensor) + if tensor_id in converted_tensors: # shared tensors + shared_tensor = torch_state_dict[converted_tensors[tensor_id]] + torch_state_dict[name] = shared_tensor + else: + converted_tensors[tensor_id] = name + if return_empty_tensor: + torch_state_dict[name] = torch.empty(tensor.shape, dtype=tensor.dtype) + else: + torch_state_dict[name] = tensor.contiguous() + return torch_state_dict + + +def get_fp32_state_dict_from_zero_checkpoint( + checkpoint_dir, tag=None, exclude_frozen_parameters=False, lazy_mode=False +): + """ + Convert ZeRO 2 or 3 checkpoint into a single fp32 consolidated state_dict that can be loaded with + ``load_state_dict()`` and used for training without DeepSpeed or shared with others, for example + via a model hub. + + Args: + - ``checkpoint_dir``: path to the desired checkpoint folder + - ``tag``: checkpoint tag used as a unique identifier for checkpoint. If not provided will attempt to load tag in 'latest' file. e.g., ``global_step14`` + - ``exclude_frozen_parameters``: exclude frozen parameters + - ``lazy_mode``: get state_dict in lazy mode. It returns a dict of pesduo tensor instead of torch tensor, which is more memory efficient. + Convert the pesduo tensor to torch tensor by ``.contiguous()`` + + Returns: + - pytorch ``state_dict`` + + A typical usage might be :: + + from deepspeed.utils.zero_to_fp32 import get_fp32_state_dict_from_zero_checkpoint + # do the training and checkpoint saving + state_dict = get_fp32_state_dict_from_zero_checkpoint(checkpoint_dir) # already on cpu + model = model.cpu() # move to cpu + model.load_state_dict(state_dict) + # submit to model hub or save the model to share with others + + In this example the ``model`` will no longer be usable in the deepspeed context of the same + application. i.e. you will need to re-initialize the deepspeed engine, since + ``model.load_state_dict(state_dict)`` will remove all the deepspeed magic from it. + + If you want it all done for you, use ``load_state_dict_from_zero_checkpoint`` instead. + + Note: the above usage may not work if your application doesn't have sufficient free CPU memory. + You may need to use the offline approach using the ``zero_to_fp32.py`` script that is saved with + the checkpoint. Or you can load state_dict in lazy mode :: + + from deepspeed.utils.zero_to_fp32 import get_fp32_state_dict_from_zero_checkpoint + state_dict = get_fp32_state_dict_from_zero_checkpoint(checkpoint_dir, lazy_mode=True) # not on cpu + for name, lazy_tensor in state_dict.item(): + tensor = lazy_tensor.contiguous() # to cpu + print(name, tensor) + # del tensor to release memory if it no longer in use + """ + if tag is None: + latest_path = os.path.join(checkpoint_dir, 'latest') + if os.path.isfile(latest_path): + with open(latest_path, 'r') as fd: + tag = fd.read().strip() + else: + raise ValueError(f"Unable to find 'latest' file at {latest_path}") + + ds_checkpoint_dir = os.path.join(checkpoint_dir, tag) + + if not os.path.isdir(ds_checkpoint_dir): + raise FileNotFoundError(f"Directory '{ds_checkpoint_dir}' doesn't exist") + + state_dict = _get_fp32_state_dict_from_zero_checkpoint( + ds_checkpoint_dir, exclude_frozen_parameters + ) + if lazy_mode: + return state_dict + else: + return to_torch_tensor(state_dict) + + +def convert_zero_checkpoint_to_fp32_state_dict( + checkpoint_dir, + output_dir, + max_shard_size="5GB", + safe_serialization=False, + tag=None, + exclude_frozen_parameters=False, +): + """ + Convert ZeRO 2 or 3 checkpoint into a single fp32 consolidated ``state_dict`` file that can be + loaded with ``torch.load(file)`` + ``load_state_dict()`` and used for training without DeepSpeed. + + Args: + - ``checkpoint_dir``: path to the desired checkpoint folder. (one that contains the tag-folder, like ``global_step14``) + - ``output_dir``: directory to the pytorch fp32 state_dict output files + - ``max_shard_size``: the maximum size for a checkpoint before being sharded, default value is 5GB + - ``safe_serialization``: whether to save the model using `safetensors` or the traditional PyTorch way (that uses `pickle`). + - ``tag``: checkpoint tag used as a unique identifier for checkpoint. If not provided will attempt to load tag in the file named ``latest`` in the checkpoint folder, e.g., ``global_step14`` + - ``exclude_frozen_parameters``: exclude frozen parameters + """ + + # Dependency pre-check + if safe_serialization: + try: + from safetensors.torch import save_file + except ImportError: + print('If you want to use `safe_serialization`, please `pip install safetensors`') + raise + if max_shard_size is not None: + try: + from huggingface_hub import split_torch_state_dict_into_shards + except ImportError: + print('If you want to use `max_shard_size`, please `pip install huggingface_hub`') + raise + + # Convert zero checkpoint to state_dict + state_dict = get_fp32_state_dict_from_zero_checkpoint( + checkpoint_dir, tag, exclude_frozen_parameters, lazy_mode=True + ) + + # Shard the model if it is too big. + weights_name = "model.safetensors" if safe_serialization else "pytorch_model.bin" + if max_shard_size is not None: + filename_pattern = weights_name.replace(".bin", "{suffix}.bin").replace( + ".safetensors", "{suffix}.safetensors" + ) + # an memory-efficient approach for sharding + empty_state_dict = to_torch_tensor(state_dict, return_empty_tensor=True) + state_dict_split = split_torch_state_dict_into_shards( + empty_state_dict, filename_pattern=filename_pattern, max_shard_size=max_shard_size + ) + else: + from collections import namedtuple + + StateDictSplit = namedtuple("StateDictSplit", ["is_sharded", "filename_to_tensors"]) + state_dict_split = StateDictSplit( + is_sharded=False, filename_to_tensors={weights_name: list(state_dict.keys())} + ) + + # Save the model by shard + os.makedirs(output_dir, exist_ok=True) + filename_to_tensors = state_dict_split.filename_to_tensors.items() + for shard_file, tensors in tqdm(filename_to_tensors, desc="Saving checkpoint shards"): + shard_state_dict = {tensor_name: state_dict[tensor_name] for tensor_name in tensors} + shard_state_dict = to_torch_tensor(shard_state_dict) + output_path = os.path.join(output_dir, shard_file) + if safe_serialization: + save_file(shard_state_dict, output_path, metadata={"format": "pt"}) + else: + torch.save(shard_state_dict, output_path) + # release the memory of current shard + for tensor_name in list(shard_state_dict.keys()): + del state_dict[tensor_name] + del shard_state_dict[tensor_name] + del shard_state_dict + gc.collect() + + # Save index if sharded + if state_dict_split.is_sharded: + index = { + "metadata": state_dict_split.metadata, + "weight_map": state_dict_split.tensor_to_filename, + } + save_index_file = ( + "model.safetensors.index.json" if safe_serialization else "pytorch_model.bin.index.json" + ) + save_index_file = os.path.join(output_dir, save_index_file) + with open(save_index_file, "w", encoding="utf-8") as f: + content = json.dumps(index, indent=2, sort_keys=True) + "\n" + f.write(content) + + +def load_state_dict_from_zero_checkpoint(model, checkpoint_dir, tag=None): + """ + 1. Put the provided model to cpu + 2. Convert ZeRO 2 or 3 checkpoint into a single fp32 consolidated ``state_dict`` + 3. Load it into the provided model + + Args: + - ``model``: the model object to update + - ``checkpoint_dir``: path to the desired checkpoint folder. (one that contains the tag-folder, like ``global_step14``) + - ``tag``: checkpoint tag used as a unique identifier for checkpoint. If not provided will attempt to load tag in the file named ``latest`` in the checkpoint folder, e.g., ``global_step14`` + + Returns: + - ``model`: modified model + + Make sure you have plenty of CPU memory available before you call this function. If you don't + have enough use the ``zero_to_fp32.py`` utility to do the conversion. You will find it + conveniently placed for you in the checkpoint folder. + + A typical usage might be :: + + from deepspeed.utils.zero_to_fp32 import load_state_dict_from_zero_checkpoint + model = load_state_dict_from_zero_checkpoint(trainer.model, checkpoint_dir) + # submit to model hub or save the model to share with others + + Note, that once this was run, the ``model`` will no longer be usable in the deepspeed context + of the same application. i.e. you will need to re-initialize the deepspeed engine, since + ``model.load_state_dict(state_dict)`` will remove all the deepspeed magic from it. + + """ + logger.info(f"Extracting fp32 weights") + state_dict = get_fp32_state_dict_from_zero_checkpoint(checkpoint_dir, tag) + + logger.info(f"Overwriting model with fp32 weights") + model = model.cpu() + model.load_state_dict(state_dict, strict=False) + + return model + + +def convert_zero_checkpoint_to_bf16_state_dict( + checkpoint_dir, + output_dir, + max_shard_size="5GB", + safe_serialization=True, + tag=None, + exclude_frozen_parameters=False, +): + """ + 将 ZeRO 2 或 ZeRO 3 格式的 DeepSpeed 检查点转换为 BF16，并输出到指定目录下，命名规则为: + - 如果只有一个分片: + diffusion_pytorch_model.safetensors + - 如果分片多于一个: + diffusion_pytorch_model-00001-of-0000X.safetensors + diffusion_pytorch_model-00002-of-0000X.safetensors + ... + diffusion_pytorch_model.safetensors.index.json + """ + + if safe_serialization: + try: + from safetensors.torch import save_file + except ImportError: + raise ImportError("You need `pip install safetensors` to use safetensors.") + if max_shard_size is not None: + try: + from huggingface_hub import split_torch_state_dict_into_shards + except ImportError: + raise ImportError("You need `pip install huggingface_hub` to use the sharding feature.") + + state_dict = get_fp32_state_dict_from_zero_checkpoint( + checkpoint_dir, tag=tag, exclude_frozen_parameters=exclude_frozen_parameters, lazy_mode=True + ) + + state_dict = to_torch_tensor(state_dict, return_empty_tensor=False) + + for key, tensor in state_dict.items(): + state_dict[key] = tensor.to(torch.bfloat16) + + if safe_serialization: + filename_pattern = "diffusion_pytorch_model{suffix}.safetensors" + else: + filename_pattern = "diffusion_pytorch_model{suffix}.bin" + + empty_state_dict = to_torch_tensor(state_dict, return_empty_tensor=True) + state_dict_split = split_torch_state_dict_into_shards( + empty_state_dict, filename_pattern=filename_pattern, max_shard_size=max_shard_size + ) + + os.makedirs(output_dir, exist_ok=True) + + filename_to_tensors = list(state_dict_split.filename_to_tensors.items()) + for shard_file, tensors in tqdm(filename_to_tensors, desc="Saving checkpoint shards"): + shard_state_dict = {t_name: state_dict[t_name] for t_name in tensors} + shard_state_dict = to_torch_tensor(shard_state_dict) + + # Save + output_path = os.path.join(output_dir, shard_file) + if safe_serialization: + save_file(shard_state_dict, output_path, metadata={"format": "pt"}) + else: + torch.save(shard_state_dict, output_path) + for t_name in shard_state_dict.keys(): + del state_dict[t_name] + del shard_state_dict + gc.collect() + + if state_dict_split.is_sharded: + index = { + "metadata": state_dict_split.metadata, + "weight_map": state_dict_split.tensor_to_filename, + } + index_path = os.path.join(output_dir, "diffusion_pytorch_model.safetensors.index.json") + with open(index_path, "w", encoding="utf-8") as f: + f.write(json.dumps(index, indent=2, sort_keys=True) + "\n") + else: + only_filename = list(state_dict_split.filename_to_tensors.keys())[0] + old_path = os.path.join(output_dir, only_filename) + new_path = os.path.join( + output_dir, + "diffusion_pytorch_model.safetensors" + if safe_serialization + else "diffusion_pytorch_model.bin", + ) + if old_path != new_path: + os.rename(old_path, new_path) + + +if __name__ == "__main__": + parser = argparse.ArgumentParser() + parser.add_argument( + "checkpoint_dir", + type=str, + help="path to the desired checkpoint folder, e.g., path/checkpoint-12", + ) + parser.add_argument( + "output_dir", + type=str, + help="directory to the pytorch fp32 state_dict output files" + "(e.g. path/checkpoint-12-output/)", + ) + parser.add_argument( + "--max_shard_size", + type=str, + default="5GB", + help="The maximum size for a checkpoint before being sharded. Checkpoints shard will then be each of size" + "lower than this size. If expressed as a string, needs to be digits followed by a unit (like `5MB`" + "We default it to 5GB in order for models to be able to run easily on free-tier google colab instances" + "without CPU OOM issues.", + ) + parser.add_argument( + "--safe_serialization", + default=False, + action='store_true', + help="Whether to save the model using `safetensors` or the traditional PyTorch way (that uses `pickle`).", + ) + parser.add_argument( + "-t", + "--tag", + type=str, + default=None, + help="checkpoint tag used as a unique identifier for checkpoint. e.g., global_step1", + ) + parser.add_argument( + "--exclude_frozen_parameters", action='store_true', help="exclude frozen parameters" + ) + parser.add_argument("-d", "--debug", action='store_true', help="enable debug") + args = parser.parse_args() + + debug = args.debug + + convert_zero_checkpoint_to_bf16_state_dict( + args.checkpoint_dir, + args.output_dir, + max_shard_size=args.max_shard_size, + safe_serialization=args.safe_serialization, + tag=args.tag, + exclude_frozen_parameters=args.exclude_frozen_parameters, + ) diff --git a/tools/convert_weight_sat2hf.py b/tools/convert_weight_sat2hf.py new file mode 100644 index 0000000..35e7fa7 --- /dev/null +++ b/tools/convert_weight_sat2hf.py @@ -0,0 +1,403 @@ +""" + +The script demonstrates how to convert the weights of the CogVideoX model from SAT to Hugging Face format. +This script supports the conversion of the following models: +- CogVideoX-2B +- CogVideoX-5B, CogVideoX-5B-I2V +- CogVideoX1.1-5B, CogVideoX1.1-5B-I2V + +Original Script: +https://github.com/huggingface/diffusers/blob/main/scripts/convert_cogvideox_to_diffusers.py + +""" + +import argparse +from typing import Any, Dict + +import torch +from transformers import T5EncoderModel, T5Tokenizer + +from diffusers import ( + AutoencoderKLCogVideoX, + CogVideoXDDIMScheduler, + CogVideoXImageToVideoPipeline, + CogVideoXPipeline, + CogVideoXTransformer3DModel, +) + + +def reassign_query_key_value_inplace(key: str, state_dict: Dict[str, Any]): + to_q_key = key.replace("query_key_value", "to_q") + to_k_key = key.replace("query_key_value", "to_k") + to_v_key = key.replace("query_key_value", "to_v") + to_q, to_k, to_v = torch.chunk(state_dict[key], chunks=3, dim=0) + state_dict[to_q_key] = to_q + state_dict[to_k_key] = to_k + state_dict[to_v_key] = to_v + state_dict.pop(key) + + +def reassign_query_key_layernorm_inplace(key: str, state_dict: Dict[str, Any]): + layer_id, weight_or_bias = key.split(".")[-2:] + + if "query" in key: + new_key = f"transformer_blocks.{layer_id}.attn1.norm_q.{weight_or_bias}" + elif "key" in key: + new_key = f"transformer_blocks.{layer_id}.attn1.norm_k.{weight_or_bias}" + + state_dict[new_key] = state_dict.pop(key) + + +def reassign_adaln_norm_inplace(key: str, state_dict: Dict[str, Any]): + layer_id, _, weight_or_bias = key.split(".")[-3:] + + weights_or_biases = state_dict[key].chunk(12, dim=0) + norm1_weights_or_biases = torch.cat(weights_or_biases[0:3] + weights_or_biases[6:9]) + norm2_weights_or_biases = torch.cat(weights_or_biases[3:6] + weights_or_biases[9:12]) + + norm1_key = f"transformer_blocks.{layer_id}.norm1.linear.{weight_or_bias}" + state_dict[norm1_key] = norm1_weights_or_biases + + norm2_key = f"transformer_blocks.{layer_id}.norm2.linear.{weight_or_bias}" + state_dict[norm2_key] = norm2_weights_or_biases + + state_dict.pop(key) + + +def remove_keys_inplace(key: str, state_dict: Dict[str, Any]): + state_dict.pop(key) + + +def replace_up_keys_inplace(key: str, state_dict: Dict[str, Any]): + key_split = key.split(".") + layer_index = int(key_split[2]) + replace_layer_index = 4 - 1 - layer_index + + key_split[1] = "up_blocks" + key_split[2] = str(replace_layer_index) + new_key = ".".join(key_split) + + state_dict[new_key] = state_dict.pop(key) + + +TRANSFORMER_KEYS_RENAME_DICT = { + "transformer.final_layernorm": "norm_final", + "transformer": "transformer_blocks", + "attention": "attn1", + "mlp": "ff.net", + "dense_h_to_4h": "0.proj", + "dense_4h_to_h": "2", + ".layers": "", + "dense": "to_out.0", + "input_layernorm": "norm1.norm", + "post_attn1_layernorm": "norm2.norm", + "time_embed.0": "time_embedding.linear_1", + "time_embed.2": "time_embedding.linear_2", + "ofs_embed.0": "ofs_embedding.linear_1", + "ofs_embed.2": "ofs_embedding.linear_2", + "mixins.patch_embed": "patch_embed", + "mixins.final_layer.norm_final": "norm_out.norm", + "mixins.final_layer.linear": "proj_out", + "mixins.final_layer.adaLN_modulation.1": "norm_out.linear", + "mixins.pos_embed.pos_embedding": "patch_embed.pos_embedding", # Specific to CogVideoX-5b-I2V +} + +TRANSFORMER_SPECIAL_KEYS_REMAP = { + "query_key_value": reassign_query_key_value_inplace, + "query_layernorm_list": reassign_query_key_layernorm_inplace, + "key_layernorm_list": reassign_query_key_layernorm_inplace, + "adaln_layer.adaLN_modulations": reassign_adaln_norm_inplace, + "embed_tokens": remove_keys_inplace, + "freqs_sin": remove_keys_inplace, + "freqs_cos": remove_keys_inplace, + "position_embedding": remove_keys_inplace, +} + +VAE_KEYS_RENAME_DICT = { + "block.": "resnets.", + "down.": "down_blocks.", + "downsample": "downsamplers.0", + "upsample": "upsamplers.0", + "nin_shortcut": "conv_shortcut", + "encoder.mid.block_1": "encoder.mid_block.resnets.0", + "encoder.mid.block_2": "encoder.mid_block.resnets.1", + "decoder.mid.block_1": "decoder.mid_block.resnets.0", + "decoder.mid.block_2": "decoder.mid_block.resnets.1", +} + +VAE_SPECIAL_KEYS_REMAP = { + "loss": remove_keys_inplace, + "up.": replace_up_keys_inplace, +} + +TOKENIZER_MAX_LENGTH = 226 + + +def get_state_dict(saved_dict: Dict[str, Any]) -> Dict[str, Any]: + state_dict = saved_dict + if "model" in saved_dict.keys(): + state_dict = state_dict["model"] + if "module" in saved_dict.keys(): + state_dict = state_dict["module"] + if "state_dict" in saved_dict.keys(): + state_dict = state_dict["state_dict"] + return state_dict + + +def update_state_dict_inplace( + state_dict: Dict[str, Any], old_key: str, new_key: str +) -> Dict[str, Any]: + state_dict[new_key] = state_dict.pop(old_key) + + +def convert_transformer( + ckpt_path: str, + num_layers: int, + num_attention_heads: int, + use_rotary_positional_embeddings: bool, + i2v: bool, + dtype: torch.dtype, + init_kwargs: Dict[str, Any], +): + PREFIX_KEY = "model.diffusion_model." + + original_state_dict = get_state_dict(torch.load(ckpt_path, map_location="cpu", mmap=True)) + transformer = CogVideoXTransformer3DModel( + in_channels=32 if i2v else 16, + num_layers=num_layers, + num_attention_heads=num_attention_heads, + use_rotary_positional_embeddings=use_rotary_positional_embeddings, + ofs_embed_dim=512 + if (i2v and init_kwargs["patch_size_t"] is not None) + else None, # CogVideoX1.5-5B-I2V + use_learned_positional_embeddings=i2v + and init_kwargs["patch_size_t"] is None, # CogVideoX-5B-I2V + **init_kwargs, + ).to(dtype=dtype) + + for key in list(original_state_dict.keys()): + new_key = key[len(PREFIX_KEY) :] + for replace_key, rename_key in TRANSFORMER_KEYS_RENAME_DICT.items(): + new_key = new_key.replace(replace_key, rename_key) + update_state_dict_inplace(original_state_dict, key, new_key) + + for key in list(original_state_dict.keys()): + for special_key, handler_fn_inplace in TRANSFORMER_SPECIAL_KEYS_REMAP.items(): + if special_key not in key: + continue + handler_fn_inplace(key, original_state_dict) + + transformer.load_state_dict(original_state_dict, strict=True) + return transformer + + +def convert_vae(ckpt_path: str, scaling_factor: float, version: str, dtype: torch.dtype): + init_kwargs = {"scaling_factor": scaling_factor} + if version == "1.5": + init_kwargs.update({"invert_scale_latents": True}) + + original_state_dict = get_state_dict(torch.load(ckpt_path, map_location="cpu", mmap=True)) + vae = AutoencoderKLCogVideoX(**init_kwargs).to(dtype=dtype) + + for key in list(original_state_dict.keys()): + new_key = key[:] + for replace_key, rename_key in VAE_KEYS_RENAME_DICT.items(): + new_key = new_key.replace(replace_key, rename_key) + update_state_dict_inplace(original_state_dict, key, new_key) + + for key in list(original_state_dict.keys()): + for special_key, handler_fn_inplace in VAE_SPECIAL_KEYS_REMAP.items(): + if special_key not in key: + continue + handler_fn_inplace(key, original_state_dict) + + vae.load_state_dict(original_state_dict, strict=True) + return vae + + +def get_transformer_init_kwargs(version: str): + if version == "1.0": + vae_scale_factor_spatial = 8 + init_kwargs = { + "patch_size": 2, + "patch_size_t": None, + "patch_bias": True, + "sample_height": 480 // vae_scale_factor_spatial, + "sample_width": 720 // vae_scale_factor_spatial, + "sample_frames": 49, + } + + elif version == "1.5": + vae_scale_factor_spatial = 8 + init_kwargs = { + "patch_size": 2, + "patch_size_t": 2, + "patch_bias": False, + "sample_height": 768 // vae_scale_factor_spatial, + "sample_width": 1360 // vae_scale_factor_spatial, + "sample_frames": 81, + } + else: + raise ValueError("Unsupported version of CogVideoX.") + + return init_kwargs + + +def get_args(): + parser = argparse.ArgumentParser() + parser.add_argument( + "--transformer_ckpt_path", + type=str, + default=None, + help="Path to original transformer checkpoint", + ) + parser.add_argument( + "--vae_ckpt_path", type=str, default=None, help="Path to original vae checkpoint" + ) + parser.add_argument( + "--output_path", type=str, required=True, help="Path where converted model should be saved" + ) + parser.add_argument( + "--fp16", + action="store_true", + default=False, + help="Whether to save the model weights in fp16", + ) + parser.add_argument( + "--bf16", + action="store_true", + default=False, + help="Whether to save the model weights in bf16", + ) + parser.add_argument( + "--push_to_hub", + action="store_true", + default=False, + help="Whether to push to HF Hub after saving", + ) + parser.add_argument( + "--text_encoder_cache_dir", + type=str, + default=None, + help="Path to text encoder cache directory", + ) + parser.add_argument( + "--typecast_text_encoder", + action="store_true", + default=False, + help="Whether or not to apply fp16/bf16 precision to text_encoder", + ) + # For CogVideoX-2B, num_layers is 30. For 5B, it is 42 + parser.add_argument("--num_layers", type=int, default=30, help="Number of transformer blocks") + # For CogVideoX-2B, num_attention_heads is 30. For 5B, it is 48 + parser.add_argument( + "--num_attention_heads", type=int, default=30, help="Number of attention heads" + ) + # For CogVideoX-2B, use_rotary_positional_embeddings is False. For 5B, it is True + parser.add_argument( + "--use_rotary_positional_embeddings", + action="store_true", + default=False, + help="Whether to use RoPE or not", + ) + # For CogVideoX-2B, scaling_factor is 1.15258426. For 5B, it is 0.7 + parser.add_argument( + "--scaling_factor", type=float, default=1.15258426, help="Scaling factor in the VAE" + ) + # For CogVideoX-2B, snr_shift_scale is 3.0. For 5B, it is 1.0 + parser.add_argument( + "--snr_shift_scale", type=float, default=3.0, help="Scaling factor in the VAE" + ) + parser.add_argument( + "--i2v", + action="store_true", + default=False, + help="Whether the model to be converted is the Image-to-Video version of CogVideoX.", + ) + parser.add_argument( + "--version", + choices=["1.0", "1.5"], + default="1.0", + help="Which version of CogVideoX to use for initializing default modeling parameters.", + ) + return parser.parse_args() + + +if __name__ == "__main__": + args = get_args() + + transformer = None + vae = None + + if args.fp16 and args.bf16: + raise ValueError("You cannot pass both --fp16 and --bf16 at the same time.") + + dtype = torch.float16 if args.fp16 else torch.bfloat16 if args.bf16 else torch.float32 + + if args.transformer_ckpt_path is not None: + init_kwargs = get_transformer_init_kwargs(args.version) + transformer = convert_transformer( + args.transformer_ckpt_path, + args.num_layers, + args.num_attention_heads, + args.use_rotary_positional_embeddings, + args.i2v, + dtype, + init_kwargs, + ) + if args.vae_ckpt_path is not None: + # Keep VAE in float32 for better quality + vae = convert_vae(args.vae_ckpt_path, args.scaling_factor, args.version, torch.float32) + + text_encoder_id = "google/t5-v1_1-xxl" + tokenizer = T5Tokenizer.from_pretrained(text_encoder_id, model_max_length=TOKENIZER_MAX_LENGTH) + text_encoder = T5EncoderModel.from_pretrained( + text_encoder_id, cache_dir=args.text_encoder_cache_dir + ) + + if args.typecast_text_encoder: + text_encoder = text_encoder.to(dtype=dtype) + + # Apparently, the conversion does not work anymore without this :shrug: + for param in text_encoder.parameters(): + param.data = param.data.contiguous() + + scheduler = CogVideoXDDIMScheduler.from_config( + { + "snr_shift_scale": args.snr_shift_scale, + "beta_end": 0.012, + "beta_schedule": "scaled_linear", + "beta_start": 0.00085, + "clip_sample": False, + "num_train_timesteps": 1000, + "prediction_type": "v_prediction", + "rescale_betas_zero_snr": True, + "set_alpha_to_one": True, + "timestep_spacing": "trailing", + } + ) + if args.i2v: + pipeline_cls = CogVideoXImageToVideoPipeline + else: + pipeline_cls = CogVideoXPipeline + + pipe = pipeline_cls( + tokenizer=tokenizer, + text_encoder=text_encoder, + vae=vae, + transformer=transformer, + scheduler=scheduler, + ) + + # We don't use variant here because the model must be run in fp16 (2B) or bf16 (5B). It would be weird + # for users to specify variant when the default is not fp32 and they want to run with the correct default (which + # is either fp16/bf16 here). + + # This is necessary This is necessary for users with insufficient memory, + # such as those using Colab and notebooks, as it can save some memory used for model loading. + pipe.save_pretrained( + args.output_path, + safe_serialization=True, + max_shard_size="5GB", + push_to_hub=args.push_to_hub, + ) diff --git a/tools/export_sat_lora_weight.py b/tools/export_sat_lora_weight.py new file mode 100644 index 0000000..4f4be2c --- /dev/null +++ b/tools/export_sat_lora_weight.py @@ -0,0 +1,82 @@ +from typing import Any, Dict +import torch +import argparse +from diffusers.loaders.lora_base import LoraBaseMixin +from diffusers.models.modeling_utils import load_state_dict + + +def get_state_dict(saved_dict: Dict[str, Any]) -> Dict[str, Any]: + state_dict = saved_dict + if "model" in saved_dict.keys(): + state_dict = state_dict["model"] + if "module" in saved_dict.keys(): + state_dict = state_dict["module"] + if "state_dict" in saved_dict.keys(): + state_dict = state_dict["state_dict"] + return state_dict + + +LORA_KEYS_RENAME = { + 'attention.query_key_value.matrix_A.0': 'attn1.to_q.lora_A.weight', + 'attention.query_key_value.matrix_A.1': 'attn1.to_k.lora_A.weight', + 'attention.query_key_value.matrix_A.2': 'attn1.to_v.lora_A.weight', + 'attention.query_key_value.matrix_B.0': 'attn1.to_q.lora_B.weight', + 'attention.query_key_value.matrix_B.1': 'attn1.to_k.lora_B.weight', + 'attention.query_key_value.matrix_B.2': 'attn1.to_v.lora_B.weight', + 'attention.dense.matrix_A.0': 'attn1.to_out.0.lora_A.weight', + 'attention.dense.matrix_B.0': 'attn1.to_out.0.lora_B.weight', +} + + +PREFIX_KEY = "model.diffusion_model." +SAT_UNIT_KEY = "layers" +LORA_PREFIX_KEY = "transformer_blocks" + + +def export_lora_weight(ckpt_path, lora_save_directory): + merge_original_state_dict = get_state_dict(torch.load(ckpt_path, map_location="cpu", mmap=True)) + + lora_state_dict = {} + for key in list(merge_original_state_dict.keys()): + new_key = key[len(PREFIX_KEY) :] + for special_key, lora_keys in LORA_KEYS_RENAME.items(): + if new_key.endswith(special_key): + new_key = new_key.replace(special_key, lora_keys) + new_key = new_key.replace(SAT_UNIT_KEY, LORA_PREFIX_KEY) + + lora_state_dict[new_key] = merge_original_state_dict[key] + + # final length should be 240 + if len(lora_state_dict) != 240: + raise ValueError("lora_state_dict length is not 240") + + lora_state_dict.keys() + + LoraBaseMixin.write_lora_layers( + state_dict=lora_state_dict, + save_directory=lora_save_directory, + is_main_process=True, + weight_name=None, + save_function=None, + safe_serialization=True, + ) + + +def get_args(): + parser = argparse.ArgumentParser() + parser.add_argument( + "--sat_pt_path", type=str, required=True, help="Path to original sat transformer checkpoint" + ) + parser.add_argument( + "--lora_save_directory", + type=str, + required=True, + help="Path where converted lora should be saved", + ) + return parser.parse_args() + + +if __name__ == "__main__": + args = get_args() + + export_lora_weight(args.sat_pt_path, args.lora_save_directory) diff --git a/tools/llm_flux_cogvideox/generate.sh b/tools/llm_flux_cogvideox/generate.sh new file mode 100644 index 0000000..ade1cbe --- /dev/null +++ b/tools/llm_flux_cogvideox/generate.sh @@ -0,0 +1,35 @@ +#!/bin/bash + +NUM_VIDEOS=10 +INFERENCE_STEPS=50 +GUIDANCE_SCALE=7.0 +OUTPUT_DIR_PREFIX="outputs/gpu_" +LOG_DIR_PREFIX="logs/gpu_" + +VIDEO_MODEL_PATH="/share/official_pretrains/hf_home/CogVideoX-5b-I2V" +LLM_MODEL_PATH="/share/home/zyx/Models/Meta-Llama-3.1-8B-Instruct" +IMAGE_MODEL_PATH = "share/home/zyx/Models/FLUX.1-dev" + +#VIDEO_MODEL_PATH="THUDM/CogVideoX-5B-I2V" +#LLM_MODEL_PATH="THUDM/glm-4-9b-chat" +#IMAGE_MODEL_PATH = "black-forest-labs/FLUX.1-dev" + +CUDA_DEVICES=${CUDA_VISIBLE_DEVICES:-"0"} + +IFS=',' read -r -a GPU_ARRAY <<< "$CUDA_DEVICES" + +for i in "${!GPU_ARRAY[@]}" +do + GPU=${GPU_ARRAY[$i]} + echo "Starting task on GPU $GPU..." + CUDA_VISIBLE_DEVICES=$GPU nohup python3 llm_flux_cogvideox.py \ + --caption_generator_model_id $LLM_MODEL_PATH \ + --image_generator_model_id $IMAGE_MODEL_PATH \ + --model_path $VIDEO_MODEL_PATH \ + --num_videos $NUM_VIDEOS \ + --image_generator_num_inference_steps $INFERENCE_STEPS \ + --guidance_scale $GUIDANCE_SCALE \ + --use_dynamic_cfg \ + --output_dir ${OUTPUT_DIR_PREFIX}${GPU} \ + > ${LOG_DIR_PREFIX}${GPU}.log 2>&1 & +done diff --git a/tools/llm_flux_cogvideox/gradio_page.py b/tools/llm_flux_cogvideox/gradio_page.py new file mode 100644 index 0000000..27fc1a0 --- /dev/null +++ b/tools/llm_flux_cogvideox/gradio_page.py @@ -0,0 +1,187 @@ +import os +import gradio as gr +import gc +import random +import torch +import numpy as np +from PIL import Image +import transformers +from diffusers import CogVideoXImageToVideoPipeline, CogVideoXDPMScheduler, DiffusionPipeline +from diffusers.utils import export_to_video +from transformers import AutoTokenizer +from datetime import datetime, timedelta +import threading +import time +from moviepy import VideoFileClip + +torch.set_float32_matmul_precision("high") + +# Set default values +caption_generator_model_id = "/share/home/zyx/Models/Meta-Llama-3.1-8B-Instruct" +image_generator_model_id = "/share/home/zyx/Models/FLUX.1-dev" +video_generator_model_id = "/share/official_pretrains/hf_home/CogVideoX-5b-I2V" +seed = 1337 + +os.makedirs("./output", exist_ok=True) +os.makedirs("./gradio_tmp", exist_ok=True) + +tokenizer = AutoTokenizer.from_pretrained(caption_generator_model_id, trust_remote_code=True) +caption_generator = transformers.pipeline( + "text-generation", + model=caption_generator_model_id, + device_map="balanced", + model_kwargs={ + "local_files_only": True, + "torch_dtype": torch.bfloat16, + }, + trust_remote_code=True, + tokenizer=tokenizer, +) + +image_generator = DiffusionPipeline.from_pretrained( + image_generator_model_id, torch_dtype=torch.bfloat16, device_map="balanced" +) +# image_generator.to("cuda") + +video_generator = CogVideoXImageToVideoPipeline.from_pretrained( + video_generator_model_id, torch_dtype=torch.bfloat16, device_map="balanced" +) + +video_generator.vae.enable_slicing() +video_generator.vae.enable_tiling() + +video_generator.scheduler = CogVideoXDPMScheduler.from_config( + video_generator.scheduler.config, timestep_spacing="trailing" +) + +# Define prompts +SYSTEM_PROMPT = """ +You are part of a team of people that create videos using generative models. You use a video-generation model that can generate a video about anything you describe. + +For example, if you respond with "A beautiful morning in the woods with the sun peaking through the trees", the video generation model will create a video of exactly as described. Your task is to summarize the descriptions of videos provided by users and create detailed prompts to feed into the generative model. + +There are a few rules to follow: +- You will only ever output a single video description per request. +- If the user mentions to summarize the prompt in [X] words, make sure not to exceed the limit. + +Your responses should just be the video generation prompt. Here are examples: +- "A detailed wooden toy ship with intricately carved masts and sails is seen gliding smoothly over a plush, blue carpet that mimics the waves of the sea. The ship's hull is painted a rich brown, with tiny windows. The carpet, soft and textured, provides a perfect backdrop, resembling an oceanic expanse. Surrounding the ship are various other toys and children's items, hinting at a playful environment. The scene captures the innocence and imagination of childhood, with the toy ship's journey symbolizing endless adventures in a whimsical, indoor setting." +- "A street artist, clad in a worn-out denim jacket and a colorful bandana, stands before a vast concrete wall in the heart of the city, holding a can of spray paint, spray-painting a colorful bird on a mottled wall." +""".strip() + +USER_PROMPT = """ +Could you generate a prompt for a video generation model? Please limit the prompt to [{0}] words. +""".strip() + + +def generate_caption(prompt): + num_words = random.choice([25, 50, 75, 100]) + user_prompt = USER_PROMPT.format(num_words) + + messages = [ + {"role": "system", "content": SYSTEM_PROMPT}, + {"role": "user", "content": prompt + "\n" + user_prompt}, + ] + + response = caption_generator(messages, max_new_tokens=226, return_full_text=False) + caption = response[0]["generated_text"] + if caption.startswith("\"") and caption.endswith("\""): + caption = caption[1:-1] + return caption + + +def generate_image(caption, progress=gr.Progress(track_tqdm=True)): + image = image_generator( + prompt=caption, + height=480, + width=720, + num_inference_steps=30, + guidance_scale=3.5, + ).images[0] + return image, image # One for output One for State + + +def generate_video(caption, image, progress=gr.Progress(track_tqdm=True)): + generator = torch.Generator().manual_seed(seed) + video_frames = video_generator( + image=image, + prompt=caption, + height=480, + width=720, + num_frames=49, + num_inference_steps=50, + guidance_scale=6, + use_dynamic_cfg=True, + generator=generator, + ).frames[0] + video_path = save_video(video_frames) + gif_path = convert_to_gif(video_path) + return video_path, gif_path + + +def save_video(tensor): + timestamp = datetime.now().strftime("%Y%m%d_%H%M%S") + video_path = f"./output/{timestamp}.mp4" + os.makedirs(os.path.dirname(video_path), exist_ok=True) + export_to_video(tensor, video_path, fps=8) + return video_path + + +def convert_to_gif(video_path): + clip = VideoFileClip(video_path) + clip = clip.with_fps(8) + clip = clip.resized(height=240) + gif_path = video_path.replace(".mp4", ".gif") + clip.write_gif(gif_path, fps=8) + return gif_path + + +def delete_old_files(): + while True: + now = datetime.now() + cutoff = now - timedelta(minutes=10) + directories = ["./output", "./gradio_tmp"] + + for directory in directories: + for filename in os.listdir(directory): + file_path = os.path.join(directory, filename) + if os.path.isfile(file_path): + file_mtime = datetime.fromtimestamp(os.path.getmtime(file_path)) + if file_mtime < cutoff: + os.remove(file_path) + time.sleep(600) + + +threading.Thread(target=delete_old_files, daemon=True).start() + +with gr.Blocks() as demo: + gr.Markdown(""" +
+ LLM + FLUX + CogVideoX-I2V Space 🤗 +
+ """) + with gr.Row(): + with gr.Column(): + prompt = gr.Textbox(label="Prompt", placeholder="Enter your prompt here", lines=5) + generate_caption_button = gr.Button("Generate Caption") + caption = gr.Textbox(label="Caption", placeholder="Caption will appear here", lines=5) + generate_image_button = gr.Button("Generate Image") + image_output = gr.Image(label="Generated Image") + state_image = gr.State() + generate_caption_button.click(fn=generate_caption, inputs=prompt, outputs=caption) + generate_image_button.click( + fn=generate_image, inputs=caption, outputs=[image_output, state_image] + ) + with gr.Column(): + video_output = gr.Video(label="Generated Video", width=720, height=480) + download_video_button = gr.File(label="📥 Download Video", visible=False) + download_gif_button = gr.File(label="📥 Download GIF", visible=False) + generate_video_button = gr.Button("Generate Video from Image") + generate_video_button.click( + fn=generate_video, + inputs=[caption, state_image], + outputs=[video_output, download_gif_button], + ) + +if __name__ == "__main__": + demo.launch() diff --git a/tools/llm_flux_cogvideox/llm_flux_cogvideox.py b/tools/llm_flux_cogvideox/llm_flux_cogvideox.py new file mode 100644 index 0000000..853f4fc --- /dev/null +++ b/tools/llm_flux_cogvideox/llm_flux_cogvideox.py @@ -0,0 +1,260 @@ +""" +The original experimental code for this project can be found at: + +https://gist.github.com/a-r-r-o-w/d070cce059ab4ceab3a9f289ff83c69c + +By using this code, description prompts will be generated through a local large language model, and images will be +generated using the black-forest-labs/FLUX.1-dev model, followed by video generation via CogVideoX. +The entire process utilizes open-source solutions, without the need for any API keys. + +You can use the generate.sh file in the same folder to automate running this code +for batch generation of videos and images. + +bash generate.sh + +""" + +import argparse +import gc +import json +import os +import pathlib +import random +from typing import Any, Dict + +from transformers import AutoTokenizer + +os.environ["TORCH_LOGS"] = "+dynamo,recompiles,graph_breaks" +os.environ["TORCHDYNAMO_VERBOSE"] = "1" + +import numpy as np +import torch +import transformers +from diffusers import CogVideoXImageToVideoPipeline, CogVideoXDPMScheduler, DiffusionPipeline +from diffusers.utils.logging import get_logger +from diffusers.utils import export_to_video + +torch.set_float32_matmul_precision("high") + +logger = get_logger(__name__) + +SYSTEM_PROMPT = """ +You are part of a team of people that create videos using generative models. You use a video-generation model that can generate a video about anything you describe. + +For example, if you respond with "A beautiful morning in the woods with the sun peaking through the trees", the video generation model will create a video of exactly as described. You task is to summarize the descriptions of videos provided to by users, and create details prompts to feed into the generative model. + +There are a few rules to follow: +- You will only ever output a single video description per request. +- If the user mentions to summarize the prompt in [X] words, make sure to not exceed the limit. + +You responses should just be the video generation prompt. Here are examples: +- “A lone figure stands on a city rooftop at night, gazing up at the full moon. The moon glows brightly, casting a gentle light over the quiet cityscape. Below, the windows of countless homes shine with warm lights, creating a contrast between the bustling life below and the peaceful solitude above. The scene captures the essence of the Mid-Autumn Festival, where despite the distance, the figure feels connected to loved ones through the shared beauty of the moonlit sky.” +- "A detailed wooden toy ship with intricately carved masts and sails is seen gliding smoothly over a plush, blue carpet that mimics the waves of the sea. The ship's hull is painted a rich brown, with tiny windows. The carpet, soft and textured, provides a perfect backdrop, resembling an oceanic expanse. Surrounding the ship are various other toys and children's items, hinting at a playful environment. The scene captures the innocence and imagination of childhood, with the toy ship's journey symbolizing endless adventures in a whimsical, indoor setting." +- "A street artist, clad in a worn-out denim jacket and a colorful banana, stands before a vast concrete wall in the heart, holding a can of spray paint, spray-painting a colorful bird on a mottled wall" +""".strip() + +USER_PROMPT = """ +Could you generate a prompt for a video generation model? +Please limit the prompt to [{0}] words. +""".strip() + + +def get_args(): + parser = argparse.ArgumentParser() + parser.add_argument( + "--num_videos", + type=int, + default=5, + help="Number of unique videos you would like to generate.", + ) + parser.add_argument( + "--model_path", + type=str, + default="THUDM/CogVideoX-5B", + help="The path of Image2Video CogVideoX-5B", + ) + parser.add_argument( + "--caption_generator_model_id", + type=str, + default="THUDM/glm-4-9b-chat", + help="Caption generation model. default GLM-4-9B", + ) + parser.add_argument( + "--caption_generator_cache_dir", + type=str, + default=None, + help="Cache directory for caption generation model.", + ) + parser.add_argument( + "--image_generator_model_id", + type=str, + default="black-forest-labs/FLUX.1-dev", + help="Image generation model.", + ) + parser.add_argument( + "--image_generator_cache_dir", + type=str, + default=None, + help="Cache directory for image generation model.", + ) + parser.add_argument( + "--image_generator_num_inference_steps", + type=int, + default=50, + help="Caption generation model.", + ) + parser.add_argument( + "--guidance_scale", type=float, default=7, help="Guidance scale to be use for generation." + ) + parser.add_argument( + "--use_dynamic_cfg", + action="store_true", + help="Whether or not to use cosine dynamic guidance for generation [Recommended].", + ) + parser.add_argument( + "--output_dir", + type=str, + default="outputs/", + help="Location where generated images and videos should be stored.", + ) + parser.add_argument( + "--compile", + action="store_true", + help="Whether or not to compile the transformer of image and video generators.", + ) + parser.add_argument( + "--enable_vae_tiling", + action="store_true", + help="Whether or not to use VAE tiling when encoding/decoding.", + ) + parser.add_argument("--seed", type=int, default=42, help="Seed for reproducibility.") + return parser.parse_args() + + +def reset_memory(): + gc.collect() + torch.cuda.empty_cache() + torch.cuda.reset_peak_memory_stats() + torch.cuda.reset_accumulated_memory_stats() + + +@torch.no_grad() +def main(args: Dict[str, Any]) -> None: + output_dir = pathlib.Path(args.output_dir) + os.makedirs(output_dir.as_posix(), exist_ok=True) + + random.seed(args.seed) + np.random.seed(args.seed) + torch.manual_seed(args.seed) + torch.cuda.manual_seed_all(args.seed) + + reset_memory() + tokenizer = AutoTokenizer.from_pretrained( + args.caption_generator_model_id, trust_remote_code=True + ) + caption_generator = transformers.pipeline( + "text-generation", + model=args.caption_generator_model_id, + device_map="auto", + model_kwargs={ + "local_files_only": True, + "cache_dir": args.caption_generator_cache_dir, + "torch_dtype": torch.bfloat16, + }, + trust_remote_code=True, + tokenizer=tokenizer, + ) + + captions = [] + for i in range(args.num_videos): + num_words = random.choice([50, 75, 100]) + user_prompt = USER_PROMPT.format(num_words) + + messages = [ + {"role": "system", "content": SYSTEM_PROMPT}, + {"role": "user", "content": user_prompt}, + ] + + outputs = caption_generator(messages, max_new_tokens=226) + caption = outputs[0]["generated_text"][-1]["content"] + if caption.startswith("\"") and caption.endswith("\""): + caption = caption[1:-1] + captions.append(caption) + logger.info(f"Generated caption: {caption}") + + with open(output_dir / "captions.json", "w") as file: + json.dump(captions, file) + + del caption_generator + reset_memory() + + image_generator = DiffusionPipeline.from_pretrained( + args.image_generator_model_id, + cache_dir=args.image_generator_cache_dir, + torch_dtype=torch.bfloat16, + ) + image_generator.to("cuda") + + if args.compile: + image_generator.transformer = torch.compile( + image_generator.transformer, mode="max-autotune", fullgraph=True + ) + + if args.enable_vae_tiling: + image_generator.vae.enable_tiling() + + images = [] + for index, caption in enumerate(captions): + image = image_generator( + prompt=caption, + height=480, + width=720, + num_inference_steps=args.image_generator_num_inference_steps, + guidance_scale=3.5, + ).images[0] + filename = ( + caption[:25].replace(".", "_").replace("'", "_").replace('"', "_").replace(",", "_") + ) + image.save(output_dir / f"{index}_{filename}.png") + images.append(image) + + del image_generator + reset_memory() + + video_generator = CogVideoXImageToVideoPipeline.from_pretrained( + args.model_path, torch_dtype=torch.bfloat16 + ).to("cuda") + video_generator.scheduler = CogVideoXDPMScheduler.from_config( + video_generator.scheduler.config, timestep_spacing="trailing" + ) + + if args.compile: + video_generator.transformer = torch.compile( + video_generator.transformer, mode="max-autotune", fullgraph=True + ) + + if args.enable_vae_tiling: + video_generator.vae.enable_tiling() + + generator = torch.Generator().manual_seed(args.seed) + for index, (caption, image) in enumerate(zip(captions, images)): + video = video_generator( + image=image, + prompt=caption, + height=480, + width=720, + num_frames=49, + num_inference_steps=50, + guidance_scale=args.guidance_scale, + use_dynamic_cfg=args.use_dynamic_cfg, + generator=generator, + ).frames[0] + filename = ( + caption[:25].replace(".", "_").replace("'", "_").replace('"', "_").replace(",", "_") + ) + export_to_video(video, output_dir / f"{index}_{filename}.mp4", fps=8) + + +if __name__ == "__main__": + args = get_args() + main(args) diff --git a/tools/load_cogvideox_lora.py b/tools/load_cogvideox_lora.py new file mode 100644 index 0000000..2adf30e --- /dev/null +++ b/tools/load_cogvideox_lora.py @@ -0,0 +1,132 @@ +# Copyright 2024 The HuggingFace Team. +# All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + + +import math +import random +import time +from diffusers.utils import export_to_video +from diffusers.image_processor import VaeImageProcessor +from datetime import datetime, timedelta +from diffusers import CogVideoXPipeline, CogVideoXDDIMScheduler, CogVideoXDPMScheduler +import os +import torch +import argparse + + +device = "cuda" if torch.cuda.is_available() else "cpu" + + +def get_args(): + parser = argparse.ArgumentParser() + parser.add_argument( + "--pretrained_model_name_or_path", + type=str, + default=None, + required=True, + help="Path to pretrained model or model identifier from huggingface.co/models.", + ) + parser.add_argument( + "--lora_weights_path", + type=str, + default=None, + required=True, + help="Path to lora weights.", + ) + parser.add_argument( + "--lora_r", + type=int, + default=128, + help="""LoRA weights have a rank parameter, with the default for 2B trans set at 128 and 5B trans set at 256. + This part is used to calculate the value for lora_scale, which is by default divided by the alpha value, + used for stable learning and to prevent underflow. In the SAT training framework, + alpha is set to 1 by default. The higher the rank, the better the expressive capability, + but it requires more memory and training time. Increasing this number blindly isn't always better. + The formula for lora_scale is: lora_r / alpha. + """, + ) + parser.add_argument( + "--lora_alpha", + type=int, + default=1, + help="""LoRA weights have a rank parameter, with the default for 2B trans set at 128 and 5B trans set at 256. + This part is used to calculate the value for lora_scale, which is by default divided by the alpha value, + used for stable learning and to prevent underflow. In the SAT training framework, + alpha is set to 1 by default. The higher the rank, the better the expressive capability, + but it requires more memory and training time. Increasing this number blindly isn't always better. + The formula for lora_scale is: lora_r / alpha. + """, + ) + parser.add_argument( + "--prompt", + type=str, + help="prompt", + ) + parser.add_argument( + "--output_dir", + type=str, + default="output", + help="The output directory where the model predictions and checkpoints will be written.", + ) + return parser.parse_args() + + +if __name__ == "__main__": + args = get_args() + pipe = CogVideoXPipeline.from_pretrained( + args.pretrained_model_name_or_path, torch_dtype=torch.bfloat16 + ).to(device) + pipe.load_lora_weights( + args.lora_weights_path, + weight_name="pytorch_lora_weights.safetensors", + adapter_name="cogvideox-lora", + ) + # pipe.fuse_lora(lora_scale=args.lora_alpha/args.lora_r, ['transformer']) + lora_scaling = args.lora_alpha / args.lora_r + pipe.set_adapters(["cogvideox-lora"], [lora_scaling]) + + pipe.scheduler = CogVideoXDPMScheduler.from_config( + pipe.scheduler.config, timestep_spacing="trailing" + ) + + os.makedirs(args.output_dir, exist_ok=True) + + latents = pipe( + prompt=args.prompt, + num_videos_per_prompt=1, + num_inference_steps=50, + num_frames=49, + use_dynamic_cfg=True, + output_type="pt", + guidance_scale=3.0, + generator=torch.Generator(device="cpu").manual_seed(42), + ).frames + batch_size = latents.shape[0] + batch_video_frames = [] + for batch_idx in range(batch_size): + pt_image = latents[batch_idx] + pt_image = torch.stack([pt_image[i] for i in range(pt_image.shape[0])]) + + image_np = VaeImageProcessor.pt_to_numpy(pt_image) + image_pil = VaeImageProcessor.numpy_to_pil(image_np) + batch_video_frames.append(image_pil) + + timestamp = datetime.now().strftime("%Y%m%d_%H%M%S") + video_path = f"{args.output_dir}/{timestamp}.mp4" + os.makedirs(os.path.dirname(video_path), exist_ok=True) + tensor = batch_video_frames[0] + fps = math.ceil((len(batch_video_frames[0]) - 1) / 6) + + export_to_video(tensor, video_path, fps=fps) diff --git a/tools/parallel_inference/parallel_inference_xdit.py b/tools/parallel_inference/parallel_inference_xdit.py new file mode 100644 index 0000000..d0b8029 --- /dev/null +++ b/tools/parallel_inference/parallel_inference_xdit.py @@ -0,0 +1,109 @@ +""" +This is a parallel inference script for CogVideo. The original script +can be found from the xDiT project at + +https://github.com/xdit-project/xDiT/blob/main/examples/cogvideox_example.py + +By using this code, the inference process is parallelized on multiple GPUs, +and thus speeded up. + +Usage: +1. pip install xfuser +2. mkdir results +3. run the following command to generate video +torchrun --nproc_per_node=4 parallel_inference_xdit.py \ + --model --ulysses_degree 1 --ring_degree 2 \ + --use_cfg_parallel --height 480 --width 720 --num_frames 9 \ + --prompt 'A small dog.' + +You can also use the run.sh file in the same folder to automate running this +code for batch generation of videos, by running: + +sh ./run.sh + +""" + +import time +import torch +import torch.distributed +from diffusers import AutoencoderKLTemporalDecoder +from xfuser import xFuserCogVideoXPipeline, xFuserArgs +from xfuser.config import FlexibleArgumentParser +from xfuser.core.distributed import ( + get_world_group, + get_data_parallel_rank, + get_data_parallel_world_size, + get_runtime_state, + is_dp_last_group, +) +from diffusers.utils import export_to_video + + +def main(): + parser = FlexibleArgumentParser(description="xFuser Arguments") + args = xFuserArgs.add_cli_args(parser).parse_args() + engine_args = xFuserArgs.from_cli_args(args) + + # Check if ulysses_degree is valid + num_heads = 30 + if engine_args.ulysses_degree > 0 and num_heads % engine_args.ulysses_degree != 0: + raise ValueError( + f"ulysses_degree ({engine_args.ulysses_degree}) must be a divisor of the number of heads ({num_heads})" + ) + + engine_config, input_config = engine_args.create_config() + local_rank = get_world_group().local_rank + + pipe = xFuserCogVideoXPipeline.from_pretrained( + pretrained_model_name_or_path=engine_config.model_config.model, + engine_config=engine_config, + torch_dtype=torch.bfloat16, + ) + if args.enable_sequential_cpu_offload: + pipe.enable_model_cpu_offload(gpu_id=local_rank) + else: + device = torch.device(f"cuda:{local_rank}") + pipe = pipe.to(device) + + # Always enable tiling and slicing to avoid VAE OOM while batch size > 1 + pipe.vae.enable_slicing() + pipe.vae.enable_tiling() + + torch.cuda.reset_peak_memory_stats() + start_time = time.time() + + output = pipe( + height=input_config.height, + width=input_config.width, + num_frames=input_config.num_frames, + prompt=input_config.prompt, + num_inference_steps=input_config.num_inference_steps, + generator=torch.Generator().manual_seed(input_config.seed), + guidance_scale=6, + use_dynamic_cfg=True, + ).frames[0] + + end_time = time.time() + elapsed_time = end_time - start_time + peak_memory = torch.cuda.max_memory_allocated(device=f"cuda:{local_rank}") + + parallel_info = ( + f"dp{engine_args.data_parallel_degree}_cfg{engine_config.parallel_config.cfg_degree}_" + f"ulysses{engine_args.ulysses_degree}_ring{engine_args.ring_degree}_" + f"tp{engine_args.tensor_parallel_degree}_" + f"pp{engine_args.pipefusion_parallel_degree}_patch{engine_args.num_pipeline_patch}" + ) + if is_dp_last_group(): + world_size = get_data_parallel_world_size() + resolution = f"{input_config.width}x{input_config.height}" + output_filename = f"results/cogvideox_{parallel_info}_{resolution}.mp4" + export_to_video(output, output_filename, fps=8) + print(f"output saved to {output_filename}") + + if get_world_group().rank == get_world_group().world_size - 1: + print(f"epoch time: {elapsed_time:.2f} sec, memory: {peak_memory/1e9} GB") + get_runtime_state().destory_distributed_env() + + +if __name__ == "__main__": + main() diff --git a/tools/parallel_inference/run.sh b/tools/parallel_inference/run.sh new file mode 100644 index 0000000..ee421a4 --- /dev/null +++ b/tools/parallel_inference/run.sh @@ -0,0 +1,51 @@ +set -x + +export PYTHONPATH=$PWD:$PYTHONPATH + +# Select the model type +# The model is downloaded to a specified location on disk, +# or you can simply use the model's ID on Hugging Face, +# which will then be downloaded to the default cache path on Hugging Face. + +export MODEL_TYPE="CogVideoX" +# Configuration for different model types +# script, model_id, inference_step +declare -A MODEL_CONFIGS=( + ["CogVideoX"]="parallel_inference_xdit.py /cfs/dit/CogVideoX-2b 20" +) + +if [[ -v MODEL_CONFIGS[$MODEL_TYPE] ]]; then + IFS=' ' read -r SCRIPT MODEL_ID INFERENCE_STEP <<< "${MODEL_CONFIGS[$MODEL_TYPE]}" + export SCRIPT MODEL_ID INFERENCE_STEP +else + echo "Invalid MODEL_TYPE: $MODEL_TYPE" + exit 1 +fi + +mkdir -p ./results + +# task args +if [ "$MODEL_TYPE" = "CogVideoX" ]; then + TASK_ARGS="--height 480 --width 720 --num_frames 9" +fi + +# CogVideoX asserts sp_degree == ulysses_degree*ring_degree <= 2. Also, do not set the pipefusion degree. +if [ "$MODEL_TYPE" = "CogVideoX" ]; then +N_GPUS=4 +PARALLEL_ARGS="--ulysses_degree 2 --ring_degree 1" +CFG_ARGS="--use_cfg_parallel" +fi + + +torchrun --nproc_per_node=$N_GPUS ./$SCRIPT \ +--model $MODEL_ID \ +$PARALLEL_ARGS \ +$TASK_ARGS \ +$PIPEFUSION_ARGS \ +$OUTPUT_ARGS \ +--num_inference_steps $INFERENCE_STEP \ +--warmup_steps 0 \ +--prompt "A small dog." \ +$CFG_ARGS \ +$PARALLLEL_VAE \ +$COMPILE_FLAG diff --git a/tools/replicate/cog.yaml b/tools/replicate/cog.yaml new file mode 100644 index 0000000..19d71ad --- /dev/null +++ b/tools/replicate/cog.yaml @@ -0,0 +1,37 @@ +# Configuration for Cog ⚙️ +# Reference: https://cog.run/yaml + +build: + # set to true if your model requires a GPU + gpu: true + + # a list of ubuntu apt packages to install + system_packages: + - "libgl1-mesa-glx" + - "libglib2.0-0" + + # python version in the form '3.11' or '3.11.4' + python_version: "3.11" + + # a list of packages in the format == + python_packages: + - diffusers>=0.30.3 + - accelerate>=0.34.2 + - transformers>=4.44.2 + - numpy==1.26.0 + - torch>=2.4.0 + - torchvision>=0.19.0 + - sentencepiece>=0.2.0 + - SwissArmyTransformer>=0.4.12 + - imageio>=2.35.1 + - imageio-ffmpeg>=0.5.1 + - openai>=1.45.0 + - moviepy>=2.0.0 + - pillow==9.5.0 + - pydantic==1.10.7 + run: + - curl -o /usr/local/bin/pget -L "https://github.com/replicate/pget/releases/download/v0.8.2/pget_linux_x86_64" && chmod +x /usr/local/bin/pget + +# predict.py defines how predictions are run on your model +predict: "predict_t2v.py:Predictor" +# predict: "predict_i2v.py:Predictor" diff --git a/tools/replicate/predict_i2v.py b/tools/replicate/predict_i2v.py new file mode 100644 index 0000000..69c294a --- /dev/null +++ b/tools/replicate/predict_i2v.py @@ -0,0 +1,83 @@ +# Prediction interface for Cog ⚙️ +# https://cog.run/python + +import os +import subprocess +import time +import torch +from diffusers import CogVideoXImageToVideoPipeline +from diffusers.utils import export_to_video, load_image +from cog import BasePredictor, Input, Path + + +MODEL_CACHE = "model_cache_i2v" +MODEL_URL = f"https://weights.replicate.delivery/default/THUDM/CogVideo/{MODEL_CACHE}.tar" +os.environ["HF_DATASETS_OFFLINE"] = "1" +os.environ["TRANSFORMERS_OFFLINE"] = "1" +os.environ["HF_HOME"] = MODEL_CACHE +os.environ["TORCH_HOME"] = MODEL_CACHE +os.environ["HF_DATASETS_CACHE"] = MODEL_CACHE +os.environ["TRANSFORMERS_CACHE"] = MODEL_CACHE +os.environ["HUGGINGFACE_HUB_CACHE"] = MODEL_CACHE + + +def download_weights(url, dest): + start = time.time() + print("downloading url: ", url) + print("downloading to: ", dest) + subprocess.check_call(["pget", "-x", url, dest], close_fds=False) + print("downloading took: ", time.time() - start) + + +class Predictor(BasePredictor): + def setup(self) -> None: + """Load the model into memory to make running multiple predictions efficient""" + + if not os.path.exists(MODEL_CACHE): + download_weights(MODEL_URL, MODEL_CACHE) + + # model_id: THUDM/CogVideoX-5b-I2V + self.pipe = CogVideoXImageToVideoPipeline.from_pretrained( + MODEL_CACHE, torch_dtype=torch.bfloat16 + ).to("cuda") + + self.pipe.enable_model_cpu_offload() + self.pipe.vae.enable_tiling() + + def predict( + self, + prompt: str = Input(description="Input prompt", default="Starry sky slowly rotating."), + image: Path = Input(description="Input image"), + num_inference_steps: int = Input( + description="Number of denoising steps", ge=1, le=500, default=50 + ), + guidance_scale: float = Input( + description="Scale for classifier-free guidance", ge=1, le=20, default=6 + ), + num_frames: int = Input(description="Number of frames for the output video", default=49), + seed: int = Input( + description="Random seed. Leave blank to randomize the seed", default=None + ), + ) -> Path: + """Run a single prediction on the model""" + + if seed is None: + seed = int.from_bytes(os.urandom(2), "big") + print(f"Using seed: {seed}") + + img = load_image(image=str(image)) + + video = self.pipe( + prompt=prompt, + image=img, + num_videos_per_prompt=1, + num_inference_steps=num_inference_steps, + num_frames=num_frames, + guidance_scale=guidance_scale, + generator=torch.Generator(device="cuda").manual_seed(seed), + ).frames[0] + + out_path = "/tmp/out.mp4" + + export_to_video(video, out_path, fps=8) + return Path(out_path) diff --git a/tools/replicate/predict_t2v.py b/tools/replicate/predict_t2v.py new file mode 100644 index 0000000..51f1106 --- /dev/null +++ b/tools/replicate/predict_t2v.py @@ -0,0 +1,83 @@ +# Prediction interface for Cog ⚙️ +# https://cog.run/python + +import os +import subprocess +import time +import torch +from diffusers import CogVideoXPipeline +from diffusers.utils import export_to_video +from cog import BasePredictor, Input, Path + + +MODEL_CACHE = "model_cache" +MODEL_URL = f"https://weights.replicate.delivery/default/THUDM/CogVideo/{MODEL_CACHE}.tar" +os.environ["HF_DATASETS_OFFLINE"] = "1" +os.environ["TRANSFORMERS_OFFLINE"] = "1" +os.environ["HF_HOME"] = MODEL_CACHE +os.environ["TORCH_HOME"] = MODEL_CACHE +os.environ["HF_DATASETS_CACHE"] = MODEL_CACHE +os.environ["TRANSFORMERS_CACHE"] = MODEL_CACHE +os.environ["HUGGINGFACE_HUB_CACHE"] = MODEL_CACHE + + +def download_weights(url, dest): + start = time.time() + print("downloading url: ", url) + print("downloading to: ", dest) + subprocess.check_call(["pget", "-x", url, dest], close_fds=False) + print("downloading took: ", time.time() - start) + + +class Predictor(BasePredictor): + def setup(self) -> None: + """Load the model into memory to make running multiple predictions efficient""" + + if not os.path.exists(MODEL_CACHE): + download_weights(MODEL_URL, MODEL_CACHE) + + # model_id: THUDM/CogVideoX-5b + self.pipe = CogVideoXPipeline.from_pretrained( + MODEL_CACHE, + torch_dtype=torch.bfloat16, + ).to("cuda") + + self.pipe.enable_model_cpu_offload() + self.pipe.vae.enable_tiling() + + def predict( + self, + prompt: str = Input( + description="Input prompt", + default="A panda, dressed in a small, red jacket and a tiny hat, sits on a wooden stool in a serene bamboo forest. The panda's fluffy paws strum a miniature acoustic guitar, producing soft, melodic tunes. Nearby, a few other pandas gather, watching curiously and some clapping in rhythm. Sunlight filters through the tall bamboo, casting a gentle glow on the scene. The panda's face is expressive, showing concentration and joy as it plays. The background includes a small, flowing stream and vibrant green foliage, enhancing the peaceful and magical atmosphere of this unique musical performance.", + ), + num_inference_steps: int = Input( + description="Number of denoising steps", ge=1, le=500, default=50 + ), + guidance_scale: float = Input( + description="Scale for classifier-free guidance", ge=1, le=20, default=6 + ), + num_frames: int = Input(description="Number of frames for the output video", default=49), + seed: int = Input( + description="Random seed. Leave blank to randomize the seed", default=None + ), + ) -> Path: + """Run a single prediction on the model""" + + if seed is None: + seed = int.from_bytes(os.urandom(2), "big") + print(f"Using seed: {seed}") + + video = self.pipe( + prompt=prompt, + num_videos_per_prompt=1, + num_inference_steps=num_inference_steps, + num_frames=num_frames, + guidance_scale=guidance_scale, + generator=torch.Generator(device="cuda").manual_seed(seed), + ).frames[0] + + out_path = "/tmp/out.mp4" + + export_to_video(video, out_path, fps=8) + return Path(out_path) diff --git a/tools/venhancer/README.md b/tools/venhancer/README.md new file mode 100644 index 0000000..05e037d --- /dev/null +++ b/tools/venhancer/README.md @@ -0,0 +1,98 @@ +# Enhance CogVideoX Generated Videos with VEnhancer + +This tutorial will guide you through using the VEnhancer tool to enhance videos generated by CogVideoX, including +achieving higher frame rates and higher resolutions. + +## Model Introduction + +VEnhancer implements spatial super-resolution, temporal super-resolution (frame interpolation), and video refinement in +a unified framework. It can flexibly adapt to different upsampling factors (e.g., 1x~8x) for spatial or temporal +super-resolution. Additionally, it provides flexible control to modify the refinement strength, enabling it to handle +diverse video artifacts. + +VEnhancer follows the design of ControlNet, copying the architecture and weights of the multi-frame encoder and middle +block from a pre-trained video diffusion model to build a trainable conditional network. This video ControlNet accepts +low-resolution keyframes and noisy full-frame latents as inputs. In addition to the time step t and prompt, our proposed +video-aware conditioning also includes noise augmentation level σ and downscaling factor s as additional network +conditioning inputs. + +## Hardware Requirements + ++ Operating System: Linux (requires xformers dependency) ++ Hardware: NVIDIA GPU with at least 60GB of VRAM per card. Machines such as H100, A100 are recommended. + +## Quick Start + +1. Clone the repository and install dependencies as per the official instructions: + +```shell +git clone https://github.com/Vchitect/VEnhancer.git +cd VEnhancer +## Torch and other dependencies can use those from CogVideoX. If you need to create a new environment, use the following commands: +pip install torch==2.0.1 torchvision==0.15.2 torchaudio==2.0.2 + +## Install required dependencies +pip install -r requirements.txt +``` + +Where: + +- `input_path` is the path to the input video +- `prompt` is the description of the video content. The prompt used by this tool should be shorter, not exceeding 77 + words. You may need to simplify the prompt used for generating the CogVideoX video. +- `target_fps` is the target frame rate for the video. Typically, 16 fps is already smooth, with 24 fps as the default + value. +- `up_scale` is recommend to be set to 2,3,4. The target resolution is limited to be around 2k and below. +- `noise_aug` value depends on the input video quality. Lower quality needs higher noise levels, which corresponds to + stronger refinement. 250~300 is for very low-quality videos. good videos: <= 200. +- `steps` if you want fewer steps, please change solver_mode to "normal" first, then decline the number of steps. " + fast" solver_mode has fixed steps (15). + The code will automatically download the required models from Hugging Face during execution. + +Typical runtime logs are as follows: + +```shell +/share/home/zyx/.conda/envs/cogvideox/lib/python3.10/site-packages/xformers/ops/fmha/flash.py:211: FutureWarning: `torch.library.impl_abstract` was renamed to `torch.library.register_fake`. Please use that instead; we will remove `torch.library.impl_abstract` in a future version of PyTorch. + @torch.library.impl_abstract("xformers_flash::flash_fwd") +/share/home/zyx/.conda/envs/cogvideox/lib/python3.10/site-packages/xformers/ops/fmha/flash.py:344: FutureWarning: `torch.library.impl_abstract` was renamed to `torch.library.register_fake`. Please use that instead; we will remove `torch.library.impl_abstract` in a future version of PyTorch. + @torch.library.impl_abstract("xformers_flash::flash_bwd") +2024-08-20 13:25:17,553 - video_to_video - INFO - checkpoint_path: ./ckpts/venhancer_paper.pt +/share/home/zyx/.conda/envs/cogvideox/lib/python3.10/site-packages/open_clip/factory.py:88: FutureWarning: You are using `torch.load` with `weights_only=False` (the current default value), which uses the default pickle module implicitly. It is possible to construct malicious pickle data which will execute arbitrary code during unpickling (See https://github.com/pytorch/pytorch/blob/main/SECURITY.md#untrusted-models for more details). In a future release, the default value for `weights_only` will be flipped to `True`. This limits the functions that could be executed during unpickling. Arbitrary objects will no longer be allowed to be loaded via this mode unless they are explicitly allowlisted by the user via `torch.serialization.add_safe_globals`. We recommend you start setting `weights_only=True` for any use case where you don't have full control of the loaded file. Please open an issue on GitHub for any issues related to this experimental feature. + checkpoint = torch.load(checkpoint_path, map_location=map_location) +2024-08-20 13:25:37,486 - video_to_video - INFO - Build encoder with FrozenOpenCLIPEmbedder +/share/home/zyx/Code/VEnhancer/video_to_video/video_to_video_model.py:35: FutureWarning: You are using `torch.load` with `weights_only=False` (the current default value), which uses the default pickle module implicitly. It is possible to construct malicious pickle data which will execute arbitrary code during unpickling (See https://github.com/pytorch/pytorch/blob/main/SECURITY.md#untrusted-models for more details). In a future release, the default value for `weights_only` will be flipped to `True`. This limits the functions that could be executed during unpickling. Arbitrary objects will no longer be allowed to be loaded via this mode unless they are explicitly allowlisted by the user via `torch.serialization.add_safe_globals`. We recommend you start setting `weights_only=True` for any use case where you don't have full control of the loaded file. Please open an issue on GitHub for any issues related to this experimental feature. + load_dict = torch.load(cfg.model_path, map_location='cpu') +2024-08-20 13:25:55,391 - video_to_video - INFO - Load model path ./ckpts/venhancer_paper.pt, with local status +2024-08-20 13:25:55,392 - video_to_video - INFO - Build diffusion with GaussianDiffusion +2024-08-20 13:26:16,092 - video_to_video - INFO - input video path: inputs/000000.mp4 +2024-08-20 13:26:16,093 - video_to_video - INFO - text: Wide-angle aerial shot at dawn,soft morning light casting long shadows,an elderly man walking his dog through a quiet,foggy park,trees and benches in the background,peaceful and serene atmosphere +2024-08-20 13:26:16,156 - video_to_video - INFO - input frames length: 49 +2024-08-20 13:26:16,156 - video_to_video - INFO - input fps: 8.0 +2024-08-20 13:26:16,156 - video_to_video - INFO - target_fps: 24.0 +2024-08-20 13:26:16,311 - video_to_video - INFO - input resolution: (480, 720) +2024-08-20 13:26:16,312 - video_to_video - INFO - target resolution: (1320, 1982) +2024-08-20 13:26:16,312 - video_to_video - INFO - noise augmentation: 250 +2024-08-20 13:26:16,312 - video_to_video - INFO - scale s is set to: 8 +2024-08-20 13:26:16,399 - video_to_video - INFO - video_data shape: torch.Size([145, 3, 1320, 1982]) +/share/home/zyx/Code/VEnhancer/video_to_video/video_to_video_model.py:108: FutureWarning: `torch.cuda.amp.autocast(args...)` is deprecated. Please use `torch.amp.autocast('cuda', args...)` instead. + with amp.autocast(enabled=True): +2024-08-20 13:27:19,605 - video_to_video - INFO - step: 0 +2024-08-20 13:30:12,020 - video_to_video - INFO - step: 1 +2024-08-20 13:33:04,956 - video_to_video - INFO - step: 2 +2024-08-20 13:35:58,691 - video_to_video - INFO - step: 3 +2024-08-20 13:38:51,254 - video_to_video - INFO - step: 4 +2024-08-20 13:41:44,150 - video_to_video - INFO - step: 5 +2024-08-20 13:44:37,017 - video_to_video - INFO - step: 6 +2024-08-20 13:47:30,037 - video_to_video - INFO - step: 7 +2024-08-20 13:50:22,838 - video_to_video - INFO - step: 8 +2024-08-20 13:53:15,844 - video_to_video - INFO - step: 9 +2024-08-20 13:56:08,657 - video_to_video - INFO - step: 10 +2024-08-20 13:59:01,648 - video_to_video - INFO - step: 11 +2024-08-20 14:01:54,541 - video_to_video - INFO - step: 12 +2024-08-20 14:04:47,488 - video_to_video - INFO - step: 13 +2024-08-20 14:10:13,637 - video_to_video - INFO - sampling, finished. + +``` + +Running on a single A100 GPU, enhancing each 6-second CogVideoX generated video with default settings will consume 60GB +of VRAM and take 40-50 minutes. diff --git a/tools/venhancer/README_ja.md b/tools/venhancer/README_ja.md new file mode 100644 index 0000000..70f2d74 --- /dev/null +++ b/tools/venhancer/README_ja.md @@ -0,0 +1,92 @@ + +# VEnhancer で CogVideoX によって生成されたビデオを強化する + +このチュートリアルでは、VEnhancer ツールを使用して、CogVideoX で生成されたビデオを強化し、より高いフレームレートと高い解像度を実現する方法を説明します。 + +## モデルの紹介 + +VEnhancer は、空間超解像、時間超解像（フレーム補間）、およびビデオのリファインメントを統一されたフレームワークで実現します。空間または時間の超解像のために、さまざまなアップサンプリング係数（例：1x〜8x）に柔軟に対応できます。さらに、多様なビデオアーティファクトを処理するために、リファインメント強度を変更する柔軟な制御を提供します。 + +VEnhancer は ControlNet の設計に従い、事前訓練されたビデオ拡散モデルのマルチフレームエンコーダーとミドルブロックのアーキテクチャとウェイトをコピーして、トレーニング可能な条件ネットワークを構築します。このビデオ ControlNet は、低解像度のキーフレームとノイズを含む完全なフレームを入力として受け取ります。さらに、タイムステップ t とプロンプトに加えて、提案されたビデオ対応条件により、ノイズ増幅レベル σ およびダウンスケーリングファクター s が追加のネットワーク条件として使用されます。 + +## ハードウェア要件 + ++ オペレーティングシステム: Linux (xformers 依存関係が必要) ++ ハードウェア: 単一カードあたり少なくとも 60GB の VRAM を持つ NVIDIA GPU。H100、A100 などのマシンを推奨します。 + +## クイックスタート + +1. 公式の指示に従ってリポジトリをクローンし、依存関係をインストールします。 + +```shell +git clone https://github.com/Vchitect/VEnhancer.git +cd VEnhancer +## Torch などの依存関係は CogVideoX の依存関係を使用できます。新しい環境を作成する必要がある場合は、以下のコマンドを使用してください。 +pip install torch==2.0.1 torchvision==0.15.2 torchaudio==2.0.2 + +## 必須の依存関係をインストールします。 +pip install -r requirements.txt +``` + +2. コードを実行します。 + +```shell +python enhance_a_video.py --up_scale 4 --target_fps 24 --noise_aug 250 --solver_mode 'fast' --steps 15 --input_path inputs/000000.mp4 --prompt 'Wide-angle aerial shot at dawn, soft morning light casting long shadows, an elderly man walking his dog through a quiet, foggy park, trees and benches in the background, peaceful and serene atmosphere' --save_dir 'results/' +``` + +次の設定を行います： + +- `input_path` 是输入视频的路径 +- `prompt` 是视频内容的描述。此工具使用的提示词应更短，不超过77个字。您可能需要简化用于生成CogVideoX视频的提示词。 +- `target_fps` 是视频的目标帧率。通常，16 fps已经很流畅，默认值为24 fps。 +- `up_scale` 推荐设置为2、3或4。目标分辨率限制在2k左右及以下。 +- `noise_aug` 的值取决于输入视频的质量。质量较低的视频需要更高的噪声级别，这对应于更强的优化。250~300适用于非常低质量的视频。对于高质量视频，设置为≤200。 +- `steps` 如果想减少步数，请先将solver_mode改为“normal”，然后减少步数。“fast”模式的步数是固定的（15步）。 + 代码在执行过程中会自动从Hugging Face下载所需的模型。 + +コードの実行中に、必要なモデルは Hugging Face から自動的にダウンロードされます。 + +```shell +/share/home/zyx/.conda/envs/cogvideox/lib/python3.10/site-packages/xformers/ops/fmha/flash.py:211: FutureWarning: `torch.library.impl_abstract` was renamed to `torch.library.register_fake`. Please use that instead; we will remove `torch.library.impl_abstract` in a future version of PyTorch. + @torch.library.impl_abstract("xformers_flash::flash_fwd") +/share/home/zyx/.conda/envs/cogvideox/lib/python3.10/site-packages/xformers/ops/fmha/flash.py:344: FutureWarning: `torch.library.impl_abstract` was renamed to `torch.library.register_fake`. Please use that instead; we will remove `torch.library.impl_abstract` in a future version of PyTorch. + @torch.library.impl_abstract("xformers_flash::flash_bwd") +2024-08-20 13:25:17,553 - video_to_video - INFO - checkpoint_path: ./ckpts/venhancer_paper.pt +/share/home/zyx/.conda/envs/cogvideox/lib/python3.10/site-packages/open_clip/factory.py:88: FutureWarning: You are using `torch.load` with `weights_only=False` (the current default value), which uses the default pickle module implicitly. It is possible to construct malicious pickle data which will execute arbitrary code during unpickling (See https://github.com/pytorch/pytorch/blob/main/SECURITY.md#untrusted-models for more details). In a future release, the default value for `weights_only` will be flipped to `True`. This limits the functions that could be executed during unpickling. Arbitrary objects will no longer be allowed to be loaded via this mode unless they are explicitly allowlisted by the user via `torch.serialization.add_safe_globals`. We recommend you start setting `weights_only=True` for any use case where you don't have full control of the loaded file. Please open an issue on GitHub for any issues related to this experimental feature. + checkpoint = torch.load(checkpoint_path, map_location=map_location) +2024-08-20 13:25:37,486 - video_to_video - INFO - Build encoder with FrozenOpenCLIPEmbedder +/share/home/zyx/Code/VEnhancer/video_to_video/video_to_video_model.py:35: FutureWarning: You are using `torch.load` with `weights_only=False` (the current default value), which uses the default pickle module implicitly. It is possible to construct malicious pickle data which will execute arbitrary code during unpickling (See https://github.com/pytorch/pytorch/blob/main/SECURITY.md#untrusted-models for more details). In a future release, the default value for `weights_only` will be flipped to `True`. This limits the functions that could be executed during unpickling. Arbitrary objects will no longer be allowed to be loaded via this mode unless they are explicitly allowlisted by the user via `torch.serialization.add_safe_globals`. We recommend you start setting `weights_only=True` for any use case where you don't have full control of the loaded file. Please open an issue on GitHub for any issues related to this experimental feature. + load_dict = torch.load(cfg.model_path, map_location='cpu') +2024-08-20 13:25:55,391 - video_to_video - INFO - Load model path ./ckpts/venhancer_paper.pt, with local status +2024-08-20 13:25:55,392 - video_to_video - INFO - Build diffusion with GaussianDiffusion +2024-08-20 13:26:16,092 - video_to_video - INFO - input video path: inputs/000000.mp4 +2024-08-20 13:26:16,093 - video_to_video - INFO - text: Wide-angle aerial shot at dawn,soft morning light casting long shadows,an elderly man walking his dog through a quiet,foggy park,trees and benches in the background,peaceful and serene atmosphere +2024-08-20 13:26:16,156 - video_to_video - INFO - input frames length: 49 +2024-08-20 13:26:16,156 - video_to_video - INFO - input fps: 8.0 +2024-08-20 13:26:16,156 - video_to_video - INFO - target_fps: 24.0 +2024-08-20 13:26:16,311 - video_to_video - INFO - input resolution: (480, 720) +2024-08-20 13:26:16,312 - video_to_video - INFO - target resolution: (1320, 1982) +2024-08-20 13:26:16,312 - video_to_video - INFO - noise augmentation: 250 +2024-08-20 13:26:16,312 - video_to_video - INFO - scale s is set to: 8 +2024-08-20 13:26:16,399 - video_to_video - INFO - video_data shape: torch.Size([145, 3, 1320, 1982]) +/share/home/zyx/Code/VEnhancer/video_to_video/video_to_video_model.py:108: FutureWarning: `torch.cuda.amp.autocast(args...)` is deprecated. Please use `torch.amp.autocast('cuda', args...)` instead. + with amp.autocast(enabled=True): +2024-08-20 13:27:19,605 - video_to_video - INFO - step: 0 +2024-08-20 13:30:12,020 - video_to_video - INFO - step: 1 +2024-08-20 13:33:04,956 - video_to_video - INFO - step: 2 +2024-08-20 13:35:58,691 - video_to_video - INFO - step: 3 +2024-08-20 13:38:51,254 - video_to_video - INFO - step: 4 +2024-08-20 13:41:44,150 - video_to_video - INFO - step: 5 +2024-08-20 13:44:37,017 - video_to_video - INFO - step: 6 +2024-08-20 13:47:30,037 - video_to_video - INFO - step: 7 +2024-08-20 13:50:22,838 - video_to_video - INFO - step: 8 +2024-08-20 13:53:15,844 - video_to_video - INFO - step: 9 +2024-08-20 13:56:08,657 - video_to_video - INFO - step: 10 +2024-08-20 13:59:01,648 - video_to_video - INFO - step: 11 +2024-08-20 14:01:54,541 - video_to_video - INFO - step: 12 +2024-08-20 14:04:47,488 - video_to_video - INFO - step: 13 +2024-08-20 14:10:13,637 - video_to_video - INFO - sampling, finished. + +``` + +A100 GPU を単一で使用している場合、CogVideoX によって生成された 6 秒間のビデオを強化するには、デフォルト設定で 60GB の VRAM を消費し、40〜50 分かかります。 diff --git a/tools/venhancer/README_zh.md b/tools/venhancer/README_zh.md new file mode 100644 index 0000000..0deabcc --- /dev/null +++ b/tools/venhancer/README_zh.md @@ -0,0 +1,101 @@ +# 使用 VEnhancer 对 CogVdieoX 生成视频进行增强 + +本教程将要使用 VEnhancer 工具 对 CogVdieoX 生成视频进行增强, 包括更高的帧率和更高的分辨率 + +## 模型介绍 + +VEnhancer 在一个统一的框架中实现了空间超分辨率、时间超分辨率（帧插值）和视频优化。它可以灵活地适应不同的上采样因子（例如，1x~ +8x）用于空间或时间超分辨率。此外，它提供了灵活的控制，以修改优化强度，从而处理多样化的视频伪影。 + +VEnhancer 遵循 ControlNet 的设计，复制了预训练的视频扩散模型的多帧编码器和中间块的架构和权重，构建了一个可训练的条件网络。这个视频 +ControlNet 接受低分辨率关键帧和包含噪声的完整帧作为输入。此外，除了时间步 t 和提示词外，我们提出的视频感知条件还将噪声增强的噪声级别 +σ 和降尺度因子 s 作为附加的网络条件输入。 + +## 硬件需求 + ++ 操作系统: Linux (需要依赖xformers) ++ 硬件: NVIDIA GPU 并至少保证单卡显存超过60G，推荐使用 H100，A100等机器。 + +## 快速上手 + +1. 按照官方指引克隆仓库并安装依赖 + +```shell +git clone https://github.com/Vchitect/VEnhancer.git +cd VEnhancer +## torch等依赖可以使用CogVideoX的依赖，如果你需要创建一个新的环境，可以使用以下命令 +pip install torch==2.0.1 torchvision==0.15.2 torchaudio==2.0.2 + +## 安装必须的依赖 +pip install -r requirements.txt +``` + +2. 运行代码 + +```shell +python enhance_a_video.py \ +--up_scale 4 --target_fps 24 --noise_aug 250 \ +--solver_mode 'fast' --steps 15 \ +--input_path inputs/000000.mp4 \ +--prompt 'Wide-angle aerial shot at dawn,soft morning light casting long shadows,an elderly man walking his dog through a quiet,foggy park,trees and benches in the background,peaceful and serene atmosphere' \ +--save_dir 'results/' +``` + +其中: + +- `input_path` 是输入视频的路径 +- `prompt` 是视频内容的描述。此工具使用的提示词应更短，不超过77个字。您可能需要简化用于生成CogVideoX视频的提示词。 +- `target_fps` 是视频的目标帧率。通常，16 fps已经很流畅，默认值为24 fps。 +- `up_scale` 推荐设置为2、3或4。目标分辨率限制在2k左右及以下。 +- `noise_aug` 的值取决于输入视频的质量。质量较低的视频需要更高的噪声级别，这对应于更强的优化。250~300适用于非常低质量的视频。对于高质量视频，设置为≤200。 +- `steps` 如果想减少步数，请先将solver_mode改为“normal”，然后减少步数。“fast”模式的步数是固定的（15步）。 + 代码在执行过程中会自动从Hugging Face下载所需的模型。 + +代码运行过程中，会自动从Huggingface拉取需要的模型 + +运行日志通常如下: + +```shell +/share/home/zyx/.conda/envs/cogvideox/lib/python3.10/site-packages/xformers/ops/fmha/flash.py:211: FutureWarning: `torch.library.impl_abstract` was renamed to `torch.library.register_fake`. Please use that instead; we will remove `torch.library.impl_abstract` in a future version of PyTorch. + @torch.library.impl_abstract("xformers_flash::flash_fwd") +/share/home/zyx/.conda/envs/cogvideox/lib/python3.10/site-packages/xformers/ops/fmha/flash.py:344: FutureWarning: `torch.library.impl_abstract` was renamed to `torch.library.register_fake`. Please use that instead; we will remove `torch.library.impl_abstract` in a future version of PyTorch. + @torch.library.impl_abstract("xformers_flash::flash_bwd") +2024-08-20 13:25:17,553 - video_to_video - INFO - checkpoint_path: ./ckpts/venhancer_paper.pt +/share/home/zyx/.conda/envs/cogvideox/lib/python3.10/site-packages/open_clip/factory.py:88: FutureWarning: You are using `torch.load` with `weights_only=False` (the current default value), which uses the default pickle module implicitly. It is possible to construct malicious pickle data which will execute arbitrary code during unpickling (See https://github.com/pytorch/pytorch/blob/main/SECURITY.md#untrusted-models for more details). In a future release, the default value for `weights_only` will be flipped to `True`. This limits the functions that could be executed during unpickling. Arbitrary objects will no longer be allowed to be loaded via this mode unless they are explicitly allowlisted by the user via `torch.serialization.add_safe_globals`. We recommend you start setting `weights_only=True` for any use case where you don't have full control of the loaded file. Please open an issue on GitHub for any issues related to this experimental feature. + checkpoint = torch.load(checkpoint_path, map_location=map_location) +2024-08-20 13:25:37,486 - video_to_video - INFO - Build encoder with FrozenOpenCLIPEmbedder +/share/home/zyx/Code/VEnhancer/video_to_video/video_to_video_model.py:35: FutureWarning: You are using `torch.load` with `weights_only=False` (the current default value), which uses the default pickle module implicitly. It is possible to construct malicious pickle data which will execute arbitrary code during unpickling (See https://github.com/pytorch/pytorch/blob/main/SECURITY.md#untrusted-models for more details). In a future release, the default value for `weights_only` will be flipped to `True`. This limits the functions that could be executed during unpickling. Arbitrary objects will no longer be allowed to be loaded via this mode unless they are explicitly allowlisted by the user via `torch.serialization.add_safe_globals`. We recommend you start setting `weights_only=True` for any use case where you don't have full control of the loaded file. Please open an issue on GitHub for any issues related to this experimental feature. + load_dict = torch.load(cfg.model_path, map_location='cpu') +2024-08-20 13:25:55,391 - video_to_video - INFO - Load model path ./ckpts/venhancer_paper.pt, with local status +2024-08-20 13:25:55,392 - video_to_video - INFO - Build diffusion with GaussianDiffusion +2024-08-20 13:26:16,092 - video_to_video - INFO - input video path: inputs/000000.mp4 +2024-08-20 13:26:16,093 - video_to_video - INFO - text: Wide-angle aerial shot at dawn,soft morning light casting long shadows,an elderly man walking his dog through a quiet,foggy park,trees and benches in the background,peaceful and serene atmosphere +2024-08-20 13:26:16,156 - video_to_video - INFO - input frames length: 49 +2024-08-20 13:26:16,156 - video_to_video - INFO - input fps: 8.0 +2024-08-20 13:26:16,156 - video_to_video - INFO - target_fps: 24.0 +2024-08-20 13:26:16,311 - video_to_video - INFO - input resolution: (480, 720) +2024-08-20 13:26:16,312 - video_to_video - INFO - target resolution: (1320, 1982) +2024-08-20 13:26:16,312 - video_to_video - INFO - noise augmentation: 250 +2024-08-20 13:26:16,312 - video_to_video - INFO - scale s is set to: 8 +2024-08-20 13:26:16,399 - video_to_video - INFO - video_data shape: torch.Size([145, 3, 1320, 1982]) +/share/home/zyx/Code/VEnhancer/video_to_video/video_to_video_model.py:108: FutureWarning: `torch.cuda.amp.autocast(args...)` is deprecated. Please use `torch.amp.autocast('cuda', args...)` instead. + with amp.autocast(enabled=True): +2024-08-20 13:27:19,605 - video_to_video - INFO - step: 0 +2024-08-20 13:30:12,020 - video_to_video - INFO - step: 1 +2024-08-20 13:33:04,956 - video_to_video - INFO - step: 2 +2024-08-20 13:35:58,691 - video_to_video - INFO - step: 3 +2024-08-20 13:38:51,254 - video_to_video - INFO - step: 4 +2024-08-20 13:41:44,150 - video_to_video - INFO - step: 5 +2024-08-20 13:44:37,017 - video_to_video - INFO - step: 6 +2024-08-20 13:47:30,037 - video_to_video - INFO - step: 7 +2024-08-20 13:50:22,838 - video_to_video - INFO - step: 8 +2024-08-20 13:53:15,844 - video_to_video - INFO - step: 9 +2024-08-20 13:56:08,657 - video_to_video - INFO - step: 10 +2024-08-20 13:59:01,648 - video_to_video - INFO - step: 11 +2024-08-20 14:01:54,541 - video_to_video - INFO - step: 12 +2024-08-20 14:04:47,488 - video_to_video - INFO - step: 13 +2024-08-20 14:10:13,637 - video_to_video - INFO - sampling, finished. + +``` + +使用A100单卡运行，对于每个CogVideoX生产的6秒视频，按照默认配置，会消耗60G显存，并用时40-50分钟。