Upload folder using huggingface_hub

.gitattributes

@@ -33,3 +33,7 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
 *.zip filter=lfs diff=lfs merge=lfs -text
 *.zst filter=lfs diff=lfs merge=lfs -text
 *tfevents* filter=lfs diff=lfs merge=lfs -text
+resources/CogVideoX.pdf filter=lfs diff=lfs merge=lfs -text
+resources/videos/2.mp4 filter=lfs diff=lfs merge=lfs -text
+resources/web_demo.png filter=lfs diff=lfs merge=lfs -text
+tools/caption/assests/cogvlm2-video-example.png filter=lfs diff=lfs merge=lfs -text

.github/ISSUE_TEMPLATE/bug_report.yaml

@@ -0,0 +1,51 @@
+name: "\U0001F41B Bug Report"
+description: Submit a bug report to help us improve CogVideoX / 提交一个 Bug 问题报告来帮助我们改进 CogVideoX 开源模型
+body:
+  - type: textarea
+    id: system-info
+    attributes:
+      label: System Info / 系統信息
+      description: Your operating environment / 您的运行环境信息
+      placeholder: Includes Cuda version, Diffusers version, Python version, operating system, hardware information (if you suspect a hardware problem)... / 包括Cuda版本，Diffusers，Python版本，操作系统，硬件信息(如果您怀疑是硬件方面的问题)...
+    validations:
+      required: true
+
+  - type: checkboxes
+    id: information-scripts-examples
+    attributes:
+      label: Information / 问题信息
+      description: 'The problem arises when using: / 问题出现在'
+      options:
+        - label: "The official example scripts / 官方的示例脚本"
+        - label: "My own modified scripts / 我自己修改的脚本和任务"
+
+  - type: textarea
+    id: reproduction
+    validations:
+      required: true
+    attributes:
+      label: Reproduction / 复现过程
+      description: |
+        Please provide a code example that reproduces the problem you encountered, preferably with a minimal reproduction unit.
+        If you have code snippets, error messages, stack traces, please provide them here as well.
+        Please format your code correctly using code tags. See https://help.github.com/en/github/writing-on-github/creating-and-highlighting-code-blocks#syntax-highlighting
+        Do not use screenshots, as they are difficult to read and (more importantly) do not allow others to copy and paste your code.
+        
+        请提供能重现您遇到的问题的代码示例,最好是最小复现单元。
+        如果您有代码片段、错误信息、堆栈跟踪，也请在此提供。
+        请使用代码标签正确格式化您的代码。请参见 https://help.github.com/en/github/writing-on-github/creating-and-highlighting-code-blocks#syntax-highlighting
+        请勿使用截图，因为截图难以阅读，而且（更重要的是）不允许他人复制粘贴您的代码。
+      placeholder: |
+        Steps to reproduce the behavior/复现Bug的步骤:
+          
+          1.
+          2.
+          3.
+
+  - type: textarea
+    id: expected-behavior
+    validations:
+      required: true
+    attributes:
+      label: Expected behavior / 期待表现
+      description: "A clear and concise description of what you would expect to happen. /简单描述您期望发生的事情。"

.github/ISSUE_TEMPLATE/feature-request.yaml

@@ -0,0 +1,34 @@
+name: "\U0001F680 Feature request"
+description: Submit a request for a new CogVideoX feature / 提交一个新的 CogVideoX开源模型的功能建议
+labels: [ "feature" ]
+body:
+  - type: textarea
+    id: feature-request
+    validations:
+      required: true
+    attributes:
+      label: Feature request  / 功能建议
+      description: |
+        A brief description of the functional proposal. Links to corresponding papers and code are desirable.
+        对功能建议的简述。最好提供对应的论文和代码链接。
+
+  - type: textarea
+    id: motivation
+    validations:
+      required: true
+    attributes:
+      label: Motivation / 动机
+      description: |
+        Your motivation for making the suggestion. If that motivation is related to another GitHub issue, link to it here.
+        您提出建议的动机。如果该动机与另一个 GitHub 问题有关，请在此处提供对应的链接。
+
+  - type: textarea
+    id: contribution
+    validations:
+      required: true
+    attributes:
+      label: Your contribution / 您的贡献
+      description: |
+        
+        Your PR link or any other link you can help with.
+        您的PR链接或者其他您能提供帮助的链接。

.github/PULL_REQUEST_TEMPLATE/pr_template.md

@@ -0,0 +1,34 @@
+#  Raise valuable PR / 提出有价值的PR
+
+## Caution / 注意事项:
+Users should keep the following points in mind when submitting PRs:
+
+1. Ensure that your code meets the requirements in the [specification](../../resources/contribute.md).
+2. the proposed PR should be relevant, if there are multiple ideas and optimizations, they should be assigned to different PRs.
+
+用户在提交PR时候应该注意以下几点:
+
+1. 确保您的代码符合 [规范](../../resources/contribute_zh.md) 中的要求。
+2. 提出的PR应该具有针对性，如果具有多个不同的想法和优化方案，应该分配到不同的PR中。
+
+## 不应该提出的PR / PRs that should not be proposed
+
+If a developer proposes a PR about any of the following, it may be closed or Rejected.
+
+1. those that don't describe improvement options.
+2. multiple issues of different types combined in one PR.
+3. The proposed PR is highly duplicative of already existing PRs.
+
+如果开发者提出关于以下方面的PR，则可能会被直接关闭或拒绝通过。
+
+1. 没有说明改进方案的。
+2. 多个不同类型的问题合并在一个PR中的。
+3. 提出的PR与已经存在的PR高度重复的。
+
+
+# 检查您的PR
+- [ ] Have you read the Contributor Guidelines, Pull Request section? / 您是否阅读了贡献者指南、Pull Request 部分？
+- [ ] Has this been discussed/approved via a Github issue or forum? If so, add a link. / 是否通过 Github 问题或论坛讨论/批准过？如果是，请添加链接。
+- [ ] Did you make sure you updated the documentation with your changes? Here are the Documentation Guidelines, and here are the Documentation Formatting Tips. /您是否确保根据您的更改更新了文档？这里是文档指南，这里是文档格式化技巧。
+- [ ] Did you write new required tests? / 您是否编写了新的必要测试？
+- [ ]  Are your PRs for only one issue / 您的PR是否仅针对一个问题

.gitignore

@@ -0,0 +1,9 @@
+output/
+*__pycache__/
+samples*/
+runs/
+checkpoints/
+master_ip
+logs/
+*.DS_Store
+.idea

LICENSE

@@ -0,0 +1,201 @@
+                                 Apache License
+                           Version 2.0, January 2004
+                        http://www.apache.org/licenses/
+
+   TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION
+
+   1. Definitions.
+
+      "License" shall mean the terms and conditions for use, reproduction,
+      and distribution as defined by Sections 1 through 9 of this document.
+
+      "Licensor" shall mean the copyright owner or entity authorized by
+      the copyright owner that is granting the License.
+
+      "Legal Entity" shall mean the union of the acting entity and all
+      other entities that control, are controlled by, or are under common
+      control with that entity. For the purposes of this definition,
+      "control" means (i) the power, direct or indirect, to cause the
+      direction or management of such entity, whether by contract or
+      otherwise, or (ii) ownership of fifty percent (50%) or more of the
+      outstanding shares, or (iii) beneficial ownership of such entity.
+
+      "You" (or "Your") shall mean an individual or Legal Entity
+      exercising permissions granted by this License.
+
+      "Source" form shall mean the preferred form for making modifications,
+      including but not limited to software source code, documentation
+      source, and configuration files.
+
+      "Object" form shall mean any form resulting from mechanical
+      transformation or translation of a Source form, including but
+      not limited to compiled object code, generated documentation,
+      and conversions to other media types.
+
+      "Work" shall mean the work of authorship, whether in Source or
+      Object form, made available under the License, as indicated by a
+      copyright notice that is included in or attached to the work
+      (an example is provided in the Appendix below).
+
+      "Derivative Works" shall mean any work, whether in Source or Object
+      form, that is based on (or derived from) the Work and for which the
+      editorial revisions, annotations, elaborations, or other modifications
+      represent, as a whole, an original work of authorship. For the purposes
+      of this License, Derivative Works shall not include works that remain
+      separable from, or merely link (or bind by name) to the interfaces of,
+      the Work and Derivative Works thereof.
+
+      "Contribution" shall mean any work of authorship, including
+      the original version of the Work and any modifications or additions
+      to that Work or Derivative Works thereof, that is intentionally
+      submitted to Licensor for inclusion in the Work by the copyright owner
+      or by an individual or Legal Entity authorized to submit on behalf of
+      the copyright owner. For the purposes of this definition, "submitted"
+      means any form of electronic, verbal, or written communication sent
+      to the Licensor or its representatives, including but not limited to
+      communication on electronic mailing lists, source code control systems,
+      and issue tracking systems that are managed by, or on behalf of, the
+      Licensor for the purpose of discussing and improving the Work, but
+      excluding communication that is conspicuously marked or otherwise
+      designated in writing by the copyright owner as "Not a Contribution."
+
+      "Contributor" shall mean Licensor and any individual or Legal Entity
+      on behalf of whom a Contribution has been received by Licensor and
+      subsequently incorporated within the Work.
+
+   2. Grant of Copyright License. Subject to the terms and conditions of
+      this License, each Contributor hereby grants to You a perpetual,
+      worldwide, non-exclusive, no-charge, royalty-free, irrevocable
+      copyright license to reproduce, prepare Derivative Works of,
+      publicly display, publicly perform, sublicense, and distribute the
+      Work and such Derivative Works in Source or Object form.
+
+   3. Grant of Patent License. Subject to the terms and conditions of
+      this License, each Contributor hereby grants to You a perpetual,
+      worldwide, non-exclusive, no-charge, royalty-free, irrevocable
+      (except as stated in this section) patent license to make, have made,
+      use, offer to sell, sell, import, and otherwise transfer the Work,
+      where such license applies only to those patent claims licensable
+      by such Contributor that are necessarily infringed by their
+      Contribution(s) alone or by combination of their Contribution(s)
+      with the Work to which such Contribution(s) was submitted. If You
+      institute patent litigation against any entity (including a
+      cross-claim or counterclaim in a lawsuit) alleging that the Work
+      or a Contribution incorporated within the Work constitutes direct
+      or contributory patent infringement, then any patent licenses
+      granted to You under this License for that Work shall terminate
+      as of the date such litigation is filed.
+
+   4. Redistribution. You may reproduce and distribute copies of the
+      Work or Derivative Works thereof in any medium, with or without
+      modifications, and in Source or Object form, provided that You
+      meet the following conditions:
+
+      (a) You must give any other recipients of the Work or
+          Derivative Works a copy of this License; and
+
+      (b) You must cause any modified files to carry prominent notices
+          stating that You changed the files; and
+
+      (c) You must retain, in the Source form of any Derivative Works
+          that You distribute, all copyright, patent, trademark, and
+          attribution notices from the Source form of the Work,
+          excluding those notices that do not pertain to any part of
+          the Derivative Works; and
+
+      (d) If the Work includes a "NOTICE" text file as part of its
+          distribution, then any Derivative Works that You distribute must
+          include a readable copy of the attribution notices contained
+          within such NOTICE file, excluding those notices that do not
+          pertain to any part of the Derivative Works, in at least one
+          of the following places: within a NOTICE text file distributed
+          as part of the Derivative Works; within the Source form or
+          documentation, if provided along with the Derivative Works; or,
+          within a display generated by the Derivative Works, if and
+          wherever such third-party notices normally appear. The contents
+          of the NOTICE file are for informational purposes only and
+          do not modify the License. You may add Your own attribution
+          notices within Derivative Works that You distribute, alongside
+          or as an addendum to the NOTICE text from the Work, provided
+          that such additional attribution notices cannot be construed
+          as modifying the License.
+
+      You may add Your own copyright statement to Your modifications and
+      may provide additional or different license terms and conditions
+      for use, reproduction, or distribution of Your modifications, or
+      for any such Derivative Works as a whole, provided Your use,
+      reproduction, and distribution of the Work otherwise complies with
+      the conditions stated in this License.
+
+   5. Submission of Contributions. Unless You explicitly state otherwise,
+      any Contribution intentionally submitted for inclusion in the Work
+      by You to the Licensor shall be under the terms and conditions of
+      this License, without any additional terms or conditions.
+      Notwithstanding the above, nothing herein shall supersede or modify
+      the terms of any separate license agreement you may have executed
+      with Licensor regarding such Contributions.
+
+   6. Trademarks. This License does not grant permission to use the trade
+      names, trademarks, service marks, or product names of the Licensor,
+      except as required for reasonable and customary use in describing the
+      origin of the Work and reproducing the content of the NOTICE file.
+
+   7. Disclaimer of Warranty. Unless required by applicable law or
+      agreed to in writing, Licensor provides the Work (and each
+      Contributor provides its Contributions) on an "AS IS" BASIS,
+      WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or
+      implied, including, without limitation, any warranties or conditions
+      of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A
+      PARTICULAR PURPOSE. You are solely responsible for determining the
+      appropriateness of using or redistributing the Work and assume any
+      risks associated with Your exercise of permissions under this License.
+
+   8. Limitation of Liability. In no event and under no legal theory,
+      whether in tort (including negligence), contract, or otherwise,
+      unless required by applicable law (such as deliberate and grossly
+      negligent acts) or agreed to in writing, shall any Contributor be
+      liable to You for damages, including any direct, indirect, special,
+      incidental, or consequential damages of any character arising as a
+      result of this License or out of the use or inability to use the
+      Work (including but not limited to damages for loss of goodwill,
+      work stoppage, computer failure or malfunction, or any and all
+      other commercial damages or losses), even if such Contributor
+      has been advised of the possibility of such damages.
+
+   9. Accepting Warranty or Additional Liability. While redistributing
+      the Work or Derivative Works thereof, You may choose to offer,
+      and charge a fee for, acceptance of support, warranty, indemnity,
+      or other liability obligations and/or rights consistent with this
+      License. However, in accepting such obligations, You may act only
+      on Your own behalf and on Your sole responsibility, not on behalf
+      of any other Contributor, and only if You agree to indemnify,
+      defend, and hold each Contributor harmless for any liability
+      incurred by, or claims asserted against, such Contributor by reason
+      of your accepting any such warranty or additional liability.
+
+   END OF TERMS AND CONDITIONS
+
+   APPENDIX: How to apply the Apache License to your work.
+
+      To apply the Apache License to your work, attach the following
+      boilerplate notice, with the fields enclosed by brackets "[]"
+      replaced with your own identifying information. (Don't include
+      the brackets!)  The text should be enclosed in the appropriate
+      comment syntax for the file format. We also recommend that a
+      file or class name and description of purpose be included on the
+      same "printed page" as the copyright notice for easier
+      identification within third-party archives.
+
+   Copyright 2024 CogVideo Model Team @ Zhipu AI
+
+   Licensed under the Apache License, Version 2.0 (the "License");
+   you may not use this file except in compliance with the License.
+   You may obtain a copy of the License at
+
+       http://www.apache.org/licenses/LICENSE-2.0
+
+   Unless required by applicable law or agreed to in writing, software
+   distributed under the License is distributed on an "AS IS" BASIS,
+   WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+   See the License for the specific language governing permissions and
+   limitations under the License.

Model_License

@@ -0,0 +1,71 @@
+The CogVideoX License
+
+1. Definitions
+
+“Licensor” means the CogVideoX Model Team that distributes its Software.
+
+“Software” means the CogVideoX model parameters made available under this license.
+
+2. License Grant
+
+Under the terms and conditions of this license, the licensor hereby grants you a non-exclusive, worldwide, non-transferable, non-sublicensable, revocable, royalty-free copyright license. The intellectual property rights of the generated content belong to the user to the extent permitted by applicable local laws.
+This license allows you to freely use all open-source models in this repository for academic research. Users who wish to use the models for commercial purposes must register and obtain a basic commercial license in https://open.bigmodel.cn/mla/form .
+Users who have registered and obtained the basic commercial license can use the models for commercial activities for free, but must comply with all terms and conditions of this license. Additionally, the number of service users (visits) for your commercial activities must not exceed 1 million visits per month.
+If the number of service users (visits) for your commercial activities exceeds 1 million visits per month, you need to contact our business team to obtain more commercial licenses.
+The above copyright statement and this license statement should be included in all copies or significant portions of this software.
+
+3. Restriction
+
+You will not use, copy, modify, merge, publish, distribute, reproduce, or create derivative works of the Software, in whole or in part, for any military, or illegal purposes.
+
+You will not use the Software for any act that may undermine China's national security and national unity, harm the public interest of society, or infringe upon the rights and interests of human beings.
+
+4. Disclaimer
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE.
+
+5. Limitation of Liability
+
+EXCEPT TO THE EXTENT PROHIBITED BY APPLICABLE LAW, IN NO EVENT AND UNDER NO LEGAL THEORY, WHETHER BASED IN TORT, NEGLIGENCE, CONTRACT, LIABILITY, OR OTHERWISE WILL ANY LICENSOR BE LIABLE TO YOU FOR ANY DIRECT, INDIRECT, SPECIAL, INCIDENTAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES, OR ANY OTHER COMMERCIAL LOSSES, EVEN IF THE LICENSOR HAS BEEN ADVISED OF THE POSSIBILITY OF SUCH DAMAGES.
+
+6. Dispute Resolution
+
+This license shall be governed and construed in accordance with the laws of People’s Republic of China. Any dispute arising from or in connection with this License shall be submitted to Haidian District People's Court in Beijing.
+
+Note that the license is subject to update to a more comprehensive version.  For any questions related to the license and copyright, please contact us at license@zhipuai.cn.
+
+1. 定义
+
+“许可方”是指分发其软件的 CogVideoX 模型团队。
+
+“软件”是指根据本许可提供的 CogVideoX 模型参数。
+
+2. 许可授予
+
+根据本许可的条款和条件，许可方特此授予您非排他性、全球性、不可转让、不可再许可、可撤销、免版税的版权许可。生成内容的知识产权所属，可根据适用当地法律的规定，在法律允许的范围内由用户享有生成内容的知识产权或其他权利。
+本许可允许您免费使用本仓库中的所有开源模型进行学术研究。对于希望将模型用于商业目的的用户，需在 https://open.bigmodel.cn/mla/form 完成登记并获得基础商用授权。
+
+经过登记并获得基础商用授权的用户可以免费使用本模型进行商业活动，但必须遵守本许可的所有条款和条件。
+在本许可证下，您的商业活动的服务用户数量（访问量）不得超过100万人次访问 / 每月。如果超过，您需要与我们的商业团队联系以获得更多的商业许可。
+上述版权声明和本许可声明应包含在本软件的所有副本或重要部分中。
+
+3.限制
+
+您不得出于任何军事或非法目的使用、复制、修改、合并、发布、分发、复制或创建本软件的全部或部分衍生作品。
+
+您不得利用本软件从事任何危害国家安全和国家统一、危害社会公共利益、侵犯人身权益的行为。
+
+4.免责声明
+
+本软件“按原样”提供，不提供任何明示或暗示的保证，包括但不限于对适销性、特定用途的适用性和非侵权性的保证。
+在任何情况下，作者或版权持有人均不对任何索赔、损害或其他责任负责，无论是在合同诉讼、侵权行为还是其他方面，由软件或软件的使用或其他交易引起、由软件引起或与之相关 软件。
+
+5. 责任限制
+
+除适用��律禁止的范围外，在任何情况下且根据任何法律理论，无论是基于侵权行为、疏忽、合同、责任或其他原因，任何许可方均不对您承担任何直接、间接、特殊、偶然、示范性、 或间接损害，或任何其他商业损失，即使许可人已被告知此类损害的可能性。
+
+6.争议解决
+
+本许可受中华人民共和国法律管辖并按其解释。 因本许可引起的或与本许可有关的任何争议应提交北京市海淀区人民法院。
+
+请注意，许可证可能会更新到更全面的版本。 有关许可和版权的任何问题，请通过 license@zhipuai.cn 与我们联系。

README.md

@@ -1,12 +1,163 @@
 ---
 title: CogVideo
-emoji: 🐨
-colorFrom: blue
-colorTo: red
+app_file: gradio_demo.py
 sdk: gradio
 sdk_version: 4.41.0
-app_file: app.py
-pinned: false
 ---
+# CogVideo && CogVideoX
 
-Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference
+[中文阅读](./README_zh.md)
+
+<div align="center">
+<img src=resources/logo.svg width="50%"/>
+</div>
+<p align="center">
+🤗 Experience on <a href="https://huggingface.co/spaces/THUDM/CogVideoX" target="_blank">CogVideoX Huggingface Space</a>
+</p>
+<p align="center">
+📚 Check here to view <a href="resources/CogVideoX.pdf" target="_blank">Paper</a>
+</p>
+<p align="center">
+    👋 Join our <a href="resources/WECHAT.md" target="_blank">WeChat</a> and <a href="https://discord.gg/Ewaabk6s" target="_blank">Discord</a> 
+</p>
+<p align="center">
+📍 Visit <a href="https://chatglm.cn/video?fr=osm_cogvideox">清影</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>
+
+## Update and News
+
+- 🔥 **News**: ``2024/8/6``: We have also open-sourced **3D Causal VAE** used in **CogVideoX-2B**, which can reconstruct
+  the video almost losslessly.
+- 🔥 **News**: ``2024/8/6``: We have open-sourced **CogVideoX-2B**，the first model in the CogVideoX series of video
+  generation models.
+- 🌱 **Source**: ```2022/5/19```: We have open-sourced CogVideo (now you can see in `CogVideo` branch)，the **first** open-sourced pretrained text-to-video model, and you can check [ICLR'23 CogVideo Paper](https://arxiv.org/abs/2205.15868) for technical details.
+
+**More powerful models with larger parameter sizes are on the way~ Stay tuned!**
+
+## CogVideoX-2B Gallery
+
+<div align="center">
+  <video src="https://github.com/user-attachments/assets/ea3af39a-3160-4999-90ec-2f7863c5b0e9" width="80%" controls autoplay></video>
+  <p>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.</p>
+</div>
+
+<div align="center">
+  <video src="https://github.com/user-attachments/assets/9de41efd-d4d1-4095-aeda-246dd834e91d" width="80%" controls autoplay></video>
+  <p>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.</p>
+</div>
+
+<div align="center">
+  <video src="https://github.com/user-attachments/assets/941d6661-6a8d-4a1b-b912-59606f0b2841" width="80%" controls autoplay></video>
+  <p>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.</p>
+</div>
+
+<div align="center">
+  <video src="https://github.com/user-attachments/assets/938529c4-91ae-4f60-b96b-3c3947fa63cb" width="80%" controls autoplay></video>
+  <p>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.</p>
+</div>
+
+## Model Introduction
+
+CogVideoX is an open-source version of the video generation model, which is homologous
+to [清影](https://chatglm.cn/video?fr=osm_cogvideox).
+
+The table below shows the list of video generation models we currently provide,
+along with related basic information:
+
+| Model Name                                | CogVideoX-2B                                                                                                                                                                                        | 
+|-------------------------------------------|-----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------|
+| Prompt Language                           | English                                                                                                                                                                                             | 
+| GPU Memory Required for Inference (FP16)  | 18GB if using [SAT](https://github.com/THUDM/SwissArmyTransformer); 36GB if using diffusers (will be optimized before the PR is merged)                                                             | 
+| GPU Memory Required for Fine-tuning(bs=1) | 40GB                                                                                                                                                                                                |
+| Prompt Max  Length                        | 226 Tokens                                                                                                                                                                                          |
+| Video Length                              | 6 seconds                                                                                                                                                                                           | 
+| Frames Per Second                         | 8 frames                                                                                                                                                                                            | 
+| Resolution                                | 720 * 480                                                                                                                                                                                           |
+| Quantized Inference                       | Not Supported                                                                                                                                                                                       |          
+| Multi-card Inference                      | Not Supported                                                                                                                                                                                       |                             
+| Download Link (HF diffusers Model)        | 🤗 [Huggingface](https://huggingface.co/THUDM/CogVideoX-2B)   [🤖 ModelScope](https://modelscope.cn/models/ZhipuAI/CogVideoX-2b)   [💫 WiseModel](https://wisemodel.cn/models/ZhipuAI/CogVideoX-2b) |
+| Download Link (SAT Model)                 | [SAT](./sat/README.md)                                                                                                                                                                              |
+
+## 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.
+
+### Inference
+
++ [cli_demo](inference/cli_demo.py): A more detailed explanation of the inference code, mentioning the significance of common parameters.
++ [cli_vae_demo](inference/cli_vae_demo.py): Executing the VAE inference code alone currently requires 71GB of memory, but it will be optimized in the future.
++ [convert_demo](inference/convert_demo.py): How to convert user input into a format suitable for CogVideoX. Because CogVideoX is trained on long caption, we need to convert the input text to be consistent with the training distribution using a LLM. By default, the script uses GLM4, but it can also be replaced with any other LLM such as GPT, Gemini, etc.
++ [web_demo](inference/web_demo.py): A simple streamlit web application demonstrating how to use the CogVideoX-2B model to generate videos.
+
+<div style="text-align: center;">
+    <img src="resources/web_demo.png" style="width: 100%; height: auto;" />
+</div>
+
+### 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): Convert SAT model weights to Huggingface model weights. 
++ [caption_demo](tools/caption): Caption tool, a model that understands videos and outputs them in text.
+
+## Project Plan
+
+- [x] Open source CogVideoX model
+    - [x] Open source 3D Causal VAE used in CogVideoX.
+    - [x] CogVideoX model inference example (CLI / Web Demo)
+    - [x] CogVideoX online experience demo (Huggingface Space)
+    - [x] CogVideoX open source model API interface example (Huggingface)
+    - [x] CogVideoX model fine-tuning example (SAT)
+    - [ ] CogVideoX model fine-tuning example (Huggingface / SAT)
+    - [ ] Open source CogVideoX-Pro (adapted for CogVideoX-2B suite)
+    - [x] Release CogVideoX technical report
+
+We welcome your contributions. You can click [here](resources/contribute.md) for more information.
+
+## Model License
+
+The code in this repository is released under the [Apache 2.0 License](LICENSE).
+
+The model weights and implementation code are released under the [CogVideoX LICENSE](MODEL_LICENSE).
+
+## CogVideo(ICLR'23)
+The official repo for the paper: [CogVideo: Large-scale Pretraining for Text-to-Video Generation via Transformers](https://arxiv.org/abs/2205.15868) is on the [CogVideo branch](https://github.com/THUDM/CogVideo/tree/CogVideo)
+
+**CogVideo is able to generate relatively high-frame-rate videos.**
+A 4-second clip of 32 frames is shown below. 
+
+![High-frame-rate sample](https://raw.githubusercontent.com/THUDM/CogVideo/CogVideo/assets/appendix-sample-highframerate.png)
+
+![Intro images](https://raw.githubusercontent.com/THUDM/CogVideo/CogVideo/assets/intro-image.png)
+<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={Zhuoyi Yang and Jiayan Teng and Wendi Zheng and Ming Ding and Shiyu Huang and JiaZheng Xu and Yuanming Yang and Xiaohan Zhang and Xiaotao Gu and Guanyu Feng and Da Yin and Wenyi Hong and Weihan Wang and Yean Cheng and Yuxuan Zhang and Ting Liu and Bin Xu and Yuxiao Dong and Jie Tang},
+      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}
+}
+```

README_zh.md

@@ -0,0 +1,149 @@
+# CogVideo && CogVideoX
+
+[Read this in English.](./README_zh)
+
+
+<div align="center">
+<img src=resources/logo.svg width="50%"/>
+</div>
+<p align="center">
+🤗 在 <a href="https://huggingface.co/spaces/THUDM/CogVideoX" target="_blank">CogVideoX Huggingface Space</a> 体验视频生成模型
+</p>
+<p align="center">
+📚 查看 <a href="resources/CogVideoX.pdf" target="_blank">论文</a>
+</p>
+<p align="center">
+    👋 加入我们的 <a href="resources/WECHAT.md" target="_blank">微信</a> 和  <a href="https://discord.gg/Ewaabk6s" 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**: ``2024/8/6``: 我们开源 **3D Causal VAE**，用于 **CogVideoX-2B**，可以几乎无损地重构视频。
+- 🔥 **News**: ``2024/8/6``: 我们开源 CogVideoX 系列视频生成模型的第一个模型, **CogVideoX-2B**。
+- 🌱 **Source**: ```2022/5/19```: 我们开源了 CogVideo 视频生成模型（现在你可以在 `CogVideo` 分支中看到），这是首个开源的基于 Transformer 的大型文本生成视频模型，您可以访问 [ICLR'23 论文](https://arxiv.org/abs/2205.15868) 查看技术细节。
+**性能更强，参数量更大的模型正在到来的路上～，欢迎关注**
+
+## CogVideoX-2B 视频作品
+
+<div align="center">
+  <video src="https://github.com/user-attachments/assets/ea3af39a-3160-4999-90ec-2f7863c5b0e9" width="80%" controls autoplay></video>
+  <p>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.</p>
+</div>
+
+<div align="center">
+  <video src="https://github.com/user-attachments/assets/9de41efd-d4d1-4095-aeda-246dd834e91d" width="80%" controls autoplay></video>
+  <p>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.</p>
+</div>
+
+<div align="center">
+  <video src="https://github.com/user-attachments/assets/941d6661-6a8d-4a1b-b912-59606f0b2841" width="80%" controls autoplay></video>
+  <p>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.</p>
+</div>
+
+<div align="center">
+  <video src="https://github.com/user-attachments/assets/938529c4-91ae-4f60-b96b-3c3947fa63cb" width="80%" controls autoplay></video>
+  <p>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.</p>
+</div>
+
+## 模型介绍
+
+CogVideoX是 [清影](https://chatglm.cn/video?fr=osm_cogvideox) 同源的开源版本视频生成模型。
+
+下表战展示目前我们提供的视频生成模型列表，以及相关基础信息:
+
+| 模型名字                | CogVideoX-2B                                                                                                                         | 
+|---------------------|--------------------------------------------------------------------------------------------------------------------------------------|
+| 提示词语言               | English                                                                                                                              | 
+| 推理显存消耗 (FP-16)      | 36GB using diffusers (will be optimized before the PR is merged) and 18GB using [SAT](https://github.com/THUDM/SwissArmyTransformer) | 
+| 微调显存消耗 (bs=1)       | 42GB                                                                                                                                 |
+| 提示词长度上限             | 226 Tokens                                                                                                                           |
+| 视频长度                | 6 seconds                                                                                                                            | 
+| 帧率（每秒）              | 8 frames                                                                                                                             | 
+| 视频分辨率               | 720 * 480                                                                                                                            |
+| 量化推理                | 不支持                                                                                                                                  |          
+| 多卡推理                | 不支持                                                                                                                                  |                             
+| 下载地址 (Diffusers 模型) | 🤗 [Huggingface](https://huggingface.co/THUDM/CogVideoX-2B)  [🤖 ModelScope](https://modelscope.cn/models/ZhipuAI/CogVideoX-2b)      |
+| 下载地址 (SAT 模型)       | [SAT](./sat/README_zh.md)                                                                                                            |
+
+## 项目结构
+
+本开源仓库将带领开发者快速上手 **CogVideoX** 开源模型的基础调用方式、微调示例。
+
+### inference
+
++ [cli_demo](inference/cli_demo.py): 更详细的推理代码讲解，常见参数的意义，在这里都会提及。
++ [cli_vae_demo](inference/cli_vae_demo.py): 单独执行VAE的推理代码，目前需要71GB显存，将来会优化。
++ [convert_demo](inference/convert_demo.py): 如何将用户的输入转换成适合 CogVideoX的长输入。因为CogVideoX是在长文本上训练的，所以我们需要把输入文本的分布通过LLM转换为和训练一致的长文本。脚本中默认使用GLM4，也可以替换为GPT、Gemini等任意大语言模型。
++ [web_demo](inference/web_demo.py): 一个简单的streamlit网页应用，展示如何使用 CogVideoX-2B 模型生成视频。
+
+<div style="text-align: center;">
+    <img src="resources/web_demo.png" style="width: 100%; height: auto;" />
+</div>
+
+### 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 工具，对视频理解并用文字输出的模型。
+
+## 项目规划
+
+- [x] CogVideoX 模型开源
+    - [x] CogVideoX 模型推理示例 (CLI / Web Demo)
+    - [x] CogVideoX 在线体验示例 (Huggingface Space)
+    - [x] CogVideoX 开源模型API接口示例 (Huggingface)
+    - [x] CogVideoX 模型微调示例 (SAT)
+    - [ ] CogVideoX 模型微调示例 (Huggingface / SAT)
+    - [ ] CogVideoX-Pro 开源(适配 CogVideoX-2B 套件)
+    - [ ] CogVideoX 技术报告公开
+
+我们欢迎您的贡献，您可以点击[这里](resources/contribute_zh.md)查看更多信息。
+
+## 模型协议
+
+本仓库代码使用 [Apache 2.0 协议](LICENSE) 发布。
+
+本模型权重和模型实现代码根据 [CogVideoX LICENSE](MODEL_LICENSE) 许可证发布。
+
+## CogVideo(ICLR'23) 
+ [CogVideo: Large-scale Pretraining for Text-to-Video Generation via Transformers](https://arxiv.org/abs/2205.15868) 的官方repo位于[CogVideo branch](https://github.com/THUDM/CogVideo/tree/CogVideo)。
+
+**CogVideo可以生成高帧率视频，下面展示了一个32帧的4秒视频。**
+
+![High-frame-rate sample](https://raw.githubusercontent.com/THUDM/CogVideo/CogVideo/assets/appendix-sample-highframerate.png)
+
+![Intro images](https://raw.githubusercontent.com/THUDM/CogVideo/CogVideo/assets/intro-image.png)
+
+
+<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={Zhuoyi Yang and Jiayan Teng and Wendi Zheng and Ming Ding and Shiyu Huang and JiaZheng Xu and Yuanming Yang and Xiaohan Zhang and Xiaotao Gu and Guanyu Feng and Da Yin and Wenyi Hong and Weihan Wang and Yean Cheng and Yuxuan Zhang and Ting Liu and Bin Xu and Yuxiao Dong and Jie Tang},
+      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}
+}
+```

gradio_demo.py

@@ -0,0 +1,254 @@
+import os
+import tempfile
+import threading
+import time
+
+import gradio as gr
+import numpy as np
+import torch
+from diffusers import CogVideoXPipeline
+from datetime import datetime, timedelta
+from openai import OpenAI
+import spaces
+import imageio
+import moviepy.editor as mp
+from typing import List, Union
+import PIL
+
+dtype = torch.bfloat16
+device = "cuda" if torch.cuda.is_available() else "cpu"
+pipe = CogVideoXPipeline.from_pretrained("THUDM/CogVideoX-2b", torch_dtype=dtype).to(device)
+
+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 export_to_video_imageio(
+        video_frames: Union[List[np.ndarray], List[PIL.Image.Image]], output_video_path: str = None, fps: int = 8
+) -> str:
+    """
+    Export the video frames to a video file using imageio lib to Avoid "green screen" issue (for example CogVideoX)
+    """
+    if output_video_path is None:
+        output_video_path = tempfile.NamedTemporaryFile(suffix=".mp4").name
+
+    if isinstance(video_frames[0], PIL.Image.Image):
+        video_frames = [np.array(frame) for frame in video_frames]
+
+    with imageio.get_writer(output_video_path, fps=fps) as writer:
+        for frame in video_frames:
+            writer.append_data(frame)
+
+    return output_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-0520",
+            temperature=0.01,
+            top_p=0.7,
+            stream=False,
+            max_tokens=250,
+        )
+        if response.choices:
+            return response.choices[0].message.content
+    return prompt
+
+
+@spaces.GPU(duration=240)
+def infer(
+        prompt: str,
+        num_inference_steps: int,
+        guidance_scale: float,
+        progress=gr.Progress(track_tqdm=True)
+):
+    torch.cuda.empty_cache()
+
+    prompt_embeds, _ = pipe.encode_prompt(
+        prompt=prompt,
+        negative_prompt=None,
+        do_classifier_free_guidance=True,
+        num_videos_per_prompt=1,
+        max_sequence_length=226,
+        device=device,
+        dtype=dtype,
+    )
+
+    video = pipe(
+        num_inference_steps=num_inference_steps,
+        guidance_scale=guidance_scale,
+        prompt_embeds=prompt_embeds,
+        negative_prompt_embeds=torch.zeros_like(prompt_embeds),
+    ).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_imageio(tensor[1:], video_path)
+    return video_path
+
+def convert_to_gif(video_path):
+    clip = mp.VideoFileClip(video_path)
+    clip = clip.set_fps(8)
+    clip = clip.resize(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)
+        output_dir = './output'
+        for filename in os.listdir(output_dir):
+            file_path = os.path.join(output_dir, 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)  # Sleep for 10 minutes
+
+
+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-2B Huggingface Space🤗
+           </div>
+           <div style="text-align: center;">
+               <a href="https://huggingface.co/THUDM/CogVideoX-2b">🤗 Model Hub</a> |
+               <a href="https://github.com/THUDM/CogVideo">🌐 Github</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 experiential use only. 
+            Users should strictly adhere to local laws and ethics.
+            </div>
+           """)
+    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>"
+                            "Turn Inference Steps larger if you want more detailed video, but it will be slower.<br>"
+                            "50 steps are recommended for most cases. will cause 120 seconds for inference.<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)
+
+    gr.Markdown("""
+        <table border="1" style="width: 100%; text-align: left; margin-top: 20px;">
+            <tr>
+                <th>Prompt</th>
+                <th>Video URL</th>
+                <th>Inference Steps</th>
+                <th>Guidance Scale</th>
+            </tr>
+            <tr>
+                <td>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.</td>
+                <td><a href="https://github.com/THUDM/CogVideo/raw/main/resources/videos/1.mp4">Video 1</a></td>
+                <td>50</td>
+                <td>6</td>
+            </tr>
+            <tr>
+                <td>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 it’s 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.</td>
+                <td><a href="https://github.com/THUDM/CogVideo/raw/main/resources/videos/2.mp4">Video 2</a></td>
+                <td>50</td>
+                <td>6</td>
+            </tr>
+            <tr>
+                <td>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.</td>
+                <td><a href="https://github.com/THUDM/CogVideo/raw/main/resources/videos/3.mp4">Video 3</a></td>
+                <td>50</td>
+                <td>6</td>
+            </tr>
+            <tr>
+                <td>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.</td>
+                <td><a href="https://github.com/THUDM/CogVideo/raw/main/resources/videos/4.mp4">Video 4</a></td>
+                <td>50</td>
+                <td>6</td>
+            </tr>
+        </table>
+    """)
+
+
+    def generate(prompt, num_inference_steps, guidance_scale, 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(server_name="127.0.0.1", server_port=7870, share=True)

inference/cli_demo.py

@@ -0,0 +1,127 @@
+"""
+This script demonstrates how to generate a video from a text prompt using CogVideoX with 🤗Huggingface Diffusers Pipeline.
+
+Note:
+    This script requires the `diffusers>=0.30.0` library to be installed.
+    If the video exported using OpenCV appears “completely green” and cannot be viewed, lease switch to a different player to watch it. This is a normal phenomenon.
+
+Run the script:
+    $ python cli_demo.py --prompt "A girl ridding a bike." --model_path THUDM/CogVideoX-2b
+
+"""
+
+import argparse
+import tempfile
+from typing import Union, List
+
+import PIL
+import imageio
+import numpy as np
+import torch
+from diffusers import CogVideoXPipeline
+
+
+def export_to_video_imageio(
+    video_frames: Union[List[np.ndarray], List[PIL.Image.Image]], output_video_path: str = None, fps: int = 8
+) -> str:
+    """
+    Export the video frames to a video file using imageio lib to Avoid "green screen" issue (for example CogVideoX)
+    """
+    if output_video_path is None:
+        output_video_path = tempfile.NamedTemporaryFile(suffix=".mp4").name
+    if isinstance(video_frames[0], PIL.Image.Image):
+        video_frames = [np.array(frame) for frame in video_frames]
+    with imageio.get_writer(output_video_path, fps=fps) as writer:
+        for frame in video_frames:
+            writer.append_data(frame)
+    return output_video_path
+
+
+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,
+    device: str = "cuda",
+    dtype: torch.dtype = torch.float16,
+):
+    """
+    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.
+    - device (str): The device to use for computation (e.g., "cuda" or "cpu").
+    - dtype (torch.dtype): The data type for computation (default is torch.float16).
+    """
+
+    # Load the pre-trained CogVideoX pipeline with the specified precision (float16) and move it to the specified device
+    pipe = CogVideoXPipeline.from_pretrained(model_path, torch_dtype=dtype).to(device)
+
+    # Encode the prompt to get the prompt embeddings
+    prompt_embeds, _ = pipe.encode_prompt(
+        prompt=prompt,  # The textual description for video generation
+        negative_prompt=None,  # The negative prompt to guide the video generation
+        do_classifier_free_guidance=True,  # Whether to use classifier-free guidance
+        num_videos_per_prompt=num_videos_per_prompt,  # Number of videos to generate per prompt
+        max_sequence_length=226,  # Maximum length of the sequence, must be 226
+        device=device,  # Device to use for computation
+        dtype=dtype,  # Data type for computation
+    )
+
+    # Generate the video frames using the pipeline
+    video = pipe(
+        num_inference_steps=num_inference_steps,  # Number of inference steps
+        guidance_scale=guidance_scale,  # Guidance scale for classifier-free guidance
+        prompt_embeds=prompt_embeds,  # Encoded prompt embeddings
+        negative_prompt_embeds=torch.zeros_like(prompt_embeds),  # Not Supported negative prompt
+    ).frames[0]
+
+    # Export the generated frames to a video file. fps must be 8
+    export_to_video_imageio(video, output_path, fps=8)
+
+
+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-2b", help="The path of the pre-trained model to be used"
+    )
+    parser.add_argument(
+        "--output_path", type=str, default="./output.mp4", help="The path where the generated video will be saved"
+    )
+    parser.add_argument(
+        "--num_inference_steps", type=int, default=50, help="Number of steps for the inference process"
+    )
+    parser.add_argument("--guidance_scale", type=float, default=6.0, help="The scale for classifier-free guidance")
+    parser.add_argument("--num_videos_per_prompt", type=int, default=1, help="Number of videos to generate per prompt")
+    parser.add_argument(
+        "--device", type=str, default="cuda", help="The device to use for computation (e.g., 'cuda' or 'cpu')"
+    )
+
+    parser.add_argument(
+        "--dtype", type=str, default="float16", help="The data type for computation (e.g., 'float16' or 'float32')"
+    )
+
+    args = parser.parse_args()
+
+    # Convert dtype argument to torch.dtype, NOT suggest BF16.
+    dtype = torch.float16 if args.dtype == "float16" else torch.float32
+
+    # main function to generate video.
+    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,
+        device=args.device,
+        dtype=dtype,
+    )

inference/cli_vae_demo.py

@@ -0,0 +1,103 @@
+"""
+This script demonstrates how to encode video frames using a pre-trained CogVideoX model with 🤗 Huggingface Diffusers.
+
+Note:
+    This script requires the `diffusers>=0.30.0` library to be installed.
+    If the video appears “completely green” and cannot be viewed, please switch to a different player to watch it. This is a normal phenomenon.
+    Cost 71GB of GPU memory for encoding a 6s video at 720p resolution.
+
+Run the script:
+    $ python cli_demo.py --model_path THUDM/CogVideoX-2b --video_path path/to/video.mp4 --output_path path/to/output
+
+"""
+
+import argparse
+import torch
+import imageio
+import numpy as np
+from diffusers import AutoencoderKLCogVideoX
+from torchvision import transforms
+
+
+def vae_demo(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.
+    """
+    # Load the pre-trained model
+    model = AutoencoderKLCogVideoX.from_pretrained(model_path, torch_dtype=dtype).to(device)
+
+    # Load video frames
+    video_reader = imageio.get_reader(video_path, "ffmpeg")
+    frames = []
+    for frame in video_reader:
+        frames.append(frame)
+    video_reader.close()
+
+    # Transform frames to Tensor
+    transform = transforms.Compose(
+        [
+            transforms.ToTensor(),
+        ]
+    )
+    frames_tensor = torch.stack([transform(frame) for frame in frames]).to(device)
+
+    # Add batch dimension and reshape to [1, 3, 49, 480, 720]
+    frames_tensor = frames_tensor.permute(1, 0, 2, 3).unsqueeze(0).to(dtype).to(device)
+
+    # Run the model with Encoder and Decoder
+    with torch.no_grad():
+        output = model(frames_tensor)
+
+    return output
+
+
+def save_video(tensor, output_path):
+    """
+    Saves the encoded video frames to a video file.
+
+    Parameters:
+    - tensor (torch.Tensor): The encoded video frames.
+    - output_path (str): The path to save the output video.
+    """
+    # Remove batch dimension and permute back to [49, 480, 720, 3]
+    frames = tensor[0].squeeze(0).permute(1, 2, 3, 0).cpu().numpy()
+
+    # Clip values to [0, 1] and convert to uint8
+    frames = np.clip(frames, 0, 1)
+    frames = (frames * 255).astype(np.uint8)
+
+    # Save frames to video
+    writer = imageio.get_writer(output_path + "/output.mp4", fps=30)
+    for frame in frames:
+        writer.append_data(frame)
+    writer.close()
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser(description="Convert a CogVideoX model to Diffusers")
+    parser.add_argument("--model_path", type=str, required=True, help="The path to the CogVideoX model")
+    parser.add_argument("--video_path", type=str, required=True, help="The path to the video file")
+    parser.add_argument("--output_path", type=str, default="./", help="The path to save the output video")
+    parser.add_argument(
+        "--dtype", type=str, default="float16", help="The data type for computation (e.g., 'float16' or 'float32')"
+    )
+    parser.add_argument(
+        "--device", type=str, default="cuda", help="The device to use for computation (e.g., 'cuda' or 'cpu')"
+    )
+    args = parser.parse_args()
+
+    # Set device and dtype
+    device = torch.device(args.device)
+    dtype = torch.float16 if args.dtype == "float16" else torch.float32
+
+    output = vae_demo(args.model_path, args.video_path, dtype, device)
+    save_video(output, args.output_path)

inference/convert_demo.py

@@ -0,0 +1,92 @@
+"""
+
+The CogVideoX model is pre-trained and fine-tuned using longer and more detailed prompts.Therefore, it requires highly granular and detailed prompts as input.This script aims to transform user inputs into executable inputs for CogVideoX, enabling superior video generation.
+
+This step is not mandatory; the model will still function correctly and without errors even if the prompts are not refined using this script. However, we strongly recommend using it to ensure the generation of high-quality videos.
+
+Note:
+Please set the OPENAI_API_KEY and OPENAI_BASE_URL(if needed) environment variable to your OpenAI API key before running this script.
+
+Run the script:
+    $ python convert_demo.py --prompt "A girl ridding a bike."  # Using with OpenAI's API
+"""
+
+import argparse
+
+from openai import OpenAI
+
+
+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):
+    """
+    Convert a prompt to a format that can be used by the model for inference
+    """
+
+    client = OpenAI()
+    text = prompt.strip()
+
+    for i in range(retry_times):
+        response = client.chat.completions.create(
+            messages=[
+                {"role": "system", "content": f"{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-0520",  # glm-4-0520 and gpt-4o have be tested
+            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")
+    args = parser.parse_args()
+
+    converted_prompt = convert_prompt(args.prompt, args.retry_times)
+    print(converted_prompt)

inference/web_demo.py

@@ -0,0 +1,214 @@
+"""
+This script is used to create a Streamlit web application for generating videos using the CogVideoX model.
+
+Run the script using Streamlit:
+    $ export OPENAI_API_KEY=your OpenAI Key or ZhiupAI Key
+    $ export OPENAI_BASE_URL=https://open.bigmodel.cn/api/paas/v4/  # using with ZhipuAI, Not using this when using OpenAI
+    $ streamlit run web_demo.py
+"""
+
+import base64
+import json
+import os
+import time
+from datetime import datetime
+from typing import List
+
+import imageio
+import numpy as np
+import streamlit as st
+import torch
+from convert_demo import convert_prompt
+from diffusers import CogVideoXPipeline
+
+
+model_path: str = "THUDM/CogVideoX-2b"
+
+
+# Load the model at the start
+@st.cache_resource
+def load_model(model_path: str, dtype: torch.dtype, device: str) -> CogVideoXPipeline:
+    """
+    Load the CogVideoX model.
+
+    Args:
+    - model_path (str): Path to the model.
+    - dtype (torch.dtype): Data type for model.
+    - device (str): Device to load the model on.
+
+    Returns:
+    - CogVideoXPipeline: Loaded model pipeline.
+    """
+    return CogVideoXPipeline.from_pretrained(model_path, torch_dtype=dtype).to(device)
+
+
+# Define a function to generate video based on the provided prompt and model path
+def generate_video(
+    pipe: CogVideoXPipeline,
+    prompt: str,
+    num_inference_steps: int = 50,
+    guidance_scale: float = 6.0,
+    num_videos_per_prompt: int = 1,
+    device: str = "cuda",
+    dtype: torch.dtype = torch.float16,
+) -> List[np.ndarray]:
+    """
+    Generate a video based on the provided prompt and model path.
+
+    Args:
+    - pipe (CogVideoXPipeline): The pipeline for generating videos.
+    - prompt (str): Text prompt for video generation.
+    - num_inference_steps (int): Number of inference steps.
+    - guidance_scale (float): Guidance scale for generation.
+    - num_videos_per_prompt (int): Number of videos to generate per prompt.
+    - device (str): Device to run the generation on.
+    - dtype (torch.dtype): Data type for the model.
+
+    Returns:
+    - List[np.ndarray]: Generated video frames.
+    """
+    prompt_embeds, _ = pipe.encode_prompt(
+        prompt=prompt,
+        negative_prompt=None,
+        do_classifier_free_guidance=True,
+        num_videos_per_prompt=num_videos_per_prompt,
+        max_sequence_length=226,
+        device=device,
+        dtype=dtype,
+    )
+
+    # Generate video
+    video = pipe(
+        num_inference_steps=num_inference_steps,
+        guidance_scale=guidance_scale,
+        prompt_embeds=prompt_embeds,
+        negative_prompt_embeds=torch.zeros_like(prompt_embeds),
+    ).frames[0]
+    return video
+
+
+def save_video(video: List[np.ndarray], path: str, fps: int = 8) -> None:
+    """
+    Save the generated video to a file.
+
+    Args:
+    - video (List[np.ndarray]): Video frames.
+    - path (str): Path to save the video.
+    - fps (int): Frames per second for the video.
+    """
+    # Remove the first frame
+    video = video[1:]
+
+    writer = imageio.get_writer(path, fps=fps, codec="libx264")
+    for frame in video:
+        np_frame = np.array(frame)
+        writer.append_data(np_frame)
+
+    writer.close()
+
+
+def save_metadata(
+    prompt: str,
+    converted_prompt: str,
+    num_inference_steps: int,
+    guidance_scale: float,
+    num_videos_per_prompt: int,
+    path: str,
+) -> None:
+    """
+    Save metadata to a JSON file.
+
+    Args:
+    - prompt (str): Original prompt.
+    - converted_prompt (str): Converted prompt.
+    - num_inference_steps (int): Number of inference steps.
+    - guidance_scale (float): Guidance scale.
+    - num_videos_per_prompt (int): Number of videos per prompt.
+    - path (str): Path to save the metadata.
+    """
+    metadata = {
+        "prompt": prompt,
+        "converted_prompt": converted_prompt,
+        "num_inference_steps": num_inference_steps,
+        "guidance_scale": guidance_scale,
+        "num_videos_per_prompt": num_videos_per_prompt,
+    }
+    with open(path, "w") as f:
+        json.dump(metadata, f, indent=4)
+
+
+def main() -> None:
+    """
+    Main function to run the Streamlit web application.
+    """
+    st.set_page_config(page_title="CogVideoX-Demo", page_icon="🎥", layout="wide")
+    st.write("# CogVideoX 🎥")
+    dtype: torch.dtype = torch.float16
+    device: str = "cuda"
+
+    global pipe
+    pipe = load_model(model_path, dtype, device)
+
+    with st.sidebar:
+        st.info("It will take some time to generate a video (~90 seconds per videos in 50 steps).", icon="ℹ️")
+        num_inference_steps: int = st.number_input("Inference Steps", min_value=1, max_value=100, value=50)
+        guidance_scale: float = st.number_input("Guidance Scale", min_value=0.0, max_value=20.0, value=6.0)
+        num_videos_per_prompt: int = st.number_input("Videos per Prompt", min_value=1, max_value=10, value=1)
+
+        share_links_container = st.empty()
+
+    prompt: str = st.chat_input("Prompt")
+
+    if prompt:
+        # Not Necessary, Suggestions
+        with st.spinner("Refining prompts..."):
+            converted_prompt = convert_prompt(prompt=prompt, retry_times=1)
+            if converted_prompt is None:
+                st.error("Failed to Refining the prompt, Using origin one.")
+
+        st.info(f"**Origin prompt:**  \n{prompt}  \n  \n**Convert prompt:**  \n{converted_prompt}")
+        torch.cuda.empty_cache()
+
+        with st.spinner("Generating Video..."):
+            start_time = time.time()
+            video_paths = []
+
+            timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
+            output_dir = f"./output/{timestamp}"
+            os.makedirs(output_dir, exist_ok=True)
+
+            metadata_path = os.path.join(output_dir, "config.json")
+            save_metadata(
+                prompt, converted_prompt, num_inference_steps, guidance_scale, num_videos_per_prompt, metadata_path
+            )
+
+            for i in range(num_videos_per_prompt):
+                video_path = os.path.join(output_dir, f"output_{i + 1}.mp4")
+
+                video = generate_video(
+                    pipe, converted_prompt or prompt, num_inference_steps, guidance_scale, 1, device, dtype
+                )
+                save_video(video, video_path, fps=8)
+                video_paths.append(video_path)
+                with open(video_path, "rb") as video_file:
+                    video_bytes: bytes = video_file.read()
+                    st.video(video_bytes, autoplay=True, loop=True, format="video/mp4")
+                torch.cuda.empty_cache()
+
+            used_time: float = time.time() - start_time
+            st.success(f"Videos generated in {used_time:.2f} seconds.")
+
+            # Create download links in the sidebar
+            with share_links_container:
+                st.sidebar.write("### Download Links:")
+                for video_path in video_paths:
+                    video_name = os.path.basename(video_path)
+                    with open(video_path, "rb") as f:
+                        video_bytes: bytes = f.read()
+                    b64_video = base64.b64encode(video_bytes).decode()
+                    href = f'<a href="data:video/mp4;base64,{b64_video}" download="{video_name}">Download {video_name}</a>'
+                    st.sidebar.markdown(href, unsafe_allow_html=True)
+
+
+if __name__ == "__main__":
+    main()

pyproject.toml

@@ -0,0 +1,27 @@
+[tool.ruff]
+line-length = 119
+
+[tool.ruff.lint]
+# Never enforce `E501` (line length violations).
+ignore = ["C901", "E501", "E741", "F402", "F823"]
+select = ["C", "E", "F", "I", "W"]
+
+# Ignore import violations in all `__init__.py` files.
+[tool.ruff.lint.per-file-ignores]
+"__init__.py" = ["E402", "F401", "F403", "F811"]
+
+[tool.ruff.lint.isort]
+lines-after-imports = 2
+
+[tool.ruff.format]
+# Like Black, use double quotes for strings.
+quote-style = "double"
+
+# Like Black, indent with spaces, rather than tabs.
+indent-style = "space"
+
+# Like Black, respect magic trailing commas.
+skip-magic-trailing-comma = false
+
+# Like Black, automatically detect the appropriate line ending.
+line-ending = "auto"

requirements.txt

@@ -0,0 +1,11 @@
+git+https://github.com/huggingface/diffusers.git@d1c575ad7ee0390c2735f50cc59a79aae666567a#egg=diffusers
+torch==2.4.0
+torchvision==0.19.0
+streamlit==1.37.0
+opencv-python
+imageio-ffmpeg==0.5.1
+openai==1.38.0
+transformers==4.43.4
+accelerate==0.33.0
+sentencepiece==0.2.0
+pillow==9.5.0

resources/CogVideoX.pdf

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:25ba30aafcd9604178c6d7adbd17f2bf1b251f3d29d1d29498e576075cb67c4e
+size 31028426

resources/WECHAT.md

@@ -0,0 +1,7 @@
+<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>
+

resources/contribute.md

@@ -0,0 +1,50 @@
+# Contribution Guide
+
+There may still be many incomplete aspects in this project.
+
+We look forward to your contributions to the repository in the following areas. If you complete the work mentioned above
+and are willing to submit a PR and share it with the community, upon review, we
+will acknowledge your contribution on the project homepage.
+
+## Model Algorithms
+
+- Support for model quantization inference (Int4, Int8, etc. quantization engineering)
+- Support for multi-card inference / model inference concurrency engineering
+- Support for non-CUDA architecture inference devices
+
+## Model Engineering / Secondary Development
+
+- Model fine-tuning examples / best prompt practices
+- Video super-resolution/frame interpolation for enhancing video generation quality.
+- Any peripheral tools for the model
+- Any minimal complete open-source projects using the CogVideoX open-source model
+
+## Code Standards
+
+Good code style is an art. We have prepared a `pyproject.toml` configuration file for the project to standardize code
+style. You can organize the code according to the following specifications:
+
+1. Install the `ruff` tool
+
+```shell
+pip install ruff
+```
+
+Then, run the `ruff` tool
+
+```shell
+ruff check tools sat inference
+```
+
+Check the code style. If there are issues, you can automatically fix them using the `ruff format` command.
+
+```shell
+ruff format tools sat inference
+```
+
+Once your code meets the standard, there should be no errors.
+
+## Naming Conventions
+1. Please use English names, do not use Pinyin or other language names. All comments should be in English.
+2. Please strictly follow the PEP8 specification and use underscores to separate words. Do not use names like a, b, c.
+

resources/contribute_zh.md

@@ -0,0 +1,45 @@
+# 贡献指南
+
+本项目可能还存在很多不完善的内容。 我们期待您在以下方面与我们共建仓库, 如果您完成了上述工作并愿意PR和分享到社区，在通过审核后，我们将在项目首页感谢您的贡献。
+
+## 模型算法
+
+- 模型量化推理支持 (Int4,Int8等量化工程)
+- 模型多卡推理支持 / 模型推理并发工程
+- 非 CUDA 架构 推理设备支持
+
+## 模型工程 / 模型二次开发
+
+- 模型微调示例 / 最佳提示词实践
+- 视频超分/插帧，用于美化视频生成效果。
+- 任何模型周边工具
+- 任何使用CogVideoX开源模型制作的最小完整开源项目
+
+## 代码规范
+
+良好的代码风格是一种艺术，我们已经为项目准备好了`pyproject.toml`配置文件，用于规范代码风格。您可以按照以下规范梳理代码:
+
+1. 安装`ruff`工具
+
+```shell
+pip install ruff
+```
+
+接着，运行`ruff`工具
+
+```shell
+ruff check tools sat inference
+```
+
+检查代码风格，如果有问题，您可以通过`ruff formate`命令自动修复。
+
+```shell
+ruff formate tools sat inference
+```
+
+如果您的代码符合规范，应该不会出现任何的错误。
+
+## 命名规范
+
+- 请使用英文命名，不要使用拼音或者其他语言命名。所有的注释均使用英文。
+- 请严格遵循 PEP8 规范，使用下划线分割单词。请勿使用 a,b,c 这样的命名。

resources/videos/2.mp4

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:6e738926f262a28b3b9af6573987905457ea82cdcadb0ec04ad9ab134324f5cc
+size 1683616

sat/README.md

@@ -0,0 +1,182 @@
+# SAT CogVideoX-2B
+
+This folder contains the inference code using [SAT](https://github.com/THUDM/SwissArmyTransformer) weights and the
+fine-tuning code for SAT weights.
+
+This code is the framework used by the team to train the model. It has few comments and requires careful study.
+
+## Inference Model
+
+1. Ensure that you have correctly installed the dependencies required by this folder.
+
+```shell
+pip install -r requirements.txt
+```
+
+2. Download the model weights
+
+First, go to the SAT mirror to download the dependencies.
+
+```shell
+mkdir CogVideoX-2b-sat
+cd CogVideoX-2b-sat
+wget https://cloud.tsinghua.edu.cn/f/fdba7608a49c463ba754/?dl=1
+mv 'index.html?dl=1' vae.zip
+unzip vae.zip
+wget https://cloud.tsinghua.edu.cn/f/556a3e1329e74f1bac45/?dl=1
+mv 'index.html?dl=1' transformer.zip
+unzip transformer.zip
+```
+
+Then unzip, the model structure should look like this:
+
+```
+.
+├── transformer
+│   ├── 1000
+│   │   └── mp_rank_00_model_states.pt
+│   └── latest
+└── vae
+    └── 3d-vae.pt
+```
+
+Next, clone the T5 model, which is not used for training and fine-tuning, but must be used.
+
+```shell
+git lfs install 
+git clone https://huggingface.co/google/t5-v1_1-xxl.git
+```
+
+**We don't need the tf_model.h5** file. This file can be deleted.
+
+3. Modify the file `configs/cogvideox_2b_infer.yaml`.
+
+```yaml
+load: "{your_CogVideoX-2b-sat_path}/transformer" ## Transformer model path
+
+conditioner_config:
+  target: sgm.modules.GeneralConditioner
+  params:
+    emb_models:
+      - is_trainable: false
+        input_key: txt
+        ucg_rate: 0.1
+        target: sgm.modules.encoders.modules.FrozenT5Embedder
+        params:
+          model_dir: "google/t5-v1_1-xxl" ## T5 model path
+          max_length: 226
+
+first_stage_config:
+  target: sgm.models.autoencoder.VideoAutoencoderInferenceWrapper
+  params:
+    cp_size: 1
+    ckpt_path: "{your_CogVideoX-2b-sat_path}/vae/3d-vae.pt" ## VAE model path
+```
+
++ If using txt to save multiple prompts, please refer to `configs/test.txt` for modification. One prompt per line. If
+  you don't know how to write prompts, you can first use [this code](../inference/convert_demo.py) to call LLM for
+  refinement.
++ If using the command line as input, modify
+
+```yaml
+input_type: cli
+```
+
+so that prompts can be entered from the command line.
+
+If you want to change the output video directory, you can modify:
+
+```yaml
+output_dir: outputs/
+```
+
+The default is saved in the `.outputs/` folder.
+
+4. Run the inference code to start inference
+
+```shell
+bash inference.sh
+```
+
+## Fine-Tuning the Model
+
+### Preparing the Dataset
+
+The dataset format should be as follows:
+
+```
+.
+├── labels
+│   ├── 1.txt
+│   ├── 2.txt
+│   ├── ...
+└── videos
+    ├── 1.mp4
+    ├── 2.mp4
+    ├── ...
+```
+
+Each txt file should have the same name as its corresponding video file and contain the labels for that video. Each
+video should have a one-to-one correspondence with a label. Typically, a video should not have multiple labels.
+
+For style fine-tuning, please prepare at least 50 videos and labels with similar styles to facilitate fitting.
+
+### Modifying the Configuration File
+
+We support both `Lora` and `full-parameter fine-tuning` methods. Please note that both fine-tuning methods only apply to the `transformer` part. The `VAE part` is not modified. `T5` is only used as an Encoder.
+
+the `configs/cogvideox_2b_sft.yaml` (for full fine-tuning) as follows.
+
+```yaml
+  # checkpoint_activations: True ## using gradient checkpointing (both checkpoint_activations in the configuration file need to be set to True)
+  model_parallel_size: 1 # Model parallel size
+  experiment_name: lora-disney  # Experiment name (do not change)
+  mode: finetune # Mode (do not change)
+  load: "{your_CogVideoX-2b-sat_path}/transformer" # Transformer model path
+  no_load_rng: True # Whether to load the random seed
+  train_iters: 1000 # Number of training iterations
+  eval_iters: 1 # Number of evaluation iterations
+  eval_interval: 100 # Evaluation interval
+  eval_batch_size: 1 # Batch size for evaluation
+  save: ckpts # Model save path
+  save_interval: 100 # Model save interval
+  log_interval: 20 # Log output interval
+  train_data: [ "your train data path" ]
+  valid_data: [ "your val data path" ] # Training and validation sets can be the same
+  split: 1,0,0 # Ratio of training, validation, and test sets
+  num_workers: 8 # Number of worker threads for data loading
+```
+
+If you wish to use Lora fine-tuning, you also need to modify:
+
+```yaml
+model:
+  scale_factor: 1.15258426
+  disable_first_stage_autocast: true
+  not_trainable_prefixes: [ 'all' ] ## Uncomment
+  log_keys:
+    - txt'
+
+  lora_config: ## Uncomment
+    target: sat.model.finetune.lora2.LoraMixin
+    params:
+      r: 256
+```
+
+### Fine-Tuning and Validation
+
+1. Run the inference code to start fine-tuning.
+
+```shell
+bash finetune.sh
+```
+
+### Converting to Huggingface Diffusers Supported Weights
+
+The SAT weight format is different from Huggingface's weight format and needs to be converted. Please run:
+
+```shell
+python ../tools/convert_weight_sat2hf.py
+```
+
+**Note**: This content has not yet been tested with LORA fine-tuning models.

sat/README_zh.md

@@ -0,0 +1,180 @@
+# SAT CogVideoX-2B
+
+本文件夹包含了使用 [SAT](https://github.com/THUDM/SwissArmyTransformer) 权重的推理代码，以及 SAT 权重的微调代码。
+
+该代码是团队训练模型时使用的框架。注释较少，需要认真研究。
+
+## 推理模型
+
+1. 确保你已经正确安装本文件夹中的要求的依赖
+
+```shell
+pip install -r requirements.txt
+```
+
+2. 下载模型权重
+
+首先，前往 SAT 镜像下载依赖。
+
+```shell
+mkdir CogVideoX-2b-sat
+cd CogVideoX-2b-sat
+wget https://cloud.tsinghua.edu.cn/f/fdba7608a49c463ba754/?dl=1
+mv 'index.html?dl=1' vae.zip
+unzip vae.zip
+wget https://cloud.tsinghua.edu.cn/f/556a3e1329e74f1bac45/?dl=1
+mv 'index.html?dl=1' transformer.zip
+unzip transformer.zip
+```
+
+然后，解压文件，模型结构应该如下
+
+```
+.
+├── transformer
+│   ├── 1000
+│   │   └── mp_rank_00_model_states.pt
+│   └── latest
+└── vae
+    └── 3d-vae.pt
+```
+
+接着，克隆 T5 模型，该模型不用做训练和微调，但是必须使用。
+
+```shell
+git lfs install 
+git clone https://huggingface.co/google/t5-v1_1-xxl.git
+```
+
+**我们不需要使用tf_model.h5**文件。该文件可以删除。
+
+3. 修改`configs/cogvideox_2b_infer.yaml`中的文件。
+
+```yaml
+load: "{your_CogVideoX-2b-sat_path}/transformer" ## Transformer 模型路径
+
+conditioner_config:
+  target: sgm.modules.GeneralConditioner
+  params:
+    emb_models:
+      - is_trainable: false
+        input_key: txt
+        ucg_rate: 0.1
+        target: sgm.modules.encoders.modules.FrozenT5Embedder
+        params:
+          model_dir: "google/t5-v1_1-xxl" ## T5 模型路径
+          max_length: 226
+
+first_stage_config:
+  target: sgm.models.autoencoder.VideoAutoencoderInferenceWrapper
+  params:
+    cp_size: 1
+    ckpt_path: "{your_CogVideoX-2b-sat_path}/vae/3d-vae.pt" ## VAE 模型路径
+
+```
+
++ 如果使用 txt 保存多个提示词，请参考`configs/test.txt`
+  进行修改。每一行一个提示词。如果您不知道如何书写提示词，可以先使用[此代码](../inference/convert_demo.py)调用 LLM进行润色。
++ 如果使用命令行作为输入，请修改
+
+```yaml
+input_type: cli
+```
+
+这样就可以从命令行输入提示词。
+
+如果你希望修改输出视频的地址，你可以修改:
+
+```yaml
+output_dir: outputs/
+```
+
+默认保存在`.outputs/`文件夹下。
+
+4. 运行推理代码,即可推理
+
+```shell
+bash inference.sh
+```
+
+## 微调模型
+
+### 准备数据集
+
+数据集格式应该如下：
+
+```
+.
+├── labels
+│   ├── 1.txt
+│   ├── 2.txt
+│   ├── ...
+└── videos
+    ├── 1.mp4
+    ├── 2.mp4
+    ├── ...
+```
+
+每个 txt 与视频同名，为视频的标签。视频与标签应该一一对应。通常情况下，不使用一个视频对应多个标签。
+
+如果为风格微调，清准备至少50条风格相似的视频和标签，以利于拟合。
+
+### 修改配置文件
+
+我们支持 `Lora` 和 全参数微调两种方式。请注意，两种微调方式都仅仅对 `transformer` 部分进行微调。不改动 `VAE` 部分。`T5`仅作为
+Encoder 使用。
+部分。 请按照以下方式修改`configs/cogvideox_2b_sft.yaml`(全量微调) 中的文件。
+
+```yaml
+  # checkpoint_activations: True ## using gradient checkpointing (配置文件中的两个checkpoint_activations都需要设置为True)
+  model_parallel_size: 1 # 模型并行大小
+  experiment_name: lora-disney  # 实验名称(不要改动)
+  mode: finetune # 模式(不要改动)
+  load: "{your_CogVideoX-2b-sat_path}/transformer" ## Transformer 模型路径
+  no_load_rng: True # 是否加载随机数种子
+  train_iters: 1000 # 训练迭代次数
+  eval_iters: 1 # 验证迭代次数
+  eval_interval: 100    # 验证间隔
+  eval_batch_size: 1  # 验证集 batch size
+  save: ckpts # 模型保存路径 
+  save_interval: 100 # 模型保存间隔
+  log_interval: 20 # 日志输出间隔
+  train_data: [ "your train data path" ]
+  valid_data: [ "your val data path" ] # 训练集和验证集可以相同
+  split: 1,0,0 # 训练集，验证集，测试集比例
+  num_workers: 8 # 数据加载器的工作线程数
+```
+
+如果你希望使用 Lora 微调，你还需要修改：
+
+```yaml
+model:
+  scale_factor: 1.15258426
+  disable_first_stage_autocast: true
+  not_trainable_prefixes: [ 'all' ] ## 解除注释
+  log_keys:
+    - txt'
+
+  lora_config: ##  解除注释
+    target: sat.model.finetune.lora2.LoraMixin
+    params:
+      r: 256
+```
+
+### 微调和验证
+
+1. 运行推理代码,即可开始微调。
+
+```shell
+bash finetune.sh
+```
+
+### 转换到 Huggingface Diffusers 库支持的权重
+
+SAT 权重格式与 Huggingface 的权重格式不同，需要转换。请运行
+
+```shell
+python ../tools/convert_weight_sat2hf.py
+```
+
+**注意** 本内容暂未测试 LORA 微调模型。

sat/arguments.py

@@ -0,0 +1,281 @@
+import argparse
+import os
+import torch
+import json
+import warnings
+import omegaconf
+from omegaconf import OmegaConf
+from sat.helpers import print_rank0
+from sat import mpu
+from sat.arguments import set_random_seed
+from sat.arguments import add_training_args, add_evaluation_args, add_data_args
+import torch.distributed
+
+
+def add_model_config_args(parser):
+    """Model arguments"""
+
+    group = parser.add_argument_group("model", "model configuration")
+    group.add_argument("--base", type=str, nargs="*", help="config for input and saving")
+    group.add_argument(
+        "--model-parallel-size", type=int, default=1, help="size of the model parallel. only use if you are an expert."
+    )
+    group.add_argument("--force-pretrain", action="store_true")
+    group.add_argument("--device", type=int, default=-1)
+    group.add_argument("--debug", action="store_true")
+    group.add_argument("--log-image", type=bool, default=True)
+
+    return parser
+
+
+def add_sampling_config_args(parser):
+    """Sampling configurations"""
+
+    group = parser.add_argument_group("sampling", "Sampling Configurations")
+    group.add_argument("--output-dir", type=str, default="samples")
+    group.add_argument("--input-dir", type=str, default=None)
+    group.add_argument("--input-type", type=str, default="cli")
+    group.add_argument("--input-file", type=str, default="input.txt")
+    group.add_argument("--final-size", type=int, default=2048)
+    group.add_argument("--sdedit", action="store_true")
+    group.add_argument("--grid-num-rows", type=int, default=1)
+    group.add_argument("--force-inference", action="store_true")
+    group.add_argument("--lcm_steps", type=int, default=None)
+    group.add_argument("--sampling-num-frames", type=int, default=32)
+    group.add_argument("--sampling-fps", type=int, default=8)
+    group.add_argument("--only-save-latents", type=bool, default=False)
+    group.add_argument("--only-log-video-latents", type=bool, default=False)
+    group.add_argument("--latent-channels", type=int, default=32)
+    group.add_argument("--image2video", action="store_true")
+
+    return parser
+
+
+def get_args(args_list=None, parser=None):
+    """Parse all the args."""
+    if parser is None:
+        parser = argparse.ArgumentParser(description="sat")
+    else:
+        assert isinstance(parser, argparse.ArgumentParser)
+    parser = add_model_config_args(parser)
+    parser = add_sampling_config_args(parser)
+    parser = add_training_args(parser)
+    parser = add_evaluation_args(parser)
+    parser = add_data_args(parser)
+
+    import deepspeed
+
+    parser = deepspeed.add_config_arguments(parser)
+
+    args = parser.parse_args(args_list)
+    args = process_config_to_args(args)
+
+    if not args.train_data:
+        print_rank0("No training data specified", level="WARNING")
+
+    assert (args.train_iters is None) or (args.epochs is None), "only one of train_iters and epochs should be set."
+    if args.train_iters is None and args.epochs is None:
+        args.train_iters = 10000  # default 10k iters
+        print_rank0("No train_iters (recommended) or epochs specified, use default 10k iters.", level="WARNING")
+
+    args.cuda = torch.cuda.is_available()
+
+    args.rank = int(os.getenv("RANK", "0"))
+    args.world_size = int(os.getenv("WORLD_SIZE", "1"))
+    if args.local_rank is None:
+        args.local_rank = int(os.getenv("LOCAL_RANK", "0"))  # torchrun
+
+    if args.device == -1:
+        if torch.cuda.device_count() == 0:
+            args.device = "cpu"
+        elif args.local_rank is not None:
+            args.device = args.local_rank
+        else:
+            args.device = args.rank % torch.cuda.device_count()
+
+    if args.local_rank != args.device and args.mode != "inference":
+        raise ValueError(
+            "LOCAL_RANK (default 0) and args.device inconsistent. "
+            "This can only happens in inference mode. "
+            "Please use CUDA_VISIBLE_DEVICES=x for single-GPU training. "
+        )
+
+    if args.rank == 0:
+        print_rank0("using world size: {}".format(args.world_size))
+
+    if args.train_data_weights is not None:
+        assert len(args.train_data_weights) == len(args.train_data)
+
+    if args.mode != "inference":  # training with deepspeed
+        args.deepspeed = True
+        if args.deepspeed_config is None:  # not specified
+            deepspeed_config_path = os.path.join(
+                os.path.dirname(__file__), "training", f"deepspeed_zero{args.zero_stage}.json"
+            )
+            with open(deepspeed_config_path) as file:
+                args.deepspeed_config = json.load(file)
+            override_deepspeed_config = True
+        else:
+            override_deepspeed_config = False
+
+    assert not (args.fp16 and args.bf16), "cannot specify both fp16 and bf16."
+
+    if args.zero_stage > 0 and not args.fp16 and not args.bf16:
+        print_rank0("Automatically set fp16=True to use ZeRO.")
+        args.fp16 = True
+        args.bf16 = False
+
+    if args.deepspeed:
+        if args.checkpoint_activations:
+            args.deepspeed_activation_checkpointing = True
+        else:
+            args.deepspeed_activation_checkpointing = False
+        if args.deepspeed_config is not None:
+            deepspeed_config = args.deepspeed_config
+
+        if override_deepspeed_config:  # not specify deepspeed_config, use args
+            if args.fp16:
+                deepspeed_config["fp16"]["enabled"] = True
+            elif args.bf16:
+                deepspeed_config["bf16"]["enabled"] = True
+                deepspeed_config["fp16"]["enabled"] = False
+            else:
+                deepspeed_config["fp16"]["enabled"] = False
+            deepspeed_config["train_micro_batch_size_per_gpu"] = args.batch_size
+            deepspeed_config["gradient_accumulation_steps"] = args.gradient_accumulation_steps
+            optimizer_params_config = deepspeed_config["optimizer"]["params"]
+            optimizer_params_config["lr"] = args.lr
+            optimizer_params_config["weight_decay"] = args.weight_decay
+        else:  # override args with values in deepspeed_config
+            if args.rank == 0:
+                print_rank0("Will override arguments with manually specified deepspeed_config!")
+            if "fp16" in deepspeed_config and deepspeed_config["fp16"]["enabled"]:
+                args.fp16 = True
+            else:
+                args.fp16 = False
+            if "bf16" in deepspeed_config and deepspeed_config["bf16"]["enabled"]:
+                args.bf16 = True
+            else:
+                args.bf16 = False
+            if "train_micro_batch_size_per_gpu" in deepspeed_config:
+                args.batch_size = deepspeed_config["train_micro_batch_size_per_gpu"]
+            if "gradient_accumulation_steps" in deepspeed_config:
+                args.gradient_accumulation_steps = deepspeed_config["gradient_accumulation_steps"]
+            else:
+                args.gradient_accumulation_steps = None
+            if "optimizer" in deepspeed_config:
+                optimizer_params_config = deepspeed_config["optimizer"].get("params", {})
+                args.lr = optimizer_params_config.get("lr", args.lr)
+                args.weight_decay = optimizer_params_config.get("weight_decay", args.weight_decay)
+        args.deepspeed_config = deepspeed_config
+
+    # initialize distributed and random seed because it always seems to be necessary.
+    initialize_distributed(args)
+    args.seed = args.seed + mpu.get_data_parallel_rank()
+    set_random_seed(args.seed)
+    return args
+
+
+def initialize_distributed(args):
+    """Initialize torch.distributed."""
+    if torch.distributed.is_initialized():
+        if mpu.model_parallel_is_initialized():
+            if args.model_parallel_size != mpu.get_model_parallel_world_size():
+                raise ValueError(
+                    "model_parallel_size is inconsistent with prior configuration."
+                    "We currently do not support changing model_parallel_size."
+                )
+            return False
+        else:
+            if args.model_parallel_size > 1:
+                warnings.warn(
+                    "model_parallel_size > 1 but torch.distributed is not initialized via SAT."
+                    "Please carefully make sure the correctness on your own."
+                )
+            mpu.initialize_model_parallel(args.model_parallel_size)
+        return True
+    # the automatic assignment of devices has been moved to arguments.py
+    if args.device == "cpu":
+        pass
+    else:
+        torch.cuda.set_device(args.device)
+    # Call the init process
+    init_method = "tcp://"
+    args.master_ip = os.getenv("MASTER_ADDR", "localhost")
+
+    if args.world_size == 1:
+        from sat.helpers import get_free_port
+
+        default_master_port = str(get_free_port())
+    else:
+        default_master_port = "6000"
+    args.master_port = os.getenv("MASTER_PORT", default_master_port)
+    init_method += args.master_ip + ":" + args.master_port
+    torch.distributed.init_process_group(
+        backend=args.distributed_backend, world_size=args.world_size, rank=args.rank, init_method=init_method
+    )
+
+    # Set the model-parallel / data-parallel communicators.
+    mpu.initialize_model_parallel(args.model_parallel_size)
+
+    # Set vae context parallel group equal to model parallel group
+    from sgm.util import set_context_parallel_group, initialize_context_parallel
+
+    if args.model_parallel_size <= 2:
+        set_context_parallel_group(args.model_parallel_size, mpu.get_model_parallel_group())
+    else:
+        initialize_context_parallel(2)
+    # mpu.initialize_model_parallel(1)
+    # Optional DeepSpeed Activation Checkpointing Features
+    if args.deepspeed:
+        import deepspeed
+
+        deepspeed.init_distributed(
+            dist_backend=args.distributed_backend, world_size=args.world_size, rank=args.rank, init_method=init_method
+        )
+        # # It seems that it has no negative influence to configure it even without using checkpointing.
+        # deepspeed.checkpointing.configure(mpu, deepspeed_config=args.deepspeed_config, num_checkpoints=args.num_layers)
+    else:
+        # in model-only mode, we don't want to init deepspeed, but we still need to init the rng tracker for model_parallel, just because we save the seed by default when dropout.
+        try:
+            import deepspeed
+            from deepspeed.runtime.activation_checkpointing.checkpointing import (
+                _CUDA_RNG_STATE_TRACKER,
+                _MODEL_PARALLEL_RNG_TRACKER_NAME,
+            )
+
+            _CUDA_RNG_STATE_TRACKER.add(_MODEL_PARALLEL_RNG_TRACKER_NAME, 1)  # default seed 1
+        except Exception as e:
+            from sat.helpers import print_rank0
+
+            print_rank0(str(e), level="DEBUG")
+
+    return True
+
+
+def process_config_to_args(args):
+    """Fetch args from only --base"""
+
+    configs = [OmegaConf.load(cfg) for cfg in args.base]
+    config = OmegaConf.merge(*configs)
+
+    args_config = config.pop("args", OmegaConf.create())
+    for key in args_config:
+        if isinstance(args_config[key], omegaconf.DictConfig) or isinstance(args_config[key], omegaconf.ListConfig):
+            arg = OmegaConf.to_object(args_config[key])
+        else:
+            arg = args_config[key]
+        if hasattr(args, key):
+            setattr(args, key, arg)
+
+    if "model" in config:
+        model_config = config.pop("model", OmegaConf.create())
+        args.model_config = model_config
+    if "deepspeed" in config:
+        deepspeed_config = config.pop("deepspeed", OmegaConf.create())
+        args.deepspeed_config = OmegaConf.to_object(deepspeed_config)
+    if "data" in config:
+        data_config = config.pop("data", OmegaConf.create())
+        args.data_config = data_config
+
+    return args

sat/configs/cogvideox_2b_infer.yaml

@@ -0,0 +1,166 @@
+args:
+  latent_channels: 16
+  mode: inference
+  load: "CogVideoX-2b-sat/transformer"
+  batch_size: 1
+  input_type: txt
+  input_file: test.txt
+  sampling_num_frames: 13  # Must be 13, 11 or 9
+  sampling_fps: 8
+  fp16: True
+  output_dir: outputs/
+  force_inference: True
+
+model:
+  scale_factor: 1.15258426
+  disable_first_stage_autocast: true
+  log_keys:
+    - txt
+
+  denoiser_config:
+    target: sgm.modules.diffusionmodules.denoiser.DiscreteDenoiser
+    params:
+      num_idx: 1000
+      quantize_c_noise: False
+
+      weighting_config:
+        target: sgm.modules.diffusionmodules.denoiser_weighting.EpsWeighting
+      scaling_config:
+        target: sgm.modules.diffusionmodules.denoiser_scaling.VideoScaling
+      discretization_config:
+        target: sgm.modules.diffusionmodules.discretizer.ZeroSNRDDPMDiscretization
+        params:
+          shift_scale: 3.0
+
+  network_config:
+    target: dit_video_concat.DiffusionTransformer
+    params:
+      time_embed_dim: 512
+      elementwise_affine: True
+      num_frames: 49
+      time_compressed_rate: 4
+      latent_width: 90
+      latent_height: 60
+      num_layers: 30
+      patch_size: 2
+      in_channels: 16
+      out_channels: 16
+      hidden_size: 1920
+      adm_in_channels: 256
+      num_attention_heads: 30
+
+      transformer_args:
+        vocab_size: 1
+        max_sequence_length: 64
+        layernorm_order: pre
+        skip_init: false
+        model_parallel_size: 1
+        is_decoder: false
+
+      modules:
+        pos_embed_config:
+          target: dit_video_concat.Basic3DPositionEmbeddingMixin
+          params:
+            text_length: 226
+            height_interpolation: 1.875
+            width_interpolation: 1.875
+
+        patch_embed_config:
+          target: dit_video_concat.ImagePatchEmbeddingMixin
+          params:
+            text_hidden_size: 4096
+
+        adaln_layer_config:
+          target: dit_video_concat.AdaLNMixin
+          params:
+            qk_ln: True
+
+        final_layer_config:
+          target: dit_video_concat.FinalLayerMixin
+
+  conditioner_config:
+    target: sgm.modules.GeneralConditioner
+    params:
+      emb_models:
+        - is_trainable: false
+          input_key: txt
+          ucg_rate: 0.1
+          target: sgm.modules.encoders.modules.FrozenT5Embedder
+          params:
+            model_dir: "google/t5-v1_1-xxl"
+            max_length: 226
+
+  first_stage_config:
+    target: vae_modules.autoencoder.VideoAutoencoderInferenceWrapper
+    params:
+      cp_size: 1
+      ckpt_path: "CogVideoX-2b-sat/vae/3d-vae.pt"
+      ignore_keys: [ 'loss' ]
+
+      loss_config:
+        target: torch.nn.Identity
+
+      regularizer_config:
+        target: vae_modules.regularizers.DiagonalGaussianRegularizer
+
+      encoder_config:
+        target: vae_modules.cp_enc_dec.ContextParallelEncoder3D
+        params:
+          double_z: true
+          z_channels: 16
+          resolution: 256
+          in_channels: 3
+          out_ch: 3
+          ch: 128
+          ch_mult: [ 1, 2, 2, 4 ]
+          attn_resolutions: [ ]
+          num_res_blocks: 3
+          dropout: 0.0
+          gather_norm: True
+
+      decoder_config:
+        target: vae_modules.cp_enc_dec.ContextParallelDecoder3D
+        params:
+          double_z: True
+          z_channels: 16
+          resolution: 256
+          in_channels: 3
+          out_ch: 3
+          ch: 128
+          ch_mult: [ 1, 2, 2, 4 ]
+          attn_resolutions: [ ]
+          num_res_blocks: 3
+          dropout: 0.0
+          gather_norm: false
+
+  loss_fn_config:
+    target: sgm.modules.diffusionmodules.loss.VideoDiffusionLoss
+    params:
+      offset_noise_level: 0
+      sigma_sampler_config:
+        target: sgm.modules.diffusionmodules.sigma_sampling.DiscreteSampling
+        params:
+          uniform_sampling: True
+          num_idx: 1000
+          discretization_config:
+            target: sgm.modules.diffusionmodules.discretizer.ZeroSNRDDPMDiscretization
+            params:
+              shift_scale: 3.0
+
+  sampler_config:
+    target: sgm.modules.diffusionmodules.sampling.VPSDEDPMPP2MSampler
+    params:
+      num_steps: 50
+      verbose: True
+
+      discretization_config:
+        target: sgm.modules.diffusionmodules.discretizer.ZeroSNRDDPMDiscretization
+        params:
+          shift_scale: 3.0
+
+      guider_config:
+        target: sgm.modules.diffusionmodules.guiders.DynamicCFG
+        params:
+          scale: 6
+          exp: 5
+          num_steps: 50

sat/configs/cogvideox_2b_sft.yaml

@@ -0,0 +1,225 @@
+args:
+  checkpoint_activations: True ## using gradient checkpointing
+  model_parallel_size: 1
+  experiment_name: lora-disney
+  mode: finetune
+  load: "CogVideoX-2b-sat/transformer"
+  no_load_rng: True
+  train_iters: 1000
+  eval_iters: 1
+  eval_interval: 100
+  eval_batch_size: 1
+  save: ckpts
+  save_interval: 100
+  log_interval: 20
+  train_data: ["disney"]
+  valid_data: ["disney"]
+  split: 1,0,0
+  num_workers: 8
+  force_train: True
+  only_log_video_latents: True
+
+data:
+    target: data_video.SFTDataset
+    params:
+      video_size: [480, 720]
+      fps: 8
+      max_num_frames: 49
+      skip_frms_num: 3.
+
+deepspeed:
+  train_micro_batch_size_per_gpu: 1
+  gradient_accumulation_steps: 1
+  steps_per_print: 50
+  gradient_clipping: 0.1
+  zero_optimization:
+    stage: 2
+    cpu_offload: false
+    contiguous_gradients: false
+    overlap_comm: true
+    reduce_scatter: true
+    reduce_bucket_size: 1000000000
+    allgather_bucket_size: 1000000000
+    load_from_fp32_weights: false
+  zero_allow_untested_optimizer: true
+  bf16:
+    enabled: False
+  fp16:
+    enabled: True
+    loss_scale: 0
+    loss_scale_window: 400
+    hysteresis: 2
+    min_loss_scale: 1
+  optimizer:
+    type: sat.ops.FusedEmaAdam
+    params:
+      lr: 0.0002
+      betas: [0.9, 0.95]
+      eps: 1e-8
+      weight_decay: 1e-4
+  activation_checkpointing:
+    partition_activations: false
+    contiguous_memory_optimization: false
+  wall_clock_breakdown: false
+
+
+model:
+  scale_factor: 1.15258426
+  disable_first_stage_autocast: true
+  not_trainable_prefixes: ['all'] ## Using Lora
+  log_keys:
+    - txt
+  
+  denoiser_config:
+    target: sgm.modules.diffusionmodules.denoiser.DiscreteDenoiser
+    params:
+      num_idx: 1000
+      quantize_c_noise: False
+
+      weighting_config:
+        target: sgm.modules.diffusionmodules.denoiser_weighting.EpsWeighting
+      scaling_config:
+        target: sgm.modules.diffusionmodules.denoiser_scaling.VideoScaling
+      discretization_config:
+        target: sgm.modules.diffusionmodules.discretizer.ZeroSNRDDPMDiscretization
+        params:
+          shift_scale: 3.0
+
+  network_config:
+    target: dit_video_concat.DiffusionTransformer
+    params:
+      time_embed_dim: 512
+      elementwise_affine: True
+      num_frames: 49
+      time_compressed_rate: 4
+      latent_width: 90
+      latent_height: 60
+      num_layers: 30
+      patch_size: 2
+      in_channels: 16
+      out_channels: 16
+      hidden_size: 1920
+      adm_in_channels: 256
+      num_attention_heads: 30
+
+      transformer_args:
+        checkpoint_activations: True ## using gradient checkpointing
+        vocab_size: 1
+        max_sequence_length: 64
+        layernorm_order: pre
+        skip_init: false
+        model_parallel_size: 1
+        is_decoder: false
+
+      modules:
+        pos_embed_config:
+          target: dit_video_concat.Basic3DPositionEmbeddingMixin
+          params:
+            text_length: 226
+            height_interpolation: 1.875
+            width_interpolation: 1.875
+
+        lora_config: ## Using Lora
+          target: sat.model.finetune.lora2.LoraMixin
+          params:
+            r: 128
+
+        patch_embed_config:
+          target: dit_video_concat.ImagePatchEmbeddingMixin
+          params:
+            text_hidden_size: 4096
+
+        adaln_layer_config:
+          target: dit_video_concat.AdaLNMixin
+          params:
+            qk_ln: True
+
+        final_layer_config:
+          target: dit_video_concat.FinalLayerMixin
+
+  conditioner_config:
+    target: sgm.modules.GeneralConditioner
+    params:
+      emb_models:
+        - is_trainable: false
+          input_key: txt
+          ucg_rate: 0.1
+          target: sgm.modules.encoders.modules.FrozenT5Embedder
+          params:
+            model_dir: "google/t5-v1_1-xxl"
+            max_length: 226
+
+  first_stage_config:
+    target: vae_modules.autoencoder.VideoAutoencoderInferenceWrapper
+    params:
+      cp_size: 1
+      ckpt_path: "CogVideoX-2b-sat/vae/3d-vae.pt"
+      ignore_keys: [ 'loss' ]
+
+      loss_config:
+        target: torch.nn.Identity
+
+      regularizer_config:
+        target: vae_modules.regularizers.DiagonalGaussianRegularizer
+
+      encoder_config:
+        target: vae_modules.cp_enc_dec.ContextParallelEncoder3D
+        params:
+          double_z: true
+          z_channels: 16
+          resolution: 256
+          in_channels: 3
+          out_ch: 3
+          ch: 128
+          ch_mult: [ 1, 2, 2, 4 ]
+          attn_resolutions: [ ]
+          num_res_blocks: 3
+          dropout: 0.0
+          gather_norm: True
+
+      decoder_config:
+        target: vae_modules.cp_enc_dec.ContextParallelDecoder3D
+        params:
+          double_z: True
+          z_channels: 16
+          resolution: 256
+          in_channels: 3
+          out_ch: 3
+          ch: 128
+          ch_mult: [ 1, 2, 2, 4 ]
+          attn_resolutions: [ ]
+          num_res_blocks: 3
+          dropout: 0.0
+          gather_norm: false
+
+  loss_fn_config:
+    target: sgm.modules.diffusionmodules.loss.VideoDiffusionLoss
+    params:
+      offset_noise_level: 0
+      sigma_sampler_config:
+        target: sgm.modules.diffusionmodules.sigma_sampling.DiscreteSampling
+        params:
+          uniform_sampling: True
+          num_idx: 1000
+          discretization_config:
+            target: sgm.modules.diffusionmodules.discretizer.ZeroSNRDDPMDiscretization
+            params:
+              shift_scale: 3.0
+
+  sampler_config:
+    target: sgm.modules.diffusionmodules.sampling.VPSDEDPMPP2MSampler
+    params:
+      num_steps: 50
+      verbose: True
+
+      discretization_config:
+        target: sgm.modules.diffusionmodules.discretizer.ZeroSNRDDPMDiscretization
+        params:
+          shift_scale: 3.0
+
+      guider_config:
+        target: sgm.modules.diffusionmodules.guiders.DynamicCFG
+        params:
+          scale: 6
+          exp: 5
+          num_steps: 50

sat/configs/test.txt

@@ -0,0 +1,3 @@
+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.
+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.
+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.

sat/data_video.py

@@ -0,0 +1,451 @@
+import io
+import os
+import sys
+from functools import partial
+import math
+import torchvision.transforms as TT
+from sgm.webds import MetaDistributedWebDataset
+import random
+from fractions import Fraction
+from typing import Union, Optional, Dict, Any, Tuple
+from torchvision.io.video import av
+import numpy as np
+import torch
+from torchvision.io import _video_opt
+from torchvision.io.video import _check_av_available, _read_from_stream, _align_audio_frames
+from torchvision.transforms.functional import center_crop, resize
+from torchvision.transforms import InterpolationMode
+import decord
+from decord import VideoReader
+from torch.utils.data import Dataset
+
+
+def read_video(
+    filename: str,
+    start_pts: Union[float, Fraction] = 0,
+    end_pts: Optional[Union[float, Fraction]] = None,
+    pts_unit: str = "pts",
+    output_format: str = "THWC",
+) -> Tuple[torch.Tensor, torch.Tensor, Dict[str, Any]]:
+    """
+    Reads a video from a file, returning both the video frames and the audio frames
+
+    Args:
+        filename (str): path to the video file
+        start_pts (int if pts_unit = 'pts', float / Fraction if pts_unit = 'sec', optional):
+            The start presentation time of the video
+        end_pts (int if pts_unit = 'pts', float / Fraction if pts_unit = 'sec', optional):
+            The end presentation time
+        pts_unit (str, optional): unit in which start_pts and end_pts values will be interpreted,
+            either 'pts' or 'sec'. Defaults to 'pts'.
+        output_format (str, optional): The format of the output video tensors. Can be either "THWC" (default) or "TCHW".
+
+    Returns:
+        vframes (Tensor[T, H, W, C] or Tensor[T, C, H, W]): the `T` video frames
+        aframes (Tensor[K, L]): the audio frames, where `K` is the number of channels and `L` is the number of points
+        info (Dict): metadata for the video and audio. Can contain the fields video_fps (float) and audio_fps (int)
+    """
+
+    output_format = output_format.upper()
+    if output_format not in ("THWC", "TCHW"):
+        raise ValueError(f"output_format should be either 'THWC' or 'TCHW', got {output_format}.")
+
+    _check_av_available()
+
+    if end_pts is None:
+        end_pts = float("inf")
+
+    if end_pts < start_pts:
+        raise ValueError(f"end_pts should be larger than start_pts, got start_pts={start_pts} and end_pts={end_pts}")
+
+    info = {}
+    audio_frames = []
+    audio_timebase = _video_opt.default_timebase
+
+    with av.open(filename, metadata_errors="ignore") as container:
+        if container.streams.audio:
+            audio_timebase = container.streams.audio[0].time_base
+        if container.streams.video:
+            video_frames = _read_from_stream(
+                container,
+                start_pts,
+                end_pts,
+                pts_unit,
+                container.streams.video[0],
+                {"video": 0},
+            )
+            video_fps = container.streams.video[0].average_rate
+            # guard against potentially corrupted files
+            if video_fps is not None:
+                info["video_fps"] = float(video_fps)
+
+        if container.streams.audio:
+            audio_frames = _read_from_stream(
+                container,
+                start_pts,
+                end_pts,
+                pts_unit,
+                container.streams.audio[0],
+                {"audio": 0},
+            )
+            info["audio_fps"] = container.streams.audio[0].rate
+
+    aframes_list = [frame.to_ndarray() for frame in audio_frames]
+
+    vframes = torch.empty((0, 1, 1, 3), dtype=torch.uint8)
+
+    if aframes_list:
+        aframes = np.concatenate(aframes_list, 1)
+        aframes = torch.as_tensor(aframes)
+        if pts_unit == "sec":
+            start_pts = int(math.floor(start_pts * (1 / audio_timebase)))
+            if end_pts != float("inf"):
+                end_pts = int(math.ceil(end_pts * (1 / audio_timebase)))
+        aframes = _align_audio_frames(aframes, audio_frames, start_pts, end_pts)
+    else:
+        aframes = torch.empty((1, 0), dtype=torch.float32)
+
+    if output_format == "TCHW":
+        # [T,H,W,C] --> [T,C,H,W]
+        vframes = vframes.permute(0, 3, 1, 2)
+
+    return vframes, aframes, info
+
+
+def resize_for_rectangle_crop(arr, image_size, reshape_mode="random"):
+    if arr.shape[3] / arr.shape[2] > image_size[1] / image_size[0]:
+        arr = resize(
+            arr,
+            size=[image_size[0], int(arr.shape[3] * image_size[0] / arr.shape[2])],
+            interpolation=InterpolationMode.BICUBIC,
+        )
+    else:
+        arr = resize(
+            arr,
+            size=[int(arr.shape[2] * image_size[1] / arr.shape[3]), image_size[1]],
+            interpolation=InterpolationMode.BICUBIC,
+        )
+
+    h, w = arr.shape[2], arr.shape[3]
+    arr = arr.squeeze(0)
+
+    delta_h = h - image_size[0]
+    delta_w = w - image_size[1]
+
+    if reshape_mode == "random" or reshape_mode == "none":
+        top = np.random.randint(0, delta_h + 1)
+        left = np.random.randint(0, delta_w + 1)
+    elif reshape_mode == "center":
+        top, left = delta_h // 2, delta_w // 2
+    else:
+        raise NotImplementedError
+    arr = TT.functional.crop(arr, top=top, left=left, height=image_size[0], width=image_size[1])
+    return arr
+
+
+def pad_last_frame(tensor, num_frames):
+    # T, H, W, C
+    if tensor.shape[0] < num_frames:
+        last_frame = tensor[-int(num_frames - tensor.shape[1]) :]
+        padded_tensor = torch.cat([tensor, last_frame], dim=0)
+        return padded_tensor
+    else:
+        return tensor[:num_frames]
+
+
+def load_video(
+    video_data,
+    sampling="uniform",
+    duration=None,
+    num_frames=4,
+    wanted_fps=None,
+    actual_fps=None,
+    skip_frms_num=0.0,
+    nb_read_frames=None,
+):
+    decord.bridge.set_bridge("torch")
+    vr = VideoReader(uri=video_data, height=-1, width=-1)
+    if nb_read_frames is not None:
+        ori_vlen = nb_read_frames
+    else:
+        ori_vlen = min(int(duration * actual_fps) - 1, len(vr))
+
+    max_seek = int(ori_vlen - skip_frms_num - num_frames / wanted_fps * actual_fps)
+    start = random.randint(skip_frms_num, max_seek + 1)
+    end = int(start + num_frames / wanted_fps * actual_fps)
+    n_frms = num_frames
+
+    if sampling == "uniform":
+        indices = np.arange(start, end, (end - start) / n_frms).astype(int)
+    else:
+        raise NotImplementedError
+
+    # get_batch -> T, H, W, C
+    temp_frms = vr.get_batch(np.arange(start, end))
+    assert temp_frms is not None
+    tensor_frms = torch.from_numpy(temp_frms) if type(temp_frms) is not torch.Tensor else temp_frms
+    tensor_frms = tensor_frms[torch.tensor((indices - start).tolist())]
+
+    return pad_last_frame(tensor_frms, num_frames)
+
+
+import threading
+
+
+def load_video_with_timeout(*args, **kwargs):
+    video_container = {}
+
+    def target_function():
+        video = load_video(*args, **kwargs)
+        video_container["video"] = video
+
+    thread = threading.Thread(target=target_function)
+    thread.start()
+    timeout = 20
+    thread.join(timeout)
+
+    if thread.is_alive():
+        print("Loading video timed out")
+        raise TimeoutError
+    return video_container.get("video", None).contiguous()
+
+
+def process_video(
+    video_path,
+    image_size=None,
+    duration=None,
+    num_frames=4,
+    wanted_fps=None,
+    actual_fps=None,
+    skip_frms_num=0.0,
+    nb_read_frames=None,
+):
+    """
+    video_path: str or io.BytesIO
+    image_size: .
+    duration: preknow the duration to speed up by seeking to sampled start. TODO by_pass if unknown.
+    num_frames: wanted num_frames.
+    wanted_fps: .
+    skip_frms_num: ignore the first and the last xx frames, avoiding transitions.
+    """
+
+    video = load_video_with_timeout(
+        video_path,
+        duration=duration,
+        num_frames=num_frames,
+        wanted_fps=wanted_fps,
+        actual_fps=actual_fps,
+        skip_frms_num=skip_frms_num,
+        nb_read_frames=nb_read_frames,
+    )
+
+    # --- copy and modify the image process ---
+    video = video.permute(0, 3, 1, 2)  # [T, C, H, W]
+
+    # resize
+    if image_size is not None:
+        video = resize_for_rectangle_crop(video, image_size, reshape_mode="center")
+
+    return video
+
+
+def process_fn_video(src, image_size, fps, num_frames, skip_frms_num=0.0, txt_key="caption"):
+    while True:
+        r = next(src)
+        if "mp4" in r:
+            video_data = r["mp4"]
+        elif "avi" in r:
+            video_data = r["avi"]
+        else:
+            print("No video data found")
+            continue
+
+        if txt_key not in r:
+            txt = ""
+        else:
+            txt = r[txt_key]
+
+        if isinstance(txt, bytes):
+            txt = txt.decode("utf-8")
+        else:
+            txt = str(txt)
+
+        duration = r.get("duration", None)
+        if duration is not None:
+            duration = float(duration)
+        else:
+            continue
+
+        actual_fps = r.get("fps", None)
+        if actual_fps is not None:
+            actual_fps = float(actual_fps)
+        else:
+            continue
+
+        required_frames = num_frames / fps * actual_fps + 2 * skip_frms_num
+        required_duration = num_frames / fps + 2 * skip_frms_num / actual_fps
+
+        if duration is not None and duration < required_duration:
+            continue
+
+        try:
+            frames = process_video(
+                io.BytesIO(video_data),
+                num_frames=num_frames,
+                wanted_fps=fps,
+                image_size=image_size,
+                duration=duration,
+                actual_fps=actual_fps,
+                skip_frms_num=skip_frms_num,
+            )
+            frames = (frames - 127.5) / 127.5
+        except Exception as e:
+            print(e)
+            continue
+
+        item = {
+            "mp4": frames,
+            "txt": txt,
+            "num_frames": num_frames,
+            "fps": fps,
+        }
+
+        yield item
+
+
+class VideoDataset(MetaDistributedWebDataset):
+    def __init__(
+        self,
+        path,
+        image_size,
+        num_frames,
+        fps,
+        skip_frms_num=0.0,
+        nshards=sys.maxsize,
+        seed=1,
+        meta_names=None,
+        shuffle_buffer=1000,
+        include_dirs=None,
+        txt_key="caption",
+        **kwargs,
+    ):
+        if seed == -1:
+            seed = random.randint(0, 1000000)
+        if meta_names is None:
+            meta_names = []
+
+        if path.startswith(";"):
+            path, include_dirs = path.split(";", 1)
+        super().__init__(
+            path,
+            partial(
+                process_fn_video, num_frames=num_frames, image_size=image_size, fps=fps, skip_frms_num=skip_frms_num
+            ),
+            seed,
+            meta_names=meta_names,
+            shuffle_buffer=shuffle_buffer,
+            nshards=nshards,
+            include_dirs=include_dirs,
+        )
+
+    @classmethod
+    def create_dataset_function(cls, path, args, **kwargs):
+        return cls(path, **kwargs)
+
+
+class SFTDataset(Dataset):
+    def __init__(self, data_dir, video_size, fps, max_num_frames, skip_frms_num=3):
+        """
+        skip_frms_num: ignore the first and the last xx frames, avoiding transitions.
+        """
+        super(SFTDataset, self).__init__()
+
+        self.videos_list = []
+        self.captions_list = []
+        self.num_frames_list = []
+        self.fps_list = []
+
+        decord.bridge.set_bridge("torch")
+        for root, dirnames, filenames in os.walk(data_dir):
+            for filename in filenames:
+                if filename.endswith(".mp4"):
+                    video_path = os.path.join(root, filename)
+                    vr = VideoReader(uri=video_path, height=-1, width=-1)
+                    actual_fps = vr.get_avg_fps()
+                    ori_vlen = len(vr)
+
+                    if ori_vlen / actual_fps * fps > max_num_frames:
+                        num_frames = max_num_frames
+                        start = int(skip_frms_num)
+                        end = int(start + num_frames / fps * actual_fps)
+                        indices = np.arange(start, end, (end - start) / num_frames).astype(int)
+                        temp_frms = vr.get_batch(np.arange(start, end))
+                        assert temp_frms is not None
+                        tensor_frms = torch.from_numpy(temp_frms) if type(temp_frms) is not torch.Tensor else temp_frms
+                        tensor_frms = tensor_frms[torch.tensor((indices - start).tolist())]
+                    else:
+                        if ori_vlen > max_num_frames:
+                            num_frames = max_num_frames
+                            start = int(skip_frms_num)
+                            end = int(ori_vlen - skip_frms_num)
+                            indices = np.arange(start, end, (end - start) / num_frames).astype(int)
+                            temp_frms = vr.get_batch(np.arange(start, end))
+                            assert temp_frms is not None
+                            tensor_frms = (
+                                torch.from_numpy(temp_frms) if type(temp_frms) is not torch.Tensor else temp_frms
+                            )
+                            tensor_frms = tensor_frms[torch.tensor((indices - start).tolist())]
+                        else:
+
+                            def nearest_smaller_4k_plus_1(n):
+                                remainder = n % 4
+                                if remainder == 0:
+                                    return n - 3
+                                else:
+                                    return n - remainder + 1
+
+                            start = int(skip_frms_num)
+                            end = int(ori_vlen - skip_frms_num)
+                            num_frames = nearest_smaller_4k_plus_1(
+                                end - start
+                            )  # 3D VAE requires the number of frames to be 4k+1
+                            end = int(start + num_frames)
+                            temp_frms = vr.get_batch(np.arange(start, end))
+                            assert temp_frms is not None
+                            tensor_frms = (
+                                torch.from_numpy(temp_frms) if type(temp_frms) is not torch.Tensor else temp_frms
+                            )
+
+                    tensor_frms = pad_last_frame(
+                        tensor_frms, num_frames
+                    )  # the len of indices may be less than num_frames, due to round error
+                    tensor_frms = tensor_frms.permute(0, 3, 1, 2)  # [T, H, W, C] -> [T, C, H, W]
+                    tensor_frms = resize_for_rectangle_crop(tensor_frms, video_size, reshape_mode="center")
+                    tensor_frms = (tensor_frms - 127.5) / 127.5
+                    self.videos_list.append(tensor_frms)
+
+                    # caption
+                    caption_path = os.path.join(root, filename.replace("videos", "labels").replace(".mp4", ".txt"))
+                    if os.path.exists(caption_path):
+                        caption = open(caption_path, "r").read().splitlines()[0]
+                    else:
+                        caption = ""
+                    self.captions_list.append(caption)
+                    self.num_frames_list.append(num_frames)
+                    self.fps_list.append(fps)
+
+    def __getitem__(self, index):
+        item = {
+            "mp4": self.videos_list[index],
+            "txt": self.captions_list[index],
+            "num_frames": self.num_frames_list[index],
+            "fps": self.fps_list[index],
+        }
+        return item
+
+    def __len__(self):
+        return len(self.fps_list)
+
+    @classmethod
+    def create_dataset_function(cls, path, args, **kwargs):
+        return cls(data_dir=path, **kwargs)

sat/diffusion_video.py

@@ -0,0 +1,318 @@
+import math
+from contextlib import contextmanager
+from typing import Any, Dict, List, Tuple, Union, Optional
+from omegaconf import ListConfig, OmegaConf
+from copy import deepcopy
+import torch.nn.functional as F
+
+from sat.helpers import print_rank0
+import torch
+from torch import nn
+
+from sgm.modules import UNCONDITIONAL_CONFIG
+from sgm.modules.autoencoding.temporal_ae import VideoDecoder
+from sgm.modules.diffusionmodules.wrappers import OPENAIUNETWRAPPER
+from sgm.util import (
+    default,
+    disabled_train,
+    get_obj_from_str,
+    instantiate_from_config,
+    log_txt_as_img,
+)
+import gc
+from sat import mpu
+import random
+
+
+class SATVideoDiffusionEngine(nn.Module):
+    def __init__(self, args, **kwargs):
+        super().__init__()
+
+        model_config = args.model_config
+        # model args preprocess
+        log_keys = model_config.get("log_keys", None)
+        input_key = model_config.get("input_key", "mp4")
+        network_config = model_config.get("network_config", None)
+        network_wrapper = model_config.get("network_wrapper", None)
+        denoiser_config = model_config.get("denoiser_config", None)
+        sampler_config = model_config.get("sampler_config", None)
+        conditioner_config = model_config.get("conditioner_config", None)
+        first_stage_config = model_config.get("first_stage_config", None)
+        loss_fn_config = model_config.get("loss_fn_config", None)
+        scale_factor = model_config.get("scale_factor", 1.0)
+        latent_input = model_config.get("latent_input", False)
+        disable_first_stage_autocast = model_config.get("disable_first_stage_autocast", False)
+        no_cond_log = model_config.get("disable_first_stage_autocast", False)
+        not_trainable_prefixes = model_config.get("not_trainable_prefixes", ["first_stage_model", "conditioner"])
+        compile_model = model_config.get("compile_model", False)
+        en_and_decode_n_samples_a_time = model_config.get("en_and_decode_n_samples_a_time", None)
+        lr_scale = model_config.get("lr_scale", None)
+        lora_train = model_config.get("lora_train", False)
+        self.use_pd = model_config.get("use_pd", False)  # progressive distillation
+
+        self.log_keys = log_keys
+        self.input_key = input_key
+        self.not_trainable_prefixes = not_trainable_prefixes
+        self.en_and_decode_n_samples_a_time = en_and_decode_n_samples_a_time
+        self.lr_scale = lr_scale
+        self.lora_train = lora_train
+        self.noised_image_input = model_config.get("noised_image_input", False)
+        self.noised_image_all_concat = model_config.get("noised_image_all_concat", False)
+        self.noised_image_dropout = model_config.get("noised_image_dropout", 0.0)
+        if args.fp16:
+            dtype = torch.float16
+            dtype_str = "fp16"
+        elif args.bf16:
+            dtype = torch.bfloat16
+            dtype_str = "bf16"
+        else:
+            dtype = torch.float32
+            dtype_str = "fp32"
+        self.dtype = dtype
+        self.dtype_str = dtype_str
+
+        network_config["params"]["dtype"] = dtype_str
+        model = instantiate_from_config(network_config)
+        self.model = get_obj_from_str(default(network_wrapper, OPENAIUNETWRAPPER))(
+            model, compile_model=compile_model, dtype=dtype
+        )
+
+        self.denoiser = instantiate_from_config(denoiser_config)
+        self.sampler = instantiate_from_config(sampler_config) if sampler_config is not None else None
+        self.conditioner = instantiate_from_config(default(conditioner_config, UNCONDITIONAL_CONFIG))
+
+        self._init_first_stage(first_stage_config)
+
+        self.loss_fn = instantiate_from_config(loss_fn_config) if loss_fn_config is not None else None
+
+        self.latent_input = latent_input
+        self.scale_factor = scale_factor
+        self.disable_first_stage_autocast = disable_first_stage_autocast
+        self.no_cond_log = no_cond_log
+        self.device = args.device
+
+    def disable_untrainable_params(self):
+        total_trainable = 0
+        for n, p in self.named_parameters():
+            if p.requires_grad == False:
+                continue
+            flag = False
+            for prefix in self.not_trainable_prefixes:
+                if n.startswith(prefix) or prefix == "all":
+                    flag = True
+                    break
+
+            lora_prefix = ["matrix_A", "matrix_B"]
+            for prefix in lora_prefix:
+                if prefix in n:
+                    flag = False
+                    break
+
+            if flag:
+                p.requires_grad_(False)
+            else:
+                total_trainable += p.numel()
+
+        print_rank0("***** Total trainable parameters: " + str(total_trainable) + " *****")
+
+    def reinit(self, parent_model=None):
+        # reload the initial params from previous trained modules
+        # you can also get access to other mixins through parent_model.get_mixin().
+        pass
+
+    def _init_first_stage(self, config):
+        model = instantiate_from_config(config).eval()
+        model.train = disabled_train
+        for param in model.parameters():
+            param.requires_grad = False
+        self.first_stage_model = model
+
+    def forward(self, x, batch):
+        loss = self.loss_fn(self.model, self.denoiser, self.conditioner, x, batch)
+        loss_mean = loss.mean()
+        loss_dict = {"loss": loss_mean}
+        return loss_mean, loss_dict
+
+    def shared_step(self, batch: Dict) -> Any:
+        x = self.get_input(batch)
+        if self.lr_scale is not None:
+            lr_x = F.interpolate(x, scale_factor=1 / self.lr_scale, mode="bilinear", align_corners=False)
+            lr_x = F.interpolate(lr_x, scale_factor=self.lr_scale, mode="bilinear", align_corners=False)
+            lr_z = self.encode_first_stage(lr_x, batch)
+            batch["lr_input"] = lr_z
+
+        x = x.permute(0, 2, 1, 3, 4).contiguous()
+        x = self.encode_first_stage(x, batch)
+        x = x.permute(0, 2, 1, 3, 4).contiguous()
+
+        gc.collect()
+        torch.cuda.empty_cache()
+        loss, loss_dict = self(x, batch)
+        return loss, loss_dict
+
+    def get_input(self, batch):
+        return batch[self.input_key].to(self.dtype)
+
+    @torch.no_grad()
+    def decode_first_stage(self, z):
+        z = 1.0 / self.scale_factor * z
+        n_samples = default(self.en_and_decode_n_samples_a_time, z.shape[0])
+        n_rounds = math.ceil(z.shape[0] / n_samples)
+        all_out = []
+        with torch.autocast("cuda", enabled=not self.disable_first_stage_autocast):
+            for n in range(n_rounds):
+                if isinstance(self.first_stage_model.decoder, VideoDecoder):
+                    kwargs = {"timesteps": len(z[n * n_samples : (n + 1) * n_samples])}
+                else:
+                    kwargs = {}
+                use_cp = False
+                out = self.first_stage_model.decode(z[n * n_samples : (n + 1) * n_samples], **kwargs)
+                all_out.append(out)
+        out = torch.cat(all_out, dim=0)
+        return out
+
+    @torch.no_grad()
+    def encode_first_stage(self, x, batch):
+        frame = x.shape[2]
+
+        if frame > 1 and self.latent_input:
+            x = x.permute(0, 2, 1, 3, 4).contiguous()
+            return x * self.scale_factor  # already encoded
+
+        use_cp = False
+
+        n_samples = default(self.en_and_decode_n_samples_a_time, x.shape[0])
+        n_rounds = math.ceil(x.shape[0] / n_samples)
+        all_out = []
+        with torch.autocast("cuda", enabled=not self.disable_first_stage_autocast):
+            for n in range(n_rounds):
+                out = self.first_stage_model.encode(x[n * n_samples : (n + 1) * n_samples])
+                all_out.append(out)
+        z = torch.cat(all_out, dim=0)
+        z = self.scale_factor * z
+        return z
+
+    @torch.no_grad()
+    def sample(
+        self,
+        cond: Dict,
+        uc: Union[Dict, None] = None,
+        batch_size: int = 16,
+        shape: Union[None, Tuple, List] = None,
+        prefix=None,
+        concat_images=None,
+        **kwargs,
+    ):
+        randn = torch.randn(batch_size, *shape).to(torch.float32).to(self.device)
+        if hasattr(self, "seeded_noise"):
+            randn = self.seeded_noise(randn)
+
+        if prefix is not None:
+            randn = torch.cat([prefix, randn[:, prefix.shape[1] :]], dim=1)
+
+        # broadcast noise
+        mp_size = mpu.get_model_parallel_world_size()
+        if mp_size > 1:
+            global_rank = torch.distributed.get_rank() // mp_size
+            src = global_rank * mp_size
+            torch.distributed.broadcast(randn, src=src, group=mpu.get_model_parallel_group())
+
+        scale = None
+        scale_emb = None
+
+        denoiser = lambda input, sigma, c, **addtional_model_inputs: self.denoiser(
+            self.model, input, sigma, c, concat_images=concat_images, **addtional_model_inputs
+        )
+
+        samples = self.sampler(denoiser, randn, cond, uc=uc, scale=scale, scale_emb=scale_emb)
+        samples = samples.to(self.dtype)
+        return samples
+
+    @torch.no_grad()
+    def log_conditionings(self, batch: Dict, n: int) -> Dict:
+        """
+        Defines heuristics to log different conditionings.
+        These can be lists of strings (text-to-image), tensors, ints, ...
+        """
+        image_h, image_w = batch[self.input_key].shape[3:]
+        log = dict()
+
+        for embedder in self.conditioner.embedders:
+            if ((self.log_keys is None) or (embedder.input_key in self.log_keys)) and not self.no_cond_log:
+                x = batch[embedder.input_key][:n]
+                if isinstance(x, torch.Tensor):
+                    if x.dim() == 1:
+                        # class-conditional, convert integer to string
+                        x = [str(x[i].item()) for i in range(x.shape[0])]
+                        xc = log_txt_as_img((image_h, image_w), x, size=image_h // 4)
+                    elif x.dim() == 2:
+                        # size and crop cond and the like
+                        x = ["x".join([str(xx) for xx in x[i].tolist()]) for i in range(x.shape[0])]
+                        xc = log_txt_as_img((image_h, image_w), x, size=image_h // 20)
+                    else:
+                        raise NotImplementedError()
+                elif isinstance(x, (List, ListConfig)):
+                    if isinstance(x[0], str):
+                        xc = log_txt_as_img((image_h, image_w), x, size=image_h // 20)
+                    else:
+                        raise NotImplementedError()
+                else:
+                    raise NotImplementedError()
+                log[embedder.input_key] = xc
+        return log
+
+    @torch.no_grad()
+    def log_video(
+        self,
+        batch: Dict,
+        N: int = 8,
+        ucg_keys: List[str] = None,
+        only_log_video_latents=False,
+        **kwargs,
+    ) -> Dict:
+        conditioner_input_keys = [e.input_key for e in self.conditioner.embedders]
+        if ucg_keys:
+            assert all(map(lambda x: x in conditioner_input_keys, ucg_keys)), (
+                "Each defined ucg key for sampling must be in the provided conditioner input keys,"
+                f"but we have {ucg_keys} vs. {conditioner_input_keys}"
+            )
+        else:
+            ucg_keys = conditioner_input_keys
+        log = dict()
+
+        x = self.get_input(batch)
+
+        c, uc = self.conditioner.get_unconditional_conditioning(
+            batch,
+            force_uc_zero_embeddings=ucg_keys if len(self.conditioner.embedders) > 0 else [],
+        )
+
+        sampling_kwargs = {}
+
+        N = min(x.shape[0], N)
+        x = x.to(self.device)[:N]
+        if not self.latent_input:
+            log["inputs"] = x.to(torch.float32)
+        x = x.permute(0, 2, 1, 3, 4).contiguous()
+        z = self.encode_first_stage(x, batch)
+        if not only_log_video_latents:
+            log["reconstructions"] = self.decode_first_stage(z).to(torch.float32)
+            log["reconstructions"] = log["reconstructions"].permute(0, 2, 1, 3, 4).contiguous()
+        z = z.permute(0, 2, 1, 3, 4).contiguous()
+
+        log.update(self.log_conditionings(batch, N))
+
+        for k in c:
+            if isinstance(c[k], torch.Tensor):
+                c[k], uc[k] = map(lambda y: y[k][:N].to(self.device), (c, uc))
+
+        samples = self.sample(c, shape=z.shape[1:], uc=uc, batch_size=N, **sampling_kwargs)  # b t c h w
+        samples = samples.permute(0, 2, 1, 3, 4).contiguous()
+        if only_log_video_latents:
+            latents = 1.0 / self.scale_factor * samples
+            log["latents"] = latents
+        else:
+            samples = self.decode_first_stage(samples).to(torch.float32)
+            samples = samples.permute(0, 2, 1, 3, 4).contiguous()
+            log["samples"] = samples
+        return log

sat/dit_video_concat.py

@@ -0,0 +1,858 @@
+from functools import partial
+from einops import rearrange, repeat
+import numpy as np
+
+import torch
+from torch import nn
+import torch.nn.functional as F
+
+from sat.model.base_model import BaseModel, non_conflict
+from sat.model.mixins import BaseMixin
+from sat.transformer_defaults import HOOKS_DEFAULT, attention_fn_default
+from sat.mpu.layers import ColumnParallelLinear
+from sgm.util import instantiate_from_config
+
+from sgm.modules.diffusionmodules.openaimodel import Timestep
+from sgm.modules.diffusionmodules.util import (
+    linear,
+    timestep_embedding,
+)
+from sat.ops.layernorm import LayerNorm, RMSNorm
+
+
+class ImagePatchEmbeddingMixin(BaseMixin):
+    def __init__(
+        self,
+        in_channels,
+        hidden_size,
+        patch_size,
+        bias=True,
+        text_hidden_size=None,
+    ):
+        super().__init__()
+        self.proj = nn.Conv2d(in_channels, hidden_size, kernel_size=patch_size, stride=patch_size, bias=bias)
+        if text_hidden_size is not None:
+            self.text_proj = nn.Linear(text_hidden_size, hidden_size)
+        else:
+            self.text_proj = None
+
+    def word_embedding_forward(self, input_ids, **kwargs):
+        # now is 3d patch
+        images = kwargs["images"]  # (b,t,c,h,w)
+        B, T = images.shape[:2]
+        emb = images.view(-1, *images.shape[2:])
+        emb = self.proj(emb)  # ((b t),d,h/2,w/2)
+        emb = emb.view(B, T, *emb.shape[1:])
+        emb = emb.flatten(3).transpose(2, 3)  # (b,t,n,d)
+        emb = rearrange(emb, "b t n d -> b (t n) d")
+
+        if self.text_proj is not None:
+            text_emb = self.text_proj(kwargs["encoder_outputs"])
+            emb = torch.cat((text_emb, emb), dim=1)  # (b,n_t+t*n_i,d)
+
+        emb = emb.contiguous()
+        return emb  # (b,n_t+t*n_i,d)
+
+    def reinit(self, parent_model=None):
+        w = self.proj.weight.data
+        nn.init.xavier_uniform_(w.view([w.shape[0], -1]))
+        nn.init.constant_(self.proj.bias, 0)
+        del self.transformer.word_embeddings
+
+
+def get_3d_sincos_pos_embed(
+    embed_dim,
+    grid_height,
+    grid_width,
+    t_size,
+    cls_token=False,
+    height_interpolation=1.0,
+    width_interpolation=1.0,
+    time_interpolation=1.0,
+):
+    """
+    grid_size: int of the grid height and width
+    t_size: int of the temporal size
+    return:
+    pos_embed: [t_size*grid_size*grid_size, embed_dim] or [1+t_size*grid_size*grid_size, embed_dim] (w/ or w/o cls_token)
+    """
+    assert embed_dim % 4 == 0
+    embed_dim_spatial = embed_dim // 4 * 3
+    embed_dim_temporal = embed_dim // 4
+
+    # spatial
+    grid_h = np.arange(grid_height, dtype=np.float32) / height_interpolation
+    grid_w = np.arange(grid_width, dtype=np.float32) / width_interpolation
+    grid = np.meshgrid(grid_w, grid_h)  # here w goes first
+    grid = np.stack(grid, axis=0)
+
+    grid = grid.reshape([2, 1, grid_height, grid_width])
+    pos_embed_spatial = get_2d_sincos_pos_embed_from_grid(embed_dim_spatial, grid)
+
+    # temporal
+    grid_t = np.arange(t_size, dtype=np.float32) / time_interpolation
+    pos_embed_temporal = get_1d_sincos_pos_embed_from_grid(embed_dim_temporal, grid_t)
+
+    # concate: [T, H, W] order
+    pos_embed_temporal = pos_embed_temporal[:, np.newaxis, :]
+    pos_embed_temporal = np.repeat(pos_embed_temporal, grid_height * grid_width, axis=1)  # [T, H*W, D // 4]
+    pos_embed_spatial = pos_embed_spatial[np.newaxis, :, :]
+    pos_embed_spatial = np.repeat(pos_embed_spatial, t_size, axis=0)  # [T, H*W, D // 4 * 3]
+
+    pos_embed = np.concatenate([pos_embed_temporal, pos_embed_spatial], axis=-1)
+    # pos_embed = pos_embed.reshape([-1, embed_dim])  # [T*H*W, D]
+
+    return pos_embed  # [T, H*W, D]
+
+
+def get_2d_sincos_pos_embed(embed_dim, grid_height, grid_width, cls_token=False, extra_tokens=0):
+    """
+    grid_size: int of the grid height and width
+    return:
+    pos_embed: [grid_size*grid_size, embed_dim] or [1+grid_size*grid_size, embed_dim] (w/ or w/o cls_token)
+    """
+    grid_h = np.arange(grid_height, dtype=np.float32)
+    grid_w = np.arange(grid_width, dtype=np.float32)
+    grid = np.meshgrid(grid_w, grid_h)  # here w goes first
+    grid = np.stack(grid, axis=0)
+
+    grid = grid.reshape([2, 1, grid_height, grid_width])
+    pos_embed = get_2d_sincos_pos_embed_from_grid(embed_dim, grid)
+    if cls_token and extra_tokens > 0:
+        pos_embed = np.concatenate([np.zeros([extra_tokens, embed_dim]), pos_embed], axis=0)
+    return pos_embed
+
+
+def get_2d_sincos_pos_embed_from_grid(embed_dim, grid):
+    assert embed_dim % 2 == 0
+
+    # use half of dimensions to encode grid_h
+    emb_h = get_1d_sincos_pos_embed_from_grid(embed_dim // 2, grid[0])  # (H*W, D/2)
+    emb_w = get_1d_sincos_pos_embed_from_grid(embed_dim // 2, grid[1])  # (H*W, D/2)
+
+    emb = np.concatenate([emb_h, emb_w], axis=1)  # (H*W, D)
+    return emb
+
+
+def get_1d_sincos_pos_embed_from_grid(embed_dim, pos):
+    """
+    embed_dim: output dimension for each position
+    pos: a list of positions to be encoded: size (M,)
+    out: (M, D)
+    """
+    assert embed_dim % 2 == 0
+    omega = np.arange(embed_dim // 2, dtype=np.float64)
+    omega /= embed_dim / 2.0
+    omega = 1.0 / 10000**omega  # (D/2,)
+
+    pos = pos.reshape(-1)  # (M,)
+    out = np.einsum("m,d->md", pos, omega)  # (M, D/2), outer product
+
+    emb_sin = np.sin(out)  # (M, D/2)
+    emb_cos = np.cos(out)  # (M, D/2)
+
+    emb = np.concatenate([emb_sin, emb_cos], axis=1)  # (M, D)
+    return emb
+
+
+class Basic3DPositionEmbeddingMixin(BaseMixin):
+    def __init__(
+        self,
+        height,
+        width,
+        compressed_num_frames,
+        hidden_size,
+        text_length=0,
+        height_interpolation=1.0,
+        width_interpolation=1.0,
+        time_interpolation=1.0,
+    ):
+        super().__init__()
+        self.height = height
+        self.width = width
+        self.text_length = text_length
+        self.compressed_num_frames = compressed_num_frames
+        self.spatial_length = height * width
+        self.num_patches = height * width * compressed_num_frames
+        self.pos_embedding = nn.Parameter(
+            torch.zeros(1, int(text_length + self.num_patches), int(hidden_size)), requires_grad=False
+        )
+        self.height_interpolation = height_interpolation
+        self.width_interpolation = width_interpolation
+        self.time_interpolation = time_interpolation
+
+    def position_embedding_forward(self, position_ids, **kwargs):
+        if kwargs["images"].shape[1] == 1:
+            return self.pos_embedding[:, : self.text_length + self.spatial_length]
+
+        return self.pos_embedding[:, : self.text_length + kwargs["seq_length"]]
+
+    def reinit(self, parent_model=None):
+        del self.transformer.position_embeddings
+        pos_embed = get_3d_sincos_pos_embed(
+            self.pos_embedding.shape[-1],
+            self.height,
+            self.width,
+            self.compressed_num_frames,
+            height_interpolation=self.height_interpolation,
+            width_interpolation=self.width_interpolation,
+            time_interpolation=self.time_interpolation,
+        )
+        pos_embed = torch.from_numpy(pos_embed).float()
+        pos_embed = rearrange(pos_embed, "t n d -> (t n) d")
+        self.pos_embedding.data[:, -self.num_patches :].copy_(pos_embed)
+
+
+def broadcat(tensors, dim=-1):
+    num_tensors = len(tensors)
+    shape_lens = set(list(map(lambda t: len(t.shape), tensors)))
+    assert len(shape_lens) == 1, "tensors must all have the same number of dimensions"
+    shape_len = list(shape_lens)[0]
+    dim = (dim + shape_len) if dim < 0 else dim
+    dims = list(zip(*map(lambda t: list(t.shape), tensors)))
+    expandable_dims = [(i, val) for i, val in enumerate(dims) if i != dim]
+    assert all(
+        [*map(lambda t: len(set(t[1])) <= 2, expandable_dims)]
+    ), "invalid dimensions for broadcastable concatentation"
+    max_dims = list(map(lambda t: (t[0], max(t[1])), expandable_dims))
+    expanded_dims = list(map(lambda t: (t[0], (t[1],) * num_tensors), max_dims))
+    expanded_dims.insert(dim, (dim, dims[dim]))
+    expandable_shapes = list(zip(*map(lambda t: t[1], expanded_dims)))
+    tensors = list(map(lambda t: t[0].expand(*t[1]), zip(tensors, expandable_shapes)))
+    return torch.cat(tensors, dim=dim)
+
+
+def rotate_half(x):
+    x = rearrange(x, "... (d r) -> ... d r", r=2)
+    x1, x2 = x.unbind(dim=-1)
+    x = torch.stack((-x2, x1), dim=-1)
+    return rearrange(x, "... d r -> ... (d r)")
+
+
+class Rotary3DPositionEmbeddingMixin(BaseMixin):
+    def __init__(
+        self,
+        height,
+        width,
+        compressed_num_frames,
+        hidden_size,
+        hidden_size_head,
+        text_length,
+        theta=10000,
+        rot_v=False,
+        pnp=False,
+        learnable_pos_embed=False,
+    ):
+        super().__init__()
+        self.rot_v = rot_v
+
+        dim_t = hidden_size_head // 4
+        dim_h = hidden_size_head // 8 * 3
+        dim_w = hidden_size_head // 8 * 3
+
+        # 'lang':
+        freqs_t = 1.0 / (theta ** (torch.arange(0, dim_t, 2)[: (dim_t // 2)].float() / dim_t))
+        freqs_h = 1.0 / (theta ** (torch.arange(0, dim_h, 2)[: (dim_h // 2)].float() / dim_h))
+        freqs_w = 1.0 / (theta ** (torch.arange(0, dim_w, 2)[: (dim_w // 2)].float() / dim_w))
+
+        grid_t = torch.arange(compressed_num_frames, dtype=torch.float32)
+        grid_h = torch.arange(height, dtype=torch.float32)
+        grid_w = torch.arange(width, dtype=torch.float32)
+
+        freqs_t = torch.einsum("..., f -> ... f", grid_t, freqs_t)
+        freqs_h = torch.einsum("..., f -> ... f", grid_h, freqs_h)
+        freqs_w = torch.einsum("..., f -> ... f", grid_w, freqs_w)
+
+        freqs_t = repeat(freqs_t, "... n -> ... (n r)", r=2)
+        freqs_h = repeat(freqs_h, "... n -> ... (n r)", r=2)
+        freqs_w = repeat(freqs_w, "... n -> ... (n r)", r=2)
+
+        freqs = broadcat((freqs_t[:, None, None, :], freqs_h[None, :, None, :], freqs_w[None, None, :, :]), dim=-1)
+        # (T H W D)
+
+        self.pnp = pnp
+
+        if not self.pnp:
+            freqs = rearrange(freqs, "t h w d -> (t h w) d")
+
+        freqs = freqs.contiguous()
+        freqs_sin = freqs.sin()
+        freqs_cos = freqs.cos()
+        self.register_buffer("freqs_sin", freqs_sin)
+        self.register_buffer("freqs_cos", freqs_cos)
+
+        self.text_length = text_length
+        if learnable_pos_embed:
+            num_patches = height * width * compressed_num_frames + text_length
+            self.pos_embedding = nn.Parameter(torch.zeros(1, num_patches, int(hidden_size)), requires_grad=True)
+        else:
+            self.pos_embedding = None
+
+    def rotary(self, t, **kwargs):
+        if self.pnp:
+            t_coords = kwargs["rope_position_ids"][:, :, 0]
+            x_coords = kwargs["rope_position_ids"][:, :, 1]
+            y_coords = kwargs["rope_position_ids"][:, :, 2]
+            mask = (x_coords != -1) & (y_coords != -1) & (t_coords != -1)
+            freqs = torch.zeros([t.shape[0], t.shape[2], t.shape[3]], dtype=t.dtype, device=t.device)
+            freqs[mask] = self.freqs[t_coords[mask], x_coords[mask], y_coords[mask]]
+
+        else:
+
+            def reshape_freq(freqs):
+                frame = t.shape[2]
+                freqs = freqs[:frame].contiguous()
+                freqs = freqs.unsqueeze(0).unsqueeze(0)
+                return freqs
+
+            freqs_cos = reshape_freq(self.freqs_cos)
+            freqs_sin = reshape_freq(self.freqs_sin)
+
+        return t * freqs_cos + rotate_half(t) * freqs_sin
+
+    def position_embedding_forward(self, position_ids, **kwargs):
+        if self.pos_embedding is not None:
+            return self.pos_embedding[:, : self.text_length + kwargs["seq_length"]]
+        else:
+            return None
+
+    def attention_fn(
+        self,
+        query_layer,
+        key_layer,
+        value_layer,
+        attention_mask,
+        attention_dropout=None,
+        log_attention_weights=None,
+        scaling_attention_score=True,
+        **kwargs,
+    ):
+        attention_fn_default = HOOKS_DEFAULT["attention_fn"]
+
+        if self.pnp:
+            query_layer = self.rotary(query_layer, **kwargs)
+            key_layer = self.rotary(key_layer, **kwargs)
+            if self.rot_v:
+                value_layer = self.rotary(value_layer)
+        else:
+            query_layer = torch.cat(
+                (
+                    query_layer[
+                        :,
+                        :,
+                        : kwargs["text_length"],
+                    ],
+                    self.rotary(
+                        query_layer[
+                            :,
+                            :,
+                            kwargs["text_length"] :,
+                        ]
+                    ),
+                ),
+                dim=2,
+            )
+            key_layer = torch.cat(
+                (
+                    key_layer[
+                        :,
+                        :,
+                        : kwargs["text_length"],
+                    ],
+                    self.rotary(
+                        key_layer[
+                            :,
+                            :,
+                            kwargs["text_length"] :,
+                        ]
+                    ),
+                ),
+                dim=2,
+            )
+            if self.rot_v:
+                value_layer = torch.cat(
+                    (
+                        value_layer[
+                            :,
+                            :,
+                            : kwargs["text_length"],
+                        ],
+                        self.rotary(
+                            value_layer[
+                                :,
+                                :,
+                                kwargs["text_length"] :,
+                            ]
+                        ),
+                    ),
+                    dim=2,
+                )
+
+        return attention_fn_default(
+            query_layer,
+            key_layer,
+            value_layer,
+            attention_mask,
+            attention_dropout=attention_dropout,
+            log_attention_weights=log_attention_weights,
+            scaling_attention_score=scaling_attention_score,
+            **kwargs,
+        )
+
+
+def modulate(x, shift, scale):
+    return x * (1 + scale.unsqueeze(1)) + shift.unsqueeze(1)
+
+
+def unpatchify(x, c, p, w, h, rope_position_ids=None, **kwargs):
+    """
+    x: (N, T/2 * S, patch_size**3 * C)
+    imgs: (N, T, H, W, C)
+    """
+    if rope_position_ids is not None:
+        assert NotImplementedError
+        # do pix2struct unpatchify
+        L = x.shape[1]
+        x = x.reshape(shape=(x.shape[0], L, p, p, c))
+        x = torch.einsum("nlpqc->ncplq", x)
+        imgs = x.reshape(shape=(x.shape[0], c, p, L * p))
+    else:
+        b = x.shape[0]
+        imgs = rearrange(x, "b (t h w) (c p q) -> b t c (h p) (w q)", b=b, h=h, w=w, c=c, p=p, q=p)
+
+    return imgs
+
+
+class FinalLayerMixin(BaseMixin):
+    def __init__(
+        self,
+        hidden_size,
+        time_embed_dim,
+        patch_size,
+        out_channels,
+        latent_width,
+        latent_height,
+        elementwise_affine,
+    ):
+        super().__init__()
+        self.hidden_size = hidden_size
+        self.patch_size = patch_size
+        self.out_channels = out_channels
+        self.norm_final = nn.LayerNorm(hidden_size, elementwise_affine=elementwise_affine, eps=1e-6)
+        self.linear = nn.Linear(hidden_size, patch_size * patch_size * out_channels, bias=True)
+        self.adaLN_modulation = nn.Sequential(nn.SiLU(), nn.Linear(time_embed_dim, 2 * hidden_size, bias=True))
+
+        self.spatial_length = latent_width * latent_height // patch_size**2
+        self.latent_width = latent_width
+        self.latent_height = latent_height
+
+    def final_forward(self, logits, **kwargs):
+        x, emb = logits[:, kwargs["text_length"] :, :], kwargs["emb"]  # x:(b,(t n),d)
+
+        shift, scale = self.adaLN_modulation(emb).chunk(2, dim=1)
+        x = modulate(self.norm_final(x), shift, scale)
+        x = self.linear(x)
+
+        return unpatchify(
+            x,
+            c=self.out_channels,
+            p=self.patch_size,
+            w=self.latent_width // self.patch_size,
+            h=self.latent_height // self.patch_size,
+            rope_position_ids=kwargs.get("rope_position_ids", None),
+            **kwargs,
+        )
+
+    def reinit(self, parent_model=None):
+        nn.init.xavier_uniform_(self.linear.weight)
+        nn.init.constant_(self.linear.bias, 0)
+
+
+class SwiGLUMixin(BaseMixin):
+    def __init__(self, num_layers, in_features, hidden_features, bias=False):
+        super().__init__()
+        self.w2 = nn.ModuleList(
+            [
+                ColumnParallelLinear(
+                    in_features,
+                    hidden_features,
+                    gather_output=False,
+                    bias=bias,
+                    module=self,
+                    name="dense_h_to_4h_gate",
+                )
+                for i in range(num_layers)
+            ]
+        )
+
+    def mlp_forward(self, hidden_states, **kw_args):
+        x = hidden_states
+        origin = self.transformer.layers[kw_args["layer_id"]].mlp
+        x1 = origin.dense_h_to_4h(x)
+        x2 = self.w2[kw_args["layer_id"]](x)
+        hidden = origin.activation_func(x2) * x1
+        x = origin.dense_4h_to_h(hidden)
+        return x
+
+
+class AdaLNMixin(BaseMixin):
+    def __init__(
+        self,
+        width,
+        height,
+        hidden_size,
+        num_layers,
+        time_embed_dim,
+        compressed_num_frames,
+        qk_ln=True,
+        hidden_size_head=None,
+        elementwise_affine=True,
+    ):
+        super().__init__()
+        self.num_layers = num_layers
+        self.width = width
+        self.height = height
+        self.compressed_num_frames = compressed_num_frames
+
+        self.adaLN_modulations = nn.ModuleList(
+            [nn.Sequential(nn.SiLU(), nn.Linear(time_embed_dim, 12 * hidden_size)) for _ in range(num_layers)]
+        )
+
+        self.qk_ln = qk_ln
+        if qk_ln:
+            self.query_layernorm_list = nn.ModuleList(
+                [
+                    LayerNorm(hidden_size_head, eps=1e-6, elementwise_affine=elementwise_affine)
+                    for _ in range(num_layers)
+                ]
+            )
+            self.key_layernorm_list = nn.ModuleList(
+                [
+                    LayerNorm(hidden_size_head, eps=1e-6, elementwise_affine=elementwise_affine)
+                    for _ in range(num_layers)
+                ]
+            )
+
+    def layer_forward(
+        self,
+        hidden_states,
+        mask,
+        *args,
+        **kwargs,
+    ):
+        text_length = kwargs["text_length"]
+        # hidden_states (b,(n_t+t*n_i),d)
+        text_hidden_states = hidden_states[:, :text_length]  # (b,n,d)
+        img_hidden_states = hidden_states[:, text_length:]  # (b,(t n),d)
+        layer = self.transformer.layers[kwargs["layer_id"]]
+        adaLN_modulation = self.adaLN_modulations[kwargs["layer_id"]]
+
+        (
+            shift_msa,
+            scale_msa,
+            gate_msa,
+            shift_mlp,
+            scale_mlp,
+            gate_mlp,
+            text_shift_msa,
+            text_scale_msa,
+            text_gate_msa,
+            text_shift_mlp,
+            text_scale_mlp,
+            text_gate_mlp,
+        ) = adaLN_modulation(kwargs["emb"]).chunk(12, dim=1)
+        gate_msa, gate_mlp, text_gate_msa, text_gate_mlp = (
+            gate_msa.unsqueeze(1),
+            gate_mlp.unsqueeze(1),
+            text_gate_msa.unsqueeze(1),
+            text_gate_mlp.unsqueeze(1),
+        )
+
+        # self full attention (b,(t n),d)
+        img_attention_input = layer.input_layernorm(img_hidden_states)
+        text_attention_input = layer.input_layernorm(text_hidden_states)
+        img_attention_input = modulate(img_attention_input, shift_msa, scale_msa)
+        text_attention_input = modulate(text_attention_input, text_shift_msa, text_scale_msa)
+
+        attention_input = torch.cat((text_attention_input, img_attention_input), dim=1)  # (b,n_t+t*n_i,d)
+        attention_output = layer.attention(attention_input, mask, **kwargs)
+        text_attention_output = attention_output[:, :text_length]  # (b,n,d)
+        img_attention_output = attention_output[:, text_length:]  # (b,(t n),d)
+
+        if self.transformer.layernorm_order == "sandwich":
+            text_attention_output = layer.third_layernorm(text_attention_output)
+            img_attention_output = layer.third_layernorm(img_attention_output)
+        img_hidden_states = img_hidden_states + gate_msa * img_attention_output  # (b,(t n),d)
+        text_hidden_states = text_hidden_states + text_gate_msa * text_attention_output  # (b,n,d)
+
+        # mlp (b,(t n),d)
+        img_mlp_input = layer.post_attention_layernorm(img_hidden_states)  # vision (b,(t n),d)
+        text_mlp_input = layer.post_attention_layernorm(text_hidden_states)  # language (b,n,d)
+        img_mlp_input = modulate(img_mlp_input, shift_mlp, scale_mlp)
+        text_mlp_input = modulate(text_mlp_input, text_shift_mlp, text_scale_mlp)
+        mlp_input = torch.cat((text_mlp_input, img_mlp_input), dim=1)  # (b,(n_t+t*n_i),d
+        mlp_output = layer.mlp(mlp_input, **kwargs)
+        img_mlp_output = mlp_output[:, text_length:]  # vision (b,(t n),d)
+        text_mlp_output = mlp_output[:, :text_length]  # language (b,n,d)
+        if self.transformer.layernorm_order == "sandwich":
+            text_mlp_output = layer.fourth_layernorm(text_mlp_output)
+            img_mlp_output = layer.fourth_layernorm(img_mlp_output)
+
+        img_hidden_states = img_hidden_states + gate_mlp * img_mlp_output  # vision (b,(t n),d)
+        text_hidden_states = text_hidden_states + text_gate_mlp * text_mlp_output  # language (b,n,d)
+
+        hidden_states = torch.cat((text_hidden_states, img_hidden_states), dim=1)  # (b,(n_t+t*n_i),d)
+        return hidden_states
+
+    def reinit(self, parent_model=None):
+        for layer in self.adaLN_modulations:
+            nn.init.constant_(layer[-1].weight, 0)
+            nn.init.constant_(layer[-1].bias, 0)
+
+    @non_conflict
+    def attention_fn(
+        self,
+        query_layer,
+        key_layer,
+        value_layer,
+        attention_mask,
+        attention_dropout=None,
+        log_attention_weights=None,
+        scaling_attention_score=True,
+        old_impl=attention_fn_default,
+        **kwargs,
+    ):
+        if self.qk_ln:
+            query_layernorm = self.query_layernorm_list[kwargs["layer_id"]]
+            key_layernorm = self.key_layernorm_list[kwargs["layer_id"]]
+            query_layer = query_layernorm(query_layer)
+            key_layer = key_layernorm(key_layer)
+
+        return old_impl(
+            query_layer,
+            key_layer,
+            value_layer,
+            attention_mask,
+            attention_dropout=attention_dropout,
+            log_attention_weights=log_attention_weights,
+            scaling_attention_score=scaling_attention_score,
+            **kwargs,
+        )
+
+
+str_to_dtype = {"fp32": torch.float32, "fp16": torch.float16, "bf16": torch.bfloat16}
+
+
+class DiffusionTransformer(BaseModel):
+    def __init__(
+        self,
+        transformer_args,
+        num_frames,
+        time_compressed_rate,
+        latent_width,
+        latent_height,
+        patch_size,
+        in_channels,
+        out_channels,
+        hidden_size,
+        num_layers,
+        num_attention_heads,
+        elementwise_affine,
+        time_embed_dim=None,
+        num_classes=None,
+        modules={},
+        input_time="adaln",
+        adm_in_channels=None,
+        parallel_output=True,
+        height_interpolation=1.0,
+        width_interpolation=1.0,
+        time_interpolation=1.0,
+        use_SwiGLU=False,
+        use_RMSNorm=False,
+        zero_init_y_embed=False,
+        **kwargs,
+    ):
+        self.latent_width = latent_width
+        self.latent_height = latent_height
+        self.patch_size = patch_size
+        self.num_frames = num_frames
+        self.time_compressed_rate = time_compressed_rate
+        self.spatial_length = latent_width * latent_height // patch_size**2
+        self.in_channels = in_channels
+        self.out_channels = out_channels
+        self.hidden_size = hidden_size
+        self.model_channels = hidden_size
+        self.time_embed_dim = time_embed_dim if time_embed_dim is not None else hidden_size
+        self.num_classes = num_classes
+        self.adm_in_channels = adm_in_channels
+        self.input_time = input_time
+        self.num_layers = num_layers
+        self.num_attention_heads = num_attention_heads
+        self.is_decoder = transformer_args.is_decoder
+        self.elementwise_affine = elementwise_affine
+        self.height_interpolation = height_interpolation
+        self.width_interpolation = width_interpolation
+        self.time_interpolation = time_interpolation
+        self.inner_hidden_size = hidden_size * 4
+        self.zero_init_y_embed = zero_init_y_embed
+        try:
+            self.dtype = str_to_dtype[kwargs.pop("dtype")]
+        except:
+            self.dtype = torch.float32
+
+        if use_SwiGLU:
+            kwargs["activation_func"] = F.silu
+        elif "activation_func" not in kwargs:
+            approx_gelu = nn.GELU(approximate="tanh")
+            kwargs["activation_func"] = approx_gelu
+
+        if use_RMSNorm:
+            kwargs["layernorm"] = RMSNorm
+        else:
+            kwargs["layernorm"] = partial(LayerNorm, elementwise_affine=elementwise_affine, eps=1e-6)
+
+        transformer_args.num_layers = num_layers
+        transformer_args.hidden_size = hidden_size
+        transformer_args.num_attention_heads = num_attention_heads
+        transformer_args.parallel_output = parallel_output
+        super().__init__(args=transformer_args, transformer=None, **kwargs)
+
+        module_configs = modules
+        self._build_modules(module_configs)
+
+        if use_SwiGLU:
+            self.add_mixin(
+                "swiglu", SwiGLUMixin(num_layers, hidden_size, self.inner_hidden_size, bias=False), reinit=True
+            )
+
+    def _build_modules(self, module_configs):
+        model_channels = self.hidden_size
+        # time_embed_dim = model_channels * 4
+        time_embed_dim = self.time_embed_dim
+        self.time_embed = nn.Sequential(
+            linear(model_channels, time_embed_dim),
+            nn.SiLU(),
+            linear(time_embed_dim, time_embed_dim),
+        )
+
+        if self.num_classes is not None:
+            if isinstance(self.num_classes, int):
+                self.label_emb = nn.Embedding(self.num_classes, time_embed_dim)
+            elif self.num_classes == "continuous":
+                print("setting up linear c_adm embedding layer")
+                self.label_emb = nn.Linear(1, time_embed_dim)
+            elif self.num_classes == "timestep":
+                self.label_emb = nn.Sequential(
+                    Timestep(model_channels),
+                    nn.Sequential(
+                        linear(model_channels, time_embed_dim),
+                        nn.SiLU(),
+                        linear(time_embed_dim, time_embed_dim),
+                    ),
+                )
+            elif self.num_classes == "sequential":
+                assert self.adm_in_channels is not None
+                self.label_emb = nn.Sequential(
+                    nn.Sequential(
+                        linear(self.adm_in_channels, time_embed_dim),
+                        nn.SiLU(),
+                        linear(time_embed_dim, time_embed_dim),
+                    )
+                )
+                if self.zero_init_y_embed:
+                    nn.init.constant_(self.label_emb[0][2].weight, 0)
+                    nn.init.constant_(self.label_emb[0][2].bias, 0)
+            else:
+                raise ValueError()
+
+        pos_embed_config = module_configs["pos_embed_config"]
+        self.add_mixin(
+            "pos_embed",
+            instantiate_from_config(
+                pos_embed_config,
+                height=self.latent_height // self.patch_size,
+                width=self.latent_width // self.patch_size,
+                compressed_num_frames=(self.num_frames - 1) // self.time_compressed_rate + 1,
+                hidden_size=self.hidden_size,
+            ),
+            reinit=True,
+        )
+
+        patch_embed_config = module_configs["patch_embed_config"]
+        self.add_mixin(
+            "patch_embed",
+            instantiate_from_config(
+                patch_embed_config,
+                patch_size=self.patch_size,
+                hidden_size=self.hidden_size,
+                in_channels=self.in_channels,
+            ),
+            reinit=True,
+        )
+        if self.input_time == "adaln":
+            adaln_layer_config = module_configs["adaln_layer_config"]
+            self.add_mixin(
+                "adaln_layer",
+                instantiate_from_config(
+                    adaln_layer_config,
+                    height=self.latent_height // self.patch_size,
+                    width=self.latent_width // self.patch_size,
+                    hidden_size=self.hidden_size,
+                    num_layers=self.num_layers,
+                    compressed_num_frames=(self.num_frames - 1) // self.time_compressed_rate + 1,
+                    hidden_size_head=self.hidden_size // self.num_attention_heads,
+                    time_embed_dim=self.time_embed_dim,
+                    elementwise_affine=self.elementwise_affine,
+                ),
+            )
+        else:
+            raise NotImplementedError
+
+        final_layer_config = module_configs["final_layer_config"]
+        self.add_mixin(
+            "final_layer",
+            instantiate_from_config(
+                final_layer_config,
+                hidden_size=self.hidden_size,
+                patch_size=self.patch_size,
+                out_channels=self.out_channels,
+                time_embed_dim=self.time_embed_dim,
+                latent_width=self.latent_width,
+                latent_height=self.latent_height,
+                elementwise_affine=self.elementwise_affine,
+            ),
+            reinit=True,
+        )
+
+        if "lora_config" in module_configs:
+            lora_config = module_configs["lora_config"]
+            self.add_mixin("lora", instantiate_from_config(lora_config, layer_num=self.num_layers), reinit=True)
+
+        return
+
+    def forward(self, x, timesteps=None, context=None, y=None, **kwargs):
+        b, t, d, h, w = x.shape
+        if x.dtype != self.dtype:
+            x = x.to(self.dtype)
+        assert (y is not None) == (
+            self.num_classes is not None
+        ), "must specify y if and only if the model is class-conditional"
+        t_emb = timestep_embedding(timesteps, self.model_channels, repeat_only=False, dtype=self.dtype)
+        emb = self.time_embed(t_emb)
+
+        if self.num_classes is not None:
+            # assert y.shape[0] == x.shape[0]
+            assert x.shape[0] % y.shape[0] == 0
+            y = y.repeat_interleave(x.shape[0] // y.shape[0], dim=0)
+            emb = emb + self.label_emb(y)
+
+        kwargs["seq_length"] = t * h * w // (self.patch_size**2)
+        kwargs["images"] = x
+        kwargs["emb"] = emb
+        kwargs["encoder_outputs"] = context
+        kwargs["text_length"] = context.shape[1]
+
+        kwargs["input_ids"] = kwargs["position_ids"] = kwargs["attention_mask"] = torch.ones((1, 1)).to(x.dtype)
+        output = super().forward(**kwargs)[0]
+
+        return output

sat/finetune.sh

@@ -0,0 +1,12 @@
+#! /bin/bash
+
+echo "RUN on `hostname`, CUDA_VISIBLE_DEVICES=$CUDA_VISIBLE_DEVICES"
+
+environs="WORLD_SIZE=1 RANK=0 LOCAL_RANK=0 LOCAL_WORLD_SIZE=1"
+
+run_cmd="$environs python train_video.py --base configs/cogvideox_2b_sft.yaml --seed $RANDOM"
+
+echo ${run_cmd}
+eval ${run_cmd}
+
+echo "DONE on `hostname`"

sat/inference.sh

@@ -0,0 +1,12 @@
+#! /bin/bash
+
+echo "CUDA_VISIBLE_DEVICES=$CUDA_VISIBLE_DEVICES"
+
+environs="WORLD_SIZE=1 RANK=0 LOCAL_RANK=0 LOCAL_WORLD_SIZE=1"
+
+run_cmd="$environs python sample_video.py --base configs/cogvideox_2b_infer.yaml"
+
+echo ${run_cmd}
+eval ${run_cmd}
+
+echo "DONE on `hostname`"

sat/requirements.txt

@@ -0,0 +1,17 @@
+git+https://github.com/spacegoing/SwissArmyTransformer.git
+diffusers>=0.29.2
+omegaconf>=2.3.0
+torch>=2.3.1
+torchvision>=0.19.0
+pytorch_lightning>=2.3.3
+kornia>=0.7.3
+beartype>=0.18.5
+numpy>=2.0.1
+fsspec>=2024.5.0
+safetensors>=0.4.3
+imageio-ffmpeg>=0.5.1
+imageio>=2.34.2
+scipy>=1.14.0
+decord>=0.6.0
+wandb>=0.17.5
+deepspeed>=0.14.4

sat/sample_video.py

@@ -0,0 +1,236 @@
+import os
+import math
+import argparse
+from typing import List, Union
+from tqdm import tqdm
+from omegaconf import ListConfig
+import imageio
+
+import torch
+import numpy as np
+from einops import rearrange
+import torchvision.transforms as TT
+
+from sat.model.base_model import get_model
+from sat.training.model_io import load_checkpoint
+from sat import mpu
+
+from diffusion_video import SATVideoDiffusionEngine
+from arguments import get_args
+from torchvision.transforms.functional import center_crop, resize
+from torchvision.transforms import InterpolationMode
+
+
+def read_from_cli():
+    cnt = 0
+    try:
+        while True:
+            x = input("Please input English text (Ctrl-D quit): ")
+            yield x.strip(), cnt
+            cnt += 1
+    except EOFError as e:
+        pass
+
+
+def read_from_file(p, rank=0, world_size=1):
+    with open(p, "r") as fin:
+        cnt = -1
+        for l in fin:
+            cnt += 1
+            if cnt % world_size != rank:
+                continue
+            yield l.strip(), cnt
+
+
+def get_unique_embedder_keys_from_conditioner(conditioner):
+    return list(set([x.input_key for x in conditioner.embedders]))
+
+
+def get_batch(keys, value_dict, N: Union[List, ListConfig], T=None, device="cuda"):
+    batch = {}
+    batch_uc = {}
+
+    for key in keys:
+        if key == "txt":
+            batch["txt"] = np.repeat([value_dict["prompt"]], repeats=math.prod(N)).reshape(N).tolist()
+            batch_uc["txt"] = np.repeat([value_dict["negative_prompt"]], repeats=math.prod(N)).reshape(N).tolist()
+        else:
+            batch[key] = value_dict[key]
+
+    if T is not None:
+        batch["num_video_frames"] = T
+
+    for key in batch.keys():
+        if key not in batch_uc and isinstance(batch[key], torch.Tensor):
+            batch_uc[key] = torch.clone(batch[key])
+    return batch, batch_uc
+
+
+def save_video_as_grid_and_mp4(video_batch: torch.Tensor, save_path: str, fps: int = 5, args=None, key=None):
+    os.makedirs(save_path, exist_ok=True)
+
+    for i, vid in enumerate(video_batch):
+        gif_frames = []
+        for frame in vid:
+            frame = rearrange(frame, "c h w -> h w c")
+            frame = (255.0 * frame).cpu().numpy().astype(np.uint8)
+            gif_frames.append(frame)
+        now_save_path = os.path.join(save_path, f"{i:06d}.mp4")
+        with imageio.get_writer(now_save_path, fps=fps) as writer:
+            for frame in gif_frames:
+                writer.append_data(frame)
+
+
+def resize_for_rectangle_crop(arr, image_size, reshape_mode="random"):
+    if arr.shape[3] / arr.shape[2] > image_size[1] / image_size[0]:
+        arr = resize(
+            arr,
+            size=[image_size[0], int(arr.shape[3] * image_size[0] / arr.shape[2])],
+            interpolation=InterpolationMode.BICUBIC,
+        )
+    else:
+        arr = resize(
+            arr,
+            size=[int(arr.shape[2] * image_size[1] / arr.shape[3]), image_size[1]],
+            interpolation=InterpolationMode.BICUBIC,
+        )
+
+    h, w = arr.shape[2], arr.shape[3]
+    arr = arr.squeeze(0)
+
+    delta_h = h - image_size[0]
+    delta_w = w - image_size[1]
+
+    if reshape_mode == "random" or reshape_mode == "none":
+        top = np.random.randint(0, delta_h + 1)
+        left = np.random.randint(0, delta_w + 1)
+    elif reshape_mode == "center":
+        top, left = delta_h // 2, delta_w // 2
+    else:
+        raise NotImplementedError
+    arr = TT.functional.crop(arr, top=top, left=left, height=image_size[0], width=image_size[1])
+    return arr
+
+
+def sampling_main(args, model_cls):
+    if isinstance(model_cls, type):
+        model = get_model(args, model_cls)
+    else:
+        model = model_cls
+
+    load_checkpoint(model, args)
+    model.eval()
+
+    if args.input_type == "cli":
+        data_iter = read_from_cli()
+    elif args.input_type == "txt":
+        rank, world_size = mpu.get_data_parallel_rank(), mpu.get_data_parallel_world_size()
+        print("rank and world_size", rank, world_size)
+        data_iter = read_from_file(args.input_file, rank=rank, world_size=world_size)
+    else:
+        raise NotImplementedError
+
+    image_size = [480, 720]
+
+    sample_func = model.sample
+    T, H, W, C, F = args.sampling_num_frames, image_size[0], image_size[1], args.latent_channels, 8
+    num_samples = [1]
+    force_uc_zero_embeddings = ["txt"]
+    device = model.device
+    with torch.no_grad():
+        for text, cnt in tqdm(data_iter):
+            # reload model on GPU
+            model.to(device)
+            print("rank:", rank, "start to process", text, cnt)
+            # TODO: broadcast image2video
+            value_dict = {
+                "prompt": text,
+                "negative_prompt": "",
+                "num_frames": torch.tensor(T).unsqueeze(0),
+            }
+
+            batch, batch_uc = get_batch(
+                get_unique_embedder_keys_from_conditioner(model.conditioner), value_dict, num_samples
+            )
+            for key in batch:
+                if isinstance(batch[key], torch.Tensor):
+                    print(key, batch[key].shape)
+                elif isinstance(batch[key], list):
+                    print(key, [len(l) for l in batch[key]])
+                else:
+                    print(key, batch[key])
+            c, uc = model.conditioner.get_unconditional_conditioning(
+                batch,
+                batch_uc=batch_uc,
+                force_uc_zero_embeddings=force_uc_zero_embeddings,
+            )
+
+            for k in c:
+                if not k == "crossattn":
+                    c[k], uc[k] = map(lambda y: y[k][: math.prod(num_samples)].to("cuda"), (c, uc))
+            for index in range(args.batch_size):
+                # reload model on GPU
+                model.to(device)
+                samples_z = sample_func(
+                    c,
+                    uc=uc,
+                    batch_size=1,
+                    shape=(T, C, H // F, W // F),
+                )
+                samples_z = samples_z.permute(0, 2, 1, 3, 4).contiguous()
+
+                # Unload the model from GPU to save GPU memory
+                model.to("cpu")
+                torch.cuda.empty_cache()
+                first_stage_model = model.first_stage_model
+                first_stage_model = first_stage_model.to(device)
+
+                latent = 1.0 / model.scale_factor * samples_z
+
+                # Decode latent serial to save GPU memory
+                recons = []
+                loop_num = (T - 1) // 2
+                for i in range(loop_num):
+                    if i == 0:
+                        start_frame, end_frame = 0, 3
+                    else:
+                        start_frame, end_frame = i * 2 + 1, i * 2 + 3
+                    if i == loop_num - 1:
+                        clear_fake_cp_cache = True
+                    else:
+                        clear_fake_cp_cache = False
+                    with torch.no_grad():
+                        recon = first_stage_model.decode(
+                            latent[:, :, start_frame:end_frame].contiguous(), clear_fake_cp_cache=clear_fake_cp_cache
+                        )
+
+                    recons.append(recon)
+
+                recon = torch.cat(recons, dim=2).to(torch.float32)
+                samples_x = recon.permute(0, 2, 1, 3, 4).contiguous()
+                samples = torch.clamp((samples_x + 1.0) / 2.0, min=0.0, max=1.0).cpu()
+
+                save_path = os.path.join(
+                    args.output_dir, str(cnt) + "_" + text.replace(" ", "_").replace("/", "")[:120], str(index)
+                )
+                if mpu.get_model_parallel_rank() == 0:
+                    save_video_as_grid_and_mp4(samples, save_path, fps=args.sampling_fps)
+
+
+if __name__ == "__main__":
+    if "OMPI_COMM_WORLD_LOCAL_RANK" in os.environ:
+        os.environ["LOCAL_RANK"] = os.environ["OMPI_COMM_WORLD_LOCAL_RANK"]
+        os.environ["WORLD_SIZE"] = os.environ["OMPI_COMM_WORLD_SIZE"]
+        os.environ["RANK"] = os.environ["OMPI_COMM_WORLD_RANK"]
+    py_parser = argparse.ArgumentParser(add_help=False)
+    known, args_list = py_parser.parse_known_args()
+
+    args = get_args(args_list)
+    args = argparse.Namespace(**vars(args), **vars(known))
+    del args.deepspeed_config
+    args.model_config.first_stage_config.params.cp_size = 1
+    args.model_config.network_config.params.transformer_args.model_parallel_size = 1
+    args.model_config.network_config.params.transformer_args.checkpoint_activations = False
+    args.model_config.loss_fn_config.params.sigma_sampler_config.params.uniform_sampling = False
+
+    sampling_main(args, model_cls=SATVideoDiffusionEngine)

sat/sgm/__init__.py

@@ -0,0 +1,4 @@
+from .models import AutoencodingEngine
+from .util import get_configs_path, instantiate_from_config
+
+__version__ = "0.1.0"

sat/sgm/lr_scheduler.py

@@ -0,0 +1,110 @@
+import numpy as np
+
+
+class LambdaWarmUpCosineScheduler:
+    """
+    note: use with a base_lr of 1.0
+    """
+
+    def __init__(
+        self,
+        warm_up_steps,
+        lr_min,
+        lr_max,
+        lr_start,
+        max_decay_steps,
+        verbosity_interval=0,
+    ):
+        self.lr_warm_up_steps = warm_up_steps
+        self.lr_start = lr_start
+        self.lr_min = lr_min
+        self.lr_max = lr_max
+        self.lr_max_decay_steps = max_decay_steps
+        self.last_lr = 0.0
+        self.verbosity_interval = verbosity_interval
+
+    def schedule(self, n, **kwargs):
+        if self.verbosity_interval > 0:
+            if n % self.verbosity_interval == 0:
+                print(f"current step: {n}, recent lr-multiplier: {self.last_lr}")
+        if n < self.lr_warm_up_steps:
+            lr = (self.lr_max - self.lr_start) / self.lr_warm_up_steps * n + self.lr_start
+            self.last_lr = lr
+            return lr
+        else:
+            t = (n - self.lr_warm_up_steps) / (self.lr_max_decay_steps - self.lr_warm_up_steps)
+            t = min(t, 1.0)
+            lr = self.lr_min + 0.5 * (self.lr_max - self.lr_min) * (1 + np.cos(t * np.pi))
+            self.last_lr = lr
+            return lr
+
+    def __call__(self, n, **kwargs):
+        return self.schedule(n, **kwargs)
+
+
+class LambdaWarmUpCosineScheduler2:
+    """
+    supports repeated iterations, configurable via lists
+    note: use with a base_lr of 1.0.
+    """
+
+    def __init__(self, warm_up_steps, f_min, f_max, f_start, cycle_lengths, verbosity_interval=0):
+        assert len(warm_up_steps) == len(f_min) == len(f_max) == len(f_start) == len(cycle_lengths)
+        self.lr_warm_up_steps = warm_up_steps
+        self.f_start = f_start
+        self.f_min = f_min
+        self.f_max = f_max
+        self.cycle_lengths = cycle_lengths
+        self.cum_cycles = np.cumsum([0] + list(self.cycle_lengths))
+        self.last_f = 0.0
+        self.verbosity_interval = verbosity_interval
+
+    def find_in_interval(self, n):
+        interval = 0
+        for cl in self.cum_cycles[1:]:
+            if n <= cl:
+                return interval
+            interval += 1
+
+    def schedule(self, n, **kwargs):
+        cycle = self.find_in_interval(n)
+        n = n - self.cum_cycles[cycle]
+        if self.verbosity_interval > 0:
+            if n % self.verbosity_interval == 0:
+                print(f"current step: {n}, recent lr-multiplier: {self.last_f}, " f"current cycle {cycle}")
+        if n < self.lr_warm_up_steps[cycle]:
+            f = (self.f_max[cycle] - self.f_start[cycle]) / self.lr_warm_up_steps[cycle] * n + self.f_start[cycle]
+            self.last_f = f
+            return f
+        else:
+            t = (n - self.lr_warm_up_steps[cycle]) / (self.cycle_lengths[cycle] - self.lr_warm_up_steps[cycle])
+            t = min(t, 1.0)
+            f = self.f_min[cycle] + 0.5 * (self.f_max[cycle] - self.f_min[cycle]) * (1 + np.cos(t * np.pi))
+            self.last_f = f
+            return f
+
+    def __call__(self, n, **kwargs):
+        return self.schedule(n, **kwargs)
+
+
+class LambdaLinearScheduler(LambdaWarmUpCosineScheduler2):
+    def schedule(self, n, **kwargs):
+        cycle = self.find_in_interval(n)
+        n = n - self.cum_cycles[cycle]
+        if self.verbosity_interval > 0:
+            if n % self.verbosity_interval == 0:
+                print(f"current step: {n}, recent lr-multiplier: {self.last_f}, " f"current cycle {cycle}")
+
+        if n < self.lr_warm_up_steps[cycle]:
+            f = (self.f_max[cycle] - self.f_start[cycle]) / self.lr_warm_up_steps[cycle] * n + self.f_start[cycle]
+            self.last_f = f
+            return f
+        else:
+            f = (
+                self.f_min[cycle]
+                + (self.f_max[cycle] - self.f_min[cycle])
+                * (self.cycle_lengths[cycle] - n)
+                / (self.cycle_lengths[cycle])
+            )
+            self.last_f = f
+            return f

sat/sgm/models/__init__.py

@@ -0,0 +1 @@
+from .autoencoder import AutoencodingEngine

sat/sgm/models/autoencoder.py

@@ -0,0 +1,630 @@
+import logging
+import math
+import re
+import random
+from abc import abstractmethod
+from contextlib import contextmanager
+from typing import Any, Dict, List, Optional, Tuple, Union
+
+import numpy as np
+import pytorch_lightning as pl
+import torch
+import torch.distributed
+import torch.nn as nn
+from einops import rearrange
+from packaging import version
+
+from ..modules.autoencoding.regularizers import AbstractRegularizer
+from ..modules.ema import LitEma
+from ..util import (
+    default,
+    get_nested_attribute,
+    get_obj_from_str,
+    instantiate_from_config,
+    initialize_context_parallel,
+    get_context_parallel_group,
+    get_context_parallel_group_rank,
+    is_context_parallel_initialized,
+)
+from ..modules.cp_enc_dec import _conv_split, _conv_gather
+
+logpy = logging.getLogger(__name__)
+
+
+class AbstractAutoencoder(pl.LightningModule):
+    """
+    This is the base class for all autoencoders, including image autoencoders, image autoencoders with discriminators,
+    unCLIP models, etc. Hence, it is fairly general, and specific features
+    (e.g. discriminator training, encoding, decoding) must be implemented in subclasses.
+    """
+
+    def __init__(
+        self,
+        ema_decay: Union[None, float] = None,
+        monitor: Union[None, str] = None,
+        input_key: str = "jpg",
+    ):
+        super().__init__()
+
+        self.input_key = input_key
+        self.use_ema = ema_decay is not None
+        if monitor is not None:
+            self.monitor = monitor
+
+        if self.use_ema:
+            self.model_ema = LitEma(self, decay=ema_decay)
+            logpy.info(f"Keeping EMAs of {len(list(self.model_ema.buffers()))}.")
+
+        if version.parse(torch.__version__) >= version.parse("2.0.0"):
+            self.automatic_optimization = False
+
+    def apply_ckpt(self, ckpt: Union[None, str, dict]):
+        if ckpt is None:
+            return
+        if isinstance(ckpt, str):
+            ckpt = {
+                "target": "sgm.modules.checkpoint.CheckpointEngine",
+                "params": {"ckpt_path": ckpt},
+            }
+        engine = instantiate_from_config(ckpt)
+        engine(self)
+
+    @abstractmethod
+    def get_input(self, batch) -> Any:
+        raise NotImplementedError()
+
+    def on_train_batch_end(self, *args, **kwargs):
+        # for EMA computation
+        if self.use_ema:
+            self.model_ema(self)
+
+    @contextmanager
+    def ema_scope(self, context=None):
+        if self.use_ema:
+            self.model_ema.store(self.parameters())
+            self.model_ema.copy_to(self)
+            if context is not None:
+                logpy.info(f"{context}: Switched to EMA weights")
+        try:
+            yield None
+        finally:
+            if self.use_ema:
+                self.model_ema.restore(self.parameters())
+                if context is not None:
+                    logpy.info(f"{context}: Restored training weights")
+
+    @abstractmethod
+    def encode(self, *args, **kwargs) -> torch.Tensor:
+        raise NotImplementedError("encode()-method of abstract base class called")
+
+    @abstractmethod
+    def decode(self, *args, **kwargs) -> torch.Tensor:
+        raise NotImplementedError("decode()-method of abstract base class called")
+
+    def instantiate_optimizer_from_config(self, params, lr, cfg):
+        logpy.info(f"loading >>> {cfg['target']} <<< optimizer from config")
+        return get_obj_from_str(cfg["target"])(params, lr=lr, **cfg.get("params", dict()))
+
+    def configure_optimizers(self) -> Any:
+        raise NotImplementedError()
+
+
+class AutoencodingEngine(AbstractAutoencoder):
+    """
+    Base class for all image autoencoders that we train, like VQGAN or AutoencoderKL
+    (we also restore them explicitly as special cases for legacy reasons).
+    Regularizations such as KL or VQ are moved to the regularizer class.
+    """
+
+    def __init__(
+        self,
+        *args,
+        encoder_config: Dict,
+        decoder_config: Dict,
+        loss_config: Dict,
+        regularizer_config: Dict,
+        optimizer_config: Union[Dict, None] = None,
+        lr_g_factor: float = 1.0,
+        trainable_ae_params: Optional[List[List[str]]] = None,
+        ae_optimizer_args: Optional[List[dict]] = None,
+        trainable_disc_params: Optional[List[List[str]]] = None,
+        disc_optimizer_args: Optional[List[dict]] = None,
+        disc_start_iter: int = 0,
+        diff_boost_factor: float = 3.0,
+        ckpt_engine: Union[None, str, dict] = None,
+        ckpt_path: Optional[str] = None,
+        additional_decode_keys: Optional[List[str]] = None,
+        **kwargs,
+    ):
+        super().__init__(*args, **kwargs)
+        self.automatic_optimization = False  # pytorch lightning
+
+        self.encoder: torch.nn.Module = instantiate_from_config(encoder_config)
+        self.decoder: torch.nn.Module = instantiate_from_config(decoder_config)
+        self.loss: torch.nn.Module = instantiate_from_config(loss_config)
+        self.regularization: AbstractRegularizer = instantiate_from_config(regularizer_config)
+        self.optimizer_config = default(optimizer_config, {"target": "torch.optim.Adam"})
+        self.diff_boost_factor = diff_boost_factor
+        self.disc_start_iter = disc_start_iter
+        self.lr_g_factor = lr_g_factor
+        self.trainable_ae_params = trainable_ae_params
+        if self.trainable_ae_params is not None:
+            self.ae_optimizer_args = default(
+                ae_optimizer_args,
+                [{} for _ in range(len(self.trainable_ae_params))],
+            )
+            assert len(self.ae_optimizer_args) == len(self.trainable_ae_params)
+        else:
+            self.ae_optimizer_args = [{}]  # makes type consitent
+
+        self.trainable_disc_params = trainable_disc_params
+        if self.trainable_disc_params is not None:
+            self.disc_optimizer_args = default(
+                disc_optimizer_args,
+                [{} for _ in range(len(self.trainable_disc_params))],
+            )
+            assert len(self.disc_optimizer_args) == len(self.trainable_disc_params)
+        else:
+            self.disc_optimizer_args = [{}]  # makes type consitent
+
+        if ckpt_path is not None:
+            assert ckpt_engine is None, "Can't set ckpt_engine and ckpt_path"
+            logpy.warn("Checkpoint path is deprecated, use `checkpoint_egnine` instead")
+        self.apply_ckpt(default(ckpt_path, ckpt_engine))
+        self.additional_decode_keys = set(default(additional_decode_keys, []))
+
+    def get_input(self, batch: Dict) -> torch.Tensor:
+        # assuming unified data format, dataloader returns a dict.
+        # image tensors should be scaled to -1 ... 1 and in channels-first
+        # format (e.g., bchw instead if bhwc)
+        return batch[self.input_key]
+
+    def get_autoencoder_params(self) -> list:
+        params = []
+        if hasattr(self.loss, "get_trainable_autoencoder_parameters"):
+            params += list(self.loss.get_trainable_autoencoder_parameters())
+        if hasattr(self.regularization, "get_trainable_parameters"):
+            params += list(self.regularization.get_trainable_parameters())
+        params = params + list(self.encoder.parameters())
+        params = params + list(self.decoder.parameters())
+        return params
+
+    def get_discriminator_params(self) -> list:
+        if hasattr(self.loss, "get_trainable_parameters"):
+            params = list(self.loss.get_trainable_parameters())  # e.g., discriminator
+        else:
+            params = []
+        return params
+
+    def get_last_layer(self):
+        return self.decoder.get_last_layer()
+
+    def encode(
+        self,
+        x: torch.Tensor,
+        return_reg_log: bool = False,
+        unregularized: bool = False,
+        **kwargs,
+    ) -> Union[torch.Tensor, Tuple[torch.Tensor, dict]]:
+        z = self.encoder(x, **kwargs)
+        if unregularized:
+            return z, dict()
+        z, reg_log = self.regularization(z)
+        if return_reg_log:
+            return z, reg_log
+        return z
+
+    def decode(self, z: torch.Tensor, **kwargs) -> torch.Tensor:
+        x = self.decoder(z, **kwargs)
+        return x
+
+    def forward(self, x: torch.Tensor, **additional_decode_kwargs) -> Tuple[torch.Tensor, torch.Tensor, dict]:
+        z, reg_log = self.encode(x, return_reg_log=True)
+        dec = self.decode(z, **additional_decode_kwargs)
+        return z, dec, reg_log
+
+    def inner_training_step(self, batch: dict, batch_idx: int, optimizer_idx: int = 0) -> torch.Tensor:
+        x = self.get_input(batch)
+        additional_decode_kwargs = {key: batch[key] for key in self.additional_decode_keys.intersection(batch)}
+        z, xrec, regularization_log = self(x, **additional_decode_kwargs)
+        if hasattr(self.loss, "forward_keys"):
+            extra_info = {
+                "z": z,
+                "optimizer_idx": optimizer_idx,
+                "global_step": self.global_step,
+                "last_layer": self.get_last_layer(),
+                "split": "train",
+                "regularization_log": regularization_log,
+                "autoencoder": self,
+            }
+            extra_info = {k: extra_info[k] for k in self.loss.forward_keys}
+        else:
+            extra_info = dict()
+
+        if optimizer_idx == 0:
+            # autoencode
+            out_loss = self.loss(x, xrec, **extra_info)
+            if isinstance(out_loss, tuple):
+                aeloss, log_dict_ae = out_loss
+            else:
+                # simple loss function
+                aeloss = out_loss
+                log_dict_ae = {"train/loss/rec": aeloss.detach()}
+
+            self.log_dict(
+                log_dict_ae,
+                prog_bar=False,
+                logger=True,
+                on_step=True,
+                on_epoch=True,
+                sync_dist=False,
+            )
+            self.log(
+                "loss",
+                aeloss.mean().detach(),
+                prog_bar=True,
+                logger=False,
+                on_epoch=False,
+                on_step=True,
+            )
+            return aeloss
+        elif optimizer_idx == 1:
+            # discriminator
+            discloss, log_dict_disc = self.loss(x, xrec, **extra_info)
+            # -> discriminator always needs to return a tuple
+            self.log_dict(log_dict_disc, prog_bar=False, logger=True, on_step=True, on_epoch=True)
+            return discloss
+        else:
+            raise NotImplementedError(f"Unknown optimizer {optimizer_idx}")
+
+    def training_step(self, batch: dict, batch_idx: int):
+        opts = self.optimizers()
+        if not isinstance(opts, list):
+            # Non-adversarial case
+            opts = [opts]
+        optimizer_idx = batch_idx % len(opts)
+        if self.global_step < self.disc_start_iter:
+            optimizer_idx = 0
+        opt = opts[optimizer_idx]
+        opt.zero_grad()
+        with opt.toggle_model():
+            loss = self.inner_training_step(batch, batch_idx, optimizer_idx=optimizer_idx)
+            self.manual_backward(loss)
+        opt.step()
+
+    def validation_step(self, batch: dict, batch_idx: int) -> Dict:
+        log_dict = self._validation_step(batch, batch_idx)
+        with self.ema_scope():
+            log_dict_ema = self._validation_step(batch, batch_idx, postfix="_ema")
+            log_dict.update(log_dict_ema)
+        return log_dict
+
+    def _validation_step(self, batch: dict, batch_idx: int, postfix: str = "") -> Dict:
+        x = self.get_input(batch)
+
+        z, xrec, regularization_log = self(x)
+        if hasattr(self.loss, "forward_keys"):
+            extra_info = {
+                "z": z,
+                "optimizer_idx": 0,
+                "global_step": self.global_step,
+                "last_layer": self.get_last_layer(),
+                "split": "val" + postfix,
+                "regularization_log": regularization_log,
+                "autoencoder": self,
+            }
+            extra_info = {k: extra_info[k] for k in self.loss.forward_keys}
+        else:
+            extra_info = dict()
+        out_loss = self.loss(x, xrec, **extra_info)
+        if isinstance(out_loss, tuple):
+            aeloss, log_dict_ae = out_loss
+        else:
+            # simple loss function
+            aeloss = out_loss
+            log_dict_ae = {f"val{postfix}/loss/rec": aeloss.detach()}
+        full_log_dict = log_dict_ae
+
+        if "optimizer_idx" in extra_info:
+            extra_info["optimizer_idx"] = 1
+            discloss, log_dict_disc = self.loss(x, xrec, **extra_info)
+            full_log_dict.update(log_dict_disc)
+        self.log(
+            f"val{postfix}/loss/rec",
+            log_dict_ae[f"val{postfix}/loss/rec"],
+            sync_dist=True,
+        )
+        self.log_dict(full_log_dict, sync_dist=True)
+        return full_log_dict
+
+    def get_param_groups(
+        self, parameter_names: List[List[str]], optimizer_args: List[dict]
+    ) -> Tuple[List[Dict[str, Any]], int]:
+        groups = []
+        num_params = 0
+        for names, args in zip(parameter_names, optimizer_args):
+            params = []
+            for pattern_ in names:
+                pattern_params = []
+                pattern = re.compile(pattern_)
+                for p_name, param in self.named_parameters():
+                    if re.match(pattern, p_name):
+                        pattern_params.append(param)
+                        num_params += param.numel()
+                if len(pattern_params) == 0:
+                    logpy.warn(f"Did not find parameters for pattern {pattern_}")
+                params.extend(pattern_params)
+            groups.append({"params": params, **args})
+        return groups, num_params
+
+    def configure_optimizers(self) -> List[torch.optim.Optimizer]:
+        if self.trainable_ae_params is None:
+            ae_params = self.get_autoencoder_params()
+        else:
+            ae_params, num_ae_params = self.get_param_groups(self.trainable_ae_params, self.ae_optimizer_args)
+            logpy.info(f"Number of trainable autoencoder parameters: {num_ae_params:,}")
+        if self.trainable_disc_params is None:
+            disc_params = self.get_discriminator_params()
+        else:
+            disc_params, num_disc_params = self.get_param_groups(self.trainable_disc_params, self.disc_optimizer_args)
+            logpy.info(f"Number of trainable discriminator parameters: {num_disc_params:,}")
+        opt_ae = self.instantiate_optimizer_from_config(
+            ae_params,
+            default(self.lr_g_factor, 1.0) * self.learning_rate,
+            self.optimizer_config,
+        )
+        opts = [opt_ae]
+        if len(disc_params) > 0:
+            opt_disc = self.instantiate_optimizer_from_config(disc_params, self.learning_rate, self.optimizer_config)
+            opts.append(opt_disc)
+
+        return opts
+
+    @torch.no_grad()
+    def log_images(self, batch: dict, additional_log_kwargs: Optional[Dict] = None, **kwargs) -> dict:
+        log = dict()
+        additional_decode_kwargs = {}
+        x = self.get_input(batch)
+        additional_decode_kwargs.update({key: batch[key] for key in self.additional_decode_keys.intersection(batch)})
+
+        _, xrec, _ = self(x, **additional_decode_kwargs)
+        log["inputs"] = x
+        log["reconstructions"] = xrec
+        diff = 0.5 * torch.abs(torch.clamp(xrec, -1.0, 1.0) - x)
+        diff.clamp_(0, 1.0)
+        log["diff"] = 2.0 * diff - 1.0
+        # diff_boost shows location of small errors, by boosting their
+        # brightness.
+        log["diff_boost"] = 2.0 * torch.clamp(self.diff_boost_factor * diff, 0.0, 1.0) - 1
+        if hasattr(self.loss, "log_images"):
+            log.update(self.loss.log_images(x, xrec))
+        with self.ema_scope():
+            _, xrec_ema, _ = self(x, **additional_decode_kwargs)
+            log["reconstructions_ema"] = xrec_ema
+            diff_ema = 0.5 * torch.abs(torch.clamp(xrec_ema, -1.0, 1.0) - x)
+            diff_ema.clamp_(0, 1.0)
+            log["diff_ema"] = 2.0 * diff_ema - 1.0
+            log["diff_boost_ema"] = 2.0 * torch.clamp(self.diff_boost_factor * diff_ema, 0.0, 1.0) - 1
+        if additional_log_kwargs:
+            additional_decode_kwargs.update(additional_log_kwargs)
+            _, xrec_add, _ = self(x, **additional_decode_kwargs)
+            log_str = "reconstructions-" + "-".join(
+                [f"{key}={additional_log_kwargs[key]}" for key in additional_log_kwargs]
+            )
+            log[log_str] = xrec_add
+        return log
+
+
+class AutoencodingEngineLegacy(AutoencodingEngine):
+    def __init__(self, embed_dim: int, **kwargs):
+        self.max_batch_size = kwargs.pop("max_batch_size", None)
+        ddconfig = kwargs.pop("ddconfig")
+        ckpt_path = kwargs.pop("ckpt_path", None)
+        ckpt_engine = kwargs.pop("ckpt_engine", None)
+        super().__init__(
+            encoder_config={
+                "target": "sgm.modules.diffusionmodules.model.Encoder",
+                "params": ddconfig,
+            },
+            decoder_config={
+                "target": "sgm.modules.diffusionmodules.model.Decoder",
+                "params": ddconfig,
+            },
+            **kwargs,
+        )
+        self.quant_conv = torch.nn.Conv2d(
+            (1 + ddconfig["double_z"]) * ddconfig["z_channels"],
+            (1 + ddconfig["double_z"]) * embed_dim,
+            1,
+        )
+        self.post_quant_conv = torch.nn.Conv2d(embed_dim, ddconfig["z_channels"], 1)
+        self.embed_dim = embed_dim
+
+        self.apply_ckpt(default(ckpt_path, ckpt_engine))
+
+    def get_autoencoder_params(self) -> list:
+        params = super().get_autoencoder_params()
+        return params
+
+    def encode(self, x: torch.Tensor, return_reg_log: bool = False) -> Union[torch.Tensor, Tuple[torch.Tensor, dict]]:
+        if self.max_batch_size is None:
+            z = self.encoder(x)
+            z = self.quant_conv(z)
+        else:
+            N = x.shape[0]
+            bs = self.max_batch_size
+            n_batches = int(math.ceil(N / bs))
+            z = list()
+            for i_batch in range(n_batches):
+                z_batch = self.encoder(x[i_batch * bs : (i_batch + 1) * bs])
+                z_batch = self.quant_conv(z_batch)
+                z.append(z_batch)
+            z = torch.cat(z, 0)
+
+        z, reg_log = self.regularization(z)
+        if return_reg_log:
+            return z, reg_log
+        return z
+
+    def decode(self, z: torch.Tensor, **decoder_kwargs) -> torch.Tensor:
+        if self.max_batch_size is None:
+            dec = self.post_quant_conv(z)
+            dec = self.decoder(dec, **decoder_kwargs)
+        else:
+            N = z.shape[0]
+            bs = self.max_batch_size
+            n_batches = int(math.ceil(N / bs))
+            dec = list()
+            for i_batch in range(n_batches):
+                dec_batch = self.post_quant_conv(z[i_batch * bs : (i_batch + 1) * bs])
+                dec_batch = self.decoder(dec_batch, **decoder_kwargs)
+                dec.append(dec_batch)
+            dec = torch.cat(dec, 0)
+
+        return dec
+
+
+class IdentityFirstStage(AbstractAutoencoder):
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+
+    def get_input(self, x: Any) -> Any:
+        return x
+
+    def encode(self, x: Any, *args, **kwargs) -> Any:
+        return x
+
+    def decode(self, x: Any, *args, **kwargs) -> Any:
+        return
+
+
+class VideoAutoencodingEngine(AutoencodingEngine):
+    def __init__(
+        self,
+        ckpt_path: Union[None, str] = None,
+        ignore_keys: Union[Tuple, list] = (),
+        image_video_weights=[1, 1],
+        only_train_decoder=False,
+        context_parallel_size=0,
+        **kwargs,
+    ):
+        super().__init__(**kwargs)
+        self.context_parallel_size = context_parallel_size
+        if ckpt_path is not None:
+            self.init_from_ckpt(ckpt_path, ignore_keys=ignore_keys)
+
+    def log_videos(self, batch: dict, additional_log_kwargs: Optional[Dict] = None, **kwargs) -> dict:
+        return self.log_images(batch, additional_log_kwargs, **kwargs)
+
+    def get_input(self, batch: dict) -> torch.Tensor:
+        if self.context_parallel_size > 0:
+            if not is_context_parallel_initialized():
+                initialize_context_parallel(self.context_parallel_size)
+
+            batch = batch[self.input_key]
+
+            global_src_rank = get_context_parallel_group_rank() * self.context_parallel_size
+            torch.distributed.broadcast(batch, src=global_src_rank, group=get_context_parallel_group())
+
+            batch = _conv_split(batch, dim=2, kernel_size=1)
+            return batch
+
+        return batch[self.input_key]
+
+    def apply_ckpt(self, ckpt: Union[None, str, dict]):
+        if ckpt is None:
+            return
+        self.init_from_ckpt(ckpt)
+
+    def init_from_ckpt(self, path, ignore_keys=list()):
+        sd = torch.load(path, map_location="cpu")["state_dict"]
+        keys = list(sd.keys())
+        for k in keys:
+            for ik in ignore_keys:
+                if k.startswith(ik):
+                    del sd[k]
+        missing_keys, unexpected_keys = self.load_state_dict(sd, strict=False)
+        print("Missing keys: ", missing_keys)
+        print("Unexpected keys: ", unexpected_keys)
+        print(f"Restored from {path}")
+
+
+class VideoAutoencoderInferenceWrapper(VideoAutoencodingEngine):
+    def __init__(
+        self,
+        cp_size=0,
+        *args,
+        **kwargs,
+    ):
+        self.cp_size = cp_size
+        return super().__init__(*args, **kwargs)
+
+    def encode(
+        self,
+        x: torch.Tensor,
+        return_reg_log: bool = False,
+        unregularized: bool = False,
+        input_cp: bool = False,
+        output_cp: bool = False,
+        use_cp: bool = True,
+    ) -> Union[torch.Tensor, Tuple[torch.Tensor, dict]]:
+        if self.cp_size <= 1:
+            use_cp = False
+        if self.cp_size > 0 and use_cp and not input_cp:
+            if not is_context_parallel_initialized:
+                initialize_context_parallel(self.cp_size)
+
+            global_src_rank = get_context_parallel_group_rank() * self.cp_size
+            torch.distributed.broadcast(x, src=global_src_rank, group=get_context_parallel_group())
+
+            x = _conv_split(x, dim=2, kernel_size=1)
+
+        if return_reg_log:
+            z, reg_log = super().encode(x, return_reg_log, unregularized, use_cp=use_cp)
+        else:
+            z = super().encode(x, return_reg_log, unregularized, use_cp=use_cp)
+
+        if self.cp_size > 0 and use_cp and not output_cp:
+            z = _conv_gather(z, dim=2, kernel_size=1)
+
+        if return_reg_log:
+            return z, reg_log
+        return z
+
+    def decode(
+        self,
+        z: torch.Tensor,
+        input_cp: bool = False,
+        output_cp: bool = False,
+        use_cp: bool = True,
+        **kwargs,
+    ):
+        if self.cp_size <= 1:
+            use_cp = False
+        if self.cp_size > 0 and use_cp and not input_cp:
+            if not is_context_parallel_initialized:
+                initialize_context_parallel(self.cp_size)
+
+            global_src_rank = get_context_parallel_group_rank() * self.cp_size
+            torch.distributed.broadcast(z, src=global_src_rank, group=get_context_parallel_group())
+
+            z = _conv_split(z, dim=2, kernel_size=1)
+
+        x = super().decode(z, use_cp=use_cp, **kwargs)
+
+        if self.cp_size > 0 and use_cp and not output_cp:
+            x = _conv_gather(x, dim=2, kernel_size=1)
+
+        return x
+
+    def forward(
+        self,
+        x: torch.Tensor,
+        input_cp: bool = False,
+        latent_cp: bool = False,
+        output_cp: bool = False,
+        **additional_decode_kwargs,
+    ) -> Tuple[torch.Tensor, torch.Tensor, dict]:
+        z, reg_log = self.encode(x, return_reg_log=True, input_cp=input_cp, output_cp=latent_cp)
+        dec = self.decode(z, input_cp=latent_cp, output_cp=output_cp, **additional_decode_kwargs)
+        return z, dec, reg_log

sat/sgm/modules/__init__.py

@@ -0,0 +1,6 @@
+from .encoders.modules import GeneralConditioner
+
+UNCONDITIONAL_CONFIG = {
+    "target": "sgm.modules.GeneralConditioner",
+    "params": {"emb_models": []},
+}

sat/sgm/modules/attention.py

@@ -0,0 +1,572 @@
+import math
+from inspect import isfunction
+from typing import Any, Optional
+
+import torch
+import torch.nn.functional as F
+from einops import rearrange, repeat
+from packaging import version
+from torch import nn
+
+if version.parse(torch.__version__) >= version.parse("2.0.0"):
+    SDP_IS_AVAILABLE = True
+    from torch.backends.cuda import SDPBackend, sdp_kernel
+
+    BACKEND_MAP = {
+        SDPBackend.MATH: {
+            "enable_math": True,
+            "enable_flash": False,
+            "enable_mem_efficient": False,
+        },
+        SDPBackend.FLASH_ATTENTION: {
+            "enable_math": False,
+            "enable_flash": True,
+            "enable_mem_efficient": False,
+        },
+        SDPBackend.EFFICIENT_ATTENTION: {
+            "enable_math": False,
+            "enable_flash": False,
+            "enable_mem_efficient": True,
+        },
+        None: {"enable_math": True, "enable_flash": True, "enable_mem_efficient": True},
+    }
+else:
+    from contextlib import nullcontext
+
+    SDP_IS_AVAILABLE = False
+    sdp_kernel = nullcontext
+    BACKEND_MAP = {}
+    print(
+        f"No SDP backend available, likely because you are running in pytorch versions < 2.0. In fact, "
+        f"you are using PyTorch {torch.__version__}. You might want to consider upgrading."
+    )
+
+try:
+    import xformers
+    import xformers.ops
+
+    XFORMERS_IS_AVAILABLE = True
+except:
+    XFORMERS_IS_AVAILABLE = False
+    print("no module 'xformers'. Processing without...")
+
+from .diffusionmodules.util import checkpoint
+
+
+def exists(val):
+    return val is not None
+
+
+def uniq(arr):
+    return {el: True for el in arr}.keys()
+
+
+def default(val, d):
+    if exists(val):
+        return val
+    return d() if isfunction(d) else d
+
+
+def max_neg_value(t):
+    return -torch.finfo(t.dtype).max
+
+
+def init_(tensor):
+    dim = tensor.shape[-1]
+    std = 1 / math.sqrt(dim)
+    tensor.uniform_(-std, std)
+    return tensor
+
+
+# feedforward
+class GEGLU(nn.Module):
+    def __init__(self, dim_in, dim_out):
+        super().__init__()
+        self.proj = nn.Linear(dim_in, dim_out * 2)
+
+    def forward(self, x):
+        x, gate = self.proj(x).chunk(2, dim=-1)
+        return x * F.gelu(gate)
+
+
+class FeedForward(nn.Module):
+    def __init__(self, dim, dim_out=None, mult=4, glu=False, dropout=0.0):
+        super().__init__()
+        inner_dim = int(dim * mult)
+        dim_out = default(dim_out, dim)
+        project_in = nn.Sequential(nn.Linear(dim, inner_dim), nn.GELU()) if not glu else GEGLU(dim, inner_dim)
+
+        self.net = nn.Sequential(project_in, nn.Dropout(dropout), nn.Linear(inner_dim, dim_out))
+
+    def forward(self, x):
+        return self.net(x)
+
+
+def zero_module(module):
+    """
+    Zero out the parameters of a module and return it.
+    """
+    for p in module.parameters():
+        p.detach().zero_()
+    return module
+
+
+def Normalize(in_channels):
+    return torch.nn.GroupNorm(num_groups=32, num_channels=in_channels, eps=1e-6, affine=True)
+
+
+class LinearAttention(nn.Module):
+    def __init__(self, dim, heads=4, dim_head=32):
+        super().__init__()
+        self.heads = heads
+        hidden_dim = dim_head * heads
+        self.to_qkv = nn.Conv2d(dim, hidden_dim * 3, 1, bias=False)
+        self.to_out = nn.Conv2d(hidden_dim, dim, 1)
+
+    def forward(self, x):
+        b, c, h, w = x.shape
+        qkv = self.to_qkv(x)
+        q, k, v = rearrange(qkv, "b (qkv heads c) h w -> qkv b heads c (h w)", heads=self.heads, qkv=3)
+        k = k.softmax(dim=-1)
+        context = torch.einsum("bhdn,bhen->bhde", k, v)
+        out = torch.einsum("bhde,bhdn->bhen", context, q)
+        out = rearrange(out, "b heads c (h w) -> b (heads c) h w", heads=self.heads, h=h, w=w)
+        return self.to_out(out)
+
+
+class SpatialSelfAttention(nn.Module):
+    def __init__(self, in_channels):
+        super().__init__()
+        self.in_channels = in_channels
+
+        self.norm = Normalize(in_channels)
+        self.q = torch.nn.Conv2d(in_channels, in_channels, kernel_size=1, stride=1, padding=0)
+        self.k = torch.nn.Conv2d(in_channels, in_channels, kernel_size=1, stride=1, padding=0)
+        self.v = torch.nn.Conv2d(in_channels, in_channels, kernel_size=1, stride=1, padding=0)
+        self.proj_out = torch.nn.Conv2d(in_channels, in_channels, kernel_size=1, stride=1, padding=0)
+
+    def forward(self, x):
+        h_ = x
+        h_ = self.norm(h_)
+        q = self.q(h_)
+        k = self.k(h_)
+        v = self.v(h_)
+
+        # compute attention
+        b, c, h, w = q.shape
+        q = rearrange(q, "b c h w -> b (h w) c")
+        k = rearrange(k, "b c h w -> b c (h w)")
+        w_ = torch.einsum("bij,bjk->bik", q, k)
+
+        w_ = w_ * (int(c) ** (-0.5))
+        w_ = torch.nn.functional.softmax(w_, dim=2)
+
+        # attend to values
+        v = rearrange(v, "b c h w -> b c (h w)")
+        w_ = rearrange(w_, "b i j -> b j i")
+        h_ = torch.einsum("bij,bjk->bik", v, w_)
+        h_ = rearrange(h_, "b c (h w) -> b c h w", h=h)
+        h_ = self.proj_out(h_)
+
+        return x + h_
+
+
+class CrossAttention(nn.Module):
+    def __init__(
+        self,
+        query_dim,
+        context_dim=None,
+        heads=8,
+        dim_head=64,
+        dropout=0.0,
+        backend=None,
+    ):
+        super().__init__()
+        inner_dim = dim_head * heads
+        context_dim = default(context_dim, query_dim)
+
+        self.scale = dim_head**-0.5
+        self.heads = heads
+
+        self.to_q = nn.Linear(query_dim, inner_dim, bias=False)
+        self.to_k = nn.Linear(context_dim, inner_dim, bias=False)
+        self.to_v = nn.Linear(context_dim, inner_dim, bias=False)
+
+        self.to_out = nn.Sequential(nn.Linear(inner_dim, query_dim), nn.Dropout(dropout))
+        self.backend = backend
+
+    def forward(
+        self,
+        x,
+        context=None,
+        mask=None,
+        additional_tokens=None,
+        n_times_crossframe_attn_in_self=0,
+    ):
+        h = self.heads
+
+        if additional_tokens is not None:
+            # get the number of masked tokens at the beginning of the output sequence
+            n_tokens_to_mask = additional_tokens.shape[1]
+            # add additional token
+            x = torch.cat([additional_tokens, x], dim=1)
+
+        q = self.to_q(x)
+        context = default(context, x)
+        k = self.to_k(context)
+        v = self.to_v(context)
+
+        if n_times_crossframe_attn_in_self:
+            # reprogramming cross-frame attention as in https://arxiv.org/abs/2303.13439
+            assert x.shape[0] % n_times_crossframe_attn_in_self == 0
+            n_cp = x.shape[0] // n_times_crossframe_attn_in_self
+            k = repeat(k[::n_times_crossframe_attn_in_self], "b ... -> (b n) ...", n=n_cp)
+            v = repeat(v[::n_times_crossframe_attn_in_self], "b ... -> (b n) ...", n=n_cp)
+
+        q, k, v = map(lambda t: rearrange(t, "b n (h d) -> b h n d", h=h), (q, k, v))
+
+        ## old
+        """
+        sim = einsum('b i d, b j d -> b i j', q, k) * self.scale
+        del q, k
+
+        if exists(mask):
+            mask = rearrange(mask, 'b ... -> b (...)')
+            max_neg_value = -torch.finfo(sim.dtype).max
+            mask = repeat(mask, 'b j -> (b h) () j', h=h)
+            sim.masked_fill_(~mask, max_neg_value)
+
+        # attention, what we cannot get enough of
+        sim = sim.softmax(dim=-1)
+
+        out = einsum('b i j, b j d -> b i d', sim, v)
+        """
+        ## new
+        with sdp_kernel(**BACKEND_MAP[self.backend]):
+            # print("dispatching into backend", self.backend, "q/k/v shape: ", q.shape, k.shape, v.shape)
+            out = F.scaled_dot_product_attention(q, k, v, attn_mask=mask)  # scale is dim_head ** -0.5 per default
+
+        del q, k, v
+        out = rearrange(out, "b h n d -> b n (h d)", h=h)
+
+        if additional_tokens is not None:
+            # remove additional token
+            out = out[:, n_tokens_to_mask:]
+        return self.to_out(out)
+
+
+class MemoryEfficientCrossAttention(nn.Module):
+    # https://github.com/MatthieuTPHR/diffusers/blob/d80b531ff8060ec1ea982b65a1b8df70f73aa67c/src/diffusers/models/attention.py#L223
+    def __init__(self, query_dim, context_dim=None, heads=8, dim_head=64, dropout=0.0, **kwargs):
+        super().__init__()
+        print(
+            f"Setting up {self.__class__.__name__}. Query dim is {query_dim}, context_dim is {context_dim} and using "
+            f"{heads} heads with a dimension of {dim_head}."
+        )
+        inner_dim = dim_head * heads
+        context_dim = default(context_dim, query_dim)
+
+        self.heads = heads
+        self.dim_head = dim_head
+
+        self.to_q = nn.Linear(query_dim, inner_dim, bias=False)
+        self.to_k = nn.Linear(context_dim, inner_dim, bias=False)
+        self.to_v = nn.Linear(context_dim, inner_dim, bias=False)
+
+        self.to_out = nn.Sequential(nn.Linear(inner_dim, query_dim), nn.Dropout(dropout))
+        self.attention_op: Optional[Any] = None
+
+    def forward(
+        self,
+        x,
+        context=None,
+        mask=None,
+        additional_tokens=None,
+        n_times_crossframe_attn_in_self=0,
+    ):
+        if additional_tokens is not None:
+            # get the number of masked tokens at the beginning of the output sequence
+            n_tokens_to_mask = additional_tokens.shape[1]
+            # add additional token
+            x = torch.cat([additional_tokens, x], dim=1)
+        q = self.to_q(x)
+        context = default(context, x)
+        k = self.to_k(context)
+        v = self.to_v(context)
+
+        if n_times_crossframe_attn_in_self:
+            # reprogramming cross-frame attention as in https://arxiv.org/abs/2303.13439
+            assert x.shape[0] % n_times_crossframe_attn_in_self == 0
+            # n_cp = x.shape[0]//n_times_crossframe_attn_in_self
+            k = repeat(
+                k[::n_times_crossframe_attn_in_self],
+                "b ... -> (b n) ...",
+                n=n_times_crossframe_attn_in_self,
+            )
+            v = repeat(
+                v[::n_times_crossframe_attn_in_self],
+                "b ... -> (b n) ...",
+                n=n_times_crossframe_attn_in_self,
+            )
+
+        b, _, _ = q.shape
+        q, k, v = map(
+            lambda t: t.unsqueeze(3)
+            .reshape(b, t.shape[1], self.heads, self.dim_head)
+            .permute(0, 2, 1, 3)
+            .reshape(b * self.heads, t.shape[1], self.dim_head)
+            .contiguous(),
+            (q, k, v),
+        )
+
+        # actually compute the attention, what we cannot get enough of
+        out = xformers.ops.memory_efficient_attention(q, k, v, attn_bias=None, op=self.attention_op)
+
+        # TODO: Use this directly in the attention operation, as a bias
+        if exists(mask):
+            raise NotImplementedError
+        out = (
+            out.unsqueeze(0)
+            .reshape(b, self.heads, out.shape[1], self.dim_head)
+            .permute(0, 2, 1, 3)
+            .reshape(b, out.shape[1], self.heads * self.dim_head)
+        )
+        if additional_tokens is not None:
+            # remove additional token
+            out = out[:, n_tokens_to_mask:]
+        return self.to_out(out)
+
+
+class BasicTransformerBlock(nn.Module):
+    ATTENTION_MODES = {
+        "softmax": CrossAttention,  # vanilla attention
+        "softmax-xformers": MemoryEfficientCrossAttention,  # ampere
+    }
+
+    def __init__(
+        self,
+        dim,
+        n_heads,
+        d_head,
+        dropout=0.0,
+        context_dim=None,
+        gated_ff=True,
+        checkpoint=True,
+        disable_self_attn=False,
+        attn_mode="softmax",
+        sdp_backend=None,
+    ):
+        super().__init__()
+        assert attn_mode in self.ATTENTION_MODES
+        if attn_mode != "softmax" and not XFORMERS_IS_AVAILABLE:
+            print(
+                f"Attention mode '{attn_mode}' is not available. Falling back to native attention. "
+                f"This is not a problem in Pytorch >= 2.0. FYI, you are running with PyTorch version {torch.__version__}"
+            )
+            attn_mode = "softmax"
+        elif attn_mode == "softmax" and not SDP_IS_AVAILABLE:
+            print("We do not support vanilla attention anymore, as it is too expensive. Sorry.")
+            if not XFORMERS_IS_AVAILABLE:
+                assert False, "Please install xformers via e.g. 'pip install xformers==0.0.16'"
+            else:
+                print("Falling back to xformers efficient attention.")
+                attn_mode = "softmax-xformers"
+        attn_cls = self.ATTENTION_MODES[attn_mode]
+        if version.parse(torch.__version__) >= version.parse("2.0.0"):
+            assert sdp_backend is None or isinstance(sdp_backend, SDPBackend)
+        else:
+            assert sdp_backend is None
+        self.disable_self_attn = disable_self_attn
+        self.attn1 = attn_cls(
+            query_dim=dim,
+            heads=n_heads,
+            dim_head=d_head,
+            dropout=dropout,
+            context_dim=context_dim if self.disable_self_attn else None,
+            backend=sdp_backend,
+        )  # is a self-attention if not self.disable_self_attn
+        self.ff = FeedForward(dim, dropout=dropout, glu=gated_ff)
+        self.attn2 = attn_cls(
+            query_dim=dim,
+            context_dim=context_dim,
+            heads=n_heads,
+            dim_head=d_head,
+            dropout=dropout,
+            backend=sdp_backend,
+        )  # is self-attn if context is none
+        self.norm1 = nn.LayerNorm(dim)
+        self.norm2 = nn.LayerNorm(dim)
+        self.norm3 = nn.LayerNorm(dim)
+        self.checkpoint = checkpoint
+        if self.checkpoint:
+            print(f"{self.__class__.__name__} is using checkpointing")
+
+    def forward(self, x, context=None, additional_tokens=None, n_times_crossframe_attn_in_self=0):
+        kwargs = {"x": x}
+
+        if context is not None:
+            kwargs.update({"context": context})
+
+        if additional_tokens is not None:
+            kwargs.update({"additional_tokens": additional_tokens})
+
+        if n_times_crossframe_attn_in_self:
+            kwargs.update({"n_times_crossframe_attn_in_self": n_times_crossframe_attn_in_self})
+
+        # return mixed_checkpoint(self._forward, kwargs, self.parameters(), self.checkpoint)
+        return checkpoint(self._forward, (x, context), self.parameters(), self.checkpoint)
+
+    def _forward(self, x, context=None, additional_tokens=None, n_times_crossframe_attn_in_self=0):
+        x = (
+            self.attn1(
+                self.norm1(x),
+                context=context if self.disable_self_attn else None,
+                additional_tokens=additional_tokens,
+                n_times_crossframe_attn_in_self=n_times_crossframe_attn_in_self if not self.disable_self_attn else 0,
+            )
+            + x
+        )
+        x = self.attn2(self.norm2(x), context=context, additional_tokens=additional_tokens) + x
+        x = self.ff(self.norm3(x)) + x
+        return x
+
+
+class BasicTransformerSingleLayerBlock(nn.Module):
+    ATTENTION_MODES = {
+        "softmax": CrossAttention,  # vanilla attention
+        "softmax-xformers": MemoryEfficientCrossAttention,  # on the A100s not quite as fast as the above version
+        # (todo might depend on head_dim, check, falls back to semi-optimized kernels for dim!=[16,32,64,128])
+    }
+
+    def __init__(
+        self,
+        dim,
+        n_heads,
+        d_head,
+        dropout=0.0,
+        context_dim=None,
+        gated_ff=True,
+        checkpoint=True,
+        attn_mode="softmax",
+    ):
+        super().__init__()
+        assert attn_mode in self.ATTENTION_MODES
+        attn_cls = self.ATTENTION_MODES[attn_mode]
+        self.attn1 = attn_cls(
+            query_dim=dim,
+            heads=n_heads,
+            dim_head=d_head,
+            dropout=dropout,
+            context_dim=context_dim,
+        )
+        self.ff = FeedForward(dim, dropout=dropout, glu=gated_ff)
+        self.norm1 = nn.LayerNorm(dim)
+        self.norm2 = nn.LayerNorm(dim)
+        self.checkpoint = checkpoint
+
+    def forward(self, x, context=None):
+        return checkpoint(self._forward, (x, context), self.parameters(), self.checkpoint)
+
+    def _forward(self, x, context=None):
+        x = self.attn1(self.norm1(x), context=context) + x
+        x = self.ff(self.norm2(x)) + x
+        return x
+
+
+class SpatialTransformer(nn.Module):
+    """
+    Transformer block for image-like data.
+    First, project the input (aka embedding)
+    and reshape to b, t, d.
+    Then apply standard transformer action.
+    Finally, reshape to image
+    NEW: use_linear for more efficiency instead of the 1x1 convs
+    """
+
+    def __init__(
+        self,
+        in_channels,
+        n_heads,
+        d_head,
+        depth=1,
+        dropout=0.0,
+        context_dim=None,
+        disable_self_attn=False,
+        use_linear=False,
+        attn_type="softmax",
+        use_checkpoint=True,
+        # sdp_backend=SDPBackend.FLASH_ATTENTION
+        sdp_backend=None,
+    ):
+        super().__init__()
+        print(f"constructing {self.__class__.__name__} of depth {depth} w/ {in_channels} channels and {n_heads} heads")
+        from omegaconf import ListConfig
+
+        if exists(context_dim) and not isinstance(context_dim, (list, ListConfig)):
+            context_dim = [context_dim]
+        if exists(context_dim) and isinstance(context_dim, list):
+            if depth != len(context_dim):
+                print(
+                    f"WARNING: {self.__class__.__name__}: Found context dims {context_dim} of depth {len(context_dim)}, "
+                    f"which does not match the specified 'depth' of {depth}. Setting context_dim to {depth * [context_dim[0]]} now."
+                )
+                # depth does not match context dims.
+                assert all(
+                    map(lambda x: x == context_dim[0], context_dim)
+                ), "need homogenous context_dim to match depth automatically"
+                context_dim = depth * [context_dim[0]]
+        elif context_dim is None:
+            context_dim = [None] * depth
+        self.in_channels = in_channels
+        inner_dim = n_heads * d_head
+        self.norm = Normalize(in_channels)
+        if not use_linear:
+            self.proj_in = nn.Conv2d(in_channels, inner_dim, kernel_size=1, stride=1, padding=0)
+        else:
+            self.proj_in = nn.Linear(in_channels, inner_dim)
+
+        self.transformer_blocks = nn.ModuleList(
+            [
+                BasicTransformerBlock(
+                    inner_dim,
+                    n_heads,
+                    d_head,
+                    dropout=dropout,
+                    context_dim=context_dim[d],
+                    disable_self_attn=disable_self_attn,
+                    attn_mode=attn_type,
+                    checkpoint=use_checkpoint,
+                    sdp_backend=sdp_backend,
+                )
+                for d in range(depth)
+            ]
+        )
+        if not use_linear:
+            self.proj_out = zero_module(nn.Conv2d(inner_dim, in_channels, kernel_size=1, stride=1, padding=0))
+        else:
+            # self.proj_out = zero_module(nn.Linear(in_channels, inner_dim))
+            self.proj_out = zero_module(nn.Linear(inner_dim, in_channels))
+        self.use_linear = use_linear
+
+    def forward(self, x, context=None):
+        # note: if no context is given, cross-attention defaults to self-attention
+        if not isinstance(context, list):
+            context = [context]
+        b, c, h, w = x.shape
+        x_in = x
+        x = self.norm(x)
+        if not self.use_linear:
+            x = self.proj_in(x)
+        x = rearrange(x, "b c h w -> b (h w) c").contiguous()
+        if self.use_linear:
+            x = self.proj_in(x)
+        for i, block in enumerate(self.transformer_blocks):
+            if i > 0 and len(context) == 1:
+                i = 0  # use same context for each block
+            x = block(x, context=context[i])
+        if self.use_linear:
+            x = self.proj_out(x)
+        x = rearrange(x, "b (h w) c -> b c h w", h=h, w=w).contiguous()
+        if not self.use_linear:
+            x = self.proj_out(x)
+        return x + x_in

sat/sgm/modules/autoencoding/losses/__init__.py

@@ -0,0 +1,8 @@
+__all__ = [
+    "GeneralLPIPSWithDiscriminator",
+    "LatentLPIPS",
+]
+
+from .discriminator_loss import GeneralLPIPSWithDiscriminator
+from .lpips import LatentLPIPS
+from .video_loss import VideoAutoencoderLoss

sat/sgm/modules/autoencoding/losses/discriminator_loss.py

@@ -0,0 +1,301 @@
+from typing import Dict, Iterator, List, Optional, Tuple, Union
+
+import numpy as np
+import torch
+import torch.nn as nn
+import torchvision
+from einops import rearrange
+from matplotlib import colormaps
+from matplotlib import pyplot as plt
+
+from ....util import default, instantiate_from_config
+from ..lpips.loss.lpips import LPIPS
+from ..lpips.model.model import weights_init
+from ..lpips.vqperceptual import hinge_d_loss, vanilla_d_loss
+
+
+class GeneralLPIPSWithDiscriminator(nn.Module):
+    def __init__(
+        self,
+        disc_start: int,
+        logvar_init: float = 0.0,
+        disc_num_layers: int = 3,
+        disc_in_channels: int = 3,
+        disc_factor: float = 1.0,
+        disc_weight: float = 1.0,
+        perceptual_weight: float = 1.0,
+        disc_loss: str = "hinge",
+        scale_input_to_tgt_size: bool = False,
+        dims: int = 2,
+        learn_logvar: bool = False,
+        regularization_weights: Union[None, Dict[str, float]] = None,
+        additional_log_keys: Optional[List[str]] = None,
+        discriminator_config: Optional[Dict] = None,
+    ):
+        super().__init__()
+        self.dims = dims
+        if self.dims > 2:
+            print(
+                f"running with dims={dims}. This means that for perceptual loss "
+                f"calculation, the LPIPS loss will be applied to each frame "
+                f"independently."
+            )
+        self.scale_input_to_tgt_size = scale_input_to_tgt_size
+        assert disc_loss in ["hinge", "vanilla"]
+        self.perceptual_loss = LPIPS().eval()
+        self.perceptual_weight = perceptual_weight
+        # output log variance
+        self.logvar = nn.Parameter(torch.full((), logvar_init), requires_grad=learn_logvar)
+        self.learn_logvar = learn_logvar
+
+        discriminator_config = default(
+            discriminator_config,
+            {
+                "target": "sgm.modules.autoencoding.lpips.model.model.NLayerDiscriminator",
+                "params": {
+                    "input_nc": disc_in_channels,
+                    "n_layers": disc_num_layers,
+                    "use_actnorm": False,
+                },
+            },
+        )
+
+        self.discriminator = instantiate_from_config(discriminator_config).apply(weights_init)
+        self.discriminator_iter_start = disc_start
+        self.disc_loss = hinge_d_loss if disc_loss == "hinge" else vanilla_d_loss
+        self.disc_factor = disc_factor
+        self.discriminator_weight = disc_weight
+        self.regularization_weights = default(regularization_weights, {})
+
+        self.forward_keys = [
+            "optimizer_idx",
+            "global_step",
+            "last_layer",
+            "split",
+            "regularization_log",
+        ]
+
+        self.additional_log_keys = set(default(additional_log_keys, []))
+        self.additional_log_keys.update(set(self.regularization_weights.keys()))
+
+    def get_trainable_parameters(self) -> Iterator[nn.Parameter]:
+        return self.discriminator.parameters()
+
+    def get_trainable_autoencoder_parameters(self) -> Iterator[nn.Parameter]:
+        if self.learn_logvar:
+            yield self.logvar
+        yield from ()
+
+    @torch.no_grad()
+    def log_images(self, inputs: torch.Tensor, reconstructions: torch.Tensor) -> Dict[str, torch.Tensor]:
+        # calc logits of real/fake
+        logits_real = self.discriminator(inputs.contiguous().detach())
+        if len(logits_real.shape) < 4:
+            # Non patch-discriminator
+            return dict()
+        logits_fake = self.discriminator(reconstructions.contiguous().detach())
+        # -> (b, 1, h, w)
+
+        # parameters for colormapping
+        high = max(logits_fake.abs().max(), logits_real.abs().max()).item()
+        cmap = colormaps["PiYG"]  # diverging colormap
+
+        def to_colormap(logits: torch.Tensor) -> torch.Tensor:
+            """(b, 1, ...) -> (b, 3, ...)"""
+            logits = (logits + high) / (2 * high)
+            logits_np = cmap(logits.cpu().numpy())[..., :3]  # truncate alpha channel
+            # -> (b, 1, ..., 3)
+            logits = torch.from_numpy(logits_np).to(logits.device)
+            return rearrange(logits, "b 1 ... c -> b c ...")
+
+        logits_real = torch.nn.functional.interpolate(
+            logits_real,
+            size=inputs.shape[-2:],
+            mode="nearest",
+            antialias=False,
+        )
+        logits_fake = torch.nn.functional.interpolate(
+            logits_fake,
+            size=reconstructions.shape[-2:],
+            mode="nearest",
+            antialias=False,
+        )
+
+        # alpha value of logits for overlay
+        alpha_real = torch.abs(logits_real) / high
+        alpha_fake = torch.abs(logits_fake) / high
+        # -> (b, 1, h, w) in range [0, 0.5]
+        # alpha value of lines don't really matter, since the values are the same
+        # for both images and logits anyway
+        grid_alpha_real = torchvision.utils.make_grid(alpha_real, nrow=4)
+        grid_alpha_fake = torchvision.utils.make_grid(alpha_fake, nrow=4)
+        grid_alpha = 0.8 * torch.cat((grid_alpha_real, grid_alpha_fake), dim=1)
+        # -> (1, h, w)
+        # blend logits and images together
+
+        # prepare logits for plotting
+        logits_real = to_colormap(logits_real)
+        logits_fake = to_colormap(logits_fake)
+        # resize logits
+        # -> (b, 3, h, w)
+
+        # make some grids
+        # add all logits to one plot
+        logits_real = torchvision.utils.make_grid(logits_real, nrow=4)
+        logits_fake = torchvision.utils.make_grid(logits_fake, nrow=4)
+        # I just love how torchvision calls the number of columns `nrow`
+        grid_logits = torch.cat((logits_real, logits_fake), dim=1)
+        # -> (3, h, w)
+
+        grid_images_real = torchvision.utils.make_grid(0.5 * inputs + 0.5, nrow=4)
+        grid_images_fake = torchvision.utils.make_grid(0.5 * reconstructions + 0.5, nrow=4)
+        grid_images = torch.cat((grid_images_real, grid_images_fake), dim=1)
+        # -> (3, h, w) in range [0, 1]
+
+        grid_blend = grid_alpha * grid_logits + (1 - grid_alpha) * grid_images
+
+        # Create labeled colorbar
+        dpi = 100
+        height = 128 / dpi
+        width = grid_logits.shape[2] / dpi
+        fig, ax = plt.subplots(figsize=(width, height), dpi=dpi)
+        img = ax.imshow(np.array([[-high, high]]), cmap=cmap)
+        plt.colorbar(
+            img,
+            cax=ax,
+            orientation="horizontal",
+            fraction=0.9,
+            aspect=width / height,
+            pad=0.0,
+        )
+        img.set_visible(False)
+        fig.tight_layout()
+        fig.canvas.draw()
+        # manually convert figure to numpy
+        cbar_np = np.frombuffer(fig.canvas.tostring_rgb(), dtype=np.uint8)
+        cbar_np = cbar_np.reshape(fig.canvas.get_width_height()[::-1] + (3,))
+        cbar = torch.from_numpy(cbar_np.copy()).to(grid_logits.dtype) / 255.0
+        cbar = rearrange(cbar, "h w c -> c h w").to(grid_logits.device)
+
+        # Add colorbar to plot
+        annotated_grid = torch.cat((grid_logits, cbar), dim=1)
+        blended_grid = torch.cat((grid_blend, cbar), dim=1)
+        return {
+            "vis_logits": 2 * annotated_grid[None, ...] - 1,
+            "vis_logits_blended": 2 * blended_grid[None, ...] - 1,
+        }
+
+    def calculate_adaptive_weight(
+        self, nll_loss: torch.Tensor, g_loss: torch.Tensor, last_layer: torch.Tensor
+    ) -> torch.Tensor:
+        nll_grads = torch.autograd.grad(nll_loss, last_layer, retain_graph=True)[0]
+        g_grads = torch.autograd.grad(g_loss, last_layer, retain_graph=True)[0]
+
+        d_weight = torch.norm(nll_grads) / (torch.norm(g_grads) + 1e-4)
+        d_weight = torch.clamp(d_weight, 0.0, 1e4).detach()
+        d_weight = d_weight * self.discriminator_weight
+        return d_weight
+
+    def forward(
+        self,
+        inputs: torch.Tensor,
+        reconstructions: torch.Tensor,
+        *,  # added because I changed the order here
+        regularization_log: Dict[str, torch.Tensor],
+        optimizer_idx: int,
+        global_step: int,
+        last_layer: torch.Tensor,
+        split: str = "train",
+        weights: Union[None, float, torch.Tensor] = None,
+    ) -> Tuple[torch.Tensor, dict]:
+        if self.scale_input_to_tgt_size:
+            inputs = torch.nn.functional.interpolate(inputs, reconstructions.shape[2:], mode="bicubic", antialias=True)
+
+        if self.dims > 2:
+            inputs, reconstructions = map(
+                lambda x: rearrange(x, "b c t h w -> (b t) c h w"),
+                (inputs, reconstructions),
+            )
+
+        rec_loss = torch.abs(inputs.contiguous() - reconstructions.contiguous())
+        if self.perceptual_weight > 0:
+            frame_indices = torch.randn((inputs.shape[0], inputs.shape[2])).topk(1, dim=-1).indices
+
+            from sgm.modules.autoencoding.losses.video_loss import pick_video_frame
+
+            input_frames = pick_video_frame(inputs, frame_indices)
+            recon_frames = pick_video_frame(reconstructions, frame_indices)
+
+            p_loss = self.perceptual_loss(input_frames.contiguous(), recon_frames.contiguous()).mean()
+            rec_loss = rec_loss + self.perceptual_weight * p_loss
+
+        nll_loss, weighted_nll_loss = self.get_nll_loss(rec_loss, weights)
+
+        # now the GAN part
+        if optimizer_idx == 0:
+            # generator update
+            if global_step >= self.discriminator_iter_start or not self.training:
+                logits_fake = self.discriminator(reconstructions.contiguous())
+                g_loss = -torch.mean(logits_fake)
+                if self.training:
+                    d_weight = self.calculate_adaptive_weight(nll_loss, g_loss, last_layer=last_layer)
+                else:
+                    d_weight = torch.tensor(1.0)
+            else:
+                d_weight = torch.tensor(0.0)
+                g_loss = torch.tensor(0.0, requires_grad=True)
+
+            loss = weighted_nll_loss + d_weight * self.disc_factor * g_loss
+            log = dict()
+            for k in regularization_log:
+                if k in self.regularization_weights:
+                    loss = loss + self.regularization_weights[k] * regularization_log[k]
+                if k in self.additional_log_keys:
+                    log[f"{split}/{k}"] = regularization_log[k].detach().float().mean()
+
+            log.update(
+                {
+                    f"{split}/loss/total": loss.clone().detach().mean(),
+                    f"{split}/loss/nll": nll_loss.detach().mean(),
+                    f"{split}/loss/rec": rec_loss.detach().mean(),
+                    f"{split}/loss/percep": p_loss.detach().mean(),
+                    f"{split}/loss/rec": rec_loss.detach().mean(),
+                    f"{split}/loss/g": g_loss.detach().mean(),
+                    f"{split}/scalars/logvar": self.logvar.detach(),
+                    f"{split}/scalars/d_weight": d_weight.detach(),
+                }
+            )
+
+            return loss, log
+        elif optimizer_idx == 1:
+            # second pass for discriminator update
+            logits_real = self.discriminator(inputs.contiguous().detach())
+            logits_fake = self.discriminator(reconstructions.contiguous().detach())
+
+            if global_step >= self.discriminator_iter_start or not self.training:
+                d_loss = self.disc_factor * self.disc_loss(logits_real, logits_fake)
+            else:
+                d_loss = torch.tensor(0.0, requires_grad=True)
+
+            log = {
+                f"{split}/loss/disc": d_loss.clone().detach().mean(),
+                f"{split}/logits/real": logits_real.detach().mean(),
+                f"{split}/logits/fake": logits_fake.detach().mean(),
+            }
+            return d_loss, log
+        else:
+            raise NotImplementedError(f"Unknown optimizer_idx {optimizer_idx}")
+
+    def get_nll_loss(
+        self,
+        rec_loss: torch.Tensor,
+        weights: Optional[Union[float, torch.Tensor]] = None,
+    ) -> Tuple[torch.Tensor, torch.Tensor]:
+        nll_loss = rec_loss / torch.exp(self.logvar) + self.logvar
+        weighted_nll_loss = nll_loss
+        if weights is not None:
+            weighted_nll_loss = weights * nll_loss
+        weighted_nll_loss = torch.sum(weighted_nll_loss) / weighted_nll_loss.shape[0]
+        nll_loss = torch.sum(nll_loss) / nll_loss.shape[0]
+
+        return nll_loss, weighted_nll_loss

sat/sgm/modules/autoencoding/losses/lpips.py

@@ -0,0 +1,64 @@
+import torch
+import torch.nn as nn
+
+from ....util import default, instantiate_from_config
+from ..lpips.loss.lpips import LPIPS
+
+
+class LatentLPIPS(nn.Module):
+    def __init__(
+        self,
+        decoder_config,
+        perceptual_weight=1.0,
+        latent_weight=1.0,
+        scale_input_to_tgt_size=False,
+        scale_tgt_to_input_size=False,
+        perceptual_weight_on_inputs=0.0,
+    ):
+        super().__init__()
+        self.scale_input_to_tgt_size = scale_input_to_tgt_size
+        self.scale_tgt_to_input_size = scale_tgt_to_input_size
+        self.init_decoder(decoder_config)
+        self.perceptual_loss = LPIPS().eval()
+        self.perceptual_weight = perceptual_weight
+        self.latent_weight = latent_weight
+        self.perceptual_weight_on_inputs = perceptual_weight_on_inputs
+
+    def init_decoder(self, config):
+        self.decoder = instantiate_from_config(config)
+        if hasattr(self.decoder, "encoder"):
+            del self.decoder.encoder
+
+    def forward(self, latent_inputs, latent_predictions, image_inputs, split="train"):
+        log = dict()
+        loss = (latent_inputs - latent_predictions) ** 2
+        log[f"{split}/latent_l2_loss"] = loss.mean().detach()
+        image_reconstructions = None
+        if self.perceptual_weight > 0.0:
+            image_reconstructions = self.decoder.decode(latent_predictions)
+            image_targets = self.decoder.decode(latent_inputs)
+            perceptual_loss = self.perceptual_loss(image_targets.contiguous(), image_reconstructions.contiguous())
+            loss = self.latent_weight * loss.mean() + self.perceptual_weight * perceptual_loss.mean()
+            log[f"{split}/perceptual_loss"] = perceptual_loss.mean().detach()
+
+        if self.perceptual_weight_on_inputs > 0.0:
+            image_reconstructions = default(image_reconstructions, self.decoder.decode(latent_predictions))
+            if self.scale_input_to_tgt_size:
+                image_inputs = torch.nn.functional.interpolate(
+                    image_inputs,
+                    image_reconstructions.shape[2:],
+                    mode="bicubic",
+                    antialias=True,
+                )
+            elif self.scale_tgt_to_input_size:
+                image_reconstructions = torch.nn.functional.interpolate(
+                    image_reconstructions,
+                    image_inputs.shape[2:],
+                    mode="bicubic",
+                    antialias=True,
+                )
+
+            perceptual_loss2 = self.perceptual_loss(image_inputs.contiguous(), image_reconstructions.contiguous())
+            loss = loss + self.perceptual_weight_on_inputs * perceptual_loss2.mean()
+            log[f"{split}/perceptual_loss_on_inputs"] = perceptual_loss2.mean().detach()
+        return loss, log