chore: import upstream snapshot with attribution
This commit is contained in:
wehub-resource-sync
@@ -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`.
|
||||
@@ -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. /简单描述您期望发生的事情。"
|
||||
@@ -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链接或者其他您能提供帮助的链接。
|
||||
@@ -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
|
||||
@@ -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.
|
||||
@@ -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 与我们联系。
|
||||
@@ -0,0 +1,458 @@
|
||||
# CogVideo & CogVideoX
|
||||
|
||||
[中文阅读](./README_zh.md)
|
||||
|
||||
[日本語で読む](./README_ja.md)
|
||||
|
||||
<div align="center">
|
||||
<img src=resources/logo.svg width="50%"/>
|
||||
</div>
|
||||
<p align="center">
|
||||
Experience the CogVideoX-5B model online at <a href="https://huggingface.co/spaces/THUDM/CogVideoX-5B" target="_blank"> 🤗 Huggingface Space</a> or <a href="https://modelscope.cn/studios/ZhipuAI/CogVideoX-5b-demo" target="_blank"> 🤖 ModelScope Space</a>
|
||||
</p>
|
||||
<p align="center">
|
||||
📚 View the <a href="https://arxiv.org/abs/2408.06072" target="_blank">paper</a> and <a href="https://zhipu-ai.feishu.cn/wiki/DHCjw1TrJiTyeukfc9RceoSRnCh" target="_blank">user guide</a>
|
||||
</p>
|
||||
<p align="center">
|
||||
👋 Join our <a href="resources/WECHAT.md" target="_blank">WeChat</a> and <a href="https://discord.gg/dCGfUsagrD" target="_blank">Discord</a>
|
||||
</p>
|
||||
<p align="center">
|
||||
📍 Visit <a href="https://chatglm.cn/video?lang=en?fr=osm_cogvideo">QingYing</a> and <a href="https://open.bigmodel.cn/?utm_campaign=open&_channel_track_key=OWTVNma9">API Platform</a> to experience larger-scale commercial video generation models.
|
||||
</p>
|
||||
|
||||
## 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
|
||||
|
||||
<table border="0" style="width: 100%; text-align: left; margin-top: 20px;">
|
||||
<tr>
|
||||
<td>
|
||||
<video src="https://github.com/user-attachments/assets/cf5953ea-96d3-48fd-9907-c4708752c714" width="100%" controls autoplay loop></video>
|
||||
</td>
|
||||
<td>
|
||||
<video src="https://github.com/user-attachments/assets/fe0a78e6-b669-4800-8cf0-b5f9b5145b52" width="100%" controls autoplay loop></video>
|
||||
</td>
|
||||
<td>
|
||||
<video src="https://github.com/user-attachments/assets/c182f606-8f8c-421d-b414-8487070fcfcb" width="100%" controls autoplay loop></video>
|
||||
</td>
|
||||
<td>
|
||||
<video src="https://github.com/user-attachments/assets/7db2bbce-194d-434d-a605-350254b6c298" width="100%" controls autoplay loop></video>
|
||||
</td>
|
||||
</tr>
|
||||
<tr>
|
||||
<td>
|
||||
<video src="https://github.com/user-attachments/assets/62b01046-8cab-44cc-bd45-4d965bb615ec" width="100%" controls autoplay loop></video>
|
||||
</td>
|
||||
<td>
|
||||
<video src="https://github.com/user-attachments/assets/d78e552a-4b3f-4b81-ac3f-3898079554f6" width="100%" controls autoplay loop></video>
|
||||
</td>
|
||||
<td>
|
||||
<video src="https://github.com/user-attachments/assets/30894f12-c741-44a2-9e6e-ddcacc231e5b" width="100%" controls autoplay loop></video>
|
||||
</td>
|
||||
<td>
|
||||
<video src="https://github.com/user-attachments/assets/926575ca-7150-435b-a0ff-4900a963297b" width="100%" controls autoplay loop></video>
|
||||
</td>
|
||||
</tr>
|
||||
</table>
|
||||
|
||||
### CogVideoX-2B
|
||||
|
||||
<table border="0" style="width: 100%; text-align: left; margin-top: 20px;">
|
||||
<tr>
|
||||
<td>
|
||||
<video src="https://github.com/user-attachments/assets/ea3af39a-3160-4999-90ec-2f7863c5b0e9" width="100%" controls autoplay loop></video>
|
||||
</td>
|
||||
<td>
|
||||
<video src="https://github.com/user-attachments/assets/9de41efd-d4d1-4095-aeda-246dd834e91d" width="100%" controls autoplay loop></video>
|
||||
</td>
|
||||
<td>
|
||||
<video src="https://github.com/user-attachments/assets/941d6661-6a8d-4a1b-b912-59606f0b2841" width="100%" controls autoplay loop></video>
|
||||
</td>
|
||||
<td>
|
||||
<video src="https://github.com/user-attachments/assets/938529c4-91ae-4f60-b96b-3c3947fa63cb" width="100%" controls autoplay loop></video>
|
||||
</td>
|
||||
</tr>
|
||||
</table>
|
||||
|
||||
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.
|
||||
|
||||
<table style="border-collapse: collapse; width: 100%;">
|
||||
<tr>
|
||||
<th style="text-align: center;">Model Name</th>
|
||||
<th style="text-align: center;">CogVideoX1.5-5B (Latest)</th>
|
||||
<th style="text-align: center;">CogVideoX1.5-5B-I2V (Latest)</th>
|
||||
<th style="text-align: center;">CogVideoX-2B</th>
|
||||
<th style="text-align: center;">CogVideoX-5B</th>
|
||||
<th style="text-align: center;">CogVideoX-5B-I2V</th>
|
||||
</tr>
|
||||
<tr>
|
||||
<td style="text-align: center;">Release Date</td>
|
||||
<th style="text-align: center;">November 8, 2024</th>
|
||||
<th style="text-align: center;">November 8, 2024</th>
|
||||
<th style="text-align: center;">August 6, 2024</th>
|
||||
<th style="text-align: center;">August 27, 2024</th>
|
||||
<th style="text-align: center;">September 19, 2024</th>
|
||||
</tr>
|
||||
<tr>
|
||||
<td style="text-align: center;">Video Resolution</td>
|
||||
<td colspan="1" style="text-align: center;">1360 * 768</td>
|
||||
<td colspan="1" style="text-align: center;"> Min(W, H) = 768 <br> 768 ≤ Max(W, H) ≤ 1360 <br> Max(W, H) % 16 = 0 </td>
|
||||
<td colspan="3" style="text-align: center;">720 * 480</td>
|
||||
</tr>
|
||||
<tr>
|
||||
<td style="text-align: center;">Number of Frames</td>
|
||||
<td colspan="2" style="text-align: center;">Should be <b>16N + 1</b> where N <= 10 (default 81)</td>
|
||||
<td colspan="3" style="text-align: center;">Should be <b>8N + 1</b> where N <= 6 (default 49)</td>
|
||||
</tr>
|
||||
<tr>
|
||||
<td style="text-align: center;">Inference Precision</td>
|
||||
<td colspan="2" style="text-align: center;"><b>BF16 (Recommended)</b>, FP16, FP32, FP8*, INT8, Not supported: INT4</td>
|
||||
<td style="text-align: center;"><b>FP16*(Recommended)</b>, BF16, FP32, FP8*, INT8, Not supported: INT4</td>
|
||||
<td colspan="2" style="text-align: center;"><b>BF16 (Recommended)</b>, FP16, FP32, FP8*, INT8, Not supported: INT4</td>
|
||||
</tr>
|
||||
<tr>
|
||||
<td style="text-align: center;">Single GPU Memory Usage<br></td>
|
||||
<td colspan="2" style="text-align: center;"><a href="https://github.com/THUDM/SwissArmyTransformer">SAT</a> BF16: 76GB <br><b>diffusers BF16: from 10GB*</b><br><b>diffusers INT8(torchao): from 7GB*</b></td>
|
||||
<td style="text-align: center;"><a href="https://github.com/THUDM/SwissArmyTransformer">SAT</a> FP16: 18GB <br><b>diffusers FP16: 4GB minimum* </b><br><b>diffusers INT8 (torchao): 3.6GB minimum*</b></td>
|
||||
<td colspan="2" style="text-align: center;"><a href="https://github.com/THUDM/SwissArmyTransformer">SAT</a> BF16: 26GB <br><b>diffusers BF16 : 5GB minimum* </b><br><b>diffusers INT8 (torchao): 4.4GB minimum* </b></td>
|
||||
</tr>
|
||||
<tr>
|
||||
<td style="text-align: center;">Multi-GPU Memory Usage</td>
|
||||
<td colspan="2" style="text-align: center;"><b>BF16: 24GB* using diffusers</b><br></td>
|
||||
<td style="text-align: center;"><b>FP16: 10GB* using diffusers</b><br></td>
|
||||
<td colspan="2" style="text-align: center;"><b>BF16: 15GB* using diffusers</b><br></td>
|
||||
</tr>
|
||||
<tr>
|
||||
<td style="text-align: center;">Inference Speed<br>(Step = 50, FP/BF16)</td>
|
||||
<td colspan="2" style="text-align: center;">Single A100: ~1000 seconds (5-second video)<br>Single H100: ~550 seconds (5-second video)</td>
|
||||
<td style="text-align: center;">Single A100: ~90 seconds<br>Single H100: ~45 seconds</td>
|
||||
<td colspan="2" style="text-align: center;">Single A100: ~180 seconds<br>Single H100: ~90 seconds</td>
|
||||
</tr>
|
||||
<tr>
|
||||
<td style="text-align: center;">Prompt Language</td>
|
||||
<td colspan="5" style="text-align: center;">English*</td>
|
||||
</tr>
|
||||
<tr>
|
||||
<td style="text-align: center;">Prompt Token Limit</td>
|
||||
<td colspan="2" style="text-align: center;">224 Tokens</td>
|
||||
<td colspan="3" style="text-align: center;">226 Tokens</td>
|
||||
</tr>
|
||||
<tr>
|
||||
<td style="text-align: center;">Video Length</td>
|
||||
<td colspan="2" style="text-align: center;">5 seconds or 10 seconds</td>
|
||||
<td colspan="3" style="text-align: center;">6 seconds</td>
|
||||
</tr>
|
||||
<tr>
|
||||
<td style="text-align: center;">Frame Rate</td>
|
||||
<td colspan="2" style="text-align: center;">16 frames / second </td>
|
||||
<td colspan="3" style="text-align: center;">8 frames / second </td>
|
||||
</tr>
|
||||
<tr>
|
||||
<td style="text-align: center;">Position Encoding</td>
|
||||
<td colspan="2" style="text-align: center;">3d_rope_pos_embed</td>
|
||||
<td style="text-align: center;">3d_sincos_pos_embed</td>
|
||||
<td style="text-align: center;">3d_rope_pos_embed</td>
|
||||
<td style="text-align: center;">3d_rope_pos_embed + learnable_pos_embed</td>
|
||||
</tr>
|
||||
<tr>
|
||||
<td style="text-align: center;">Download Link (Diffusers)</td>
|
||||
<td style="text-align: center;"><a href="https://huggingface.co/THUDM/CogVideoX1.5-5B">🤗 HuggingFace</a><br><a href="https://modelscope.cn/models/ZhipuAI/CogVideoX1.5-5B">🤖 ModelScope</a><br><a href="https://wisemodel.cn/models/ZhipuAI/CogVideoX1.5-5B">🟣 WiseModel</a></td>
|
||||
<td style="text-align: center;"><a href="https://huggingface.co/THUDM/CogVideoX1.5-5B-I2V">🤗 HuggingFace</a><br><a href="https://modelscope.cn/models/ZhipuAI/CogVideoX1.5-5B-I2V">🤖 ModelScope</a><br><a href="https://wisemodel.cn/models/ZhipuAI/CogVideoX1.5-5B-I2V">🟣 WiseModel</a></td>
|
||||
<td style="text-align: center;"><a href="https://huggingface.co/THUDM/CogVideoX-2b">🤗 HuggingFace</a><br><a href="https://modelscope.cn/models/ZhipuAI/CogVideoX-2b">🤖 ModelScope</a><br><a href="https://wisemodel.cn/models/ZhipuAI/CogVideoX-2b">🟣 WiseModel</a></td>
|
||||
<td style="text-align: center;"><a href="https://huggingface.co/THUDM/CogVideoX-5b">🤗 HuggingFace</a><br><a href="https://modelscope.cn/models/ZhipuAI/CogVideoX-5b">🤖 ModelScope</a><br><a href="https://wisemodel.cn/models/ZhipuAI/CogVideoX-5b">🟣 WiseModel</a></td>
|
||||
<td style="text-align: center;"><a href="https://huggingface.co/THUDM/CogVideoX-5b-I2V">🤗 HuggingFace</a><br><a href="https://modelscope.cn/models/ZhipuAI/CogVideoX-5b-I2V">🤖 ModelScope</a><br><a href="https://wisemodel.cn/models/ZhipuAI/CogVideoX-5b-I2V">🟣 WiseModel</a></td>
|
||||
</tr>
|
||||
<tr>
|
||||
<td style="text-align: center;">Download Link (SAT)</td>
|
||||
<td colspan="2" style="text-align: center;"><a href="https://huggingface.co/THUDM/CogVideoX1.5-5b-SAT">🤗 HuggingFace</a><br><a href="https://modelscope.cn/models/ZhipuAI/CogVideoX1.5-5b-SAT">🤖 ModelScope</a><br><a href="https://wisemodel.cn/models/ZhipuAI/CogVideoX1.5-5b-SAT">🟣 WiseModel</a></td>
|
||||
<td colspan="3" style="text-align: center;"><a href="./sat/README_zh.md">SAT</a></td>
|
||||
</tr>
|
||||
</table>
|
||||
|
||||
**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.
|
||||
|
||||
<div style="text-align: center;">
|
||||
<img src="resources/web_demo.png" style="width: 100%; height: auto;" />
|
||||
</div>
|
||||
|
||||
+ [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.
|
||||
|
||||
|
||||
|
||||
|
||||
<div align="center">
|
||||
<video src="https://github.com/user-attachments/assets/2fa19651-e925-4a2a-b8d6-b3f216d490ba" width="80%" controls autoplay></video>
|
||||
</div>
|
||||
|
||||
|
||||
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).
|
||||
@@ -0,0 +1,7 @@
|
||||
# WeHub 来源说明
|
||||
|
||||
- 原始项目：`zai-org/CogVideo`
|
||||
- 原始仓库：https://github.com/zai-org/CogVideo
|
||||
- 导入方式：上游默认分支的最新快照
|
||||
- 原作者、版权和许可证信息以原始仓库及本仓库 LICENSE 为准
|
||||
- 本文件仅用于记录来源，不代表 WeHub 是原项目作者
|
||||
+432
@@ -0,0 +1,432 @@
|
||||
# CogVideo & CogVideoX
|
||||
|
||||
[Read this in English](./README.md)
|
||||
|
||||
[中文阅读](./README_zh.md)
|
||||
|
||||
<div align="center">
|
||||
<img src=resources/logo.svg width="50%"/>
|
||||
</div>
|
||||
<p align="center">
|
||||
<a href="https://huggingface.co/spaces/THUDM/CogVideoX-5B" target="_blank"> 🤗 Huggingface Space</a> または <a href="https://modelscope.cn/studios/ZhipuAI/CogVideoX-5b-demo" target="_blank"> 🤖 ModelScope Space</a> で CogVideoX-5B モデルをオンラインで体験してください
|
||||
</p>
|
||||
<p align="center">
|
||||
📚 <a href="https://arxiv.org/abs/2408.06072" target="_blank">論文</a>と<a href="https://zhipu-ai.feishu.cn/wiki/DHCjw1TrJiTyeukfc9RceoSRnCh" target="_blank">使用ドキュメント</a>を表示します。
|
||||
</p>
|
||||
<p align="center">
|
||||
👋 <a href="resources/WECHAT.md" target="_blank">WeChat</a> と <a href="https://discord.gg/dCGfUsagrD" target="_blank">Discord</a> に参加
|
||||
</p>
|
||||
<p align="center">
|
||||
📍 <a href="https://chatglm.cn/video?lang=en?fr=osm_cogvideo">清影</a> と <a href="https://open.bigmodel.cn/?utm_campaign=open&_channel_track_key=OWTVNma9">APIプラットフォーム</a> を訪問して、より大規模な商用ビデオ生成モデルを体験.
|
||||
</p>
|
||||
|
||||
## 更新とニュース
|
||||
- 🔥🔥 ```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
|
||||
|
||||
<table border="0" style="width: 100%; text-align: left; margin-top: 20px;">
|
||||
<tr>
|
||||
<td>
|
||||
<video src="https://github.com/user-attachments/assets/cf5953ea-96d3-48fd-9907-c4708752c714" width="100%" controls autoplay loop></video>
|
||||
</td>
|
||||
<td>
|
||||
<video src="https://github.com/user-attachments/assets/fe0a78e6-b669-4800-8cf0-b5f9b5145b52" width="100%" controls autoplay loop></video>
|
||||
</td>
|
||||
<td>
|
||||
<video src="https://github.com/user-attachments/assets/c182f606-8f8c-421d-b414-8487070fcfcb" width="100%" controls autoplay loop></video>
|
||||
</td>
|
||||
<td>
|
||||
<video src="https://github.com/user-attachments/assets/7db2bbce-194d-434d-a605-350254b6c298" width="100%" controls autoplay loop></video>
|
||||
</td>
|
||||
</tr>
|
||||
<tr>
|
||||
<td>
|
||||
<video src="https://github.com/user-attachments/assets/62b01046-8cab-44cc-bd45-4d965bb615ec" width="100%" controls autoplay loop></video>
|
||||
</td>
|
||||
<td>
|
||||
<video src="https://github.com/user-attachments/assets/d78e552a-4b3f-4b81-ac3f-3898079554f6" width="100%" controls autoplay loop></video>
|
||||
</td>
|
||||
<td>
|
||||
<video src="https://github.com/user-attachments/assets/30894f12-c741-44a2-9e6e-ddcacc231e5b" width="100%" controls autoplay loop></video>
|
||||
</td>
|
||||
<td>
|
||||
<video src="https://github.com/user-attachments/assets/926575ca-7150-435b-a0ff-4900a963297b" width="100%" controls autoplay loop></video>
|
||||
</td>
|
||||
</tr>
|
||||
</table>
|
||||
|
||||
### CogVideoX-2B
|
||||
|
||||
<table border="0" style="width: 100%; text-align: left; margin-top: 20px;">
|
||||
<tr>
|
||||
<td>
|
||||
<video src="https://github.com/user-attachments/assets/ea3af39a-3160-4999-90ec-2f7863c5b0e9" width="100%" controls autoplay loop></video>
|
||||
</td>
|
||||
<td>
|
||||
<video src="https://github.com/user-attachments/assets/9de41efd-d4d1-4095-aeda-246dd834e91d" width="100%" controls autoplay loop></video>
|
||||
</td>
|
||||
<td>
|
||||
<video src="https://github.com/user-attachments/assets/941d6661-6a8d-4a1b-b912-59606f0b2841" width="100%" controls autoplay loop></video>
|
||||
</td>
|
||||
<td>
|
||||
<video src="https://github.com/user-attachments/assets/938529c4-91ae-4f60-b96b-3c3947fa63cb" width="100%" controls autoplay loop></video>
|
||||
</td>
|
||||
</tr>
|
||||
</table>
|
||||
|
||||
ギャラリーの対応するプロンプトワードを表示するには、[こちら](resources/galary_prompt.md)をクリックしてください
|
||||
|
||||
## モデル紹介
|
||||
|
||||
CogVideoXは、[清影](https://chatglm.cn/video?fr=osm_cogvideox) と同源のオープンソース版ビデオ生成モデルです。
|
||||
以下の表に、提供しているビデオ生成モデルの基本情報を示します:
|
||||
|
||||
<table style="border-collapse: collapse; width: 100%;">
|
||||
<tr>
|
||||
<th style="text-align: center;">モデル名</th>
|
||||
<th style="text-align: center;">CogVideoX1.5-5B (最新)</th>
|
||||
<th style="text-align: center;">CogVideoX1.5-5B-I2V (最新)</th>
|
||||
<th style="text-align: center;">CogVideoX-2B</th>
|
||||
<th style="text-align: center;">CogVideoX-5B</th>
|
||||
<th style="text-align: center;">CogVideoX-5B-I2V</th>
|
||||
</tr>
|
||||
<tr>
|
||||
<td style="text-align: center;">公開日</td>
|
||||
<th style="text-align: center;">2024年11月8日</th>
|
||||
<th style="text-align: center;">2024年11月8日</th>
|
||||
<th style="text-align: center;">2024年8月6日</th>
|
||||
<th style="text-align: center;">2024年8月27日</th>
|
||||
<th style="text-align: center;">2024年9月19日</th>
|
||||
</tr>
|
||||
<tr>
|
||||
<td style="text-align: center;">ビデオ解像度</td>
|
||||
<td colspan="1" style="text-align: center;">1360 * 768</td>
|
||||
<td colspan="1" style="text-align: center;"> Min(W, H) = 768 <br> 768 ≤ Max(W, H) ≤ 1360 <br> Max(W, H) % 16 = 0 </td>
|
||||
<td colspan="3" style="text-align: center;">720 * 480</td>
|
||||
</tr>
|
||||
<tr>
|
||||
<td style="text-align: center;">フレーム数</td>
|
||||
<td colspan="2" style="text-align: center;"><b>16N + 1</b> (N <= 10) である必要があります (デフォルト 81)</td>
|
||||
<td colspan="3" style="text-align: center;"><b>8N + 1</b> (N <= 6) である必要があります (デフォルト 49)</td>
|
||||
</tr>
|
||||
<tr>
|
||||
<td style="text-align: center;">推論精度</td>
|
||||
<td colspan="2" style="text-align: center;"><b>BF16(推奨)</b>, FP16, FP32，FP8*，INT8，INT4非対応</td>
|
||||
<td style="text-align: center;"><b>FP16*(推奨)</b>, BF16, FP32，FP8*，INT8，INT4非対応</td>
|
||||
<td colspan="2" style="text-align: center;"><b>BF16(推奨)</b>, FP16, FP32，FP8*，INT8，INT4非対応</td>
|
||||
</tr>
|
||||
<tr>
|
||||
<td style="text-align: center;">単一GPUメモリ消費量<br></td>
|
||||
<td colspan="2" style="text-align: center;"><a href="https://github.com/THUDM/SwissArmyTransformer">SAT</a> BF16: 76GB <br><b>diffusers BF16：10GBから*</b><br><b>diffusers INT8(torchao)：7GBから*</b></td>
|
||||
<td style="text-align: center;"><a href="https://github.com/THUDM/SwissArmyTransformer">SAT</a> FP16: 18GB <br><b>diffusers FP16: 4GB以上* </b><br><b>diffusers INT8(torchao): 3.6GB以上*</b></td>
|
||||
<td colspan="2" style="text-align: center;"><a href="https://github.com/THUDM/SwissArmyTransformer">SAT</a> BF16: 26GB <br><b>diffusers BF16 : 5GB以上* </b><br><b>diffusers INT8(torchao): 4.4GB以上* </b></td>
|
||||
</tr>
|
||||
<tr>
|
||||
<td style="text-align: center;">複数GPU推論メモリ消費量</td>
|
||||
<td colspan="2" style="text-align: center;"><b>BF16: 24GB* using diffusers</b><br></td>
|
||||
<td style="text-align: center;"><b>FP16: 10GB* diffusers使用</b><br></td>
|
||||
<td colspan="2" style="text-align: center;"><b>BF16: 15GB* diffusers使用</b><br></td>
|
||||
</tr>
|
||||
<tr>
|
||||
<td style="text-align: center;">推論速度<br>(Step = 50, FP/BF16)</td>
|
||||
<td colspan="2" style="text-align: center;">シングルA100: ~1000秒(5秒ビデオ)<br>シングルH100: ~550秒(5秒ビデオ)</td>
|
||||
<td style="text-align: center;">シングルA100: ~90秒<br>シングルH100: ~45秒</td>
|
||||
<td colspan="2" style="text-align: center;">シングルA100: ~180秒<br>シングルH100: ~90秒</td>
|
||||
</tr>
|
||||
<tr>
|
||||
<td style="text-align: center;">プロンプト言語</td>
|
||||
<td colspan="5" style="text-align: center;">英語*</td>
|
||||
</tr>
|
||||
<tr>
|
||||
<td style="text-align: center;">プロンプト長さの上限</td>
|
||||
<td colspan="2" style="text-align: center;">224トークン</td>
|
||||
<td colspan="3" style="text-align: center;">226トークン</td>
|
||||
</tr>
|
||||
<tr>
|
||||
<td style="text-align: center;">ビデオ長さ</td>
|
||||
<td colspan="2" style="text-align: center;">5秒または10秒</td>
|
||||
<td colspan="3" style="text-align: center;">6秒</td>
|
||||
</tr>
|
||||
<tr>
|
||||
<td style="text-align: center;">フレームレート</td>
|
||||
<td colspan="2" style="text-align: center;">16フレーム/秒</td>
|
||||
<td colspan="3" style="text-align: center;">8フレーム/秒</td>
|
||||
</tr>
|
||||
<tr>
|
||||
<td style="text-align: center;">位置エンコーディング</td>
|
||||
<td colspan="2" style="text-align: center;">3d_rope_pos_embed</td>
|
||||
<td style="text-align: center;">3d_sincos_pos_embed</td>
|
||||
<td style="text-align: center;">3d_rope_pos_embed</td>
|
||||
<td style="text-align: center;">3d_rope_pos_embed + learnable_pos_embed</td>
|
||||
</tr>
|
||||
<tr>
|
||||
<td style="text-align: center;">ダウンロードリンク (Diffusers)</td>
|
||||
<td style="text-align: center;"><a href="https://huggingface.co/THUDM/CogVideoX1.5-5B">🤗 HuggingFace</a><br><a href="https://modelscope.cn/models/ZhipuAI/CogVideoX1.5-5B">🤖 ModelScope</a><br><a href="https://wisemodel.cn/models/ZhipuAI/CogVideoX1.5-5B">🟣 WiseModel</a></td>
|
||||
<td style="text-align: center;"><a href="https://huggingface.co/THUDM/CogVideoX1.5-5B-I2V">🤗 HuggingFace</a><br><a href="https://modelscope.cn/models/ZhipuAI/CogVideoX1.5-5B-I2V">🤖 ModelScope</a><br><a href="https://wisemodel.cn/models/ZhipuAI/CogVideoX1.5-5B-I2V">🟣 WiseModel</a></td>
|
||||
<td style="text-align: center;"><a href="https://huggingface.co/THUDM/CogVideoX-2b">🤗 HuggingFace</a><br><a href="https://modelscope.cn/models/ZhipuAI/CogVideoX-2b">🤖 ModelScope</a><br><a href="https://wisemodel.cn/models/ZhipuAI/CogVideoX-2b">🟣 WiseModel</a></td>
|
||||
<td style="text-align: center;"><a href="https://huggingface.co/THUDM/CogVideoX-5b">🤗 HuggingFace</a><br><a href="https://modelscope.cn/models/ZhipuAI/CogVideoX-5b">🤖 ModelScope</a><br><a href="https://wisemodel.cn/models/ZhipuAI/CogVideoX-5b">🟣 WiseModel</a></td>
|
||||
<td style="text-align: center;"><a href="https://huggingface.co/THUDM/CogVideoX-5b-I2V">🤗 HuggingFace</a><br><a href="https://modelscope.cn/models/ZhipuAI/CogVideoX-5b-I2V">🤖 ModelScope</a><br><a href="https://wisemodel.cn/models/ZhipuAI/CogVideoX-5b-I2V">🟣 WiseModel</a></td>
|
||||
</tr>
|
||||
<tr>
|
||||
<td style="text-align: center;">ダウンロードリンク (SAT)</td>
|
||||
<td colspan="2" style="text-align: center;"><a href="https://huggingface.co/THUDM/CogVideoX1.5-5b-SAT">🤗 HuggingFace</a><br><a href="https://modelscope.cn/models/ZhipuAI/CogVideoX1.5-5b-SAT">🤖 ModelScope</a><br><a href="https://wisemodel.cn/models/ZhipuAI/CogVideoX1.5-5b-SAT">🟣 WiseModel</a></td>
|
||||
<td colspan="3" style="text-align: center;"><a href="./sat/README_zh.md">SAT</a></td>
|
||||
</tr>
|
||||
</table>
|
||||
|
||||
**データ解説**
|
||||
|
||||
+ 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コードで、フレーム補間や超解像ツールが組み込まれています。
|
||||
|
||||
<div style="text-align: center;">
|
||||
<img src="resources/web_demo.png" style="width: 100%; height: auto;" />
|
||||
</div>
|
||||
|
||||
+ [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秒間のクリップが以下に示されています。
|
||||
|
||||
|
||||
|
||||
|
||||
<div align="center">
|
||||
<video src="https://github.com/user-attachments/assets/2fa19651-e925-4a2a-b8d6-b3f216d490ba" width="80%" controls autoplay></video>
|
||||
</div>
|
||||
|
||||
|
||||
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) の下で公開されています。
|
||||
+412
@@ -0,0 +1,412 @@
|
||||
# CogVideo & CogVideoX
|
||||
|
||||
[Read this in English](./README.md)
|
||||
|
||||
[日本語で読む](./README_ja.md)
|
||||
|
||||
<div align="center">
|
||||
<img src=resources/logo.svg width="50%"/>
|
||||
</div>
|
||||
<p align="center">
|
||||
在 <a href="https://huggingface.co/spaces/THUDM/CogVideoX-5B" target="_blank"> 🤗 Huggingface Space</a> 或 <a href="https://modelscope.cn/studios/ZhipuAI/CogVideoX-5b-demo" target="_blank"> 🤖 ModelScope Space</a> 在线体验 CogVideoX-5B 模型
|
||||
</p>
|
||||
<p align="center">
|
||||
📚 查看 <a href="https://arxiv.org/abs/2408.06072" target="_blank">论文</a> 和 <a href="https://zhipu-ai.feishu.cn/wiki/DHCjw1TrJiTyeukfc9RceoSRnCh" target="_blank">使用文档</a>
|
||||
</p>
|
||||
<p align="center">
|
||||
👋 加入我们的 <a href="resources/WECHAT.md" target="_blank">微信</a> 和 <a href="https://discord.gg/dCGfUsagrD" target="_blank">Discord</a>
|
||||
</p>
|
||||
<p align="center">
|
||||
📍 前往<a href="https://chatglm.cn/video?fr=osm_cogvideox"> 清影</a> 和 <a href="https://open.bigmodel.cn/?utm_campaign=open&_channel_track_key=OWTVNma9"> API平台</a> 体验更大规模的商业版视频生成模型。
|
||||
</p>
|
||||
|
||||
## 项目更新
|
||||
|
||||
- 🔥🔥 **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
|
||||
|
||||
<table border="0" style="width: 100%; text-align: left; margin-top: 20px;">
|
||||
<tr>
|
||||
<td>
|
||||
<video src="https://github.com/user-attachments/assets/cf5953ea-96d3-48fd-9907-c4708752c714" width="100%" controls autoplay loop></video>
|
||||
</td>
|
||||
<td>
|
||||
<video src="https://github.com/user-attachments/assets/fe0a78e6-b669-4800-8cf0-b5f9b5145b52" width="100%" controls autoplay loop></video>
|
||||
</td>
|
||||
<td>
|
||||
<video src="https://github.com/user-attachments/assets/c182f606-8f8c-421d-b414-8487070fcfcb" width="100%" controls autoplay loop></video>
|
||||
</td>
|
||||
<td>
|
||||
<video src="https://github.com/user-attachments/assets/7db2bbce-194d-434d-a605-350254b6c298" width="100%" controls autoplay loop></video>
|
||||
</td>
|
||||
</tr>
|
||||
<tr>
|
||||
<td>
|
||||
<video src="https://github.com/user-attachments/assets/62b01046-8cab-44cc-bd45-4d965bb615ec" width="100%" controls autoplay loop></video>
|
||||
</td>
|
||||
<td>
|
||||
<video src="https://github.com/user-attachments/assets/d78e552a-4b3f-4b81-ac3f-3898079554f6" width="100%" controls autoplay loop></video>
|
||||
</td>
|
||||
<td>
|
||||
<video src="https://github.com/user-attachments/assets/30894f12-c741-44a2-9e6e-ddcacc231e5b" width="100%" controls autoplay loop></video>
|
||||
</td>
|
||||
<td>
|
||||
<video src="https://github.com/user-attachments/assets/926575ca-7150-435b-a0ff-4900a963297b" width="100%" controls autoplay loop></video>
|
||||
</td>
|
||||
</tr>
|
||||
</table>
|
||||
|
||||
### CogVideoX-2B
|
||||
|
||||
<table border="0" style="width: 100%; text-align: left; margin-top: 20px;">
|
||||
<tr>
|
||||
<td>
|
||||
<video src="https://github.com/user-attachments/assets/ea3af39a-3160-4999-90ec-2f7863c5b0e9" width="100%" controls autoplay loop></video>
|
||||
</td>
|
||||
<td>
|
||||
<video src="https://github.com/user-attachments/assets/9de41efd-d4d1-4095-aeda-246dd834e91d" width="100%" controls autoplay loop></video>
|
||||
</td>
|
||||
<td>
|
||||
<video src="https://github.com/user-attachments/assets/941d6661-6a8d-4a1b-b912-59606f0b2841" width="100%" controls autoplay loop></video>
|
||||
</td>
|
||||
<td>
|
||||
<video src="https://github.com/user-attachments/assets/938529c4-91ae-4f60-b96b-3c3947fa63cb" width="100%" controls autoplay loop></video>
|
||||
</td>
|
||||
</tr>
|
||||
</table>
|
||||
|
||||
|
||||
查看画廊的对应提示词，请点击[这里](resources/galary_prompt.md)
|
||||
|
||||
## 模型介绍
|
||||
|
||||
CogVideoX是 [清影](https://chatglm.cn/video?fr=osm_cogvideox) 同源的开源版本视频生成模型。
|
||||
下表展示我们提供的视频生成模型相关基础信息:
|
||||
|
||||
<table style="border-collapse: collapse; width: 100%;">
|
||||
<tr>
|
||||
<th style="text-align: center;">模型名</th>
|
||||
<th style="text-align: center;">CogVideoX1.5-5B (最新)</th>
|
||||
<th style="text-align: center;">CogVideoX1.5-5B-I2V (最新)</th>
|
||||
<th style="text-align: center;">CogVideoX-2B</th>
|
||||
<th style="text-align: center;">CogVideoX-5B</th>
|
||||
<th style="text-align: center;">CogVideoX-5B-I2V </th>
|
||||
</tr>
|
||||
<tr>
|
||||
<td style="text-align: center;">发布时间</td>
|
||||
<th style="text-align: center;">2024年11月8日</th>
|
||||
<th style="text-align: center;">2024年11月8日</th>
|
||||
<th style="text-align: center;">2024年8月6日</th>
|
||||
<th style="text-align: center;">2024年8月27日</th>
|
||||
<th style="text-align: center;">2024年9月19日</th>
|
||||
</tr>
|
||||
<tr>
|
||||
<td style="text-align: center;">视频分辨率</td>
|
||||
<td colspan="1" style="text-align: center;">1360 * 768</td>
|
||||
<td colspan="1" style="text-align: center;"> Min(W, H) = 768 <br> 768 ≤ Max(W, H) ≤ 1360 <br> Max(W, H) % 16 = 0 </td>
|
||||
<td colspan="3" style="text-align: center;">720 * 480</td>
|
||||
</tr>
|
||||
<tr>
|
||||
<td style="text-align: center;">帧数</td>
|
||||
<td colspan="2" style="text-align: center;">必须为 <b>16N + 1</b> 其中 N <= 10 (默认 81)</td>
|
||||
<td colspan="3" style="text-align: center;">必须为 <b>8N + 1</b> 其中 N <= 6 (默认 49)</td>
|
||||
</tr>
|
||||
<tr>
|
||||
<td style="text-align: center;">推理精度</td>
|
||||
<td colspan="2" style="text-align: center;"><b>BF16(推荐)</b>, FP16, FP32，FP8*，INT8，不支持INT4</td>
|
||||
<td style="text-align: center;"><b>FP16*(推荐)</b>, BF16, FP32，FP8*，INT8，不支持INT4</td>
|
||||
<td colspan="2" style="text-align: center;"><b>BF16(推荐)</b>, FP16, FP32，FP8*，INT8，不支持INT4</td>
|
||||
</tr>
|
||||
<tr>
|
||||
<td style="text-align: center;">单GPU显存消耗<br></td>
|
||||
<td colspan="2" style="text-align: center;"><a href="https://github.com/THUDM/SwissArmyTransformer">SAT</a> BF16: 76GB <br><b>diffusers BF16 : 10GB起* </b><br><b>diffusers INT8(torchao): 7G起* </b></td>
|
||||
<td style="text-align: center;"><a href="https://github.com/THUDM/SwissArmyTransformer">SAT</a> FP16: 18GB <br><b>diffusers FP16: 4GB起* </b><br><b>diffusers INT8(torchao): 3.6G起*</b></td>
|
||||
<td colspan="2" style="text-align: center;"><a href="https://github.com/THUDM/SwissArmyTransformer">SAT</a> BF16: 26GB <br><b>diffusers BF16 : 5GB起* </b><br><b>diffusers INT8(torchao): 4.4G起* </b></td>
|
||||
</tr>
|
||||
<tr>
|
||||
<td style="text-align: center;">多GPU推理显存消耗</td>
|
||||
<td colspan="2" style="text-align: center;"><b>BF16: 24GB* using diffusers</b><br></td>
|
||||
<td style="text-align: center;"><b>FP16: 10GB* using diffusers</b><br></td>
|
||||
<td colspan="2" style="text-align: center;"><b>BF16: 15GB* using diffusers</b><br></td>
|
||||
</tr>
|
||||
<tr>
|
||||
<td style="text-align: center;">推理速度<br>(Step = 50, FP/BF16)</td>
|
||||
<td colspan="2" style="text-align: center;">单卡A100: ~1000秒(5秒视频)<br>单卡H100: ~550秒(5秒视频)</td>
|
||||
<td style="text-align: center;">单卡A100: ~90秒<br>单卡H100: ~45秒</td>
|
||||
<td colspan="2" style="text-align: center;">单卡A100: ~180秒<br>单卡H100: ~90秒</td>
|
||||
</tr>
|
||||
<tr>
|
||||
<td style="text-align: center;">提示词语言</td>
|
||||
<td colspan="5" style="text-align: center;">English*</td>
|
||||
</tr>
|
||||
<tr>
|
||||
<td style="text-align: center;">提示词长度上限</td>
|
||||
<td colspan="2" style="text-align: center;">224 Tokens</td>
|
||||
<td colspan="3" style="text-align: center;">226 Tokens</td>
|
||||
</tr>
|
||||
<tr>
|
||||
<td style="text-align: center;">视频长度</td>
|
||||
<td colspan="2" style="text-align: center;">5 秒 或 10 秒</td>
|
||||
<td colspan="3" style="text-align: center;">6 秒</td>
|
||||
</tr>
|
||||
<tr>
|
||||
<td style="text-align: center;">帧率</td>
|
||||
<td colspan="2" style="text-align: center;">16 帧 / 秒 </td>
|
||||
<td colspan="3" style="text-align: center;">8 帧 / 秒 </td>
|
||||
</tr>
|
||||
<tr>
|
||||
<td style="text-align: center;">位置编码</td>
|
||||
<td colspan="2" style="text-align: center;">3d_rope_pos_embed</td>
|
||||
<td style="text-align: center;">3d_sincos_pos_embed</td>
|
||||
<td style="text-align: center;">3d_rope_pos_embed</td>
|
||||
<td style="text-align: center;">3d_rope_pos_embed + learnable_pos_embed</td>
|
||||
</tr>
|
||||
<tr>
|
||||
<td style="text-align: center;">下载链接 (Diffusers)</td>
|
||||
<td style="text-align: center;"><a href="https://huggingface.co/THUDM/CogVideoX1.5-5B">🤗 HuggingFace</a><br><a href="https://modelscope.cn/models/ZhipuAI/CogVideoX1.5-5B">🤖 ModelScope</a><br><a href="https://wisemodel.cn/models/ZhipuAI/CogVideoX1.5-5B">🟣 WiseModel</a></td>
|
||||
<td style="text-align: center;"><a href="https://huggingface.co/THUDM/CogVideoX1.5-5B-I2V">🤗 HuggingFace</a><br><a href="https://modelscope.cn/models/ZhipuAI/CogVideoX1.5-5B-I2V">🤖 ModelScope</a><br><a href="https://wisemodel.cn/models/ZhipuAI/CogVideoX1.5-5B-I2V">🟣 WiseModel</a></td>
|
||||
<td style="text-align: center;"><a href="https://huggingface.co/THUDM/CogVideoX-2b">🤗 HuggingFace</a><br><a href="https://modelscope.cn/models/ZhipuAI/CogVideoX-2b">🤖 ModelScope</a><br><a href="https://wisemodel.cn/models/ZhipuAI/CogVideoX-2b">🟣 WiseModel</a></td>
|
||||
<td style="text-align: center;"><a href="https://huggingface.co/THUDM/CogVideoX-5b">🤗 HuggingFace</a><br><a href="https://modelscope.cn/models/ZhipuAI/CogVideoX-5b">🤖 ModelScope</a><br><a href="https://wisemodel.cn/models/ZhipuAI/CogVideoX-5b">🟣 WiseModel</a></td>
|
||||
<td style="text-align: center;"><a href="https://huggingface.co/THUDM/CogVideoX-5b-I2V">🤗 HuggingFace</a><br><a href="https://modelscope.cn/models/ZhipuAI/CogVideoX-5b-I2V">🤖 ModelScope</a><br><a href="https://wisemodel.cn/models/ZhipuAI/CogVideoX-5b-I2V">🟣 WiseModel</a></td>
|
||||
</tr>
|
||||
<tr>
|
||||
<td style="text-align: center;">下载链接 (SAT)</td>
|
||||
<td colspan="2" style="text-align: center;"><a href="https://huggingface.co/THUDM/CogVideoX1.5-5b-SAT">🤗 HuggingFace</a><br><a href="https://modelscope.cn/models/ZhipuAI/CogVideoX1.5-5b-SAT">🤖 ModelScope</a><br><a href="https://wisemodel.cn/models/ZhipuAI/CogVideoX1.5-5b-SAT">🟣 WiseModel</a></td>
|
||||
<td colspan="3" style="text-align: center;"><a href="./sat/README_zh.md">SAT</a></td>
|
||||
</tr>
|
||||
</table>
|
||||
|
||||
**数据解释**
|
||||
|
||||
+ 使用 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 代码，植入了插帧，超分工具。
|
||||
|
||||
<div style="text-align: center;">
|
||||
<img src="resources/web_demo.png" style="width: 100%; height: auto;" />
|
||||
</div>
|
||||
|
||||
+ [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秒视频。**
|
||||
|
||||
|
||||
|
||||
|
||||
|
||||
|
||||
<div align="center">
|
||||
<video src="https://github.com/user-attachments/assets/ea3af39a-3160-4999-90ec-2f7863c5b0e9" width="80%" controls autoplay></video>
|
||||
</div>
|
||||
|
||||
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)
|
||||
许可证发布。
|
||||
@@ -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.
|
||||
@@ -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以上に設定することをお勧めします。
|
||||
@@ -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 及以上的设置。
|
||||
@@ -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
|
||||
@@ -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
|
||||
}
|
||||
@@ -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
|
||||
}
|
||||
@@ -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
|
||||
}
|
||||
@@ -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
|
||||
}
|
||||
@@ -0,0 +1,2 @@
|
||||
LOG_NAME = "trainer"
|
||||
LOG_LEVEL = "INFO"
|
||||
@@ -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",
|
||||
]
|
||||
@@ -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] = []
|
||||
@@ -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)
|
||||
@@ -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)
|
||||
@@ -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
|
||||
@@ -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__)
|
||||
@@ -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)
|
||||
@@ -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)
|
||||
@@ -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)
|
||||
@@ -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)
|
||||
@@ -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)
|
||||
@@ -0,0 +1,9 @@
|
||||
from ..cogvideox_i2v.lora_trainer import CogVideoXI2VLoraTrainer
|
||||
from ..utils import register
|
||||
|
||||
|
||||
class CogVideoXI2VSftTrainer(CogVideoXI2VLoraTrainer):
|
||||
pass
|
||||
|
||||
|
||||
register("cogvideox-i2v", "sft", CogVideoXI2VSftTrainer)
|
||||
@@ -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)
|
||||
@@ -0,0 +1,9 @@
|
||||
from ..cogvideox_t2v.lora_trainer import CogVideoXT2VLoraTrainer
|
||||
from ..utils import register
|
||||
|
||||
|
||||
class CogVideoXT2VSftTrainer(CogVideoXT2VLoraTrainer):
|
||||
pass
|
||||
|
||||
|
||||
register("cogvideox-t2v", "sft", CogVideoXT2VSftTrainer)
|
||||
@@ -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]
|
||||
@@ -0,0 +1,6 @@
|
||||
from .args import Args
|
||||
from .components import Components
|
||||
from .state import State
|
||||
|
||||
|
||||
__all__ = ["Args", "State", "Components"]
|
||||
@@ -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))
|
||||
@@ -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
|
||||
@@ -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
|
||||
@@ -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")
|
||||
@@ -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()
|
||||
@@ -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[@]}"
|
||||
@@ -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[@]}"
|
||||
@@ -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[@]}"
|
||||
@@ -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[@]}"
|
||||
@@ -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)
|
||||
@@ -0,0 +1,5 @@
|
||||
from .checkpointing import *
|
||||
from .file_utils import *
|
||||
from .memory_utils import *
|
||||
from .optimizer_utils import *
|
||||
from .torch_utils import *
|
||||
@@ -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
|
||||
@@ -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("?", "-")
|
||||
)
|
||||
@@ -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
|
||||
@@ -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
|
||||
@@ -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)
|
||||
@@ -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,
|
||||
)
|
||||
@@ -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,
|
||||
)
|
||||
@@ -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)
|
||||
@@ -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)
|
||||
@@ -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 <littlenyima[at]163[dot]com>
|
||||
"""
|
||||
|
||||
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)
|
||||
Binary file not shown.
@@ -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
|
||||
```
|
||||
@@ -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("""
|
||||
<div style="text-align: center; font-size: 32px; font-weight: bold; margin-bottom: 20px;">
|
||||
CogVideoX-5B Huggingface Space🤗
|
||||
</div>
|
||||
<div style="text-align: center;">
|
||||
<a href="https://huggingface.co/THUDM/CogVideoX-5B">🤗 5B(T2V) Model Hub</a> |
|
||||
<a href="https://huggingface.co/THUDM/CogVideoX-5B-I2V">🤗 5B(I2V) Model Hub</a> |
|
||||
<a href="https://github.com/THUDM/CogVideo">🌐 Github</a> |
|
||||
<a href="https://arxiv.org/pdf/2408.06072">📜 arxiv </a>
|
||||
</div>
|
||||
<div style="text-align: center;display: flex;justify-content: center;align-items: center;margin-top: 1em;margin-bottom: .5em;">
|
||||
<span>If the Space is too busy, duplicate it to use privately</span>
|
||||
<a href="https://huggingface.co/spaces/THUDM/CogVideoX-5B-Space?duplicate=true"><img src="https://huggingface.co/datasets/huggingface/badges/resolve/main/duplicate-this-space-lg.svg" width="160" style="
|
||||
margin-left: .75em;
|
||||
"></a>
|
||||
</div>
|
||||
<div style="text-align: center; font-size: 15px; font-weight: bold; color: red; margin-bottom: 20px;">
|
||||
⚠️ This demo is for academic research and experimental use only.
|
||||
</div>
|
||||
""")
|
||||
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).<br> 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("""
|
||||
<table border="0" style="width: 100%; text-align: left; margin-top: 20px;">
|
||||
<div style="text-align: center; font-size: 32px; font-weight: bold; margin-bottom: 20px;">
|
||||
🎥 Video Gallery(For 5B)
|
||||
</div>
|
||||
<tr>
|
||||
<td style="width: 25%; vertical-align: top; font-size: 0.9em;">
|
||||
<p>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.</p>
|
||||
</td>
|
||||
<td style="width: 25%; vertical-align: top;">
|
||||
<video src="https://github.com/user-attachments/assets/cf5953ea-96d3-48fd-9907-c4708752c714" width="100%" controls autoplay loop></video>
|
||||
</td>
|
||||
<td style="width: 25%; vertical-align: top; font-size: 0.9em;">
|
||||
<p>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.</p>
|
||||
</td>
|
||||
<td style="width: 25%; vertical-align: top;">
|
||||
<video src="https://github.com/user-attachments/assets/fe0a78e6-b669-4800-8cf0-b5f9b5145b52" width="100%" controls autoplay loop></video>
|
||||
</td>
|
||||
</tr>
|
||||
<tr>
|
||||
<td style="width: 25%; vertical-align: top; font-size: 0.9em;">
|
||||
<p>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.</p>
|
||||
</td>
|
||||
<td style="width: 25%; vertical-align: top;">
|
||||
<video src="https://github.com/user-attachments/assets/c182f606-8f8c-421d-b414-8487070fcfcb" width="100%" controls autoplay loop></video>
|
||||
</td>
|
||||
<td style="width: 25%; vertical-align: top; font-size: 0.9em;">
|
||||
<p>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.</p>
|
||||
</td>
|
||||
<td style="width: 25%; vertical-align: top;">
|
||||
<video src="https://github.com/user-attachments/assets/7db2bbce-194d-434d-a605-350254b6c298" width="100%" controls autoplay loop></video>
|
||||
</td>
|
||||
</tr>
|
||||
<tr>
|
||||
<td style="width: 25%; vertical-align: top; font-size: 0.9em;">
|
||||
<p>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.</p>
|
||||
</td>
|
||||
<td style="width: 25%; vertical-align: top;">
|
||||
<video src="https://github.com/user-attachments/assets/62b01046-8cab-44cc-bd45-4d965bb615ec" width="100%" controls autoplay loop></video>
|
||||
</td>
|
||||
<td style="width: 25%; vertical-align: top; font-size: 0.9em;">
|
||||
<p>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.</p>
|
||||
</td>
|
||||
<td style="width: 25%; vertical-align: top;">
|
||||
<video src="https://github.com/user-attachments/assets/d78e552a-4b3f-4b81-ac3f-3898079554f6" width="100%" controls autoplay loop></video>
|
||||
</td>
|
||||
</tr>
|
||||
<tr>
|
||||
<td style="width: 25%; vertical-align: top; font-size: 0.9em;">
|
||||
<p>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.</p>
|
||||
</td>
|
||||
<td style="width: 25%; vertical-align: top;">
|
||||
<video src="https://github.com/user-attachments/assets/30894f12-c741-44a2-9e6e-ddcacc231e5b" width="100%" controls autoplay loop></video>
|
||||
</td>
|
||||
<td style="width: 25%; vertical-align: top; font-size: 0.9em;">
|
||||
<p>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.</p>
|
||||
</td>
|
||||
<td style="width: 25%; vertical-align: top;">
|
||||
<video src="https://github.com/user-attachments/assets/926575ca-7150-435b-a0ff-4900a963297b" width="100%" controls autoplay loop></video>
|
||||
</td>
|
||||
</tr>
|
||||
</table>
|
||||
""")
|
||||
|
||||
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()
|
||||
Binary file not shown.
|
After Width: | Height: | Size: 376 KiB |
Binary file not shown.
|
After Width: | Height: | Size: 473 KiB |
Binary file not shown.
|
After Width: | Height: | Size: 467 KiB |
@@ -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
|
||||
@@ -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
|
||||
@@ -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
|
||||
@@ -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
|
||||
@@ -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
|
||||
@@ -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,
|
||||
}
|
||||
@@ -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,
|
||||
}
|
||||
@@ -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))
|
||||
@@ -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)
|
||||
@@ -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)
|
||||
@@ -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)
|
||||
@@ -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)
|
||||
@@ -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
|
||||
)
|
||||
@@ -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)
|
||||
@@ -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)
|
||||
@@ -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("""
|
||||
<div style="text-align: center; font-size: 32px; font-weight: bold; margin-bottom: 20px;">
|
||||
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)<br>"
|
||||
"Increasing the number of inference steps will produce more detailed videos, but it will slow down the process.<br>"
|
||||
"50 steps are recommended for most cases.<br>"
|
||||
)
|
||||
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()
|
||||
@@ -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",
|
||||
]
|
||||
@@ -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
|
||||
@@ -0,0 +1,6 @@
|
||||
<div align="center">
|
||||
<img src=wechat.jpg width="60%"/>
|
||||
|
||||
<p> 扫码关注公众号，加入「 CogVideoX 交流群」 </p>
|
||||
<p> Scan the QR code to follow the official account and join the "CogVLM Discussion Group" </p>
|
||||
</div>
|
||||
@@ -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.
|
||||
Binary file not shown.
|
After Width: | Height: | Size: 60 KiB |
@@ -0,0 +1,142 @@
|
||||
<svg width="737" height="307" viewBox="0 0 737 307" fill="none" xmlns="http://www.w3.org/2000/svg">
|
||||
<path d="M242.629 203.026C237.829 200.426 234.029 196.826 231.429 192.126C228.729 187.426 227.429 182.026 227.429 175.826V149.226C227.429 143.126 228.729 137.826 231.429 133.126C234.129 128.426 237.829 124.926 242.629 122.326C247.429 119.726 253.029 118.526 259.429 118.526C265.829 118.526 271.429 119.726 276.229 122.126C281.029 124.526 284.829 127.926 287.429 132.326C290.129 136.726 291.429 141.826 291.429 147.526C291.429 147.926 291.329 148.226 291.029 148.526C290.729 148.726 290.429 148.926 290.029 148.926L273.029 150.026C272.029 150.026 271.629 149.526 271.629 148.626C271.629 144.726 270.529 141.726 268.229 139.426C265.929 137.126 263.029 136.026 259.329 136.026C255.629 136.026 252.729 137.226 250.429 139.526C248.229 141.926 247.029 144.926 247.029 148.726V176.726C247.029 180.526 248.129 183.526 250.429 185.826C252.629 188.126 255.629 189.226 259.329 189.226C263.029 189.226 265.929 188.126 268.229 185.826C270.429 183.526 271.629 180.526 271.629 176.726C271.629 175.826 272.129 175.326 273.029 175.326L290.029 176.226C290.429 176.226 290.729 176.326 291.029 176.626C291.329 176.826 291.429 177.126 291.429 177.526C291.429 183.326 290.129 188.526 287.429 192.926C284.729 197.326 281.029 200.826 276.229 203.226C271.429 205.626 265.829 206.826 259.429 206.826C253.029 206.826 247.429 205.526 242.629 203.026Z" fill="#1F08A3"/>
|
||||
<path d="M306.929 200.926C301.929 197.026 298.529 191.626 296.929 184.926C295.929 181.526 295.529 177.826 295.529 173.826C295.529 169.326 296.029 165.326 297.129 161.826C298.929 155.326 302.329 150.226 307.329 146.626C312.329 143.026 318.329 141.126 325.429 141.126C332.429 141.126 338.429 142.926 343.229 146.626C348.129 150.226 351.429 155.326 353.429 161.726C354.529 165.426 355.129 169.326 355.129 173.526C355.129 177.326 354.729 180.926 353.829 184.426C352.129 191.326 348.829 196.826 343.829 200.826C338.829 204.826 332.629 206.926 325.429 206.926C318.029 206.826 311.929 204.826 306.929 200.926ZM330.929 187.026C332.429 185.526 333.529 183.426 334.229 180.826C334.729 178.726 334.929 176.326 334.929 173.826C334.929 171.326 334.629 168.926 334.129 166.726C333.529 164.126 332.429 162.226 330.929 160.826C329.429 159.426 327.529 158.726 325.229 158.726C320.629 158.726 317.629 161.426 316.129 166.726C315.629 168.726 315.429 171.126 315.429 173.826C315.429 176.426 315.629 178.726 316.129 180.826C316.729 183.426 317.829 185.526 319.429 187.026C321.029 188.526 322.929 189.326 325.229 189.326C327.529 189.326 329.429 188.526 330.929 187.026Z" fill="#1F08A3"/>
|
||||
<path d="M397.929 142.526C398.229 142.226 398.529 142.126 398.929 142.126H415.929C416.329 142.126 416.629 142.226 416.929 142.526C417.229 142.826 417.329 143.126 417.329 143.626V198.726C417.329 210.026 414.229 218.026 407.829 222.926C401.529 227.726 393.529 230.226 383.929 230.226C380.329 230.226 376.529 229.926 372.529 229.226C371.729 229.126 371.329 228.626 371.329 227.626L371.929 212.326C371.929 211.726 372.129 211.326 372.429 211.126C372.729 210.926 373.129 210.926 373.629 211.026C376.929 211.726 379.929 212.126 382.729 212.126C387.229 212.126 390.829 211.026 393.529 208.926C396.229 206.826 397.529 203.426 397.529 198.826L396.729 199.726C393.829 202.726 389.729 204.126 384.329 204.126C379.029 204.126 374.229 202.926 369.829 200.526C365.429 198.126 362.329 193.926 360.329 187.926C359.029 184.026 358.429 179.126 358.429 173.326C358.429 167.026 359.229 161.826 360.729 157.826C362.529 152.726 365.529 148.626 369.529 145.626C373.529 142.626 378.229 141.026 383.529 141.026C389.229 141.026 393.729 142.826 396.829 146.426C397.029 146.626 397.129 146.626 397.329 146.626C397.529 146.526 397.529 146.426 397.529 146.226V143.526C397.529 143.126 397.629 142.826 397.929 142.526ZM397.529 173.326C397.529 171.126 397.429 169.326 397.329 168.126C397.129 166.826 396.829 165.626 396.329 164.526C395.729 162.726 394.729 161.326 393.229 160.326C391.829 159.326 390.129 158.826 388.129 158.826C384.329 158.826 381.729 160.726 380.129 164.626C378.929 166.926 378.329 169.926 378.329 173.726C378.329 177.726 378.829 180.726 379.929 182.626C380.629 184.326 381.729 185.726 383.129 186.826C384.529 187.926 386.229 188.426 388.229 188.426C392.329 188.426 395.029 186.526 396.329 182.726C397.129 180.626 397.529 177.526 397.529 173.326Z" fill="#1F08A3"/>
|
||||
<path d="M442.629 204.626L418.529 121.226L418.429 120.726C418.429 119.926 418.829 119.526 419.729 119.526H438.029C438.929 119.526 439.429 119.926 439.729 120.726L453.529 176.826C453.629 177.026 453.729 177.226 453.929 177.226C454.129 177.226 454.229 177.126 454.329 176.826L467.829 120.726C468.029 119.926 468.629 119.526 469.529 119.526H487.429C487.929 119.526 488.229 119.726 488.529 120.026C488.729 120.326 488.829 120.726 488.629 121.226L463.929 204.626C463.729 205.426 463.129 205.826 462.329 205.826H444.129C443.429 205.826 442.929 205.426 442.629 204.626Z" fill="#1F08A3"/>
|
||||
<path d="M495.529 133.426C493.529 131.426 492.629 128.826 492.629 125.726C492.629 122.526 493.629 119.926 495.529 117.926C497.529 115.926 499.929 114.926 503.029 114.926C506.129 114.926 508.529 115.926 510.529 117.926C512.429 119.926 513.429 122.526 513.429 125.726C513.429 128.726 512.429 131.326 510.529 133.326C508.529 135.426 506.129 136.426 503.029 136.426C499.929 136.426 497.529 135.426 495.529 133.426ZM493.429 205.426C493.129 205.126 493.029 204.826 493.029 204.426V143.626C493.029 143.226 493.129 142.826 493.429 142.526C493.729 142.226 494.029 142.126 494.429 142.126H511.429C511.829 142.126 512.129 142.226 512.429 142.526C512.729 142.826 512.829 143.126 512.829 143.626V204.426C512.829 204.826 512.729 205.226 512.429 205.426C512.129 205.726 511.829 205.826 511.429 205.826H494.429C494.029 205.826 493.729 205.726 493.429 205.426Z" fill="#1F08A3"/>
|
||||
<path d="M642.029 178.126C641.929 179.126 641.429 179.626 640.429 179.626H603.929C603.729 179.626 603.629 179.726 603.429 179.826C603.229 179.926 603.229 180.126 603.329 180.226C603.529 181.126 603.829 182.326 604.529 183.826C605.529 185.526 606.929 186.926 608.829 188.026C610.729 189.126 613.129 189.626 615.929 189.626C620.929 189.626 624.929 187.926 627.729 184.526C628.029 184.126 628.429 183.926 628.829 183.926C629.229 183.926 629.529 184.126 629.829 184.426L639.029 195.326C639.329 195.526 639.529 195.926 639.529 196.326C639.529 196.626 639.329 197.026 639.029 197.326C636.229 200.426 632.829 202.826 628.729 204.426C624.629 206.026 620.129 206.926 615.229 206.926C607.929 206.926 601.829 205.326 596.729 202.026C591.629 198.826 588.029 194.226 585.729 188.426C584.029 184.326 583.229 179.226 583.229 173.026C583.229 168.726 583.829 164.626 585.129 160.726C587.229 154.626 590.629 149.826 595.329 146.326C600.029 142.826 605.729 141.126 612.229 141.126C617.529 141.126 622.229 142.326 626.329 144.626C630.429 147.026 633.829 150.226 636.429 154.426C639.029 158.626 640.729 163.326 641.529 168.526C642.129 170.826 642.229 174.026 642.029 178.126ZM604.229 164.626C603.929 165.526 603.729 166.426 603.629 167.226C603.429 167.626 603.629 167.826 604.129 167.826H621.129C621.429 167.826 621.629 167.626 621.629 167.326C621.629 166.626 621.429 165.826 621.129 164.826C620.629 162.726 619.529 161.226 618.029 160.126C616.529 159.026 614.729 158.526 612.429 158.526C608.229 158.626 605.529 160.626 604.229 164.626Z" fill="#1F08A3"/>
|
||||
<path d="M656.629 200.926C651.629 197.026 648.229 191.626 646.629 184.926C645.629 181.526 645.229 177.826 645.229 173.826C645.229 169.326 645.729 165.326 646.829 161.826C648.629 155.326 652.029 150.226 657.029 146.626C662.029 143.026 668.029 141.126 675.129 141.126C682.129 141.126 688.129 142.926 692.929 146.626C697.829 150.226 701.129 155.326 703.129 161.726C704.229 165.426 704.829 169.326 704.829 173.526C704.829 177.326 704.429 180.926 703.529 184.426C701.829 191.326 698.529 196.826 693.529 200.826C688.529 204.826 682.329 206.926 675.129 206.926C667.729 206.826 661.629 204.826 656.629 200.926ZM680.729 187.026C682.229 185.526 683.329 183.426 684.029 180.826C684.529 178.726 684.729 176.326 684.729 173.826C684.729 171.326 684.429 168.926 683.929 166.726C683.229 164.126 682.229 162.226 680.729 160.826C679.229 159.426 677.329 158.726 675.029 158.726C670.429 158.726 667.329 161.426 665.929 166.726C665.429 168.726 665.229 171.126 665.229 173.826C665.229 176.426 665.429 178.726 665.929 180.826C666.529 183.426 667.629 185.526 669.229 187.026C670.829 188.526 672.729 189.326 675.029 189.326C677.229 189.326 679.129 188.526 680.729 187.026Z" fill="#1F08A3"/>
|
||||
<path d="M577.629 120.026C577.329 119.726 577.029 119.626 576.629 119.626H559.629C559.229 119.626 558.929 119.726 558.629 120.026C558.329 120.326 558.229 120.626 558.229 121.126V146.726C558.229 146.926 558.129 147.126 558.029 147.226C557.929 147.326 557.729 147.226 557.529 147.026C554.529 143.126 550.229 141.226 544.529 141.226C538.829 141.226 533.929 142.826 529.829 145.926C525.829 149.026 522.829 153.126 521.029 158.226C519.329 162.826 518.529 168.026 518.529 173.726C518.529 178.826 519.129 183.626 520.429 188.026C522.229 193.826 525.129 198.426 529.329 201.826C533.529 205.226 538.329 206.926 543.729 206.926C549.729 206.926 554.329 204.626 557.529 200.126C557.729 199.926 557.829 199.926 558.029 199.926C558.229 200.026 558.229 200.126 558.229 200.326V204.426C558.229 204.826 558.329 205.226 558.629 205.426C558.929 205.726 559.229 205.826 559.629 205.826H576.629C577.029 205.826 577.329 205.726 577.629 205.426C577.929 205.126 578.029 204.826 578.029 204.426V121.026C578.029 120.626 577.929 120.326 577.629 120.026ZM557.829 179.026L545.029 186.626C540.929 189.126 535.729 186.026 535.729 181.026V165.726C535.729 160.826 540.929 157.726 545.029 160.126L557.829 167.726C561.929 170.326 561.929 176.526 557.829 179.026Z" fill="#1F08A3"/>
|
||||
<path d="M151.429 241.226H109.329C108.429 241.226 107.629 240.826 107.129 240.026L32.6294 76.7258C31.4294 74.9258 32.7294 72.5258 34.8294 72.5258H76.9294C77.8294 72.5258 78.6294 72.9258 79.1294 73.7258L153.629 236.926C154.829 238.826 153.529 241.226 151.429 241.226Z" fill="url(#paint0_linear_4629_775)"/>
|
||||
<path d="M226.629 72.6257H183.929C183.129 72.6257 182.329 73.0257 181.929 73.7257L107.129 237.426C106.029 239.026 107.229 241.226 109.129 241.226H151.829C152.629 241.226 153.429 240.826 153.829 240.126L228.629 76.4257C229.729 74.8257 228.529 72.6257 226.629 72.6257Z" fill="url(#paint1_linear_4629_775)"/>
|
||||
<path d="M209.329 81.6258C202.129 81.2258 194.129 85.9258 182.529 99.9258C179.829 103.226 169.629 103.226 170.929 110.226C177.029 144.826 145.029 217.326 152.129 217.726C159.229 218.126 177.229 204.426 183.529 196.426C189.829 188.426 194.029 181.826 198.929 171.926C203.829 162.026 208.829 147.726 213.229 132.326C217.629 116.926 217.829 101.726 218.129 93.1258C218.329 90.5258 216.529 82.0258 209.329 81.6258Z" fill="#191935"/>
|
||||
<path d="M179.929 106.926C178.229 97.6258 171.029 78.8258 166.129 72.8258C161.229 66.8258 156.529 64.0258 148.829 66.9258C144.629 68.5258 134.729 74.2258 127.529 96.2258C125.229 103.226 121.429 106.726 118.029 111.426C114.629 116.026 109.229 122.426 109.229 131.226C109.229 140.126 112.829 147.626 117.729 153.326C122.629 159.026 143.529 175.826 147.929 180.426C152.429 185.026 164.629 197.226 167.829 202.726C171.029 208.226 171.029 208.226 171.029 208.226C171.029 208.226 177.029 194.326 178.929 189.826C180.829 185.426 184.329 162.026 184.629 148.726C184.929 135.526 181.129 113.626 179.929 106.926Z" fill="black"/>
|
||||
<path d="M154.829 225.526C154.029 225.926 153.529 223.626 156.229 219.726C158.929 215.826 164.329 208.326 168.629 205.026C173.029 201.726 183.929 192.926 188.029 192.026C192.129 191.126 192.129 194.726 190.129 196.226C186.329 199.126 185.629 198.426 180.829 202.826C176.029 207.226 170.129 209.526 166.529 212.826C162.929 216.126 157.029 224.326 154.829 225.526Z" fill="black"/>
|
||||
<path d="M115.629 159.226C103.129 155.626 82.8293 151.626 75.1293 180.126C75.1293 180.126 69.6293 182.526 65.5293 191.926C58.9293 207.026 73.4293 247.326 116.329 236.826C136.729 231.826 153.729 221.026 158.729 216.126C166.329 208.726 171.229 208.326 171.229 208.326C167.829 197.526 143.429 167.326 115.629 159.226Z" fill="black"/>
|
||||
<path d="M187.729 194.826C191.829 191.026 197.129 185.326 199.429 181.326C201.729 177.326 204.029 171.726 204.429 170.126C204.829 168.526 203.829 163.526 206.229 163.826C208.629 164.126 207.129 168.026 206.129 169.926C204.729 172.426 203.029 176.326 201.929 179.026C200.829 181.826 196.829 186.826 193.729 190.426C190.529 194.026 187.429 196.526 187.429 196.526L187.729 194.826Z" fill="black"/>
|
||||
<path d="M186.529 196.126C186.529 196.126 191.629 193.126 195.729 190.226C199.829 187.326 208.129 179.826 210.129 177.126C212.229 174.426 212.529 171.926 213.629 172.526C215.529 173.726 212.329 176.526 211.029 177.826C209.629 179.126 204.929 183.726 201.329 186.726C197.629 189.726 188.129 196.626 186.529 196.126Z" fill="black"/>
|
||||
<path d="M173.729 198.526C173.829 198.926 174.129 199.226 174.129 198.726C175.029 190.026 178.529 179.726 178.629 154.326C178.729 127.426 178.629 115.126 159.029 110.126C157.729 109.826 156.529 110.126 156.729 111.426C157.229 114.426 164.029 129.226 167.429 148.626C171.529 172.526 172.929 194.526 173.729 198.526Z" fill="url(#paint2_linear_4629_775)"/>
|
||||
<path d="M88.8293 186.526C82.1293 188.426 98.2293 191.726 113.629 194.726C124.129 196.826 157.229 201.926 168.429 202.526C168.629 202.526 168.729 202.226 168.529 202.126C157.329 194.826 132.529 178.926 125.429 178.926C112.629 178.826 93.8293 185.126 88.8293 186.526Z" fill="url(#paint3_linear_4629_775)"/>
|
||||
<path d="M127.229 134.626C127.229 134.626 132.529 145.726 140.829 152.626C148.029 158.526 154.329 162.326 156.629 165.126C164.529 174.426 167.629 190.726 171.529 198.426C172.129 199.526 142.829 170.526 127.929 145.726C124.529 140.026 126.329 133.126 127.229 134.626Z" fill="url(#paint4_linear_4629_775)"/>
|
||||
<path d="M146.029 85.9258C149.129 85.3258 157.329 90.0258 157.329 92.4258C157.329 93.5258 145.529 91.3258 143.529 90.3258C140.829 88.9258 144.029 86.3258 146.029 85.9258Z" fill="url(#paint5_linear_4629_775)"/>
|
||||
<path d="M140.629 94.5258C142.129 93.2258 148.429 94.6258 149.629 95.7258C150.829 97.0258 143.429 97.8258 141.029 98.1258C138.929 98.3258 139.129 95.8258 140.629 94.5258Z" fill="url(#paint6_linear_4629_775)"/>
|
||||
<path d="M155.13 75.1258C159.73 75.5258 168.93 89.0258 166.93 90.6258C165.83 91.4258 151.53 82.6258 149.43 80.1258C147.63 77.9258 151.63 74.8258 155.13 75.1258Z" fill="url(#paint7_linear_4629_775)"/>
|
||||
<path d="M200.829 96.0258C203.129 96.3258 207.73 101.426 207.23 103.226C207.03 104.026 199.23 99.7258 198.03 98.5258C196.43 96.8258 199.329 95.8258 200.829 96.0258Z" fill="url(#paint8_linear_4629_775)"/>
|
||||
<path d="M195.029 101.326C196.129 100.626 200.629 103.126 200.729 104.326C200.829 105.326 196.329 104.226 194.529 103.926C193.029 103.726 193.829 102.126 195.029 101.326Z" fill="url(#paint9_linear_4629_775)"/>
|
||||
<path d="M208.429 90.5258C211.629 91.8258 214.329 100.826 213.329 101.726C212.229 102.726 204.229 95.0258 203.329 92.8258C202.529 90.9258 206.029 89.6258 208.429 90.5258Z" fill="url(#paint10_linear_4629_775)"/>
|
||||
<path d="M164.029 166.426C163.029 161.426 161.429 140.726 156.229 128.326C152.429 119.326 142.129 102.626 140.029 108.326C136.929 116.726 129.929 124.526 132.329 130.726C134.729 136.926 142.829 145.426 148.929 150.726C154.829 155.926 164.129 166.926 164.029 166.426Z" fill="url(#paint11_linear_4629_775)"/>
|
||||
<path d="M108.129 199.126C100.229 198.026 88.1291 195.526 80.6291 194.026C75.9291 193.126 81.4291 206.126 87.8291 209.526C93.3291 212.426 118.129 210.826 134.529 209.226C148.129 207.826 168.929 204.226 161.829 204.126C150.229 203.726 131.129 202.426 108.129 199.126Z" fill="url(#paint12_linear_4629_775)"/>
|
||||
<path d="M163.629 207.426C164.029 207.226 164.529 206.426 164.229 206.626C158.129 208.626 135.629 214.826 111.029 217.526C92.2292 219.526 88.2292 217.926 86.0292 218.926C83.4292 220.026 90.6292 229.826 106.729 228.926C125.129 227.826 153.129 214.226 163.629 207.426Z" fill="url(#paint13_linear_4629_775)"/>
|
||||
<path d="M90.5292 177.226C89.3292 177.226 89.3292 173.726 93.2292 172.026C93.2292 172.026 97.6292 166.426 106.429 168.126C112.529 169.326 122.229 174.826 120.929 174.526C116.329 173.426 96.3292 177.226 90.5292 177.226Z" fill="url(#paint14_linear_4629_775)"/>
|
||||
<path d="M179.629 194.626C179.629 194.226 183.529 188.126 188.629 161.026C191.829 144.226 191.029 129.426 190.529 126.526C190.229 124.426 197.429 128.626 197.529 136.526C197.629 142.226 196.329 165.026 186.029 185.526C181.829 194.026 179.629 195.126 179.629 194.626Z" fill="url(#paint15_linear_4629_775)"/>
|
||||
<path d="M194.529 173.526C194.429 173.726 194.029 173.626 194.129 173.426C196.229 168.826 201.929 149.226 203.129 134.326C203.629 128.526 203.629 123.826 204.329 122.326C205.029 120.526 209.429 115.526 210.429 114.926C212.529 113.726 211.529 126.026 206.829 141.326C204.029 150.826 198.229 167.026 194.529 173.526Z" fill="url(#paint16_linear_4629_775)"/>
|
||||
<defs>
|
||||
<linearGradient id="paint0_linear_4629_775" x1="49.9997" y1="84.6314" x2="126.703" y2="213.187" gradientUnits="userSpaceOnUse">
|
||||
<stop stop-color="#19D46A"/>
|
||||
<stop offset="1" stop-color="#0D00B5"/>
|
||||
</linearGradient>
|
||||
<linearGradient id="paint1_linear_4629_775" x1="137.115" y1="226.067" x2="190.755" y2="105.47" gradientUnits="userSpaceOnUse">
|
||||
<stop offset="0.000425658" stop-color="#D70066"/>
|
||||
<stop offset="1" stop-color="#0D00B5"/>
|
||||
</linearGradient>
|
||||
<linearGradient id="paint2_linear_4629_775" x1="165.812" y1="110.885" x2="178.143" y2="199.088" gradientUnits="userSpaceOnUse">
|
||||
<stop stop-color="#8890F4"/>
|
||||
<stop offset="0.2835" stop-color="#FFCFFD"/>
|
||||
<stop offset="0.6062" stop-color="#F4F0FE"/>
|
||||
<stop offset="0.7338" stop-color="#EFFFFF"/>
|
||||
<stop offset="1" stop-color="#75D9DD"/>
|
||||
</linearGradient>
|
||||
<linearGradient id="paint3_linear_4629_775" x1="82.5454" y1="177.066" x2="173.477" y2="202.043" gradientUnits="userSpaceOnUse">
|
||||
<stop offset="0.00507186" stop-color="#6F69F4"/>
|
||||
<stop offset="0.4148" stop-color="#FFCEAD"/>
|
||||
<stop offset="0.5961" stop-color="#EFEFB9"/>
|
||||
<stop offset="0.6727" stop-color="#E7FFBF"/>
|
||||
<stop offset="1" stop-color="#75D9DD"/>
|
||||
</linearGradient>
|
||||
<linearGradient id="paint4_linear_4629_775" x1="120.656" y1="131.195" x2="183.797" y2="206.698" gradientUnits="userSpaceOnUse">
|
||||
<stop offset="0.00507186" stop-color="#6F69F4"/>
|
||||
<stop offset="0.3821" stop-color="#FFAFFF"/>
|
||||
<stop offset="0.5127" stop-color="#F5D1E4"/>
|
||||
<stop offset="0.6727" stop-color="#E7FFBF"/>
|
||||
<stop offset="0.9718" stop-color="#75D9DD"/>
|
||||
</linearGradient>
|
||||
<linearGradient id="paint5_linear_4629_775" x1="155.235" y1="97.0899" x2="149.664" y2="90.3277" gradientUnits="userSpaceOnUse">
|
||||
<stop stop-color="#FFA2FF"/>
|
||||
<stop offset="0.5569" stop-color="#FFFFB0"/>
|
||||
<stop offset="0.7133" stop-color="#DDFFCA"/>
|
||||
<stop offset="1" stop-color="#97FFFF"/>
|
||||
</linearGradient>
|
||||
<linearGradient id="paint6_linear_4629_775" x1="147.529" y1="100.455" x2="142.542" y2="94.02" gradientUnits="userSpaceOnUse">
|
||||
<stop stop-color="#FFA2FF"/>
|
||||
<stop offset="0.5569" stop-color="#FFFFB0"/>
|
||||
<stop offset="0.7133" stop-color="#DDFFCA"/>
|
||||
<stop offset="1" stop-color="#97FFFF"/>
|
||||
</linearGradient>
|
||||
<linearGradient id="paint7_linear_4629_775" x1="165.834" y1="95.5891" x2="158.095" y2="83.7515" gradientUnits="userSpaceOnUse">
|
||||
<stop stop-color="#FFA2FF"/>
|
||||
<stop offset="0.5569" stop-color="#FFFFB0"/>
|
||||
<stop offset="0.7133" stop-color="#DDFFCA"/>
|
||||
<stop offset="1" stop-color="#97FFFF"/>
|
||||
</linearGradient>
|
||||
<linearGradient id="paint8_linear_4629_775" x1="204.577" y1="106.047" x2="202.324" y2="100.006" gradientUnits="userSpaceOnUse">
|
||||
<stop stop-color="#FFA2FF"/>
|
||||
<stop offset="0.5569" stop-color="#FFFFB0"/>
|
||||
<stop offset="0.7133" stop-color="#DDFFCA"/>
|
||||
<stop offset="1" stop-color="#97FFFF"/>
|
||||
</linearGradient>
|
||||
<linearGradient id="paint9_linear_4629_775" x1="198.54" y1="107.056" x2="196.25" y2="101.429" gradientUnits="userSpaceOnUse">
|
||||
<stop stop-color="#FFA2FF"/>
|
||||
<stop offset="0.5569" stop-color="#FFFFB0"/>
|
||||
<stop offset="0.7133" stop-color="#DDFFCA"/>
|
||||
<stop offset="1" stop-color="#97FFFF"/>
|
||||
</linearGradient>
|
||||
<linearGradient id="paint10_linear_4629_775" x1="210.937" y1="104.861" x2="208.401" y2="96.2553" gradientUnits="userSpaceOnUse">
|
||||
<stop stop-color="#FFA2FF"/>
|
||||
<stop offset="0.5569" stop-color="#FFFFB0"/>
|
||||
<stop offset="0.7133" stop-color="#DDFFCA"/>
|
||||
<stop offset="1" stop-color="#97FFFF"/>
|
||||
</linearGradient>
|
||||
<linearGradient id="paint11_linear_4629_775" x1="167.808" y1="155.267" x2="133.72" y2="115.546" gradientUnits="userSpaceOnUse">
|
||||
<stop offset="0.00507186" stop-color="#6F69F4"/>
|
||||
<stop offset="0.3291" stop-color="#FFD9AD"/>
|
||||
<stop offset="0.6346" stop-color="#EAFABD"/>
|
||||
<stop offset="0.6727" stop-color="#E7FFBF"/>
|
||||
<stop offset="1" stop-color="#75D9DD"/>
|
||||
</linearGradient>
|
||||
<linearGradient id="paint12_linear_4629_775" x1="165.376" y1="206.257" x2="79.2543" y2="201.981" gradientUnits="userSpaceOnUse">
|
||||
<stop offset="0.00507186" stop-color="#6F69F4"/>
|
||||
<stop offset="0.3302" stop-color="#FFECEC"/>
|
||||
<stop offset="0.5907" stop-color="#EDFACB"/>
|
||||
<stop offset="0.6727" stop-color="#E7FFBF"/>
|
||||
<stop offset="0.9718" stop-color="#75D9DD"/>
|
||||
</linearGradient>
|
||||
<linearGradient id="paint13_linear_4629_775" x1="85.4478" y1="220.341" x2="164.817" y2="216.145" gradientUnits="userSpaceOnUse">
|
||||
<stop offset="0.00507186" stop-color="#6F69F4"/>
|
||||
<stop offset="0.3302" stop-color="#FFECEC"/>
|
||||
<stop offset="0.5907" stop-color="#EDFACB"/>
|
||||
<stop offset="0.6727" stop-color="#E7FFBF"/>
|
||||
<stop offset="0.9718" stop-color="#75D9DD"/>
|
||||
</linearGradient>
|
||||
<linearGradient id="paint14_linear_4629_775" x1="153.232" y1="186.551" x2="86.9228" y2="169.651" gradientUnits="userSpaceOnUse">
|
||||
<stop offset="0.00507186" stop-color="#6F69F4"/>
|
||||
<stop offset="0.601" stop-color="#FFCEAD"/>
|
||||
<stop offset="0.707" stop-color="#EFEFB9"/>
|
||||
<stop offset="0.7518" stop-color="#E7FFBF"/>
|
||||
<stop offset="0.976" stop-color="#75D9DD"/>
|
||||
</linearGradient>
|
||||
<linearGradient id="paint15_linear_4629_775" x1="196.161" y1="125.694" x2="180.939" y2="202.727" gradientUnits="userSpaceOnUse">
|
||||
<stop stop-color="#8890F4"/>
|
||||
<stop offset="0.4248" stop-color="#FFB8FD"/>
|
||||
<stop offset="0.5488" stop-color="#F4D9FE"/>
|
||||
<stop offset="0.6727" stop-color="#E7FFFF"/>
|
||||
<stop offset="1" stop-color="#75D9DD"/>
|
||||
</linearGradient>
|
||||
<linearGradient id="paint16_linear_4629_775" x1="194.758" y1="172.089" x2="209.302" y2="118.443" gradientUnits="userSpaceOnUse">
|
||||
<stop offset="0.00507186" stop-color="#6F69F4"/>
|
||||
<stop offset="0.3291" stop-color="#FFCEAD"/>
|
||||
<stop offset="0.5707" stop-color="#EFEFB9"/>
|
||||
<stop offset="0.6727" stop-color="#E7FFBF"/>
|
||||
<stop offset="1" stop-color="#75D9DD"/>
|
||||
</linearGradient>
|
||||
</defs>
|
||||
</svg>
|
||||
|
After Width: | Height: | Size: 22 KiB |
Binary file not shown.
|
After Width: | Height: | Size: 605 KiB |
Binary file not shown.
|
After Width: | Height: | Size: 195 KiB |
@@ -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:
|
||||
<div align="center">
|
||||
<img width="600px" height="auto" src="./assests/CogVLM2-Caption-example.png">
|
||||
</div>
|
||||
|
||||
## 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:
|
||||
<div align="center">
|
||||
<a href="https://cogvlm2-video.github.io/"><img width="600px" height="auto" src="./assests/cogvlm2-video-example.png"></a>
|
||||
</div>
|
||||
|
||||
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}
|
||||
}
|
||||
```
|
||||
@@ -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
|
||||
```
|
||||
|
||||
例:
|
||||
<div align="center">
|
||||
<img width="600px" height="auto" src="./assests/CogVLM2-Caption-example.png">
|
||||
</div>
|
||||
|
||||
|
||||
|
||||
## 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 は、タイムスタンプベースの質問応答機能を備えた多機能なビデオ理解モデルです。ユーザーは `このビデオを詳細に説明してください。` などのプロンプトをモデルに入力して、詳細なビデオキャプションを取得できます：
|
||||
<div align="center">
|
||||
<a href="https://cogvlm2-video.github.io/"><img width="600px" height="auto" src="./assests/cogvlm2-video-example.png"></a>
|
||||
</div>
|
||||
|
||||
ユーザーは提供された[コード](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}
|
||||
}
|
||||
```
|
||||
@@ -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
|
||||
```
|
||||
|
||||
示例：
|
||||
<div align="center">
|
||||
<img width="600px" height="auto" src="./assests/CogVLM2-Caption-example.png">
|
||||
</div>
|
||||
|
||||
## 通过 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：
|
||||
|
||||
|
||||
<div align="center">
|
||||
<a href="https://cogvlm2-video.github.io/"><img width="600px" height="auto" src="./assests/cogvlm2-video-example.png"></a>
|
||||
</div>
|
||||
|
||||
用户可以使用提供的[代码](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}
|
||||
}
|
||||
```
|
||||
Binary file not shown.
|
After Width: | Height: | Size: 1.1 MiB |
Binary file not shown.
|
After Width: | Height: | Size: 1.4 MiB |
@@ -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
|
||||
@@ -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()
|
||||
@@ -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,
|
||||
)
|
||||
@@ -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,
|
||||
)
|
||||
@@ -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)
|
||||
@@ -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
|
||||
@@ -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("""
|
||||
<div style="text-align: center; font-size: 32px; font-weight: bold; margin-bottom: 20px;">
|
||||
LLM + FLUX + CogVideoX-I2V Space 🤗
|
||||
</div>
|
||||
""")
|
||||
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()
|
||||
@@ -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)
|
||||
Some files were not shown because too many files have changed in this diff Show More
Reference in New Issue
Block a user