[ { "path": "CogVideo/.github/ISSUE_TEMPLATE/bug_report.yaml", "content": "name: \"\\U0001F41B Bug Report\"\ndescription: Submit a bug report to help us improve CogVideoX / 提交一个 Bug 问题报告来帮助我们改进 CogVideoX 开源模型\nbody:\n - type: textarea\n id: system-info\n attributes:\n label: System Info / 系統信息\n description: Your operating environment / 您的运行环境信息\n placeholder: Includes Cuda version, Diffusers version, Python version, operating system, hardware information (if you suspect a hardware problem)... / 包括Cuda版本，Diffusers，Python版本，操作系统，硬件信息(如果您怀疑是硬件方面的问题)...\n validations:\n required: true\n\n - type: checkboxes\n id: information-scripts-examples\n attributes:\n label: Information / 问题信息\n description: 'The problem arises when using: / 问题出现在'\n options:\n - label: \"The official example scripts / 官方的示例脚本\"\n - label: \"My own modified scripts / 我自己修改的脚本和任务\"\n\n - type: textarea\n id: reproduction\n validations:\n required: true\n attributes:\n label: Reproduction / 复现过程\n description: |\n Please provide a code example that reproduces the problem you encountered, preferably with a minimal reproduction unit.\n If you have code snippets, error messages, stack traces, please provide them here as well.\n 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\n Do not use screenshots, as they are difficult to read and (more importantly) do not allow others to copy and paste your code.\n \n 请提供能重现您遇到的问题的代码示例,最好是最小复现单元。\n 如果您有代码片段、错误信息、堆栈跟踪，也请在此提供。\n 请使用代码标签正确格式化您的代码。请参见 https://help.github.com/en/github/writing-on-github/creating-and-highlighting-code-blocks#syntax-highlighting\n 请勿使用截图，因为截图难以阅读，而且（更重要的是）不允许他人复制粘贴您的代码。\n placeholder: |\n Steps to reproduce the behavior/复现Bug的步骤:\n \n 1.\n 2.\n 3.\n\n - type: textarea\n id: expected-behavior\n validations:\n required: true\n attributes:\n label: Expected behavior / 期待表现\n description: \"A clear and concise description of what you would expect to happen. /简单描述您期望发生的事情。\"" }, { "path": "CogVideo/.github/ISSUE_TEMPLATE/feature-request.yaml", "content": "name: \"\\U0001F680 Feature request\"\ndescription: Submit a request for a new CogVideoX feature / 提交一个新的 CogVideoX开源模型的功能建议\nlabels: [ \"feature\" ]\nbody:\n - type: textarea\n id: feature-request\n validations:\n required: true\n attributes:\n label: Feature request / 功能建议\n description: |\n A brief description of the functional proposal. Links to corresponding papers and code are desirable.\n 对功能建议的简述。最好提供对应的论文和代码链接。\n\n - type: textarea\n id: motivation\n validations:\n required: true\n attributes:\n label: Motivation / 动机\n description: |\n Your motivation for making the suggestion. If that motivation is related to another GitHub issue, link to it here.\n 您提出建议的动机。如果该动机与另一个 GitHub 问题有关，请在此处提供对应的链接。\n\n - type: textarea\n id: contribution\n validations:\n required: true\n attributes:\n label: Your contribution / 您的贡献\n description: |\n \n Your PR link or any other link you can help with.\n 您的PR链接或者其他您能提供帮助的链接。" }, { "path": "CogVideo/.github/PULL_REQUEST_TEMPLATE/pr_template.md", "content": "# Raise valuable PR / 提出有价值的PR\n\n## Caution / 注意事项:\nUsers should keep the following points in mind when submitting PRs:\n\n1. Ensure that your code meets the requirements in the [specification](../../resources/contribute.md).\n2. the proposed PR should be relevant, if there are multiple ideas and optimizations, they should be assigned to different PRs.\n\n用户在提交PR时候应该注意以下几点:\n\n1. 确保您的代码符合 [规范](../../resources/contribute_zh.md) 中的要求。\n2. 提出的PR应该具有针对性，如果具有多个不同的想法和优化方案，应该分配到不同的PR中。\n\n## 不应该提出的PR / PRs that should not be proposed\n\nIf a developer proposes a PR about any of the following, it may be closed or Rejected.\n\n1. those that don't describe improvement options.\n2. multiple issues of different types combined in one PR.\n3. The proposed PR is highly duplicative of already existing PRs.\n\n如果开发者提出关于以下方面的PR，则可能会被直接关闭或拒绝通过。\n\n1. 没有说明改进方案的。\n2. 多个不同类型的问题合并在一个PR中的。\n3. 提出的PR与已经存在的PR高度重复的。\n\n\n# 检查您的PR\n- [ ] Have you read the Contributor Guidelines, Pull Request section? / 您是否阅读了贡献者指南、Pull Request 部分？\n- [ ] Has this been discussed/approved via a Github issue or forum? If so, add a link. / 是否通过 Github 问题或论坛讨论/批准过？如果是，请添加链接。\n- [ ] Did you make sure you updated the documentation with your changes? Here are the Documentation Guidelines, and here are the Documentation Formatting Tips. /您是否确保根据您的更改更新了文档？这里是文档指南，这里是文档格式化技巧。\n- [ ] Did you write new required tests? / 您是否编写了新的必要测试？\n- [ ] Are your PRs for only one issue / 您的PR是否仅针对一个问题" }, { "path": "CogVideo/.gitignore", "content": "*__pycache__/\nsamples*/\nruns/\ncheckpoints/\nmaster_ip\nlogs/\n*.DS_Store\n.idea\noutput*\ntest*" }, { "path": "CogVideo/LICENSE", "content": " Apache License\n Version 2.0, January 2004\n http://www.apache.org/licenses/\n\n TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION\n\n 1. Definitions.\n\n \"License\" shall mean the terms and conditions for use, reproduction,\n and distribution as defined by Sections 1 through 9 of this document.\n\n \"Licensor\" shall mean the copyright owner or entity authorized by\n the copyright owner that is granting the License.\n\n \"Legal Entity\" shall mean the union of the acting entity and all\n other entities that control, are controlled by, or are under common\n control with that entity. For the purposes of this definition,\n \"control\" means (i) the power, direct or indirect, to cause the\n direction or management of such entity, whether by contract or\n otherwise, or (ii) ownership of fifty percent (50%) or more of the\n outstanding shares, or (iii) beneficial ownership of such entity.\n\n \"You\" (or \"Your\") shall mean an individual or Legal Entity\n exercising permissions granted by this License.\n\n \"Source\" form shall mean the preferred form for making modifications,\n including but not limited to software source code, documentation\n source, and configuration files.\n\n \"Object\" form shall mean any form resulting from mechanical\n transformation or translation of a Source form, including but\n not limited to compiled object code, generated documentation,\n and conversions to other media types.\n\n \"Work\" shall mean the work of authorship, whether in Source or\n Object form, made available under the License, as indicated by a\n copyright notice that is included in or attached to the work\n (an example is provided in the Appendix below).\n\n \"Derivative Works\" shall mean any work, whether in Source or Object\n form, that is based on (or derived from) the Work and for which the\n editorial revisions, annotations, elaborations, or other modifications\n represent, as a whole, an original work of authorship. For the purposes\n of this License, Derivative Works shall not include works that remain\n separable from, or merely link (or bind by name) to the interfaces of,\n the Work and Derivative Works thereof.\n\n \"Contribution\" shall mean any work of authorship, including\n the original version of the Work and any modifications or additions\n to that Work or Derivative Works thereof, that is intentionally\n submitted to Licensor for inclusion in the Work by the copyright owner\n or by an individual or Legal Entity authorized to submit on behalf of\n the copyright owner. For the purposes of this definition, \"submitted\"\n means any form of electronic, verbal, or written communication sent\n to the Licensor or its representatives, including but not limited to\n communication on electronic mailing lists, source code control systems,\n and issue tracking systems that are managed by, or on behalf of, the\n Licensor for the purpose of discussing and improving the Work, but\n excluding communication that is conspicuously marked or otherwise\n designated in writing by the copyright owner as \"Not a Contribution.\"\n\n \"Contributor\" shall mean Licensor and any individual or Legal Entity\n on behalf of whom a Contribution has been received by Licensor and\n subsequently incorporated within the Work.\n\n 2. Grant of Copyright License. Subject to the terms and conditions of\n this License, each Contributor hereby grants to You a perpetual,\n worldwide, non-exclusive, no-charge, royalty-free, irrevocable\n copyright license to reproduce, prepare Derivative Works of,\n publicly display, publicly perform, sublicense, and distribute the\n Work and such Derivative Works in Source or Object form.\n\n 3. Grant of Patent License. Subject to the terms and conditions of\n this License, each Contributor hereby grants to You a perpetual,\n worldwide, non-exclusive, no-charge, royalty-free, irrevocable\n (except as stated in this section) patent license to make, have made,\n use, offer to sell, sell, import, and otherwise transfer the Work,\n where such license applies only to those patent claims licensable\n by such Contributor that are necessarily infringed by their\n Contribution(s) alone or by combination of their Contribution(s)\n with the Work to which such Contribution(s) was submitted. If You\n institute patent litigation against any entity (including a\n cross-claim or counterclaim in a lawsuit) alleging that the Work\n or a Contribution incorporated within the Work constitutes direct\n or contributory patent infringement, then any patent licenses\n granted to You under this License for that Work shall terminate\n as of the date such litigation is filed.\n\n 4. Redistribution. You may reproduce and distribute copies of the\n Work or Derivative Works thereof in any medium, with or without\n modifications, and in Source or Object form, provided that You\n meet the following conditions:\n\n (a) You must give any other recipients of the Work or\n Derivative Works a copy of this License; and\n\n (b) You must cause any modified files to carry prominent notices\n stating that You changed the files; and\n\n (c) You must retain, in the Source form of any Derivative Works\n that You distribute, all copyright, patent, trademark, and\n attribution notices from the Source form of the Work,\n excluding those notices that do not pertain to any part of\n the Derivative Works; and\n\n (d) If the Work includes a \"NOTICE\" text file as part of its\n distribution, then any Derivative Works that You distribute must\n include a readable copy of the attribution notices contained\n within such NOTICE file, excluding those notices that do not\n pertain to any part of the Derivative Works, in at least one\n of the following places: within a NOTICE text file distributed\n as part of the Derivative Works; within the Source form or\n documentation, if provided along with the Derivative Works; or,\n within a display generated by the Derivative Works, if and\n wherever such third-party notices normally appear. The contents\n of the NOTICE file are for informational purposes only and\n do not modify the License. You may add Your own attribution\n notices within Derivative Works that You distribute, alongside\n or as an addendum to the NOTICE text from the Work, provided\n that such additional attribution notices cannot be construed\n as modifying the License.\n\n You may add Your own copyright statement to Your modifications and\n may provide additional or different license terms and conditions\n for use, reproduction, or distribution of Your modifications, or\n for any such Derivative Works as a whole, provided Your use,\n reproduction, and distribution of the Work otherwise complies with\n the conditions stated in this License.\n\n 5. Submission of Contributions. Unless You explicitly state otherwise,\n any Contribution intentionally submitted for inclusion in the Work\n by You to the Licensor shall be under the terms and conditions of\n this License, without any additional terms or conditions.\n Notwithstanding the above, nothing herein shall supersede or modify\n the terms of any separate license agreement you may have executed\n with Licensor regarding such Contributions.\n\n 6. Trademarks. This License does not grant permission to use the trade\n names, trademarks, service marks, or product names of the Licensor,\n except as required for reasonable and customary use in describing the\n origin of the Work and reproducing the content of the NOTICE file.\n\n 7. Disclaimer of Warranty. Unless required by applicable law or\n agreed to in writing, Licensor provides the Work (and each\n Contributor provides its Contributions) on an \"AS IS\" BASIS,\n WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or\n implied, including, without limitation, any warranties or conditions\n of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A\n PARTICULAR PURPOSE. You are solely responsible for determining the\n appropriateness of using or redistributing the Work and assume any\n risks associated with Your exercise of permissions under this License.\n\n 8. Limitation of Liability. In no event and under no legal theory,\n whether in tort (including negligence), contract, or otherwise,\n unless required by applicable law (such as deliberate and grossly\n negligent acts) or agreed to in writing, shall any Contributor be\n liable to You for damages, including any direct, indirect, special,\n incidental, or consequential damages of any character arising as a\n result of this License or out of the use or inability to use the\n Work (including but not limited to damages for loss of goodwill,\n work stoppage, computer failure or malfunction, or any and all\n other commercial damages or losses), even if such Contributor\n has been advised of the possibility of such damages.\n\n 9. Accepting Warranty or Additional Liability. While redistributing\n the Work or Derivative Works thereof, You may choose to offer,\n and charge a fee for, acceptance of support, warranty, indemnity,\n or other liability obligations and/or rights consistent with this\n License. However, in accepting such obligations, You may act only\n on Your own behalf and on Your sole responsibility, not on behalf\n of any other Contributor, and only if You agree to indemnify,\n defend, and hold each Contributor harmless for any liability\n incurred by, or claims asserted against, such Contributor by reason\n of your accepting any such warranty or additional liability.\n\n END OF TERMS AND CONDITIONS\n\n APPENDIX: How to apply the Apache License to your work.\n\n To apply the Apache License to your work, attach the following\n boilerplate notice, with the fields enclosed by brackets \"[]\"\n replaced with your own identifying information. (Don't include\n the brackets!) The text should be enclosed in the appropriate\n comment syntax for the file format. We also recommend that a\n file or class name and description of purpose be included on the\n same \"printed page\" as the copyright notice for easier\n identification within third-party archives.\n\n Copyright 2024 CogVideo Model Team @ Zhipu AI\n\n Licensed under the Apache License, Version 2.0 (the \"License\");\n you may not use this file except in compliance with the License.\n You may obtain a copy of the License at\n\n http://www.apache.org/licenses/LICENSE-2.0\n\n Unless required by applicable law or agreed to in writing, software\n distributed under the License is distributed on an \"AS IS\" BASIS,\n WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n See the License for the specific language governing permissions and\n limitations under the License.\n" }, { "path": "CogVideo/MODEL_LICENSE", "content": "The CogVideoX License\n\n1. Definitions\n\n“Licensor” means the CogVideoX Model Team that distributes its Software.\n\n“Software” means the CogVideoX model parameters made available under this license.\n\n2. License Grant\n\nUnder 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.\nThis 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 .\nUsers 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.\nIf 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.\nThe above copyright statement and this license statement should be included in all copies or significant portions of this software.\n\n3. Restriction\n\nYou 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.\n\nYou 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.\n\n4. Disclaimer\n\nTHE 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.\n\n5. Limitation of Liability\n\nEXCEPT 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.\n\n6. Dispute Resolution\n\nThis 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.\n\nNote 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.\n\n1. 定义\n\n“许可方”是指分发其软件的 CogVideoX 模型团队。\n\n“软件”是指根据本许可提供的 CogVideoX 模型参数。\n\n2. 许可授予\n\n根据本许可的条款和条件，许可方特此授予您非排他性、全球性、不可转让、不可再许可、可撤销、免版税的版权许可。生成内容的知识产权所属，可根据适用当地法律的规定，在法律允许的范围内由用户享有生成内容的知识产权或其他权利。\n本许可允许您免费使用本仓库中的所有开源模型进行学术研究。对于希望将模型用于商业目的的用户，需在 https://open.bigmodel.cn/mla/form 完成登记并获得基础商用授权。\n\n经过登记并获得基础商用授权的用户可以免费使用本模型进行商业活动，但必须遵守本许可的所有条款和条件。\n在本许可证下，您的商业活动的服务用户数量（访问量）不得超过100万人次访问 / 每月。如果超过，您需要与我们的商业团队联系以获得更多的商业许可。\n上述版权声明和本许可声明应包含在本软件的所有副本或重要部分中。\n\n3.限制\n\n您不得出于任何军事或非法目的使用、复制、修改、合并、发布、分发、复制或创建本软件的全部或部分衍生作品。\n\n您不得利用本软件从事任何危害国家安全和国家统一、危害社会公共利益、侵犯人身权益的行为。\n\n4.免责声明\n\n本软件“按原样”提供，不提供任何明示或暗示的保证，包括但不限于对适销性、特定用途的适用性和非侵权性的保证。\n在任何情况下，作者或版权持有人均不对任何索赔、损害或其他责任负责，无论是在合同诉讼、侵权行为还是其他方面，由软件或软件的使用或其他交易引起、由软件引起或与之相关 软件。\n\n5. 责任限制\n\n除适用法律禁止的范围外，在任何情况下且根据任何法律理论，无论是基于侵权行为、疏忽、合同、责任或其他原因，任何许可方均不对您承担任何直接、间接、特殊、偶然、示范性、 或间接损害，或任何其他商业损失，即使许可人已被告知此类损害的可能性。\n\n6.争议解决\n\n本许可受中华人民共和国法律管辖并按其解释。 因本许可引起的或与本许可有关的任何争议应提交北京市海淀区人民法院。\n\n请注意，许可证可能会更新到更全面的版本。 有关许可和版权的任何问题，请通过 license@zhipuai.cn 与我们联系。" }, { "path": "CogVideo/README.md", "content": "# CogVideo & CogVideoX\n\n[中文阅读](./README_zh.md)\n\n[日本語で読む](./README_ja.md)\n\n
\n\n
\n

\nExperience the CogVideoX-5B model online at 🤗 Huggingface Space or 🤖 ModelScope Space\n

\n

\n📚 View the paper and user guide\n

\n

\n 👋 Join our WeChat and Discord \n

\n

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

\n\n## Project Updates\n\n- 🔥🔥 **News**: ```2024/10/13```: A more cost-effective fine-tuning framework for `CogVideoX-5B` that works with a single\n 4090 GPU, [cogvideox-factory](https://github.com/a-r-r-o-w/cogvideox-factory), has been released. It supports\n fine-tuning with multiple resolutions. Feel free to use it!\n- 🔥 **News**: ```2024/10/10```: We have updated our technical report. Please\n click [here](https://arxiv.org/pdf/2408.06072) to view it. More training details and a demo have been added. To see\n the demo, click [here](https://yzy-thu.github.io/CogVideoX-demo/).- 🔥 **News**: ```2024/10/09```: We have publicly\n released the [technical documentation](https://zhipu-ai.feishu.cn/wiki/DHCjw1TrJiTyeukfc9RceoSRnCh) for CogVideoX\n fine-tuning on Feishu, further increasing distribution flexibility. All examples in the public documentation can be\n fully reproduced.\n- 🔥 **News**: ```2024/9/19```: We have open-sourced the CogVideoX series image-to-video model **CogVideoX-5B-I2V**.\n This model can take an image as a background input and generate a video combined with prompt words, offering greater\n controllability. With this, the CogVideoX series models now support three tasks: text-to-video generation, video\n continuation, and image-to-video generation. Welcome to try it online\n at [Experience](https://huggingface.co/spaces/THUDM/CogVideoX-5B-Space).\n- 🔥 ```2024/9/19```: The Caption\n model [CogVLM2-Caption](https://huggingface.co/THUDM/cogvlm2-llama3-caption), used in the training process of\n CogVideoX to convert video data into text descriptions, has been open-sourced. Welcome to download and use it.\n- 🔥 ```2024/8/27```: We have open-sourced a larger model in the CogVideoX series, **CogVideoX-5B**. We have\n significantly optimized the model's inference performance, greatly lowering the inference threshold. You can run *\n *CogVideoX-2B** on older GPUs like `GTX 1080TI`, and **CogVideoX-5B** on desktop GPUs like `RTX 3060`. Please strictly\n follow the [requirements](requirements.txt) to update and install dependencies, and refer\n to [cli_demo](inference/cli_demo.py) for inference code. Additionally, the open-source license for the **CogVideoX-2B\n ** model has been changed to the **Apache 2.0 License**.\n- 🔥 ```2024/8/6```: We have open-sourced **3D Causal VAE**, used for **CogVideoX-2B**, which can reconstruct videos with\n almost no loss.\n- 🔥 ```2024/8/6```: We have open-sourced the first model of the CogVideoX series video generation models, **CogVideoX-2B\n **.\n- 🌱 **Source**: ```2022/5/19```: We have open-sourced the CogVideo video generation model (now you can see it in\n the `CogVideo` branch). This is the first open-source large Transformer-based text-to-video generation model. You can\n access the [ICLR'23 paper](https://arxiv.org/abs/2205.15868) for technical details.\n\n## Table of Contents\n\nJump to a specific section:\n\n- [Quick Start](#Quick-Start)\n - [SAT](#sat)\n - [Diffusers](#Diffusers)\n- [CogVideoX-2B Video Works](#cogvideox-2b-gallery)\n- [Introduction to the CogVideoX Model](#Model-Introduction)\n- [Full Project Structure](#project-structure)\n - [Inference](#inference)\n - [SAT](#sat)\n - [Tools](#tools)\n- [Introduction to CogVideo(ICLR'23) Model](#cogvideoiclr23)\n- [Citations](#Citation)\n- [Open Source Project Plan](#Open-Source-Project-Plan)\n- [Model License](#Model-License)\n\n## Quick Start\n\n### Prompt Optimization\n\nBefore running the model, please refer to [this guide](inference/convert_demo.py) to see how we use large models like\nGLM-4 (or other comparable products, such as GPT-4) to optimize the model. This is crucial because the model is trained\nwith long prompts, and a good prompt directly impacts the quality of the video generation.\n\n### SAT\n\n**Please make sure your Python version is between 3.10 and 3.12, inclusive of both 3.10 and 3.12.**\n\nFollow instructions in [sat_demo](sat/README.md): Contains the inference code and fine-tuning code of SAT weights. It is\nrecommended to improve based on the CogVideoX model structure. Innovative researchers use this code to better perform\nrapid stacking and development.\n\n### Diffusers\n\n**Please make sure your Python version is between 3.10 and 3.12, inclusive of both 3.10 and 3.12.**\n\n```\npip install -r requirements.txt\n```\n\nThen follow [diffusers_demo](inference/cli_demo.py): A more detailed explanation of the inference code, mentioning the\nsignificance of common parameters.\n\nFor more details on quantized inference, please refer\nto [diffusers-torchao](https://github.com/sayakpaul/diffusers-torchao/). With Diffusers and TorchAO, quantized inference\nis also possible leading to memory-efficient inference as well as speedup in some cases when compiled. A full list of\nmemory and time benchmarks with various settings on A100 and H100 has been published\nat [diffusers-torchao](https://github.com/sayakpaul/diffusers-torchao).\n\n## Gallery\n\n### CogVideoX-5B\n\n\n \n \n \n \n \n \n \n \n \n \n \n \n
\n \n \n \n \n \n \n \n
\n \n \n \n \n \n \n \n
\n\n### CogVideoX-2B\n\n\n \n \n \n \n \n \n
\n \n \n \n \n \n \n \n
\n\nTo view the corresponding prompt words for the gallery, please click [here](resources/galary_prompt.md)\n\n## Model Introduction\n\nCogVideoX is an open-source version of the video generation model originating\nfrom [QingYing](https://chatglm.cn/video?lang=en?fr=osm_cogvideo). The table below displays the list of video generation\nmodels we currently offer, along with their foundational information.\n\n\n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n
Model NameCogVideoX-2BCogVideoX-5BCogVideoX-5B-I2V
Model DescriptionEntry-level model, balancing compatibility. Low cost for running and secondary development.Larger model with higher video generation quality and better visual effects.CogVideoX-5B image-to-video version.
Inference PrecisionFP16*(recommended), BF16, FP32, FP8*, INT8, not supported: INT4BF16 (recommended), FP16, FP32, FP8*, INT8, not supported: INT4
Single GPU Memory Usage
SAT FP16: 18GB
diffusers FP16: from 4GB*
diffusers INT8 (torchao): from 3.6GB*		SAT BF16: 26GB
diffusers BF16: from 5GB*
diffusers INT8 (torchao): from 4.4GB*
Multi-GPU Inference Memory UsageFP16: 10GB* using diffusers
BF16: 15GB* using diffusers
Inference Speed
(Step = 50, FP/BF16)		Single A100: ~90 seconds
Single H100: ~45 seconds		Single A100: ~180 seconds
Single H100: ~90 seconds
Fine-tuning PrecisionFP16BF16
Fine-tuning Memory Usage47 GB (bs=1, LORA)
61 GB (bs=2, LORA)
62GB (bs=1, SFT)		63 GB (bs=1, LORA)
80 GB (bs=2, LORA)
75GB (bs=1, SFT)
78 GB (bs=1, LORA)
75GB (bs=1, SFT, 16GPU)
Prompt LanguageEnglish*
Maximum Prompt Length226 Tokens
Video Length6 Seconds
Frame Rate8 Frames / Second
Video Resolution720 x 480, no support for other resolutions (including fine-tuning)
Position Encoding3d_sincos_pos_embed3d_sincos_pos_embed3d_rope_pos_embed + learnable_pos_embed
Download Link (Diffusers)🤗 HuggingFace
🤖 ModelScope
🟣 WiseModel		🤗 HuggingFace
🤖 ModelScope
🟣 WiseModel		🤗 HuggingFace
🤖 ModelScope
🟣 WiseModel
Download Link (SAT)SAT
\n\n**Data Explanation**\n\n+ While testing using the diffusers library, all optimizations included in the diffusers library were enabled. This\n scheme has not been tested for actual memory usage on devices outside of **NVIDIA A100 / H100** architectures.\n Generally, this scheme can be adapted to all **NVIDIA Ampere architecture** and above devices. If optimizations are\n disabled, memory consumption will multiply, with peak memory usage being about 3 times the value in the table.\n However, speed will increase by about 3-4 times. You can selectively disable some optimizations, including:\n\n```\npipe.enable_sequential_cpu_offload()\npipe.vae.enable_slicing()\npipe.vae.enable_tiling()\n```\n\n+ For multi-GPU inference, the `enable_sequential_cpu_offload()` optimization needs to be disabled.\n+ Using INT8 models will slow down inference, which is done to accommodate lower-memory GPUs while maintaining minimal\n video quality loss, though inference speed will significantly decrease.\n+ The CogVideoX-2B model was trained in `FP16` precision, and all CogVideoX-5B models were trained in `BF16` precision.\n We recommend using the precision in which the model was trained for inference.\n+ [PytorchAO](https://github.com/pytorch/ao) and [Optimum-quanto](https://github.com/huggingface/optimum-quanto/) can be\n used to quantize the text encoder, transformer, and VAE modules to reduce the memory requirements of CogVideoX. This\n allows the model to run on free T4 Colabs or GPUs with smaller memory! Also, note that TorchAO quantization is fully\n compatible with `torch.compile`, which can significantly improve inference speed. FP8 precision must be used on\n devices with NVIDIA H100 and above, requiring source installation of `torch`, `torchao`, `diffusers`, and `accelerate`\n Python packages. CUDA 12.4 is recommended.\n+ The inference speed tests also used the above memory optimization scheme. Without memory optimization, inference speed\n increases by about 10%. Only the `diffusers` version of the model supports quantization.\n+ The model only supports English input; other languages can be translated into English for use via large model\n refinement.\n+ The memory usage of model fine-tuning is tested in an `8 * H100` environment, and the program automatically\n uses `Zero 2` optimization. If a specific number of GPUs is marked in the table, that number or more GPUs must be used\n for fine-tuning.\n\n## Friendly Links\n\nWe highly welcome contributions from the community and actively contribute to the open-source community. The following\nworks have already been adapted for CogVideoX, and we invite everyone to use them:\n\n+ [CogVideoX-Fun](https://github.com/aigc-apps/CogVideoX-Fun): CogVideoX-Fun is a modified pipeline based on the\n CogVideoX architecture, supporting flexible resolutions and multiple launch methods.\n+ [CogStudio](https://github.com/pinokiofactory/cogstudio): A separate repository for CogVideo's Gradio Web UI, which\n supports more functional Web UIs.\n+ [Xorbits Inference](https://github.com/xorbitsai/inference): A powerful and comprehensive distributed inference\n framework, allowing you to easily deploy your own models or the latest cutting-edge open-source models with just one\n click.\n+ [ComfyUI-CogVideoXWrapper](https://github.com/kijai/ComfyUI-CogVideoXWrapper) Use the ComfyUI framework to integrate\n CogVideoX into your workflow.\n+ [VideoSys](https://github.com/NUS-HPC-AI-Lab/VideoSys): VideoSys provides a user-friendly, high-performance\n infrastructure for video generation, with full pipeline support and continuous integration of the latest models and\n techniques.\n+ [AutoDL Space](https://www.codewithgpu.com/i/THUDM/CogVideo/CogVideoX-5b-demo): A one-click deployment Huggingface\n Space image provided by community members.\n+ [Interior Design Fine-Tuning Model](https://huggingface.co/collections/bertjiazheng/koolcogvideox-66e4762f53287b7f39f8f3ba):\n is a fine-tuned model based on CogVideoX, specifically designed for interior design.\n+ [xDiT](https://github.com/xdit-project/xDiT): xDiT is a scalable inference engine for Diffusion Transformers (DiTs)\n on multiple GPU Clusters. xDiT supports real-time image and video generations services.\n [cogvideox-factory](https://github.com/a-r-r-o-w/cogvideox-factory): A cost-effective\n fine-tuning framework for CogVideoX, compatible with the `diffusers` version model. Supports more resolutions, and\n fine-tuning CogVideoX-5B can be done with a single 4090 GPU.\n+ [CogVideoX-Interpolation](https://github.com/feizc/CogvideX-Interpolation): A pipeline based on the modified CogVideoX\n structure, aimed at providing greater flexibility for keyframe interpolation generation.\n+ [DiffSynth-Studio](https://github.com/modelscope/DiffSynth-Studio): DiffSynth Studio is a diffusion engine. It has\n restructured the architecture, including text encoders, UNet, VAE, etc., enhancing computational performance while\n maintaining compatibility with open-source community models. The framework has been adapted for CogVideoX.\n\n## Project Structure\n\nThis open-source repository will guide developers to quickly get started with the basic usage and fine-tuning examples\nof the **CogVideoX** open-source model.\n\n### Quick Start with Colab\n\nHere provide three projects that can be run directly on free Colab T4 instances:\n\n+ [CogVideoX-5B-T2V-Colab.ipynb](https://colab.research.google.com/drive/1pCe5s0bC_xuXbBlpvIH1z0kfdTLQPzCS?usp=sharing):\n CogVideoX-5B Text-to-Video Colab code.\n+ [CogVideoX-5B-T2V-Int8-Colab.ipynb](https://colab.research.google.com/drive/1DUffhcjrU-uz7_cpuJO3E_D4BaJT7OPa?usp=sharing):\n CogVideoX-5B Quantized Text-to-Video Inference Colab code, which takes about 30 minutes per run.\n+ [CogVideoX-5B-I2V-Colab.ipynb](https://colab.research.google.com/drive/17CqYCqSwz39nZAX2YyonDxosVKUZGzcX?usp=sharing):\n CogVideoX-5B Image-to-Video Colab code.\n+ [CogVideoX-5B-V2V-Colab.ipynb](https://colab.research.google.com/drive/1comfGAUJnChl5NwPuO8Ox5_6WCy4kbNN?usp=sharing):\n CogVideoX-5B Video-to-Video Colab code.\n\n### Inference\n\n+ [dcli_demo](inference/cli_demo.py): A more detailed inference code explanation, including the significance of\n common parameters. All of this is covered here.\n+ [cli_demo_quantization](inference/cli_demo_quantization.py):\n Quantized model inference code that can run on devices with lower memory. You can also modify this code to support\n running CogVideoX models in FP8 precision.\n+ [diffusers_vae_demo](inference/cli_vae_demo.py): Code for running VAE inference separately.\n+ [space demo](inference/gradio_composite_demo): The same GUI code as used in the Huggingface Space, with frame\n interpolation and super-resolution tools integrated.\n\n
\n \n
\n\n+ [convert_demo](inference/convert_demo.py): How to convert user input into long-form input suitable for CogVideoX.\n Since CogVideoX is trained on long texts, we need to transform the input text distribution to match the training data\n using an LLM. The script defaults to using GLM-4, but it can be replaced with GPT, Gemini, or any other large language\n model.\n+ [gradio_web_demo](inference/gradio_composite_demo): A simple Gradio web application demonstrating how to use the\n CogVideoX-2B / 5B model to generate videos. Similar to our Huggingface Space, you can use this script to run a simple\n web application for video generation.\n\n### finetune\n\n+ [finetune_demo](finetune/README.md): Fine-tuning scheme and details of the diffusers version of the CogVideoX model.\n\n### sat\n\n+ [sat_demo](sat/README.md): Contains the inference code and fine-tuning code of SAT weights. It is recommended to\n improve based on the CogVideoX model structure. Innovative researchers use this code to better perform rapid stacking\n and development.\n\n### Tools\n\nThis folder contains some tools for model conversion / caption generation, etc.\n\n+ [convert_weight_sat2hf](tools/convert_weight_sat2hf.py): Converts SAT model weights to Huggingface model weights.\n+ [caption_demo](tools/caption/README.md): Caption tool, a model that understands videos and outputs descriptions in\n text.\n+ [export_sat_lora_weight](tools/export_sat_lora_weight.py): SAT fine-tuning model export tool, exports the SAT Lora\n Adapter in diffusers format.\n+ [load_cogvideox_lora](tools/load_cogvideox_lora.py): Tool code for loading the diffusers version of fine-tuned Lora\n Adapter.\n+ [llm_flux_cogvideox](tools/llm_flux_cogvideox/llm_flux_cogvideox.py): Automatically generate videos using an\n open-source local large language model + Flux + CogVideoX.\n+ [parallel_inference_xdit](tools/parallel_inference/parallel_inference_xdit.py):\n Supported by [xDiT](https://github.com/xdit-project/xDiT), parallelize the\n video generation process on multiple GPUs.\n\n## CogVideo(ICLR'23)\n\nThe official repo for the\npaper: [CogVideo: Large-scale Pretraining for Text-to-Video Generation via Transformers](https://arxiv.org/abs/2205.15868)\nis on the [CogVideo branch](https://github.com/THUDM/CogVideo/tree/CogVideo)\n\n**CogVideo is able to generate relatively high-frame-rate videos.**\nA 4-second clip of 32 frames is shown below.\n\n![High-frame-rate sample](https://raw.githubusercontent.com/THUDM/CogVideo/CogVideo/assets/appendix-sample-highframerate.png)\n\n![Intro images](https://raw.githubusercontent.com/THUDM/CogVideo/CogVideo/assets/intro-image.png)\n
\n \n
\n\n\nThe demo for CogVideo is at [https://models.aminer.cn/cogvideo](https://models.aminer.cn/cogvideo/), where you can get\nhands-on practice on text-to-video generation. *The original input is in Chinese.*\n\n## Citation\n\n🌟 If you find our work helpful, please leave us a star and cite our paper.\n\n```\n@article{yang2024cogvideox,\n title={CogVideoX: Text-to-Video Diffusion Models with An Expert Transformer},\n author={Yang, Zhuoyi and Teng, Jiayan and Zheng, Wendi and Ding, Ming and Huang, Shiyu and Xu, Jiazheng and Yang, Yuanming and Hong, Wenyi and Zhang, Xiaohan and Feng, Guanyu and others},\n journal={arXiv preprint arXiv:2408.06072},\n year={2024}\n}\n@article{hong2022cogvideo,\n title={CogVideo: Large-scale Pretraining for Text-to-Video Generation via Transformers},\n author={Hong, Wenyi and Ding, Ming and Zheng, Wendi and Liu, Xinghan and Tang, Jie},\n journal={arXiv preprint arXiv:2205.15868},\n year={2022}\n}\n```\n\nWe welcome your contributions! You can click [here](resources/contribute.md) for more information.\n\n## License Agreement\n\nThe code in this repository is released under the [Apache 2.0 License](LICENSE).\n\nThe CogVideoX-2B model (including its corresponding Transformers module and VAE module) is released under\nthe [Apache 2.0 License](LICENSE).\n\nThe CogVideoX-5B model (Transformers module, include I2V and T2V) is released under\nthe [CogVideoX LICENSE](https://huggingface.co/THUDM/CogVideoX-5b/blob/main/LICENSE).\n" }, { "path": "CogVideo/README_ja.md", "content": "# CogVideo & CogVideoX\n\n[Read this in English](./README_zh.md)\n\n[中文阅读](./README_zh.md)\n\n
\n\n
\n

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

\n

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

\n

\n 👋 WeChat と Discord に参加\n

\n

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

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

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

\n

\n📚 查看 论文 和 使用文档\n

\n

\n 👋 加入我们的 微信 和 Discord \n

\n

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

\n\n## 项目更新\n\n- 🔥🔥 **News**: ```2024/10/13```: 成本更低，单卡4090可微调`CogVideoX-5B`\n 的微调框架[cogvideox-factory](https://github.com/a-r-r-o-w/cogvideox-factory)已经推出，多种分辨率微调，欢迎使用。\n- 🔥 **News**: ```2024/10/10```: 我们更新了我们的技术报告,请点击 [这里](https://arxiv.org/pdf/2408.06072)\n 查看，附上了更多的训练细节和demo，关于demo，点击[这里](https://yzy-thu.github.io/CogVideoX-demo/) 查看。\n- 🔥 **News**: ```2024/10/09```: 我们在飞书[技术文档](https://zhipu-ai.feishu.cn/wiki/DHCjw1TrJiTyeukfc9RceoSRnCh\")\n 公开CogVideoX微调指导，以进一步增加分发自由度，公开文档中所有示例可以完全复现\n- 🔥 **News**: ```2024/9/19```: 我们开源 CogVideoX 系列图生视频模型 **CogVideoX-5B-I2V**\n 。该模型可以将一张图像作为背景输入，结合提示词一起生成视频，具有更强的可控性。\n 至此，CogVideoX系列模型已经支持文本生成视频，视频续写，图片生成视频三种任务。欢迎前往在线[体验](https://huggingface.co/spaces/THUDM/CogVideoX-5B-Space)。\n- 🔥 **News**: ```2024/9/19```: CogVideoX 训练过程中用于将视频数据转换为文本描述的 Caption\n 模型 [CogVLM2-Caption](https://huggingface.co/THUDM/cogvlm2-llama3-caption)\n 已经开源。欢迎前往下载并使用。\n- 🔥 ```2024/8/27```: 我们开源 CogVideoX 系列更大的模型 **CogVideoX-5B**\n 。我们大幅度优化了模型的推理性能，推理门槛大幅降低，您可以在 `GTX 1080TI` 等早期显卡运行 **CogVideoX-2B**，在 `RTX 3060`\n 等桌面端甜品卡运行 **CogVideoX-5B** 模型。 请严格按照[要求](requirements.txt)\n 更新安装依赖，推理代码请查看 [cli_demo](inference/cli_demo.py)。同时，**CogVideoX-2B** 模型开源协议已经修改为**Apache 2.0\n 协议**。\n- 🔥 ```2024/8/6```: 我们开源 **3D Causal VAE**，用于 **CogVideoX-2B**，可以几乎无损地重构视频。\n- 🔥 ```2024/8/6```: 我们开源 CogVideoX 系列视频生成模型的第一个模型, **CogVideoX-2B**。\n- 🌱 **Source**: ```2022/5/19```: 我们开源了 CogVideo 视频生成模型（现在你可以在 `CogVideo` 分支中看到），这是首个开源的基于\n Transformer 的大型文本生成视频模型，您可以访问 [ICLR'23 论文](https://arxiv.org/abs/2205.15868) 查看技术细节。\n\n## 目录\n\n跳转到指定部分：\n\n- [快速开始](#快速开始)\n - [SAT](#sat)\n - [Diffusers](#Diffusers)\n- [CogVideoX-2B 视频作品](#cogvideox-2b-视频作品)\n- [CogVideoX模型介绍](#模型介绍)\n- [完整项目代码结构](#完整项目代码结构)\n - [Inference](#inference)\n - [SAT](#sat)\n - [Tools](#tools)\n- [开源项目规划](#开源项目规划)\n- [模型协议](#模型协议)\n- [CogVideo(ICLR'23)模型介绍](#cogvideoiclr23)\n- [引用](#引用)\n\n## 快速开始\n\n### 提示词优化\n\n在开始运行模型之前，请参考 [这里](inference/convert_demo.py) 查看我们是怎么使用GLM-4(或者同级别的其他产品，例如GPT-4)\n大模型对模型进行优化的，这很重要，\n由于模型是在长提示词下训练的，一个好的提示词直接影响了视频生成的质量。\n\n### SAT\n\n查看sat文件夹下的 [sat_demo](sat/README.md)：包含了 SAT 权重的推理代码和微调代码，推荐基于此代码进行 CogVideoX\n模型结构的改进，研究者使用该代码可以更好的进行快速的迭代和开发。\n\n### Diffusers\n\n```\npip install -r requirements.txt\n```\n\n查看[diffusers_demo](inference/cli_demo.py)：包含对推理代码更详细的解释，包括各种关键的参数。\n\n欲了解更多关于量化推理的细节，请参考 [diffusers-torchao](https://github.com/sayakpaul/diffusers-torchao/)。使用 Diffusers\n和 TorchAO，量化推理也是可能的，这可以实现内存高效的推理，并且在某些情况下编译后速度有所提升。有关在 A100 和 H100\n上使用各种设置的内存和时间基准测试的完整列表，已发布在 [diffusers-torchao](https://github.com/sayakpaul/diffusers-torchao)\n上。\n\n## 视频作品\n\n### CogVideoX-5B\n\n\n \n \n \n \n \n \n \n \n \n \n \n \n
\n \n \n \n \n \n \n \n
\n \n \n \n \n \n \n \n
\n\n### CogVideoX-2B\n\n\n \n \n \n \n \n \n
\n \n \n \n \n \n \n \n
\n\n\n查看画廊的对应提示词，请点击[这里](resources/galary_prompt.md)\n\n## 模型介绍\n\nCogVideoX是 [清影](https://chatglm.cn/video?fr=osm_cogvideox) 同源的开源版本视频生成模型。\n下表展示我们提供的视频生成模型相关基础信息:\n\n\n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n
模型名CogVideoX-2BCogVideoX-5BCogVideoX-5B-I2V
推理精度FP16*(推荐), BF16, FP32，FP8*，INT8，不支持INT4BF16(推荐), FP16, FP32，FP8*，INT8，不支持INT4
单GPU显存消耗
SAT FP16: 18GB
diffusers FP16: 4GB起*
diffusers INT8(torchao): 3.6G起*		SAT BF16: 26GB
diffusers BF16 : 5GB起*
diffusers INT8(torchao): 4.4G起*
多GPU推理显存消耗FP16: 10GB* using diffusers
BF16: 15GB* using diffusers
推理速度
(Step = 50, FP/BF16)		单卡A100: ~90秒
单卡H100: ~45秒		单卡A100: ~180秒
单卡H100: ~90秒
微调精度FP16BF16
微调显存消耗47 GB (bs=1, LORA)
61 GB (bs=2, LORA)
62GB (bs=1, SFT)		63 GB (bs=1, LORA)
80 GB (bs=2, LORA)
75GB (bs=1, SFT)
78 GB (bs=1, LORA)
75GB (bs=1, SFT, 16GPU)
提示词语言English*
提示词长度上限226 Tokens
视频长度6 秒
帧率8 帧 / 秒
视频分辨率720 * 480，不支持其他分辨率(含微调)
位置编码3d_sincos_pos_embed3d_sincos_pos_embed3d_rope_pos_embed + learnable_pos_embed
下载链接 (Diffusers)🤗 HuggingFace
🤖 ModelScope
🟣 WiseModel		🤗 HuggingFace
🤖 ModelScope
🟣 WiseModel		🤗 HuggingFace
🤖 ModelScope
🟣 WiseModel
下载链接 (SAT)SAT
\n\n**数据解释**\n\n+ 使用 diffusers 库进行测试时，启用了全部`diffusers`库自带的优化，该方案未测试在非**NVIDIA A100 / H100**\n 外的设备上的实际显存 / 内存占用。通常，该方案可以适配于所有 **NVIDIA 安培架构**\n 以上的设备。若关闭优化，显存占用会成倍增加，峰值显存约为表格的3倍。但速度提升3-4倍左右。你可以选择性的关闭部分优化，这些优化包括:\n\n```\npipe.enable_sequential_cpu_offload()\npipe.vae.enable_slicing()\npipe.vae.enable_tiling()\n```\n\n+ 多GPU推理时，需要关闭 `enable_sequential_cpu_offload()` 优化。\n+ 使用 INT8 模型会导致推理速度降低，此举是为了满足显存较低的显卡能正常推理并保持较少的视频质量损失，推理速度大幅降低。\n+ CogVideoX-2B 模型采用 `FP16` 精度训练， 搜有 CogVideoX-5B 模型采用 `BF16` 精度训练。我们推荐使用模型训练的精度进行推理。\n+ [PytorchAO](https://github.com/pytorch/ao) 和 [Optimum-quanto](https://github.com/huggingface/optimum-quanto/)\n 可以用于量化文本编码器、Transformer 和 VAE 模块，以降低 CogVideoX 的内存需求。这使得在免费的 T4 Colab 或更小显存的 GPU\n 上运行模型成为可能！同样值得注意的是，TorchAO 量化完全兼容 `torch.compile`，这可以显著提高推理速度。在 `NVIDIA H100`\n 及以上设备上必须使用 `FP8` 精度，这需要源码安装 `torch`、`torchao`、`diffusers` 和 `accelerate` Python\n 包。建议使用 `CUDA 12.4`。\n+ 推理速度测试同样采用了上述显存优化方案，不采用显存优化的情况下，推理速度提升约10%。 只有`diffusers`版本模型支持量化。\n+ 模型仅支持英语输入，其他语言可以通过大模型润色时翻译为英语。\n+ 模型微调所占用的显存是在 `8 * H100` 环境下进行测试，程序已经自动使用`Zero 2` 优化。表格中若有标注具体GPU数量则必须使用大于等于该数量的GPU进行微调。\n\n## 友情链接\n\n我们非常欢迎来自社区的贡献，并积极的贡献开源社区。以下作品已经对CogVideoX进行了适配，欢迎大家使用:\n\n+ [CogVideoX-Fun](https://github.com/aigc-apps/CogVideoX-Fun):\n CogVideoX-Fun是一个基于CogVideoX结构修改后的的pipeline，支持自由的分辨率，多种启动方式。\n+ [CogStudio](https://github.com/pinokiofactory/cogstudio): CogVideo 的 Gradio Web UI单独实现仓库，支持更多功能的 Web UI。\n+ [Xorbits Inference](https://github.com/xorbitsai/inference): 性能强大且功能全面的分布式推理框架，轻松一键部署你自己的模型或内置的前沿开源模型。\n+ [ComfyUI-CogVideoXWrapper](https://github.com/kijai/ComfyUI-CogVideoXWrapper) 使用ComfyUI框架，将CogVideoX加入到你的工作流中。\n+ [VideoSys](https://github.com/NUS-HPC-AI-Lab/VideoSys): VideoSys 提供了易用且高性能的视频生成基础设施，支持完整的管道，并持续集成最新的模型和技术。\n+ [AutoDL镜像](https://www.codewithgpu.com/i/THUDM/CogVideo/CogVideoX-5b-demo): 由社区成员提供的一键部署Huggingface\n Space镜像。\n+ [室内设计微调模型](https://huggingface.co/collections/bertjiazheng/koolcogvideox-66e4762f53287b7f39f8f3ba) 基于\n CogVideoX的微调模型，它专为室内设计而设计\n+ [xDiT](https://github.com/xdit-project/xDiT): xDiT是一个用于在多GPU集群上对DiTs并行推理的引擎。xDiT支持实时图像和视频生成服务。\n+ [CogVideoX-Interpolation](https://github.com/feizc/CogvideX-Interpolation): 基于 CogVideoX 结构修改的管道，旨在为关键帧插值生成提供更大的灵活性。\n+ [DiffSynth-Studio](https://github.com/modelscope/DiffSynth-Studio): DiffSynth 工作室是一款扩散引擎。重构了架构，包括文本编码器、UNet、VAE\n 等，在保持与开源社区模型兼容性的同时，提升了计算性能。该框架已经适配 CogVideoX。\n\n## 完整项目代码结构\n\n本开源仓库将带领开发者快速上手 **CogVideoX** 开源模型的基础调用方式、微调示例。\n\n### Colab 快速使用\n\n这里提供了三个能直接在免费的 Colab T4上 运行的项目\n\n+ [CogVideoX-5B-T2V-Colab.ipynb](https://colab.research.google.com/drive/1pCe5s0bC_xuXbBlpvIH1z0kfdTLQPzCS?usp=sharing):\n CogVideoX-5B 文字生成视频 Colab 代码。\n+ [CogVideoX-5B-T2V-Int8-Colab.ipynb](https://colab.research.google.com/drive/1DUffhcjrU-uz7_cpuJO3E_D4BaJT7OPa?usp=sharing):\n CogVideoX-5B 文字生成视频量化推理 Colab 代码，运行一次大约需要30分钟。\n+ [CogVideoX-5B-I2V-Colab.ipynb](https://colab.research.google.com/drive/17CqYCqSwz39nZAX2YyonDxosVKUZGzcX?usp=sharing):\n CogVideoX-5B 图片生成视频 Colab 代码。\n+ [CogVideoX-5B-V2V-Colab.ipynb](https://colab.research.google.com/drive/1comfGAUJnChl5NwPuO8Ox5_6WCy4kbNN?usp=sharing):\n CogVideoX-5B 视频生成视频 Colab 代码。\n\n### inference\n\n+ [cli_demo](inference/cli_demo.py): 更详细的推理代码讲解，常见参数的意义，在这里都会提及。\n+ [cli_demo_quantization](inference/cli_demo_quantization.py):\n 量化模型推理代码，可以在显存较低的设备上运行，也可以基于此代码修改，以支持运行FP8等精度的CogVideoX模型。请注意，FP8\n 仅测试通过，且必须将 `torch-nightly`,`torchao`源代码安装，不建议在生产环境中使用。\n+ [diffusers_vae_demo](inference/cli_vae_demo.py): 单独执行VAE的推理代码。\n+ [space demo](inference/gradio_composite_demo): Huggingface Space同款的 GUI 代码，植入了插帧，超分工具。\n\n
\n \n
\n\n+ [convert_demo](inference/convert_demo.py): 如何将用户的输入转换成适合\n CogVideoX的长输入。因为CogVideoX是在长文本上训练的，所以我们需要把输入文本的分布通过LLM转换为和训练一致的长文本。脚本中默认使用GLM-4，也可以替换为GPT、Gemini等任意大语言模型。\n+ [gradio_web_demo](inference/gradio_composite_demo/app.py): 与 Huggingface Space 完全相同的代码实现，快速部署 CogVideoX\n GUI体验。\n\n### finetune\n\n+ [train_cogvideox_lora](finetune/README_zh.md): diffusers版本 CogVideoX 模型微调方案和细节。\n\n### sat\n\n+ [sat_demo](sat/README_zh.md): 包含了 SAT 权重的推理代码和微调代码，推荐基于 CogVideoX\n 模型结构进行改进，创新的研究者使用改代码以更好的进行快速的堆叠和开发。\n\n### tools\n\n本文件夹包含了一些工具，用于模型的转换 / Caption 等工作。\n\n+ [convert_weight_sat2hf](tools/convert_weight_sat2hf.py): 将 SAT 模型权重转换为 Huggingface 模型权重。\n+ [caption_demo](tools/caption/README_zh.md): Caption 工具，对视频理解并用文字输出的模型。\n+ [export_sat_lora_weight](tools/export_sat_lora_weight.py): SAT微调模型导出工具，将\n SAT Lora Adapter 导出为 diffusers 格式。\n+ [load_cogvideox_lora](tools/load_cogvideox_lora.py): 载入diffusers版微调Lora Adapter的工具代码。\n+ [llm_flux_cogvideox](tools/llm_flux_cogvideox/llm_flux_cogvideox.py): 使用开源本地大语言模型 + Flux +\n CogVideoX实现自动化生成视频。\n+ [parallel_inference_xdit](tools/parallel_inference/parallel_inference_xdit.py):\n 在多个 GPU 上并行化视频生成过程，\n 由[xDiT](https://github.com/xdit-project/xDiT)提供支持。\n+ [cogvideox-factory](https://github.com/a-r-r-o-w/cogvideox-factory): CogVideoX低成文微调框架，适配`diffusers`\n 版本模型。支持更多分辨率，单卡4090即可微调 CogVideoX-5B 。\n\n## CogVideo(ICLR'23)\n\n[CogVideo: Large-scale Pretraining for Text-to-Video Generation via Transformers](https://arxiv.org/abs/2205.15868)\n的官方repo位于[CogVideo branch](https://github.com/THUDM/CogVideo/tree/CogVideo)。\n\n**CogVideo可以生成高帧率视频，下面展示了一个32帧的4秒视频。**\n\n![High-frame-rate sample](https://raw.githubusercontent.com/THUDM/CogVideo/CogVideo/assets/appendix-sample-highframerate.png)\n\n![Intro images](https://raw.githubusercontent.com/THUDM/CogVideo/CogVideo/assets/intro-image.png)\n\n\n
\n \n
\n\nCogVideo的demo网站在[https://models.aminer.cn/cogvideo](https://models.aminer.cn/cogvideo/)。您可以在这里体验文本到视频生成。\n*原始输入为中文。*\n\n## 引用\n\n🌟 如果您发现我们的工作有所帮助，欢迎引用我们的文章，留下宝贵的stars\n\n```\n@article{yang2024cogvideox,\n title={CogVideoX: Text-to-Video Diffusion Models with An Expert Transformer},\n author={Yang, Zhuoyi and Teng, Jiayan and Zheng, Wendi and Ding, Ming and Huang, Shiyu and Xu, Jiazheng and Yang, Yuanming and Hong, Wenyi and Zhang, Xiaohan and Feng, Guanyu and others},\n journal={arXiv preprint arXiv:2408.06072},\n year={2024}\n}\n@article{hong2022cogvideo,\n title={CogVideo: Large-scale Pretraining for Text-to-Video Generation via Transformers},\n author={Hong, Wenyi and Ding, Ming and Zheng, Wendi and Liu, Xinghan and Tang, Jie},\n journal={arXiv preprint arXiv:2205.15868},\n year={2022}\n}\n```\n\n我们欢迎您的贡献，您可以点击[这里](resources/contribute_zh.md)查看更多信息。\n\n## 模型协议\n\n本仓库代码使用 [Apache 2.0 协议](LICENSE) 发布。\n\nCogVideoX-2B 模型 (包括其对应的Transformers模块，VAE模块) 根据 [Apache 2.0 协议](LICENSE) 许可证发布。\n\nCogVideoX-5B 模型 (Transformers 模块，包括图生视频，文生视频版本)\n根据 [CogVideoX LICENSE](https://huggingface.co/THUDM/CogVideoX-5b/blob/main/LICENSE)\n许可证发布。\n" }, { "path": "CogVideo/download.sh", "content": "mkdir CogVideoX-2b-sat\ncd CogVideoX-2b-sat\nwget https://cloud.tsinghua.edu.cn/f/fdba7608a49c463ba754/?dl=1\nmv 'index.html?dl=1' vae.zip\nunzip vae.zip\nwget https://cloud.tsinghua.edu.cn/f/556a3e1329e74f1bac45/?dl=1\nmv 'index.html?dl=1' transformer.zip\nunzip transformer.zip" }, { "path": "CogVideo/finetune/README.md", "content": "# CogVideoX diffusers Fine-tuning Guide\n\n[中文阅读](./README_zh.md)\n\n[日本語で読む](./README_ja.md)\n\nThis feature is not fully complete yet. If you want to check the fine-tuning for the SAT version, please\nsee [here](../sat/README_zh.md). The dataset format is different from this version.\n\n## Hardware Requirements\n\n+ CogVideoX-2B / 5B LoRA: 1 * A100 (5B need to use `--use_8bit_adam`)\n+ CogVideoX-2B SFT: 8 * A100 (Working)\n+ CogVideoX-5B-I2V is not supported yet.\n\n## Install Dependencies\n\nSince the related code has not been merged into the diffusers release, you need to base your fine-tuning on the\ndiffusers branch. Please follow the steps below to install dependencies:\n\n```shell\ngit clone https://github.com/huggingface/diffusers.git\ncd diffusers # Now in Main branch\npip install -e .\n```\n\n## Prepare the Dataset\n\nFirst, you need to prepare the dataset. The dataset format should be as follows, with `videos.txt` containing the list\nof videos in the `videos` directory:\n\n```\n.\n├── prompts.txt\n├── videos\n└── videos.txt\n```\n\nYou can download\nthe [Disney Steamboat Willie](https://huggingface.co/datasets/Wild-Heart/Disney-VideoGeneration-Dataset) dataset from\nhere.\n\nThis video fine-tuning dataset is used as a test for fine-tuning.\n\n## Configuration Files and Execution\n\nThe `accelerate` configuration files are as follows:\n\n+ `accelerate_config_machine_multi.yaml`: Suitable for multi-GPU use\n+ `accelerate_config_machine_single.yaml`: Suitable for single-GPU use\n\nThe configuration for the `finetune` script is as follows:\n\n```\naccelerate launch --config_file accelerate_config_machine_single.yaml --multi_gpu \\ # Use accelerate to launch multi-GPU training with the config file accelerate_config_machine_single.yaml\n train_cogvideox_lora.py \\ # Training script train_cogvideox_lora.py for LoRA fine-tuning on CogVideoX model\n --gradient_checkpointing \\ # Enable gradient checkpointing to reduce memory usage\n --pretrained_model_name_or_path $MODEL_PATH \\ # Path to the pretrained model, specified by $MODEL_PATH\n --cache_dir $CACHE_PATH \\ # Cache directory for model files, specified by $CACHE_PATH\n --enable_tiling \\ # Enable tiling technique to process videos in chunks, saving memory\n --enable_slicing \\ # Enable slicing to further optimize memory by slicing inputs\n --instance_data_root $DATASET_PATH \\ # Dataset path specified by $DATASET_PATH\n --caption_column prompts.txt \\ # Specify the file prompts.txt for video descriptions used in training\n --video_column videos.txt \\ # Specify the file videos.txt for video paths used in training\n --validation_prompt \"\" \\ # Prompt used for generating validation videos during training\n --validation_prompt_separator ::: \\ # Set ::: as the separator for validation prompts\n --num_validation_videos 1 \\ # Generate 1 validation video per validation round\n --validation_epochs 100 \\ # Perform validation every 100 training epochs\n --seed 42 \\ # Set random seed to 42 for reproducibility\n --rank 128 \\ # Set the rank for LoRA parameters to 128\n --lora_alpha 64 \\ # Set the alpha parameter for LoRA to 64, adjusting LoRA learning rate\n --mixed_precision bf16 \\ # Use bf16 mixed precision for training to save memory\n --output_dir $OUTPUT_PATH \\ # Specify the output directory for the model, defined by $OUTPUT_PATH\n --height 480 \\ # Set video height to 480 pixels\n --width 720 \\ # Set video width to 720 pixels\n --fps 8 \\ # Set video frame rate to 8 frames per second\n --max_num_frames 49 \\ # Set the maximum number of frames per video to 49\n --skip_frames_start 0 \\ # Skip 0 frames at the start of the video\n --skip_frames_end 0 \\ # Skip 0 frames at the end of the video\n --train_batch_size 4 \\ # Set training batch size to 4\n --num_train_epochs 30 \\ # Total number of training epochs set to 30\n --checkpointing_steps 1000 \\ # Save model checkpoint every 1000 steps\n --gradient_accumulation_steps 1 \\ # Accumulate gradients for 1 step, updating after each batch\n --learning_rate 1e-3 \\ # Set learning rate to 0.001\n --lr_scheduler cosine_with_restarts \\ # Use cosine learning rate scheduler with restarts\n --lr_warmup_steps 200 \\ # Warm up the learning rate for the first 200 steps\n --lr_num_cycles 1 \\ # Set the number of learning rate cycles to 1\n --optimizer AdamW \\ # Use the AdamW optimizer\n --adam_beta1 0.9 \\ # Set Adam optimizer beta1 parameter to 0.9\n --adam_beta2 0.95 \\ # Set Adam optimizer beta2 parameter to 0.95\n --max_grad_norm 1.0 \\ # Set maximum gradient clipping value to 1.0\n --allow_tf32 \\ # Enable TF32 to speed up training\n --report_to wandb # Use Weights and Biases (wandb) for logging and monitoring the training\n```\n\n## Running the Script to Start Fine-tuning\n\nSingle Node (One GPU or Multi GPU) fine-tuning:\n\n```shell\nbash finetune_single_rank.sh\n```\n\nMulti-Node fine-tuning:\n\n```shell\nbash finetune_multi_rank.sh # Needs to be run on each node\n```\n\n## Loading the Fine-tuned Model\n\n+ Please refer to [cli_demo.py](../inference/cli_demo.py) for how to load the fine-tuned model.\n\n## Best Practices\n\n+ Includes 70 training videos with a resolution of `200 x 480 x 720` (frames x height x width). By skipping frames in\n the data preprocessing, we created two smaller datasets with 49 and 16 frames to speed up experimentation, as the\n maximum frame limit recommended by the CogVideoX team is 49 frames. We split the 70 videos into three groups of 10,\n 25, and 50 videos, with similar conceptual nature.\n+ Using 25 or more videos works best when training new concepts and styles.\n+ It works better to train using identifier tokens specified with `--id_token`. This is similar to Dreambooth training,\n but regular fine-tuning without such tokens also works.\n+ The original repository used `lora_alpha` set to 1. We found this value ineffective across multiple runs, likely due\n to differences in the backend and training setup. Our recommendation is to set `lora_alpha` equal to rank or rank //\n 2.\n+ We recommend using a rank of 64 or higher.\n" }, { "path": "CogVideo/finetune/README_ja.md", "content": "# CogVideoX diffusers 微調整方法\n\n[Read this in English.](./README_zh)\n\n[中文阅读](./README_zh.md)\n\n\nこの機能はまだ完全に完成していません。SATバージョンの微調整を確認したい場合は、[こちら](../sat/README_ja.md)を参照してください。本バージョンとは異なるデータセット形式を使用しています。\n\n## ハードウェア要件\n\n+ CogVideoX-2B / 5B T2V LORA: 1 * A100 (5B need to use `--use_8bit_adam`)\n+ CogVideoX-2B SFT: 8 * A100 （動作確認済み）\n+ CogVideoX-5B-I2V まだサポートしていません\n\n## 依存関係のインストール\n\n関連コードはまだdiffusersのリリース版に統合されていないため、diffusersブランチを使用して微調整を行う必要があります。以下の手順に従って依存関係をインストールしてください：\n\n```shell\ngit clone https://github.com/huggingface/diffusers.git\ncd diffusers # Now in Main branch\npip install -e .\n```\n\n## データセットの準備\n\nまず、データセットを準備する必要があります。データセットの形式は以下のようになります。\n\n```\n.\n├── prompts.txt\n├── videos\n└── videos.txt\n```\n\n[ディズニースチームボートウィリー](https://huggingface.co/datasets/Wild-Heart/Disney-VideoGeneration-Dataset)をここからダウンロードできます。\n\nビデオ微調整データセットはテスト用として使用されます。\n\n## 設定ファイルと実行\n\n`accelerate` 設定ファイルは以下の通りです:\n\n+ accelerate_config_machine_multi.yaml 複数GPU向け\n+ accelerate_config_machine_single.yaml 単一GPU向け\n\n`finetune` スクリプト設定ファイルの例：\n\n```\naccelerate launch --config_file accelerate_config_machine_single.yaml --multi_gpu \\ # accelerateを使用してmulti-GPUトレーニングを起動、設定ファイルはaccelerate_config_machine_single.yaml\n train_cogvideox_lora.py \\ # LoRAの微調整用のトレーニングスクリプトtrain_cogvideox_lora.pyを実行\n --gradient_checkpointing \\ # メモリ使用量を減らすためにgradient checkpointingを有効化\n --pretrained_model_name_or_path $MODEL_PATH \\ # 事前学習済みモデルのパスを$MODEL_PATHで指定\n --cache_dir $CACHE_PATH \\ # モデルファイルのキャッシュディレクトリを$CACHE_PATHで指定\n --enable_tiling \\ # メモリ節約のためにタイル処理を有効化し、動画をチャンク分けして処理\n --enable_slicing \\ # 入力をスライスしてさらにメモリ最適化\n --instance_data_root $DATASET_PATH \\ # データセットのパスを$DATASET_PATHで指定\n --caption_column prompts.txt \\ # トレーニングで使用する動画の説明ファイルをprompts.txtで指定\n --video_column videos.txt \\ # トレーニングで使用する動画のパスファイルをvideos.txtで指定\n --validation_prompt \"\" \\ # トレーニング中に検証用の動画を生成する際のプロンプト\n --validation_prompt_separator ::: \\ # 検証プロンプトの区切り文字を:::に設定\n --num_validation_videos 1 \\ # 各検証ラウンドで1本の動画を生成\n --validation_epochs 100 \\ # 100エポックごとに検証を実施\n --seed 42 \\ # 再現性を保証するためにランダムシードを42に設定\n --rank 128 \\ # LoRAのパラメータのランクを128に設定\n --lora_alpha 64 \\ # LoRAのalphaパラメータを64に設定し、LoRAの学習率を調整\n --mixed_precision bf16 \\ # bf16混合精度でトレーニングし、メモリを節約\n --output_dir $OUTPUT_PATH \\ # モデルの出力ディレクトリを$OUTPUT_PATHで指定\n --height 480 \\ # 動画の高さを480ピクセルに設定\n --width 720 \\ # 動画の幅を720ピクセルに設定\n --fps 8 \\ # 動画のフレームレートを1秒あたり8フレームに設定\n --max_num_frames 49 \\ # 各動画の最大フレーム数を49に設定\n --skip_frames_start 0 \\ # 動画の最初のフレームを0スキップ\n --skip_frames_end 0 \\ # 動画の最後のフレームを0スキップ\n --train_batch_size 4 \\ # トレーニングのバッチサイズを4に設定\n --num_train_epochs 30 \\ # 総トレーニングエポック数を30に設定\n --checkpointing_steps 1000 \\ # 1000ステップごとにモデルのチェックポイントを保存\n --gradient_accumulation_steps 1 \\ # 1ステップの勾配累積を行い、各バッチ後に更新\n --learning_rate 1e-3 \\ # 学習率を0.001に設定\n --lr_scheduler cosine_with_restarts \\ # リスタート付きのコサイン学習率スケジューラを使用\n --lr_warmup_steps 200 \\ # トレーニングの最初の200ステップで学習率をウォームアップ\n --lr_num_cycles 1 \\ # 学習率のサイクル数を1に設定\n --optimizer AdamW \\ # AdamWオプティマイザーを使用\n --adam_beta1 0.9 \\ # Adamオプティマイザーのbeta1パラメータを0.9に設定\n --adam_beta2 0.95 \\ # Adamオプティマイザーのbeta2パラメータを0.95に設定\n --max_grad_norm 1.0 \\ # 勾配クリッピングの最大値を1.0に設定\n --allow_tf32 \\ # トレーニングを高速化するためにTF32を有効化\n --report_to wandb # Weights and Biasesを使用してトレーニングの記録とモニタリングを行う\n```\n\n## 微調整を開始\n\n単一マシン (シングルGPU、マルチGPU) での微調整:\n\n```shell\nbash finetune_single_rank.sh\n```\n\n複数マシン・マルチGPUでの微調整：\n\n```shell\nbash finetune_multi_rank.sh # 各ノードで実行する必要があります。\n```\n\n## 微調整済みモデルのロード\n\n+ 微調整済みのモデルをロードする方法については、[cli_demo.py](../inference/cli_demo.py) を参照してください。\n\n## ベストプラクティス\n\n+ 解像度が `200 x 480 x 720`（フレーム数 x 高さ x 幅）のトレーニングビデオが70本含まれています。データ前処理でフレームをスキップすることで、49フレームと16フレームの小さなデータセットを作成しました。これは実験を加速するためのもので、CogVideoXチームが推奨する最大フレーム数制限は49フレームです。\n+ 25本以上のビデオが新しい概念やスタイルのトレーニングに最適です。\n+ 現在、`--id_token` を指定して識別トークンを使用してトレーニングする方が効果的です。これはDreamboothトレーニングに似ていますが、通常の微調整でも機能します。\n+ 元のリポジトリでは `lora_alpha` を1に設定していましたが、複数の実行でこの値が効果的でないことがわかりました。モデルのバックエンドやトレーニング設定によるかもしれません。私たちの提案は、lora_alphaをrankと同じか、rank // 2に設定することです。\n+ Rank 64以上の設定を推奨します。\n" }, { "path": "CogVideo/finetune/README_zh.md", "content": "# CogVideoX diffusers 微调方案\n\n[Read this in English](./README_zh.md)\n\n[日本語で読む](./README_ja.md)\n\n本功能尚未完全完善，如果您想查看SAT版本微调，请查看[这里](../sat/README_zh.md)。其数据集格式与本版本不同。\n\n## 硬件要求\n\n+ CogVideoX-2B / 5B T2V LORA: 1 * A100 (5B need to use `--use_8bit_adam`)\n+ CogVideoX-2B SFT: 8 * A100 (制作中)\n+ CogVideoX-5B-I2V 暂未支持\n\n## 安装依赖\n\n由于相关代码还没有被合并到diffusers发行版，你需要基于diffusers分支进行微调。请按照以下步骤安装依赖：\n\n```shell\ngit clone https://github.com/huggingface/diffusers.git\ncd diffusers # Now in Main branch\npip install -e .\n```\n\n## 准备数据集\n\n首先，你需要准备数据集，数据集格式如下，其中，videos.txt 存放 videos 中的视频。\n\n```\n.\n├── prompts.txt\n├── videos\n└── videos.txt\n```\n\n你可以从这里下载 [迪士尼汽船威利号](https://huggingface.co/datasets/Wild-Heart/Disney-VideoGeneration-Dataset)\n\n视频微调数据集作为测试微调。\n\n## 配置文件和运行\n\n`accelerate` 配置文件如下:\n\n+ accelerate_config_machine_multi.yaml 适合多GPU使用\n+ accelerate_config_machine_single.yaml 适合单GPU使用\n\n`finetune` 脚本配置文件如下：\n\n```shell\n\naccelerate launch --config_file accelerate_config_machine_single.yaml --multi_gpu \\ # 使用 accelerate 启动多GPU训练，配置文件为 accelerate_config_machine_single.yaml\n train_cogvideox_lora.py \\ # 运行的训练脚本为 train_cogvideox_lora.py，用于在 CogVideoX 模型上进行 LoRA 微调\n --gradient_checkpointing \\ # 启用梯度检查点功能，以减少显存使用\n --pretrained_model_name_or_path $MODEL_PATH \\ # 预训练模型路径，通过 $MODEL_PATH 指定\n --cache_dir $CACHE_PATH \\ # 模型缓存路径，由 $CACHE_PATH 指定\n --enable_tiling \\ # 启用tiling技术，以分片处理视频，节省显存\n --enable_slicing \\ # 启用slicing技术，将输入切片，以进一步优化内存\n --instance_data_root $DATASET_PATH \\ # 数据集路径，由 $DATASET_PATH 指定\n --caption_column prompts.txt \\ # 指定用于训练的视频描述文件，文件名为 prompts.txt\n --video_column videos.txt \\ # 指定用于训练的视频路径文件，文件名为 videos.txt\n --validation_prompt \"\" \\ # 验证集的提示语 (prompt)，用于在训练期间生成验证视频\n --validation_prompt_separator ::: \\ # 设置验证提示语的分隔符为 :::\n --num_validation_videos 1 \\ # 每个验证回合生成 1 个视频\n --validation_epochs 100 \\ # 每 100 个训练epoch进行一次验证\n --seed 42 \\ # 设置随机种子为 42，以保证结果的可复现性\n --rank 128 \\ # 设置 LoRA 参数的秩 (rank) 为 128\n --lora_alpha 64 \\ # 设置 LoRA 的 alpha 参数为 64，用于调整LoRA的学习率\n --mixed_precision bf16 \\ # 使用 bf16 混合精度进行训练，减少显存使用\n --output_dir $OUTPUT_PATH \\ # 指定模型输出目录，由 $OUTPUT_PATH 定义\n --height 480 \\ # 视频高度为 480 像素\n --width 720 \\ # 视频宽度为 720 像素\n --fps 8 \\ # 视频帧率设置为 8 帧每秒\n --max_num_frames 49 \\ # 每个视频的最大帧数为 49 帧\n --skip_frames_start 0 \\ # 跳过视频开头的帧数为 0\n --skip_frames_end 0 \\ # 跳过视频结尾的帧数为 0\n --train_batch_size 4 \\ # 训练时的 batch size 设置为 4\n --num_train_epochs 30 \\ # 总训练epoch数为 30\n --checkpointing_steps 1000 \\ # 每 1000 步保存一次模型检查点\n --gradient_accumulation_steps 1 \\ # 梯度累计步数为 1，即每个 batch 后都会更新梯度\n --learning_rate 1e-3 \\ # 学习率设置为 0.001\n --lr_scheduler cosine_with_restarts \\ # 使用带重启的余弦学习率调度器\n --lr_warmup_steps 200 \\ # 在训练的前 200 步进行学习率预热\n --lr_num_cycles 1 \\ # 学习率周期设置为 1\n --optimizer AdamW \\ # 使用 AdamW 优化器\n --adam_beta1 0.9 \\ # 设置 Adam 优化器的 beta1 参数为 0.9\n --adam_beta2 0.95 \\ # 设置 Adam 优化器的 beta2 参数为 0.95\n --max_grad_norm 1.0 \\ # 最大梯度裁剪值设置为 1.0\n --allow_tf32 \\ # 启用 TF32 以加速训练\n --report_to wandb # 使用 Weights and Biases 进行训练记录与监控\n```\n\n## 运行脚本，开始微调\n\n单机(单卡，多卡)微调：\n\n```shell\nbash finetune_single_rank.sh\n```\n\n多机多卡微调：\n\n```shell\nbash finetune_multi_rank.sh #需要在每个节点运行\n```\n\n## 载入微调的模型\n\n+ 请关注[cli_demo.py](../inference/cli_demo.py) 以了解如何加载微调的模型。\n\n## 最佳实践\n\n+ 包含70个分辨率为 `200 x 480 x 720`（帧数 x 高 x\n 宽）的训练视频。通过数据预处理中的帧跳过，我们创建了两个较小的49帧和16帧数据集，以加快实验速度，因为CogVideoX团队建议的最大帧数限制是49帧。我们将70个视频分成三组，分别为10、25和50个视频。这些视频的概念性质相似。\n+ 25个及以上的视频在训练新概念和风格时效果最佳。\n+ 现使用可以通过 `--id_token` 指定的标识符token进行训练效果更好。这类似于 Dreambooth 训练，但不使用这种token的常规微调也可以工作。\n+ 原始仓库使用 `lora_alpha` 设置为 1。我们发现这个值在多次运行中效果不佳，可能是因为模型后端和训练设置的不同。我们的建议是将\n lora_alpha 设置为与 rank 相同或 rank // 2。\n+ 建议使用 rank 为 64 及以上的设置。\n\n" }, { "path": "CogVideo/finetune/accelerate_config_machine_single.yaml", "content": "compute_environment: LOCAL_MACHINE\ndebug: false\ndeepspeed_config:\n gradient_accumulation_steps: 1\n gradient_clipping: 1.0\n offload_optimizer_device: none\n offload_param_device: none\n zero3_init_flag: false\n zero_stage: 2\ndistributed_type: DEEPSPEED\ndowncast_bf16: 'no'\nenable_cpu_affinity: false\nmachine_rank: 0\nmain_training_function: main\ndynamo_backend: 'no'\nmixed_precision: 'no'\nnum_machines: 1\nnum_processes: 8\nrdzv_backend: static\nsame_network: true\ntpu_env: []\ntpu_use_cluster: false\ntpu_use_sudo: false\nuse_cpu: false" }, { "path": "CogVideo/finetune/accelerate_config_machine_single_debug.yaml", "content": "compute_environment: LOCAL_MACHINE\ndebug: false\ndeepspeed_config:\n gradient_accumulation_steps: 1\n gradient_clipping: 1.0\n offload_optimizer_device: none\n offload_param_device: none\n zero3_init_flag: false\n zero_stage: 2\ndistributed_type: DEEPSPEED\ndowncast_bf16: 'no'\nenable_cpu_affinity: false\nmachine_rank: 0\nmain_training_function: main\ndynamo_backend: 'no'\nmixed_precision: 'no'\nnum_machines: 1\nnum_processes: 1\nrdzv_backend: static\nsame_network: true\ntpu_env: []\ntpu_use_cluster: false\ntpu_use_sudo: false\nuse_cpu: false" }, { "path": "CogVideo/finetune/finetune_single_rank_injector.sh", "content": "#!/bin/bash\n\nexport MODEL_PATH=\"/m2v_intern/fuxiao/CogVideo-release/weights/cogvideox-5b\" # Change it to CogVideoX-5B path\nexport TRANSFORMER_PATH=\"\" # Resume from pretrained injector checkpoint\nexport LORA_PATH=\"/m2v_intern/fuxiao/CogVideo-release/weights/lora\" # Change it to pretrained lora path\nexport CACHE_PATH=\"~/.cache\"\nexport DATASET_PATH=\"/ytech_m2v2_hdd/fuxiao/360Motion-Dataset\" # Change it to 360-Motion Dataset path\nexport OUTPUT_PATH=\"injector\"\nexport PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True\nexport CUDA_VISIBLE_DEVICES=\"0,1,2,3,4,5,6,7,\"\n\n# if you are not using wth 8 gus, change `accelerate_config_machine_single_debug.yaml` num_processes as your gpu number\naccelerate launch --config_file accelerate_config_machine_single.yaml --multi_gpu \\\n train_cogvideox_injector.py \\\n --gradient_checkpointing \\\n --pretrained_model_name_or_path $MODEL_PATH \\\n --lora_path $LORA_PATH \\\n --cache_dir $CACHE_PATH \\\n --enable_tiling \\\n --enable_slicing \\\n --finetune_init \\\n --instance_data_root $DATASET_PATH \\\n --validation_prompt \"a woman with short black wavy hair, lean figure, a green and yellow plaid shirt, dark brown pants, and black suede shoes and a robotic gazelle with a sturdy aluminum frame, an agile build, articulated legs and curved, metallic horns are moving in the city\" \\\n --validation_prompt_separator ::: \\\n --num_validation_videos 1 \\\n --validation_epochs 1 \\\n --block_interval 2 \\\n --seed 42 \\\n --lora_scale 1.0 \\\n --mixed_precision bf16 \\\n --output_dir $OUTPUT_PATH \\\n --height 480 \\\n --width 720 \\\n --fps 8 \\\n --max_num_frames 49 \\\n --skip_frames_start 0 \\\n --skip_frames_end 0 \\\n --train_batch_size 1 \\\n --num_train_epochs 1000 \\\n --checkpointing_steps 4000 \\\n --gradient_accumulation_steps 1 \\\n --learning_rate 1e-4 \\\n --lr_scheduler cosine_with_restarts \\\n --lr_warmup_steps 200 \\\n --lr_num_cycles 1 \\\n --enable_slicing \\\n --enable_tiling \\\n --gradient_checkpointing \\\n --optimizer AdamW \\\n --adam_beta1 0.9 \\\n --adam_beta2 0.95 \\\n --max_grad_norm 1.0 \\\n --allow_tf32 \\\n --report_to wandb\n\n # --resume_from_checkpoint $TRANSFORMER_PATH \\ " }, { "path": "CogVideo/finetune/finetune_single_rank_lora.sh", "content": "#!/bin/bash\n\nexport MODEL_PATH=\"/m2v_intern/fuxiao/CogVideo-release/weights/cogvideox-5b\" # Change it to CogVideoX-5B path\nexport CACHE_PATH=\"~/.cache\"\nexport DATASET_PATH=\"/ytech_m2v2_hdd/fuxiao/360Motion-Dataset\" # Change it to 360-Motion Dataset path\nexport OUTPUT_PATH=\"lora\"\nexport PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True\nexport CUDA_VISIBLE_DEVICES=\"0,1,2,3,4,5,6,7,\"\n\n# if you are not using wth 1 gpu, change `accelerate_config_machine_single_debug.yaml` num_processes as your gpu number\naccelerate launch --config_file accelerate_config_machine_single.yaml --multi_gpu \\\n train_cogvideox_lora.py \\\n --gradient_checkpointing \\\n --pretrained_model_name_or_path $MODEL_PATH \\\n --cache_dir $CACHE_PATH \\\n --enable_tiling \\\n --enable_slicing \\\n --instance_data_root $DATASET_PATH \\\n --validation_prompt \"a woman with short black wavy hair, lean figure, a green and yellow plaid shirt, dark brown pants, and black suede shoes and a robotic gazelle with a sturdy aluminum frame, an agile build, articulated legs and curved, metallic horns are moving in the city\" \\\n --validation_prompt_separator ::: \\\n --num_validation_videos 1 \\\n --validation_epochs 1 \\\n --seed 42 \\\n --rank 32 \\\n --lora_alpha 32 \\\n --mixed_precision bf16 \\\n --output_dir $OUTPUT_PATH \\\n --height 480 \\\n --width 720 \\\n --fps 8 \\\n --max_num_frames 49 \\\n --skip_frames_start 0 \\\n --skip_frames_end 0 \\\n --train_batch_size 2 \\\n --num_train_epochs 1000 \\\n --checkpointing_steps 1000 \\\n --gradient_accumulation_steps 1 \\\n --learning_rate 3e-4 \\\n --lr_scheduler cosine_with_restarts \\\n --lr_warmup_steps 200 \\\n --lr_num_cycles 1 \\\n --enable_slicing \\\n --enable_tiling \\\n --gradient_checkpointing \\\n --optimizer AdamW \\\n --adam_beta1 0.9 \\\n --adam_beta2 0.95 \\\n --max_grad_norm 1.0 \\\n --allow_tf32 \\\n --report_to wandb\n" }, { "path": "CogVideo/finetune/hostfile.txt", "content": "node1 slots=8\nnode2 slots=8" }, { "path": "CogVideo/finetune/models/attention.py", "content": "# Copyright 2024 The HuggingFace Team. All rights reserved.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n# http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\nfrom typing import Any, Dict, List, Optional, Tuple\n\nimport torch\nimport torch.nn.functional as F\nfrom torch import nn\n\nfrom diffusers.utils import deprecate, logging\nfrom diffusers.utils.torch_utils import maybe_allow_in_graph\nfrom diffusers.models.activations import GEGLU, GELU, ApproximateGELU, FP32SiLU, SwiGLU\nfrom models.attention_processor import Attention, JointAttnProcessor2_0\nfrom diffusers.models.embeddings import SinusoidalPositionalEmbedding\nfrom diffusers.models.normalization import AdaLayerNorm, AdaLayerNormContinuous, AdaLayerNormZero, RMSNorm, SD35AdaLayerNormZeroX\n\nlogger = logging.get_logger(__name__)\n\ndef _chunked_feed_forward(ff: nn.Module, hidden_states: torch.Tensor, chunk_dim: int, chunk_size: int):\n # \"feed_forward_chunk_size\" can be used to save memory\n if hidden_states.shape[chunk_dim] % chunk_size != 0:\n raise ValueError(\n f\"`hidden_states` dimension to be chunked: {hidden_states.shape[chunk_dim]} has to be divisible by chunk size: {chunk_size}. Make sure to set an appropriate `chunk_size` when calling `unet.enable_forward_chunking`.\"\n )\n\n num_chunks = hidden_states.shape[chunk_dim] // chunk_size\n ff_output = torch.cat(\n [ff(hid_slice) for hid_slice in hidden_states.chunk(num_chunks, dim=chunk_dim)],\n dim=chunk_dim,\n )\n return ff_output\n\n\n@maybe_allow_in_graph\nclass GatedSelfAttentionDense(nn.Module):\n r\"\"\"\n A gated self-attention dense layer that combines visual features and object features.\n\n Parameters:\n query_dim (`int`): The number of channels in the query.\n context_dim (`int`): The number of channels in the context.\n n_heads (`int`): The number of heads to use for attention.\n d_head (`int`): The number of channels in each head.\n \"\"\"\n\n def __init__(self, query_dim: int, context_dim: int, n_heads: int, d_head: int):\n super().__init__()\n\n # we need a linear projection since we need cat visual feature and obj feature\n self.linear = nn.Linear(context_dim, query_dim)\n\n self.attn = Attention(query_dim=query_dim, heads=n_heads, dim_head=d_head)\n self.ff = FeedForward(query_dim, activation_fn=\"geglu\")\n\n self.norm1 = nn.LayerNorm(query_dim)\n self.norm2 = nn.LayerNorm(query_dim)\n\n self.register_parameter(\"alpha_attn\", nn.Parameter(torch.tensor(0.0)))\n self.register_parameter(\"alpha_dense\", nn.Parameter(torch.tensor(0.0)))\n\n self.enabled = True\n\n def forward(self, x: torch.Tensor, objs: torch.Tensor) -> torch.Tensor:\n if not self.enabled:\n return x\n\n n_visual = x.shape[1]\n objs = self.linear(objs)\n\n x = x + self.alpha_attn.tanh() * self.attn(self.norm1(torch.cat([x, objs], dim=1)))[:, :n_visual, :]\n x = x + self.alpha_dense.tanh() * self.ff(self.norm2(x))\n\n return x\n\n\n@maybe_allow_in_graph\nclass JointTransformerBlock(nn.Module):\n r\"\"\"\n A Transformer block following the MMDiT architecture, introduced in Stable Diffusion 3.\n\n Reference: https://arxiv.org/abs/2403.03206\n\n Parameters:\n dim (`int`): The number of channels in the input and output.\n num_attention_heads (`int`): The number of heads to use for multi-head attention.\n attention_head_dim (`int`): The number of channels in each head.\n context_pre_only (`bool`): Boolean to determine if we should add some blocks associated with the\n processing of `context` conditions.\n \"\"\"\n\n def __init__(\n self,\n dim: int,\n num_attention_heads: int,\n attention_head_dim: int,\n context_pre_only: bool = False,\n qk_norm: Optional[str] = None,\n use_dual_attention: bool = False,\n ):\n super().__init__()\n\n self.use_dual_attention = use_dual_attention\n self.context_pre_only = context_pre_only\n context_norm_type = \"ada_norm_continous\" if context_pre_only else \"ada_norm_zero\"\n\n if use_dual_attention:\n self.norm1 = SD35AdaLayerNormZeroX(dim)\n else:\n self.norm1 = AdaLayerNormZero(dim)\n\n if context_norm_type == \"ada_norm_continous\":\n self.norm1_context = AdaLayerNormContinuous(\n dim, dim, elementwise_affine=False, eps=1e-6, bias=True, norm_type=\"layer_norm\"\n )\n elif context_norm_type == \"ada_norm_zero\":\n self.norm1_context = AdaLayerNormZero(dim)\n else:\n raise ValueError(\n f\"Unknown context_norm_type: {context_norm_type}, currently only support `ada_norm_continous`, `ada_norm_zero`\"\n )\n\n if hasattr(F, \"scaled_dot_product_attention\"):\n processor = JointAttnProcessor2_0()\n else:\n raise ValueError(\n \"The current PyTorch version does not support the `scaled_dot_product_attention` function.\"\n )\n\n self.attn = Attention(\n query_dim=dim,\n cross_attention_dim=None,\n added_kv_proj_dim=dim,\n dim_head=attention_head_dim,\n heads=num_attention_heads,\n out_dim=dim,\n context_pre_only=context_pre_only,\n bias=True,\n processor=processor,\n qk_norm=qk_norm,\n eps=1e-6,\n )\n\n if use_dual_attention:\n self.attn2 = Attention(\n query_dim=dim,\n cross_attention_dim=None,\n dim_head=attention_head_dim,\n heads=num_attention_heads,\n out_dim=dim,\n bias=True,\n processor=processor,\n qk_norm=qk_norm,\n eps=1e-6,\n )\n else:\n self.attn2 = None\n\n self.norm2 = nn.LayerNorm(dim, elementwise_affine=False, eps=1e-6)\n self.ff = FeedForward(dim=dim, dim_out=dim, activation_fn=\"gelu-approximate\")\n\n if not context_pre_only:\n self.norm2_context = nn.LayerNorm(dim, elementwise_affine=False, eps=1e-6)\n self.ff_context = FeedForward(dim=dim, dim_out=dim, activation_fn=\"gelu-approximate\")\n else:\n self.norm2_context = None\n self.ff_context = None\n\n # let chunk size default to None\n self._chunk_size = None\n self._chunk_dim = 0\n\n # Copied from diffusers.models.attention.BasicTransformerBlock.set_chunk_feed_forward\n def set_chunk_feed_forward(self, chunk_size: Optional[int], dim: int = 0):\n # Sets chunk feed-forward\n self._chunk_size = chunk_size\n self._chunk_dim = dim\n\n def forward(\n self, hidden_states: torch.FloatTensor, encoder_hidden_states: torch.FloatTensor, temb: torch.FloatTensor\n ):\n if self.use_dual_attention:\n norm_hidden_states, gate_msa, shift_mlp, scale_mlp, gate_mlp, norm_hidden_states2, gate_msa2 = self.norm1(\n hidden_states, emb=temb\n )\n else:\n norm_hidden_states, gate_msa, shift_mlp, scale_mlp, gate_mlp = self.norm1(hidden_states, emb=temb)\n\n if self.context_pre_only:\n norm_encoder_hidden_states = self.norm1_context(encoder_hidden_states, temb)\n else:\n norm_encoder_hidden_states, c_gate_msa, c_shift_mlp, c_scale_mlp, c_gate_mlp = self.norm1_context(\n encoder_hidden_states, emb=temb\n )\n\n # Attention.\n attn_output, context_attn_output = self.attn(\n hidden_states=norm_hidden_states, encoder_hidden_states=norm_encoder_hidden_states\n )\n\n # Process attention outputs for the `hidden_states`.\n attn_output = gate_msa.unsqueeze(1) * attn_output\n hidden_states = hidden_states + attn_output\n\n if self.use_dual_attention:\n attn_output2 = self.attn2(hidden_states=norm_hidden_states2)\n attn_output2 = gate_msa2.unsqueeze(1) * attn_output2\n hidden_states = hidden_states + attn_output2\n\n norm_hidden_states = self.norm2(hidden_states)\n norm_hidden_states = norm_hidden_states * (1 + scale_mlp[:, None]) + shift_mlp[:, None]\n if self._chunk_size is not None:\n # \"feed_forward_chunk_size\" can be used to save memory\n ff_output = _chunked_feed_forward(self.ff, norm_hidden_states, self._chunk_dim, self._chunk_size)\n else:\n ff_output = self.ff(norm_hidden_states)\n ff_output = gate_mlp.unsqueeze(1) * ff_output\n\n hidden_states = hidden_states + ff_output\n\n # Process attention outputs for the `encoder_hidden_states`.\n if self.context_pre_only:\n encoder_hidden_states = None\n else:\n context_attn_output = c_gate_msa.unsqueeze(1) * context_attn_output\n encoder_hidden_states = encoder_hidden_states + context_attn_output\n\n norm_encoder_hidden_states = self.norm2_context(encoder_hidden_states)\n norm_encoder_hidden_states = norm_encoder_hidden_states * (1 + c_scale_mlp[:, None]) + c_shift_mlp[:, None]\n if self._chunk_size is not None:\n # \"feed_forward_chunk_size\" can be used to save memory\n context_ff_output = _chunked_feed_forward(\n self.ff_context, norm_encoder_hidden_states, self._chunk_dim, self._chunk_size\n )\n else:\n context_ff_output = self.ff_context(norm_encoder_hidden_states)\n encoder_hidden_states = encoder_hidden_states + c_gate_mlp.unsqueeze(1) * context_ff_output\n\n return encoder_hidden_states, hidden_states\n\n\n@maybe_allow_in_graph\nclass BasicTransformerBlock(nn.Module):\n r\"\"\"\n A basic Transformer block.\n\n Parameters:\n dim (`int`): The number of channels in the input and output.\n num_attention_heads (`int`): The number of heads to use for multi-head attention.\n attention_head_dim (`int`): The number of channels in each head.\n dropout (`float`, *optional*, defaults to 0.0): The dropout probability to use.\n cross_attention_dim (`int`, *optional*): The size of the encoder_hidden_states vector for cross attention.\n activation_fn (`str`, *optional*, defaults to `\"geglu\"`): Activation function to be used in feed-forward.\n num_embeds_ada_norm (:\n obj: `int`, *optional*): The number of diffusion steps used during training. See `Transformer2DModel`.\n attention_bias (:\n obj: `bool`, *optional*, defaults to `False`): Configure if the attentions should contain a bias parameter.\n only_cross_attention (`bool`, *optional*):\n Whether to use only cross-attention layers. In this case two cross attention layers are used.\n double_self_attention (`bool`, *optional*):\n Whether to use two self-attention layers. In this case no cross attention layers are used.\n upcast_attention (`bool`, *optional*):\n Whether to upcast the attention computation to float32. This is useful for mixed precision training.\n norm_elementwise_affine (`bool`, *optional*, defaults to `True`):\n Whether to use learnable elementwise affine parameters for normalization.\n norm_type (`str`, *optional*, defaults to `\"layer_norm\"`):\n The normalization layer to use. Can be `\"layer_norm\"`, `\"ada_norm\"` or `\"ada_norm_zero\"`.\n final_dropout (`bool` *optional*, defaults to False):\n Whether to apply a final dropout after the last feed-forward layer.\n attention_type (`str`, *optional*, defaults to `\"default\"`):\n The type of attention to use. Can be `\"default\"` or `\"gated\"` or `\"gated-text-image\"`.\n positional_embeddings (`str`, *optional*, defaults to `None`):\n The type of positional embeddings to apply to.\n num_positional_embeddings (`int`, *optional*, defaults to `None`):\n The maximum number of positional embeddings to apply.\n \"\"\"\n\n def __init__(\n self,\n dim: int,\n num_attention_heads: int,\n attention_head_dim: int,\n dropout=0.0,\n cross_attention_dim: Optional[int] = None,\n activation_fn: str = \"geglu\",\n num_embeds_ada_norm: Optional[int] = None,\n attention_bias: bool = False,\n only_cross_attention: bool = False,\n double_self_attention: bool = False,\n upcast_attention: bool = False,\n norm_elementwise_affine: bool = True,\n norm_type: str = \"layer_norm\", # 'layer_norm', 'ada_norm', 'ada_norm_zero', 'ada_norm_single', 'ada_norm_continuous', 'layer_norm_i2vgen'\n norm_eps: float = 1e-5,\n final_dropout: bool = False,\n attention_type: str = \"default\",\n positional_embeddings: Optional[str] = None,\n num_positional_embeddings: Optional[int] = None,\n ada_norm_continous_conditioning_embedding_dim: Optional[int] = None,\n ada_norm_bias: Optional[int] = None,\n ff_inner_dim: Optional[int] = None,\n ff_bias: bool = True,\n attention_out_bias: bool = True,\n ):\n super().__init__()\n self.dim = dim\n self.num_attention_heads = num_attention_heads\n self.attention_head_dim = attention_head_dim\n self.dropout = dropout\n self.cross_attention_dim = cross_attention_dim\n self.activation_fn = activation_fn\n self.attention_bias = attention_bias\n self.double_self_attention = double_self_attention\n self.norm_elementwise_affine = norm_elementwise_affine\n self.positional_embeddings = positional_embeddings\n self.num_positional_embeddings = num_positional_embeddings\n self.only_cross_attention = only_cross_attention\n\n # We keep these boolean flags for backward-compatibility.\n self.use_ada_layer_norm_zero = (num_embeds_ada_norm is not None) and norm_type == \"ada_norm_zero\"\n self.use_ada_layer_norm = (num_embeds_ada_norm is not None) and norm_type == \"ada_norm\"\n self.use_ada_layer_norm_single = norm_type == \"ada_norm_single\"\n self.use_layer_norm = norm_type == \"layer_norm\"\n self.use_ada_layer_norm_continuous = norm_type == \"ada_norm_continuous\"\n\n if norm_type in (\"ada_norm\", \"ada_norm_zero\") and num_embeds_ada_norm is None:\n raise ValueError(\n f\"`norm_type` is set to {norm_type}, but `num_embeds_ada_norm` is not defined. Please make sure to\"\n f\" define `num_embeds_ada_norm` if setting `norm_type` to {norm_type}.\"\n )\n\n self.norm_type = norm_type\n self.num_embeds_ada_norm = num_embeds_ada_norm\n\n if positional_embeddings and (num_positional_embeddings is None):\n raise ValueError(\n \"If `positional_embedding` type is defined, `num_positition_embeddings` must also be defined.\"\n )\n\n if positional_embeddings == \"sinusoidal\":\n self.pos_embed = SinusoidalPositionalEmbedding(dim, max_seq_length=num_positional_embeddings)\n else:\n self.pos_embed = None\n\n # Define 3 blocks. Each block has its own normalization layer.\n # 1. Self-Attn\n if norm_type == \"ada_norm\":\n self.norm1 = AdaLayerNorm(dim, num_embeds_ada_norm)\n elif norm_type == \"ada_norm_zero\":\n self.norm1 = AdaLayerNormZero(dim, num_embeds_ada_norm)\n elif norm_type == \"ada_norm_continuous\":\n self.norm1 = AdaLayerNormContinuous(\n dim,\n ada_norm_continous_conditioning_embedding_dim,\n norm_elementwise_affine,\n norm_eps,\n ada_norm_bias,\n \"rms_norm\",\n )\n else:\n self.norm1 = nn.LayerNorm(dim, elementwise_affine=norm_elementwise_affine, eps=norm_eps)\n\n self.attn1 = Attention(\n query_dim=dim,\n heads=num_attention_heads,\n dim_head=attention_head_dim,\n dropout=dropout,\n bias=attention_bias,\n cross_attention_dim=cross_attention_dim if only_cross_attention else None,\n upcast_attention=upcast_attention,\n out_bias=attention_out_bias,\n )\n\n # 2. Cross-Attn\n if cross_attention_dim is not None or double_self_attention:\n # We currently only use AdaLayerNormZero for self attention where there will only be one attention block.\n # I.e. the number of returned modulation chunks from AdaLayerZero would not make sense if returned during\n # the second cross attention block.\n if norm_type == \"ada_norm\":\n self.norm2 = AdaLayerNorm(dim, num_embeds_ada_norm)\n elif norm_type == \"ada_norm_continuous\":\n self.norm2 = AdaLayerNormContinuous(\n dim,\n ada_norm_continous_conditioning_embedding_dim,\n norm_elementwise_affine,\n norm_eps,\n ada_norm_bias,\n \"rms_norm\",\n )\n else:\n self.norm2 = nn.LayerNorm(dim, norm_eps, norm_elementwise_affine)\n\n self.attn2 = Attention(\n query_dim=dim,\n cross_attention_dim=cross_attention_dim if not double_self_attention else None,\n heads=num_attention_heads,\n dim_head=attention_head_dim,\n dropout=dropout,\n bias=attention_bias,\n upcast_attention=upcast_attention,\n out_bias=attention_out_bias,\n ) # is self-attn if encoder_hidden_states is none\n else:\n if norm_type == \"ada_norm_single\": # For Latte\n self.norm2 = nn.LayerNorm(dim, norm_eps, norm_elementwise_affine)\n else:\n self.norm2 = None\n self.attn2 = None\n\n # 3. Feed-forward\n if norm_type == \"ada_norm_continuous\":\n self.norm3 = AdaLayerNormContinuous(\n dim,\n ada_norm_continous_conditioning_embedding_dim,\n norm_elementwise_affine,\n norm_eps,\n ada_norm_bias,\n \"layer_norm\",\n )\n\n elif norm_type in [\"ada_norm_zero\", \"ada_norm\", \"layer_norm\"]:\n self.norm3 = nn.LayerNorm(dim, norm_eps, norm_elementwise_affine)\n elif norm_type == \"layer_norm_i2vgen\":\n self.norm3 = None\n\n self.ff = FeedForward(\n dim,\n dropout=dropout,\n activation_fn=activation_fn,\n final_dropout=final_dropout,\n inner_dim=ff_inner_dim,\n bias=ff_bias,\n )\n\n # 4. Fuser\n if attention_type == \"gated\" or attention_type == \"gated-text-image\":\n self.fuser = GatedSelfAttentionDense(dim, cross_attention_dim, num_attention_heads, attention_head_dim)\n\n # 5. Scale-shift for PixArt-Alpha.\n if norm_type == \"ada_norm_single\":\n self.scale_shift_table = nn.Parameter(torch.randn(6, dim) / dim**0.5)\n\n # let chunk size default to None\n self._chunk_size = None\n self._chunk_dim = 0\n\n def set_chunk_feed_forward(self, chunk_size: Optional[int], dim: int = 0):\n # Sets chunk feed-forward\n self._chunk_size = chunk_size\n self._chunk_dim = dim\n\n def forward(\n self,\n hidden_states: torch.Tensor,\n attention_mask: Optional[torch.Tensor] = None,\n encoder_hidden_states: Optional[torch.Tensor] = None,\n encoder_attention_mask: Optional[torch.Tensor] = None,\n timestep: Optional[torch.LongTensor] = None,\n cross_attention_kwargs: Dict[str, Any] = None,\n class_labels: Optional[torch.LongTensor] = None,\n added_cond_kwargs: Optional[Dict[str, torch.Tensor]] = None,\n ) -> torch.Tensor:\n if cross_attention_kwargs is not None:\n if cross_attention_kwargs.get(\"scale\", None) is not None:\n logger.warning(\"Passing `scale` to `cross_attention_kwargs` is deprecated. `scale` will be ignored.\")\n\n # Notice that normalization is always applied before the real computation in the following blocks.\n # 0. Self-Attention\n batch_size = hidden_states.shape[0]\n\n if self.norm_type == \"ada_norm\":\n norm_hidden_states = self.norm1(hidden_states, timestep)\n elif self.norm_type == \"ada_norm_zero\":\n norm_hidden_states, gate_msa, shift_mlp, scale_mlp, gate_mlp = self.norm1(\n hidden_states, timestep, class_labels, hidden_dtype=hidden_states.dtype\n )\n elif self.norm_type in [\"layer_norm\", \"layer_norm_i2vgen\"]:\n norm_hidden_states = self.norm1(hidden_states)\n elif self.norm_type == \"ada_norm_continuous\":\n norm_hidden_states = self.norm1(hidden_states, added_cond_kwargs[\"pooled_text_emb\"])\n elif self.norm_type == \"ada_norm_single\":\n shift_msa, scale_msa, gate_msa, shift_mlp, scale_mlp, gate_mlp = (\n self.scale_shift_table[None] + timestep.reshape(batch_size, 6, -1)\n ).chunk(6, dim=1)\n norm_hidden_states = self.norm1(hidden_states)\n norm_hidden_states = norm_hidden_states * (1 + scale_msa) + shift_msa\n else:\n raise ValueError(\"Incorrect norm used\")\n\n if self.pos_embed is not None:\n norm_hidden_states = self.pos_embed(norm_hidden_states)\n\n # 1. Prepare GLIGEN inputs\n cross_attention_kwargs = cross_attention_kwargs.copy() if cross_attention_kwargs is not None else {}\n gligen_kwargs = cross_attention_kwargs.pop(\"gligen\", None)\n\n attn_output = self.attn1(\n norm_hidden_states,\n encoder_hidden_states=encoder_hidden_states if self.only_cross_attention else None,\n attention_mask=attention_mask,\n **cross_attention_kwargs,\n )\n\n if self.norm_type == \"ada_norm_zero\":\n attn_output = gate_msa.unsqueeze(1) * attn_output\n elif self.norm_type == \"ada_norm_single\":\n attn_output = gate_msa * attn_output\n\n hidden_states = attn_output + hidden_states\n if hidden_states.ndim == 4:\n hidden_states = hidden_states.squeeze(1)\n\n # 1.2 GLIGEN Control\n if gligen_kwargs is not None:\n hidden_states = self.fuser(hidden_states, gligen_kwargs[\"objs\"])\n\n # 3. Cross-Attention\n if self.attn2 is not None:\n if self.norm_type == \"ada_norm\":\n norm_hidden_states = self.norm2(hidden_states, timestep)\n elif self.norm_type in [\"ada_norm_zero\", \"layer_norm\", \"layer_norm_i2vgen\"]:\n norm_hidden_states = self.norm2(hidden_states)\n elif self.norm_type == \"ada_norm_single\":\n # For PixArt norm2 isn't applied here:\n # https://github.com/PixArt-alpha/PixArt-alpha/blob/0f55e922376d8b797edd44d25d0e7464b260dcab/diffusion/model/nets/PixArtMS.py#L70C1-L76C103\n norm_hidden_states = hidden_states\n elif self.norm_type == \"ada_norm_continuous\":\n norm_hidden_states = self.norm2(hidden_states, added_cond_kwargs[\"pooled_text_emb\"])\n else:\n raise ValueError(\"Incorrect norm\")\n\n if self.pos_embed is not None and self.norm_type != \"ada_norm_single\":\n norm_hidden_states = self.pos_embed(norm_hidden_states)\n\n attn_output = self.attn2(\n norm_hidden_states,\n encoder_hidden_states=encoder_hidden_states,\n attention_mask=encoder_attention_mask,\n **cross_attention_kwargs,\n )\n hidden_states = attn_output + hidden_states\n\n # 4. Feed-forward\n # i2vgen doesn't have this norm 🤷‍♂️\n if self.norm_type == \"ada_norm_continuous\":\n norm_hidden_states = self.norm3(hidden_states, added_cond_kwargs[\"pooled_text_emb\"])\n elif not self.norm_type == \"ada_norm_single\":\n norm_hidden_states = self.norm3(hidden_states)\n\n if self.norm_type == \"ada_norm_zero\":\n norm_hidden_states = norm_hidden_states * (1 + scale_mlp[:, None]) + shift_mlp[:, None]\n\n if self.norm_type == \"ada_norm_single\":\n norm_hidden_states = self.norm2(hidden_states)\n norm_hidden_states = norm_hidden_states * (1 + scale_mlp) + shift_mlp\n\n if self._chunk_size is not None:\n # \"feed_forward_chunk_size\" can be used to save memory\n ff_output = _chunked_feed_forward(self.ff, norm_hidden_states, self._chunk_dim, self._chunk_size)\n else:\n ff_output = self.ff(norm_hidden_states)\n\n if self.norm_type == \"ada_norm_zero\":\n ff_output = gate_mlp.unsqueeze(1) * ff_output\n elif self.norm_type == \"ada_norm_single\":\n ff_output = gate_mlp * ff_output\n\n hidden_states = ff_output + hidden_states\n if hidden_states.ndim == 4:\n hidden_states = hidden_states.squeeze(1)\n\n return hidden_states\n\n\nclass LuminaFeedForward(nn.Module):\n r\"\"\"\n A feed-forward layer.\n\n Parameters:\n hidden_size (`int`):\n The dimensionality of the hidden layers in the model. This parameter determines the width of the model's\n hidden representations.\n intermediate_size (`int`): The intermediate dimension of the feedforward layer.\n multiple_of (`int`, *optional*): Value to ensure hidden dimension is a multiple\n of this value.\n ffn_dim_multiplier (float, *optional*): Custom multiplier for hidden\n dimension. Defaults to None.\n \"\"\"\n\n def __init__(\n self,\n dim: int,\n inner_dim: int,\n multiple_of: Optional[int] = 256,\n ffn_dim_multiplier: Optional[float] = None,\n ):\n super().__init__()\n inner_dim = int(2 * inner_dim / 3)\n # custom hidden_size factor multiplier\n if ffn_dim_multiplier is not None:\n inner_dim = int(ffn_dim_multiplier * inner_dim)\n inner_dim = multiple_of * ((inner_dim + multiple_of - 1) // multiple_of)\n\n self.linear_1 = nn.Linear(\n dim,\n inner_dim,\n bias=False,\n )\n self.linear_2 = nn.Linear(\n inner_dim,\n dim,\n bias=False,\n )\n self.linear_3 = nn.Linear(\n dim,\n inner_dim,\n bias=False,\n )\n self.silu = FP32SiLU()\n\n def forward(self, x):\n return self.linear_2(self.silu(self.linear_1(x)) * self.linear_3(x))\n\n\n@maybe_allow_in_graph\nclass TemporalBasicTransformerBlock(nn.Module):\n r\"\"\"\n A basic Transformer block for video like data.\n\n Parameters:\n dim (`int`): The number of channels in the input and output.\n time_mix_inner_dim (`int`): The number of channels for temporal attention.\n num_attention_heads (`int`): The number of heads to use for multi-head attention.\n attention_head_dim (`int`): The number of channels in each head.\n cross_attention_dim (`int`, *optional*): The size of the encoder_hidden_states vector for cross attention.\n \"\"\"\n\n def __init__(\n self,\n dim: int,\n time_mix_inner_dim: int,\n num_attention_heads: int,\n attention_head_dim: int,\n cross_attention_dim: Optional[int] = None,\n ):\n super().__init__()\n self.is_res = dim == time_mix_inner_dim\n\n self.norm_in = nn.LayerNorm(dim)\n\n # Define 3 blocks. Each block has its own normalization layer.\n # 1. Self-Attn\n self.ff_in = FeedForward(\n dim,\n dim_out=time_mix_inner_dim,\n activation_fn=\"geglu\",\n )\n\n self.norm1 = nn.LayerNorm(time_mix_inner_dim)\n self.attn1 = Attention(\n query_dim=time_mix_inner_dim,\n heads=num_attention_heads,\n dim_head=attention_head_dim,\n cross_attention_dim=None,\n )\n\n # 2. Cross-Attn\n if cross_attention_dim is not None:\n # We currently only use AdaLayerNormZero for self attention where there will only be one attention block.\n # I.e. the number of returned modulation chunks from AdaLayerZero would not make sense if returned during\n # the second cross attention block.\n self.norm2 = nn.LayerNorm(time_mix_inner_dim)\n self.attn2 = Attention(\n query_dim=time_mix_inner_dim,\n cross_attention_dim=cross_attention_dim,\n heads=num_attention_heads,\n dim_head=attention_head_dim,\n ) # is self-attn if encoder_hidden_states is none\n else:\n self.norm2 = None\n self.attn2 = None\n\n # 3. Feed-forward\n self.norm3 = nn.LayerNorm(time_mix_inner_dim)\n self.ff = FeedForward(time_mix_inner_dim, activation_fn=\"geglu\")\n\n # let chunk size default to None\n self._chunk_size = None\n self._chunk_dim = None\n\n def set_chunk_feed_forward(self, chunk_size: Optional[int], **kwargs):\n # Sets chunk feed-forward\n self._chunk_size = chunk_size\n # chunk dim should be hardcoded to 1 to have better speed vs. memory trade-off\n self._chunk_dim = 1\n\n def forward(\n self,\n hidden_states: torch.Tensor,\n num_frames: int,\n encoder_hidden_states: Optional[torch.Tensor] = None,\n ) -> torch.Tensor:\n # Notice that normalization is always applied before the real computation in the following blocks.\n # 0. Self-Attention\n batch_size = hidden_states.shape[0]\n\n batch_frames, seq_length, channels = hidden_states.shape\n batch_size = batch_frames // num_frames\n\n hidden_states = hidden_states[None, :].reshape(batch_size, num_frames, seq_length, channels)\n hidden_states = hidden_states.permute(0, 2, 1, 3)\n hidden_states = hidden_states.reshape(batch_size * seq_length, num_frames, channels)\n\n residual = hidden_states\n hidden_states = self.norm_in(hidden_states)\n\n if self._chunk_size is not None:\n hidden_states = _chunked_feed_forward(self.ff_in, hidden_states, self._chunk_dim, self._chunk_size)\n else:\n hidden_states = self.ff_in(hidden_states)\n\n if self.is_res:\n hidden_states = hidden_states + residual\n\n norm_hidden_states = self.norm1(hidden_states)\n attn_output = self.attn1(norm_hidden_states, encoder_hidden_states=None)\n hidden_states = attn_output + hidden_states\n\n # 3. Cross-Attention\n if self.attn2 is not None:\n norm_hidden_states = self.norm2(hidden_states)\n attn_output = self.attn2(norm_hidden_states, encoder_hidden_states=encoder_hidden_states)\n hidden_states = attn_output + hidden_states\n\n # 4. Feed-forward\n norm_hidden_states = self.norm3(hidden_states)\n\n if self._chunk_size is not None:\n ff_output = _chunked_feed_forward(self.ff, norm_hidden_states, self._chunk_dim, self._chunk_size)\n else:\n ff_output = self.ff(norm_hidden_states)\n\n if self.is_res:\n hidden_states = ff_output + hidden_states\n else:\n hidden_states = ff_output\n\n hidden_states = hidden_states[None, :].reshape(batch_size, seq_length, num_frames, channels)\n hidden_states = hidden_states.permute(0, 2, 1, 3)\n hidden_states = hidden_states.reshape(batch_size * num_frames, seq_length, channels)\n\n return hidden_states\n\n\nclass SkipFFTransformerBlock(nn.Module):\n def __init__(\n self,\n dim: int,\n num_attention_heads: int,\n attention_head_dim: int,\n kv_input_dim: int,\n kv_input_dim_proj_use_bias: bool,\n dropout=0.0,\n cross_attention_dim: Optional[int] = None,\n attention_bias: bool = False,\n attention_out_bias: bool = True,\n ):\n super().__init__()\n if kv_input_dim != dim:\n self.kv_mapper = nn.Linear(kv_input_dim, dim, kv_input_dim_proj_use_bias)\n else:\n self.kv_mapper = None\n\n self.norm1 = RMSNorm(dim, 1e-06)\n\n self.attn1 = Attention(\n query_dim=dim,\n heads=num_attention_heads,\n dim_head=attention_head_dim,\n dropout=dropout,\n bias=attention_bias,\n cross_attention_dim=cross_attention_dim,\n out_bias=attention_out_bias,\n )\n\n self.norm2 = RMSNorm(dim, 1e-06)\n\n self.attn2 = Attention(\n query_dim=dim,\n cross_attention_dim=cross_attention_dim,\n heads=num_attention_heads,\n dim_head=attention_head_dim,\n dropout=dropout,\n bias=attention_bias,\n out_bias=attention_out_bias,\n )\n\n def forward(self, hidden_states, encoder_hidden_states, cross_attention_kwargs):\n cross_attention_kwargs = cross_attention_kwargs.copy() if cross_attention_kwargs is not None else {}\n\n if self.kv_mapper is not None:\n encoder_hidden_states = self.kv_mapper(F.silu(encoder_hidden_states))\n\n norm_hidden_states = self.norm1(hidden_states)\n\n attn_output = self.attn1(\n norm_hidden_states,\n encoder_hidden_states=encoder_hidden_states,\n **cross_attention_kwargs,\n )\n\n hidden_states = attn_output + hidden_states\n\n norm_hidden_states = self.norm2(hidden_states)\n\n attn_output = self.attn2(\n norm_hidden_states,\n encoder_hidden_states=encoder_hidden_states,\n **cross_attention_kwargs,\n )\n\n hidden_states = attn_output + hidden_states\n\n return hidden_states\n\n\n@maybe_allow_in_graph\nclass FreeNoiseTransformerBlock(nn.Module):\n r\"\"\"\n A FreeNoise Transformer block.\n\n Parameters:\n dim (`int`):\n The number of channels in the input and output.\n num_attention_heads (`int`):\n The number of heads to use for multi-head attention.\n attention_head_dim (`int`):\n The number of channels in each head.\n dropout (`float`, *optional*, defaults to 0.0):\n The dropout probability to use.\n cross_attention_dim (`int`, *optional*):\n The size of the encoder_hidden_states vector for cross attention.\n activation_fn (`str`, *optional*, defaults to `\"geglu\"`):\n Activation function to be used in feed-forward.\n num_embeds_ada_norm (`int`, *optional*):\n The number of diffusion steps used during training. See `Transformer2DModel`.\n attention_bias (`bool`, defaults to `False`):\n Configure if the attentions should contain a bias parameter.\n only_cross_attention (`bool`, defaults to `False`):\n Whether to use only cross-attention layers. In this case two cross attention layers are used.\n double_self_attention (`bool`, defaults to `False`):\n Whether to use two self-attention layers. In this case no cross attention layers are used.\n upcast_attention (`bool`, defaults to `False`):\n Whether to upcast the attention computation to float32. This is useful for mixed precision training.\n norm_elementwise_affine (`bool`, defaults to `True`):\n Whether to use learnable elementwise affine parameters for normalization.\n norm_type (`str`, defaults to `\"layer_norm\"`):\n The normalization layer to use. Can be `\"layer_norm\"`, `\"ada_norm\"` or `\"ada_norm_zero\"`.\n final_dropout (`bool` defaults to `False`):\n Whether to apply a final dropout after the last feed-forward layer.\n attention_type (`str`, defaults to `\"default\"`):\n The type of attention to use. Can be `\"default\"` or `\"gated\"` or `\"gated-text-image\"`.\n positional_embeddings (`str`, *optional*):\n The type of positional embeddings to apply to.\n num_positional_embeddings (`int`, *optional*, defaults to `None`):\n The maximum number of positional embeddings to apply.\n ff_inner_dim (`int`, *optional*):\n Hidden dimension of feed-forward MLP.\n ff_bias (`bool`, defaults to `True`):\n Whether or not to use bias in feed-forward MLP.\n attention_out_bias (`bool`, defaults to `True`):\n Whether or not to use bias in attention output project layer.\n context_length (`int`, defaults to `16`):\n The maximum number of frames that the FreeNoise block processes at once.\n context_stride (`int`, defaults to `4`):\n The number of frames to be skipped before starting to process a new batch of `context_length` frames.\n weighting_scheme (`str`, defaults to `\"pyramid\"`):\n The weighting scheme to use for weighting averaging of processed latent frames. As described in the\n Equation 9. of the [FreeNoise](https://arxiv.org/abs/2310.15169) paper, \"pyramid\" is the default setting\n used.\n \"\"\"\n\n def __init__(\n self,\n dim: int,\n num_attention_heads: int,\n attention_head_dim: int,\n dropout: float = 0.0,\n cross_attention_dim: Optional[int] = None,\n activation_fn: str = \"geglu\",\n num_embeds_ada_norm: Optional[int] = None,\n attention_bias: bool = False,\n only_cross_attention: bool = False,\n double_self_attention: bool = False,\n upcast_attention: bool = False,\n norm_elementwise_affine: bool = True,\n norm_type: str = \"layer_norm\",\n norm_eps: float = 1e-5,\n final_dropout: bool = False,\n positional_embeddings: Optional[str] = None,\n num_positional_embeddings: Optional[int] = None,\n ff_inner_dim: Optional[int] = None,\n ff_bias: bool = True,\n attention_out_bias: bool = True,\n context_length: int = 16,\n context_stride: int = 4,\n weighting_scheme: str = \"pyramid\",\n ):\n super().__init__()\n self.dim = dim\n self.num_attention_heads = num_attention_heads\n self.attention_head_dim = attention_head_dim\n self.dropout = dropout\n self.cross_attention_dim = cross_attention_dim\n self.activation_fn = activation_fn\n self.attention_bias = attention_bias\n self.double_self_attention = double_self_attention\n self.norm_elementwise_affine = norm_elementwise_affine\n self.positional_embeddings = positional_embeddings\n self.num_positional_embeddings = num_positional_embeddings\n self.only_cross_attention = only_cross_attention\n\n self.set_free_noise_properties(context_length, context_stride, weighting_scheme)\n\n # We keep these boolean flags for backward-compatibility.\n self.use_ada_layer_norm_zero = (num_embeds_ada_norm is not None) and norm_type == \"ada_norm_zero\"\n self.use_ada_layer_norm = (num_embeds_ada_norm is not None) and norm_type == \"ada_norm\"\n self.use_ada_layer_norm_single = norm_type == \"ada_norm_single\"\n self.use_layer_norm = norm_type == \"layer_norm\"\n self.use_ada_layer_norm_continuous = norm_type == \"ada_norm_continuous\"\n\n if norm_type in (\"ada_norm\", \"ada_norm_zero\") and num_embeds_ada_norm is None:\n raise ValueError(\n f\"`norm_type` is set to {norm_type}, but `num_embeds_ada_norm` is not defined. Please make sure to\"\n f\" define `num_embeds_ada_norm` if setting `norm_type` to {norm_type}.\"\n )\n\n self.norm_type = norm_type\n self.num_embeds_ada_norm = num_embeds_ada_norm\n\n if positional_embeddings and (num_positional_embeddings is None):\n raise ValueError(\n \"If `positional_embedding` type is defined, `num_positition_embeddings` must also be defined.\"\n )\n\n if positional_embeddings == \"sinusoidal\":\n self.pos_embed = SinusoidalPositionalEmbedding(dim, max_seq_length=num_positional_embeddings)\n else:\n self.pos_embed = None\n\n # Define 3 blocks. Each block has its own normalization layer.\n # 1. Self-Attn\n self.norm1 = nn.LayerNorm(dim, elementwise_affine=norm_elementwise_affine, eps=norm_eps)\n\n self.attn1 = Attention(\n query_dim=dim,\n heads=num_attention_heads,\n dim_head=attention_head_dim,\n dropout=dropout,\n bias=attention_bias,\n cross_attention_dim=cross_attention_dim if only_cross_attention else None,\n upcast_attention=upcast_attention,\n out_bias=attention_out_bias,\n )\n\n # 2. Cross-Attn\n if cross_attention_dim is not None or double_self_attention:\n self.norm2 = nn.LayerNorm(dim, norm_eps, norm_elementwise_affine)\n\n self.attn2 = Attention(\n query_dim=dim,\n cross_attention_dim=cross_attention_dim if not double_self_attention else None,\n heads=num_attention_heads,\n dim_head=attention_head_dim,\n dropout=dropout,\n bias=attention_bias,\n upcast_attention=upcast_attention,\n out_bias=attention_out_bias,\n ) # is self-attn if encoder_hidden_states is none\n\n # 3. Feed-forward\n self.ff = FeedForward(\n dim,\n dropout=dropout,\n activation_fn=activation_fn,\n final_dropout=final_dropout,\n inner_dim=ff_inner_dim,\n bias=ff_bias,\n )\n\n self.norm3 = nn.LayerNorm(dim, norm_eps, norm_elementwise_affine)\n\n # let chunk size default to None\n self._chunk_size = None\n self._chunk_dim = 0\n\n def _get_frame_indices(self, num_frames: int) -> List[Tuple[int, int]]:\n frame_indices = []\n for i in range(0, num_frames - self.context_length + 1, self.context_stride):\n window_start = i\n window_end = min(num_frames, i + self.context_length)\n frame_indices.append((window_start, window_end))\n return frame_indices\n\n def _get_frame_weights(self, num_frames: int, weighting_scheme: str = \"pyramid\") -> List[float]:\n if weighting_scheme == \"flat\":\n weights = [1.0] * num_frames\n\n elif weighting_scheme == \"pyramid\":\n if num_frames % 2 == 0:\n # num_frames = 4 => [1, 2, 2, 1]\n mid = num_frames // 2\n weights = list(range(1, mid + 1))\n weights = weights + weights[::-1]\n else:\n # num_frames = 5 => [1, 2, 3, 2, 1]\n mid = (num_frames + 1) // 2\n weights = list(range(1, mid))\n weights = weights + [mid] + weights[::-1]\n\n elif weighting_scheme == \"delayed_reverse_sawtooth\":\n if num_frames % 2 == 0:\n # num_frames = 4 => [0.01, 2, 2, 1]\n mid = num_frames // 2\n weights = [0.01] * (mid - 1) + [mid]\n weights = weights + list(range(mid, 0, -1))\n else:\n # num_frames = 5 => [0.01, 0.01, 3, 2, 1]\n mid = (num_frames + 1) // 2\n weights = [0.01] * mid\n weights = weights + list(range(mid, 0, -1))\n else:\n raise ValueError(f\"Unsupported value for weighting_scheme={weighting_scheme}\")\n\n return weights\n\n def set_free_noise_properties(\n self, context_length: int, context_stride: int, weighting_scheme: str = \"pyramid\"\n ) -> None:\n self.context_length = context_length\n self.context_stride = context_stride\n self.weighting_scheme = weighting_scheme\n\n def set_chunk_feed_forward(self, chunk_size: Optional[int], dim: int = 0) -> None:\n # Sets chunk feed-forward\n self._chunk_size = chunk_size\n self._chunk_dim = dim\n\n def forward(\n self,\n hidden_states: torch.Tensor,\n attention_mask: Optional[torch.Tensor] = None,\n encoder_hidden_states: Optional[torch.Tensor] = None,\n encoder_attention_mask: Optional[torch.Tensor] = None,\n cross_attention_kwargs: Dict[str, Any] = None,\n *args,\n **kwargs,\n ) -> torch.Tensor:\n if cross_attention_kwargs is not None:\n if cross_attention_kwargs.get(\"scale\", None) is not None:\n logger.warning(\"Passing `scale` to `cross_attention_kwargs` is deprecated. `scale` will be ignored.\")\n\n cross_attention_kwargs = cross_attention_kwargs.copy() if cross_attention_kwargs is not None else {}\n\n # hidden_states: [B x H x W, F, C]\n device = hidden_states.device\n dtype = hidden_states.dtype\n\n num_frames = hidden_states.size(1)\n frame_indices = self._get_frame_indices(num_frames)\n frame_weights = self._get_frame_weights(self.context_length, self.weighting_scheme)\n frame_weights = torch.tensor(frame_weights, device=device, dtype=dtype).unsqueeze(0).unsqueeze(-1)\n is_last_frame_batch_complete = frame_indices[-1][1] == num_frames\n\n # Handle out-of-bounds case if num_frames isn't perfectly divisible by context_length\n # For example, num_frames=25, context_length=16, context_stride=4, then we expect the ranges:\n # [(0, 16), (4, 20), (8, 24), (10, 26)]\n if not is_last_frame_batch_complete:\n if num_frames < self.context_length:\n raise ValueError(f\"Expected {num_frames=} to be greater or equal than {self.context_length=}\")\n last_frame_batch_length = num_frames - frame_indices[-1][1]\n frame_indices.append((num_frames - self.context_length, num_frames))\n\n num_times_accumulated = torch.zeros((1, num_frames, 1), device=device)\n accumulated_values = torch.zeros_like(hidden_states)\n\n for i, (frame_start, frame_end) in enumerate(frame_indices):\n # The reason for slicing here is to ensure that if (frame_end - frame_start) is to handle\n # cases like frame_indices=[(0, 16), (16, 20)], if the user provided a video with 19 frames, or\n # essentially a non-multiple of `context_length`.\n weights = torch.ones_like(num_times_accumulated[:, frame_start:frame_end])\n weights *= frame_weights\n\n hidden_states_chunk = hidden_states[:, frame_start:frame_end]\n\n # Notice that normalization is always applied before the real computation in the following blocks.\n # 1. Self-Attention\n norm_hidden_states = self.norm1(hidden_states_chunk)\n\n if self.pos_embed is not None:\n norm_hidden_states = self.pos_embed(norm_hidden_states)\n\n attn_output = self.attn1(\n norm_hidden_states,\n encoder_hidden_states=encoder_hidden_states if self.only_cross_attention else None,\n attention_mask=attention_mask,\n **cross_attention_kwargs,\n )\n\n hidden_states_chunk = attn_output + hidden_states_chunk\n if hidden_states_chunk.ndim == 4:\n hidden_states_chunk = hidden_states_chunk.squeeze(1)\n\n # 2. Cross-Attention\n if self.attn2 is not None:\n norm_hidden_states = self.norm2(hidden_states_chunk)\n\n if self.pos_embed is not None and self.norm_type != \"ada_norm_single\":\n norm_hidden_states = self.pos_embed(norm_hidden_states)\n\n attn_output = self.attn2(\n norm_hidden_states,\n encoder_hidden_states=encoder_hidden_states,\n attention_mask=encoder_attention_mask,\n **cross_attention_kwargs,\n )\n hidden_states_chunk = attn_output + hidden_states_chunk\n\n if i == len(frame_indices) - 1 and not is_last_frame_batch_complete:\n accumulated_values[:, -last_frame_batch_length:] += (\n hidden_states_chunk[:, -last_frame_batch_length:] * weights[:, -last_frame_batch_length:]\n )\n num_times_accumulated[:, -last_frame_batch_length:] += weights[:, -last_frame_batch_length]\n else:\n accumulated_values[:, frame_start:frame_end] += hidden_states_chunk * weights\n num_times_accumulated[:, frame_start:frame_end] += weights\n\n # TODO(aryan): Maybe this could be done in a better way.\n #\n # Previously, this was:\n # hidden_states = torch.where(\n # num_times_accumulated > 0, accumulated_values / num_times_accumulated, accumulated_values\n # )\n #\n # The reasoning for the change here is `torch.where` became a bottleneck at some point when golfing memory\n # spikes. It is particularly noticeable when the number of frames is high. My understanding is that this comes\n # from tensors being copied - which is why we resort to spliting and concatenating here. I've not particularly\n # looked into this deeply because other memory optimizations led to more pronounced reductions.\n hidden_states = torch.cat(\n [\n torch.where(num_times_split > 0, accumulated_split / num_times_split, accumulated_split)\n for accumulated_split, num_times_split in zip(\n accumulated_values.split(self.context_length, dim=1),\n num_times_accumulated.split(self.context_length, dim=1),\n )\n ],\n dim=1,\n ).to(dtype)\n\n # 3. Feed-forward\n norm_hidden_states = self.norm3(hidden_states)\n\n if self._chunk_size is not None:\n ff_output = _chunked_feed_forward(self.ff, norm_hidden_states, self._chunk_dim, self._chunk_size)\n else:\n ff_output = self.ff(norm_hidden_states)\n\n hidden_states = ff_output + hidden_states\n if hidden_states.ndim == 4:\n hidden_states = hidden_states.squeeze(1)\n\n return hidden_states\n\n\nclass FeedForward(nn.Module):\n r\"\"\"\n A feed-forward layer.\n\n Parameters:\n dim (`int`): The number of channels in the input.\n dim_out (`int`, *optional*): The number of channels in the output. If not given, defaults to `dim`.\n mult (`int`, *optional*, defaults to 4): The multiplier to use for the hidden dimension.\n dropout (`float`, *optional*, defaults to 0.0): The dropout probability to use.\n activation_fn (`str`, *optional*, defaults to `\"geglu\"`): Activation function to be used in feed-forward.\n final_dropout (`bool` *optional*, defaults to False): Apply a final dropout.\n bias (`bool`, defaults to True): Whether to use a bias in the linear layer.\n \"\"\"\n\n def __init__(\n self,\n dim: int,\n dim_out: Optional[int] = None,\n mult: int = 4,\n dropout: float = 0.0,\n activation_fn: str = \"geglu\",\n final_dropout: bool = False,\n inner_dim=None,\n bias: bool = True,\n ):\n super().__init__()\n if inner_dim is None:\n inner_dim = int(dim * mult)\n dim_out = dim_out if dim_out is not None else dim\n\n if activation_fn == \"gelu\":\n act_fn = GELU(dim, inner_dim, bias=bias)\n if activation_fn == \"gelu-approximate\":\n act_fn = GELU(dim, inner_dim, approximate=\"tanh\", bias=bias)\n elif activation_fn == \"geglu\":\n act_fn = GEGLU(dim, inner_dim, bias=bias)\n elif activation_fn == \"geglu-approximate\":\n act_fn = ApproximateGELU(dim, inner_dim, bias=bias)\n elif activation_fn == \"swiglu\":\n act_fn = SwiGLU(dim, inner_dim, bias=bias)\n\n self.net = nn.ModuleList([])\n # project in\n self.net.append(act_fn)\n # project dropout\n self.net.append(nn.Dropout(dropout))\n # project out\n self.net.append(nn.Linear(inner_dim, dim_out, bias=bias))\n # FF as used in Vision Transformer, MLP-Mixer, etc. have a final dropout\n if final_dropout:\n self.net.append(nn.Dropout(dropout))\n\n def forward(self, hidden_states: torch.Tensor, *args, **kwargs) -> torch.Tensor:\n if len(args) > 0 or kwargs.get(\"scale\", None) is not None:\n deprecation_message = \"The `scale` argument is deprecated and will be ignored. Please remove it, as passing it will raise an error in the future. `scale` should directly be passed while calling the underlying pipeline component i.e., via `cross_attention_kwargs`.\"\n deprecate(\"scale\", \"1.0.0\", deprecation_message)\n for module in self.net:\n hidden_states = module(hidden_states)\n return hidden_states\n" }, { "path": "CogVideo/finetune/models/attention_processor.py", "content": "# Copyright 2024 The HuggingFace Team. All rights reserved.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n# http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\nimport inspect\nimport math\nfrom typing import Callable, List, Optional, Tuple, Union\nfrom einops import rearrange, repeat\n\nimport torch\nimport torch.nn.functional as F\nfrom torch import nn\n\nfrom diffusers.image_processor import IPAdapterMaskProcessor\nfrom diffusers.utils import deprecate, logging\nfrom diffusers.utils.import_utils import is_torch_npu_available, is_xformers_available\nfrom diffusers.utils.torch_utils import is_torch_version, maybe_allow_in_graph\n\nlogger = logging.get_logger(__name__) # pylint: disable=invalid-name\n\nif is_torch_npu_available():\n import torch_npu\n\nif is_xformers_available():\n import xformers\n import xformers.ops\nelse:\n xformers = None\n\n\n@maybe_allow_in_graph\nclass Attention(nn.Module):\n r\"\"\"\n A cross attention layer.\n\n Parameters:\n query_dim (`int`):\n The number of channels in the query.\n cross_attention_dim (`int`, *optional*):\n The number of channels in the encoder_hidden_states. If not given, defaults to `query_dim`.\n heads (`int`, *optional*, defaults to 8):\n The number of heads to use for multi-head attention.\n kv_heads (`int`, *optional*, defaults to `None`):\n The number of key and value heads to use for multi-head attention. Defaults to `heads`. If\n `kv_heads=heads`, the model will use Multi Head Attention (MHA), if `kv_heads=1` the model will use Multi\n Query Attention (MQA) otherwise GQA is used.\n dim_head (`int`, *optional*, defaults to 64):\n The number of channels in each head.\n dropout (`float`, *optional*, defaults to 0.0):\n The dropout probability to use.\n bias (`bool`, *optional*, defaults to False):\n Set to `True` for the query, key, and value linear layers to contain a bias parameter.\n upcast_attention (`bool`, *optional*, defaults to False):\n Set to `True` to upcast the attention computation to `float32`.\n upcast_softmax (`bool`, *optional*, defaults to False):\n Set to `True` to upcast the softmax computation to `float32`.\n cross_attention_norm (`str`, *optional*, defaults to `None`):\n The type of normalization to use for the cross attention. Can be `None`, `layer_norm`, or `group_norm`.\n cross_attention_norm_num_groups (`int`, *optional*, defaults to 32):\n The number of groups to use for the group norm in the cross attention.\n added_kv_proj_dim (`int`, *optional*, defaults to `None`):\n The number of channels to use for the added key and value projections. If `None`, no projection is used.\n norm_num_groups (`int`, *optional*, defaults to `None`):\n The number of groups to use for the group norm in the attention.\n spatial_norm_dim (`int`, *optional*, defaults to `None`):\n The number of channels to use for the spatial normalization.\n out_bias (`bool`, *optional*, defaults to `True`):\n Set to `True` to use a bias in the output linear layer.\n scale_qk (`bool`, *optional*, defaults to `True`):\n Set to `True` to scale the query and key by `1 / sqrt(dim_head)`.\n only_cross_attention (`bool`, *optional*, defaults to `False`):\n Set to `True` to only use cross attention and not added_kv_proj_dim. Can only be set to `True` if\n `added_kv_proj_dim` is not `None`.\n eps (`float`, *optional*, defaults to 1e-5):\n An additional value added to the denominator in group normalization that is used for numerical stability.\n rescale_output_factor (`float`, *optional*, defaults to 1.0):\n A factor to rescale the output by dividing it with this value.\n residual_connection (`bool`, *optional*, defaults to `False`):\n Set to `True` to add the residual connection to the output.\n _from_deprecated_attn_block (`bool`, *optional*, defaults to `False`):\n Set to `True` if the attention block is loaded from a deprecated state dict.\n processor (`AttnProcessor`, *optional*, defaults to `None`):\n The attention processor to use. If `None`, defaults to `AttnProcessor2_0` if `torch 2.x` is used and\n `AttnProcessor` otherwise.\n \"\"\"\n\n def __init__(\n self,\n query_dim: int,\n cross_attention_dim: Optional[int] = None,\n heads: int = 8,\n kv_heads: Optional[int] = None,\n dim_head: int = 64,\n dropout: float = 0.0,\n bias: bool = False,\n upcast_attention: bool = False,\n upcast_softmax: bool = False,\n cross_attention_norm: Optional[str] = None,\n cross_attention_norm_num_groups: int = 32,\n qk_norm: Optional[str] = None,\n added_kv_proj_dim: Optional[int] = None,\n added_proj_bias: Optional[bool] = True,\n norm_num_groups: Optional[int] = None,\n spatial_norm_dim: Optional[int] = None,\n out_bias: bool = True,\n scale_qk: bool = True,\n only_cross_attention: bool = False,\n eps: float = 1e-5,\n rescale_output_factor: float = 1.0,\n residual_connection: bool = False,\n _from_deprecated_attn_block: bool = False,\n processor: Optional[\"AttnProcessor\"] = None,\n out_dim: int = None,\n context_pre_only=None,\n pre_only=False,\n elementwise_affine: bool = True,\n ):\n super().__init__()\n\n # To prevent circular import.\n from diffusers.models.normalization import FP32LayerNorm, RMSNorm\n\n self.inner_dim = out_dim if out_dim is not None else dim_head * heads\n self.inner_kv_dim = self.inner_dim if kv_heads is None else dim_head * kv_heads\n self.query_dim = query_dim\n self.use_bias = bias\n self.is_cross_attention = cross_attention_dim is not None\n self.cross_attention_dim = cross_attention_dim if cross_attention_dim is not None else query_dim\n self.upcast_attention = upcast_attention\n self.upcast_softmax = upcast_softmax\n self.rescale_output_factor = rescale_output_factor\n self.residual_connection = residual_connection\n self.dropout = dropout\n self.fused_projections = False\n self.out_dim = out_dim if out_dim is not None else query_dim\n self.context_pre_only = context_pre_only\n self.pre_only = pre_only\n\n # we make use of this private variable to know whether this class is loaded\n # with an deprecated state dict so that we can convert it on the fly\n self._from_deprecated_attn_block = _from_deprecated_attn_block\n\n self.scale_qk = scale_qk\n self.scale = dim_head**-0.5 if self.scale_qk else 1.0\n\n self.heads = out_dim // dim_head if out_dim is not None else heads\n # for slice_size > 0 the attention score computation\n # is split across the batch axis to save memory\n # You can set slice_size with `set_attention_slice`\n self.sliceable_head_dim = heads\n\n self.added_kv_proj_dim = added_kv_proj_dim\n self.only_cross_attention = only_cross_attention\n\n if self.added_kv_proj_dim is None and self.only_cross_attention:\n raise ValueError(\n \"`only_cross_attention` can only be set to True if `added_kv_proj_dim` is not None. Make sure to set either `only_cross_attention=False` or define `added_kv_proj_dim`.\"\n )\n\n if norm_num_groups is not None:\n self.group_norm = nn.GroupNorm(num_channels=query_dim, num_groups=norm_num_groups, eps=eps, affine=True)\n else:\n self.group_norm = None\n\n if spatial_norm_dim is not None:\n self.spatial_norm = SpatialNorm(f_channels=query_dim, zq_channels=spatial_norm_dim)\n else:\n self.spatial_norm = None\n\n if qk_norm is None:\n self.norm_q = None\n self.norm_k = None\n elif qk_norm == \"layer_norm\":\n self.norm_q = nn.LayerNorm(dim_head, eps=eps, elementwise_affine=elementwise_affine)\n self.norm_k = nn.LayerNorm(dim_head, eps=eps, elementwise_affine=elementwise_affine)\n elif qk_norm == \"fp32_layer_norm\":\n self.norm_q = FP32LayerNorm(dim_head, elementwise_affine=False, bias=False, eps=eps)\n self.norm_k = FP32LayerNorm(dim_head, elementwise_affine=False, bias=False, eps=eps)\n elif qk_norm == \"layer_norm_across_heads\":\n # Lumina applys qk norm across all heads\n self.norm_q = nn.LayerNorm(dim_head * heads, eps=eps)\n self.norm_k = nn.LayerNorm(dim_head * kv_heads, eps=eps)\n elif qk_norm == \"rms_norm\":\n self.norm_q = RMSNorm(dim_head, eps=eps)\n self.norm_k = RMSNorm(dim_head, eps=eps)\n else:\n raise ValueError(f\"unknown qk_norm: {qk_norm}. Should be None,'layer_norm','fp32_layer_norm','rms_norm'\")\n\n if cross_attention_norm is None:\n self.norm_cross = None\n elif cross_attention_norm == \"layer_norm\":\n self.norm_cross = nn.LayerNorm(self.cross_attention_dim)\n elif cross_attention_norm == \"group_norm\":\n if self.added_kv_proj_dim is not None:\n # The given `encoder_hidden_states` are initially of shape\n # (batch_size, seq_len, added_kv_proj_dim) before being projected\n # to (batch_size, seq_len, cross_attention_dim). The norm is applied\n # before the projection, so we need to use `added_kv_proj_dim` as\n # the number of channels for the group norm.\n norm_cross_num_channels = added_kv_proj_dim\n else:\n norm_cross_num_channels = self.cross_attention_dim\n\n self.norm_cross = nn.GroupNorm(\n num_channels=norm_cross_num_channels, num_groups=cross_attention_norm_num_groups, eps=1e-5, affine=True\n )\n else:\n raise ValueError(\n f\"unknown cross_attention_norm: {cross_attention_norm}. Should be None, 'layer_norm' or 'group_norm'\"\n )\n\n self.to_q = nn.Linear(query_dim, self.inner_dim, bias=bias)\n\n if not self.only_cross_attention:\n # only relevant for the `AddedKVProcessor` classes\n self.to_k = nn.Linear(self.cross_attention_dim, self.inner_kv_dim, bias=bias)\n self.to_v = nn.Linear(self.cross_attention_dim, self.inner_kv_dim, bias=bias)\n else:\n self.to_k = None\n self.to_v = None\n\n self.added_proj_bias = added_proj_bias\n if self.added_kv_proj_dim is not None:\n self.add_k_proj = nn.Linear(added_kv_proj_dim, self.inner_kv_dim, bias=added_proj_bias)\n self.add_v_proj = nn.Linear(added_kv_proj_dim, self.inner_kv_dim, bias=added_proj_bias)\n if self.context_pre_only is not None:\n self.add_q_proj = nn.Linear(added_kv_proj_dim, self.inner_dim, bias=added_proj_bias)\n\n if not self.pre_only:\n self.to_out = nn.ModuleList([])\n self.to_out.append(nn.Linear(self.inner_dim, self.out_dim, bias=out_bias))\n self.to_out.append(nn.Dropout(dropout))\n\n if self.context_pre_only is not None and not self.context_pre_only:\n self.to_add_out = nn.Linear(self.inner_dim, self.out_dim, bias=out_bias)\n\n if qk_norm is not None and added_kv_proj_dim is not None:\n if qk_norm == \"fp32_layer_norm\":\n self.norm_added_q = FP32LayerNorm(dim_head, elementwise_affine=False, bias=False, eps=eps)\n self.norm_added_k = FP32LayerNorm(dim_head, elementwise_affine=False, bias=False, eps=eps)\n elif qk_norm == \"rms_norm\":\n self.norm_added_q = RMSNorm(dim_head, eps=eps)\n self.norm_added_k = RMSNorm(dim_head, eps=eps)\n else:\n raise ValueError(\n f\"unknown qk_norm: {qk_norm}. Should be one of `None,'layer_norm','fp32_layer_norm','rms_norm'`\"\n )\n else:\n self.norm_added_q = None\n self.norm_added_k = None\n\n # set attention processor\n # We use the AttnProcessor2_0 by default when torch 2.x is used which uses\n # torch.nn.functional.scaled_dot_product_attention for native Flash/memory_efficient_attention\n # but only if it has the default `scale` argument. TODO remove scale_qk check when we move to torch 2.1\n if processor is None:\n processor = (\n AttnProcessor2_0() if hasattr(F, \"scaled_dot_product_attention\") and self.scale_qk else AttnProcessor()\n )\n self.set_processor(processor)\n\n def set_use_npu_flash_attention(self, use_npu_flash_attention: bool) -> None:\n r\"\"\"\n Set whether to use npu flash attention from `torch_npu` or not.\n\n \"\"\"\n if use_npu_flash_attention:\n processor = AttnProcessorNPU()\n else:\n # set attention processor\n # We use the AttnProcessor2_0 by default when torch 2.x is used which uses\n # torch.nn.functional.scaled_dot_product_attention for native Flash/memory_efficient_attention\n # but only if it has the default `scale` argument. TODO remove scale_qk check when we move to torch 2.1\n processor = (\n AttnProcessor2_0() if hasattr(F, \"scaled_dot_product_attention\") and self.scale_qk else AttnProcessor()\n )\n self.set_processor(processor)\n\n def set_use_memory_efficient_attention_xformers(\n self, use_memory_efficient_attention_xformers: bool, attention_op: Optional[Callable] = None\n ) -> None:\n r\"\"\"\n Set whether to use memory efficient attention from `xformers` or not.\n\n Args:\n use_memory_efficient_attention_xformers (`bool`):\n Whether to use memory efficient attention from `xformers` or not.\n attention_op (`Callable`, *optional*):\n The attention operation to use. Defaults to `None` which uses the default attention operation from\n `xformers`.\n \"\"\"\n is_custom_diffusion = hasattr(self, \"processor\") and isinstance(\n self.processor,\n (CustomDiffusionAttnProcessor, CustomDiffusionXFormersAttnProcessor, CustomDiffusionAttnProcessor2_0),\n )\n is_added_kv_processor = hasattr(self, \"processor\") and isinstance(\n self.processor,\n (\n AttnAddedKVProcessor,\n AttnAddedKVProcessor2_0,\n SlicedAttnAddedKVProcessor,\n XFormersAttnAddedKVProcessor,\n ),\n )\n\n if use_memory_efficient_attention_xformers:\n if is_added_kv_processor and is_custom_diffusion:\n raise NotImplementedError(\n f\"Memory efficient attention is currently not supported for custom diffusion for attention processor type {self.processor}\"\n )\n if not is_xformers_available():\n raise ModuleNotFoundError(\n (\n \"Refer to https://github.com/facebookresearch/xformers for more information on how to install\"\n \" xformers\"\n ),\n name=\"xformers\",\n )\n elif not torch.cuda.is_available():\n raise ValueError(\n \"torch.cuda.is_available() should be True but is False. xformers' memory efficient attention is\"\n \" only available for GPU \"\n )\n else:\n try:\n # Make sure we can run the memory efficient attention\n _ = xformers.ops.memory_efficient_attention(\n torch.randn((1, 2, 40), device=\"cuda\"),\n torch.randn((1, 2, 40), device=\"cuda\"),\n torch.randn((1, 2, 40), device=\"cuda\"),\n )\n except Exception as e:\n raise e\n\n if is_custom_diffusion:\n processor = CustomDiffusionXFormersAttnProcessor(\n train_kv=self.processor.train_kv,\n train_q_out=self.processor.train_q_out,\n hidden_size=self.processor.hidden_size,\n cross_attention_dim=self.processor.cross_attention_dim,\n attention_op=attention_op,\n )\n processor.load_state_dict(self.processor.state_dict())\n if hasattr(self.processor, \"to_k_custom_diffusion\"):\n processor.to(self.processor.to_k_custom_diffusion.weight.device)\n elif is_added_kv_processor:\n # TODO(Patrick, Suraj, William) - currently xformers doesn't work for UnCLIP\n # which uses this type of cross attention ONLY because the attention mask of format\n # [0, ..., -10.000, ..., 0, ...,] is not supported\n # throw warning\n logger.info(\n \"Memory efficient attention with `xformers` might currently not work correctly if an attention mask is required for the attention operation.\"\n )\n processor = XFormersAttnAddedKVProcessor(attention_op=attention_op)\n else:\n processor = XFormersAttnProcessor(attention_op=attention_op)\n else:\n if is_custom_diffusion:\n attn_processor_class = (\n CustomDiffusionAttnProcessor2_0\n if hasattr(F, \"scaled_dot_product_attention\")\n else CustomDiffusionAttnProcessor\n )\n processor = attn_processor_class(\n train_kv=self.processor.train_kv,\n train_q_out=self.processor.train_q_out,\n hidden_size=self.processor.hidden_size,\n cross_attention_dim=self.processor.cross_attention_dim,\n )\n processor.load_state_dict(self.processor.state_dict())\n if hasattr(self.processor, \"to_k_custom_diffusion\"):\n processor.to(self.processor.to_k_custom_diffusion.weight.device)\n else:\n # set attention processor\n # We use the AttnProcessor2_0 by default when torch 2.x is used which uses\n # torch.nn.functional.scaled_dot_product_attention for native Flash/memory_efficient_attention\n # but only if it has the default `scale` argument. TODO remove scale_qk check when we move to torch 2.1\n processor = (\n AttnProcessor2_0()\n if hasattr(F, \"scaled_dot_product_attention\") and self.scale_qk\n else AttnProcessor()\n )\n\n self.set_processor(processor)\n\n def set_attention_slice(self, slice_size: int) -> None:\n r\"\"\"\n Set the slice size for attention computation.\n\n Args:\n slice_size (`int`):\n The slice size for attention computation.\n \"\"\"\n if slice_size is not None and slice_size > self.sliceable_head_dim:\n raise ValueError(f\"slice_size {slice_size} has to be smaller or equal to {self.sliceable_head_dim}.\")\n\n if slice_size is not None and self.added_kv_proj_dim is not None:\n processor = SlicedAttnAddedKVProcessor(slice_size)\n elif slice_size is not None:\n processor = SlicedAttnProcessor(slice_size)\n elif self.added_kv_proj_dim is not None:\n processor = AttnAddedKVProcessor()\n else:\n # set attention processor\n # We use the AttnProcessor2_0 by default when torch 2.x is used which uses\n # torch.nn.functional.scaled_dot_product_attention for native Flash/memory_efficient_attention\n # but only if it has the default `scale` argument. TODO remove scale_qk check when we move to torch 2.1\n processor = (\n AttnProcessor2_0() if hasattr(F, \"scaled_dot_product_attention\") and self.scale_qk else AttnProcessor()\n )\n\n self.set_processor(processor)\n\n def set_processor(self, processor: \"AttnProcessor\") -> None:\n r\"\"\"\n Set the attention processor to use.\n\n Args:\n processor (`AttnProcessor`):\n The attention processor to use.\n \"\"\"\n # if current processor is in `self._modules` and if passed `processor` is not, we need to\n # pop `processor` from `self._modules`\n if (\n hasattr(self, \"processor\")\n and isinstance(self.processor, torch.nn.Module)\n and not isinstance(processor, torch.nn.Module)\n ):\n logger.info(f\"You are removing possibly trained weights of {self.processor} with {processor}\")\n self._modules.pop(\"processor\")\n\n self.processor = processor\n\n def get_processor(self, return_deprecated_lora: bool = False) -> \"AttentionProcessor\":\n r\"\"\"\n Get the attention processor in use.\n\n Args:\n return_deprecated_lora (`bool`, *optional*, defaults to `False`):\n Set to `True` to return the deprecated LoRA attention processor.\n\n Returns:\n \"AttentionProcessor\": The attention processor in use.\n \"\"\"\n if not return_deprecated_lora:\n return self.processor\n\n def forward(\n self,\n hidden_states: torch.Tensor,\n encoder_hidden_states: Optional[torch.Tensor] = None,\n attention_mask: Optional[torch.Tensor] = None,\n **cross_attention_kwargs,\n ) -> torch.Tensor:\n r\"\"\"\n The forward method of the `Attention` class.\n\n Args:\n hidden_states (`torch.Tensor`):\n The hidden states of the query.\n encoder_hidden_states (`torch.Tensor`, *optional*):\n The hidden states of the encoder.\n attention_mask (`torch.Tensor`, *optional*):\n The attention mask to use. If `None`, no mask is applied.\n **cross_attention_kwargs:\n Additional keyword arguments to pass along to the cross attention.\n\n Returns:\n `torch.Tensor`: The output of the attention layer.\n \"\"\"\n # The `Attention` class can call different attention processors / attention functions\n # here we simply pass along all tensors to the selected processor class\n # For standard processors that are defined here, `**cross_attention_kwargs` is empty\n\n attn_parameters = set(inspect.signature(self.processor.__call__).parameters.keys())\n quiet_attn_parameters = {\"ip_adapter_masks\"}\n unused_kwargs = [\n k for k, _ in cross_attention_kwargs.items() if k not in attn_parameters and k not in quiet_attn_parameters\n ]\n if len(unused_kwargs) > 0:\n logger.warning(\n f\"cross_attention_kwargs {unused_kwargs} are not expected by {self.processor.__class__.__name__} and will be ignored.\"\n )\n cross_attention_kwargs = {k: w for k, w in cross_attention_kwargs.items() if k in attn_parameters}\n\n return self.processor(\n self,\n hidden_states,\n encoder_hidden_states=encoder_hidden_states,\n attention_mask=attention_mask,\n **cross_attention_kwargs,\n )\n\n def batch_to_head_dim(self, tensor: torch.Tensor) -> torch.Tensor:\n r\"\"\"\n Reshape the tensor from `[batch_size, seq_len, dim]` to `[batch_size // heads, seq_len, dim * heads]`. `heads`\n is the number of heads initialized while constructing the `Attention` class.\n\n Args:\n tensor (`torch.Tensor`): The tensor to reshape.\n\n Returns:\n `torch.Tensor`: The reshaped tensor.\n \"\"\"\n head_size = self.heads\n batch_size, seq_len, dim = tensor.shape\n tensor = tensor.reshape(batch_size // head_size, head_size, seq_len, dim)\n tensor = tensor.permute(0, 2, 1, 3).reshape(batch_size // head_size, seq_len, dim * head_size)\n return tensor\n\n def head_to_batch_dim(self, tensor: torch.Tensor, out_dim: int = 3) -> torch.Tensor:\n r\"\"\"\n Reshape the tensor from `[batch_size, seq_len, dim]` to `[batch_size, seq_len, heads, dim // heads]` `heads` is\n the number of heads initialized while constructing the `Attention` class.\n\n Args:\n tensor (`torch.Tensor`): The tensor to reshape.\n out_dim (`int`, *optional*, defaults to `3`): The output dimension of the tensor. If `3`, the tensor is\n reshaped to `[batch_size * heads, seq_len, dim // heads]`.\n\n Returns:\n `torch.Tensor`: The reshaped tensor.\n \"\"\"\n head_size = self.heads\n if tensor.ndim == 3:\n batch_size, seq_len, dim = tensor.shape\n extra_dim = 1\n else:\n batch_size, extra_dim, seq_len, dim = tensor.shape\n tensor = tensor.reshape(batch_size, seq_len * extra_dim, head_size, dim // head_size)\n tensor = tensor.permute(0, 2, 1, 3)\n\n if out_dim == 3:\n tensor = tensor.reshape(batch_size * head_size, seq_len * extra_dim, dim // head_size)\n\n return tensor\n\n def get_attention_scores(\n self, query: torch.Tensor, key: torch.Tensor, attention_mask: Optional[torch.Tensor] = None\n ) -> torch.Tensor:\n r\"\"\"\n Compute the attention scores.\n\n Args:\n query (`torch.Tensor`): The query tensor.\n key (`torch.Tensor`): The key tensor.\n attention_mask (`torch.Tensor`, *optional*): The attention mask to use. If `None`, no mask is applied.\n\n Returns:\n `torch.Tensor`: The attention probabilities/scores.\n \"\"\"\n dtype = query.dtype\n if self.upcast_attention:\n query = query.float()\n key = key.float()\n\n if attention_mask is None:\n baddbmm_input = torch.empty(\n query.shape[0], query.shape[1], key.shape[1], dtype=query.dtype, device=query.device\n )\n beta = 0\n else:\n baddbmm_input = attention_mask\n beta = 1\n\n attention_scores = torch.baddbmm(\n baddbmm_input,\n query,\n key.transpose(-1, -2),\n beta=beta,\n alpha=self.scale,\n )\n del baddbmm_input\n\n if self.upcast_softmax:\n attention_scores = attention_scores.float()\n\n attention_probs = attention_scores.softmax(dim=-1)\n del attention_scores\n\n attention_probs = attention_probs.to(dtype)\n\n return attention_probs\n\n def prepare_attention_mask(\n self, attention_mask: torch.Tensor, target_length: int, batch_size: int, out_dim: int = 3\n ) -> torch.Tensor:\n r\"\"\"\n Prepare the attention mask for the attention computation.\n\n Args:\n attention_mask (`torch.Tensor`):\n The attention mask to prepare.\n target_length (`int`):\n The target length of the attention mask. This is the length of the attention mask after padding.\n batch_size (`int`):\n The batch size, which is used to repeat the attention mask.\n out_dim (`int`, *optional*, defaults to `3`):\n The output dimension of the attention mask. Can be either `3` or `4`.\n\n Returns:\n `torch.Tensor`: The prepared attention mask.\n \"\"\"\n head_size = self.heads\n if attention_mask is None:\n return attention_mask\n\n current_length: int = attention_mask.shape[-1]\n if current_length != target_length:\n if attention_mask.device.type == \"mps\":\n # HACK: MPS: Does not support padding by greater than dimension of input tensor.\n # Instead, we can manually construct the padding tensor.\n padding_shape = (attention_mask.shape[0], attention_mask.shape[1], target_length)\n padding = torch.zeros(padding_shape, dtype=attention_mask.dtype, device=attention_mask.device)\n attention_mask = torch.cat([attention_mask, padding], dim=2)\n else:\n # TODO: for pipelines such as stable-diffusion, padding cross-attn mask:\n # we want to instead pad by (0, remaining_length), where remaining_length is:\n # remaining_length: int = target_length - current_length\n # TODO: re-enable tests/models/test_models_unet_2d_condition.py#test_model_xattn_padding\n attention_mask = F.pad(attention_mask, (0, target_length), value=0.0)\n\n if out_dim == 3:\n if attention_mask.shape[0] < batch_size * head_size:\n attention_mask = attention_mask.repeat_interleave(head_size, dim=0)\n elif out_dim == 4:\n attention_mask = attention_mask.unsqueeze(1)\n attention_mask = attention_mask.repeat_interleave(head_size, dim=1)\n\n return attention_mask\n\n def norm_encoder_hidden_states(self, encoder_hidden_states: torch.Tensor) -> torch.Tensor:\n r\"\"\"\n Normalize the encoder hidden states. Requires `self.norm_cross` to be specified when constructing the\n `Attention` class.\n\n Args:\n encoder_hidden_states (`torch.Tensor`): Hidden states of the encoder.\n\n Returns:\n `torch.Tensor`: The normalized encoder hidden states.\n \"\"\"\n assert self.norm_cross is not None, \"self.norm_cross must be defined to call self.norm_encoder_hidden_states\"\n\n if isinstance(self.norm_cross, nn.LayerNorm):\n encoder_hidden_states = self.norm_cross(encoder_hidden_states)\n elif isinstance(self.norm_cross, nn.GroupNorm):\n # Group norm norms along the channels dimension and expects\n # input to be in the shape of (N, C, *). In this case, we want\n # to norm along the hidden dimension, so we need to move\n # (batch_size, sequence_length, hidden_size) ->\n # (batch_size, hidden_size, sequence_length)\n encoder_hidden_states = encoder_hidden_states.transpose(1, 2)\n encoder_hidden_states = self.norm_cross(encoder_hidden_states)\n encoder_hidden_states = encoder_hidden_states.transpose(1, 2)\n else:\n assert False\n\n return encoder_hidden_states\n\n @torch.no_grad()\n def fuse_projections(self, fuse=True):\n device = self.to_q.weight.data.device\n dtype = self.to_q.weight.data.dtype\n\n if not self.is_cross_attention:\n # fetch weight matrices.\n concatenated_weights = torch.cat([self.to_q.weight.data, self.to_k.weight.data, self.to_v.weight.data])\n in_features = concatenated_weights.shape[1]\n out_features = concatenated_weights.shape[0]\n\n # create a new single projection layer and copy over the weights.\n self.to_qkv = nn.Linear(in_features, out_features, bias=self.use_bias, device=device, dtype=dtype)\n self.to_qkv.weight.copy_(concatenated_weights)\n if self.use_bias:\n concatenated_bias = torch.cat([self.to_q.bias.data, self.to_k.bias.data, self.to_v.bias.data])\n self.to_qkv.bias.copy_(concatenated_bias)\n\n else:\n concatenated_weights = torch.cat([self.to_k.weight.data, self.to_v.weight.data])\n in_features = concatenated_weights.shape[1]\n out_features = concatenated_weights.shape[0]\n\n self.to_kv = nn.Linear(in_features, out_features, bias=self.use_bias, device=device, dtype=dtype)\n self.to_kv.weight.copy_(concatenated_weights)\n if self.use_bias:\n concatenated_bias = torch.cat([self.to_k.bias.data, self.to_v.bias.data])\n self.to_kv.bias.copy_(concatenated_bias)\n\n # handle added projections for SD3 and others.\n if hasattr(self, \"add_q_proj\") and hasattr(self, \"add_k_proj\") and hasattr(self, \"add_v_proj\"):\n concatenated_weights = torch.cat(\n [self.add_q_proj.weight.data, self.add_k_proj.weight.data, self.add_v_proj.weight.data]\n )\n in_features = concatenated_weights.shape[1]\n out_features = concatenated_weights.shape[0]\n\n self.to_added_qkv = nn.Linear(\n in_features, out_features, bias=self.added_proj_bias, device=device, dtype=dtype\n )\n self.to_added_qkv.weight.copy_(concatenated_weights)\n if self.added_proj_bias:\n concatenated_bias = torch.cat(\n [self.add_q_proj.bias.data, self.add_k_proj.bias.data, self.add_v_proj.bias.data]\n )\n self.to_added_qkv.bias.copy_(concatenated_bias)\n\n self.fused_projections = fuse\n\n\nclass AttnProcessor:\n r\"\"\"\n Default processor for performing attention-related computations.\n \"\"\"\n\n def __call__(\n self,\n attn: Attention,\n hidden_states: torch.Tensor,\n encoder_hidden_states: Optional[torch.Tensor] = None,\n attention_mask: Optional[torch.Tensor] = None,\n temb: Optional[torch.Tensor] = None,\n *args,\n **kwargs,\n ) -> torch.Tensor:\n if len(args) > 0 or kwargs.get(\"scale\", None) is not None:\n deprecation_message = \"The `scale` argument is deprecated and will be ignored. Please remove it, as passing it will raise an error in the future. `scale` should directly be passed while calling the underlying pipeline component i.e., via `cross_attention_kwargs`.\"\n deprecate(\"scale\", \"1.0.0\", deprecation_message)\n\n residual = hidden_states\n\n if attn.spatial_norm is not None:\n hidden_states = attn.spatial_norm(hidden_states, temb)\n\n input_ndim = hidden_states.ndim\n\n if input_ndim == 4:\n batch_size, channel, height, width = hidden_states.shape\n hidden_states = hidden_states.view(batch_size, channel, height * width).transpose(1, 2)\n\n batch_size, sequence_length, _ = (\n hidden_states.shape if encoder_hidden_states is None else encoder_hidden_states.shape\n )\n attention_mask = attn.prepare_attention_mask(attention_mask, sequence_length, batch_size)\n\n if attn.group_norm is not None:\n hidden_states = attn.group_norm(hidden_states.transpose(1, 2)).transpose(1, 2)\n\n query = attn.to_q(hidden_states)\n\n if encoder_hidden_states is None:\n encoder_hidden_states = hidden_states\n elif attn.norm_cross:\n encoder_hidden_states = attn.norm_encoder_hidden_states(encoder_hidden_states)\n\n key = attn.to_k(encoder_hidden_states)\n value = attn.to_v(encoder_hidden_states)\n\n query = attn.head_to_batch_dim(query)\n key = attn.head_to_batch_dim(key)\n value = attn.head_to_batch_dim(value)\n\n attention_probs = attn.get_attention_scores(query, key, attention_mask)\n hidden_states = torch.bmm(attention_probs, value)\n hidden_states = attn.batch_to_head_dim(hidden_states)\n\n # linear proj\n hidden_states = attn.to_out[0](hidden_states)\n # dropout\n hidden_states = attn.to_out[1](hidden_states)\n\n if input_ndim == 4:\n hidden_states = hidden_states.transpose(-1, -2).reshape(batch_size, channel, height, width)\n\n if attn.residual_connection:\n hidden_states = hidden_states + residual\n\n hidden_states = hidden_states / attn.rescale_output_factor\n\n return hidden_states\n\n\nclass CustomDiffusionAttnProcessor(nn.Module):\n r\"\"\"\n Processor for implementing attention for the Custom Diffusion method.\n\n Args:\n train_kv (`bool`, defaults to `True`):\n Whether to newly train the key and value matrices corresponding to the text features.\n train_q_out (`bool`, defaults to `True`):\n Whether to newly train query matrices corresponding to the latent image features.\n hidden_size (`int`, *optional*, defaults to `None`):\n The hidden size of the attention layer.\n cross_attention_dim (`int`, *optional*, defaults to `None`):\n The number of channels in the `encoder_hidden_states`.\n out_bias (`bool`, defaults to `True`):\n Whether to include the bias parameter in `train_q_out`.\n dropout (`float`, *optional*, defaults to 0.0):\n The dropout probability to use.\n \"\"\"\n\n def __init__(\n self,\n train_kv: bool = True,\n train_q_out: bool = True,\n hidden_size: Optional[int] = None,\n cross_attention_dim: Optional[int] = None,\n out_bias: bool = True,\n dropout: float = 0.0,\n ):\n super().__init__()\n self.train_kv = train_kv\n self.train_q_out = train_q_out\n\n self.hidden_size = hidden_size\n self.cross_attention_dim = cross_attention_dim\n\n # `_custom_diffusion` id for easy serialization and loading.\n if self.train_kv:\n self.to_k_custom_diffusion = nn.Linear(cross_attention_dim or hidden_size, hidden_size, bias=False)\n self.to_v_custom_diffusion = nn.Linear(cross_attention_dim or hidden_size, hidden_size, bias=False)\n if self.train_q_out:\n self.to_q_custom_diffusion = nn.Linear(hidden_size, hidden_size, bias=False)\n self.to_out_custom_diffusion = nn.ModuleList([])\n self.to_out_custom_diffusion.append(nn.Linear(hidden_size, hidden_size, bias=out_bias))\n self.to_out_custom_diffusion.append(nn.Dropout(dropout))\n\n def __call__(\n self,\n attn: Attention,\n hidden_states: torch.Tensor,\n encoder_hidden_states: Optional[torch.Tensor] = None,\n attention_mask: Optional[torch.Tensor] = None,\n ) -> torch.Tensor:\n batch_size, sequence_length, _ = hidden_states.shape\n attention_mask = attn.prepare_attention_mask(attention_mask, sequence_length, batch_size)\n if self.train_q_out:\n query = self.to_q_custom_diffusion(hidden_states).to(attn.to_q.weight.dtype)\n else:\n query = attn.to_q(hidden_states.to(attn.to_q.weight.dtype))\n\n if encoder_hidden_states is None:\n crossattn = False\n encoder_hidden_states = hidden_states\n else:\n crossattn = True\n if attn.norm_cross:\n encoder_hidden_states = attn.norm_encoder_hidden_states(encoder_hidden_states)\n\n if self.train_kv:\n key = self.to_k_custom_diffusion(encoder_hidden_states.to(self.to_k_custom_diffusion.weight.dtype))\n value = self.to_v_custom_diffusion(encoder_hidden_states.to(self.to_v_custom_diffusion.weight.dtype))\n key = key.to(attn.to_q.weight.dtype)\n value = value.to(attn.to_q.weight.dtype)\n else:\n key = attn.to_k(encoder_hidden_states)\n value = attn.to_v(encoder_hidden_states)\n\n if crossattn:\n detach = torch.ones_like(key)\n detach[:, :1, :] = detach[:, :1, :] * 0.0\n key = detach * key + (1 - detach) * key.detach()\n value = detach * value + (1 - detach) * value.detach()\n\n query = attn.head_to_batch_dim(query)\n key = attn.head_to_batch_dim(key)\n value = attn.head_to_batch_dim(value)\n\n attention_probs = attn.get_attention_scores(query, key, attention_mask)\n hidden_states = torch.bmm(attention_probs, value)\n hidden_states = attn.batch_to_head_dim(hidden_states)\n\n if self.train_q_out:\n # linear proj\n hidden_states = self.to_out_custom_diffusion[0](hidden_states)\n # dropout\n hidden_states = self.to_out_custom_diffusion[1](hidden_states)\n else:\n # linear proj\n hidden_states = attn.to_out[0](hidden_states)\n # dropout\n hidden_states = attn.to_out[1](hidden_states)\n\n return hidden_states\n\n\nclass AttnAddedKVProcessor:\n r\"\"\"\n Processor for performing attention-related computations with extra learnable key and value matrices for the text\n encoder.\n \"\"\"\n\n def __call__(\n self,\n attn: Attention,\n hidden_states: torch.Tensor,\n encoder_hidden_states: Optional[torch.Tensor] = None,\n attention_mask: Optional[torch.Tensor] = None,\n *args,\n **kwargs,\n ) -> torch.Tensor:\n if len(args) > 0 or kwargs.get(\"scale\", None) is not None:\n deprecation_message = \"The `scale` argument is deprecated and will be ignored. Please remove it, as passing it will raise an error in the future. `scale` should directly be passed while calling the underlying pipeline component i.e., via `cross_attention_kwargs`.\"\n deprecate(\"scale\", \"1.0.0\", deprecation_message)\n\n residual = hidden_states\n\n hidden_states = hidden_states.view(hidden_states.shape[0], hidden_states.shape[1], -1).transpose(1, 2)\n batch_size, sequence_length, _ = hidden_states.shape\n\n attention_mask = attn.prepare_attention_mask(attention_mask, sequence_length, batch_size)\n\n if encoder_hidden_states is None:\n encoder_hidden_states = hidden_states\n elif attn.norm_cross:\n encoder_hidden_states = attn.norm_encoder_hidden_states(encoder_hidden_states)\n\n hidden_states = attn.group_norm(hidden_states.transpose(1, 2)).transpose(1, 2)\n\n query = attn.to_q(hidden_states)\n query = attn.head_to_batch_dim(query)\n\n encoder_hidden_states_key_proj = attn.add_k_proj(encoder_hidden_states)\n encoder_hidden_states_value_proj = attn.add_v_proj(encoder_hidden_states)\n encoder_hidden_states_key_proj = attn.head_to_batch_dim(encoder_hidden_states_key_proj)\n encoder_hidden_states_value_proj = attn.head_to_batch_dim(encoder_hidden_states_value_proj)\n\n if not attn.only_cross_attention:\n key = attn.to_k(hidden_states)\n value = attn.to_v(hidden_states)\n key = attn.head_to_batch_dim(key)\n value = attn.head_to_batch_dim(value)\n key = torch.cat([encoder_hidden_states_key_proj, key], dim=1)\n value = torch.cat([encoder_hidden_states_value_proj, value], dim=1)\n else:\n key = encoder_hidden_states_key_proj\n value = encoder_hidden_states_value_proj\n\n attention_probs = attn.get_attention_scores(query, key, attention_mask)\n hidden_states = torch.bmm(attention_probs, value)\n hidden_states = attn.batch_to_head_dim(hidden_states)\n\n # linear proj\n hidden_states = attn.to_out[0](hidden_states)\n # dropout\n hidden_states = attn.to_out[1](hidden_states)\n\n hidden_states = hidden_states.transpose(-1, -2).reshape(residual.shape)\n hidden_states = hidden_states + residual\n\n return hidden_states\n\n\nclass AttnAddedKVProcessor2_0:\n r\"\"\"\n Processor for performing scaled dot-product attention (enabled by default if you're using PyTorch 2.0), with extra\n learnable key and value matrices for the text encoder.\n \"\"\"\n\n def __init__(self):\n if not hasattr(F, \"scaled_dot_product_attention\"):\n raise ImportError(\n \"AttnAddedKVProcessor2_0 requires PyTorch 2.0, to use it, please upgrade PyTorch to 2.0.\"\n )\n\n def __call__(\n self,\n attn: Attention,\n hidden_states: torch.Tensor,\n encoder_hidden_states: Optional[torch.Tensor] = None,\n attention_mask: Optional[torch.Tensor] = None,\n *args,\n **kwargs,\n ) -> torch.Tensor:\n if len(args) > 0 or kwargs.get(\"scale\", None) is not None:\n deprecation_message = \"The `scale` argument is deprecated and will be ignored. Please remove it, as passing it will raise an error in the future. `scale` should directly be passed while calling the underlying pipeline component i.e., via `cross_attention_kwargs`.\"\n deprecate(\"scale\", \"1.0.0\", deprecation_message)\n\n residual = hidden_states\n\n hidden_states = hidden_states.view(hidden_states.shape[0], hidden_states.shape[1], -1).transpose(1, 2)\n batch_size, sequence_length, _ = hidden_states.shape\n\n attention_mask = attn.prepare_attention_mask(attention_mask, sequence_length, batch_size, out_dim=4)\n\n if encoder_hidden_states is None:\n encoder_hidden_states = hidden_states\n elif attn.norm_cross:\n encoder_hidden_states = attn.norm_encoder_hidden_states(encoder_hidden_states)\n\n hidden_states = attn.group_norm(hidden_states.transpose(1, 2)).transpose(1, 2)\n\n query = attn.to_q(hidden_states)\n query = attn.head_to_batch_dim(query, out_dim=4)\n\n encoder_hidden_states_key_proj = attn.add_k_proj(encoder_hidden_states)\n encoder_hidden_states_value_proj = attn.add_v_proj(encoder_hidden_states)\n encoder_hidden_states_key_proj = attn.head_to_batch_dim(encoder_hidden_states_key_proj, out_dim=4)\n encoder_hidden_states_value_proj = attn.head_to_batch_dim(encoder_hidden_states_value_proj, out_dim=4)\n\n if not attn.only_cross_attention:\n key = attn.to_k(hidden_states)\n value = attn.to_v(hidden_states)\n key = attn.head_to_batch_dim(key, out_dim=4)\n value = attn.head_to_batch_dim(value, out_dim=4)\n key = torch.cat([encoder_hidden_states_key_proj, key], dim=2)\n value = torch.cat([encoder_hidden_states_value_proj, value], dim=2)\n else:\n key = encoder_hidden_states_key_proj\n value = encoder_hidden_states_value_proj\n\n # the output of sdp = (batch, num_heads, seq_len, head_dim)\n # TODO: add support for attn.scale when we move to Torch 2.1\n hidden_states = F.scaled_dot_product_attention(\n query, key, value, attn_mask=attention_mask, dropout_p=0.0, is_causal=False\n )\n hidden_states = hidden_states.transpose(1, 2).reshape(batch_size, -1, residual.shape[1])\n\n # linear proj\n hidden_states = attn.to_out[0](hidden_states)\n # dropout\n hidden_states = attn.to_out[1](hidden_states)\n\n hidden_states = hidden_states.transpose(-1, -2).reshape(residual.shape)\n hidden_states = hidden_states + residual\n\n return hidden_states\n\n\nclass JointAttnProcessor2_0:\n \"\"\"Attention processor used typically in processing the SD3-like self-attention projections.\"\"\"\n\n def __init__(self):\n if not hasattr(F, \"scaled_dot_product_attention\"):\n raise ImportError(\"AttnProcessor2_0 requires PyTorch 2.0, to use it, please upgrade PyTorch to 2.0.\")\n\n def __call__(\n self,\n attn: Attention,\n hidden_states: torch.FloatTensor,\n encoder_hidden_states: torch.FloatTensor = None,\n attention_mask: Optional[torch.FloatTensor] = None,\n *args,\n **kwargs,\n ) -> torch.FloatTensor:\n residual = hidden_states\n\n batch_size = hidden_states.shape[0]\n\n # `sample` projections.\n query = attn.to_q(hidden_states)\n key = attn.to_k(hidden_states)\n value = attn.to_v(hidden_states)\n\n inner_dim = key.shape[-1]\n head_dim = inner_dim // attn.heads\n\n query = query.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2)\n key = key.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2)\n value = value.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2)\n\n if attn.norm_q is not None:\n query = attn.norm_q(query)\n if attn.norm_k is not None:\n key = attn.norm_k(key)\n\n # `context` projections.\n if encoder_hidden_states is not None:\n encoder_hidden_states_query_proj = attn.add_q_proj(encoder_hidden_states)\n encoder_hidden_states_key_proj = attn.add_k_proj(encoder_hidden_states)\n encoder_hidden_states_value_proj = attn.add_v_proj(encoder_hidden_states)\n\n encoder_hidden_states_query_proj = encoder_hidden_states_query_proj.view(\n batch_size, -1, attn.heads, head_dim\n ).transpose(1, 2)\n encoder_hidden_states_key_proj = encoder_hidden_states_key_proj.view(\n batch_size, -1, attn.heads, head_dim\n ).transpose(1, 2)\n encoder_hidden_states_value_proj = encoder_hidden_states_value_proj.view(\n batch_size, -1, attn.heads, head_dim\n ).transpose(1, 2)\n\n if attn.norm_added_q is not None:\n encoder_hidden_states_query_proj = attn.norm_added_q(encoder_hidden_states_query_proj)\n if attn.norm_added_k is not None:\n encoder_hidden_states_key_proj = attn.norm_added_k(encoder_hidden_states_key_proj)\n\n query = torch.cat([query, encoder_hidden_states_query_proj], dim=2)\n key = torch.cat([key, encoder_hidden_states_key_proj], dim=2)\n value = torch.cat([value, encoder_hidden_states_value_proj], dim=2)\n\n hidden_states = F.scaled_dot_product_attention(query, key, value, dropout_p=0.0, is_causal=False)\n hidden_states = hidden_states.transpose(1, 2).reshape(batch_size, -1, attn.heads * head_dim)\n hidden_states = hidden_states.to(query.dtype)\n\n if encoder_hidden_states is not None:\n # Split the attention outputs.\n hidden_states, encoder_hidden_states = (\n hidden_states[:, : residual.shape[1]],\n hidden_states[:, residual.shape[1] :],\n )\n if not attn.context_pre_only:\n encoder_hidden_states = attn.to_add_out(encoder_hidden_states)\n\n # linear proj\n hidden_states = attn.to_out[0](hidden_states)\n # dropout\n hidden_states = attn.to_out[1](hidden_states)\n\n if encoder_hidden_states is not None:\n return hidden_states, encoder_hidden_states\n else:\n return hidden_states\n\n\nclass PAGJointAttnProcessor2_0:\n \"\"\"Attention processor used typically in processing the SD3-like self-attention projections.\"\"\"\n\n def __init__(self):\n if not hasattr(F, \"scaled_dot_product_attention\"):\n raise ImportError(\n \"PAGJointAttnProcessor2_0 requires PyTorch 2.0, to use it, please upgrade PyTorch to 2.0.\"\n )\n\n def __call__(\n self,\n attn: Attention,\n hidden_states: torch.FloatTensor,\n encoder_hidden_states: torch.FloatTensor = None,\n ) -> torch.FloatTensor:\n residual = hidden_states\n\n input_ndim = hidden_states.ndim\n if input_ndim == 4:\n batch_size, channel, height, width = hidden_states.shape\n hidden_states = hidden_states.view(batch_size, channel, height * width).transpose(1, 2)\n context_input_ndim = encoder_hidden_states.ndim\n if context_input_ndim == 4:\n batch_size, channel, height, width = encoder_hidden_states.shape\n encoder_hidden_states = encoder_hidden_states.view(batch_size, channel, height * width).transpose(1, 2)\n\n # store the length of image patch sequences to create a mask that prevents interaction between patches\n # similar to making the self-attention map an identity matrix\n identity_block_size = hidden_states.shape[1]\n\n # chunk\n hidden_states_org, hidden_states_ptb = hidden_states.chunk(2)\n encoder_hidden_states_org, encoder_hidden_states_ptb = encoder_hidden_states.chunk(2)\n\n ################## original path ##################\n batch_size = encoder_hidden_states_org.shape[0]\n\n # `sample` projections.\n query_org = attn.to_q(hidden_states_org)\n key_org = attn.to_k(hidden_states_org)\n value_org = attn.to_v(hidden_states_org)\n\n # `context` projections.\n encoder_hidden_states_org_query_proj = attn.add_q_proj(encoder_hidden_states_org)\n encoder_hidden_states_org_key_proj = attn.add_k_proj(encoder_hidden_states_org)\n encoder_hidden_states_org_value_proj = attn.add_v_proj(encoder_hidden_states_org)\n\n # attention\n query_org = torch.cat([query_org, encoder_hidden_states_org_query_proj], dim=1)\n key_org = torch.cat([key_org, encoder_hidden_states_org_key_proj], dim=1)\n value_org = torch.cat([value_org, encoder_hidden_states_org_value_proj], dim=1)\n\n inner_dim = key_org.shape[-1]\n head_dim = inner_dim // attn.heads\n query_org = query_org.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2)\n key_org = key_org.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2)\n value_org = value_org.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2)\n\n hidden_states_org = F.scaled_dot_product_attention(\n query_org, key_org, value_org, dropout_p=0.0, is_causal=False\n )\n hidden_states_org = hidden_states_org.transpose(1, 2).reshape(batch_size, -1, attn.heads * head_dim)\n hidden_states_org = hidden_states_org.to(query_org.dtype)\n\n # Split the attention outputs.\n hidden_states_org, encoder_hidden_states_org = (\n hidden_states_org[:, : residual.shape[1]],\n hidden_states_org[:, residual.shape[1] :],\n )\n\n # linear proj\n hidden_states_org = attn.to_out[0](hidden_states_org)\n # dropout\n hidden_states_org = attn.to_out[1](hidden_states_org)\n if not attn.context_pre_only:\n encoder_hidden_states_org = attn.to_add_out(encoder_hidden_states_org)\n\n if input_ndim == 4:\n hidden_states_org = hidden_states_org.transpose(-1, -2).reshape(batch_size, channel, height, width)\n if context_input_ndim == 4:\n encoder_hidden_states_org = encoder_hidden_states_org.transpose(-1, -2).reshape(\n batch_size, channel, height, width\n )\n\n ################## perturbed path ##################\n\n batch_size = encoder_hidden_states_ptb.shape[0]\n\n # `sample` projections.\n query_ptb = attn.to_q(hidden_states_ptb)\n key_ptb = attn.to_k(hidden_states_ptb)\n value_ptb = attn.to_v(hidden_states_ptb)\n\n # `context` projections.\n encoder_hidden_states_ptb_query_proj = attn.add_q_proj(encoder_hidden_states_ptb)\n encoder_hidden_states_ptb_key_proj = attn.add_k_proj(encoder_hidden_states_ptb)\n encoder_hidden_states_ptb_value_proj = attn.add_v_proj(encoder_hidden_states_ptb)\n\n # attention\n query_ptb = torch.cat([query_ptb, encoder_hidden_states_ptb_query_proj], dim=1)\n key_ptb = torch.cat([key_ptb, encoder_hidden_states_ptb_key_proj], dim=1)\n value_ptb = torch.cat([value_ptb, encoder_hidden_states_ptb_value_proj], dim=1)\n\n inner_dim = key_ptb.shape[-1]\n head_dim = inner_dim // attn.heads\n query_ptb = query_ptb.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2)\n key_ptb = key_ptb.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2)\n value_ptb = value_ptb.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2)\n\n # create a full mask with all entries set to 0\n seq_len = query_ptb.size(2)\n full_mask = torch.zeros((seq_len, seq_len), device=query_ptb.device, dtype=query_ptb.dtype)\n\n # set the attention value between image patches to -inf\n full_mask[:identity_block_size, :identity_block_size] = float(\"-inf\")\n\n # set the diagonal of the attention value between image patches to 0\n full_mask[:identity_block_size, :identity_block_size].fill_diagonal_(0)\n\n # expand the mask to match the attention weights shape\n full_mask = full_mask.unsqueeze(0).unsqueeze(0) # Add batch and num_heads dimensions\n\n hidden_states_ptb = F.scaled_dot_product_attention(\n query_ptb, key_ptb, value_ptb, attn_mask=full_mask, dropout_p=0.0, is_causal=False\n )\n hidden_states_ptb = hidden_states_ptb.transpose(1, 2).reshape(batch_size, -1, attn.heads * head_dim)\n hidden_states_ptb = hidden_states_ptb.to(query_ptb.dtype)\n\n # split the attention outputs.\n hidden_states_ptb, encoder_hidden_states_ptb = (\n hidden_states_ptb[:, : residual.shape[1]],\n hidden_states_ptb[:, residual.shape[1] :],\n )\n\n # linear proj\n hidden_states_ptb = attn.to_out[0](hidden_states_ptb)\n # dropout\n hidden_states_ptb = attn.to_out[1](hidden_states_ptb)\n if not attn.context_pre_only:\n encoder_hidden_states_ptb = attn.to_add_out(encoder_hidden_states_ptb)\n\n if input_ndim == 4:\n hidden_states_ptb = hidden_states_ptb.transpose(-1, -2).reshape(batch_size, channel, height, width)\n if context_input_ndim == 4:\n encoder_hidden_states_ptb = encoder_hidden_states_ptb.transpose(-1, -2).reshape(\n batch_size, channel, height, width\n )\n\n ################ concat ###############\n hidden_states = torch.cat([hidden_states_org, hidden_states_ptb])\n encoder_hidden_states = torch.cat([encoder_hidden_states_org, encoder_hidden_states_ptb])\n\n return hidden_states, encoder_hidden_states\n\n\nclass PAGCFGJointAttnProcessor2_0:\n \"\"\"Attention processor used typically in processing the SD3-like self-attention projections.\"\"\"\n\n def __init__(self):\n if not hasattr(F, \"scaled_dot_product_attention\"):\n raise ImportError(\n \"PAGCFGJointAttnProcessor2_0 requires PyTorch 2.0, to use it, please upgrade PyTorch to 2.0.\"\n )\n\n def __call__(\n self,\n attn: Attention,\n hidden_states: torch.FloatTensor,\n encoder_hidden_states: torch.FloatTensor = None,\n attention_mask: Optional[torch.FloatTensor] = None,\n *args,\n **kwargs,\n ) -> torch.FloatTensor:\n residual = hidden_states\n\n input_ndim = hidden_states.ndim\n if input_ndim == 4:\n batch_size, channel, height, width = hidden_states.shape\n hidden_states = hidden_states.view(batch_size, channel, height * width).transpose(1, 2)\n context_input_ndim = encoder_hidden_states.ndim\n if context_input_ndim == 4:\n batch_size, channel, height, width = encoder_hidden_states.shape\n encoder_hidden_states = encoder_hidden_states.view(batch_size, channel, height * width).transpose(1, 2)\n\n identity_block_size = hidden_states.shape[\n 1\n ] # patch embeddings width * height (correspond to self-attention map width or height)\n\n # chunk\n hidden_states_uncond, hidden_states_org, hidden_states_ptb = hidden_states.chunk(3)\n hidden_states_org = torch.cat([hidden_states_uncond, hidden_states_org])\n\n (\n encoder_hidden_states_uncond,\n encoder_hidden_states_org,\n encoder_hidden_states_ptb,\n ) = encoder_hidden_states.chunk(3)\n encoder_hidden_states_org = torch.cat([encoder_hidden_states_uncond, encoder_hidden_states_org])\n\n ################## original path ##################\n batch_size = encoder_hidden_states_org.shape[0]\n\n # `sample` projections.\n query_org = attn.to_q(hidden_states_org)\n key_org = attn.to_k(hidden_states_org)\n value_org = attn.to_v(hidden_states_org)\n\n # `context` projections.\n encoder_hidden_states_org_query_proj = attn.add_q_proj(encoder_hidden_states_org)\n encoder_hidden_states_org_key_proj = attn.add_k_proj(encoder_hidden_states_org)\n encoder_hidden_states_org_value_proj = attn.add_v_proj(encoder_hidden_states_org)\n\n # attention\n query_org = torch.cat([query_org, encoder_hidden_states_org_query_proj], dim=1)\n key_org = torch.cat([key_org, encoder_hidden_states_org_key_proj], dim=1)\n value_org = torch.cat([value_org, encoder_hidden_states_org_value_proj], dim=1)\n\n inner_dim = key_org.shape[-1]\n head_dim = inner_dim // attn.heads\n query_org = query_org.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2)\n key_org = key_org.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2)\n value_org = value_org.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2)\n\n hidden_states_org = F.scaled_dot_product_attention(\n query_org, key_org, value_org, dropout_p=0.0, is_causal=False\n )\n hidden_states_org = hidden_states_org.transpose(1, 2).reshape(batch_size, -1, attn.heads * head_dim)\n hidden_states_org = hidden_states_org.to(query_org.dtype)\n\n # Split the attention outputs.\n hidden_states_org, encoder_hidden_states_org = (\n hidden_states_org[:, : residual.shape[1]],\n hidden_states_org[:, residual.shape[1] :],\n )\n\n # linear proj\n hidden_states_org = attn.to_out[0](hidden_states_org)\n # dropout\n hidden_states_org = attn.to_out[1](hidden_states_org)\n if not attn.context_pre_only:\n encoder_hidden_states_org = attn.to_add_out(encoder_hidden_states_org)\n\n if input_ndim == 4:\n hidden_states_org = hidden_states_org.transpose(-1, -2).reshape(batch_size, channel, height, width)\n if context_input_ndim == 4:\n encoder_hidden_states_org = encoder_hidden_states_org.transpose(-1, -2).reshape(\n batch_size, channel, height, width\n )\n\n ################## perturbed path ##################\n\n batch_size = encoder_hidden_states_ptb.shape[0]\n\n # `sample` projections.\n query_ptb = attn.to_q(hidden_states_ptb)\n key_ptb = attn.to_k(hidden_states_ptb)\n value_ptb = attn.to_v(hidden_states_ptb)\n\n # `context` projections.\n encoder_hidden_states_ptb_query_proj = attn.add_q_proj(encoder_hidden_states_ptb)\n encoder_hidden_states_ptb_key_proj = attn.add_k_proj(encoder_hidden_states_ptb)\n encoder_hidden_states_ptb_value_proj = attn.add_v_proj(encoder_hidden_states_ptb)\n\n # attention\n query_ptb = torch.cat([query_ptb, encoder_hidden_states_ptb_query_proj], dim=1)\n key_ptb = torch.cat([key_ptb, encoder_hidden_states_ptb_key_proj], dim=1)\n value_ptb = torch.cat([value_ptb, encoder_hidden_states_ptb_value_proj], dim=1)\n\n inner_dim = key_ptb.shape[-1]\n head_dim = inner_dim // attn.heads\n query_ptb = query_ptb.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2)\n key_ptb = key_ptb.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2)\n value_ptb = value_ptb.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2)\n\n # create a full mask with all entries set to 0\n seq_len = query_ptb.size(2)\n full_mask = torch.zeros((seq_len, seq_len), device=query_ptb.device, dtype=query_ptb.dtype)\n\n # set the attention value between image patches to -inf\n full_mask[:identity_block_size, :identity_block_size] = float(\"-inf\")\n\n # set the diagonal of the attention value between image patches to 0\n full_mask[:identity_block_size, :identity_block_size].fill_diagonal_(0)\n\n # expand the mask to match the attention weights shape\n full_mask = full_mask.unsqueeze(0).unsqueeze(0) # Add batch and num_heads dimensions\n\n hidden_states_ptb = F.scaled_dot_product_attention(\n query_ptb, key_ptb, value_ptb, attn_mask=full_mask, dropout_p=0.0, is_causal=False\n )\n hidden_states_ptb = hidden_states_ptb.transpose(1, 2).reshape(batch_size, -1, attn.heads * head_dim)\n hidden_states_ptb = hidden_states_ptb.to(query_ptb.dtype)\n\n # split the attention outputs.\n hidden_states_ptb, encoder_hidden_states_ptb = (\n hidden_states_ptb[:, : residual.shape[1]],\n hidden_states_ptb[:, residual.shape[1] :],\n )\n\n # linear proj\n hidden_states_ptb = attn.to_out[0](hidden_states_ptb)\n # dropout\n hidden_states_ptb = attn.to_out[1](hidden_states_ptb)\n if not attn.context_pre_only:\n encoder_hidden_states_ptb = attn.to_add_out(encoder_hidden_states_ptb)\n\n if input_ndim == 4:\n hidden_states_ptb = hidden_states_ptb.transpose(-1, -2).reshape(batch_size, channel, height, width)\n if context_input_ndim == 4:\n encoder_hidden_states_ptb = encoder_hidden_states_ptb.transpose(-1, -2).reshape(\n batch_size, channel, height, width\n )\n\n ################ concat ###############\n hidden_states = torch.cat([hidden_states_org, hidden_states_ptb])\n encoder_hidden_states = torch.cat([encoder_hidden_states_org, encoder_hidden_states_ptb])\n\n return hidden_states, encoder_hidden_states\n\n\nclass FusedJointAttnProcessor2_0:\n \"\"\"Attention processor used typically in processing the SD3-like self-attention projections.\"\"\"\n\n def __init__(self):\n if not hasattr(F, \"scaled_dot_product_attention\"):\n raise ImportError(\"AttnProcessor2_0 requires PyTorch 2.0, to use it, please upgrade PyTorch to 2.0.\")\n\n def __call__(\n self,\n attn: Attention,\n hidden_states: torch.FloatTensor,\n encoder_hidden_states: torch.FloatTensor = None,\n attention_mask: Optional[torch.FloatTensor] = None,\n *args,\n **kwargs,\n ) -> torch.FloatTensor:\n residual = hidden_states\n\n input_ndim = hidden_states.ndim\n if input_ndim == 4:\n batch_size, channel, height, width = hidden_states.shape\n hidden_states = hidden_states.view(batch_size, channel, height * width).transpose(1, 2)\n context_input_ndim = encoder_hidden_states.ndim\n if context_input_ndim == 4:\n batch_size, channel, height, width = encoder_hidden_states.shape\n encoder_hidden_states = encoder_hidden_states.view(batch_size, channel, height * width).transpose(1, 2)\n\n batch_size = encoder_hidden_states.shape[0]\n\n # `sample` projections.\n qkv = attn.to_qkv(hidden_states)\n split_size = qkv.shape[-1] // 3\n query, key, value = torch.split(qkv, split_size, dim=-1)\n\n # `context` projections.\n encoder_qkv = attn.to_added_qkv(encoder_hidden_states)\n split_size = encoder_qkv.shape[-1] // 3\n (\n encoder_hidden_states_query_proj,\n encoder_hidden_states_key_proj,\n encoder_hidden_states_value_proj,\n ) = torch.split(encoder_qkv, split_size, dim=-1)\n\n # attention\n query = torch.cat([query, encoder_hidden_states_query_proj], dim=1)\n key = torch.cat([key, encoder_hidden_states_key_proj], dim=1)\n value = torch.cat([value, encoder_hidden_states_value_proj], dim=1)\n\n inner_dim = key.shape[-1]\n head_dim = inner_dim // attn.heads\n query = query.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2)\n key = key.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2)\n value = value.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2)\n\n hidden_states = F.scaled_dot_product_attention(query, key, value, dropout_p=0.0, is_causal=False)\n hidden_states = hidden_states.transpose(1, 2).reshape(batch_size, -1, attn.heads * head_dim)\n hidden_states = hidden_states.to(query.dtype)\n\n # Split the attention outputs.\n hidden_states, encoder_hidden_states = (\n hidden_states[:, : residual.shape[1]],\n hidden_states[:, residual.shape[1] :],\n )\n\n # linear proj\n hidden_states = attn.to_out[0](hidden_states)\n # dropout\n hidden_states = attn.to_out[1](hidden_states)\n if not attn.context_pre_only:\n encoder_hidden_states = attn.to_add_out(encoder_hidden_states)\n\n if input_ndim == 4:\n hidden_states = hidden_states.transpose(-1, -2).reshape(batch_size, channel, height, width)\n if context_input_ndim == 4:\n encoder_hidden_states = encoder_hidden_states.transpose(-1, -2).reshape(batch_size, channel, height, width)\n\n return hidden_states, encoder_hidden_states\n\n\nclass AuraFlowAttnProcessor2_0:\n \"\"\"Attention processor used typically in processing Aura Flow.\"\"\"\n\n def __init__(self):\n if not hasattr(F, \"scaled_dot_product_attention\") and is_torch_version(\"<\", \"2.1\"):\n raise ImportError(\n \"AuraFlowAttnProcessor2_0 requires PyTorch 2.0, to use it, please upgrade PyTorch to at least 2.1 or above as we use `scale` in `F.scaled_dot_product_attention()`. \"\n )\n\n def __call__(\n self,\n attn: Attention,\n hidden_states: torch.FloatTensor,\n encoder_hidden_states: torch.FloatTensor = None,\n *args,\n **kwargs,\n ) -> torch.FloatTensor:\n batch_size = hidden_states.shape[0]\n\n # `sample` projections.\n query = attn.to_q(hidden_states)\n key = attn.to_k(hidden_states)\n value = attn.to_v(hidden_states)\n\n # `context` projections.\n if encoder_hidden_states is not None:\n encoder_hidden_states_query_proj = attn.add_q_proj(encoder_hidden_states)\n encoder_hidden_states_key_proj = attn.add_k_proj(encoder_hidden_states)\n encoder_hidden_states_value_proj = attn.add_v_proj(encoder_hidden_states)\n\n # Reshape.\n inner_dim = key.shape[-1]\n head_dim = inner_dim // attn.heads\n query = query.view(batch_size, -1, attn.heads, head_dim)\n key = key.view(batch_size, -1, attn.heads, head_dim)\n value = value.view(batch_size, -1, attn.heads, head_dim)\n\n # Apply QK norm.\n if attn.norm_q is not None:\n query = attn.norm_q(query)\n if attn.norm_k is not None:\n key = attn.norm_k(key)\n\n # Concatenate the projections.\n if encoder_hidden_states is not None:\n encoder_hidden_states_query_proj = encoder_hidden_states_query_proj.view(\n batch_size, -1, attn.heads, head_dim\n )\n encoder_hidden_states_key_proj = encoder_hidden_states_key_proj.view(batch_size, -1, attn.heads, head_dim)\n encoder_hidden_states_value_proj = encoder_hidden_states_value_proj.view(\n batch_size, -1, attn.heads, head_dim\n )\n\n if attn.norm_added_q is not None:\n encoder_hidden_states_query_proj = attn.norm_added_q(encoder_hidden_states_query_proj)\n if attn.norm_added_k is not None:\n encoder_hidden_states_key_proj = attn.norm_added_q(encoder_hidden_states_key_proj)\n\n query = torch.cat([encoder_hidden_states_query_proj, query], dim=1)\n key = torch.cat([encoder_hidden_states_key_proj, key], dim=1)\n value = torch.cat([encoder_hidden_states_value_proj, value], dim=1)\n\n query = query.transpose(1, 2)\n key = key.transpose(1, 2)\n value = value.transpose(1, 2)\n\n # Attention.\n hidden_states = F.scaled_dot_product_attention(\n query, key, value, dropout_p=0.0, scale=attn.scale, is_causal=False\n )\n hidden_states = hidden_states.transpose(1, 2).reshape(batch_size, -1, attn.heads * head_dim)\n hidden_states = hidden_states.to(query.dtype)\n\n # Split the attention outputs.\n if encoder_hidden_states is not None:\n hidden_states, encoder_hidden_states = (\n hidden_states[:, encoder_hidden_states.shape[1] :],\n hidden_states[:, : encoder_hidden_states.shape[1]],\n )\n\n # linear proj\n hidden_states = attn.to_out[0](hidden_states)\n # dropout\n hidden_states = attn.to_out[1](hidden_states)\n if encoder_hidden_states is not None:\n encoder_hidden_states = attn.to_add_out(encoder_hidden_states)\n\n if encoder_hidden_states is not None:\n return hidden_states, encoder_hidden_states\n else:\n return hidden_states\n\n\nclass FusedAuraFlowAttnProcessor2_0:\n \"\"\"Attention processor used typically in processing Aura Flow with fused projections.\"\"\"\n\n def __init__(self):\n if not hasattr(F, \"scaled_dot_product_attention\") and is_torch_version(\"<\", \"2.1\"):\n raise ImportError(\n \"FusedAuraFlowAttnProcessor2_0 requires PyTorch 2.0, to use it, please upgrade PyTorch to at least 2.1 or above as we use `scale` in `F.scaled_dot_product_attention()`. \"\n )\n\n def __call__(\n self,\n attn: Attention,\n hidden_states: torch.FloatTensor,\n encoder_hidden_states: torch.FloatTensor = None,\n *args,\n **kwargs,\n ) -> torch.FloatTensor:\n batch_size = hidden_states.shape[0]\n\n # `sample` projections.\n qkv = attn.to_qkv(hidden_states)\n split_size = qkv.shape[-1] // 3\n query, key, value = torch.split(qkv, split_size, dim=-1)\n\n # `context` projections.\n if encoder_hidden_states is not None:\n encoder_qkv = attn.to_added_qkv(encoder_hidden_states)\n split_size = encoder_qkv.shape[-1] // 3\n (\n encoder_hidden_states_query_proj,\n encoder_hidden_states_key_proj,\n encoder_hidden_states_value_proj,\n ) = torch.split(encoder_qkv, split_size, dim=-1)\n\n # Reshape.\n inner_dim = key.shape[-1]\n head_dim = inner_dim // attn.heads\n query = query.view(batch_size, -1, attn.heads, head_dim)\n key = key.view(batch_size, -1, attn.heads, head_dim)\n value = value.view(batch_size, -1, attn.heads, head_dim)\n\n # Apply QK norm.\n if attn.norm_q is not None:\n query = attn.norm_q(query)\n if attn.norm_k is not None:\n key = attn.norm_k(key)\n\n # Concatenate the projections.\n if encoder_hidden_states is not None:\n encoder_hidden_states_query_proj = encoder_hidden_states_query_proj.view(\n batch_size, -1, attn.heads, head_dim\n )\n encoder_hidden_states_key_proj = encoder_hidden_states_key_proj.view(batch_size, -1, attn.heads, head_dim)\n encoder_hidden_states_value_proj = encoder_hidden_states_value_proj.view(\n batch_size, -1, attn.heads, head_dim\n )\n\n if attn.norm_added_q is not None:\n encoder_hidden_states_query_proj = attn.norm_added_q(encoder_hidden_states_query_proj)\n if attn.norm_added_k is not None:\n encoder_hidden_states_key_proj = attn.norm_added_q(encoder_hidden_states_key_proj)\n\n query = torch.cat([encoder_hidden_states_query_proj, query], dim=1)\n key = torch.cat([encoder_hidden_states_key_proj, key], dim=1)\n value = torch.cat([encoder_hidden_states_value_proj, value], dim=1)\n\n query = query.transpose(1, 2)\n key = key.transpose(1, 2)\n value = value.transpose(1, 2)\n\n # Attention.\n hidden_states = F.scaled_dot_product_attention(\n query, key, value, dropout_p=0.0, scale=attn.scale, is_causal=False\n )\n hidden_states = hidden_states.transpose(1, 2).reshape(batch_size, -1, attn.heads * head_dim)\n hidden_states = hidden_states.to(query.dtype)\n\n # Split the attention outputs.\n if encoder_hidden_states is not None:\n hidden_states, encoder_hidden_states = (\n hidden_states[:, encoder_hidden_states.shape[1] :],\n hidden_states[:, : encoder_hidden_states.shape[1]],\n )\n\n # linear proj\n hidden_states = attn.to_out[0](hidden_states)\n # dropout\n hidden_states = attn.to_out[1](hidden_states)\n if encoder_hidden_states is not None:\n encoder_hidden_states = attn.to_add_out(encoder_hidden_states)\n\n if encoder_hidden_states is not None:\n return hidden_states, encoder_hidden_states\n else:\n return hidden_states\n\n\nclass FluxAttnProcessor2_0:\n \"\"\"Attention processor used typically in processing the SD3-like self-attention projections.\"\"\"\n\n def __init__(self):\n if not hasattr(F, \"scaled_dot_product_attention\"):\n raise ImportError(\"FluxAttnProcessor2_0 requires PyTorch 2.0, to use it, please upgrade PyTorch to 2.0.\")\n\n def __call__(\n self,\n attn: Attention,\n hidden_states: torch.FloatTensor,\n encoder_hidden_states: torch.FloatTensor = None,\n attention_mask: Optional[torch.FloatTensor] = None,\n image_rotary_emb: Optional[torch.Tensor] = None,\n ) -> torch.FloatTensor:\n batch_size, _, _ = hidden_states.shape if encoder_hidden_states is None else encoder_hidden_states.shape\n\n # `sample` projections.\n query = attn.to_q(hidden_states)\n key = attn.to_k(hidden_states)\n value = attn.to_v(hidden_states)\n\n inner_dim = key.shape[-1]\n head_dim = inner_dim // attn.heads\n\n query = query.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2)\n key = key.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2)\n value = value.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2)\n\n if attn.norm_q is not None:\n query = attn.norm_q(query)\n if attn.norm_k is not None:\n key = attn.norm_k(key)\n\n # the attention in FluxSingleTransformerBlock does not use `encoder_hidden_states`\n if encoder_hidden_states is not None:\n # `context` projections.\n encoder_hidden_states_query_proj = attn.add_q_proj(encoder_hidden_states)\n encoder_hidden_states_key_proj = attn.add_k_proj(encoder_hidden_states)\n encoder_hidden_states_value_proj = attn.add_v_proj(encoder_hidden_states)\n\n encoder_hidden_states_query_proj = encoder_hidden_states_query_proj.view(\n batch_size, -1, attn.heads, head_dim\n ).transpose(1, 2)\n encoder_hidden_states_key_proj = encoder_hidden_states_key_proj.view(\n batch_size, -1, attn.heads, head_dim\n ).transpose(1, 2)\n encoder_hidden_states_value_proj = encoder_hidden_states_value_proj.view(\n batch_size, -1, attn.heads, head_dim\n ).transpose(1, 2)\n\n if attn.norm_added_q is not None:\n encoder_hidden_states_query_proj = attn.norm_added_q(encoder_hidden_states_query_proj)\n if attn.norm_added_k is not None:\n encoder_hidden_states_key_proj = attn.norm_added_k(encoder_hidden_states_key_proj)\n\n # attention\n query = torch.cat([encoder_hidden_states_query_proj, query], dim=2)\n key = torch.cat([encoder_hidden_states_key_proj, key], dim=2)\n value = torch.cat([encoder_hidden_states_value_proj, value], dim=2)\n\n if image_rotary_emb is not None:\n from diffusers.models.embeddings import apply_rotary_emb\n\n query = apply_rotary_emb(query, image_rotary_emb)\n key = apply_rotary_emb(key, image_rotary_emb)\n\n hidden_states = F.scaled_dot_product_attention(query, key, value, dropout_p=0.0, is_causal=False)\n hidden_states = hidden_states.transpose(1, 2).reshape(batch_size, -1, attn.heads * head_dim)\n hidden_states = hidden_states.to(query.dtype)\n\n if encoder_hidden_states is not None:\n encoder_hidden_states, hidden_states = (\n hidden_states[:, : encoder_hidden_states.shape[1]],\n hidden_states[:, encoder_hidden_states.shape[1] :],\n )\n\n # linear proj\n hidden_states = attn.to_out[0](hidden_states)\n # dropout\n hidden_states = attn.to_out[1](hidden_states)\n encoder_hidden_states = attn.to_add_out(encoder_hidden_states)\n\n return hidden_states, encoder_hidden_states\n else:\n return hidden_states\n\n\nclass FusedFluxAttnProcessor2_0:\n \"\"\"Attention processor used typically in processing the SD3-like self-attention projections.\"\"\"\n\n def __init__(self):\n if not hasattr(F, \"scaled_dot_product_attention\"):\n raise ImportError(\n \"FusedFluxAttnProcessor2_0 requires PyTorch 2.0, to use it, please upgrade PyTorch to 2.0.\"\n )\n\n def __call__(\n self,\n attn: Attention,\n hidden_states: torch.FloatTensor,\n encoder_hidden_states: torch.FloatTensor = None,\n attention_mask: Optional[torch.FloatTensor] = None,\n image_rotary_emb: Optional[torch.Tensor] = None,\n ) -> torch.FloatTensor:\n batch_size, _, _ = hidden_states.shape if encoder_hidden_states is None else encoder_hidden_states.shape\n\n # `sample` projections.\n qkv = attn.to_qkv(hidden_states)\n split_size = qkv.shape[-1] // 3\n query, key, value = torch.split(qkv, split_size, dim=-1)\n\n inner_dim = key.shape[-1]\n head_dim = inner_dim // attn.heads\n\n query = query.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2)\n key = key.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2)\n value = value.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2)\n\n if attn.norm_q is not None:\n query = attn.norm_q(query)\n if attn.norm_k is not None:\n key = attn.norm_k(key)\n\n # the attention in FluxSingleTransformerBlock does not use `encoder_hidden_states`\n # `context` projections.\n if encoder_hidden_states is not None:\n encoder_qkv = attn.to_added_qkv(encoder_hidden_states)\n split_size = encoder_qkv.shape[-1] // 3\n (\n encoder_hidden_states_query_proj,\n encoder_hidden_states_key_proj,\n encoder_hidden_states_value_proj,\n ) = torch.split(encoder_qkv, split_size, dim=-1)\n\n encoder_hidden_states_query_proj = encoder_hidden_states_query_proj.view(\n batch_size, -1, attn.heads, head_dim\n ).transpose(1, 2)\n encoder_hidden_states_key_proj = encoder_hidden_states_key_proj.view(\n batch_size, -1, attn.heads, head_dim\n ).transpose(1, 2)\n encoder_hidden_states_value_proj = encoder_hidden_states_value_proj.view(\n batch_size, -1, attn.heads, head_dim\n ).transpose(1, 2)\n\n if attn.norm_added_q is not None:\n encoder_hidden_states_query_proj = attn.norm_added_q(encoder_hidden_states_query_proj)\n if attn.norm_added_k is not None:\n encoder_hidden_states_key_proj = attn.norm_added_k(encoder_hidden_states_key_proj)\n\n # attention\n query = torch.cat([encoder_hidden_states_query_proj, query], dim=2)\n key = torch.cat([encoder_hidden_states_key_proj, key], dim=2)\n value = torch.cat([encoder_hidden_states_value_proj, value], dim=2)\n\n if image_rotary_emb is not None:\n from diffusers.models.embeddings import apply_rotary_emb\n\n query = apply_rotary_emb(query, image_rotary_emb)\n key = apply_rotary_emb(key, image_rotary_emb)\n\n hidden_states = F.scaled_dot_product_attention(query, key, value, dropout_p=0.0, is_causal=False)\n hidden_states = hidden_states.transpose(1, 2).reshape(batch_size, -1, attn.heads * head_dim)\n hidden_states = hidden_states.to(query.dtype)\n\n if encoder_hidden_states is not None:\n encoder_hidden_states, hidden_states = (\n hidden_states[:, : encoder_hidden_states.shape[1]],\n hidden_states[:, encoder_hidden_states.shape[1] :],\n )\n\n # linear proj\n hidden_states = attn.to_out[0](hidden_states)\n # dropout\n hidden_states = attn.to_out[1](hidden_states)\n encoder_hidden_states = attn.to_add_out(encoder_hidden_states)\n\n return hidden_states, encoder_hidden_states\n else:\n return hidden_states\n\n\nclass CogVideoXAttnProcessor2_0:\n r\"\"\"\n Processor for implementing scaled dot-product attention for the CogVideoX model. It applies a rotary embedding on\n query and key vectors, but does not include spatial normalization.\n \"\"\"\n\n def __init__(self):\n if not hasattr(F, \"scaled_dot_product_attention\"):\n raise ImportError(\"CogVideoXAttnProcessor requires PyTorch 2.0, to use it, please upgrade PyTorch to 2.0.\")\n\n def __call__(\n self,\n attn: Attention,\n hidden_states: torch.Tensor,\n encoder_hidden_states: torch.Tensor,\n attention_mask: Optional[torch.Tensor] = None,\n image_rotary_emb: Optional[torch.Tensor] = None,\n ) -> torch.Tensor:\n text_seq_length = encoder_hidden_states.size(1)\n\n hidden_states = torch.cat([encoder_hidden_states, hidden_states], dim=1)\n\n batch_size, sequence_length, _ = (\n hidden_states.shape if encoder_hidden_states is None else encoder_hidden_states.shape\n )\n\n if attention_mask is not None:\n attention_mask = attn.prepare_attention_mask(attention_mask, sequence_length, batch_size)\n attention_mask = attention_mask.view(batch_size, attn.heads, -1, attention_mask.shape[-1])\n\n query = attn.to_q(hidden_states)\n key = attn.to_k(hidden_states)\n value = attn.to_v(hidden_states)\n\n inner_dim = key.shape[-1]\n head_dim = inner_dim // attn.heads\n\n query = query.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2)\n key = key.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2)\n value = value.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2)\n\n if attn.norm_q is not None:\n query = attn.norm_q(query)\n if attn.norm_k is not None:\n key = attn.norm_k(key)\n\n # Apply RoPE if needed\n if image_rotary_emb is not None:\n from diffusers.models.embeddings import apply_rotary_emb\n\n if query.shape[-2] == 19726:\n \n num_frames = 13\n\n text_temp = query[:, :, :text_seq_length]\n query_temp = rearrange(query[:, :, text_seq_length:], \"b l (n t) d -> b l n t d\",n=num_frames)[:,:,:,:-150,:].flatten(2,3)\n traj_temp = rearrange(query[:, :, text_seq_length:], \"b l (n t) d -> b l n t d\",n=num_frames)[:,:,:,-150:,:].flatten(2,3)\n query_temp = apply_rotary_emb(query_temp, image_rotary_emb)\n query = torch.cat((text_temp, query_temp, traj_temp), dim=-2)\n\n if not attn.is_cross_attention:\n text_temp = key[:, :, :text_seq_length]\n key_temp = rearrange(key[:, :, text_seq_length:], \"b l (n t) d -> b l n t d\",n=num_frames)[:,:,:,:-150,:].flatten(2,3)\n traj_temp = rearrange(key[:, :, text_seq_length:], \"b l (n t) d -> b l n t d\",n=num_frames)[:,:,:,-150:,:].flatten(2,3)\n key_temp = apply_rotary_emb(key_temp, image_rotary_emb)\n key = torch.cat((text_temp, key_temp, traj_temp), dim=-2)\n\n else:\n query[:, :, text_seq_length:] = apply_rotary_emb(query[:, :, text_seq_length:], image_rotary_emb)\n if not attn.is_cross_attention:\n key[:, :, text_seq_length:] = apply_rotary_emb(key[:, :, text_seq_length:], image_rotary_emb)\n\n hidden_states = F.scaled_dot_product_attention(\n query, key, value, attn_mask=attention_mask, dropout_p=0.0, is_causal=False\n )\n\n hidden_states = hidden_states.transpose(1, 2).reshape(batch_size, -1, attn.heads * head_dim)\n\n # linear proj\n hidden_states = attn.to_out[0](hidden_states)\n # dropout\n hidden_states = attn.to_out[1](hidden_states)\n\n encoder_hidden_states, hidden_states = hidden_states.split(\n [text_seq_length, hidden_states.size(1) - text_seq_length], dim=1\n )\n return hidden_states, encoder_hidden_states\n\n\nclass FusedCogVideoXAttnProcessor2_0:\n r\"\"\"\n Processor for implementing scaled dot-product attention for the CogVideoX model. It applies a rotary embedding on\n query and key vectors, but does not include spatial normalization.\n \"\"\"\n\n def __init__(self):\n if not hasattr(F, \"scaled_dot_product_attention\"):\n raise ImportError(\"CogVideoXAttnProcessor requires PyTorch 2.0, to use it, please upgrade PyTorch to 2.0.\")\n\n def __call__(\n self,\n attn: Attention,\n hidden_states: torch.Tensor,\n encoder_hidden_states: torch.Tensor,\n attention_mask: Optional[torch.Tensor] = None,\n image_rotary_emb: Optional[torch.Tensor] = None,\n ) -> torch.Tensor:\n text_seq_length = encoder_hidden_states.size(1)\n\n hidden_states = torch.cat([encoder_hidden_states, hidden_states], dim=1)\n\n batch_size, sequence_length, _ = (\n hidden_states.shape if encoder_hidden_states is None else encoder_hidden_states.shape\n )\n\n if attention_mask is not None:\n attention_mask = attn.prepare_attention_mask(attention_mask, sequence_length, batch_size)\n attention_mask = attention_mask.view(batch_size, attn.heads, -1, attention_mask.shape[-1])\n\n qkv = attn.to_qkv(hidden_states)\n split_size = qkv.shape[-1] // 3\n query, key, value = torch.split(qkv, split_size, dim=-1)\n\n inner_dim = key.shape[-1]\n head_dim = inner_dim // attn.heads\n\n query = query.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2)\n key = key.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2)\n value = value.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2)\n\n if attn.norm_q is not None:\n query = attn.norm_q(query)\n if attn.norm_k is not None:\n key = attn.norm_k(key)\n\n # Apply RoPE if needed\n if image_rotary_emb is not None:\n from diffusers.models.embeddings import apply_rotary_emb\n\n query[:, :, text_seq_length:] = apply_rotary_emb(query[:, :, text_seq_length:], image_rotary_emb)\n if not attn.is_cross_attention:\n key[:, :, text_seq_length:] = apply_rotary_emb(key[:, :, text_seq_length:], image_rotary_emb)\n\n hidden_states = F.scaled_dot_product_attention(\n query, key, value, attn_mask=attention_mask, dropout_p=0.0, is_causal=False\n )\n\n hidden_states = hidden_states.transpose(1, 2).reshape(batch_size, -1, attn.heads * head_dim)\n\n # linear proj\n hidden_states = attn.to_out[0](hidden_states)\n # dropout\n hidden_states = attn.to_out[1](hidden_states)\n\n encoder_hidden_states, hidden_states = hidden_states.split(\n [text_seq_length, hidden_states.size(1) - text_seq_length], dim=1\n )\n return hidden_states, encoder_hidden_states\n\n\nclass XFormersAttnAddedKVProcessor:\n r\"\"\"\n Processor for implementing memory efficient attention using xFormers.\n\n Args:\n attention_op (`Callable`, *optional*, defaults to `None`):\n The base\n [operator](https://facebookresearch.github.io/xformers/components/ops.html#xformers.ops.AttentionOpBase) to\n use as the attention operator. It is recommended to set to `None`, and allow xFormers to choose the best\n operator.\n \"\"\"\n\n def __init__(self, attention_op: Optional[Callable] = None):\n self.attention_op = attention_op\n\n def __call__(\n self,\n attn: Attention,\n hidden_states: torch.Tensor,\n encoder_hidden_states: Optional[torch.Tensor] = None,\n attention_mask: Optional[torch.Tensor] = None,\n ) -> torch.Tensor:\n residual = hidden_states\n hidden_states = hidden_states.view(hidden_states.shape[0], hidden_states.shape[1], -1).transpose(1, 2)\n batch_size, sequence_length, _ = hidden_states.shape\n\n attention_mask = attn.prepare_attention_mask(attention_mask, sequence_length, batch_size)\n\n if encoder_hidden_states is None:\n encoder_hidden_states = hidden_states\n elif attn.norm_cross:\n encoder_hidden_states = attn.norm_encoder_hidden_states(encoder_hidden_states)\n\n hidden_states = attn.group_norm(hidden_states.transpose(1, 2)).transpose(1, 2)\n\n query = attn.to_q(hidden_states)\n query = attn.head_to_batch_dim(query)\n\n encoder_hidden_states_key_proj = attn.add_k_proj(encoder_hidden_states)\n encoder_hidden_states_value_proj = attn.add_v_proj(encoder_hidden_states)\n encoder_hidden_states_key_proj = attn.head_to_batch_dim(encoder_hidden_states_key_proj)\n encoder_hidden_states_value_proj = attn.head_to_batch_dim(encoder_hidden_states_value_proj)\n\n if not attn.only_cross_attention:\n key = attn.to_k(hidden_states)\n value = attn.to_v(hidden_states)\n key = attn.head_to_batch_dim(key)\n value = attn.head_to_batch_dim(value)\n key = torch.cat([encoder_hidden_states_key_proj, key], dim=1)\n value = torch.cat([encoder_hidden_states_value_proj, value], dim=1)\n else:\n key = encoder_hidden_states_key_proj\n value = encoder_hidden_states_value_proj\n\n hidden_states = xformers.ops.memory_efficient_attention(\n query, key, value, attn_bias=attention_mask, op=self.attention_op, scale=attn.scale\n )\n hidden_states = hidden_states.to(query.dtype)\n hidden_states = attn.batch_to_head_dim(hidden_states)\n\n # linear proj\n hidden_states = attn.to_out[0](hidden_states)\n # dropout\n hidden_states = attn.to_out[1](hidden_states)\n\n hidden_states = hidden_states.transpose(-1, -2).reshape(residual.shape)\n hidden_states = hidden_states + residual\n\n return hidden_states\n\n\nclass XFormersAttnProcessor:\n r\"\"\"\n Processor for implementing memory efficient attention using xFormers.\n\n Args:\n attention_op (`Callable`, *optional*, defaults to `None`):\n The base\n [operator](https://facebookresearch.github.io/xformers/components/ops.html#xformers.ops.AttentionOpBase) to\n use as the attention operator. It is recommended to set to `None`, and allow xFormers to choose the best\n operator.\n \"\"\"\n\n def __init__(self, attention_op: Optional[Callable] = None):\n self.attention_op = attention_op\n\n def __call__(\n self,\n attn: Attention,\n hidden_states: torch.Tensor,\n encoder_hidden_states: Optional[torch.Tensor] = None,\n attention_mask: Optional[torch.Tensor] = None,\n temb: Optional[torch.Tensor] = None,\n *args,\n **kwargs,\n ) -> torch.Tensor:\n if len(args) > 0 or kwargs.get(\"scale\", None) is not None:\n deprecation_message = \"The `scale` argument is deprecated and will be ignored. Please remove it, as passing it will raise an error in the future. `scale` should directly be passed while calling the underlying pipeline component i.e., via `cross_attention_kwargs`.\"\n deprecate(\"scale\", \"1.0.0\", deprecation_message)\n\n residual = hidden_states\n\n if attn.spatial_norm is not None:\n hidden_states = attn.spatial_norm(hidden_states, temb)\n\n input_ndim = hidden_states.ndim\n\n if input_ndim == 4:\n batch_size, channel, height, width = hidden_states.shape\n hidden_states = hidden_states.view(batch_size, channel, height * width).transpose(1, 2)\n\n batch_size, key_tokens, _ = (\n hidden_states.shape if encoder_hidden_states is None else encoder_hidden_states.shape\n )\n\n attention_mask = attn.prepare_attention_mask(attention_mask, key_tokens, batch_size)\n if attention_mask is not None:\n # expand our mask's singleton query_tokens dimension:\n # [batch*heads, 1, key_tokens] ->\n # [batch*heads, query_tokens, key_tokens]\n # so that it can be added as a bias onto the attention scores that xformers computes:\n # [batch*heads, query_tokens, key_tokens]\n # we do this explicitly because xformers doesn't broadcast the singleton dimension for us.\n _, query_tokens, _ = hidden_states.shape\n attention_mask = attention_mask.expand(-1, query_tokens, -1)\n\n if attn.group_norm is not None:\n hidden_states = attn.group_norm(hidden_states.transpose(1, 2)).transpose(1, 2)\n\n query = attn.to_q(hidden_states)\n\n if encoder_hidden_states is None:\n encoder_hidden_states = hidden_states\n elif attn.norm_cross:\n encoder_hidden_states = attn.norm_encoder_hidden_states(encoder_hidden_states)\n\n key = attn.to_k(encoder_hidden_states)\n value = attn.to_v(encoder_hidden_states)\n\n query = attn.head_to_batch_dim(query).contiguous()\n key = attn.head_to_batch_dim(key).contiguous()\n value = attn.head_to_batch_dim(value).contiguous()\n\n hidden_states = xformers.ops.memory_efficient_attention(\n query, key, value, attn_bias=attention_mask, op=self.attention_op, scale=attn.scale\n )\n hidden_states = hidden_states.to(query.dtype)\n hidden_states = attn.batch_to_head_dim(hidden_states)\n\n # linear proj\n hidden_states = attn.to_out[0](hidden_states)\n # dropout\n hidden_states = attn.to_out[1](hidden_states)\n\n if input_ndim == 4:\n hidden_states = hidden_states.transpose(-1, -2).reshape(batch_size, channel, height, width)\n\n if attn.residual_connection:\n hidden_states = hidden_states + residual\n\n hidden_states = hidden_states / attn.rescale_output_factor\n\n return hidden_states\n\n\nclass AttnProcessorNPU:\n r\"\"\"\n Processor for implementing flash attention using torch_npu. Torch_npu supports only fp16 and bf16 data types. If\n fp32 is used, F.scaled_dot_product_attention will be used for computation, but the acceleration effect on NPU is\n not significant.\n\n \"\"\"\n\n def __init__(self):\n if not is_torch_npu_available():\n raise ImportError(\"AttnProcessorNPU requires torch_npu extensions and is supported only on npu devices.\")\n\n def __call__(\n self,\n attn: Attention,\n hidden_states: torch.Tensor,\n encoder_hidden_states: Optional[torch.Tensor] = None,\n attention_mask: Optional[torch.Tensor] = None,\n temb: Optional[torch.Tensor] = None,\n *args,\n **kwargs,\n ) -> torch.Tensor:\n if len(args) > 0 or kwargs.get(\"scale\", None) is not None:\n deprecation_message = \"The `scale` argument is deprecated and will be ignored. Please remove it, as passing it will raise an error in the future. `scale` should directly be passed while calling the underlying pipeline component i.e., via `cross_attention_kwargs`.\"\n deprecate(\"scale\", \"1.0.0\", deprecation_message)\n\n residual = hidden_states\n if attn.spatial_norm is not None:\n hidden_states = attn.spatial_norm(hidden_states, temb)\n\n input_ndim = hidden_states.ndim\n\n if input_ndim == 4:\n batch_size, channel, height, width = hidden_states.shape\n hidden_states = hidden_states.view(batch_size, channel, height * width).transpose(1, 2)\n\n batch_size, sequence_length, _ = (\n hidden_states.shape if encoder_hidden_states is None else encoder_hidden_states.shape\n )\n\n if attention_mask is not None:\n attention_mask = attn.prepare_attention_mask(attention_mask, sequence_length, batch_size)\n # scaled_dot_product_attention expects attention_mask shape to be\n # (batch, heads, source_length, target_length)\n attention_mask = attention_mask.view(batch_size, attn.heads, -1, attention_mask.shape[-1])\n\n if attn.group_norm is not None:\n hidden_states = attn.group_norm(hidden_states.transpose(1, 2)).transpose(1, 2)\n\n query = attn.to_q(hidden_states)\n\n if encoder_hidden_states is None:\n encoder_hidden_states = hidden_states\n elif attn.norm_cross:\n encoder_hidden_states = attn.norm_encoder_hidden_states(encoder_hidden_states)\n\n key = attn.to_k(encoder_hidden_states)\n value = attn.to_v(encoder_hidden_states)\n\n inner_dim = key.shape[-1]\n head_dim = inner_dim // attn.heads\n\n query = query.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2)\n\n key = key.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2)\n value = value.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2)\n\n # the output of sdp = (batch, num_heads, seq_len, head_dim)\n if query.dtype in (torch.float16, torch.bfloat16):\n hidden_states = torch_npu.npu_fusion_attention(\n query,\n key,\n value,\n attn.heads,\n input_layout=\"BNSD\",\n pse=None,\n atten_mask=attention_mask,\n scale=1.0 / math.sqrt(query.shape[-1]),\n pre_tockens=65536,\n next_tockens=65536,\n keep_prob=1.0,\n sync=False,\n inner_precise=0,\n )[0]\n else:\n # TODO: add support for attn.scale when we move to Torch 2.1\n hidden_states = F.scaled_dot_product_attention(\n query, key, value, attn_mask=attention_mask, dropout_p=0.0, is_causal=False\n )\n\n hidden_states = hidden_states.transpose(1, 2).reshape(batch_size, -1, attn.heads * head_dim)\n hidden_states = hidden_states.to(query.dtype)\n\n # linear proj\n hidden_states = attn.to_out[0](hidden_states)\n # dropout\n hidden_states = attn.to_out[1](hidden_states)\n\n if input_ndim == 4:\n hidden_states = hidden_states.transpose(-1, -2).reshape(batch_size, channel, height, width)\n\n if attn.residual_connection:\n hidden_states = hidden_states + residual\n\n hidden_states = hidden_states / attn.rescale_output_factor\n\n return hidden_states\n\n\nclass AttnProcessor2_0:\n r\"\"\"\n Processor for implementing scaled dot-product attention (enabled by default if you're using PyTorch 2.0).\n \"\"\"\n\n def __init__(self):\n if not hasattr(F, \"scaled_dot_product_attention\"):\n raise ImportError(\"AttnProcessor2_0 requires PyTorch 2.0, to use it, please upgrade PyTorch to 2.0.\")\n\n def __call__(\n self,\n attn: Attention,\n hidden_states: torch.Tensor,\n encoder_hidden_states: Optional[torch.Tensor] = None,\n attention_mask: Optional[torch.Tensor] = None,\n temb: Optional[torch.Tensor] = None,\n *args,\n **kwargs,\n ) -> torch.Tensor:\n if len(args) > 0 or kwargs.get(\"scale\", None) is not None:\n deprecation_message = \"The `scale` argument is deprecated and will be ignored. Please remove it, as passing it will raise an error in the future. `scale` should directly be passed while calling the underlying pipeline component i.e., via `cross_attention_kwargs`.\"\n deprecate(\"scale\", \"1.0.0\", deprecation_message)\n\n residual = hidden_states\n if attn.spatial_norm is not None:\n hidden_states = attn.spatial_norm(hidden_states, temb)\n\n input_ndim = hidden_states.ndim\n\n if input_ndim == 4:\n batch_size, channel, height, width = hidden_states.shape\n hidden_states = hidden_states.view(batch_size, channel, height * width).transpose(1, 2)\n\n batch_size, sequence_length, _ = (\n hidden_states.shape if encoder_hidden_states is None else encoder_hidden_states.shape\n )\n\n if attention_mask is not None:\n attention_mask = attn.prepare_attention_mask(attention_mask, sequence_length, batch_size)\n # scaled_dot_product_attention expects attention_mask shape to be\n # (batch, heads, source_length, target_length)\n attention_mask = attention_mask.view(batch_size, attn.heads, -1, attention_mask.shape[-1])\n\n if attn.group_norm is not None:\n hidden_states = attn.group_norm(hidden_states.transpose(1, 2)).transpose(1, 2)\n\n query = attn.to_q(hidden_states)\n\n if encoder_hidden_states is None:\n encoder_hidden_states = hidden_states\n elif attn.norm_cross:\n encoder_hidden_states = attn.norm_encoder_hidden_states(encoder_hidden_states)\n\n key = attn.to_k(encoder_hidden_states)\n value = attn.to_v(encoder_hidden_states)\n\n inner_dim = key.shape[-1]\n head_dim = inner_dim // attn.heads\n\n query = query.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2)\n\n key = key.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2)\n value = value.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2)\n\n if attn.norm_q is not None:\n query = attn.norm_q(query)\n if attn.norm_k is not None:\n key = attn.norm_k(key)\n\n # the output of sdp = (batch, num_heads, seq_len, head_dim)\n # TODO: add support for attn.scale when we move to Torch 2.1\n hidden_states = F.scaled_dot_product_attention(\n query, key, value, attn_mask=attention_mask, dropout_p=0.0, is_causal=False\n )\n\n hidden_states = hidden_states.transpose(1, 2).reshape(batch_size, -1, attn.heads * head_dim)\n hidden_states = hidden_states.to(query.dtype)\n\n # linear proj\n hidden_states = attn.to_out[0](hidden_states)\n # dropout\n hidden_states = attn.to_out[1](hidden_states)\n\n if input_ndim == 4:\n hidden_states = hidden_states.transpose(-1, -2).reshape(batch_size, channel, height, width)\n\n if attn.residual_connection:\n hidden_states = hidden_states + residual\n\n hidden_states = hidden_states / attn.rescale_output_factor\n\n return hidden_states\n\n\nclass StableAudioAttnProcessor2_0:\n r\"\"\"\n Processor for implementing scaled dot-product attention (enabled by default if you're using PyTorch 2.0). This is\n used in the Stable Audio model. It applies rotary embedding on query and key vector, and allows MHA, GQA or MQA.\n \"\"\"\n\n def __init__(self):\n if not hasattr(F, \"scaled_dot_product_attention\"):\n raise ImportError(\n \"StableAudioAttnProcessor2_0 requires PyTorch 2.0, to use it, please upgrade PyTorch to 2.0.\"\n )\n\n def apply_partial_rotary_emb(\n self,\n x: torch.Tensor,\n freqs_cis: Tuple[torch.Tensor],\n ) -> torch.Tensor:\n from diffusers.models.embeddings import apply_rotary_emb\n\n rot_dim = freqs_cis[0].shape[-1]\n x_to_rotate, x_unrotated = x[..., :rot_dim], x[..., rot_dim:]\n\n x_rotated = apply_rotary_emb(x_to_rotate, freqs_cis, use_real=True, use_real_unbind_dim=-2)\n\n out = torch.cat((x_rotated, x_unrotated), dim=-1)\n return out\n\n def __call__(\n self,\n attn: Attention,\n hidden_states: torch.Tensor,\n encoder_hidden_states: Optional[torch.Tensor] = None,\n attention_mask: Optional[torch.Tensor] = None,\n rotary_emb: Optional[torch.Tensor] = None,\n ) -> torch.Tensor:\n from diffusers.models.embeddings import apply_rotary_emb\n\n residual = hidden_states\n\n input_ndim = hidden_states.ndim\n\n if input_ndim == 4:\n batch_size, channel, height, width = hidden_states.shape\n hidden_states = hidden_states.view(batch_size, channel, height * width).transpose(1, 2)\n\n batch_size, sequence_length, _ = (\n hidden_states.shape if encoder_hidden_states is None else encoder_hidden_states.shape\n )\n\n if attention_mask is not None:\n attention_mask = attn.prepare_attention_mask(attention_mask, sequence_length, batch_size)\n # scaled_dot_product_attention expects attention_mask shape to be\n # (batch, heads, source_length, target_length)\n attention_mask = attention_mask.view(batch_size, attn.heads, -1, attention_mask.shape[-1])\n\n query = attn.to_q(hidden_states)\n\n if encoder_hidden_states is None:\n encoder_hidden_states = hidden_states\n elif attn.norm_cross:\n encoder_hidden_states = attn.norm_encoder_hidden_states(encoder_hidden_states)\n\n key = attn.to_k(encoder_hidden_states)\n value = attn.to_v(encoder_hidden_states)\n\n head_dim = query.shape[-1] // attn.heads\n kv_heads = key.shape[-1] // head_dim\n\n query = query.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2)\n\n key = key.view(batch_size, -1, kv_heads, head_dim).transpose(1, 2)\n value = value.view(batch_size, -1, kv_heads, head_dim).transpose(1, 2)\n\n if kv_heads != attn.heads:\n # if GQA or MQA, repeat the key/value heads to reach the number of query heads.\n heads_per_kv_head = attn.heads // kv_heads\n key = torch.repeat_interleave(key, heads_per_kv_head, dim=1)\n value = torch.repeat_interleave(value, heads_per_kv_head, dim=1)\n\n if attn.norm_q is not None:\n query = attn.norm_q(query)\n if attn.norm_k is not None:\n key = attn.norm_k(key)\n\n # Apply RoPE if needed\n if rotary_emb is not None:\n query_dtype = query.dtype\n key_dtype = key.dtype\n query = query.to(torch.float32)\n key = key.to(torch.float32)\n\n rot_dim = rotary_emb[0].shape[-1]\n query_to_rotate, query_unrotated = query[..., :rot_dim], query[..., rot_dim:]\n query_rotated = apply_rotary_emb(query_to_rotate, rotary_emb, use_real=True, use_real_unbind_dim=-2)\n\n query = torch.cat((query_rotated, query_unrotated), dim=-1)\n\n if not attn.is_cross_attention:\n key_to_rotate, key_unrotated = key[..., :rot_dim], key[..., rot_dim:]\n key_rotated = apply_rotary_emb(key_to_rotate, rotary_emb, use_real=True, use_real_unbind_dim=-2)\n\n key = torch.cat((key_rotated, key_unrotated), dim=-1)\n\n query = query.to(query_dtype)\n key = key.to(key_dtype)\n\n # the output of sdp = (batch, num_heads, seq_len, head_dim)\n # TODO: add support for attn.scale when we move to Torch 2.1\n hidden_states = F.scaled_dot_product_attention(\n query, key, value, attn_mask=attention_mask, dropout_p=0.0, is_causal=False\n )\n\n hidden_states = hidden_states.transpose(1, 2).reshape(batch_size, -1, attn.heads * head_dim)\n hidden_states = hidden_states.to(query.dtype)\n\n # linear proj\n hidden_states = attn.to_out[0](hidden_states)\n # dropout\n hidden_states = attn.to_out[1](hidden_states)\n\n if input_ndim == 4:\n hidden_states = hidden_states.transpose(-1, -2).reshape(batch_size, channel, height, width)\n\n if attn.residual_connection:\n hidden_states = hidden_states + residual\n\n hidden_states = hidden_states / attn.rescale_output_factor\n\n return hidden_states\n\n\nclass HunyuanAttnProcessor2_0:\n r\"\"\"\n Processor for implementing scaled dot-product attention (enabled by default if you're using PyTorch 2.0). This is\n used in the HunyuanDiT model. It applies a s normalization layer and rotary embedding on query and key vector.\n \"\"\"\n\n def __init__(self):\n if not hasattr(F, \"scaled_dot_product_attention\"):\n raise ImportError(\"AttnProcessor2_0 requires PyTorch 2.0, to use it, please upgrade PyTorch to 2.0.\")\n\n def __call__(\n self,\n attn: Attention,\n hidden_states: torch.Tensor,\n encoder_hidden_states: Optional[torch.Tensor] = None,\n attention_mask: Optional[torch.Tensor] = None,\n temb: Optional[torch.Tensor] = None,\n image_rotary_emb: Optional[torch.Tensor] = None,\n ) -> torch.Tensor:\n from diffusers.models.embeddings import apply_rotary_emb\n\n residual = hidden_states\n if attn.spatial_norm is not None:\n hidden_states = attn.spatial_norm(hidden_states, temb)\n\n input_ndim = hidden_states.ndim\n\n if input_ndim == 4:\n batch_size, channel, height, width = hidden_states.shape\n hidden_states = hidden_states.view(batch_size, channel, height * width).transpose(1, 2)\n\n batch_size, sequence_length, _ = (\n hidden_states.shape if encoder_hidden_states is None else encoder_hidden_states.shape\n )\n\n if attention_mask is not None:\n attention_mask = attn.prepare_attention_mask(attention_mask, sequence_length, batch_size)\n # scaled_dot_product_attention expects attention_mask shape to be\n # (batch, heads, source_length, target_length)\n attention_mask = attention_mask.view(batch_size, attn.heads, -1, attention_mask.shape[-1])\n\n if attn.group_norm is not None:\n hidden_states = attn.group_norm(hidden_states.transpose(1, 2)).transpose(1, 2)\n\n query = attn.to_q(hidden_states)\n\n if encoder_hidden_states is None:\n encoder_hidden_states = hidden_states\n elif attn.norm_cross:\n encoder_hidden_states = attn.norm_encoder_hidden_states(encoder_hidden_states)\n\n key = attn.to_k(encoder_hidden_states)\n value = attn.to_v(encoder_hidden_states)\n\n inner_dim = key.shape[-1]\n head_dim = inner_dim // attn.heads\n\n query = query.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2)\n\n key = key.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2)\n value = value.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2)\n\n if attn.norm_q is not None:\n query = attn.norm_q(query)\n if attn.norm_k is not None:\n key = attn.norm_k(key)\n\n # Apply RoPE if needed\n if image_rotary_emb is not None:\n query = apply_rotary_emb(query, image_rotary_emb)\n if not attn.is_cross_attention:\n key = apply_rotary_emb(key, image_rotary_emb)\n\n # the output of sdp = (batch, num_heads, seq_len, head_dim)\n # TODO: add support for attn.scale when we move to Torch 2.1\n hidden_states = F.scaled_dot_product_attention(\n query, key, value, attn_mask=attention_mask, dropout_p=0.0, is_causal=False\n )\n\n hidden_states = hidden_states.transpose(1, 2).reshape(batch_size, -1, attn.heads * head_dim)\n hidden_states = hidden_states.to(query.dtype)\n\n # linear proj\n hidden_states = attn.to_out[0](hidden_states)\n # dropout\n hidden_states = attn.to_out[1](hidden_states)\n\n if input_ndim == 4:\n hidden_states = hidden_states.transpose(-1, -2).reshape(batch_size, channel, height, width)\n\n if attn.residual_connection:\n hidden_states = hidden_states + residual\n\n hidden_states = hidden_states / attn.rescale_output_factor\n\n return hidden_states\n\n\nclass FusedHunyuanAttnProcessor2_0:\n r\"\"\"\n Processor for implementing scaled dot-product attention (enabled by default if you're using PyTorch 2.0) with fused\n projection layers. This is used in the HunyuanDiT model. It applies a s normalization layer and rotary embedding on\n query and key vector.\n \"\"\"\n\n def __init__(self):\n if not hasattr(F, \"scaled_dot_product_attention\"):\n raise ImportError(\n \"FusedHunyuanAttnProcessor2_0 requires PyTorch 2.0, to use it, please upgrade PyTorch to 2.0.\"\n )\n\n def __call__(\n self,\n attn: Attention,\n hidden_states: torch.Tensor,\n encoder_hidden_states: Optional[torch.Tensor] = None,\n attention_mask: Optional[torch.Tensor] = None,\n temb: Optional[torch.Tensor] = None,\n image_rotary_emb: Optional[torch.Tensor] = None,\n ) -> torch.Tensor:\n from diffusers.models.embeddings import apply_rotary_emb\n\n residual = hidden_states\n if attn.spatial_norm is not None:\n hidden_states = attn.spatial_norm(hidden_states, temb)\n\n input_ndim = hidden_states.ndim\n\n if input_ndim == 4:\n batch_size, channel, height, width = hidden_states.shape\n hidden_states = hidden_states.view(batch_size, channel, height * width).transpose(1, 2)\n\n batch_size, sequence_length, _ = (\n hidden_states.shape if encoder_hidden_states is None else encoder_hidden_states.shape\n )\n\n if attention_mask is not None:\n attention_mask = attn.prepare_attention_mask(attention_mask, sequence_length, batch_size)\n # scaled_dot_product_attention expects attention_mask shape to be\n # (batch, heads, source_length, target_length)\n attention_mask = attention_mask.view(batch_size, attn.heads, -1, attention_mask.shape[-1])\n\n if attn.group_norm is not None:\n hidden_states = attn.group_norm(hidden_states.transpose(1, 2)).transpose(1, 2)\n\n if encoder_hidden_states is None:\n qkv = attn.to_qkv(hidden_states)\n split_size = qkv.shape[-1] // 3\n query, key, value = torch.split(qkv, split_size, dim=-1)\n else:\n if attn.norm_cross:\n encoder_hidden_states = attn.norm_encoder_hidden_states(encoder_hidden_states)\n query = attn.to_q(hidden_states)\n\n kv = attn.to_kv(encoder_hidden_states)\n split_size = kv.shape[-1] // 2\n key, value = torch.split(kv, split_size, dim=-1)\n\n inner_dim = key.shape[-1]\n head_dim = inner_dim // attn.heads\n\n query = query.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2)\n key = key.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2)\n value = value.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2)\n\n if attn.norm_q is not None:\n query = attn.norm_q(query)\n if attn.norm_k is not None:\n key = attn.norm_k(key)\n\n # Apply RoPE if needed\n if image_rotary_emb is not None:\n query = apply_rotary_emb(query, image_rotary_emb)\n if not attn.is_cross_attention:\n key = apply_rotary_emb(key, image_rotary_emb)\n\n # the output of sdp = (batch, num_heads, seq_len, head_dim)\n # TODO: add support for attn.scale when we move to Torch 2.1\n hidden_states = F.scaled_dot_product_attention(\n query, key, value, attn_mask=attention_mask, dropout_p=0.0, is_causal=False\n )\n\n hidden_states = hidden_states.transpose(1, 2).reshape(batch_size, -1, attn.heads * head_dim)\n hidden_states = hidden_states.to(query.dtype)\n\n # linear proj\n hidden_states = attn.to_out[0](hidden_states)\n # dropout\n hidden_states = attn.to_out[1](hidden_states)\n\n if input_ndim == 4:\n hidden_states = hidden_states.transpose(-1, -2).reshape(batch_size, channel, height, width)\n\n if attn.residual_connection:\n hidden_states = hidden_states + residual\n\n hidden_states = hidden_states / attn.rescale_output_factor\n\n return hidden_states\n\n\nclass PAGHunyuanAttnProcessor2_0:\n r\"\"\"\n Processor for implementing scaled dot-product attention (enabled by default if you're using PyTorch 2.0). This is\n used in the HunyuanDiT model. It applies a normalization layer and rotary embedding on query and key vector. This\n variant of the processor employs [Pertubed Attention Guidance](https://arxiv.org/abs/2403.17377).\n \"\"\"\n\n def __init__(self):\n if not hasattr(F, \"scaled_dot_product_attention\"):\n raise ImportError(\n \"PAGHunyuanAttnProcessor2_0 requires PyTorch 2.0, to use it, please upgrade PyTorch to 2.0.\"\n )\n\n def __call__(\n self,\n attn: Attention,\n hidden_states: torch.Tensor,\n encoder_hidden_states: Optional[torch.Tensor] = None,\n attention_mask: Optional[torch.Tensor] = None,\n temb: Optional[torch.Tensor] = None,\n image_rotary_emb: Optional[torch.Tensor] = None,\n ) -> torch.Tensor:\n from diffusers.models.embeddings import apply_rotary_emb\n\n residual = hidden_states\n if attn.spatial_norm is not None:\n hidden_states = attn.spatial_norm(hidden_states, temb)\n\n input_ndim = hidden_states.ndim\n\n if input_ndim == 4:\n batch_size, channel, height, width = hidden_states.shape\n hidden_states = hidden_states.view(batch_size, channel, height * width).transpose(1, 2)\n\n # chunk\n hidden_states_org, hidden_states_ptb = hidden_states.chunk(2)\n\n # 1. Original Path\n batch_size, sequence_length, _ = (\n hidden_states_org.shape if encoder_hidden_states is None else encoder_hidden_states.shape\n )\n\n if attention_mask is not None:\n attention_mask = attn.prepare_attention_mask(attention_mask, sequence_length, batch_size)\n # scaled_dot_product_attention expects attention_mask shape to be\n # (batch, heads, source_length, target_length)\n attention_mask = attention_mask.view(batch_size, attn.heads, -1, attention_mask.shape[-1])\n\n if attn.group_norm is not None:\n hidden_states_org = attn.group_norm(hidden_states_org.transpose(1, 2)).transpose(1, 2)\n\n query = attn.to_q(hidden_states_org)\n\n if encoder_hidden_states is None:\n encoder_hidden_states = hidden_states_org\n elif attn.norm_cross:\n encoder_hidden_states = attn.norm_encoder_hidden_states(encoder_hidden_states)\n\n key = attn.to_k(encoder_hidden_states)\n value = attn.to_v(encoder_hidden_states)\n\n inner_dim = key.shape[-1]\n head_dim = inner_dim // attn.heads\n\n query = query.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2)\n\n key = key.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2)\n value = value.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2)\n\n if attn.norm_q is not None:\n query = attn.norm_q(query)\n if attn.norm_k is not None:\n key = attn.norm_k(key)\n\n # Apply RoPE if needed\n if image_rotary_emb is not None:\n query = apply_rotary_emb(query, image_rotary_emb)\n if not attn.is_cross_attention:\n key = apply_rotary_emb(key, image_rotary_emb)\n\n # the output of sdp = (batch, num_heads, seq_len, head_dim)\n # TODO: add support for attn.scale when we move to Torch 2.1\n hidden_states_org = F.scaled_dot_product_attention(\n query, key, value, attn_mask=attention_mask, dropout_p=0.0, is_causal=False\n )\n\n hidden_states_org = hidden_states_org.transpose(1, 2).reshape(batch_size, -1, attn.heads * head_dim)\n hidden_states_org = hidden_states_org.to(query.dtype)\n\n # linear proj\n hidden_states_org = attn.to_out[0](hidden_states_org)\n # dropout\n hidden_states_org = attn.to_out[1](hidden_states_org)\n\n if input_ndim == 4:\n hidden_states_org = hidden_states_org.transpose(-1, -2).reshape(batch_size, channel, height, width)\n\n # 2. Perturbed Path\n if attn.group_norm is not None:\n hidden_states_ptb = attn.group_norm(hidden_states_ptb.transpose(1, 2)).transpose(1, 2)\n\n hidden_states_ptb = attn.to_v(hidden_states_ptb)\n hidden_states_ptb = hidden_states_ptb.to(query.dtype)\n\n # linear proj\n hidden_states_ptb = attn.to_out[0](hidden_states_ptb)\n # dropout\n hidden_states_ptb = attn.to_out[1](hidden_states_ptb)\n\n if input_ndim == 4:\n hidden_states_ptb = hidden_states_ptb.transpose(-1, -2).reshape(batch_size, channel, height, width)\n\n # cat\n hidden_states = torch.cat([hidden_states_org, hidden_states_ptb])\n\n if attn.residual_connection:\n hidden_states = hidden_states + residual\n\n hidden_states = hidden_states / attn.rescale_output_factor\n\n return hidden_states\n\n\nclass PAGCFGHunyuanAttnProcessor2_0:\n r\"\"\"\n Processor for implementing scaled dot-product attention (enabled by default if you're using PyTorch 2.0). This is\n used in the HunyuanDiT model. It applies a normalization layer and rotary embedding on query and key vector. This\n variant of the processor employs [Pertubed Attention Guidance](https://arxiv.org/abs/2403.17377).\n \"\"\"\n\n def __init__(self):\n if not hasattr(F, \"scaled_dot_product_attention\"):\n raise ImportError(\n \"PAGCFGHunyuanAttnProcessor2_0 requires PyTorch 2.0, to use it, please upgrade PyTorch to 2.0.\"\n )\n\n def __call__(\n self,\n attn: Attention,\n hidden_states: torch.Tensor,\n encoder_hidden_states: Optional[torch.Tensor] = None,\n attention_mask: Optional[torch.Tensor] = None,\n temb: Optional[torch.Tensor] = None,\n image_rotary_emb: Optional[torch.Tensor] = None,\n ) -> torch.Tensor:\n from diffusers.models.embeddings import apply_rotary_emb\n\n residual = hidden_states\n if attn.spatial_norm is not None:\n hidden_states = attn.spatial_norm(hidden_states, temb)\n\n input_ndim = hidden_states.ndim\n\n if input_ndim == 4:\n batch_size, channel, height, width = hidden_states.shape\n hidden_states = hidden_states.view(batch_size, channel, height * width).transpose(1, 2)\n\n # chunk\n hidden_states_uncond, hidden_states_org, hidden_states_ptb = hidden_states.chunk(3)\n hidden_states_org = torch.cat([hidden_states_uncond, hidden_states_org])\n\n # 1. Original Path\n batch_size, sequence_length, _ = (\n hidden_states_org.shape if encoder_hidden_states is None else encoder_hidden_states.shape\n )\n\n if attention_mask is not None:\n attention_mask = attn.prepare_attention_mask(attention_mask, sequence_length, batch_size)\n # scaled_dot_product_attention expects attention_mask shape to be\n # (batch, heads, source_length, target_length)\n attention_mask = attention_mask.view(batch_size, attn.heads, -1, attention_mask.shape[-1])\n\n if attn.group_norm is not None:\n hidden_states_org = attn.group_norm(hidden_states_org.transpose(1, 2)).transpose(1, 2)\n\n query = attn.to_q(hidden_states_org)\n\n if encoder_hidden_states is None:\n encoder_hidden_states = hidden_states_org\n elif attn.norm_cross:\n encoder_hidden_states = attn.norm_encoder_hidden_states(encoder_hidden_states)\n\n key = attn.to_k(encoder_hidden_states)\n value = attn.to_v(encoder_hidden_states)\n\n inner_dim = key.shape[-1]\n head_dim = inner_dim // attn.heads\n\n query = query.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2)\n\n key = key.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2)\n value = value.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2)\n\n if attn.norm_q is not None:\n query = attn.norm_q(query)\n if attn.norm_k is not None:\n key = attn.norm_k(key)\n\n # Apply RoPE if needed\n if image_rotary_emb is not None:\n query = apply_rotary_emb(query, image_rotary_emb)\n if not attn.is_cross_attention:\n key = apply_rotary_emb(key, image_rotary_emb)\n\n # the output of sdp = (batch, num_heads, seq_len, head_dim)\n # TODO: add support for attn.scale when we move to Torch 2.1\n hidden_states_org = F.scaled_dot_product_attention(\n query, key, value, attn_mask=attention_mask, dropout_p=0.0, is_causal=False\n )\n\n hidden_states_org = hidden_states_org.transpose(1, 2).reshape(batch_size, -1, attn.heads * head_dim)\n hidden_states_org = hidden_states_org.to(query.dtype)\n\n # linear proj\n hidden_states_org = attn.to_out[0](hidden_states_org)\n # dropout\n hidden_states_org = attn.to_out[1](hidden_states_org)\n\n if input_ndim == 4:\n hidden_states_org = hidden_states_org.transpose(-1, -2).reshape(batch_size, channel, height, width)\n\n # 2. Perturbed Path\n if attn.group_norm is not None:\n hidden_states_ptb = attn.group_norm(hidden_states_ptb.transpose(1, 2)).transpose(1, 2)\n\n hidden_states_ptb = attn.to_v(hidden_states_ptb)\n hidden_states_ptb = hidden_states_ptb.to(query.dtype)\n\n # linear proj\n hidden_states_ptb = attn.to_out[0](hidden_states_ptb)\n # dropout\n hidden_states_ptb = attn.to_out[1](hidden_states_ptb)\n\n if input_ndim == 4:\n hidden_states_ptb = hidden_states_ptb.transpose(-1, -2).reshape(batch_size, channel, height, width)\n\n # cat\n hidden_states = torch.cat([hidden_states_org, hidden_states_ptb])\n\n if attn.residual_connection:\n hidden_states = hidden_states + residual\n\n hidden_states = hidden_states / attn.rescale_output_factor\n\n return hidden_states\n\n\nclass LuminaAttnProcessor2_0:\n r\"\"\"\n Processor for implementing scaled dot-product attention (enabled by default if you're using PyTorch 2.0). This is\n used in the LuminaNextDiT model. It applies a s normalization layer and rotary embedding on query and key vector.\n \"\"\"\n\n def __init__(self):\n if not hasattr(F, \"scaled_dot_product_attention\"):\n raise ImportError(\"AttnProcessor2_0 requires PyTorch 2.0, to use it, please upgrade PyTorch to 2.0.\")\n\n def __call__(\n self,\n attn: Attention,\n hidden_states: torch.Tensor,\n encoder_hidden_states: torch.Tensor,\n attention_mask: Optional[torch.Tensor] = None,\n query_rotary_emb: Optional[torch.Tensor] = None,\n key_rotary_emb: Optional[torch.Tensor] = None,\n base_sequence_length: Optional[int] = None,\n ) -> torch.Tensor:\n from diffusers.models.embeddings import apply_rotary_emb\n\n input_ndim = hidden_states.ndim\n\n if input_ndim == 4:\n batch_size, channel, height, width = hidden_states.shape\n hidden_states = hidden_states.view(batch_size, channel, height * width).transpose(1, 2)\n\n batch_size, sequence_length, _ = hidden_states.shape\n\n # Get Query-Key-Value Pair\n query = attn.to_q(hidden_states)\n key = attn.to_k(encoder_hidden_states)\n value = attn.to_v(encoder_hidden_states)\n\n query_dim = query.shape[-1]\n inner_dim = key.shape[-1]\n head_dim = query_dim // attn.heads\n dtype = query.dtype\n\n # Get key-value heads\n kv_heads = inner_dim // head_dim\n\n # Apply Query-Key Norm if needed\n if attn.norm_q is not None:\n query = attn.norm_q(query)\n if attn.norm_k is not None:\n key = attn.norm_k(key)\n\n query = query.view(batch_size, -1, attn.heads, head_dim)\n\n key = key.view(batch_size, -1, kv_heads, head_dim)\n value = value.view(batch_size, -1, kv_heads, head_dim)\n\n # Apply RoPE if needed\n if query_rotary_emb is not None:\n query = apply_rotary_emb(query, query_rotary_emb, use_real=False)\n if key_rotary_emb is not None:\n key = apply_rotary_emb(key, key_rotary_emb, use_real=False)\n\n query, key = query.to(dtype), key.to(dtype)\n\n # Apply proportional attention if true\n if key_rotary_emb is None:\n softmax_scale = None\n else:\n if base_sequence_length is not None:\n softmax_scale = math.sqrt(math.log(sequence_length, base_sequence_length)) * attn.scale\n else:\n softmax_scale = attn.scale\n\n # perform Grouped-qurey Attention (GQA)\n n_rep = attn.heads // kv_heads\n if n_rep >= 1:\n key = key.unsqueeze(3).repeat(1, 1, 1, n_rep, 1).flatten(2, 3)\n value = value.unsqueeze(3).repeat(1, 1, 1, n_rep, 1).flatten(2, 3)\n\n # scaled_dot_product_attention expects attention_mask shape to be\n # (batch, heads, source_length, target_length)\n attention_mask = attention_mask.bool().view(batch_size, 1, 1, -1)\n attention_mask = attention_mask.expand(-1, attn.heads, sequence_length, -1)\n\n query = query.transpose(1, 2)\n key = key.transpose(1, 2)\n value = value.transpose(1, 2)\n\n # the output of sdp = (batch, num_heads, seq_len, head_dim)\n # TODO: add support for attn.scale when we move to Torch 2.1\n hidden_states = F.scaled_dot_product_attention(\n query, key, value, attn_mask=attention_mask, scale=softmax_scale\n )\n hidden_states = hidden_states.transpose(1, 2).to(dtype)\n\n return hidden_states\n\n\nclass FusedAttnProcessor2_0:\n r\"\"\"\n Processor for implementing scaled dot-product attention (enabled by default if you're using PyTorch 2.0). It uses\n fused projection layers. For self-attention modules, all projection matrices (i.e., query, key, value) are fused.\n For cross-attention modules, key and value projection matrices are fused.\n\n \n\n This API is currently 🧪 experimental in nature and can change in future.\n\n \n \"\"\"\n\n def __init__(self):\n if not hasattr(F, \"scaled_dot_product_attention\"):\n raise ImportError(\n \"FusedAttnProcessor2_0 requires at least PyTorch 2.0, to use it. Please upgrade PyTorch to > 2.0.\"\n )\n\n def __call__(\n self,\n attn: Attention,\n hidden_states: torch.Tensor,\n encoder_hidden_states: Optional[torch.Tensor] = None,\n attention_mask: Optional[torch.Tensor] = None,\n temb: Optional[torch.Tensor] = None,\n *args,\n **kwargs,\n ) -> torch.Tensor:\n if len(args) > 0 or kwargs.get(\"scale\", None) is not None:\n deprecation_message = \"The `scale` argument is deprecated and will be ignored. Please remove it, as passing it will raise an error in the future. `scale` should directly be passed while calling the underlying pipeline component i.e., via `cross_attention_kwargs`.\"\n deprecate(\"scale\", \"1.0.0\", deprecation_message)\n\n residual = hidden_states\n if attn.spatial_norm is not None:\n hidden_states = attn.spatial_norm(hidden_states, temb)\n\n input_ndim = hidden_states.ndim\n\n if input_ndim == 4:\n batch_size, channel, height, width = hidden_states.shape\n hidden_states = hidden_states.view(batch_size, channel, height * width).transpose(1, 2)\n\n batch_size, sequence_length, _ = (\n hidden_states.shape if encoder_hidden_states is None else encoder_hidden_states.shape\n )\n\n if attention_mask is not None:\n attention_mask = attn.prepare_attention_mask(attention_mask, sequence_length, batch_size)\n # scaled_dot_product_attention expects attention_mask shape to be\n # (batch, heads, source_length, target_length)\n attention_mask = attention_mask.view(batch_size, attn.heads, -1, attention_mask.shape[-1])\n\n if attn.group_norm is not None:\n hidden_states = attn.group_norm(hidden_states.transpose(1, 2)).transpose(1, 2)\n\n if encoder_hidden_states is None:\n qkv = attn.to_qkv(hidden_states)\n split_size = qkv.shape[-1] // 3\n query, key, value = torch.split(qkv, split_size, dim=-1)\n else:\n if attn.norm_cross:\n encoder_hidden_states = attn.norm_encoder_hidden_states(encoder_hidden_states)\n query = attn.to_q(hidden_states)\n\n kv = attn.to_kv(encoder_hidden_states)\n split_size = kv.shape[-1] // 2\n key, value = torch.split(kv, split_size, dim=-1)\n\n inner_dim = key.shape[-1]\n head_dim = inner_dim // attn.heads\n\n query = query.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2)\n key = key.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2)\n value = value.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2)\n\n if attn.norm_q is not None:\n query = attn.norm_q(query)\n if attn.norm_k is not None:\n key = attn.norm_k(key)\n\n # the output of sdp = (batch, num_heads, seq_len, head_dim)\n # TODO: add support for attn.scale when we move to Torch 2.1\n hidden_states = F.scaled_dot_product_attention(\n query, key, value, attn_mask=attention_mask, dropout_p=0.0, is_causal=False\n )\n\n hidden_states = hidden_states.transpose(1, 2).reshape(batch_size, -1, attn.heads * head_dim)\n hidden_states = hidden_states.to(query.dtype)\n\n # linear proj\n hidden_states = attn.to_out[0](hidden_states)\n # dropout\n hidden_states = attn.to_out[1](hidden_states)\n\n if input_ndim == 4:\n hidden_states = hidden_states.transpose(-1, -2).reshape(batch_size, channel, height, width)\n\n if attn.residual_connection:\n hidden_states = hidden_states + residual\n\n hidden_states = hidden_states / attn.rescale_output_factor\n\n return hidden_states\n\n\nclass CustomDiffusionXFormersAttnProcessor(nn.Module):\n r\"\"\"\n Processor for implementing memory efficient attention using xFormers for the Custom Diffusion method.\n\n Args:\n train_kv (`bool`, defaults to `True`):\n Whether to newly train the key and value matrices corresponding to the text features.\n train_q_out (`bool`, defaults to `True`):\n Whether to newly train query matrices corresponding to the latent image features.\n hidden_size (`int`, *optional*, defaults to `None`):\n The hidden size of the attention layer.\n cross_attention_dim (`int`, *optional*, defaults to `None`):\n The number of channels in the `encoder_hidden_states`.\n out_bias (`bool`, defaults to `True`):\n Whether to include the bias parameter in `train_q_out`.\n dropout (`float`, *optional*, defaults to 0.0):\n The dropout probability to use.\n attention_op (`Callable`, *optional*, defaults to `None`):\n The base\n [operator](https://facebookresearch.github.io/xformers/components/ops.html#xformers.ops.AttentionOpBase) to use\n as the attention operator. It is recommended to set to `None`, and allow xFormers to choose the best operator.\n \"\"\"\n\n def __init__(\n self,\n train_kv: bool = True,\n train_q_out: bool = False,\n hidden_size: Optional[int] = None,\n cross_attention_dim: Optional[int] = None,\n out_bias: bool = True,\n dropout: float = 0.0,\n attention_op: Optional[Callable] = None,\n ):\n super().__init__()\n self.train_kv = train_kv\n self.train_q_out = train_q_out\n\n self.hidden_size = hidden_size\n self.cross_attention_dim = cross_attention_dim\n self.attention_op = attention_op\n\n # `_custom_diffusion` id for easy serialization and loading.\n if self.train_kv:\n self.to_k_custom_diffusion = nn.Linear(cross_attention_dim or hidden_size, hidden_size, bias=False)\n self.to_v_custom_diffusion = nn.Linear(cross_attention_dim or hidden_size, hidden_size, bias=False)\n if self.train_q_out:\n self.to_q_custom_diffusion = nn.Linear(hidden_size, hidden_size, bias=False)\n self.to_out_custom_diffusion = nn.ModuleList([])\n self.to_out_custom_diffusion.append(nn.Linear(hidden_size, hidden_size, bias=out_bias))\n self.to_out_custom_diffusion.append(nn.Dropout(dropout))\n\n def __call__(\n self,\n attn: Attention,\n hidden_states: torch.Tensor,\n encoder_hidden_states: Optional[torch.Tensor] = None,\n attention_mask: Optional[torch.Tensor] = None,\n ) -> torch.Tensor:\n batch_size, sequence_length, _ = (\n hidden_states.shape if encoder_hidden_states is None else encoder_hidden_states.shape\n )\n\n attention_mask = attn.prepare_attention_mask(attention_mask, sequence_length, batch_size)\n\n if self.train_q_out:\n query = self.to_q_custom_diffusion(hidden_states).to(attn.to_q.weight.dtype)\n else:\n query = attn.to_q(hidden_states.to(attn.to_q.weight.dtype))\n\n if encoder_hidden_states is None:\n crossattn = False\n encoder_hidden_states = hidden_states\n else:\n crossattn = True\n if attn.norm_cross:\n encoder_hidden_states = attn.norm_encoder_hidden_states(encoder_hidden_states)\n\n if self.train_kv:\n key = self.to_k_custom_diffusion(encoder_hidden_states.to(self.to_k_custom_diffusion.weight.dtype))\n value = self.to_v_custom_diffusion(encoder_hidden_states.to(self.to_v_custom_diffusion.weight.dtype))\n key = key.to(attn.to_q.weight.dtype)\n value = value.to(attn.to_q.weight.dtype)\n else:\n key = attn.to_k(encoder_hidden_states)\n value = attn.to_v(encoder_hidden_states)\n\n if crossattn:\n detach = torch.ones_like(key)\n detach[:, :1, :] = detach[:, :1, :] * 0.0\n key = detach * key + (1 - detach) * key.detach()\n value = detach * value + (1 - detach) * value.detach()\n\n query = attn.head_to_batch_dim(query).contiguous()\n key = attn.head_to_batch_dim(key).contiguous()\n value = attn.head_to_batch_dim(value).contiguous()\n\n hidden_states = xformers.ops.memory_efficient_attention(\n query, key, value, attn_bias=attention_mask, op=self.attention_op, scale=attn.scale\n )\n hidden_states = hidden_states.to(query.dtype)\n hidden_states = attn.batch_to_head_dim(hidden_states)\n\n if self.train_q_out:\n # linear proj\n hidden_states = self.to_out_custom_diffusion[0](hidden_states)\n # dropout\n hidden_states = self.to_out_custom_diffusion[1](hidden_states)\n else:\n # linear proj\n hidden_states = attn.to_out[0](hidden_states)\n # dropout\n hidden_states = attn.to_out[1](hidden_states)\n\n return hidden_states\n\n\nclass CustomDiffusionAttnProcessor2_0(nn.Module):\n r\"\"\"\n Processor for implementing attention for the Custom Diffusion method using PyTorch 2.0’s memory-efficient scaled\n dot-product attention.\n\n Args:\n train_kv (`bool`, defaults to `True`):\n Whether to newly train the key and value matrices corresponding to the text features.\n train_q_out (`bool`, defaults to `True`):\n Whether to newly train query matrices corresponding to the latent image features.\n hidden_size (`int`, *optional*, defaults to `None`):\n The hidden size of the attention layer.\n cross_attention_dim (`int`, *optional*, defaults to `None`):\n The number of channels in the `encoder_hidden_states`.\n out_bias (`bool`, defaults to `True`):\n Whether to include the bias parameter in `train_q_out`.\n dropout (`float`, *optional*, defaults to 0.0):\n The dropout probability to use.\n \"\"\"\n\n def __init__(\n self,\n train_kv: bool = True,\n train_q_out: bool = True,\n hidden_size: Optional[int] = None,\n cross_attention_dim: Optional[int] = None,\n out_bias: bool = True,\n dropout: float = 0.0,\n ):\n super().__init__()\n self.train_kv = train_kv\n self.train_q_out = train_q_out\n\n self.hidden_size = hidden_size\n self.cross_attention_dim = cross_attention_dim\n\n # `_custom_diffusion` id for easy serialization and loading.\n if self.train_kv:\n self.to_k_custom_diffusion = nn.Linear(cross_attention_dim or hidden_size, hidden_size, bias=False)\n self.to_v_custom_diffusion = nn.Linear(cross_attention_dim or hidden_size, hidden_size, bias=False)\n if self.train_q_out:\n self.to_q_custom_diffusion = nn.Linear(hidden_size, hidden_size, bias=False)\n self.to_out_custom_diffusion = nn.ModuleList([])\n self.to_out_custom_diffusion.append(nn.Linear(hidden_size, hidden_size, bias=out_bias))\n self.to_out_custom_diffusion.append(nn.Dropout(dropout))\n\n def __call__(\n self,\n attn: Attention,\n hidden_states: torch.Tensor,\n encoder_hidden_states: Optional[torch.Tensor] = None,\n attention_mask: Optional[torch.Tensor] = None,\n ) -> torch.Tensor:\n batch_size, sequence_length, _ = hidden_states.shape\n attention_mask = attn.prepare_attention_mask(attention_mask, sequence_length, batch_size)\n if self.train_q_out:\n query = self.to_q_custom_diffusion(hidden_states)\n else:\n query = attn.to_q(hidden_states)\n\n if encoder_hidden_states is None:\n crossattn = False\n encoder_hidden_states = hidden_states\n else:\n crossattn = True\n if attn.norm_cross:\n encoder_hidden_states = attn.norm_encoder_hidden_states(encoder_hidden_states)\n\n if self.train_kv:\n key = self.to_k_custom_diffusion(encoder_hidden_states.to(self.to_k_custom_diffusion.weight.dtype))\n value = self.to_v_custom_diffusion(encoder_hidden_states.to(self.to_v_custom_diffusion.weight.dtype))\n key = key.to(attn.to_q.weight.dtype)\n value = value.to(attn.to_q.weight.dtype)\n\n else:\n key = attn.to_k(encoder_hidden_states)\n value = attn.to_v(encoder_hidden_states)\n\n if crossattn:\n detach = torch.ones_like(key)\n detach[:, :1, :] = detach[:, :1, :] * 0.0\n key = detach * key + (1 - detach) * key.detach()\n value = detach * value + (1 - detach) * value.detach()\n\n inner_dim = hidden_states.shape[-1]\n\n head_dim = inner_dim // attn.heads\n query = query.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2)\n key = key.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2)\n value = value.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2)\n\n # the output of sdp = (batch, num_heads, seq_len, head_dim)\n # TODO: add support for attn.scale when we move to Torch 2.1\n hidden_states = F.scaled_dot_product_attention(\n query, key, value, attn_mask=attention_mask, dropout_p=0.0, is_causal=False\n )\n\n hidden_states = hidden_states.transpose(1, 2).reshape(batch_size, -1, attn.heads * head_dim)\n hidden_states = hidden_states.to(query.dtype)\n\n if self.train_q_out:\n # linear proj\n hidden_states = self.to_out_custom_diffusion[0](hidden_states)\n # dropout\n hidden_states = self.to_out_custom_diffusion[1](hidden_states)\n else:\n # linear proj\n hidden_states = attn.to_out[0](hidden_states)\n # dropout\n hidden_states = attn.to_out[1](hidden_states)\n\n return hidden_states\n\n\nclass SlicedAttnProcessor:\n r\"\"\"\n Processor for implementing sliced attention.\n\n Args:\n slice_size (`int`, *optional*):\n The number of steps to compute attention. Uses as many slices as `attention_head_dim // slice_size`, and\n `attention_head_dim` must be a multiple of the `slice_size`.\n \"\"\"\n\n def __init__(self, slice_size: int):\n self.slice_size = slice_size\n\n def __call__(\n self,\n attn: Attention,\n hidden_states: torch.Tensor,\n encoder_hidden_states: Optional[torch.Tensor] = None,\n attention_mask: Optional[torch.Tensor] = None,\n ) -> torch.Tensor:\n residual = hidden_states\n\n input_ndim = hidden_states.ndim\n\n if input_ndim == 4:\n batch_size, channel, height, width = hidden_states.shape\n hidden_states = hidden_states.view(batch_size, channel, height * width).transpose(1, 2)\n\n batch_size, sequence_length, _ = (\n hidden_states.shape if encoder_hidden_states is None else encoder_hidden_states.shape\n )\n attention_mask = attn.prepare_attention_mask(attention_mask, sequence_length, batch_size)\n\n if attn.group_norm is not None:\n hidden_states = attn.group_norm(hidden_states.transpose(1, 2)).transpose(1, 2)\n\n query = attn.to_q(hidden_states)\n dim = query.shape[-1]\n query = attn.head_to_batch_dim(query)\n\n if encoder_hidden_states is None:\n encoder_hidden_states = hidden_states\n elif attn.norm_cross:\n encoder_hidden_states = attn.norm_encoder_hidden_states(encoder_hidden_states)\n\n key = attn.to_k(encoder_hidden_states)\n value = attn.to_v(encoder_hidden_states)\n key = attn.head_to_batch_dim(key)\n value = attn.head_to_batch_dim(value)\n\n batch_size_attention, query_tokens, _ = query.shape\n hidden_states = torch.zeros(\n (batch_size_attention, query_tokens, dim // attn.heads), device=query.device, dtype=query.dtype\n )\n\n for i in range((batch_size_attention - 1) // self.slice_size + 1):\n start_idx = i * self.slice_size\n end_idx = (i + 1) * self.slice_size\n\n query_slice = query[start_idx:end_idx]\n key_slice = key[start_idx:end_idx]\n attn_mask_slice = attention_mask[start_idx:end_idx] if attention_mask is not None else None\n\n attn_slice = attn.get_attention_scores(query_slice, key_slice, attn_mask_slice)\n\n attn_slice = torch.bmm(attn_slice, value[start_idx:end_idx])\n\n hidden_states[start_idx:end_idx] = attn_slice\n\n hidden_states = attn.batch_to_head_dim(hidden_states)\n\n # linear proj\n hidden_states = attn.to_out[0](hidden_states)\n # dropout\n hidden_states = attn.to_out[1](hidden_states)\n\n if input_ndim == 4:\n hidden_states = hidden_states.transpose(-1, -2).reshape(batch_size, channel, height, width)\n\n if attn.residual_connection:\n hidden_states = hidden_states + residual\n\n hidden_states = hidden_states / attn.rescale_output_factor\n\n return hidden_states\n\n\nclass SlicedAttnAddedKVProcessor:\n r\"\"\"\n Processor for implementing sliced attention with extra learnable key and value matrices for the text encoder.\n\n Args:\n slice_size (`int`, *optional*):\n The number of steps to compute attention. Uses as many slices as `attention_head_dim // slice_size`, and\n `attention_head_dim` must be a multiple of the `slice_size`.\n \"\"\"\n\n def __init__(self, slice_size):\n self.slice_size = slice_size\n\n def __call__(\n self,\n attn: \"Attention\",\n hidden_states: torch.Tensor,\n encoder_hidden_states: Optional[torch.Tensor] = None,\n attention_mask: Optional[torch.Tensor] = None,\n temb: Optional[torch.Tensor] = None,\n ) -> torch.Tensor:\n residual = hidden_states\n\n if attn.spatial_norm is not None:\n hidden_states = attn.spatial_norm(hidden_states, temb)\n\n hidden_states = hidden_states.view(hidden_states.shape[0], hidden_states.shape[1], -1).transpose(1, 2)\n\n batch_size, sequence_length, _ = hidden_states.shape\n\n attention_mask = attn.prepare_attention_mask(attention_mask, sequence_length, batch_size)\n\n if encoder_hidden_states is None:\n encoder_hidden_states = hidden_states\n elif attn.norm_cross:\n encoder_hidden_states = attn.norm_encoder_hidden_states(encoder_hidden_states)\n\n hidden_states = attn.group_norm(hidden_states.transpose(1, 2)).transpose(1, 2)\n\n query = attn.to_q(hidden_states)\n dim = query.shape[-1]\n query = attn.head_to_batch_dim(query)\n\n encoder_hidden_states_key_proj = attn.add_k_proj(encoder_hidden_states)\n encoder_hidden_states_value_proj = attn.add_v_proj(encoder_hidden_states)\n\n encoder_hidden_states_key_proj = attn.head_to_batch_dim(encoder_hidden_states_key_proj)\n encoder_hidden_states_value_proj = attn.head_to_batch_dim(encoder_hidden_states_value_proj)\n\n if not attn.only_cross_attention:\n key = attn.to_k(hidden_states)\n value = attn.to_v(hidden_states)\n key = attn.head_to_batch_dim(key)\n value = attn.head_to_batch_dim(value)\n key = torch.cat([encoder_hidden_states_key_proj, key], dim=1)\n value = torch.cat([encoder_hidden_states_value_proj, value], dim=1)\n else:\n key = encoder_hidden_states_key_proj\n value = encoder_hidden_states_value_proj\n\n batch_size_attention, query_tokens, _ = query.shape\n hidden_states = torch.zeros(\n (batch_size_attention, query_tokens, dim // attn.heads), device=query.device, dtype=query.dtype\n )\n\n for i in range((batch_size_attention - 1) // self.slice_size + 1):\n start_idx = i * self.slice_size\n end_idx = (i + 1) * self.slice_size\n\n query_slice = query[start_idx:end_idx]\n key_slice = key[start_idx:end_idx]\n attn_mask_slice = attention_mask[start_idx:end_idx] if attention_mask is not None else None\n\n attn_slice = attn.get_attention_scores(query_slice, key_slice, attn_mask_slice)\n\n attn_slice = torch.bmm(attn_slice, value[start_idx:end_idx])\n\n hidden_states[start_idx:end_idx] = attn_slice\n\n hidden_states = attn.batch_to_head_dim(hidden_states)\n\n # linear proj\n hidden_states = attn.to_out[0](hidden_states)\n # dropout\n hidden_states = attn.to_out[1](hidden_states)\n\n hidden_states = hidden_states.transpose(-1, -2).reshape(residual.shape)\n hidden_states = hidden_states + residual\n\n return hidden_states\n\n\nclass SpatialNorm(nn.Module):\n \"\"\"\n Spatially conditioned normalization as defined in https://arxiv.org/abs/2209.09002.\n\n Args:\n f_channels (`int`):\n The number of channels for input to group normalization layer, and output of the spatial norm layer.\n zq_channels (`int`):\n The number of channels for the quantized vector as described in the paper.\n \"\"\"\n\n def __init__(\n self,\n f_channels: int,\n zq_channels: int,\n ):\n super().__init__()\n self.norm_layer = nn.GroupNorm(num_channels=f_channels, num_groups=32, eps=1e-6, affine=True)\n self.conv_y = nn.Conv2d(zq_channels, f_channels, kernel_size=1, stride=1, padding=0)\n self.conv_b = nn.Conv2d(zq_channels, f_channels, kernel_size=1, stride=1, padding=0)\n\n def forward(self, f: torch.Tensor, zq: torch.Tensor) -> torch.Tensor:\n f_size = f.shape[-2:]\n zq = F.interpolate(zq, size=f_size, mode=\"nearest\")\n norm_f = self.norm_layer(f)\n new_f = norm_f * self.conv_y(zq) + self.conv_b(zq)\n return new_f\n\n\nclass IPAdapterAttnProcessor(nn.Module):\n r\"\"\"\n Attention processor for Multiple IP-Adapters.\n\n Args:\n hidden_size (`int`):\n The hidden size of the attention layer.\n cross_attention_dim (`int`):\n The number of channels in the `encoder_hidden_states`.\n num_tokens (`int`, `Tuple[int]` or `List[int]`, defaults to `(4,)`):\n The context length of the image features.\n scale (`float` or List[`float`], defaults to 1.0):\n the weight scale of image prompt.\n \"\"\"\n\n def __init__(self, hidden_size, cross_attention_dim=None, num_tokens=(4,), scale=1.0):\n super().__init__()\n\n self.hidden_size = hidden_size\n self.cross_attention_dim = cross_attention_dim\n\n if not isinstance(num_tokens, (tuple, list)):\n num_tokens = [num_tokens]\n self.num_tokens = num_tokens\n\n if not isinstance(scale, list):\n scale = [scale] * len(num_tokens)\n if len(scale) != len(num_tokens):\n raise ValueError(\"`scale` should be a list of integers with the same length as `num_tokens`.\")\n self.scale = scale\n\n self.to_k_ip = nn.ModuleList(\n [nn.Linear(cross_attention_dim, hidden_size, bias=False) for _ in range(len(num_tokens))]\n )\n self.to_v_ip = nn.ModuleList(\n [nn.Linear(cross_attention_dim, hidden_size, bias=False) for _ in range(len(num_tokens))]\n )\n\n def __call__(\n self,\n attn: Attention,\n hidden_states: torch.Tensor,\n encoder_hidden_states: Optional[torch.Tensor] = None,\n attention_mask: Optional[torch.Tensor] = None,\n temb: Optional[torch.Tensor] = None,\n scale: float = 1.0,\n ip_adapter_masks: Optional[torch.Tensor] = None,\n ):\n residual = hidden_states\n\n # separate ip_hidden_states from encoder_hidden_states\n if encoder_hidden_states is not None:\n if isinstance(encoder_hidden_states, tuple):\n encoder_hidden_states, ip_hidden_states = encoder_hidden_states\n else:\n deprecation_message = (\n \"You have passed a tensor as `encoder_hidden_states`. This is deprecated and will be removed in a future release.\"\n \" Please make sure to update your script to pass `encoder_hidden_states` as a tuple to suppress this warning.\"\n )\n deprecate(\"encoder_hidden_states not a tuple\", \"1.0.0\", deprecation_message, standard_warn=False)\n end_pos = encoder_hidden_states.shape[1] - self.num_tokens[0]\n encoder_hidden_states, ip_hidden_states = (\n encoder_hidden_states[:, :end_pos, :],\n [encoder_hidden_states[:, end_pos:, :]],\n )\n\n if attn.spatial_norm is not None:\n hidden_states = attn.spatial_norm(hidden_states, temb)\n\n input_ndim = hidden_states.ndim\n\n if input_ndim == 4:\n batch_size, channel, height, width = hidden_states.shape\n hidden_states = hidden_states.view(batch_size, channel, height * width).transpose(1, 2)\n\n batch_size, sequence_length, _ = (\n hidden_states.shape if encoder_hidden_states is None else encoder_hidden_states.shape\n )\n attention_mask = attn.prepare_attention_mask(attention_mask, sequence_length, batch_size)\n\n if attn.group_norm is not None:\n hidden_states = attn.group_norm(hidden_states.transpose(1, 2)).transpose(1, 2)\n\n query = attn.to_q(hidden_states)\n\n if encoder_hidden_states is None:\n encoder_hidden_states = hidden_states\n elif attn.norm_cross:\n encoder_hidden_states = attn.norm_encoder_hidden_states(encoder_hidden_states)\n\n key = attn.to_k(encoder_hidden_states)\n value = attn.to_v(encoder_hidden_states)\n\n query = attn.head_to_batch_dim(query)\n key = attn.head_to_batch_dim(key)\n value = attn.head_to_batch_dim(value)\n\n attention_probs = attn.get_attention_scores(query, key, attention_mask)\n hidden_states = torch.bmm(attention_probs, value)\n hidden_states = attn.batch_to_head_dim(hidden_states)\n\n if ip_adapter_masks is not None:\n if not isinstance(ip_adapter_masks, List):\n # for backward compatibility, we accept `ip_adapter_mask` as a tensor of shape [num_ip_adapter, 1, height, width]\n ip_adapter_masks = list(ip_adapter_masks.unsqueeze(1))\n if not (len(ip_adapter_masks) == len(self.scale) == len(ip_hidden_states)):\n raise ValueError(\n f\"Length of ip_adapter_masks array ({len(ip_adapter_masks)}) must match \"\n f\"length of self.scale array ({len(self.scale)}) and number of ip_hidden_states \"\n f\"({len(ip_hidden_states)})\"\n )\n else:\n for index, (mask, scale, ip_state) in enumerate(zip(ip_adapter_masks, self.scale, ip_hidden_states)):\n if not isinstance(mask, torch.Tensor) or mask.ndim != 4:\n raise ValueError(\n \"Each element of the ip_adapter_masks array should be a tensor with shape \"\n \"[1, num_images_for_ip_adapter, height, width].\"\n \" Please use `IPAdapterMaskProcessor` to preprocess your mask\"\n )\n if mask.shape[1] != ip_state.shape[1]:\n raise ValueError(\n f\"Number of masks ({mask.shape[1]}) does not match \"\n f\"number of ip images ({ip_state.shape[1]}) at index {index}\"\n )\n if isinstance(scale, list) and not len(scale) == mask.shape[1]:\n raise ValueError(\n f\"Number of masks ({mask.shape[1]}) does not match \"\n f\"number of scales ({len(scale)}) at index {index}\"\n )\n else:\n ip_adapter_masks = [None] * len(self.scale)\n\n # for ip-adapter\n for current_ip_hidden_states, scale, to_k_ip, to_v_ip, mask in zip(\n ip_hidden_states, self.scale, self.to_k_ip, self.to_v_ip, ip_adapter_masks\n ):\n skip = False\n if isinstance(scale, list):\n if all(s == 0 for s in scale):\n skip = True\n elif scale == 0:\n skip = True\n if not skip:\n if mask is not None:\n if not isinstance(scale, list):\n scale = [scale] * mask.shape[1]\n\n current_num_images = mask.shape[1]\n for i in range(current_num_images):\n ip_key = to_k_ip(current_ip_hidden_states[:, i, :, :])\n ip_value = to_v_ip(current_ip_hidden_states[:, i, :, :])\n\n ip_key = attn.head_to_batch_dim(ip_key)\n ip_value = attn.head_to_batch_dim(ip_value)\n\n ip_attention_probs = attn.get_attention_scores(query, ip_key, None)\n _current_ip_hidden_states = torch.bmm(ip_attention_probs, ip_value)\n _current_ip_hidden_states = attn.batch_to_head_dim(_current_ip_hidden_states)\n\n mask_downsample = IPAdapterMaskProcessor.downsample(\n mask[:, i, :, :],\n batch_size,\n _current_ip_hidden_states.shape[1],\n _current_ip_hidden_states.shape[2],\n )\n\n mask_downsample = mask_downsample.to(dtype=query.dtype, device=query.device)\n\n hidden_states = hidden_states + scale[i] * (_current_ip_hidden_states * mask_downsample)\n else:\n ip_key = to_k_ip(current_ip_hidden_states)\n ip_value = to_v_ip(current_ip_hidden_states)\n\n ip_key = attn.head_to_batch_dim(ip_key)\n ip_value = attn.head_to_batch_dim(ip_value)\n\n ip_attention_probs = attn.get_attention_scores(query, ip_key, None)\n current_ip_hidden_states = torch.bmm(ip_attention_probs, ip_value)\n current_ip_hidden_states = attn.batch_to_head_dim(current_ip_hidden_states)\n\n hidden_states = hidden_states + scale * current_ip_hidden_states\n\n # linear proj\n hidden_states = attn.to_out[0](hidden_states)\n # dropout\n hidden_states = attn.to_out[1](hidden_states)\n\n if input_ndim == 4:\n hidden_states = hidden_states.transpose(-1, -2).reshape(batch_size, channel, height, width)\n\n if attn.residual_connection:\n hidden_states = hidden_states + residual\n\n hidden_states = hidden_states / attn.rescale_output_factor\n\n return hidden_states\n\n\nclass IPAdapterAttnProcessor2_0(torch.nn.Module):\n r\"\"\"\n Attention processor for IP-Adapter for PyTorch 2.0.\n\n Args:\n hidden_size (`int`):\n The hidden size of the attention layer.\n cross_attention_dim (`int`):\n The number of channels in the `encoder_hidden_states`.\n num_tokens (`int`, `Tuple[int]` or `List[int]`, defaults to `(4,)`):\n The context length of the image features.\n scale (`float` or `List[float]`, defaults to 1.0):\n the weight scale of image prompt.\n \"\"\"\n\n def __init__(self, hidden_size, cross_attention_dim=None, num_tokens=(4,), scale=1.0):\n super().__init__()\n\n if not hasattr(F, \"scaled_dot_product_attention\"):\n raise ImportError(\n f\"{self.__class__.__name__} requires PyTorch 2.0, to use it, please upgrade PyTorch to 2.0.\"\n )\n\n self.hidden_size = hidden_size\n self.cross_attention_dim = cross_attention_dim\n\n if not isinstance(num_tokens, (tuple, list)):\n num_tokens = [num_tokens]\n self.num_tokens = num_tokens\n\n if not isinstance(scale, list):\n scale = [scale] * len(num_tokens)\n if len(scale) != len(num_tokens):\n raise ValueError(\"`scale` should be a list of integers with the same length as `num_tokens`.\")\n self.scale = scale\n\n self.to_k_ip = nn.ModuleList(\n [nn.Linear(cross_attention_dim, hidden_size, bias=False) for _ in range(len(num_tokens))]\n )\n self.to_v_ip = nn.ModuleList(\n [nn.Linear(cross_attention_dim, hidden_size, bias=False) for _ in range(len(num_tokens))]\n )\n\n def __call__(\n self,\n attn: Attention,\n hidden_states: torch.Tensor,\n encoder_hidden_states: Optional[torch.Tensor] = None,\n attention_mask: Optional[torch.Tensor] = None,\n temb: Optional[torch.Tensor] = None,\n scale: float = 1.0,\n ip_adapter_masks: Optional[torch.Tensor] = None,\n ):\n residual = hidden_states\n\n # separate ip_hidden_states from encoder_hidden_states\n if encoder_hidden_states is not None:\n if isinstance(encoder_hidden_states, tuple):\n encoder_hidden_states, ip_hidden_states = encoder_hidden_states\n else:\n deprecation_message = (\n \"You have passed a tensor as `encoder_hidden_states`. This is deprecated and will be removed in a future release.\"\n \" Please make sure to update your script to pass `encoder_hidden_states` as a tuple to suppress this warning.\"\n )\n deprecate(\"encoder_hidden_states not a tuple\", \"1.0.0\", deprecation_message, standard_warn=False)\n end_pos = encoder_hidden_states.shape[1] - self.num_tokens[0]\n encoder_hidden_states, ip_hidden_states = (\n encoder_hidden_states[:, :end_pos, :],\n [encoder_hidden_states[:, end_pos:, :]],\n )\n\n if attn.spatial_norm is not None:\n hidden_states = attn.spatial_norm(hidden_states, temb)\n\n input_ndim = hidden_states.ndim\n\n if input_ndim == 4:\n batch_size, channel, height, width = hidden_states.shape\n hidden_states = hidden_states.view(batch_size, channel, height * width).transpose(1, 2)\n\n batch_size, sequence_length, _ = (\n hidden_states.shape if encoder_hidden_states is None else encoder_hidden_states.shape\n )\n\n if attention_mask is not None:\n attention_mask = attn.prepare_attention_mask(attention_mask, sequence_length, batch_size)\n # scaled_dot_product_attention expects attention_mask shape to be\n # (batch, heads, source_length, target_length)\n attention_mask = attention_mask.view(batch_size, attn.heads, -1, attention_mask.shape[-1])\n\n if attn.group_norm is not None:\n hidden_states = attn.group_norm(hidden_states.transpose(1, 2)).transpose(1, 2)\n\n query = attn.to_q(hidden_states)\n\n if encoder_hidden_states is None:\n encoder_hidden_states = hidden_states\n elif attn.norm_cross:\n encoder_hidden_states = attn.norm_encoder_hidden_states(encoder_hidden_states)\n\n key = attn.to_k(encoder_hidden_states)\n value = attn.to_v(encoder_hidden_states)\n\n inner_dim = key.shape[-1]\n head_dim = inner_dim // attn.heads\n\n query = query.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2)\n\n key = key.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2)\n value = value.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2)\n\n # the output of sdp = (batch, num_heads, seq_len, head_dim)\n # TODO: add support for attn.scale when we move to Torch 2.1\n hidden_states = F.scaled_dot_product_attention(\n query, key, value, attn_mask=attention_mask, dropout_p=0.0, is_causal=False\n )\n\n hidden_states = hidden_states.transpose(1, 2).reshape(batch_size, -1, attn.heads * head_dim)\n hidden_states = hidden_states.to(query.dtype)\n\n if ip_adapter_masks is not None:\n if not isinstance(ip_adapter_masks, List):\n # for backward compatibility, we accept `ip_adapter_mask` as a tensor of shape [num_ip_adapter, 1, height, width]\n ip_adapter_masks = list(ip_adapter_masks.unsqueeze(1))\n if not (len(ip_adapter_masks) == len(self.scale) == len(ip_hidden_states)):\n raise ValueError(\n f\"Length of ip_adapter_masks array ({len(ip_adapter_masks)}) must match \"\n f\"length of self.scale array ({len(self.scale)}) and number of ip_hidden_states \"\n f\"({len(ip_hidden_states)})\"\n )\n else:\n for index, (mask, scale, ip_state) in enumerate(zip(ip_adapter_masks, self.scale, ip_hidden_states)):\n if not isinstance(mask, torch.Tensor) or mask.ndim != 4:\n raise ValueError(\n \"Each element of the ip_adapter_masks array should be a tensor with shape \"\n \"[1, num_images_for_ip_adapter, height, width].\"\n \" Please use `IPAdapterMaskProcessor` to preprocess your mask\"\n )\n if mask.shape[1] != ip_state.shape[1]:\n raise ValueError(\n f\"Number of masks ({mask.shape[1]}) does not match \"\n f\"number of ip images ({ip_state.shape[1]}) at index {index}\"\n )\n if isinstance(scale, list) and not len(scale) == mask.shape[1]:\n raise ValueError(\n f\"Number of masks ({mask.shape[1]}) does not match \"\n f\"number of scales ({len(scale)}) at index {index}\"\n )\n else:\n ip_adapter_masks = [None] * len(self.scale)\n\n # for ip-adapter\n for current_ip_hidden_states, scale, to_k_ip, to_v_ip, mask in zip(\n ip_hidden_states, self.scale, self.to_k_ip, self.to_v_ip, ip_adapter_masks\n ):\n skip = False\n if isinstance(scale, list):\n if all(s == 0 for s in scale):\n skip = True\n elif scale == 0:\n skip = True\n if not skip:\n if mask is not None:\n if not isinstance(scale, list):\n scale = [scale] * mask.shape[1]\n\n current_num_images = mask.shape[1]\n for i in range(current_num_images):\n ip_key = to_k_ip(current_ip_hidden_states[:, i, :, :])\n ip_value = to_v_ip(current_ip_hidden_states[:, i, :, :])\n\n ip_key = ip_key.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2)\n ip_value = ip_value.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2)\n\n # the output of sdp = (batch, num_heads, seq_len, head_dim)\n # TODO: add support for attn.scale when we move to Torch 2.1\n _current_ip_hidden_states = F.scaled_dot_product_attention(\n query, ip_key, ip_value, attn_mask=None, dropout_p=0.0, is_causal=False\n )\n\n _current_ip_hidden_states = _current_ip_hidden_states.transpose(1, 2).reshape(\n batch_size, -1, attn.heads * head_dim\n )\n _current_ip_hidden_states = _current_ip_hidden_states.to(query.dtype)\n\n mask_downsample = IPAdapterMaskProcessor.downsample(\n mask[:, i, :, :],\n batch_size,\n _current_ip_hidden_states.shape[1],\n _current_ip_hidden_states.shape[2],\n )\n\n mask_downsample = mask_downsample.to(dtype=query.dtype, device=query.device)\n hidden_states = hidden_states + scale[i] * (_current_ip_hidden_states * mask_downsample)\n else:\n ip_key = to_k_ip(current_ip_hidden_states)\n ip_value = to_v_ip(current_ip_hidden_states)\n\n ip_key = ip_key.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2)\n ip_value = ip_value.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2)\n\n # the output of sdp = (batch, num_heads, seq_len, head_dim)\n # TODO: add support for attn.scale when we move to Torch 2.1\n current_ip_hidden_states = F.scaled_dot_product_attention(\n query, ip_key, ip_value, attn_mask=None, dropout_p=0.0, is_causal=False\n )\n\n current_ip_hidden_states = current_ip_hidden_states.transpose(1, 2).reshape(\n batch_size, -1, attn.heads * head_dim\n )\n current_ip_hidden_states = current_ip_hidden_states.to(query.dtype)\n\n hidden_states = hidden_states + scale * current_ip_hidden_states\n\n # linear proj\n hidden_states = attn.to_out[0](hidden_states)\n # dropout\n hidden_states = attn.to_out[1](hidden_states)\n\n if input_ndim == 4:\n hidden_states = hidden_states.transpose(-1, -2).reshape(batch_size, channel, height, width)\n\n if attn.residual_connection:\n hidden_states = hidden_states + residual\n\n hidden_states = hidden_states / attn.rescale_output_factor\n\n return hidden_states\n\n\nclass PAGIdentitySelfAttnProcessor2_0:\n r\"\"\"\n Processor for implementing PAG using scaled dot-product attention (enabled by default if you're using PyTorch 2.0).\n PAG reference: https://arxiv.org/abs/2403.17377\n \"\"\"\n\n def __init__(self):\n if not hasattr(F, \"scaled_dot_product_attention\"):\n raise ImportError(\n \"PAGIdentitySelfAttnProcessor2_0 requires PyTorch 2.0, to use it, please upgrade PyTorch to 2.0.\"\n )\n\n def __call__(\n self,\n attn: Attention,\n hidden_states: torch.FloatTensor,\n encoder_hidden_states: Optional[torch.FloatTensor] = None,\n attention_mask: Optional[torch.FloatTensor] = None,\n temb: Optional[torch.FloatTensor] = None,\n ) -> torch.Tensor:\n residual = hidden_states\n if attn.spatial_norm is not None:\n hidden_states = attn.spatial_norm(hidden_states, temb)\n\n input_ndim = hidden_states.ndim\n if input_ndim == 4:\n batch_size, channel, height, width = hidden_states.shape\n hidden_states = hidden_states.view(batch_size, channel, height * width).transpose(1, 2)\n\n # chunk\n hidden_states_org, hidden_states_ptb = hidden_states.chunk(2)\n\n # original path\n batch_size, sequence_length, _ = hidden_states_org.shape\n\n if attention_mask is not None:\n attention_mask = attn.prepare_attention_mask(attention_mask, sequence_length, batch_size)\n # scaled_dot_product_attention expects attention_mask shape to be\n # (batch, heads, source_length, target_length)\n attention_mask = attention_mask.view(batch_size, attn.heads, -1, attention_mask.shape[-1])\n\n if attn.group_norm is not None:\n hidden_states_org = attn.group_norm(hidden_states_org.transpose(1, 2)).transpose(1, 2)\n\n query = attn.to_q(hidden_states_org)\n key = attn.to_k(hidden_states_org)\n value = attn.to_v(hidden_states_org)\n\n inner_dim = key.shape[-1]\n head_dim = inner_dim // attn.heads\n\n query = query.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2)\n key = key.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2)\n value = value.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2)\n\n # the output of sdp = (batch, num_heads, seq_len, head_dim)\n # TODO: add support for attn.scale when we move to Torch 2.1\n hidden_states_org = F.scaled_dot_product_attention(\n query, key, value, attn_mask=attention_mask, dropout_p=0.0, is_causal=False\n )\n hidden_states_org = hidden_states_org.transpose(1, 2).reshape(batch_size, -1, attn.heads * head_dim)\n hidden_states_org = hidden_states_org.to(query.dtype)\n\n # linear proj\n hidden_states_org = attn.to_out[0](hidden_states_org)\n # dropout\n hidden_states_org = attn.to_out[1](hidden_states_org)\n\n if input_ndim == 4:\n hidden_states_org = hidden_states_org.transpose(-1, -2).reshape(batch_size, channel, height, width)\n\n # perturbed path (identity attention)\n batch_size, sequence_length, _ = hidden_states_ptb.shape\n\n if attn.group_norm is not None:\n hidden_states_ptb = attn.group_norm(hidden_states_ptb.transpose(1, 2)).transpose(1, 2)\n\n hidden_states_ptb = attn.to_v(hidden_states_ptb)\n hidden_states_ptb = hidden_states_ptb.to(query.dtype)\n\n # linear proj\n hidden_states_ptb = attn.to_out[0](hidden_states_ptb)\n # dropout\n hidden_states_ptb = attn.to_out[1](hidden_states_ptb)\n\n if input_ndim == 4:\n hidden_states_ptb = hidden_states_ptb.transpose(-1, -2).reshape(batch_size, channel, height, width)\n\n # cat\n hidden_states = torch.cat([hidden_states_org, hidden_states_ptb])\n\n if attn.residual_connection:\n hidden_states = hidden_states + residual\n\n hidden_states = hidden_states / attn.rescale_output_factor\n\n return hidden_states\n\n\nclass PAGCFGIdentitySelfAttnProcessor2_0:\n r\"\"\"\n Processor for implementing PAG using scaled dot-product attention (enabled by default if you're using PyTorch 2.0).\n PAG reference: https://arxiv.org/abs/2403.17377\n \"\"\"\n\n def __init__(self):\n if not hasattr(F, \"scaled_dot_product_attention\"):\n raise ImportError(\n \"PAGCFGIdentitySelfAttnProcessor2_0 requires PyTorch 2.0, to use it, please upgrade PyTorch to 2.0.\"\n )\n\n def __call__(\n self,\n attn: Attention,\n hidden_states: torch.FloatTensor,\n encoder_hidden_states: Optional[torch.FloatTensor] = None,\n attention_mask: Optional[torch.FloatTensor] = None,\n temb: Optional[torch.FloatTensor] = None,\n ) -> torch.Tensor:\n residual = hidden_states\n if attn.spatial_norm is not None:\n hidden_states = attn.spatial_norm(hidden_states, temb)\n\n input_ndim = hidden_states.ndim\n if input_ndim == 4:\n batch_size, channel, height, width = hidden_states.shape\n hidden_states = hidden_states.view(batch_size, channel, height * width).transpose(1, 2)\n\n # chunk\n hidden_states_uncond, hidden_states_org, hidden_states_ptb = hidden_states.chunk(3)\n hidden_states_org = torch.cat([hidden_states_uncond, hidden_states_org])\n\n # original path\n batch_size, sequence_length, _ = hidden_states_org.shape\n\n if attention_mask is not None:\n attention_mask = attn.prepare_attention_mask(attention_mask, sequence_length, batch_size)\n # scaled_dot_product_attention expects attention_mask shape to be\n # (batch, heads, source_length, target_length)\n attention_mask = attention_mask.view(batch_size, attn.heads, -1, attention_mask.shape[-1])\n\n if attn.group_norm is not None:\n hidden_states_org = attn.group_norm(hidden_states_org.transpose(1, 2)).transpose(1, 2)\n\n query = attn.to_q(hidden_states_org)\n key = attn.to_k(hidden_states_org)\n value = attn.to_v(hidden_states_org)\n\n inner_dim = key.shape[-1]\n head_dim = inner_dim // attn.heads\n\n query = query.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2)\n\n key = key.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2)\n value = value.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2)\n\n # the output of sdp = (batch, num_heads, seq_len, head_dim)\n # TODO: add support for attn.scale when we move to Torch 2.1\n hidden_states_org = F.scaled_dot_product_attention(\n query, key, value, attn_mask=attention_mask, dropout_p=0.0, is_causal=False\n )\n\n hidden_states_org = hidden_states_org.transpose(1, 2).reshape(batch_size, -1, attn.heads * head_dim)\n hidden_states_org = hidden_states_org.to(query.dtype)\n\n # linear proj\n hidden_states_org = attn.to_out[0](hidden_states_org)\n # dropout\n hidden_states_org = attn.to_out[1](hidden_states_org)\n\n if input_ndim == 4:\n hidden_states_org = hidden_states_org.transpose(-1, -2).reshape(batch_size, channel, height, width)\n\n # perturbed path (identity attention)\n batch_size, sequence_length, _ = hidden_states_ptb.shape\n\n if attn.group_norm is not None:\n hidden_states_ptb = attn.group_norm(hidden_states_ptb.transpose(1, 2)).transpose(1, 2)\n\n value = attn.to_v(hidden_states_ptb)\n hidden_states_ptb = value\n hidden_states_ptb = hidden_states_ptb.to(query.dtype)\n\n # linear proj\n hidden_states_ptb = attn.to_out[0](hidden_states_ptb)\n # dropout\n hidden_states_ptb = attn.to_out[1](hidden_states_ptb)\n\n if input_ndim == 4:\n hidden_states_ptb = hidden_states_ptb.transpose(-1, -2).reshape(batch_size, channel, height, width)\n\n # cat\n hidden_states = torch.cat([hidden_states_org, hidden_states_ptb])\n\n if attn.residual_connection:\n hidden_states = hidden_states + residual\n\n hidden_states = hidden_states / attn.rescale_output_factor\n\n return hidden_states\n\n\nclass LoRAAttnProcessor:\n def __init__(self):\n pass\n\n\nclass LoRAAttnProcessor2_0:\n def __init__(self):\n pass\n\n\nclass LoRAXFormersAttnProcessor:\n def __init__(self):\n pass\n\n\nclass LoRAAttnAddedKVProcessor:\n def __init__(self):\n pass\n\n\nclass FluxSingleAttnProcessor2_0(FluxAttnProcessor2_0):\n r\"\"\"\n Processor for implementing scaled dot-product attention (enabled by default if you're using PyTorch 2.0).\n \"\"\"\n\n def __init__(self):\n deprecation_message = \"`FluxSingleAttnProcessor2_0` is deprecated and will be removed in a future version. Please use `FluxAttnProcessor2_0` instead.\"\n deprecate(\"FluxSingleAttnProcessor2_0\", \"0.32.0\", deprecation_message)\n super().__init__()\n\n\nADDED_KV_ATTENTION_PROCESSORS = (\n AttnAddedKVProcessor,\n SlicedAttnAddedKVProcessor,\n AttnAddedKVProcessor2_0,\n XFormersAttnAddedKVProcessor,\n)\n\nCROSS_ATTENTION_PROCESSORS = (\n AttnProcessor,\n AttnProcessor2_0,\n XFormersAttnProcessor,\n SlicedAttnProcessor,\n IPAdapterAttnProcessor,\n IPAdapterAttnProcessor2_0,\n)\n\nAttentionProcessor = Union[\n AttnProcessor,\n AttnProcessor2_0,\n FusedAttnProcessor2_0,\n XFormersAttnProcessor,\n SlicedAttnProcessor,\n AttnAddedKVProcessor,\n SlicedAttnAddedKVProcessor,\n AttnAddedKVProcessor2_0,\n XFormersAttnAddedKVProcessor,\n CustomDiffusionAttnProcessor,\n CustomDiffusionXFormersAttnProcessor,\n CustomDiffusionAttnProcessor2_0,\n PAGCFGIdentitySelfAttnProcessor2_0,\n PAGIdentitySelfAttnProcessor2_0,\n PAGCFGHunyuanAttnProcessor2_0,\n PAGHunyuanAttnProcessor2_0,\n]\n" }, { "path": "CogVideo/finetune/models/cogvideox_transformer_3d.py", "content": "# Copyright 2024 The CogVideoX team, Tsinghua University & ZhipuAI and The HuggingFace Team.\n# All rights reserved.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n# http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\nfrom typing import Any, Dict, Optional, Tuple, Union\nfrom einops import rearrange, repeat\n\nimport torch\nfrom torch import nn\n\nfrom diffusers.configuration_utils import ConfigMixin, register_to_config\nfrom diffusers.loaders import PeftAdapterMixin\nfrom diffusers.utils import USE_PEFT_BACKEND, is_torch_version, logging, scale_lora_layers, unscale_lora_layers\nfrom diffusers.utils.torch_utils import maybe_allow_in_graph\nfrom models.attention import Attention, FeedForward\nfrom models.attention_processor import AttentionProcessor, CogVideoXAttnProcessor2_0, FusedCogVideoXAttnProcessor2_0\nfrom models.embeddings import CogVideoXPatchEmbed, TimestepEmbedding, Timesteps\nfrom diffusers.models.modeling_outputs import Transformer2DModelOutput\nfrom diffusers.models.modeling_utils import ModelMixin\nfrom diffusers.models.normalization import AdaLayerNorm, CogVideoXLayerNormZero\n\nlogger = logging.get_logger(__name__) # pylint: disable=invalid-name\n\n@maybe_allow_in_graph\nclass CogVideoXBlock(nn.Module):\n r\"\"\"\n Transformer block used in [CogVideoX](https://github.com/THUDM/CogVideo) model.\n\n Parameters:\n dim (`int`):\n The number of channels in the input and output.\n num_attention_heads (`int`):\n The number of heads to use for multi-head attention.\n attention_head_dim (`int`):\n The number of channels in each head.\n time_embed_dim (`int`):\n The number of channels in timestep embedding.\n dropout (`float`, defaults to `0.0`):\n The dropout probability to use.\n activation_fn (`str`, defaults to `\"gelu-approximate\"`):\n Activation function to be used in feed-forward.\n attention_bias (`bool`, defaults to `False`):\n Whether or not to use bias in attention projection layers.\n qk_norm (`bool`, defaults to `True`):\n Whether or not to use normalization after query and key projections in Attention.\n norm_elementwise_affine (`bool`, defaults to `True`):\n Whether to use learnable elementwise affine parameters for normalization.\n norm_eps (`float`, defaults to `1e-5`):\n Epsilon value for normalization layers.\n final_dropout (`bool` defaults to `False`):\n Whether to apply a final dropout after the last feed-forward layer.\n ff_inner_dim (`int`, *optional*, defaults to `None`):\n Custom hidden dimension of Feed-forward layer. If not provided, `4 * dim` is used.\n ff_bias (`bool`, defaults to `True`):\n Whether or not to use bias in Feed-forward layer.\n attention_out_bias (`bool`, defaults to `True`):\n Whether or not to use bias in Attention output projection layer.\n \"\"\"\n \n def __init__(\n self,\n block_idx: int,\n dim: int,\n num_attention_heads: int,\n attention_head_dim: int,\n time_embed_dim: int,\n block_interval: int = 2,\n dropout: float = 0.0,\n activation_fn: str = \"gelu-approximate\",\n attention_bias: bool = False,\n qk_norm: bool = True,\n norm_elementwise_affine: bool = True,\n norm_eps: float = 1e-5,\n final_dropout: bool = True,\n ff_inner_dim: Optional[int] = None,\n ff_bias: bool = True,\n attention_out_bias: bool = True,\n finetune_init: bool = False,\n ):\n super().__init__()\n \n # 1. Self Attention\n self.norm1 = CogVideoXLayerNormZero(time_embed_dim, dim, norm_elementwise_affine, norm_eps, bias=True)\n \n self.attn1 = Attention(\n query_dim=dim,\n dim_head=attention_head_dim,\n heads=num_attention_heads,\n qk_norm=\"layer_norm\" if qk_norm else None,\n eps=1e-6,\n bias=attention_bias,\n out_bias=attention_out_bias,\n processor=CogVideoXAttnProcessor2_0(),\n )\n \n # if False: ## for finetuning stage w/o loading pretrained checkpoints\n if not finetune_init and (block_idx%block_interval==0):\n self.attn_injector = Attention(\n query_dim=dim,\n dim_head=attention_head_dim,\n heads=num_attention_heads,\n qk_norm=\"layer_norm\" if qk_norm else None,\n eps=1e-6,\n bias=attention_bias,\n out_bias=attention_out_bias,\n processor=CogVideoXAttnProcessor2_0(),\n )\n self.pose_fuse_layer = nn.Linear(12, dim)\n self.attn_null_feature = nn.Parameter(torch.zeros([dim]))\n\n # 2. Feed Forward\n self.norm2 = CogVideoXLayerNormZero(time_embed_dim, dim, norm_elementwise_affine, norm_eps, bias=True)\n\n self.ff = FeedForward(\n dim,\n dropout=dropout,\n activation_fn=activation_fn,\n final_dropout=final_dropout,\n inner_dim=ff_inner_dim,\n bias=ff_bias,\n )\n\n def forward(\n self,\n hidden_states: torch.Tensor,\n encoder_hidden_states: torch.Tensor,\n empty_encoder_hidden_states: torch.Tensor,\n temb: torch.Tensor,\n pose_embeds: Optional[torch.FloatTensor] = None,\n prompt_entities_embeds: Optional[torch.FloatTensor] = None,\n image_rotary_emb: Optional[Tuple[torch.Tensor, torch.Tensor]] = None,\n ) -> torch.Tensor:\n text_seq_length = encoder_hidden_states.size(1)\n \n # norm & modulate\n norm_hidden_states, norm_encoder_hidden_states, gate_msa, enc_gate_msa = self.norm1(\n hidden_states, encoder_hidden_states, temb\n )\n\n # attention\n attn_hidden_states, attn_encoder_hidden_states = self.attn1(\n hidden_states=norm_hidden_states,\n encoder_hidden_states=norm_encoder_hidden_states,\n image_rotary_emb=image_rotary_emb,\n )\n\n hidden_states = hidden_states + gate_msa * attn_hidden_states\n encoder_hidden_states = encoder_hidden_states + enc_gate_msa * attn_encoder_hidden_states\n\n if (pose_embeds is not None) and hasattr(self, \"attn_injector\") :\n\n # 1. norm & modulate\n norm_hidden_states, norm_empty_encoder_hidden_states, gate_msa, enc_gate_msa = self.norm1(hidden_states, empty_encoder_hidden_states, temb)\n bz, N_visual, dim = norm_hidden_states.shape\n max_entity_num = 3\n _, entity_num, num_frames, _ = pose_embeds.shape\n\n # 2. pair-wise fusion of trajectory and entity\n attn_input = self.attn_null_feature.repeat(bz, max_entity_num, 50, num_frames, 1)\n pose_embeds = self.pose_fuse_layer(pose_embeds)\n attn_input[:,:entity_num,:,:,:] = pose_embeds.unsqueeze(-3) + prompt_entities_embeds.unsqueeze(-2)\n attn_input = torch.cat((\n rearrange(norm_hidden_states, \"b (n t) d -> b n t d\",n=num_frames), \n rearrange(attn_input, \"b n t f d -> b f (n t) d\")),\n dim=2\n ).flatten(1,2)\n \n # 3. gated self-attention\n attn_hidden_states, attn_encoder_hidden_states = self.attn_injector(\n hidden_states=attn_input,\n encoder_hidden_states=norm_empty_encoder_hidden_states,\n image_rotary_emb=image_rotary_emb,\n )\n attn_hidden_states = attn_hidden_states[:,:N_visual,:]\n\n hidden_states = hidden_states + gate_msa * attn_hidden_states\n\n # norm & modulate\n norm_hidden_states, norm_encoder_hidden_states, gate_ff, enc_gate_ff = self.norm2(\n hidden_states, encoder_hidden_states, temb\n )\n\n # feed-forward\n norm_hidden_states = torch.cat([norm_encoder_hidden_states, norm_hidden_states], dim=1)\n ff_output = self.ff(norm_hidden_states)\n\n hidden_states = hidden_states + gate_ff * ff_output[:, text_seq_length:]\n encoder_hidden_states = encoder_hidden_states + enc_gate_ff * ff_output[:, :text_seq_length]\n\n return hidden_states, encoder_hidden_states\n\n\nclass CogVideoXTransformer3DModel(ModelMixin, ConfigMixin, PeftAdapterMixin):\n \"\"\"\n A Transformer model for video-like data in [CogVideoX](https://github.com/THUDM/CogVideo).\n\n Parameters:\n num_attention_heads (`int`, defaults to `30`):\n The number of heads to use for multi-head attention.\n attention_head_dim (`int`, defaults to `64`):\n The number of channels in each head.\n in_channels (`int`, defaults to `16`):\n The number of channels in the input.\n out_channels (`int`, *optional*, defaults to `16`):\n The number of channels in the output.\n flip_sin_to_cos (`bool`, defaults to `True`):\n Whether to flip the sin to cos in the time embedding.\n time_embed_dim (`int`, defaults to `512`):\n Output dimension of timestep embeddings.\n text_embed_dim (`int`, defaults to `4096`):\n Input dimension of text embeddings from the text encoder.\n num_layers (`int`, defaults to `30`):\n The number of layers of Transformer blocks to use.\n dropout (`float`, defaults to `0.0`):\n The dropout probability to use.\n attention_bias (`bool`, defaults to `True`):\n Whether or not to use bias in the attention projection layers.\n sample_width (`int`, defaults to `90`):\n The width of the input latents.\n sample_height (`int`, defaults to `60`):\n The height of the input latents.\n sample_frames (`int`, defaults to `49`):\n The number of frames in the input latents. Note that this parameter was incorrectly initialized to 49\n instead of 13 because CogVideoX processed 13 latent frames at once in its default and recommended settings,\n but cannot be changed to the correct value to ensure backwards compatibility. To create a transformer with\n K latent frames, the correct value to pass here would be: ((K - 1) * temporal_compression_ratio + 1).\n patch_size (`int`, defaults to `2`):\n The size of the patches to use in the patch embedding layer.\n temporal_compression_ratio (`int`, defaults to `4`):\n The compression ratio across the temporal dimension. See documentation for `sample_frames`.\n max_text_seq_length (`int`, defaults to `226`):\n The maximum sequence length of the input text embeddings.\n activation_fn (`str`, defaults to `\"gelu-approximate\"`):\n Activation function to use in feed-forward.\n timestep_activation_fn (`str`, defaults to `\"silu\"`):\n Activation function to use when generating the timestep embeddings.\n norm_elementwise_affine (`bool`, defaults to `True`):\n Whether or not to use elementwise affine in normalization layers.\n norm_eps (`float`, defaults to `1e-5`):\n The epsilon value to use in normalization layers.\n spatial_interpolation_scale (`float`, defaults to `1.875`):\n Scaling factor to apply in 3D positional embeddings across spatial dimensions.\n temporal_interpolation_scale (`float`, defaults to `1.0`):\n Scaling factor to apply in 3D positional embeddings across temporal dimensions.\n \"\"\"\n\n _supports_gradient_checkpointing = True\n\n @register_to_config\n def __init__(\n self,\n num_attention_heads: int = 30,\n attention_head_dim: int = 64,\n in_channels: int = 16,\n out_channels: Optional[int] = 16,\n flip_sin_to_cos: bool = True,\n freq_shift: int = 0,\n time_embed_dim: int = 512,\n text_embed_dim: int = 4096,\n num_layers: int = 30,\n dropout: float = 0.0,\n attention_bias: bool = True,\n sample_width: int = 90,\n sample_height: int = 60,\n sample_frames: int = 49,\n patch_size: int = 2,\n temporal_compression_ratio: int = 4,\n max_text_seq_length: int = 226,\n activation_fn: str = \"gelu-approximate\",\n timestep_activation_fn: str = \"silu\",\n norm_elementwise_affine: bool = True,\n norm_eps: float = 1e-5,\n spatial_interpolation_scale: float = 1.875,\n temporal_interpolation_scale: float = 1.0,\n use_rotary_positional_embeddings: bool = False,\n use_learned_positional_embeddings: bool = False,\n finetune_init: bool = False,\n ):\n super().__init__()\n inner_dim = num_attention_heads * attention_head_dim\n\n if not use_rotary_positional_embeddings and use_learned_positional_embeddings:\n raise ValueError(\n \"There are no CogVideoX checkpoints available with disable rotary embeddings and learned positional \"\n \"embeddings. If you're using a custom model and/or believe this should be supported, please open an \"\n \"issue at https://github.com/huggingface/diffusers/issues.\"\n )\n\n # 1. Patch embedding\n self.patch_embed = CogVideoXPatchEmbed(\n patch_size=patch_size,\n in_channels=in_channels,\n embed_dim=inner_dim,\n text_embed_dim=text_embed_dim,\n bias=True,\n sample_width=sample_width,\n sample_height=sample_height,\n sample_frames=sample_frames,\n temporal_compression_ratio=temporal_compression_ratio,\n max_text_seq_length=max_text_seq_length,\n spatial_interpolation_scale=spatial_interpolation_scale,\n temporal_interpolation_scale=temporal_interpolation_scale,\n use_positional_embeddings=not use_rotary_positional_embeddings,\n use_learned_positional_embeddings=use_learned_positional_embeddings,\n )\n self.embedding_dropout = nn.Dropout(dropout)\n\n # 2. Time embeddings\n self.time_proj = Timesteps(inner_dim, flip_sin_to_cos, freq_shift)\n self.time_embedding = TimestepEmbedding(inner_dim, time_embed_dim, timestep_activation_fn)\n\n # 3. Define spatio-temporal transformers blocks\n self.transformer_blocks = nn.ModuleList(\n [\n CogVideoXBlock(\n block_idx=idx,\n dim=inner_dim,\n num_attention_heads=num_attention_heads,\n attention_head_dim=attention_head_dim,\n time_embed_dim=time_embed_dim,\n dropout=dropout,\n activation_fn=activation_fn,\n attention_bias=attention_bias,\n norm_elementwise_affine=norm_elementwise_affine,\n norm_eps=norm_eps,\n finetune_init=finetune_init,\n )\n for idx in range(num_layers)\n ]\n )\n self.norm_final = nn.LayerNorm(inner_dim, norm_eps, norm_elementwise_affine)\n\n # 4. Output blocks\n self.norm_out = AdaLayerNorm(\n embedding_dim=time_embed_dim,\n output_dim=2 * inner_dim,\n norm_elementwise_affine=norm_elementwise_affine,\n norm_eps=norm_eps,\n chunk_dim=1,\n )\n self.proj_out = nn.Linear(inner_dim, patch_size * patch_size * out_channels)\n\n self.gradient_checkpointing = False\n \n def _set_gradient_checkpointing(self, module, value=False):\n self.gradient_checkpointing = value\n\n @property\n # Copied from diffusers.models.unets.unet_2d_condition.UNet2DConditionModel.attn_processors\n def attn_processors(self) -> Dict[str, AttentionProcessor]:\n r\"\"\"\n Returns:\n `dict` of attention processors: A dictionary containing all attention processors used in the model with\n indexed by its weight name.\n \"\"\"\n # set recursively\n processors = {}\n\n def fn_recursive_add_processors(name: str, module: torch.nn.Module, processors: Dict[str, AttentionProcessor]):\n if hasattr(module, \"get_processor\"):\n processors[f\"{name}.processor\"] = module.get_processor()\n\n for sub_name, child in module.named_children():\n fn_recursive_add_processors(f\"{name}.{sub_name}\", child, processors)\n\n return processors\n\n for name, module in self.named_children():\n fn_recursive_add_processors(name, module, processors)\n\n return processors\n\n # Copied from diffusers.models.unets.unet_2d_condition.UNet2DConditionModel.set_attn_processor\n def set_attn_processor(self, processor: Union[AttentionProcessor, Dict[str, AttentionProcessor]]):\n r\"\"\"\n Sets the attention processor to use to compute attention.\n\n Parameters:\n processor (`dict` of `AttentionProcessor` or only `AttentionProcessor`):\n The instantiated processor class or a dictionary of processor classes that will be set as the processor\n for **all** `Attention` layers.\n\n If `processor` is a dict, the key needs to define the path to the corresponding cross attention\n processor. This is strongly recommended when setting trainable attention processors.\n\n \"\"\"\n count = len(self.attn_processors.keys())\n\n if isinstance(processor, dict) and len(processor) != count:\n raise ValueError(\n f\"A dict of processors was passed, but the number of processors {len(processor)} does not match the\"\n f\" number of attention layers: {count}. Please make sure to pass {count} processor classes.\"\n )\n\n def fn_recursive_attn_processor(name: str, module: torch.nn.Module, processor):\n if hasattr(module, \"set_processor\"):\n if not isinstance(processor, dict):\n module.set_processor(processor)\n else:\n module.set_processor(processor.pop(f\"{name}.processor\"))\n\n for sub_name, child in module.named_children():\n fn_recursive_attn_processor(f\"{name}.{sub_name}\", child, processor)\n\n for name, module in self.named_children():\n fn_recursive_attn_processor(name, module, processor)\n\n # Copied from diffusers.models.unets.unet_2d_condition.UNet2DConditionModel.fuse_qkv_projections with FusedAttnProcessor2_0->FusedCogVideoXAttnProcessor2_0\n def fuse_qkv_projections(self):\n \"\"\"\n Enables fused QKV projections. For self-attention modules, all projection matrices (i.e., query, key, value)\n are fused. For cross-attention modules, key and value projection matrices are fused.\n\n \n\n This API is 🧪 experimental.\n\n \n \"\"\"\n self.original_attn_processors = None\n\n for _, attn_processor in self.attn_processors.items():\n if \"Added\" in str(attn_processor.__class__.__name__):\n raise ValueError(\"`fuse_qkv_projections()` is not supported for models having added KV projections.\")\n \n self.original_attn_processors = self.attn_processors\n\n for module in self.modules():\n if isinstance(module, Attention):\n module.fuse_projections(fuse=True)\n\n self.set_attn_processor(FusedCogVideoXAttnProcessor2_0())\n\n # Copied from diffusers.models.unets.unet_2d_condition.UNet2DConditionModel.unfuse_qkv_projections\n def unfuse_qkv_projections(self):\n \"\"\"Disables the fused QKV projection if enabled.\n\n \n\n This API is 🧪 experimental.\n\n \n\n \"\"\"\n if self.original_attn_processors is not None:\n self.set_attn_processor(self.original_attn_processors)\n\n def forward(\n self,\n hidden_states: torch.Tensor,\n encoder_hidden_states: torch.Tensor,\n empty_encoder_hidden_states: torch.Tensor,\n timestep: Union[int, float, torch.LongTensor],\n pose_embeds: Optional[torch.FloatTensor] = None,\n prompt_entities_embeds: Optional[torch.FloatTensor] = None,\n prompt_entities_attention_mask: Optional[torch.FloatTensor] = None,\n timestep_cond: Optional[torch.Tensor] = None,\n image_rotary_emb: Optional[Tuple[torch.Tensor, torch.Tensor]] = None,\n attention_kwargs: Optional[Dict[str, Any]] = None,\n return_dict: bool = True,\n ):\n if attention_kwargs is not None:\n attention_kwargs = attention_kwargs.copy()\n lora_scale = attention_kwargs.pop(\"scale\", 1.0)\n else:\n lora_scale = 1.0\n\n if USE_PEFT_BACKEND:\n # weight the lora layers by setting `lora_scale` for each PEFT layer\n scale_lora_layers(self, lora_scale)\n else:\n if attention_kwargs is not None and attention_kwargs.get(\"scale\", None) is not None:\n logger.warning(\n \"Passing `scale` via `attention_kwargs` when not using the PEFT backend is ineffective.\"\n )\n\n batch_size, num_frames, channels, height, width = hidden_states.shape\n\n # 1. Time embedding\n timesteps = timestep\n t_emb = self.time_proj(timesteps)\n\n # timesteps does not contain any weights and will always return f32 tensors\n # but time_embedding might actually be running in fp16. so we need to cast here.\n # there might be better ways to encapsulate this.\n t_emb = t_emb.to(dtype=hidden_states.dtype)\n emb = self.time_embedding(t_emb, timestep_cond)\n\n # 2. Patch embedding\n hidden_states, prompt_entities_embeds, empty_encoder_hidden_states = self.patch_embed(empty_encoder_hidden_states, encoder_hidden_states, hidden_states, prompt_entities_embeds, prompt_entities_attention_mask)\n hidden_states = self.embedding_dropout(hidden_states)\n\n text_seq_length = encoder_hidden_states.shape[1]\n encoder_hidden_states = hidden_states[:, :text_seq_length]\n hidden_states = hidden_states[:, text_seq_length:]\n\n # 3. Transformer blocks\n for i, block in enumerate(self.transformer_blocks):\n if self.training and self.gradient_checkpointing:\n\n def create_custom_forward(module):\n def custom_forward(*inputs):\n return module(*inputs)\n\n return custom_forward\n\n ckpt_kwargs: Dict[str, Any] = {\"use_reentrant\": False} if is_torch_version(\">=\", \"1.11.0\") else {}\n hidden_states, encoder_hidden_states = torch.utils.checkpoint.checkpoint(\n create_custom_forward(block),\n hidden_states,\n encoder_hidden_states,\n empty_encoder_hidden_states,\n emb,\n pose_embeds,\n prompt_entities_embeds,\n image_rotary_emb,\n **ckpt_kwargs,\n )\n else:\n hidden_states, encoder_hidden_states = block(\n hidden_states=hidden_states,\n encoder_hidden_states=encoder_hidden_states,\n empty_encoder_hidden_states=empty_encoder_hidden_states,\n temb=emb,\n pose_embeds=pose_embeds,\n prompt_entities_embeds=prompt_entities_embeds,\n image_rotary_emb=image_rotary_emb,\n )\n\n if not self.config.use_rotary_positional_embeddings:\n # CogVideoX-2B\n hidden_states = self.norm_final(hidden_states)\n else:\n # CogVideoX-5B\n hidden_states = torch.cat([encoder_hidden_states, hidden_states], dim=1)\n hidden_states = self.norm_final(hidden_states)\n hidden_states = hidden_states[:, text_seq_length:]\n\n # 4. Final block\n hidden_states = self.norm_out(hidden_states, temb=emb)\n hidden_states = self.proj_out(hidden_states)\n\n # 5. Unpatchify\n # Note: we use `-1` instead of `channels`:\n # - It is okay to `channels` use for CogVideoX-2b and CogVideoX-5b (number of input channels is equal to output channels)\n # - However, for CogVideoX-5b-I2V also takes concatenated input image latents (number of input channels is twice the output channels)\n p = self.config.patch_size\n output = hidden_states.reshape(batch_size, num_frames, height // p, width // p, -1, p, p)\n output = output.permute(0, 1, 4, 2, 5, 3, 6).flatten(5, 6).flatten(3, 4)\n\n if USE_PEFT_BACKEND:\n # remove `lora_scale` from each PEFT layer\n unscale_lora_layers(self, lora_scale)\n\n if not return_dict:\n return (output,)\n return Transformer2DModelOutput(sample=output)\n" }, { "path": "CogVideo/finetune/models/embeddings.py", "content": "# Copyright 2024 The HuggingFace Team. All rights reserved.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n# http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\nimport math\nfrom typing import List, Optional, Tuple, Union\n\nimport numpy as np\nimport torch\nimport torch.nn.functional as F\nfrom torch import nn\n\nfrom diffusers.utils import deprecate\nfrom diffusers.models.activations import FP32SiLU, get_activation\nfrom diffusers.models.attention_processor import Attention\n\nfrom einops import rearrange, repeat\n\n\ndef get_timestep_embedding(\n timesteps: torch.Tensor,\n embedding_dim: int,\n flip_sin_to_cos: bool = False,\n downscale_freq_shift: float = 1,\n scale: float = 1,\n max_period: int = 10000,\n):\n \"\"\"\n This matches the implementation in Denoising Diffusion Probabilistic Models: Create sinusoidal timestep embeddings.\n\n Args\n timesteps (torch.Tensor):\n a 1-D Tensor of N indices, one per batch element. These may be fractional.\n embedding_dim (int):\n the dimension of the output.\n flip_sin_to_cos (bool):\n Whether the embedding order should be `cos, sin` (if True) or `sin, cos` (if False)\n downscale_freq_shift (float):\n Controls the delta between frequencies between dimensions\n scale (float):\n Scaling factor applied to the embeddings.\n max_period (int):\n Controls the maximum frequency of the embeddings\n Returns\n torch.Tensor: an [N x dim] Tensor of positional embeddings.\n \"\"\"\n assert len(timesteps.shape) == 1, \"Timesteps should be a 1d-array\"\n\n half_dim = embedding_dim // 2\n exponent = -math.log(max_period) * torch.arange(\n start=0, end=half_dim, dtype=torch.float32, device=timesteps.device\n )\n exponent = exponent / (half_dim - downscale_freq_shift)\n\n emb = torch.exp(exponent)\n emb = timesteps[:, None].float() * emb[None, :]\n\n # scale embeddings\n emb = scale * emb\n\n # concat sine and cosine embeddings\n emb = torch.cat([torch.sin(emb), torch.cos(emb)], dim=-1)\n\n # flip sine and cosine embeddings\n if flip_sin_to_cos:\n emb = torch.cat([emb[:, half_dim:], emb[:, :half_dim]], dim=-1)\n\n # zero pad\n if embedding_dim % 2 == 1:\n emb = torch.nn.functional.pad(emb, (0, 1, 0, 0))\n return emb\n\n\ndef get_3d_sincos_pos_embed(\n embed_dim: int,\n spatial_size: Union[int, Tuple[int, int]],\n temporal_size: int,\n spatial_interpolation_scale: float = 1.0,\n temporal_interpolation_scale: float = 1.0,\n) -> np.ndarray:\n r\"\"\"\n Args:\n embed_dim (`int`):\n spatial_size (`int` or `Tuple[int, int]`):\n temporal_size (`int`):\n spatial_interpolation_scale (`float`, defaults to 1.0):\n temporal_interpolation_scale (`float`, defaults to 1.0):\n \"\"\"\n if embed_dim % 4 != 0:\n raise ValueError(\"`embed_dim` must be divisible by 4\")\n if isinstance(spatial_size, int):\n spatial_size = (spatial_size, spatial_size)\n\n embed_dim_spatial = 3 * embed_dim // 4\n embed_dim_temporal = embed_dim // 4\n\n # 1. Spatial\n grid_h = np.arange(spatial_size[1], dtype=np.float32) / spatial_interpolation_scale\n grid_w = np.arange(spatial_size[0], dtype=np.float32) / spatial_interpolation_scale\n grid = np.meshgrid(grid_w, grid_h) # here w goes first\n grid = np.stack(grid, axis=0)\n\n grid = grid.reshape([2, 1, spatial_size[1], spatial_size[0]])\n pos_embed_spatial = get_2d_sincos_pos_embed_from_grid(embed_dim_spatial, grid)\n\n # 2. Temporal\n grid_t = np.arange(temporal_size, dtype=np.float32) / temporal_interpolation_scale\n pos_embed_temporal = get_1d_sincos_pos_embed_from_grid(embed_dim_temporal, grid_t)\n\n # 3. Concat\n pos_embed_spatial = pos_embed_spatial[np.newaxis, :, :]\n pos_embed_spatial = np.repeat(pos_embed_spatial, temporal_size, axis=0) # [T, H*W, D // 4 * 3]\n\n pos_embed_temporal = pos_embed_temporal[:, np.newaxis, :]\n pos_embed_temporal = np.repeat(pos_embed_temporal, spatial_size[0] * spatial_size[1], axis=1) # [T, H*W, D // 4]\n\n pos_embed = np.concatenate([pos_embed_temporal, pos_embed_spatial], axis=-1) # [T, H*W, D]\n return pos_embed\n\n\ndef get_2d_sincos_pos_embed(\n embed_dim, grid_size, cls_token=False, extra_tokens=0, interpolation_scale=1.0, base_size=16\n):\n \"\"\"\n grid_size: int of the grid height and width return: pos_embed: [grid_size*grid_size, embed_dim] or\n [1+grid_size*grid_size, embed_dim] (w/ or w/o cls_token)\n \"\"\"\n if isinstance(grid_size, int):\n grid_size = (grid_size, grid_size)\n\n grid_h = np.arange(grid_size[0], dtype=np.float32) / (grid_size[0] / base_size) / interpolation_scale\n grid_w = np.arange(grid_size[1], dtype=np.float32) / (grid_size[1] / base_size) / interpolation_scale\n grid = np.meshgrid(grid_w, grid_h) # here w goes first\n grid = np.stack(grid, axis=0)\n\n grid = grid.reshape([2, 1, grid_size[1], grid_size[0]])\n pos_embed = get_2d_sincos_pos_embed_from_grid(embed_dim, grid)\n if cls_token and extra_tokens > 0:\n pos_embed = np.concatenate([np.zeros([extra_tokens, embed_dim]), pos_embed], axis=0)\n return pos_embed\n\n\ndef get_2d_sincos_pos_embed_from_grid(embed_dim, grid):\n if embed_dim % 2 != 0:\n raise ValueError(\"embed_dim must be divisible by 2\")\n\n # use half of dimensions to encode grid_h\n emb_h = get_1d_sincos_pos_embed_from_grid(embed_dim // 2, grid[0]) # (H*W, D/2)\n emb_w = get_1d_sincos_pos_embed_from_grid(embed_dim // 2, grid[1]) # (H*W, D/2)\n\n emb = np.concatenate([emb_h, emb_w], axis=1) # (H*W, D)\n return emb\n\n\ndef get_1d_sincos_pos_embed_from_grid(embed_dim, pos):\n \"\"\"\n embed_dim: output dimension for each position pos: a list of positions to be encoded: size (M,) out: (M, D)\n \"\"\"\n if embed_dim % 2 != 0:\n raise ValueError(\"embed_dim must be divisible by 2\")\n\n omega = np.arange(embed_dim // 2, dtype=np.float64)\n omega /= embed_dim / 2.0\n omega = 1.0 / 10000**omega # (D/2,)\n\n pos = pos.reshape(-1) # (M,)\n out = np.einsum(\"m,d->md\", pos, omega) # (M, D/2), outer product\n\n emb_sin = np.sin(out) # (M, D/2)\n emb_cos = np.cos(out) # (M, D/2)\n\n emb = np.concatenate([emb_sin, emb_cos], axis=1) # (M, D)\n return emb\n\n\nclass PatchEmbed(nn.Module):\n \"\"\"2D Image to Patch Embedding with support for SD3 cropping.\"\"\"\n\n def __init__(\n self,\n height=224,\n width=224,\n patch_size=16,\n in_channels=3,\n embed_dim=768,\n layer_norm=False,\n flatten=True,\n bias=True,\n interpolation_scale=1,\n pos_embed_type=\"sincos\",\n pos_embed_max_size=None, # For SD3 cropping\n ):\n super().__init__()\n\n num_patches = (height // patch_size) * (width // patch_size)\n self.flatten = flatten\n self.layer_norm = layer_norm\n self.pos_embed_max_size = pos_embed_max_size\n\n self.proj = nn.Conv2d(\n in_channels, embed_dim, kernel_size=(patch_size, patch_size), stride=patch_size, bias=bias\n )\n if layer_norm:\n self.norm = nn.LayerNorm(embed_dim, elementwise_affine=False, eps=1e-6)\n else:\n self.norm = None\n\n self.patch_size = patch_size\n self.height, self.width = height // patch_size, width // patch_size\n self.base_size = height // patch_size\n self.interpolation_scale = interpolation_scale\n\n # Calculate positional embeddings based on max size or default\n if pos_embed_max_size:\n grid_size = pos_embed_max_size\n else:\n grid_size = int(num_patches**0.5)\n\n if pos_embed_type is None:\n self.pos_embed = None\n elif pos_embed_type == \"sincos\":\n pos_embed = get_2d_sincos_pos_embed(\n embed_dim, grid_size, base_size=self.base_size, interpolation_scale=self.interpolation_scale\n )\n persistent = True if pos_embed_max_size else False\n self.register_buffer(\"pos_embed\", torch.from_numpy(pos_embed).float().unsqueeze(0), persistent=persistent)\n else:\n raise ValueError(f\"Unsupported pos_embed_type: {pos_embed_type}\")\n\n def cropped_pos_embed(self, height, width):\n \"\"\"Crops positional embeddings for SD3 compatibility.\"\"\"\n if self.pos_embed_max_size is None:\n raise ValueError(\"`pos_embed_max_size` must be set for cropping.\")\n\n height = height // self.patch_size\n width = width // self.patch_size\n if height > self.pos_embed_max_size:\n raise ValueError(\n f\"Height ({height}) cannot be greater than `pos_embed_max_size`: {self.pos_embed_max_size}.\"\n )\n if width > self.pos_embed_max_size:\n raise ValueError(\n f\"Width ({width}) cannot be greater than `pos_embed_max_size`: {self.pos_embed_max_size}.\"\n )\n\n top = (self.pos_embed_max_size - height) // 2\n left = (self.pos_embed_max_size - width) // 2\n spatial_pos_embed = self.pos_embed.reshape(1, self.pos_embed_max_size, self.pos_embed_max_size, -1)\n spatial_pos_embed = spatial_pos_embed[:, top : top + height, left : left + width, :]\n spatial_pos_embed = spatial_pos_embed.reshape(1, -1, spatial_pos_embed.shape[-1])\n return spatial_pos_embed\n\n def forward(self, latent):\n if self.pos_embed_max_size is not None:\n height, width = latent.shape[-2:]\n else:\n height, width = latent.shape[-2] // self.patch_size, latent.shape[-1] // self.patch_size\n\n latent = self.proj(latent)\n if self.flatten:\n latent = latent.flatten(2).transpose(1, 2) # BCHW -> BNC\n if self.layer_norm:\n latent = self.norm(latent)\n if self.pos_embed is None:\n return latent.to(latent.dtype)\n # Interpolate or crop positional embeddings as needed\n if self.pos_embed_max_size:\n pos_embed = self.cropped_pos_embed(height, width)\n else:\n if self.height != height or self.width != width:\n pos_embed = get_2d_sincos_pos_embed(\n embed_dim=self.pos_embed.shape[-1],\n grid_size=(height, width),\n base_size=self.base_size,\n interpolation_scale=self.interpolation_scale,\n )\n pos_embed = torch.from_numpy(pos_embed).float().unsqueeze(0).to(latent.device)\n else:\n pos_embed = self.pos_embed\n\n return (latent + pos_embed).to(latent.dtype)\n\n\nclass LuminaPatchEmbed(nn.Module):\n \"\"\"2D Image to Patch Embedding with support for Lumina-T2X\"\"\"\n\n def __init__(self, patch_size=2, in_channels=4, embed_dim=768, bias=True):\n super().__init__()\n self.patch_size = patch_size\n self.proj = nn.Linear(\n in_features=patch_size * patch_size * in_channels,\n out_features=embed_dim,\n bias=bias,\n )\n\n def forward(self, x, freqs_cis):\n \"\"\"\n Patchifies and embeds the input tensor(s).\n\n Args:\n x (List[torch.Tensor] | torch.Tensor): The input tensor(s) to be patchified and embedded.\n\n Returns:\n Tuple[torch.Tensor, torch.Tensor, List[Tuple[int, int]], torch.Tensor]: A tuple containing the patchified\n and embedded tensor(s), the mask indicating the valid patches, the original image size(s), and the\n frequency tensor(s).\n \"\"\"\n freqs_cis = freqs_cis.to(x[0].device)\n patch_height = patch_width = self.patch_size\n batch_size, channel, height, width = x.size()\n height_tokens, width_tokens = height // patch_height, width // patch_width\n\n x = x.view(batch_size, channel, height_tokens, patch_height, width_tokens, patch_width).permute(\n 0, 2, 4, 1, 3, 5\n )\n x = x.flatten(3)\n x = self.proj(x)\n x = x.flatten(1, 2)\n\n mask = torch.ones(x.shape[0], x.shape[1], dtype=torch.int32, device=x.device)\n\n return (\n x,\n mask,\n [(height, width)] * batch_size,\n freqs_cis[:height_tokens, :width_tokens].flatten(0, 1).unsqueeze(0),\n )\n\n\nclass CogVideoXPatchEmbed(nn.Module):\n def __init__(\n self,\n patch_size: int = 2,\n in_channels: int = 16,\n embed_dim: int = 1920,\n text_embed_dim: int = 4096,\n bias: bool = True,\n sample_width: int = 90,\n sample_height: int = 60,\n sample_frames: int = 49,\n temporal_compression_ratio: int = 4,\n max_text_seq_length: int = 226,\n spatial_interpolation_scale: float = 1.875,\n temporal_interpolation_scale: float = 1.0,\n use_positional_embeddings: bool = True,\n use_learned_positional_embeddings: bool = True,\n ) -> None:\n super().__init__()\n\n self.patch_size = patch_size\n self.embed_dim = embed_dim\n self.sample_height = sample_height\n self.sample_width = sample_width\n self.sample_frames = sample_frames\n self.temporal_compression_ratio = temporal_compression_ratio\n self.max_text_seq_length = max_text_seq_length\n self.spatial_interpolation_scale = spatial_interpolation_scale\n self.temporal_interpolation_scale = temporal_interpolation_scale\n self.use_positional_embeddings = use_positional_embeddings\n self.use_learned_positional_embeddings = use_learned_positional_embeddings\n\n self.proj = nn.Conv2d(\n in_channels, embed_dim, kernel_size=(patch_size, patch_size), stride=patch_size, bias=bias\n )\n self.text_proj = nn.Linear(text_embed_dim, embed_dim)\n\n if use_positional_embeddings or use_learned_positional_embeddings:\n persistent = use_learned_positional_embeddings\n pos_embedding = self._get_positional_embeddings(sample_height, sample_width, sample_frames)\n self.register_buffer(\"pos_embedding\", pos_embedding, persistent=persistent)\n\n def _get_positional_embeddings(self, sample_height: int, sample_width: int, sample_frames: int) -> torch.Tensor:\n post_patch_height = sample_height // self.patch_size\n post_patch_width = sample_width // self.patch_size\n post_time_compression_frames = (sample_frames - 1) // self.temporal_compression_ratio + 1\n num_patches = post_patch_height * post_patch_width * post_time_compression_frames\n\n pos_embedding = get_3d_sincos_pos_embed(\n self.embed_dim,\n (post_patch_width, post_patch_height),\n post_time_compression_frames,\n self.spatial_interpolation_scale,\n self.temporal_interpolation_scale,\n )\n pos_embedding = torch.from_numpy(pos_embedding).flatten(0, 1)\n joint_pos_embedding = torch.zeros(\n 1, self.max_text_seq_length + num_patches, self.embed_dim, requires_grad=False\n )\n joint_pos_embedding.data[:, self.max_text_seq_length :].copy_(pos_embedding)\n\n return joint_pos_embedding\n\n def forward(self, empty_text_embeds: torch.Tensor, text_embeds: torch.Tensor, image_embeds: torch.Tensor, prompt_entities_embeds: torch.Tensor, prompt_entities_attention_mask: torch.Tensor):\n\n r\"\"\"\n Args:\n text_embeds (`torch.Tensor`):\n Input text embeddings. Expected shape: (batch_size, seq_length, embedding_dim).\n image_embeds (`torch.Tensor`):\n Input image embeddings. Expected shape: (batch_size, num_frames, channels, height, width).\n \"\"\"\n\n batch, num_frames, channels, height, width = image_embeds.shape\n\n empty_text_embeds = self.text_proj(empty_text_embeds)\n text_embeds = self.text_proj(text_embeds)\n\n prompt_entities_embeds = self.text_proj(prompt_entities_embeds)\n prompt_entities_embeds[prompt_entities_attention_mask<1.] = 0.\n prompt_entities_embeds = prompt_entities_embeds[:,:,:50,:]\n\n image_embeds = image_embeds.reshape(-1, channels, height, width)\n image_embeds = self.proj(image_embeds)\n image_embeds = image_embeds.view(batch, num_frames, *image_embeds.shape[1:])\n image_embeds = image_embeds.flatten(3).transpose(2, 3) # [batch, num_frames, height x width, channels]\n image_embeds = image_embeds.flatten(1, 2) # [batch, num_frames x height x width, channels]\n\n embeds = torch.cat(\n [text_embeds, image_embeds], dim=1\n ).contiguous() # [batch, seq_length + num_frames x height x width, channels]\n\n if self.use_positional_embeddings or self.use_learned_positional_embeddings:\n if self.use_learned_positional_embeddings and (self.sample_width != width or self.sample_height != height):\n raise ValueError(\n \"It is currently not possible to generate videos at a different resolution that the defaults. This should only be the case with 'THUDM/CogVideoX-5b-I2V'.\"\n \"If you think this is incorrect, please open an issue at https://github.com/huggingface/diffusers/issues.\"\n )\n\n pre_time_compression_frames = (num_frames - 1) * self.temporal_compression_ratio + 1\n\n if (\n self.sample_height != height\n or self.sample_width != width\n or self.sample_frames != pre_time_compression_frames\n ):\n pos_embedding = self._get_positional_embeddings(height, width, pre_time_compression_frames)\n pos_embedding = pos_embedding.to(embeds.device, dtype=embeds.dtype)\n else:\n pos_embedding = self.pos_embedding\n\n embeds = embeds + pos_embedding\n\n return embeds, prompt_entities_embeds, empty_text_embeds\n\n\nclass CogView3PlusPatchEmbed(nn.Module):\n def __init__(\n self,\n in_channels: int = 16,\n hidden_size: int = 2560,\n patch_size: int = 2,\n text_hidden_size: int = 4096,\n pos_embed_max_size: int = 128,\n ):\n super().__init__()\n self.in_channels = in_channels\n self.hidden_size = hidden_size\n self.patch_size = patch_size\n self.text_hidden_size = text_hidden_size\n self.pos_embed_max_size = pos_embed_max_size\n # Linear projection for image patches\n self.proj = nn.Linear(in_channels * patch_size**2, hidden_size)\n\n # Linear projection for text embeddings\n self.text_proj = nn.Linear(text_hidden_size, hidden_size)\n\n pos_embed = get_2d_sincos_pos_embed(hidden_size, pos_embed_max_size, base_size=pos_embed_max_size)\n pos_embed = pos_embed.reshape(pos_embed_max_size, pos_embed_max_size, hidden_size)\n self.register_buffer(\"pos_embed\", torch.from_numpy(pos_embed).float(), persistent=False)\n\n def forward(self, hidden_states: torch.Tensor, encoder_hidden_states: torch.Tensor) -> torch.Tensor:\n batch_size, channel, height, width = hidden_states.shape\n\n if height % self.patch_size != 0 or width % self.patch_size != 0:\n raise ValueError(\"Height and width must be divisible by patch size\")\n\n height = height // self.patch_size\n width = width // self.patch_size\n hidden_states = hidden_states.view(batch_size, channel, height, self.patch_size, width, self.patch_size)\n hidden_states = hidden_states.permute(0, 2, 4, 1, 3, 5).contiguous()\n hidden_states = hidden_states.view(batch_size, height * width, channel * self.patch_size * self.patch_size)\n\n # Project the patches\n hidden_states = self.proj(hidden_states)\n encoder_hidden_states = self.text_proj(encoder_hidden_states)\n hidden_states = torch.cat([encoder_hidden_states, hidden_states], dim=1)\n\n # Calculate text_length\n text_length = encoder_hidden_states.shape[1]\n\n image_pos_embed = self.pos_embed[:height, :width].reshape(height * width, -1)\n text_pos_embed = torch.zeros(\n (text_length, self.hidden_size), dtype=image_pos_embed.dtype, device=image_pos_embed.device\n )\n pos_embed = torch.cat([text_pos_embed, image_pos_embed], dim=0)[None, ...]\n\n return (hidden_states + pos_embed).to(hidden_states.dtype)\n\n\ndef get_3d_rotary_pos_embed(\n embed_dim, crops_coords, grid_size, temporal_size, theta: int = 10000, use_real: bool = True\n) -> Union[torch.Tensor, Tuple[torch.Tensor, torch.Tensor]]:\n \"\"\"\n RoPE for video tokens with 3D structure.\n\n Args:\n embed_dim: (`int`):\n The embedding dimension size, corresponding to hidden_size_head.\n crops_coords (`Tuple[int]`):\n The top-left and bottom-right coordinates of the crop.\n grid_size (`Tuple[int]`):\n The grid size of the spatial positional embedding (height, width).\n temporal_size (`int`):\n The size of the temporal dimension.\n theta (`float`):\n Scaling factor for frequency computation.\n\n Returns:\n `torch.Tensor`: positional embedding with shape `(temporal_size * grid_size[0] * grid_size[1], embed_dim/2)`.\n \"\"\"\n if use_real is not True:\n raise ValueError(\" `use_real = False` is not currently supported for get_3d_rotary_pos_embed\")\n start, stop = crops_coords\n grid_size_h, grid_size_w = grid_size\n grid_h = np.linspace(start[0], stop[0], grid_size_h, endpoint=False, dtype=np.float32)\n grid_w = np.linspace(start[1], stop[1], grid_size_w, endpoint=False, dtype=np.float32)\n grid_t = np.linspace(0, temporal_size, temporal_size, endpoint=False, dtype=np.float32)\n\n # Compute dimensions for each axis\n dim_t = embed_dim // 4\n dim_h = embed_dim // 8 * 3\n dim_w = embed_dim // 8 * 3\n\n # Temporal frequencies\n freqs_t = get_1d_rotary_pos_embed(dim_t, grid_t, use_real=True)\n # Spatial frequencies for height and width\n freqs_h = get_1d_rotary_pos_embed(dim_h, grid_h, use_real=True)\n freqs_w = get_1d_rotary_pos_embed(dim_w, grid_w, use_real=True)\n\n # BroadCast and concatenate temporal and spaial frequencie (height and width) into a 3d tensor\n def combine_time_height_width(freqs_t, freqs_h, freqs_w):\n freqs_t = freqs_t[:, None, None, :].expand(\n -1, grid_size_h, grid_size_w, -1\n ) # temporal_size, grid_size_h, grid_size_w, dim_t\n freqs_h = freqs_h[None, :, None, :].expand(\n temporal_size, -1, grid_size_w, -1\n ) # temporal_size, grid_size_h, grid_size_2, dim_h\n freqs_w = freqs_w[None, None, :, :].expand(\n temporal_size, grid_size_h, -1, -1\n ) # temporal_size, grid_size_h, grid_size_2, dim_w\n\n freqs = torch.cat(\n [freqs_t, freqs_h, freqs_w], dim=-1\n ) # temporal_size, grid_size_h, grid_size_w, (dim_t + dim_h + dim_w)\n freqs = freqs.view(\n temporal_size * grid_size_h * grid_size_w, -1\n ) # (temporal_size * grid_size_h * grid_size_w), (dim_t + dim_h + dim_w)\n return freqs\n\n t_cos, t_sin = freqs_t # both t_cos and t_sin has shape: temporal_size, dim_t\n h_cos, h_sin = freqs_h # both h_cos and h_sin has shape: grid_size_h, dim_h\n w_cos, w_sin = freqs_w # both w_cos and w_sin has shape: grid_size_w, dim_w\n cos = combine_time_height_width(t_cos, h_cos, w_cos)\n sin = combine_time_height_width(t_sin, h_sin, w_sin)\n return cos, sin\n\n\ndef get_2d_rotary_pos_embed(embed_dim, crops_coords, grid_size, use_real=True):\n \"\"\"\n RoPE for image tokens with 2d structure.\n\n Args:\n embed_dim: (`int`):\n The embedding dimension size\n crops_coords (`Tuple[int]`)\n The top-left and bottom-right coordinates of the crop.\n grid_size (`Tuple[int]`):\n The grid size of the positional embedding.\n use_real (`bool`):\n If True, return real part and imaginary part separately. Otherwise, return complex numbers.\n\n Returns:\n `torch.Tensor`: positional embedding with shape `( grid_size * grid_size, embed_dim/2)`.\n \"\"\"\n start, stop = crops_coords\n grid_h = np.linspace(start[0], stop[0], grid_size[0], endpoint=False, dtype=np.float32)\n grid_w = np.linspace(start[1], stop[1], grid_size[1], endpoint=False, dtype=np.float32)\n grid = np.meshgrid(grid_w, grid_h) # here w goes first\n grid = np.stack(grid, axis=0) # [2, W, H]\n\n grid = grid.reshape([2, 1, *grid.shape[1:]])\n pos_embed = get_2d_rotary_pos_embed_from_grid(embed_dim, grid, use_real=use_real)\n return pos_embed\n\n\ndef get_2d_rotary_pos_embed_from_grid(embed_dim, grid, use_real=False):\n assert embed_dim % 4 == 0\n\n # use half of dimensions to encode grid_h\n emb_h = get_1d_rotary_pos_embed(\n embed_dim // 2, grid[0].reshape(-1), use_real=use_real\n ) # (H*W, D/2) if use_real else (H*W, D/4)\n emb_w = get_1d_rotary_pos_embed(\n embed_dim // 2, grid[1].reshape(-1), use_real=use_real\n ) # (H*W, D/2) if use_real else (H*W, D/4)\n\n if use_real:\n cos = torch.cat([emb_h[0], emb_w[0]], dim=1) # (H*W, D)\n sin = torch.cat([emb_h[1], emb_w[1]], dim=1) # (H*W, D)\n return cos, sin\n else:\n emb = torch.cat([emb_h, emb_w], dim=1) # (H*W, D/2)\n return emb\n\n\ndef get_2d_rotary_pos_embed_lumina(embed_dim, len_h, len_w, linear_factor=1.0, ntk_factor=1.0):\n assert embed_dim % 4 == 0\n\n emb_h = get_1d_rotary_pos_embed(\n embed_dim // 2, len_h, linear_factor=linear_factor, ntk_factor=ntk_factor\n ) # (H, D/4)\n emb_w = get_1d_rotary_pos_embed(\n embed_dim // 2, len_w, linear_factor=linear_factor, ntk_factor=ntk_factor\n ) # (W, D/4)\n emb_h = emb_h.view(len_h, 1, embed_dim // 4, 1).repeat(1, len_w, 1, 1) # (H, W, D/4, 1)\n emb_w = emb_w.view(1, len_w, embed_dim // 4, 1).repeat(len_h, 1, 1, 1) # (H, W, D/4, 1)\n\n emb = torch.cat([emb_h, emb_w], dim=-1).flatten(2) # (H, W, D/2)\n return emb\n\n\ndef get_1d_rotary_pos_embed(\n dim: int,\n pos: Union[np.ndarray, int],\n theta: float = 10000.0,\n use_real=False,\n linear_factor=1.0,\n ntk_factor=1.0,\n repeat_interleave_real=True,\n freqs_dtype=torch.float32, # torch.float32, torch.float64 (flux)\n):\n \"\"\"\n Precompute the frequency tensor for complex exponentials (cis) with given dimensions.\n\n This function calculates a frequency tensor with complex exponentials using the given dimension 'dim' and the end\n index 'end'. The 'theta' parameter scales the frequencies. The returned tensor contains complex values in complex64\n data type.\n\n Args:\n dim (`int`): Dimension of the frequency tensor.\n pos (`np.ndarray` or `int`): Position indices for the frequency tensor. [S] or scalar\n theta (`float`, *optional*, defaults to 10000.0):\n Scaling factor for frequency computation. Defaults to 10000.0.\n use_real (`bool`, *optional*):\n If True, return real part and imaginary part separately. Otherwise, return complex numbers.\n linear_factor (`float`, *optional*, defaults to 1.0):\n Scaling factor for the context extrapolation. Defaults to 1.0.\n ntk_factor (`float`, *optional*, defaults to 1.0):\n Scaling factor for the NTK-Aware RoPE. Defaults to 1.0.\n repeat_interleave_real (`bool`, *optional*, defaults to `True`):\n If `True` and `use_real`, real part and imaginary part are each interleaved with themselves to reach `dim`.\n Otherwise, they are concateanted with themselves.\n freqs_dtype (`torch.float32` or `torch.float64`, *optional*, defaults to `torch.float32`):\n the dtype of the frequency tensor.\n Returns:\n `torch.Tensor`: Precomputed frequency tensor with complex exponentials. [S, D/2]\n \"\"\"\n assert dim % 2 == 0\n\n if isinstance(pos, int):\n pos = torch.arange(pos)\n if isinstance(pos, np.ndarray):\n pos = torch.from_numpy(pos) # type: ignore # [S]\n\n theta = theta * ntk_factor\n freqs = (\n 1.0\n / (theta ** (torch.arange(0, dim, 2, dtype=freqs_dtype, device=pos.device)[: (dim // 2)] / dim))\n / linear_factor\n ) # [D/2]\n freqs = torch.outer(pos, freqs) # type: ignore # [S, D/2]\n if use_real and repeat_interleave_real:\n # flux, hunyuan-dit, cogvideox\n freqs_cos = freqs.cos().repeat_interleave(2, dim=1).float() # [S, D]\n freqs_sin = freqs.sin().repeat_interleave(2, dim=1).float() # [S, D]\n return freqs_cos, freqs_sin\n elif use_real:\n # stable audio\n freqs_cos = torch.cat([freqs.cos(), freqs.cos()], dim=-1).float() # [S, D]\n freqs_sin = torch.cat([freqs.sin(), freqs.sin()], dim=-1).float() # [S, D]\n return freqs_cos, freqs_sin\n else:\n # lumina\n freqs_cis = torch.polar(torch.ones_like(freqs), freqs) # complex64 # [S, D/2]\n return freqs_cis\n\n\ndef apply_rotary_emb(\n x: torch.Tensor,\n freqs_cis: Union[torch.Tensor, Tuple[torch.Tensor]],\n use_real: bool = True,\n use_real_unbind_dim: int = -1,\n) -> Tuple[torch.Tensor, torch.Tensor]:\n \"\"\"\n Apply rotary embeddings to input tensors using the given frequency tensor. This function applies rotary embeddings\n to the given query or key 'x' tensors using the provided frequency tensor 'freqs_cis'. The input tensors are\n reshaped as complex numbers, and the frequency tensor is reshaped for broadcasting compatibility. The resulting\n tensors contain rotary embeddings and are returned as real tensors.\n\n Args:\n x (`torch.Tensor`):\n Query or key tensor to apply rotary embeddings. [B, H, S, D] xk (torch.Tensor): Key tensor to apply\n freqs_cis (`Tuple[torch.Tensor]`): Precomputed frequency tensor for complex exponentials. ([S, D], [S, D],)\n\n Returns:\n Tuple[torch.Tensor, torch.Tensor]: Tuple of modified query tensor and key tensor with rotary embeddings.\n \"\"\"\n if use_real:\n cos, sin = freqs_cis # [S, D]\n cos = cos[None, None]\n sin = sin[None, None]\n cos, sin = cos.to(x.device), sin.to(x.device)\n\n if use_real_unbind_dim == -1:\n # Used for flux, cogvideox, hunyuan-dit\n x_real, x_imag = x.reshape(*x.shape[:-1], -1, 2).unbind(-1) # [B, S, H, D//2]\n x_rotated = torch.stack([-x_imag, x_real], dim=-1).flatten(3)\n elif use_real_unbind_dim == -2:\n # Used for Stable Audio\n x_real, x_imag = x.reshape(*x.shape[:-1], 2, -1).unbind(-2) # [B, S, H, D//2]\n x_rotated = torch.cat([-x_imag, x_real], dim=-1)\n else:\n raise ValueError(f\"`use_real_unbind_dim={use_real_unbind_dim}` but should be -1 or -2.\")\n\n out = (x.float() * cos + x_rotated.float() * sin).to(x.dtype)\n\n return out\n else:\n # used for lumina\n x_rotated = torch.view_as_complex(x.float().reshape(*x.shape[:-1], -1, 2))\n freqs_cis = freqs_cis.unsqueeze(2)\n x_out = torch.view_as_real(x_rotated * freqs_cis).flatten(3)\n\n return x_out.type_as(x)\n\n\nclass FluxPosEmbed(nn.Module):\n # modified from https://github.com/black-forest-labs/flux/blob/c00d7c60b085fce8058b9df845e036090873f2ce/src/flux/modules/layers.py#L11\n def __init__(self, theta: int, axes_dim: List[int]):\n super().__init__()\n self.theta = theta\n self.axes_dim = axes_dim\n\n def forward(self, ids: torch.Tensor) -> torch.Tensor:\n n_axes = ids.shape[-1]\n cos_out = []\n sin_out = []\n pos = ids.float()\n is_mps = ids.device.type == \"mps\"\n freqs_dtype = torch.float32 if is_mps else torch.float64\n for i in range(n_axes):\n cos, sin = get_1d_rotary_pos_embed(\n self.axes_dim[i], pos[:, i], repeat_interleave_real=True, use_real=True, freqs_dtype=freqs_dtype\n )\n cos_out.append(cos)\n sin_out.append(sin)\n freqs_cos = torch.cat(cos_out, dim=-1).to(ids.device)\n freqs_sin = torch.cat(sin_out, dim=-1).to(ids.device)\n return freqs_cos, freqs_sin\n\n\nclass TimestepEmbedding(nn.Module):\n def __init__(\n self,\n in_channels: int,\n time_embed_dim: int,\n act_fn: str = \"silu\",\n out_dim: int = None,\n post_act_fn: Optional[str] = None,\n cond_proj_dim=None,\n sample_proj_bias=True,\n ):\n super().__init__()\n\n self.linear_1 = nn.Linear(in_channels, time_embed_dim, sample_proj_bias)\n\n if cond_proj_dim is not None:\n self.cond_proj = nn.Linear(cond_proj_dim, in_channels, bias=False)\n else:\n self.cond_proj = None\n\n self.act = get_activation(act_fn)\n\n if out_dim is not None:\n time_embed_dim_out = out_dim\n else:\n time_embed_dim_out = time_embed_dim\n self.linear_2 = nn.Linear(time_embed_dim, time_embed_dim_out, sample_proj_bias)\n\n if post_act_fn is None:\n self.post_act = None\n else:\n self.post_act = get_activation(post_act_fn)\n\n def forward(self, sample, condition=None):\n if condition is not None:\n sample = sample + self.cond_proj(condition)\n sample = self.linear_1(sample)\n\n if self.act is not None:\n sample = self.act(sample)\n\n sample = self.linear_2(sample)\n\n if self.post_act is not None:\n sample = self.post_act(sample)\n return sample\n\n\nclass Timesteps(nn.Module):\n def __init__(self, num_channels: int, flip_sin_to_cos: bool, downscale_freq_shift: float, scale: int = 1):\n super().__init__()\n self.num_channels = num_channels\n self.flip_sin_to_cos = flip_sin_to_cos\n self.downscale_freq_shift = downscale_freq_shift\n self.scale = scale\n\n def forward(self, timesteps):\n t_emb = get_timestep_embedding(\n timesteps,\n self.num_channels,\n flip_sin_to_cos=self.flip_sin_to_cos,\n downscale_freq_shift=self.downscale_freq_shift,\n scale=self.scale,\n )\n return t_emb\n\n\nclass GaussianFourierProjection(nn.Module):\n \"\"\"Gaussian Fourier embeddings for noise levels.\"\"\"\n\n def __init__(\n self, embedding_size: int = 256, scale: float = 1.0, set_W_to_weight=True, log=True, flip_sin_to_cos=False\n ):\n super().__init__()\n self.weight = nn.Parameter(torch.randn(embedding_size) * scale, requires_grad=False)\n self.log = log\n self.flip_sin_to_cos = flip_sin_to_cos\n\n if set_W_to_weight:\n # to delete later\n del self.weight\n self.W = nn.Parameter(torch.randn(embedding_size) * scale, requires_grad=False)\n self.weight = self.W\n del self.W\n\n def forward(self, x):\n if self.log:\n x = torch.log(x)\n\n x_proj = x[:, None] * self.weight[None, :] * 2 * np.pi\n\n if self.flip_sin_to_cos:\n out = torch.cat([torch.cos(x_proj), torch.sin(x_proj)], dim=-1)\n else:\n out = torch.cat([torch.sin(x_proj), torch.cos(x_proj)], dim=-1)\n return out\n\n\nclass SinusoidalPositionalEmbedding(nn.Module):\n \"\"\"Apply positional information to a sequence of embeddings.\n\n Takes in a sequence of embeddings with shape (batch_size, seq_length, embed_dim) and adds positional embeddings to\n them\n\n Args:\n embed_dim: (int): Dimension of the positional embedding.\n max_seq_length: Maximum sequence length to apply positional embeddings\n\n \"\"\"\n\n def __init__(self, embed_dim: int, max_seq_length: int = 32):\n super().__init__()\n position = torch.arange(max_seq_length).unsqueeze(1)\n div_term = torch.exp(torch.arange(0, embed_dim, 2) * (-math.log(10000.0) / embed_dim))\n pe = torch.zeros(1, max_seq_length, embed_dim)\n pe[0, :, 0::2] = torch.sin(position * div_term)\n pe[0, :, 1::2] = torch.cos(position * div_term)\n self.register_buffer(\"pe\", pe)\n\n def forward(self, x):\n _, seq_length, _ = x.shape\n x = x + self.pe[:, :seq_length]\n return x\n\n\nclass ImagePositionalEmbeddings(nn.Module):\n \"\"\"\n Converts latent image classes into vector embeddings. Sums the vector embeddings with positional embeddings for the\n height and width of the latent space.\n\n For more details, see figure 10 of the dall-e paper: https://arxiv.org/abs/2102.12092\n\n For VQ-diffusion:\n\n Output vector embeddings are used as input for the transformer.\n\n Note that the vector embeddings for the transformer are different than the vector embeddings from the VQVAE.\n\n Args:\n num_embed (`int`):\n Number of embeddings for the latent pixels embeddings.\n height (`int`):\n Height of the latent image i.e. the number of height embeddings.\n width (`int`):\n Width of the latent image i.e. the number of width embeddings.\n embed_dim (`int`):\n Dimension of the produced vector embeddings. Used for the latent pixel, height, and width embeddings.\n \"\"\"\n\n def __init__(\n self,\n num_embed: int,\n height: int,\n width: int,\n embed_dim: int,\n ):\n super().__init__()\n\n self.height = height\n self.width = width\n self.num_embed = num_embed\n self.embed_dim = embed_dim\n\n self.emb = nn.Embedding(self.num_embed, embed_dim)\n self.height_emb = nn.Embedding(self.height, embed_dim)\n self.width_emb = nn.Embedding(self.width, embed_dim)\n\n def forward(self, index):\n emb = self.emb(index)\n\n height_emb = self.height_emb(torch.arange(self.height, device=index.device).view(1, self.height))\n\n # 1 x H x D -> 1 x H x 1 x D\n height_emb = height_emb.unsqueeze(2)\n\n width_emb = self.width_emb(torch.arange(self.width, device=index.device).view(1, self.width))\n\n # 1 x W x D -> 1 x 1 x W x D\n width_emb = width_emb.unsqueeze(1)\n\n pos_emb = height_emb + width_emb\n\n # 1 x H x W x D -> 1 x L xD\n pos_emb = pos_emb.view(1, self.height * self.width, -1)\n\n emb = emb + pos_emb[:, : emb.shape[1], :]\n\n return emb\n\n\nclass LabelEmbedding(nn.Module):\n \"\"\"\n Embeds class labels into vector representations. Also handles label dropout for classifier-free guidance.\n\n Args:\n num_classes (`int`): The number of classes.\n hidden_size (`int`): The size of the vector embeddings.\n dropout_prob (`float`): The probability of dropping a label.\n \"\"\"\n\n def __init__(self, num_classes, hidden_size, dropout_prob):\n super().__init__()\n use_cfg_embedding = dropout_prob > 0\n self.embedding_table = nn.Embedding(num_classes + use_cfg_embedding, hidden_size)\n self.num_classes = num_classes\n self.dropout_prob = dropout_prob\n\n def token_drop(self, labels, force_drop_ids=None):\n \"\"\"\n Drops labels to enable classifier-free guidance.\n \"\"\"\n if force_drop_ids is None:\n drop_ids = torch.rand(labels.shape[0], device=labels.device) < self.dropout_prob\n else:\n drop_ids = torch.tensor(force_drop_ids == 1)\n labels = torch.where(drop_ids, self.num_classes, labels)\n return labels\n\n def forward(self, labels: torch.LongTensor, force_drop_ids=None):\n use_dropout = self.dropout_prob > 0\n if (self.training and use_dropout) or (force_drop_ids is not None):\n labels = self.token_drop(labels, force_drop_ids)\n embeddings = self.embedding_table(labels)\n return embeddings\n\n\nclass TextImageProjection(nn.Module):\n def __init__(\n self,\n text_embed_dim: int = 1024,\n image_embed_dim: int = 768,\n cross_attention_dim: int = 768,\n num_image_text_embeds: int = 10,\n ):\n super().__init__()\n\n self.num_image_text_embeds = num_image_text_embeds\n self.image_embeds = nn.Linear(image_embed_dim, self.num_image_text_embeds * cross_attention_dim)\n self.text_proj = nn.Linear(text_embed_dim, cross_attention_dim)\n\n def forward(self, text_embeds: torch.Tensor, image_embeds: torch.Tensor):\n batch_size = text_embeds.shape[0]\n\n # image\n image_text_embeds = self.image_embeds(image_embeds)\n image_text_embeds = image_text_embeds.reshape(batch_size, self.num_image_text_embeds, -1)\n\n # text\n text_embeds = self.text_proj(text_embeds)\n\n return torch.cat([image_text_embeds, text_embeds], dim=1)\n\n\nclass ImageProjection(nn.Module):\n def __init__(\n self,\n image_embed_dim: int = 768,\n cross_attention_dim: int = 768,\n num_image_text_embeds: int = 32,\n ):\n super().__init__()\n\n self.num_image_text_embeds = num_image_text_embeds\n self.image_embeds = nn.Linear(image_embed_dim, self.num_image_text_embeds * cross_attention_dim)\n self.norm = nn.LayerNorm(cross_attention_dim)\n\n def forward(self, image_embeds: torch.Tensor):\n batch_size = image_embeds.shape[0]\n\n # image\n image_embeds = self.image_embeds(image_embeds)\n image_embeds = image_embeds.reshape(batch_size, self.num_image_text_embeds, -1)\n image_embeds = self.norm(image_embeds)\n return image_embeds\n\n\nclass IPAdapterFullImageProjection(nn.Module):\n def __init__(self, image_embed_dim=1024, cross_attention_dim=1024):\n super().__init__()\n from .attention import FeedForward\n\n self.ff = FeedForward(image_embed_dim, cross_attention_dim, mult=1, activation_fn=\"gelu\")\n self.norm = nn.LayerNorm(cross_attention_dim)\n\n def forward(self, image_embeds: torch.Tensor):\n return self.norm(self.ff(image_embeds))\n\n\nclass IPAdapterFaceIDImageProjection(nn.Module):\n def __init__(self, image_embed_dim=1024, cross_attention_dim=1024, mult=1, num_tokens=1):\n super().__init__()\n from .attention import FeedForward\n\n self.num_tokens = num_tokens\n self.cross_attention_dim = cross_attention_dim\n self.ff = FeedForward(image_embed_dim, cross_attention_dim * num_tokens, mult=mult, activation_fn=\"gelu\")\n self.norm = nn.LayerNorm(cross_attention_dim)\n\n def forward(self, image_embeds: torch.Tensor):\n x = self.ff(image_embeds)\n x = x.reshape(-1, self.num_tokens, self.cross_attention_dim)\n return self.norm(x)\n\n\nclass CombinedTimestepLabelEmbeddings(nn.Module):\n def __init__(self, num_classes, embedding_dim, class_dropout_prob=0.1):\n super().__init__()\n\n self.time_proj = Timesteps(num_channels=256, flip_sin_to_cos=True, downscale_freq_shift=1)\n self.timestep_embedder = TimestepEmbedding(in_channels=256, time_embed_dim=embedding_dim)\n self.class_embedder = LabelEmbedding(num_classes, embedding_dim, class_dropout_prob)\n\n def forward(self, timestep, class_labels, hidden_dtype=None):\n timesteps_proj = self.time_proj(timestep)\n timesteps_emb = self.timestep_embedder(timesteps_proj.to(dtype=hidden_dtype)) # (N, D)\n\n class_labels = self.class_embedder(class_labels) # (N, D)\n\n conditioning = timesteps_emb + class_labels # (N, D)\n\n return conditioning\n\n\nclass CombinedTimestepTextProjEmbeddings(nn.Module):\n def __init__(self, embedding_dim, pooled_projection_dim):\n super().__init__()\n\n self.time_proj = Timesteps(num_channels=256, flip_sin_to_cos=True, downscale_freq_shift=0)\n self.timestep_embedder = TimestepEmbedding(in_channels=256, time_embed_dim=embedding_dim)\n self.text_embedder = PixArtAlphaTextProjection(pooled_projection_dim, embedding_dim, act_fn=\"silu\")\n\n def forward(self, timestep, pooled_projection):\n timesteps_proj = self.time_proj(timestep)\n timesteps_emb = self.timestep_embedder(timesteps_proj.to(dtype=pooled_projection.dtype)) # (N, D)\n\n pooled_projections = self.text_embedder(pooled_projection)\n\n conditioning = timesteps_emb + pooled_projections\n\n return conditioning\n\n\nclass CombinedTimestepGuidanceTextProjEmbeddings(nn.Module):\n def __init__(self, embedding_dim, pooled_projection_dim):\n super().__init__()\n\n self.time_proj = Timesteps(num_channels=256, flip_sin_to_cos=True, downscale_freq_shift=0)\n self.timestep_embedder = TimestepEmbedding(in_channels=256, time_embed_dim=embedding_dim)\n self.guidance_embedder = TimestepEmbedding(in_channels=256, time_embed_dim=embedding_dim)\n self.text_embedder = PixArtAlphaTextProjection(pooled_projection_dim, embedding_dim, act_fn=\"silu\")\n\n def forward(self, timestep, guidance, pooled_projection):\n timesteps_proj = self.time_proj(timestep)\n timesteps_emb = self.timestep_embedder(timesteps_proj.to(dtype=pooled_projection.dtype)) # (N, D)\n\n guidance_proj = self.time_proj(guidance)\n guidance_emb = self.guidance_embedder(guidance_proj.to(dtype=pooled_projection.dtype)) # (N, D)\n\n time_guidance_emb = timesteps_emb + guidance_emb\n\n pooled_projections = self.text_embedder(pooled_projection)\n conditioning = time_guidance_emb + pooled_projections\n\n return conditioning\n\n\nclass CogView3CombinedTimestepSizeEmbeddings(nn.Module):\n def __init__(self, embedding_dim: int, condition_dim: int, pooled_projection_dim: int, timesteps_dim: int = 256):\n super().__init__()\n\n self.time_proj = Timesteps(num_channels=timesteps_dim, flip_sin_to_cos=True, downscale_freq_shift=0)\n self.condition_proj = Timesteps(num_channels=condition_dim, flip_sin_to_cos=True, downscale_freq_shift=0)\n self.timestep_embedder = TimestepEmbedding(in_channels=timesteps_dim, time_embed_dim=embedding_dim)\n self.condition_embedder = PixArtAlphaTextProjection(pooled_projection_dim, embedding_dim, act_fn=\"silu\")\n\n def forward(\n self,\n timestep: torch.Tensor,\n original_size: torch.Tensor,\n target_size: torch.Tensor,\n crop_coords: torch.Tensor,\n hidden_dtype: torch.dtype,\n ) -> torch.Tensor:\n timesteps_proj = self.time_proj(timestep)\n\n original_size_proj = self.condition_proj(original_size.flatten()).view(original_size.size(0), -1)\n crop_coords_proj = self.condition_proj(crop_coords.flatten()).view(crop_coords.size(0), -1)\n target_size_proj = self.condition_proj(target_size.flatten()).view(target_size.size(0), -1)\n\n # (B, 3 * condition_dim)\n condition_proj = torch.cat([original_size_proj, crop_coords_proj, target_size_proj], dim=1)\n\n timesteps_emb = self.timestep_embedder(timesteps_proj.to(dtype=hidden_dtype)) # (B, embedding_dim)\n condition_emb = self.condition_embedder(condition_proj.to(dtype=hidden_dtype)) # (B, embedding_dim)\n\n conditioning = timesteps_emb + condition_emb\n return conditioning\n\n\nclass HunyuanDiTAttentionPool(nn.Module):\n # Copied from https://github.com/Tencent/HunyuanDiT/blob/cb709308d92e6c7e8d59d0dff41b74d35088db6a/hydit/modules/poolers.py#L6\n\n def __init__(self, spacial_dim: int, embed_dim: int, num_heads: int, output_dim: int = None):\n super().__init__()\n self.positional_embedding = nn.Parameter(torch.randn(spacial_dim + 1, embed_dim) / embed_dim**0.5)\n self.k_proj = nn.Linear(embed_dim, embed_dim)\n self.q_proj = nn.Linear(embed_dim, embed_dim)\n self.v_proj = nn.Linear(embed_dim, embed_dim)\n self.c_proj = nn.Linear(embed_dim, output_dim or embed_dim)\n self.num_heads = num_heads\n\n def forward(self, x):\n x = x.permute(1, 0, 2) # NLC -> LNC\n x = torch.cat([x.mean(dim=0, keepdim=True), x], dim=0) # (L+1)NC\n x = x + self.positional_embedding[:, None, :].to(x.dtype) # (L+1)NC\n x, _ = F.multi_head_attention_forward(\n query=x[:1],\n key=x,\n value=x,\n embed_dim_to_check=x.shape[-1],\n num_heads=self.num_heads,\n q_proj_weight=self.q_proj.weight,\n k_proj_weight=self.k_proj.weight,\n v_proj_weight=self.v_proj.weight,\n in_proj_weight=None,\n in_proj_bias=torch.cat([self.q_proj.bias, self.k_proj.bias, self.v_proj.bias]),\n bias_k=None,\n bias_v=None,\n add_zero_attn=False,\n dropout_p=0,\n out_proj_weight=self.c_proj.weight,\n out_proj_bias=self.c_proj.bias,\n use_separate_proj_weight=True,\n training=self.training,\n need_weights=False,\n )\n return x.squeeze(0)\n\n\nclass HunyuanCombinedTimestepTextSizeStyleEmbedding(nn.Module):\n def __init__(\n self,\n embedding_dim,\n pooled_projection_dim=1024,\n seq_len=256,\n cross_attention_dim=2048,\n use_style_cond_and_image_meta_size=True,\n ):\n super().__init__()\n\n self.time_proj = Timesteps(num_channels=256, flip_sin_to_cos=True, downscale_freq_shift=0)\n self.timestep_embedder = TimestepEmbedding(in_channels=256, time_embed_dim=embedding_dim)\n\n self.size_proj = Timesteps(num_channels=256, flip_sin_to_cos=True, downscale_freq_shift=0)\n\n self.pooler = HunyuanDiTAttentionPool(\n seq_len, cross_attention_dim, num_heads=8, output_dim=pooled_projection_dim\n )\n\n # Here we use a default learned embedder layer for future extension.\n self.use_style_cond_and_image_meta_size = use_style_cond_and_image_meta_size\n if use_style_cond_and_image_meta_size:\n self.style_embedder = nn.Embedding(1, embedding_dim)\n extra_in_dim = 256 * 6 + embedding_dim + pooled_projection_dim\n else:\n extra_in_dim = pooled_projection_dim\n\n self.extra_embedder = PixArtAlphaTextProjection(\n in_features=extra_in_dim,\n hidden_size=embedding_dim * 4,\n out_features=embedding_dim,\n act_fn=\"silu_fp32\",\n )\n\n def forward(self, timestep, encoder_hidden_states, image_meta_size, style, hidden_dtype=None):\n timesteps_proj = self.time_proj(timestep)\n timesteps_emb = self.timestep_embedder(timesteps_proj.to(dtype=hidden_dtype)) # (N, 256)\n\n # extra condition1: text\n pooled_projections = self.pooler(encoder_hidden_states) # (N, 1024)\n\n if self.use_style_cond_and_image_meta_size:\n # extra condition2: image meta size embedding\n image_meta_size = self.size_proj(image_meta_size.view(-1))\n image_meta_size = image_meta_size.to(dtype=hidden_dtype)\n image_meta_size = image_meta_size.view(-1, 6 * 256) # (N, 1536)\n\n # extra condition3: style embedding\n style_embedding = self.style_embedder(style) # (N, embedding_dim)\n\n # Concatenate all extra vectors\n extra_cond = torch.cat([pooled_projections, image_meta_size, style_embedding], dim=1)\n else:\n extra_cond = torch.cat([pooled_projections], dim=1)\n\n conditioning = timesteps_emb + self.extra_embedder(extra_cond) # [B, D]\n\n return conditioning\n\n\nclass LuminaCombinedTimestepCaptionEmbedding(nn.Module):\n def __init__(self, hidden_size=4096, cross_attention_dim=2048, frequency_embedding_size=256):\n super().__init__()\n self.time_proj = Timesteps(\n num_channels=frequency_embedding_size, flip_sin_to_cos=True, downscale_freq_shift=0.0\n )\n\n self.timestep_embedder = TimestepEmbedding(in_channels=frequency_embedding_size, time_embed_dim=hidden_size)\n\n self.caption_embedder = nn.Sequential(\n nn.LayerNorm(cross_attention_dim),\n nn.Linear(\n cross_attention_dim,\n hidden_size,\n bias=True,\n ),\n )\n\n def forward(self, timestep, caption_feat, caption_mask):\n # timestep embedding:\n time_freq = self.time_proj(timestep)\n time_embed = self.timestep_embedder(time_freq.to(dtype=self.timestep_embedder.linear_1.weight.dtype))\n\n # caption condition embedding:\n caption_mask_float = caption_mask.float().unsqueeze(-1)\n caption_feats_pool = (caption_feat * caption_mask_float).sum(dim=1) / caption_mask_float.sum(dim=1)\n caption_feats_pool = caption_feats_pool.to(caption_feat)\n caption_embed = self.caption_embedder(caption_feats_pool)\n\n conditioning = time_embed + caption_embed\n\n return conditioning\n\n\nclass TextTimeEmbedding(nn.Module):\n def __init__(self, encoder_dim: int, time_embed_dim: int, num_heads: int = 64):\n super().__init__()\n self.norm1 = nn.LayerNorm(encoder_dim)\n self.pool = AttentionPooling(num_heads, encoder_dim)\n self.proj = nn.Linear(encoder_dim, time_embed_dim)\n self.norm2 = nn.LayerNorm(time_embed_dim)\n\n def forward(self, hidden_states):\n hidden_states = self.norm1(hidden_states)\n hidden_states = self.pool(hidden_states)\n hidden_states = self.proj(hidden_states)\n hidden_states = self.norm2(hidden_states)\n return hidden_states\n\n\nclass TextImageTimeEmbedding(nn.Module):\n def __init__(self, text_embed_dim: int = 768, image_embed_dim: int = 768, time_embed_dim: int = 1536):\n super().__init__()\n self.text_proj = nn.Linear(text_embed_dim, time_embed_dim)\n self.text_norm = nn.LayerNorm(time_embed_dim)\n self.image_proj = nn.Linear(image_embed_dim, time_embed_dim)\n\n def forward(self, text_embeds: torch.Tensor, image_embeds: torch.Tensor):\n # text\n time_text_embeds = self.text_proj(text_embeds)\n time_text_embeds = self.text_norm(time_text_embeds)\n\n # image\n time_image_embeds = self.image_proj(image_embeds)\n\n return time_image_embeds + time_text_embeds\n\n\nclass ImageTimeEmbedding(nn.Module):\n def __init__(self, image_embed_dim: int = 768, time_embed_dim: int = 1536):\n super().__init__()\n self.image_proj = nn.Linear(image_embed_dim, time_embed_dim)\n self.image_norm = nn.LayerNorm(time_embed_dim)\n\n def forward(self, image_embeds: torch.Tensor):\n # image\n time_image_embeds = self.image_proj(image_embeds)\n time_image_embeds = self.image_norm(time_image_embeds)\n return time_image_embeds\n\n\nclass ImageHintTimeEmbedding(nn.Module):\n def __init__(self, image_embed_dim: int = 768, time_embed_dim: int = 1536):\n super().__init__()\n self.image_proj = nn.Linear(image_embed_dim, time_embed_dim)\n self.image_norm = nn.LayerNorm(time_embed_dim)\n self.input_hint_block = nn.Sequential(\n nn.Conv2d(3, 16, 3, padding=1),\n nn.SiLU(),\n nn.Conv2d(16, 16, 3, padding=1),\n nn.SiLU(),\n nn.Conv2d(16, 32, 3, padding=1, stride=2),\n nn.SiLU(),\n nn.Conv2d(32, 32, 3, padding=1),\n nn.SiLU(),\n nn.Conv2d(32, 96, 3, padding=1, stride=2),\n nn.SiLU(),\n nn.Conv2d(96, 96, 3, padding=1),\n nn.SiLU(),\n nn.Conv2d(96, 256, 3, padding=1, stride=2),\n nn.SiLU(),\n nn.Conv2d(256, 4, 3, padding=1),\n )\n\n def forward(self, image_embeds: torch.Tensor, hint: torch.Tensor):\n # image\n time_image_embeds = self.image_proj(image_embeds)\n time_image_embeds = self.image_norm(time_image_embeds)\n hint = self.input_hint_block(hint)\n return time_image_embeds, hint\n\n\nclass AttentionPooling(nn.Module):\n # Copied from https://github.com/deep-floyd/IF/blob/2f91391f27dd3c468bf174be5805b4cc92980c0b/deepfloyd_if/model/nn.py#L54\n\n def __init__(self, num_heads, embed_dim, dtype=None):\n super().__init__()\n self.dtype = dtype\n self.positional_embedding = nn.Parameter(torch.randn(1, embed_dim) / embed_dim**0.5)\n self.k_proj = nn.Linear(embed_dim, embed_dim, dtype=self.dtype)\n self.q_proj = nn.Linear(embed_dim, embed_dim, dtype=self.dtype)\n self.v_proj = nn.Linear(embed_dim, embed_dim, dtype=self.dtype)\n self.num_heads = num_heads\n self.dim_per_head = embed_dim // self.num_heads\n\n def forward(self, x):\n bs, length, width = x.size()\n\n def shape(x):\n # (bs, length, width) --> (bs, length, n_heads, dim_per_head)\n x = x.view(bs, -1, self.num_heads, self.dim_per_head)\n # (bs, length, n_heads, dim_per_head) --> (bs, n_heads, length, dim_per_head)\n x = x.transpose(1, 2)\n # (bs, n_heads, length, dim_per_head) --> (bs*n_heads, length, dim_per_head)\n x = x.reshape(bs * self.num_heads, -1, self.dim_per_head)\n # (bs*n_heads, length, dim_per_head) --> (bs*n_heads, dim_per_head, length)\n x = x.transpose(1, 2)\n return x\n\n class_token = x.mean(dim=1, keepdim=True) + self.positional_embedding.to(x.dtype)\n x = torch.cat([class_token, x], dim=1) # (bs, length+1, width)\n\n # (bs*n_heads, class_token_length, dim_per_head)\n q = shape(self.q_proj(class_token))\n # (bs*n_heads, length+class_token_length, dim_per_head)\n k = shape(self.k_proj(x))\n v = shape(self.v_proj(x))\n\n # (bs*n_heads, class_token_length, length+class_token_length):\n scale = 1 / math.sqrt(math.sqrt(self.dim_per_head))\n weight = torch.einsum(\"bct,bcs->bts\", q * scale, k * scale) # More stable with f16 than dividing afterwards\n weight = torch.softmax(weight.float(), dim=-1).type(weight.dtype)\n\n # (bs*n_heads, dim_per_head, class_token_length)\n a = torch.einsum(\"bts,bcs->bct\", weight, v)\n\n # (bs, length+1, width)\n a = a.reshape(bs, -1, 1).transpose(1, 2)\n\n return a[:, 0, :] # cls_token\n\n\ndef get_fourier_embeds_from_boundingbox(embed_dim, box):\n \"\"\"\n Args:\n embed_dim: int\n box: a 3-D tensor [B x N x 4] representing the bounding boxes for GLIGEN pipeline\n Returns:\n [B x N x embed_dim] tensor of positional embeddings\n \"\"\"\n\n batch_size, num_boxes = box.shape[:2]\n\n emb = 100 ** (torch.arange(embed_dim) / embed_dim)\n emb = emb[None, None, None].to(device=box.device, dtype=box.dtype)\n emb = emb * box.unsqueeze(-1)\n\n emb = torch.stack((emb.sin(), emb.cos()), dim=-1)\n emb = emb.permute(0, 1, 3, 4, 2).reshape(batch_size, num_boxes, embed_dim * 2 * 4)\n\n return emb\n\n\nclass GLIGENTextBoundingboxProjection(nn.Module):\n def __init__(self, positive_len, out_dim, feature_type=\"text-only\", fourier_freqs=8):\n super().__init__()\n self.positive_len = positive_len\n self.out_dim = out_dim\n\n self.fourier_embedder_dim = fourier_freqs\n self.position_dim = fourier_freqs * 2 * 4 # 2: sin/cos, 4: xyxy\n\n if isinstance(out_dim, tuple):\n out_dim = out_dim[0]\n\n if feature_type == \"text-only\":\n self.linears = nn.Sequential(\n nn.Linear(self.positive_len + self.position_dim, 512),\n nn.SiLU(),\n nn.Linear(512, 512),\n nn.SiLU(),\n nn.Linear(512, out_dim),\n )\n self.null_positive_feature = torch.nn.Parameter(torch.zeros([self.positive_len]))\n\n elif feature_type == \"text-image\":\n self.linears_text = nn.Sequential(\n nn.Linear(self.positive_len + self.position_dim, 512),\n nn.SiLU(),\n nn.Linear(512, 512),\n nn.SiLU(),\n nn.Linear(512, out_dim),\n )\n self.linears_image = nn.Sequential(\n nn.Linear(self.positive_len + self.position_dim, 512),\n nn.SiLU(),\n nn.Linear(512, 512),\n nn.SiLU(),\n nn.Linear(512, out_dim),\n )\n self.null_text_feature = torch.nn.Parameter(torch.zeros([self.positive_len]))\n self.null_image_feature = torch.nn.Parameter(torch.zeros([self.positive_len]))\n\n self.null_position_feature = torch.nn.Parameter(torch.zeros([self.position_dim]))\n\n def forward(\n self,\n boxes,\n masks,\n positive_embeddings=None,\n phrases_masks=None,\n image_masks=None,\n phrases_embeddings=None,\n image_embeddings=None,\n ):\n masks = masks.unsqueeze(-1)\n\n # embedding position (it may includes padding as placeholder)\n xyxy_embedding = get_fourier_embeds_from_boundingbox(self.fourier_embedder_dim, boxes) # B*N*4 -> B*N*C\n\n # learnable null embedding\n xyxy_null = self.null_position_feature.view(1, 1, -1)\n\n # replace padding with learnable null embedding\n xyxy_embedding = xyxy_embedding * masks + (1 - masks) * xyxy_null\n\n # positionet with text only information\n if positive_embeddings is not None:\n # learnable null embedding\n positive_null = self.null_positive_feature.view(1, 1, -1)\n\n # replace padding with learnable null embedding\n positive_embeddings = positive_embeddings * masks + (1 - masks) * positive_null\n\n objs = self.linears(torch.cat([positive_embeddings, xyxy_embedding], dim=-1))\n\n # positionet with text and image information\n else:\n phrases_masks = phrases_masks.unsqueeze(-1)\n image_masks = image_masks.unsqueeze(-1)\n\n # learnable null embedding\n text_null = self.null_text_feature.view(1, 1, -1)\n image_null = self.null_image_feature.view(1, 1, -1)\n\n # replace padding with learnable null embedding\n phrases_embeddings = phrases_embeddings * phrases_masks + (1 - phrases_masks) * text_null\n image_embeddings = image_embeddings * image_masks + (1 - image_masks) * image_null\n\n objs_text = self.linears_text(torch.cat([phrases_embeddings, xyxy_embedding], dim=-1))\n objs_image = self.linears_image(torch.cat([image_embeddings, xyxy_embedding], dim=-1))\n objs = torch.cat([objs_text, objs_image], dim=1)\n\n return objs\n\n\nclass PixArtAlphaCombinedTimestepSizeEmbeddings(nn.Module):\n \"\"\"\n For PixArt-Alpha.\n\n Reference:\n https://github.com/PixArt-alpha/PixArt-alpha/blob/0f55e922376d8b797edd44d25d0e7464b260dcab/diffusion/model/nets/PixArtMS.py#L164C9-L168C29\n \"\"\"\n\n def __init__(self, embedding_dim, size_emb_dim, use_additional_conditions: bool = False):\n super().__init__()\n\n self.outdim = size_emb_dim\n self.time_proj = Timesteps(num_channels=256, flip_sin_to_cos=True, downscale_freq_shift=0)\n self.timestep_embedder = TimestepEmbedding(in_channels=256, time_embed_dim=embedding_dim)\n\n self.use_additional_conditions = use_additional_conditions\n if use_additional_conditions:\n self.additional_condition_proj = Timesteps(num_channels=256, flip_sin_to_cos=True, downscale_freq_shift=0)\n self.resolution_embedder = TimestepEmbedding(in_channels=256, time_embed_dim=size_emb_dim)\n self.aspect_ratio_embedder = TimestepEmbedding(in_channels=256, time_embed_dim=size_emb_dim)\n\n def forward(self, timestep, resolution, aspect_ratio, batch_size, hidden_dtype):\n timesteps_proj = self.time_proj(timestep)\n timesteps_emb = self.timestep_embedder(timesteps_proj.to(dtype=hidden_dtype)) # (N, D)\n\n if self.use_additional_conditions:\n resolution_emb = self.additional_condition_proj(resolution.flatten()).to(hidden_dtype)\n resolution_emb = self.resolution_embedder(resolution_emb).reshape(batch_size, -1)\n aspect_ratio_emb = self.additional_condition_proj(aspect_ratio.flatten()).to(hidden_dtype)\n aspect_ratio_emb = self.aspect_ratio_embedder(aspect_ratio_emb).reshape(batch_size, -1)\n conditioning = timesteps_emb + torch.cat([resolution_emb, aspect_ratio_emb], dim=1)\n else:\n conditioning = timesteps_emb\n\n return conditioning\n\n\nclass PixArtAlphaTextProjection(nn.Module):\n \"\"\"\n Projects caption embeddings. Also handles dropout for classifier-free guidance.\n\n Adapted from https://github.com/PixArt-alpha/PixArt-alpha/blob/master/diffusion/model/nets/PixArt_blocks.py\n \"\"\"\n\n def __init__(self, in_features, hidden_size, out_features=None, act_fn=\"gelu_tanh\"):\n super().__init__()\n if out_features is None:\n out_features = hidden_size\n self.linear_1 = nn.Linear(in_features=in_features, out_features=hidden_size, bias=True)\n if act_fn == \"gelu_tanh\":\n self.act_1 = nn.GELU(approximate=\"tanh\")\n elif act_fn == \"silu\":\n self.act_1 = nn.SiLU()\n elif act_fn == \"silu_fp32\":\n self.act_1 = FP32SiLU()\n else:\n raise ValueError(f\"Unknown activation function: {act_fn}\")\n self.linear_2 = nn.Linear(in_features=hidden_size, out_features=out_features, bias=True)\n\n def forward(self, caption):\n hidden_states = self.linear_1(caption)\n hidden_states = self.act_1(hidden_states)\n hidden_states = self.linear_2(hidden_states)\n return hidden_states\n\n\nclass IPAdapterPlusImageProjectionBlock(nn.Module):\n def __init__(\n self,\n embed_dims: int = 768,\n dim_head: int = 64,\n heads: int = 16,\n ffn_ratio: float = 4,\n ) -> None:\n super().__init__()\n from .attention import FeedForward\n\n self.ln0 = nn.LayerNorm(embed_dims)\n self.ln1 = nn.LayerNorm(embed_dims)\n self.attn = Attention(\n query_dim=embed_dims,\n dim_head=dim_head,\n heads=heads,\n out_bias=False,\n )\n self.ff = nn.Sequential(\n nn.LayerNorm(embed_dims),\n FeedForward(embed_dims, embed_dims, activation_fn=\"gelu\", mult=ffn_ratio, bias=False),\n )\n\n def forward(self, x, latents, residual):\n encoder_hidden_states = self.ln0(x)\n latents = self.ln1(latents)\n encoder_hidden_states = torch.cat([encoder_hidden_states, latents], dim=-2)\n latents = self.attn(latents, encoder_hidden_states) + residual\n latents = self.ff(latents) + latents\n return latents\n\n\nclass IPAdapterPlusImageProjection(nn.Module):\n \"\"\"Resampler of IP-Adapter Plus.\n\n Args:\n embed_dims (int): The feature dimension. Defaults to 768. output_dims (int): The number of output channels,\n that is the same\n number of the channels in the `unet.config.cross_attention_dim`. Defaults to 1024.\n hidden_dims (int):\n The number of hidden channels. Defaults to 1280. depth (int): The number of blocks. Defaults\n to 8. dim_head (int): The number of head channels. Defaults to 64. heads (int): Parallel attention heads.\n Defaults to 16. num_queries (int):\n The number of queries. Defaults to 8. ffn_ratio (float): The expansion ratio\n of feedforward network hidden\n layer channels. Defaults to 4.\n \"\"\"\n\n def __init__(\n self,\n embed_dims: int = 768,\n output_dims: int = 1024,\n hidden_dims: int = 1280,\n depth: int = 4,\n dim_head: int = 64,\n heads: int = 16,\n num_queries: int = 8,\n ffn_ratio: float = 4,\n ) -> None:\n super().__init__()\n self.latents = nn.Parameter(torch.randn(1, num_queries, hidden_dims) / hidden_dims**0.5)\n\n self.proj_in = nn.Linear(embed_dims, hidden_dims)\n\n self.proj_out = nn.Linear(hidden_dims, output_dims)\n self.norm_out = nn.LayerNorm(output_dims)\n\n self.layers = nn.ModuleList(\n [IPAdapterPlusImageProjectionBlock(hidden_dims, dim_head, heads, ffn_ratio) for _ in range(depth)]\n )\n\n def forward(self, x: torch.Tensor) -> torch.Tensor:\n \"\"\"Forward pass.\n\n Args:\n x (torch.Tensor): Input Tensor.\n Returns:\n torch.Tensor: Output Tensor.\n \"\"\"\n latents = self.latents.repeat(x.size(0), 1, 1)\n\n x = self.proj_in(x)\n\n for block in self.layers:\n residual = latents\n latents = block(x, latents, residual)\n\n latents = self.proj_out(latents)\n return self.norm_out(latents)\n\n\nclass IPAdapterFaceIDPlusImageProjection(nn.Module):\n \"\"\"FacePerceiverResampler of IP-Adapter Plus.\n\n Args:\n embed_dims (int): The feature dimension. Defaults to 768. output_dims (int): The number of output channels,\n that is the same\n number of the channels in the `unet.config.cross_attention_dim`. Defaults to 1024.\n hidden_dims (int):\n The number of hidden channels. Defaults to 1280. depth (int): The number of blocks. Defaults\n to 8. dim_head (int): The number of head channels. Defaults to 64. heads (int): Parallel attention heads.\n Defaults to 16. num_tokens (int): Number of tokens num_queries (int): The number of queries. Defaults to 8.\n ffn_ratio (float): The expansion ratio of feedforward network hidden\n layer channels. Defaults to 4.\n ffproj_ratio (float): The expansion ratio of feedforward network hidden\n layer channels (for ID embeddings). Defaults to 4.\n \"\"\"\n\n def __init__(\n self,\n embed_dims: int = 768,\n output_dims: int = 768,\n hidden_dims: int = 1280,\n id_embeddings_dim: int = 512,\n depth: int = 4,\n dim_head: int = 64,\n heads: int = 16,\n num_tokens: int = 4,\n num_queries: int = 8,\n ffn_ratio: float = 4,\n ffproj_ratio: int = 2,\n ) -> None:\n super().__init__()\n from .attention import FeedForward\n\n self.num_tokens = num_tokens\n self.embed_dim = embed_dims\n self.clip_embeds = None\n self.shortcut = False\n self.shortcut_scale = 1.0\n\n self.proj = FeedForward(id_embeddings_dim, embed_dims * num_tokens, activation_fn=\"gelu\", mult=ffproj_ratio)\n self.norm = nn.LayerNorm(embed_dims)\n\n self.proj_in = nn.Linear(hidden_dims, embed_dims)\n\n self.proj_out = nn.Linear(embed_dims, output_dims)\n self.norm_out = nn.LayerNorm(output_dims)\n\n self.layers = nn.ModuleList(\n [IPAdapterPlusImageProjectionBlock(embed_dims, dim_head, heads, ffn_ratio) for _ in range(depth)]\n )\n\n def forward(self, id_embeds: torch.Tensor) -> torch.Tensor:\n \"\"\"Forward pass.\n\n Args:\n id_embeds (torch.Tensor): Input Tensor (ID embeds).\n Returns:\n torch.Tensor: Output Tensor.\n \"\"\"\n id_embeds = id_embeds.to(self.clip_embeds.dtype)\n id_embeds = self.proj(id_embeds)\n id_embeds = id_embeds.reshape(-1, self.num_tokens, self.embed_dim)\n id_embeds = self.norm(id_embeds)\n latents = id_embeds\n\n clip_embeds = self.proj_in(self.clip_embeds)\n x = clip_embeds.reshape(-1, clip_embeds.shape[2], clip_embeds.shape[3])\n\n for block in self.layers:\n residual = latents\n latents = block(x, latents, residual)\n\n latents = self.proj_out(latents)\n out = self.norm_out(latents)\n if self.shortcut:\n out = id_embeds + self.shortcut_scale * out\n return out\n\n\nclass MultiIPAdapterImageProjection(nn.Module):\n def __init__(self, IPAdapterImageProjectionLayers: Union[List[nn.Module], Tuple[nn.Module]]):\n super().__init__()\n self.image_projection_layers = nn.ModuleList(IPAdapterImageProjectionLayers)\n\n def forward(self, image_embeds: List[torch.Tensor]):\n projected_image_embeds = []\n\n # currently, we accept `image_embeds` as\n # 1. a tensor (deprecated) with shape [batch_size, embed_dim] or [batch_size, sequence_length, embed_dim]\n # 2. list of `n` tensors where `n` is number of ip-adapters, each tensor can hae shape [batch_size, num_images, embed_dim] or [batch_size, num_images, sequence_length, embed_dim]\n if not isinstance(image_embeds, list):\n deprecation_message = (\n \"You have passed a tensor as `image_embeds`.This is deprecated and will be removed in a future release.\"\n \" Please make sure to update your script to pass `image_embeds` as a list of tensors to suppress this warning.\"\n )\n deprecate(\"image_embeds not a list\", \"1.0.0\", deprecation_message, standard_warn=False)\n image_embeds = [image_embeds.unsqueeze(1)]\n\n if len(image_embeds) != len(self.image_projection_layers):\n raise ValueError(\n f\"image_embeds must have the same length as image_projection_layers, got {len(image_embeds)} and {len(self.image_projection_layers)}\"\n )\n\n for image_embed, image_projection_layer in zip(image_embeds, self.image_projection_layers):\n batch_size, num_images = image_embed.shape[0], image_embed.shape[1]\n image_embed = image_embed.reshape((batch_size * num_images,) + image_embed.shape[2:])\n image_embed = image_projection_layer(image_embed)\n image_embed = image_embed.reshape((batch_size, num_images) + image_embed.shape[1:])\n\n projected_image_embeds.append(image_embed)\n\n return projected_image_embeds\n" }, { "path": "CogVideo/finetune/models/pipeline_cogvideox.py", "content": "# Copyright 2024 The CogVideoX team, Tsinghua University & ZhipuAI and The HuggingFace Team.\n# All rights reserved.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n# http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\nimport inspect\nimport math\nfrom typing import Any, Callable, Dict, List, Optional, Tuple, Union\n\nimport torch\nfrom transformers import T5EncoderModel, T5Tokenizer\n\nfrom diffusers.callbacks import MultiPipelineCallbacks, PipelineCallback\nfrom diffusers.loaders import CogVideoXLoraLoaderMixin\nfrom diffusers.models.embeddings import get_3d_rotary_pos_embed\nfrom diffusers.pipelines.pipeline_utils import DiffusionPipeline\nfrom diffusers.schedulers import CogVideoXDDIMScheduler, CogVideoXDPMScheduler\nfrom diffusers.utils import logging, replace_example_docstring\nfrom diffusers.utils.torch_utils import randn_tensor\nfrom diffusers.video_processor import VideoProcessor\nfrom models.pipeline_output import CogVideoXPipelineOutput\nfrom models.cogvideox_transformer_3d import CogVideoXTransformer3DModel\nfrom einops import rearrange\n\nimport diffusers\nfrom diffusers import AutoencoderKLCogVideoX, CogVideoXDPMScheduler\nfrom diffusers.models.embeddings import get_3d_rotary_pos_embed\nfrom diffusers.optimization import get_scheduler\nfrom diffusers.pipelines.cogvideo.pipeline_cogvideox import get_resize_crop_region_for_grid\n\nlogger = logging.get_logger(__name__) # pylint: disable=invalid-name\n\n\nEXAMPLE_DOC_STRING = \"\"\"\n Examples:\n ```python\n >>> import torch\n >>> from diffusers import CogVideoXPipeline\n >>> from diffusers.utils import export_to_video\n\n >>> # Models: \"THUDM/CogVideoX-2b\" or \"THUDM/CogVideoX-5b\"\n >>> pipe = CogVideoXPipeline.from_pretrained(\"THUDM/CogVideoX-2b\", torch_dtype=torch.float16).to(\"cuda\")\n >>> prompt = (\n ... \"A panda, dressed in a small, red jacket and a tiny hat, sits on a wooden stool in a serene bamboo forest. \"\n ... \"The panda's fluffy paws strum a miniature acoustic guitar, producing soft, melodic tunes. Nearby, a few other \"\n ... \"pandas gather, watching curiously and some clapping in rhythm. Sunlight filters through the tall bamboo, \"\n ... \"casting a gentle glow on the scene. The panda's face is expressive, showing concentration and joy as it plays. \"\n ... \"The background includes a small, flowing stream and vibrant green foliage, enhancing the peaceful and magical \"\n ... \"atmosphere of this unique musical performance.\"\n ... )\n >>> video = pipe(prompt=prompt, guidance_scale=6, num_inference_steps=50).frames[0]\n >>> export_to_video(video, \"output.mp4\", fps=8)\n ```\n\"\"\"\n\n\n# Similar to diffusers.pipelines.hunyuandit.pipeline_hunyuandit.get_resize_crop_region_for_grid\ndef get_resize_crop_region_for_grid(src, tgt_width, tgt_height):\n tw = tgt_width\n th = tgt_height\n h, w = src\n r = h / w\n if r > (th / tw):\n resize_height = th\n resize_width = int(round(th / h * w))\n else:\n resize_width = tw\n resize_height = int(round(tw / w * h))\n\n crop_top = int(round((th - resize_height) / 2.0))\n crop_left = int(round((tw - resize_width) / 2.0))\n\n return (crop_top, crop_left), (crop_top + resize_height, crop_left + resize_width)\n\n\n# Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.retrieve_timesteps\ndef retrieve_timesteps(\n scheduler,\n num_inference_steps: Optional[int] = None,\n device: Optional[Union[str, torch.device]] = None,\n timesteps: Optional[List[int]] = None,\n sigmas: Optional[List[float]] = None,\n **kwargs,\n):\n r\"\"\"\n Calls the scheduler's `set_timesteps` method and retrieves timesteps from the scheduler after the call. Handles\n custom timesteps. Any kwargs will be supplied to `scheduler.set_timesteps`.\n\n Args:\n scheduler (`SchedulerMixin`):\n The scheduler to get timesteps from.\n num_inference_steps (`int`):\n The number of diffusion steps used when generating samples with a pre-trained model. If used, `timesteps`\n must be `None`.\n device (`str` or `torch.device`, *optional*):\n The device to which the timesteps should be moved to. If `None`, the timesteps are not moved.\n timesteps (`List[int]`, *optional*):\n Custom timesteps used to override the timestep spacing strategy of the scheduler. If `timesteps` is passed,\n `num_inference_steps` and `sigmas` must be `None`.\n sigmas (`List[float]`, *optional*):\n Custom sigmas used to override the timestep spacing strategy of the scheduler. If `sigmas` is passed,\n `num_inference_steps` and `timesteps` must be `None`.\n\n Returns:\n `Tuple[torch.Tensor, int]`: A tuple where the first element is the timestep schedule from the scheduler and the\n second element is the number of inference steps.\n \"\"\"\n if timesteps is not None and sigmas is not None:\n raise ValueError(\"Only one of `timesteps` or `sigmas` can be passed. Please choose one to set custom values\")\n if timesteps is not None:\n accepts_timesteps = \"timesteps\" in set(inspect.signature(scheduler.set_timesteps).parameters.keys())\n if not accepts_timesteps:\n raise ValueError(\n f\"The current scheduler class {scheduler.__class__}'s `set_timesteps` does not support custom\"\n f\" timestep schedules. Please check whether you are using the correct scheduler.\"\n )\n scheduler.set_timesteps(timesteps=timesteps, device=device, **kwargs)\n timesteps = scheduler.timesteps\n num_inference_steps = len(timesteps)\n elif sigmas is not None:\n accept_sigmas = \"sigmas\" in set(inspect.signature(scheduler.set_timesteps).parameters.keys())\n if not accept_sigmas:\n raise ValueError(\n f\"The current scheduler class {scheduler.__class__}'s `set_timesteps` does not support custom\"\n f\" sigmas schedules. Please check whether you are using the correct scheduler.\"\n )\n scheduler.set_timesteps(sigmas=sigmas, device=device, **kwargs)\n timesteps = scheduler.timesteps\n num_inference_steps = len(timesteps)\n else:\n scheduler.set_timesteps(num_inference_steps, device=device, **kwargs)\n timesteps = scheduler.timesteps\n return timesteps, num_inference_steps\n\n\nclass CogVideoXPipeline(DiffusionPipeline, CogVideoXLoraLoaderMixin):\n r\"\"\"\n Pipeline for text-to-video generation using CogVideoX.\n\n This model inherits from [`DiffusionPipeline`]. Check the superclass documentation for the generic methods the\n library implements for all the pipelines (such as downloading or saving, running on a particular device, etc.)\n\n Args:\n vae ([`AutoencoderKL`]):\n Variational Auto-Encoder (VAE) Model to encode and decode videos to and from latent representations.\n text_encoder ([`T5EncoderModel`]):\n Frozen text-encoder. CogVideoX uses\n [T5](https://huggingface.co/docs/transformers/model_doc/t5#transformers.T5EncoderModel); specifically the\n [t5-v1_1-xxl](https://huggingface.co/PixArt-alpha/PixArt-alpha/tree/main/t5-v1_1-xxl) variant.\n tokenizer (`T5Tokenizer`):\n Tokenizer of class\n [T5Tokenizer](https://huggingface.co/docs/transformers/model_doc/t5#transformers.T5Tokenizer).\n transformer ([`CogVideoXTransformer3DModel`]):\n A text conditioned `CogVideoXTransformer3DModel` to denoise the encoded video latents.\n scheduler ([`SchedulerMixin`]):\n A scheduler to be used in combination with `transformer` to denoise the encoded video latents.\n \"\"\"\n\n _optional_components = []\n model_cpu_offload_seq = \"text_encoder->transformer->vae\"\n\n _callback_tensor_inputs = [\n \"latents\",\n \"prompt_embeds\",\n \"negative_prompt_embeds\",\n ]\n\n def __init__(\n self,\n tokenizer: T5Tokenizer,\n text_encoder: T5EncoderModel,\n vae: AutoencoderKLCogVideoX,\n transformer: CogVideoXTransformer3DModel,\n scheduler: Union[CogVideoXDDIMScheduler, CogVideoXDPMScheduler],\n ):\n super().__init__()\n\n self.register_modules(\n tokenizer=tokenizer, text_encoder=text_encoder, vae=vae, transformer=transformer, scheduler=scheduler\n )\n self.vae_scale_factor_spatial = (\n 2 ** (len(self.vae.config.block_out_channels) - 1) if hasattr(self, \"vae\") and self.vae is not None else 8\n )\n self.vae_scale_factor_temporal = (\n self.vae.config.temporal_compression_ratio if hasattr(self, \"vae\") and self.vae is not None else 4\n )\n self.vae_scaling_factor_image = (\n self.vae.config.scaling_factor if hasattr(self, \"vae\") and self.vae is not None else 0.7\n )\n\n self.video_processor = VideoProcessor(vae_scale_factor=self.vae_scale_factor_spatial)\n\n def _get_t5_prompt_embeds(\n self,\n prompt: Union[str, List[str]] = None,\n num_videos_per_prompt: int = 1,\n max_sequence_length: int = 226,\n device: Optional[torch.device] = None,\n dtype: Optional[torch.dtype] = None,\n ):\n device = device or self._execution_device\n dtype = dtype or self.text_encoder.dtype\n\n prompt = [prompt] if isinstance(prompt, str) else prompt\n batch_size = len(prompt)\n\n text_inputs = self.tokenizer(\n prompt,\n padding=\"max_length\",\n max_length=max_sequence_length,\n truncation=True,\n add_special_tokens=True,\n return_tensors=\"pt\",\n )\n text_input_ids = text_inputs.input_ids\n untruncated_ids = self.tokenizer(prompt, padding=\"longest\", return_tensors=\"pt\").input_ids\n\n if untruncated_ids.shape[-1] >= text_input_ids.shape[-1] and not torch.equal(text_input_ids, untruncated_ids):\n removed_text = self.tokenizer.batch_decode(untruncated_ids[:, max_sequence_length - 1 : -1])\n logger.warning(\n \"The following part of your input was truncated because `max_sequence_length` is set to \"\n f\" {max_sequence_length} tokens: {removed_text}\"\n )\n\n prompt_embeds = self.text_encoder(text_input_ids.to(device))[0]\n prompt_embeds = prompt_embeds.to(dtype=dtype, device=device)\n\n prompt_attention_mask = text_inputs.attention_mask\n prompt_attention_mask = prompt_attention_mask.to(dtype=dtype, device=device)\n\n # duplicate text embeddings for each generation per prompt, using mps friendly method\n _, seq_len, _ = prompt_embeds.shape\n prompt_embeds = prompt_embeds.repeat(1, num_videos_per_prompt, 1)\n prompt_embeds = prompt_embeds.view(batch_size * num_videos_per_prompt, seq_len, -1)\n\n return prompt_embeds, prompt_attention_mask\n\n def encode_prompt(\n self,\n prompt: Union[str, List[str]],\n negative_prompt: Optional[Union[str, List[str]]] = None,\n do_classifier_free_guidance: bool = True,\n num_videos_per_prompt: int = 1,\n prompt_embeds: Optional[torch.Tensor] = None,\n negative_prompt_embeds: Optional[torch.Tensor] = None,\n max_sequence_length: int = 226,\n device: Optional[torch.device] = None,\n dtype: Optional[torch.dtype] = None,\n ):\n r\"\"\"\n Encodes the prompt into text encoder hidden states.\n\n Args:\n prompt (`str` or `List[str]`, *optional*):\n prompt to be encoded\n negative_prompt (`str` or `List[str]`, *optional*):\n The prompt or prompts not to guide the image generation. If not defined, one has to pass\n `negative_prompt_embeds` instead. Ignored when not using guidance (i.e., ignored if `guidance_scale` is\n less than `1`).\n do_classifier_free_guidance (`bool`, *optional*, defaults to `True`):\n Whether to use classifier free guidance or not.\n num_videos_per_prompt (`int`, *optional*, defaults to 1):\n Number of videos that should be generated per prompt. torch device to place the resulting embeddings on\n prompt_embeds (`torch.Tensor`, *optional*):\n Pre-generated text embeddings. Can be used to easily tweak text inputs, *e.g.* prompt weighting. If not\n provided, text embeddings will be generated from `prompt` input argument.\n negative_prompt_embeds (`torch.Tensor`, *optional*):\n Pre-generated negative text embeddings. Can be used to easily tweak text inputs, *e.g.* prompt\n weighting. If not provided, negative_prompt_embeds will be generated from `negative_prompt` input\n argument.\n device: (`torch.device`, *optional*):\n torch device\n dtype: (`torch.dtype`, *optional*):\n torch dtype\n \"\"\"\n device = device or self._execution_device\n\n prompt = [prompt] if isinstance(prompt, str) else prompt\n if prompt is not None:\n batch_size = len(prompt)\n else:\n batch_size = prompt_embeds.shape[0]\n\n if prompt_embeds is None:\n prompt_embeds, _ = self._get_t5_prompt_embeds(\n prompt=prompt,\n num_videos_per_prompt=num_videos_per_prompt,\n max_sequence_length=max_sequence_length,\n device=device,\n dtype=dtype,\n )\n\n if do_classifier_free_guidance and negative_prompt_embeds is None:\n negative_prompt = negative_prompt or \"\"\n negative_prompt = batch_size * [negative_prompt] if isinstance(negative_prompt, str) else negative_prompt\n\n if prompt is not None and type(prompt) is not type(negative_prompt):\n raise TypeError(\n f\"`negative_prompt` should be the same type to `prompt`, but got {type(negative_prompt)} !=\"\n f\" {type(prompt)}.\"\n )\n elif batch_size != len(negative_prompt):\n raise ValueError(\n f\"`negative_prompt`: {negative_prompt} has batch size {len(negative_prompt)}, but `prompt`:\"\n f\" {prompt} has batch size {batch_size}. Please make sure that passed `negative_prompt` matches\"\n \" the batch size of `prompt`.\"\n )\n\n negative_prompt_embeds, _ = self._get_t5_prompt_embeds(\n prompt=negative_prompt,\n num_videos_per_prompt=num_videos_per_prompt,\n max_sequence_length=max_sequence_length,\n device=device,\n dtype=dtype,\n )\n\n return prompt_embeds, negative_prompt_embeds\n\n def prepare_latents(\n self, batch_size, num_channels_latents, num_frames, height, width, dtype, device, generator, latents=None\n ):\n if isinstance(generator, list) and len(generator) != batch_size:\n raise ValueError(\n f\"You have passed a list of generators of length {len(generator)}, but requested an effective batch\"\n f\" size of {batch_size}. Make sure the batch size matches the length of the generators.\"\n )\n\n shape = (\n batch_size,\n (num_frames - 1) // self.vae_scale_factor_temporal + 1,\n num_channels_latents,\n height // self.vae_scale_factor_spatial,\n width // self.vae_scale_factor_spatial,\n )\n\n if latents is None:\n latents = randn_tensor(shape, generator=generator, device=device, dtype=dtype)\n else:\n latents = latents.to(device)\n\n # scale the initial noise by the standard deviation required by the scheduler\n latents = latents * self.scheduler.init_noise_sigma\n return latents\n\n def decode_latents(self, latents: torch.Tensor) -> torch.Tensor:\n latents = latents.permute(0, 2, 1, 3, 4) # [batch_size, num_channels, num_frames, height, width]\n latents = 1 / self.vae_scaling_factor_image * latents\n\n frames = self.vae.decode(latents).sample\n return frames\n\n # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.StableDiffusionPipeline.prepare_extra_step_kwargs\n def prepare_extra_step_kwargs(self, generator, eta):\n # prepare extra kwargs for the scheduler step, since not all schedulers have the same signature\n # eta (η) is only used with the DDIMScheduler, it will be ignored for other schedulers.\n # eta corresponds to η in DDIM paper: https://arxiv.org/abs/2010.02502\n # and should be between [0, 1]\n\n accepts_eta = \"eta\" in set(inspect.signature(self.scheduler.step).parameters.keys())\n extra_step_kwargs = {}\n if accepts_eta:\n extra_step_kwargs[\"eta\"] = eta\n\n # check if the scheduler accepts generator\n accepts_generator = \"generator\" in set(inspect.signature(self.scheduler.step).parameters.keys())\n if accepts_generator:\n extra_step_kwargs[\"generator\"] = generator\n return extra_step_kwargs\n\n # Copied from diffusers.pipelines.latte.pipeline_latte.LattePipeline.check_inputs\n def check_inputs(\n self,\n prompt,\n height,\n width,\n negative_prompt,\n callback_on_step_end_tensor_inputs,\n prompt_embeds=None,\n negative_prompt_embeds=None,\n ):\n if height % 8 != 0 or width % 8 != 0:\n raise ValueError(f\"`height` and `width` have to be divisible by 8 but are {height} and {width}.\")\n\n if callback_on_step_end_tensor_inputs is not None and not all(\n k in self._callback_tensor_inputs for k in callback_on_step_end_tensor_inputs\n ):\n raise ValueError(\n f\"`callback_on_step_end_tensor_inputs` has to be in {self._callback_tensor_inputs}, but found {[k for k in callback_on_step_end_tensor_inputs if k not in self._callback_tensor_inputs]}\"\n )\n if prompt is not None and prompt_embeds is not None:\n raise ValueError(\n f\"Cannot forward both `prompt`: {prompt} and `prompt_embeds`: {prompt_embeds}. Please make sure to\"\n \" only forward one of the two.\"\n )\n elif prompt is None and prompt_embeds is None:\n raise ValueError(\n \"Provide either `prompt` or `prompt_embeds`. Cannot leave both `prompt` and `prompt_embeds` undefined.\"\n )\n elif prompt is not None and (not isinstance(prompt, str) and not isinstance(prompt, list)):\n raise ValueError(f\"`prompt` has to be of type `str` or `list` but is {type(prompt)}\")\n\n if prompt is not None and negative_prompt_embeds is not None:\n raise ValueError(\n f\"Cannot forward both `prompt`: {prompt} and `negative_prompt_embeds`:\"\n f\" {negative_prompt_embeds}. Please make sure to only forward one of the two.\"\n )\n\n if negative_prompt is not None and negative_prompt_embeds is not None:\n raise ValueError(\n f\"Cannot forward both `negative_prompt`: {negative_prompt} and `negative_prompt_embeds`:\"\n f\" {negative_prompt_embeds}. Please make sure to only forward one of the two.\"\n )\n\n if prompt_embeds is not None and negative_prompt_embeds is not None:\n if prompt_embeds.shape != negative_prompt_embeds.shape:\n raise ValueError(\n \"`prompt_embeds` and `negative_prompt_embeds` must have the same shape when passed directly, but\"\n f\" got: `prompt_embeds` {prompt_embeds.shape} != `negative_prompt_embeds`\"\n f\" {negative_prompt_embeds.shape}.\"\n )\n\n def fuse_qkv_projections(self) -> None:\n r\"\"\"Enables fused QKV projections.\"\"\"\n self.fusing_transformer = True\n self.transformer.fuse_qkv_projections()\n\n def unfuse_qkv_projections(self) -> None:\n r\"\"\"Disable QKV projection fusion if enabled.\"\"\"\n if not self.fusing_transformer:\n logger.warning(\"The Transformer was not initially fused for QKV projections. Doing nothing.\")\n else:\n self.transformer.unfuse_qkv_projections()\n self.fusing_transformer = False\n\n def _prepare_rotary_positional_embeddings(\n self,\n height: int,\n width: int,\n num_frames: int,\n device: torch.device,\n ) -> Tuple[torch.Tensor, torch.Tensor]:\n grid_height = height // (self.vae_scale_factor_spatial * self.transformer.config.patch_size)\n grid_width = width // (self.vae_scale_factor_spatial * self.transformer.config.patch_size)\n base_size_width = 720 // (self.vae_scale_factor_spatial * self.transformer.config.patch_size)\n base_size_height = 480 // (self.vae_scale_factor_spatial * self.transformer.config.patch_size)\n\n grid_crops_coords = get_resize_crop_region_for_grid(\n (grid_height, grid_width), base_size_width, base_size_height\n )\n freqs_cos, freqs_sin = get_3d_rotary_pos_embed(\n embed_dim=self.transformer.config.attention_head_dim,\n crops_coords=grid_crops_coords,\n grid_size=(grid_height, grid_width),\n temporal_size=num_frames,\n )\n\n freqs_cos = freqs_cos.to(device=device)\n freqs_sin = freqs_sin.to(device=device)\n return freqs_cos, freqs_sin\n\n @property\n def guidance_scale(self):\n return self._guidance_scale\n\n @property\n def num_timesteps(self):\n return self._num_timesteps\n\n @property\n def attention_kwargs(self):\n return self._attention_kwargs\n\n @property\n def interrupt(self):\n return self._interrupt\n\n @torch.no_grad()\n @replace_example_docstring(EXAMPLE_DOC_STRING)\n def __call__(\n self,\n prompt: Optional[Union[str, List[str]]] = None,\n negative_prompt: Optional[Union[str, List[str]]] = None,\n prompts_list: List[str] = None,\n pose_embeds: Optional[torch.FloatTensor] = None,\n height: int = 480,\n width: int = 720,\n num_frames: int = 49,\n annealed_sample_step: int = 15,\n num_inference_steps: int = 50,\n timesteps: Optional[List[int]] = None,\n guidance_scale: float = 6,\n use_dynamic_cfg: bool = False,\n num_videos_per_prompt: int = 1,\n eta: float = 0.0,\n generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None,\n latents: Optional[torch.FloatTensor] = None,\n prompt_embeds: Optional[torch.FloatTensor] = None,\n negative_prompt_embeds: Optional[torch.FloatTensor] = None,\n output_type: str = \"pil\",\n return_dict: bool = True,\n attention_kwargs: Optional[Dict[str, Any]] = None,\n callback_on_step_end: Optional[\n Union[Callable[[int, int, Dict], None], PipelineCallback, MultiPipelineCallbacks]\n ] = None,\n callback_on_step_end_tensor_inputs: List[str] = [\"latents\"],\n max_sequence_length: int = 226,\n ) -> Union[CogVideoXPipelineOutput, Tuple]:\n \"\"\"\n Function invoked when calling the pipeline for generation.\n\n Args:\n prompt (`str` or `List[str]`, *optional*):\n The prompt or prompts to guide the image generation. If not defined, one has to pass `prompt_embeds`.\n instead.\n negative_prompt (`str` or `List[str]`, *optional*):\n The prompt or prompts not to guide the image generation. If not defined, one has to pass\n `negative_prompt_embeds` instead. Ignored when not using guidance (i.e., ignored if `guidance_scale` is\n less than `1`).\n height (`int`, *optional*, defaults to self.transformer.config.sample_height * self.vae_scale_factor_spatial):\n The height in pixels of the generated image. This is set to 480 by default for the best results.\n width (`int`, *optional*, defaults to self.transformer.config.sample_height * self.vae_scale_factor_spatial):\n The width in pixels of the generated image. This is set to 720 by default for the best results.\n num_frames (`int`, defaults to `48`):\n Number of frames to generate. Must be divisible by self.vae_scale_factor_temporal. Generated video will\n contain 1 extra frame because CogVideoX is conditioned with (num_seconds * fps + 1) frames where\n num_seconds is 6 and fps is 8. However, since videos can be saved at any fps, the only condition that\n needs to be satisfied is that of divisibility mentioned above.\n annealed_sample_step (`int`, *optional*, defaults to 15):\n The number of annealed sampling steps that inject paired poses and objects motions into the base T2V models. \n In the latter steps, the object injector will be discarded into traditional T2V.\n num_inference_steps (`int`, *optional*, defaults to 50):\n The number of denoising steps. More denoising steps usually lead to a higher quality image at the\n expense of slower inference.\n timesteps (`List[int]`, *optional*):\n Custom timesteps to use for the denoising process with schedulers which support a `timesteps` argument\n in their `set_timesteps` method. If not defined, the default behavior when `num_inference_steps` is\n passed will be used. Must be in descending order.\n guidance_scale (`float`, *optional*, defaults to 7.0):\n Guidance scale as defined in [Classifier-Free Diffusion Guidance](https://arxiv.org/abs/2207.12598).\n `guidance_scale` is defined as `w` of equation 2. of [Imagen\n Paper](https://arxiv.org/pdf/2205.11487.pdf). Guidance scale is enabled by setting `guidance_scale >\n 1`. Higher guidance scale encourages to generate images that are closely linked to the text `prompt`,\n usually at the expense of lower image quality.\n num_videos_per_prompt (`int`, *optional*, defaults to 1):\n The number of videos to generate per prompt.\n generator (`torch.Generator` or `List[torch.Generator]`, *optional*):\n One or a list of [torch generator(s)](https://pytorch.org/docs/stable/generated/torch.Generator.html)\n to make generation deterministic.\n latents (`torch.FloatTensor`, *optional*):\n Pre-generated noisy latents, sampled from a Gaussian distribution, to be used as inputs for image\n generation. Can be used to tweak the same generation with different prompts. If not provided, a latents\n tensor will ge generated by sampling using the supplied random `generator`.\n prompt_embeds (`torch.FloatTensor`, *optional*):\n Pre-generated text embeddings. Can be used to easily tweak text inputs, *e.g.* prompt weighting. If not\n provided, text embeddings will be generated from `prompt` input argument.\n negative_prompt_embeds (`torch.FloatTensor`, *optional*):\n Pre-generated negative text embeddings. Can be used to easily tweak text inputs, *e.g.* prompt\n weighting. If not provided, negative_prompt_embeds will be generated from `negative_prompt` input\n argument.\n output_type (`str`, *optional*, defaults to `\"pil\"`):\n The output format of the generate image. Choose between\n [PIL](https://pillow.readthedocs.io/en/stable/): `PIL.Image.Image` or `np.array`.\n return_dict (`bool`, *optional*, defaults to `True`):\n Whether or not to return a [`~pipelines.stable_diffusion_xl.StableDiffusionXLPipelineOutput`] instead\n of a plain tuple.\n attention_kwargs (`dict`, *optional*):\n A kwargs dictionary that if specified is passed along to the `AttentionProcessor` as defined under\n `self.processor` in\n [diffusers.models.attention_processor](https://github.com/huggingface/diffusers/blob/main/src/diffusers/models/attention_processor.py).\n callback_on_step_end (`Callable`, *optional*):\n A function that calls at the end of each denoising steps during the inference. The function is called\n with the following arguments: `callback_on_step_end(self: DiffusionPipeline, step: int, timestep: int,\n callback_kwargs: Dict)`. `callback_kwargs` will include a list of all tensors as specified by\n `callback_on_step_end_tensor_inputs`.\n callback_on_step_end_tensor_inputs (`List`, *optional*):\n The list of tensor inputs for the `callback_on_step_end` function. The tensors specified in the list\n will be passed as `callback_kwargs` argument. You will only be able to include variables listed in the\n `._callback_tensor_inputs` attribute of your pipeline class.\n max_sequence_length (`int`, defaults to `226`):\n Maximum sequence length in encoded prompt. Must be consistent with\n `self.transformer.config.max_text_seq_length` otherwise may lead to poor results.\n\n Examples:\n\n Returns:\n [`~pipelines.cogvideo.pipeline_cogvideox.CogVideoXPipelineOutput`] or `tuple`:\n [`~pipelines.cogvideo.pipeline_cogvideox.CogVideoXPipelineOutput`] if `return_dict` is True, otherwise a\n `tuple`. When returning a tuple, the first element is a list with the generated images.\n \"\"\"\n\n if num_frames > 49:\n raise ValueError(\n \"The number of frames must be less than 49 for now due to static positional embeddings. This will be updated in the future to remove this limitation.\"\n )\n\n if isinstance(callback_on_step_end, (PipelineCallback, MultiPipelineCallbacks)):\n callback_on_step_end_tensor_inputs = callback_on_step_end.tensor_inputs\n\n num_videos_per_prompt = 1\n\n # 1. Check inputs. Raise error if not correct\n self.check_inputs(\n prompt,\n height,\n width,\n negative_prompt,\n callback_on_step_end_tensor_inputs,\n prompt_embeds,\n negative_prompt_embeds,\n )\n self._guidance_scale = guidance_scale\n self._attention_kwargs = attention_kwargs\n self._interrupt = False\n \n # 2. Default call parameters\n if prompt is not None and isinstance(prompt, str):\n batch_size = 1\n elif prompt is not None and isinstance(prompt, list):\n batch_size = len(prompt)\n else:\n batch_size = prompt_embeds.shape[0]\n\n device = self._execution_device\n\n # here `guidance_scale` is defined analog to the guidance weight `w` of equation (2)\n # of the Imagen paper: https://arxiv.org/pdf/2205.11487.pdf . `guidance_scale = 1`\n # corresponds to doing no classifier free guidance.\n do_classifier_free_guidance = guidance_scale > 1.0\n\n # 3. Encode input prompt\n prompt_embeds, negative_prompt_embeds = self.encode_prompt(\n prompt,\n negative_prompt,\n do_classifier_free_guidance,\n num_videos_per_prompt=num_videos_per_prompt,\n prompt_embeds=prompt_embeds,\n negative_prompt_embeds=negative_prompt_embeds,\n max_sequence_length=max_sequence_length,\n device=device,\n )\n if do_classifier_free_guidance:\n prompt_embeds = torch.cat([negative_prompt_embeds, prompt_embeds], dim=0)\n\n entity_num = len(prompts_list)\n prompt_entities_embeds = []\n prompt_entities_attention_mask = []\n\n for idx in range(entity_num):\n prompt_entity_embeds, prompt_entity_attention_mask = self._get_t5_prompt_embeds(\n prompt=prompts_list[idx],\n num_videos_per_prompt=num_videos_per_prompt,\n max_sequence_length=max_sequence_length,\n device=device,\n )\n prompt_entities_embeds.append(prompt_entity_embeds)\n prompt_entities_attention_mask.append(prompt_entity_attention_mask)\n\n prompt_entities_embeds = rearrange(torch.stack(prompt_entities_embeds), \"n b t d -> b n t d\")\n prompt_entities_attention_mask = rearrange(torch.stack(prompt_entities_attention_mask), \"n b t -> b n t\")\n\n empty_encoder_hidden_states, _ = self._get_t5_prompt_embeds(\n prompt=\"\",\n num_videos_per_prompt=num_videos_per_prompt,\n max_sequence_length=max_sequence_length,\n device=device,\n )\n\n # 4. Prepare timesteps\n timesteps, num_inference_steps = retrieve_timesteps(self.scheduler, num_inference_steps, device, timesteps)\n self._num_timesteps = len(timesteps)\n\n # 5. Prepare latents.\n latent_channels = self.transformer.config.in_channels\n latents = self.prepare_latents(\n batch_size * num_videos_per_prompt,\n latent_channels,\n num_frames,\n height,\n width,\n prompt_embeds.dtype,\n device,\n generator,\n latents,\n )\n \n # 6. Prepare extra step kwargs. TODO: Logic should ideally just be moved out of the pipeline\n extra_step_kwargs = self.prepare_extra_step_kwargs(generator, eta)\n\n # 7. Create rotary embeds if required\n image_rotary_emb = (\n self._prepare_rotary_positional_embeddings(height, width, latents.size(1), device)\n if self.transformer.config.use_rotary_positional_embeddings\n else None\n )\n\n # 8. Denoising loop\n num_warmup_steps = max(len(timesteps) - num_inference_steps * self.scheduler.order, 0)\n\n with self.progress_bar(total=num_inference_steps) as progress_bar:\n # for DPM-solver++\n old_pred_original_sample = None\n for i, t in enumerate(timesteps):\n if self.interrupt:\n continue\n\n latent_model_input = torch.cat([latents] * 2) if do_classifier_free_guidance else latents\n latent_model_input = self.scheduler.scale_model_input(latent_model_input, t)\n\n # broadcast to batch dimension in a way that's compatible with ONNX/Core ML\n timestep = t.expand(latent_model_input.shape[0])\n \n # predict noise model_output w/ annealed sampling\n if i <= annealed_sample_step:\n noise_pred = self.transformer(\n hidden_states=latent_model_input,\n encoder_hidden_states=prompt_embeds,\n empty_encoder_hidden_states=empty_encoder_hidden_states,\n prompt_entities_embeds=prompt_entities_embeds,\n prompt_entities_attention_mask=prompt_entities_attention_mask,\n pose_embeds=pose_embeds,\n timestep=timestep,\n image_rotary_emb=image_rotary_emb,\n attention_kwargs=attention_kwargs,\n return_dict=False,\n )[0]\n else:\n noise_pred = self.transformer_ori(\n hidden_states=latent_model_input,\n encoder_hidden_states=prompt_embeds,\n empty_encoder_hidden_states=empty_encoder_hidden_states,\n prompt_entities_embeds=prompt_entities_embeds,\n prompt_entities_attention_mask=prompt_entities_attention_mask,\n pose_embeds=None,\n timestep=timestep,\n image_rotary_emb=image_rotary_emb,\n attention_kwargs=attention_kwargs,\n return_dict=False,\n )[0]\n \n noise_pred = noise_pred.float()\n\n # perform guidance\n if use_dynamic_cfg:\n self._guidance_scale = 1 + guidance_scale * (\n (1 - math.cos(math.pi * ((num_inference_steps - t.item()) / num_inference_steps) ** 5.0)) / 2\n )\n if do_classifier_free_guidance:\n noise_pred_uncond, noise_pred_text = noise_pred.chunk(2)\n noise_pred = noise_pred_uncond + self.guidance_scale * (noise_pred_text - noise_pred_uncond)\n\n # compute the previous noisy sample x_t -> x_t-1\n if not isinstance(self.scheduler, CogVideoXDPMScheduler):\n latents = self.scheduler.step(noise_pred, t, latents, **extra_step_kwargs, return_dict=False)[0]\n else:\n latents, old_pred_original_sample = self.scheduler.step(\n noise_pred,\n old_pred_original_sample,\n t,\n timesteps[i - 1] if i > 0 else None,\n latents,\n **extra_step_kwargs,\n return_dict=False,\n )\n latents = latents.to(prompt_embeds.dtype)\n\n # call the callback, if provided\n if callback_on_step_end is not None:\n callback_kwargs = {}\n for k in callback_on_step_end_tensor_inputs:\n callback_kwargs[k] = locals()[k]\n callback_outputs = callback_on_step_end(self, i, t, callback_kwargs)\n\n latents = callback_outputs.pop(\"latents\", latents)\n prompt_embeds = callback_outputs.pop(\"prompt_embeds\", prompt_embeds)\n negative_prompt_embeds = callback_outputs.pop(\"negative_prompt_embeds\", negative_prompt_embeds)\n\n if i == len(timesteps) - 1 or ((i + 1) > num_warmup_steps and (i + 1) % self.scheduler.order == 0):\n progress_bar.update()\n\n if not output_type == \"latent\":\n video = self.decode_latents(latents)\n video = self.video_processor.postprocess_video(video=video, output_type=output_type)\n else:\n video = latents\n\n # Offload all models\n self.maybe_free_model_hooks()\n\n if not return_dict:\n return (video,)\n\n return CogVideoXPipelineOutput(frames=video)\n" }, { "path": "CogVideo/finetune/models/pipeline_output.py", "content": "from dataclasses import dataclass\n\nimport torch\n\nfrom diffusers.utils import BaseOutput\n\n\n@dataclass\nclass CogVideoXPipelineOutput(BaseOutput):\n r\"\"\"\n Output class for CogVideo pipelines.\n\n Args:\n frames (`torch.Tensor`, `np.ndarray`, or List[List[PIL.Image.Image]]):\n List of video outputs - It can be a nested list of length `batch_size,` with each sub-list containing\n denoised PIL image sequences of length `num_frames.` It can also be a NumPy array or Torch tensor of shape\n `(batch_size, num_frames, channels, height, width)`.\n \"\"\"\n\n frames: torch.Tensor\n" }, { "path": "CogVideo/finetune/models/utils.py", "content": "import os\nfrom typing import Callable, Dict, List, Optional, Union\n\nimport torch\nfrom huggingface_hub.utils import validate_hf_hub_args\n\nfrom diffusers.utils import (\n USE_PEFT_BACKEND,\n convert_state_dict_to_diffusers,\n convert_state_dict_to_peft,\n convert_unet_state_dict_to_peft,\n deprecate,\n get_adapter_name,\n get_peft_kwargs,\n is_peft_available,\n is_peft_version,\n is_torch_version,\n is_transformers_available,\n is_transformers_version,\n logging,\n scale_lora_layers,\n)\nfrom diffusers.loaders.lora_base import LoraBaseMixin\nfrom diffusers.loaders.lora_conversion_utils import (\n _convert_kohya_flux_lora_to_diffusers,\n _convert_non_diffusers_lora_to_diffusers,\n _convert_xlabs_flux_lora_to_diffusers,\n _maybe_map_sgm_blocks_to_diffusers,\n)\n\n\n_LOW_CPU_MEM_USAGE_DEFAULT_LORA = False\nif is_torch_version(\">=\", \"1.9.0\"):\n if (\n is_peft_available()\n and is_peft_version(\">=\", \"0.13.1\")\n and is_transformers_available()\n and is_transformers_version(\">\", \"4.45.2\")\n ):\n _LOW_CPU_MEM_USAGE_DEFAULT_LORA = True\n\n\nif is_transformers_available():\n from diffusers.models.lora import text_encoder_attn_modules, text_encoder_mlp_modules\n\nlogger = logging.get_logger(__name__)\n\nTEXT_ENCODER_NAME = \"text_encoder\"\nUNET_NAME = \"unet\"\nTRANSFORMER_NAME = \"transformer\"\n\nLORA_WEIGHT_NAME = \"pytorch_lora_weights.bin\"\nLORA_WEIGHT_NAME_SAFE = \"pytorch_lora_weights.safetensors\"\n\ndef lora_state_dict(\n cls,\n pretrained_model_name_or_path_or_dict: Union[str, Dict[str, torch.Tensor]],\n return_alphas: bool = False,\n **kwargs,\n):\n\n # Load the main state dict first which has the LoRA layers for either of\n # transformer and text encoder or both.\n cache_dir = kwargs.pop(\"cache_dir\", None)\n force_download = kwargs.pop(\"force_download\", False)\n proxies = kwargs.pop(\"proxies\", None)\n local_files_only = kwargs.pop(\"local_files_only\", None)\n token = kwargs.pop(\"token\", None)\n revision = kwargs.pop(\"revision\", None)\n subfolder = kwargs.pop(\"subfolder\", None)\n weight_name = kwargs.pop(\"weight_name\", None)\n use_safetensors = kwargs.pop(\"use_safetensors\", None)\n\n allow_pickle = False\n if use_safetensors is None:\n use_safetensors = True\n allow_pickle = True\n\n user_agent = {\n \"file_type\": \"attn_procs_weights\",\n \"framework\": \"pytorch\",\n }\n\n state_dict = cls._fetch_state_dict(\n pretrained_model_name_or_path_or_dict=pretrained_model_name_or_path_or_dict,\n weight_name=weight_name,\n use_safetensors=use_safetensors,\n local_files_only=local_files_only,\n cache_dir=cache_dir,\n force_download=force_download,\n proxies=proxies,\n token=token,\n revision=revision,\n subfolder=subfolder,\n user_agent=user_agent,\n allow_pickle=allow_pickle,\n )\n is_dora_scale_present = any(\"dora_scale\" in k for k in state_dict)\n if is_dora_scale_present:\n warn_msg = \"It seems like you are using a DoRA checkpoint that is not compatible in Diffusers at the moment. So, we are going to filter out the keys associated to 'dora_scale` from the state dict. If you think this is a mistake please open an issue https://github.com/huggingface/diffusers/issues/new.\"\n logger.warning(warn_msg)\n state_dict = {k: v for k, v in state_dict.items() if \"dora_scale\" not in k}\n\n # TODO (sayakpaul): to a follow-up to clean and try to unify the conditions.\n is_kohya = any(\".lora_down.weight\" in k for k in state_dict)\n if is_kohya:\n state_dict = _convert_kohya_flux_lora_to_diffusers(state_dict)\n # Kohya already takes care of scaling the LoRA parameters with alpha.\n return (state_dict, None) if return_alphas else state_dict\n\n is_xlabs = any(\"processor\" in k for k in state_dict)\n if is_xlabs:\n state_dict = _convert_xlabs_flux_lora_to_diffusers(state_dict)\n # xlabs doesn't use `alpha`.\n return (state_dict, None) if return_alphas else state_dict\n\n # For state dicts like\n # https://huggingface.co/TheLastBen/Jon_Snow_Flux_LoRA\n keys = list(state_dict.keys())\n network_alphas = {}\n for k in keys:\n if \"alpha\" in k:\n alpha_value = state_dict.get(k)\n if (torch.is_tensor(alpha_value) and torch.is_floating_point(alpha_value)) or isinstance(\n alpha_value, float\n ):\n network_alphas[k] = state_dict.pop(k)\n else:\n raise ValueError(\n f\"The alpha key ({k}) seems to be incorrect. If you think this error is unexpected, please open as issue.\"\n )\n\n if return_alphas:\n return state_dict, network_alphas\n else:\n return state_dict\n \npretrained_model_name_or_path_or_dict = '/ytech_m2v2_hdd/fuxiao/CogVideo/finetune/cogvideox5b-lora-single-node-r32/checkpoint-2000/pytorch_lora_weights.safetensors'\nstate_dict, network_alphas = lora_state_dict(pretrained_model_name_or_path_or_dict, return_alphas=True, **kwargs)\n\nimport ipdb; ipdb.set_trace()" }, { "path": "CogVideo/finetune/train_cogvideox_injector.py", "content": "\"\"\"\nAdapted from CogVideoX-5B: https://github.com/THUDM/CogVideo by Xiao Fu (CUHK)\n\"\"\"\n\nimport argparse\nimport logging\nimport math\nimport os\nimport shutil\nfrom pathlib import Path\nfrom typing import List, Optional, Tuple, Union\n\nfrom torch import nn\nimport copy\nimport torch\nimport json\nimport numpy as np\nimport random\nimport cv2\nimport decord\nfrom einops import rearrange\nimport transformers\nfrom accelerate import Accelerator\nfrom accelerate.logging import get_logger\nfrom accelerate.utils import DistributedDataParallelKwargs, ProjectConfiguration, set_seed\nfrom huggingface_hub import create_repo, upload_folder\nfrom peft import LoraConfig, get_peft_model_state_dict, set_peft_model_state_dict\nfrom torch.utils.data import DataLoader, Dataset\nfrom torchvision import transforms\nfrom tqdm.auto import tqdm\nfrom transformers import AutoTokenizer, T5EncoderModel, T5Tokenizer\nfrom safetensors.torch import save_file\n\nimport diffusers\nfrom diffusers import AutoencoderKLCogVideoX, CogVideoXDPMScheduler\nfrom models.pipeline_cogvideox import CogVideoXPipeline\nfrom models.cogvideox_transformer_3d import CogVideoXTransformer3DModel\nfrom diffusers.models.embeddings import get_3d_rotary_pos_embed\nfrom diffusers.optimization import get_scheduler\nfrom diffusers.pipelines.cogvideo.pipeline_cogvideox import get_resize_crop_region_for_grid\nfrom diffusers.training_utils import (\n cast_training_params,\n free_memory,\n)\nfrom diffusers.utils import check_min_version, convert_unet_state_dict_to_peft, export_to_video, is_wandb_available\nfrom diffusers.utils.hub_utils import load_or_create_model_card, populate_model_card\nfrom diffusers.utils.torch_utils import is_compiled_module\n\n\nif is_wandb_available():\n import wandb\n\n# Will error if the minimal version of diffusers is not installed. Remove at your own risks.\ncheck_min_version(\"0.31.0.dev0\")\n\nlogger = get_logger(__name__)\n\n\ndef get_args():\n parser = argparse.ArgumentParser(description=\"Simple example of a training script for CogVideoX.\")\n\n # Model information\n parser.add_argument(\n \"--pretrained_model_name_or_path\",\n type=str,\n default=None,\n required=True,\n help=\"Path to pretrained model or model identifier from huggingface.co/models.\",\n )\n parser.add_argument(\n \"--revision\",\n type=str,\n default=None,\n required=False,\n help=\"Revision of pretrained model identifier from huggingface.co/models.\",\n )\n parser.add_argument(\n \"--variant\",\n type=str,\n default=None,\n help=\"Variant of the model files of the pretrained model identifier from huggingface.co/models, 'e.g.' fp16\",\n )\n parser.add_argument(\n \"--cache_dir\",\n type=str,\n default=None,\n help=\"The directory where the downloaded models and datasets will be stored.\",\n )\n\n # Dataset information\n parser.add_argument(\n \"--dataset_name\",\n type=str,\n default=None,\n help=(\n \"The name of the Dataset (from the HuggingFace hub) containing the training data of instance images (could be your own, possibly private,\"\n \" dataset). It can also be a path pointing to a local copy of a dataset in your filesystem,\"\n \" or to a folder containing files that 🤗 Datasets can understand.\"\n ),\n )\n parser.add_argument(\n \"--dataset_config_name\",\n type=str,\n default=None,\n help=\"The config of the Dataset, leave as None if there's only one config.\",\n )\n parser.add_argument(\n \"--instance_data_root\",\n type=str,\n default=None,\n help=(\"A folder containing the training data.\"),\n )\n parser.add_argument(\n \"--id_token\", type=str, default=None, help=\"Identifier token appended to the start of each prompt if provided.\"\n )\n parser.add_argument(\n \"--dataloader_num_workers\",\n type=int,\n default=0,\n help=(\n \"Number of subprocesses to use for data loading. 0 means that the data will be loaded in the main process.\"\n ),\n )\n # Validation\n parser.add_argument(\n \"--validation_prompt\",\n type=str,\n default=None,\n help=\"One or more prompt(s) that is used during validation to verify that the model is learning. Multiple validation prompts should be separated by the '--validation_prompt_seperator' string.\",\n )\n parser.add_argument(\n \"--validation_prompt_separator\",\n type=str,\n default=\":::\",\n help=\"String that separates multiple validation prompts\",\n )\n parser.add_argument(\n \"--num_validation_videos\",\n type=int,\n default=1,\n help=\"Number of videos that should be generated during validation per `validation_prompt`.\",\n )\n parser.add_argument(\n \"--validation_epochs\",\n type=int,\n default=50,\n help=(\n \"Run validation every X epochs. Validation consists of running the prompt `args.validation_prompt` multiple times: `args.num_validation_videos`.\"\n ),\n )\n parser.add_argument(\n \"--guidance_scale\",\n type=float,\n default=6,\n help=\"The guidance scale to use while sampling validation videos.\",\n )\n parser.add_argument(\n \"--use_dynamic_cfg\",\n action=\"store_true\",\n default=False,\n help=\"Whether or not to use the default cosine dynamic guidance schedule when sampling validation videos.\",\n )\n\n # Training information\n parser.add_argument(\"--seed\", type=int, default=None, help=\"A seed for reproducible training.\")\n parser.add_argument(\n \"--lora_path\",\n type=str,\n default=None,\n help=\"Path to pretrained model or model identifier from huggingface.co/models.\",\n )\n parser.add_argument(\n \"--lora_scale\",\n type=float,\n default=1.0,\n help=(\"The scaling factor to scale LoRA weight update.`\"),\n )\n parser.add_argument(\n \"--mixed_precision\",\n type=str,\n default=None,\n choices=[\"no\", \"fp16\", \"bf16\"],\n help=(\n \"Whether to use mixed precision. Choose between fp16 and bf16 (bfloat16). Bf16 requires PyTorch >=\"\n \" 1.10.and an Nvidia Ampere GPU. Default to the value of accelerate config of the current system or the\"\n \" flag passed with the `accelerate.launch` command. Use this argument to override the accelerate config.\"\n ),\n )\n parser.add_argument(\n \"--output_dir\",\n type=str,\n default=\"cogvideox-lora\",\n help=\"The output directory where the model predictions and checkpoints will be written.\",\n )\n parser.add_argument(\n \"--height\",\n type=int,\n default=480,\n help=\"All input videos are resized to this height.\",\n )\n parser.add_argument(\n \"--width\",\n type=int,\n default=720,\n help=\"All input videos are resized to this width.\",\n )\n parser.add_argument(\"--fps\", type=int, default=8, help=\"All input videos will be used at this FPS.\")\n parser.add_argument(\n \"--max_num_frames\", type=int, default=49, help=\"All input videos will be truncated to these many frames.\"\n )\n parser.add_argument(\n \"--skip_frames_start\",\n type=int,\n default=0,\n help=\"Number of frames to skip from the beginning of each input video. Useful if training data contains intro sequences.\",\n )\n parser.add_argument(\n \"--skip_frames_end\",\n type=int,\n default=0,\n help=\"Number of frames to skip from the end of each input video. Useful if training data contains outro sequences.\",\n )\n parser.add_argument(\n \"--random_flip\",\n action=\"store_true\",\n help=\"whether to randomly flip videos horizontally\",\n )\n parser.add_argument(\n \"--train_batch_size\", type=int, default=4, help=\"Batch size (per device) for the training dataloader.\"\n )\n parser.add_argument(\"--num_train_epochs\", type=int, default=1)\n parser.add_argument(\n \"--max_train_steps\",\n type=int,\n default=None,\n help=\"Total number of training steps to perform. If provided, overrides `--num_train_epochs`.\",\n )\n parser.add_argument(\n \"--checkpointing_steps\",\n type=int,\n default=500,\n help=(\n \"Save a checkpoint of the training state every X updates. These checkpoints can be used both as final\"\n \" checkpoints in case they are better than the last checkpoint, and are also suitable for resuming\"\n \" training using `--resume_from_checkpoint`.\"\n ),\n )\n parser.add_argument(\n \"--checkpoints_total_limit\",\n type=int,\n default=None,\n help=(\"Max number of checkpoints to store.\"),\n )\n parser.add_argument(\n \"--resume_from_checkpoint\",\n type=str,\n default=None,\n help=(\n \"Whether training should be resumed from a previous checkpoint. Use a path saved by\"\n ' `--checkpointing_steps`, or `\"latest\"` to automatically select the last available checkpoint.'\n ),\n )\n parser.add_argument(\n \"--gradient_accumulation_steps\",\n type=int,\n default=1,\n help=\"Number of updates steps to accumulate before performing a backward/update pass.\",\n )\n parser.add_argument(\n \"--dim\",\n type=int,\n default=3072,\n help=\"dimension of each basic transformer block.\",\n )\n parser.add_argument(\n \"--block_interval\",\n type=int,\n default=3,\n help=\"the injector at intervals in transformer blocks to reduce training parameters and improve inference speed.\",\n )\n parser.add_argument(\n \"--gradient_checkpointing\",\n action=\"store_true\",\n help=\"Whether or not to use gradient checkpointing to save memory at the expense of slower backward pass.\",\n )\n parser.add_argument(\n \"--learning_rate\",\n type=float,\n default=1e-4,\n help=\"Initial learning rate (after the potential warmup period) to use.\",\n )\n parser.add_argument(\n \"--scale_lr\",\n action=\"store_true\",\n default=False,\n help=\"Scale the learning rate by the number of GPUs, gradient accumulation steps, and batch size.\",\n )\n parser.add_argument(\n \"--lr_scheduler\",\n type=str,\n default=\"constant\",\n help=(\n 'The scheduler type to use. Choose between [\"linear\", \"cosine\", \"cosine_with_restarts\", \"polynomial\",'\n ' \"constant\", \"constant_with_warmup\"]'\n ),\n )\n parser.add_argument(\n \"--lr_warmup_steps\", type=int, default=500, help=\"Number of steps for the warmup in the lr scheduler.\"\n )\n parser.add_argument(\n \"--lr_num_cycles\",\n type=int,\n default=1,\n help=\"Number of hard resets of the lr in cosine_with_restarts scheduler.\",\n )\n parser.add_argument(\"--lr_power\", type=float, default=1.0, help=\"Power factor of the polynomial scheduler.\")\n parser.add_argument(\n \"--enable_slicing\",\n action=\"store_true\",\n default=False,\n help=\"Whether or not to use VAE slicing for saving memory.\",\n )\n parser.add_argument(\n \"--enable_tiling\",\n action=\"store_true\",\n default=False,\n help=\"Whether or not to use VAE tiling for saving memory.\",\n )\n\n # Optimizer\n parser.add_argument(\n \"--optimizer\",\n type=lambda s: s.lower(),\n default=\"adam\",\n choices=[\"adam\", \"adamw\", \"prodigy\"],\n help=(\"The optimizer type to use.\"),\n )\n parser.add_argument(\n \"--use_8bit_adam\",\n action=\"store_true\",\n help=\"Whether or not to use 8-bit Adam from bitsandbytes. Ignored if optimizer is not set to AdamW\",\n )\n parser.add_argument(\n \"--adam_beta1\", type=float, default=0.9, help=\"The beta1 parameter for the Adam and Prodigy optimizers.\"\n )\n parser.add_argument(\n \"--adam_beta2\", type=float, default=0.95, help=\"The beta2 parameter for the Adam and Prodigy optimizers.\"\n )\n parser.add_argument(\n \"--prodigy_beta3\",\n type=float,\n default=None,\n help=\"Coefficients for computing the Prodigy optimizer's stepsize using running averages. If set to None, uses the value of square root of beta2.\",\n )\n parser.add_argument(\"--prodigy_decouple\", action=\"store_true\", help=\"Use AdamW style decoupled weight decay\")\n parser.add_argument(\"--adam_weight_decay\", type=float, default=1e-04, help=\"Weight decay to use for unet params\")\n parser.add_argument(\n \"--adam_epsilon\",\n type=float,\n default=1e-08,\n help=\"Epsilon value for the Adam optimizer and Prodigy optimizers.\",\n )\n parser.add_argument(\"--max_grad_norm\", default=1.0, type=float, help=\"Max gradient norm.\")\n parser.add_argument(\"--prodigy_use_bias_correction\", action=\"store_true\", help=\"Turn on Adam's bias correction.\")\n parser.add_argument(\n \"--prodigy_safeguard_warmup\",\n action=\"store_true\",\n help=\"Remove lr from the denominator of D estimate to avoid issues during warm-up stage.\",\n )\n parser.add_argument(\n \"--finetune_init\",\n action=\"store_true\",\n help=\"Remove the injector n the first finetune stage w/o loading pretrained ckpt.\",\n )\n\n # Other information\n parser.add_argument(\"--tracker_name\", type=str, default=None, help=\"Project tracker name\")\n parser.add_argument(\"--push_to_hub\", action=\"store_true\", help=\"Whether or not to push the model to the Hub.\")\n parser.add_argument(\"--hub_token\", type=str, default=None, help=\"The token to use to push to the Model Hub.\")\n parser.add_argument(\n \"--hub_model_id\",\n type=str,\n default=None,\n help=\"The name of the repository to keep in sync with the local `output_dir`.\",\n )\n parser.add_argument(\n \"--logging_dir\",\n type=str,\n default=\"logs\",\n help=\"Directory where logs are stored.\",\n )\n parser.add_argument(\n \"--allow_tf32\",\n action=\"store_true\",\n help=(\n \"Whether or not to allow TF32 on Ampere GPUs. Can be used to speed up training. For more information, see\"\n \" https://pytorch.org/docs/stable/notes/cuda.html#tensorfloat-32-tf32-on-ampere-devices\"\n ),\n )\n parser.add_argument(\n \"--report_to\",\n type=str,\n default=None,\n help=(\n 'The integration to report the results and logs to. Supported platforms are `\"tensorboard\"`'\n ' (default), `\"wandb\"` and `\"comet_ml\"`. Use `\"all\"` to report to all integrations.'\n ),\n )\n\n return parser.parse_args()\n\ndef parse_matrix(matrix_str):\n rows = matrix_str.strip().split('] [')\n matrix = []\n for row in rows:\n row = row.replace('[', '').replace(']', '')\n matrix.append(list(map(float, row.split())))\n return np.array(matrix)\n\nclass VideoDataset(Dataset):\n def __init__(\n self,\n instance_data_root: Optional[str] = None,\n dataset_name: Optional[str] = None,\n dataset_config_name: Optional[str] = None,\n height: int = 480,\n width: int = 720,\n fps: int = 8,\n max_num_frames: int = 49,\n skip_frames_start: int = 0,\n skip_frames_end: int = 0,\n cache_dir: Optional[str] = None,\n id_token: Optional[str] = None,\n ) -> None:\n super().__init__()\n\n self.instance_data_root = Path(instance_data_root) if instance_data_root is not None else None\n self.dataset_name = dataset_name\n self.dataset_config_name = dataset_config_name\n self.height = height\n self.width = width\n self.sample_size = (self.height, self.width)\n self.fps = fps\n self.max_num_frames = max_num_frames\n self.skip_frames_start = skip_frames_start\n self.skip_frames_end = skip_frames_end\n self.cache_dir = cache_dir\n self.id_token = id_token or \"\"\n self.pixel_transforms = [\n transforms.Resize(self.sample_size),\n transforms.Normalize(\n mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5], inplace=True\n ),\n ]\n\n self.video_names = []\n\n # ----------------------------------- Released Dataset -----------------------------------\n scenes = ['Desert', 'HDRI']\n scene_location_pair = {\n 'Desert' : 'desert',\n 'HDRI' : \n {\n 'loc1' : 'snowy street',\n 'loc2' : 'park',\n 'loc3' : 'indoor open space',\n 'loc11' : 'gymnastics room',\n 'loc13' : 'autumn forest',\n }\n }\n\n for scene in scenes:\n video_path = os.path.join(self.instance_data_root, '480_720', scene)\n video_names = os.listdir(video_path)\n locations_path = os.path.join(video_path, \"location_data.json\")\n with open(locations_path, 'r') as f: locations = json.load(f)\n locations_info = {locations[idx]['name']:locations[idx] for idx in range(len(locations))}\n for video_name in video_names:\n if video_name.endswith('Hemi12_1') == True:\n if scene != 'HDRI':\n location = scene_location_pair[scene]\n else:\n location = scene_location_pair['HDRI'][video_name.split('_')[1]]\n self.video_names.append((scene, video_name, location, locations_info))\n\n # ----------------------------------- Internal Dataset -----------------------------------\n # scenes = ['AsianTown_480_720', 'Desert_480_720', 'HDRI_480_720', 'Forest_480_720']\n # scene_location_pair = {\n # 'AsianTown_480_720' : 'asian town',\n # 'Desert_480_720' : 'desert',\n # 'Forest_480_720' : 'crossland',\n # 'MatrixCity' : 'city',\n # 'HDRI_480_720' : \n # {\n # 'loc1' : 'snowy street',\n # 'loc2' : 'park',\n # 'loc3' : 'indoor open space',\n # 'loc11' : 'gymnastics room',\n # 'loc13' : 'autumn forest',\n # }\n # }\n \n # for scene in scenes:\n # video_path = os.path.join(self.instance_data_root, scene)\n # video_names = os.listdir(video_path)\n # locations_path = os.path.join(video_path, \"location_data.json\")\n # with open(locations_path, 'r') as f: locations = json.load(f)\n # locations_info = {locations[idx]['name']:locations[idx] for idx in range(len(locations))}\n # for video_name in video_names:\n # if video_name.endswith('Hemi12_1') == True:\n # if scene != 'HDRI_480_720':\n # location = scene_location_pair[scene]\n # else:\n # location = scene_location_pair['HDRI_480_720'][video_name.split('_')[1]]\n # self.video_names.append((scene, video_name, location, locations_info))\n\n self.cam_num = 12\n self.max_objs_num = 3\n self.length = len(self.video_names)\n self.captions_path = os.path.join(self.instance_data_root, \"CharacterInfo.json\")\n with open(self.captions_path, 'r') as f: captions = json.load(f)['CharacterInfo']\n self.captions_info = {int(captions[idx]['index']):captions[idx]['eng'] for idx in range(len(captions))}\n\n self.cams_path = os.path.join(self.instance_data_root, \"Hemi12_transforms.json\")\n with open(self.cams_path, 'r') as f: self.cams_info = json.load(f)\n cam_poses = []\n for i, key in enumerate(self.cams_info.keys()):\n if \"C_\" in key:\n cam_poses.append(parse_matrix(self.cams_info[key]))\n cam_poses = np.stack(cam_poses)\n cam_poses = np.transpose(cam_poses, (0,2,1))\n cam_poses = cam_poses[:,:,[1,2,0,3]]\n cam_poses[:,:3,3] /= 100.\n self.cam_poses = cam_poses\n self.sample_n_frames = 49\n\n def __len__(self):\n return self.length\n \n def save_images2video(self, images, video_name):\n fps = 8\n format = \"mp4\"\n codec = \"libx264\" \n ffmpeg_params = [\"-crf\", str(12)]\n pixelformat = \"yuv420p\" \n video_stream = BytesIO()\n \n with imageio.get_writer(\n video_stream,\n fps=fps,\n format=format,\n codec=codec,\n ffmpeg_params=ffmpeg_params,\n pixelformat=pixelformat,\n ) as writer:\n for idx in range(len(images)):\n writer.append_data(images[idx])\n \n video_data = video_stream.getvalue()\n output_path = os.path.join(video_name + \".mp4\")\n with open(output_path, \"wb\") as f:\n f.write(video_data)\n\n def __getitem__(self, idx):\n \n while True:\n try:\n (scene, video_name, location, locations_info) = self.video_names[idx]\n\n with open(os.path.join(self.instance_data_root, '480_720', scene, video_name, video_name+'.json'), 'r') as f: objs_file = json.load(f)\n objs_num = len(objs_file['0'])\n video_index = random.randint(1, self.cam_num-1)\n\n location_name = video_name.split('_')[1]\n location_info = locations_info[location_name]\n cam_pose = self.cam_poses[video_index-1]\n obj_transl = location_info['coordinates']['CameraTarget']['position']\n\n video_caption_concat = ''\n video_caption_list = []\n obj_poses_list = []\n\n for obj_idx in range(objs_num):\n \n obj_name_index = objs_file['0'][obj_idx]['index'] \n video_caption = self.captions_info[obj_name_index]\n\n if video_caption.startswith(\" \"):\n video_caption = video_caption[1:]\n if video_caption.endswith(\".\"):\n video_caption = video_caption[:-1]\n video_caption = video_caption.lower()\n video_caption_list.append(video_caption)\n \n obj_poses = self.load_sceneposes(objs_file, obj_idx, obj_transl)\n obj_poses = np.linalg.inv(cam_pose) @ obj_poses\n obj_poses_list.append(obj_poses)\n\n for obj_idx in range(objs_num):\n video_caption = video_caption_list[obj_idx]\n if obj_idx == objs_num - 1:\n if objs_num == 1:\n video_caption_concat += video_caption + ' is moving in the ' + location\n else:\n video_caption_concat += video_caption + ' are moving in the ' + location\n else:\n video_caption_concat += video_caption + ' and '\n\n obj_poses_all = torch.from_numpy(np.array(obj_poses_list))\n\n total_frames = 99\n current_sample_stride = 1.75\n cropped_length = int(self.sample_n_frames * current_sample_stride)\n start_frame_ind = random.randint(10, max(10, total_frames - cropped_length - 1))\n end_frame_ind = min(start_frame_ind + cropped_length, total_frames)\n frame_indices = np.linspace(start_frame_ind, end_frame_ind - 1, self.sample_n_frames, dtype=int)\n \n video_frames_path = os.path.join(self.instance_data_root, '480_720', scene, video_name, 'videos', video_name+ f'_C_{video_index:02d}_35mm.mp4')\n cap = cv2.VideoCapture(video_frames_path)\n height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))\n width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))\n \n # get local rank\n ctx = decord.cpu(0)\n reader = decord.VideoReader(video_frames_path, ctx=ctx, height=height, width=width)\n assert len(reader) == total_frames or len(reader) == total_frames+1\n frame_indexes = [frame_idx for frame_idx in range(total_frames)]\n try:\n video_chunk = reader.get_batch(frame_indexes).asnumpy() \n except:\n video_chunk = reader.get_batch(frame_indexes).numpy()\n\n frame_inverse = torch.rand(1).item() > 0.5\n if frame_inverse:\n frame_indices = frame_indices[::-1]\n video_name += '_inv'\n\n pixel_values = np.array([video_chunk[indice] for indice in frame_indices])\n pixel_values = rearrange(torch.from_numpy(pixel_values) / 255.0, \"f h w c -> f c h w\")\n pixel_values = self.pixel_transforms[0](pixel_values)\n pixel_values = self.pixel_transforms[1](pixel_values)\n\n # interpolation\n trunc_frame_indices = np.zeros_like(frame_indices[::4])\n trunc_frame_indices[0] = frame_indices[0]\n trunc_frame_indices[1:] = ((frame_indices[1:][::4] + frame_indices[4:][::4]) / 2).astype(np.int64)\n\n obj_poses_all = obj_poses_all[:, trunc_frame_indices]\n pose_embeds = rearrange(obj_poses_all[:, :, :3, :], \"b f p q -> b f (q p)\").contiguous()\n\n break\n \n except Exception as e:\n \n (scene, video_name, location, locations_info) = self.video_names[idx]\n\n with open(f'invalid_scene.txt', 'a+') as f:\n f.write(f'{scene} {video_name} {location}')\n f.write('\\n')\n\n idx = random.randint(0, self.length - 1)\n\n return {\n \"prompt\": video_caption_concat, \n \"prompts_list\": video_caption_list,\n \"pose_embeds\": pose_embeds,\n \"video\": pixel_values,\n }\n\n\n def load_sceneposes(self, objs_file, obj_idx, obj_transl):\n ext_poses = []\n for i, key in enumerate(objs_file.keys()):\n ext_poses.append(parse_matrix(objs_file[key][obj_idx]['matrix']))\n ext_poses = np.stack(ext_poses)\n ext_poses = np.transpose(ext_poses, (0,2,1))\n ext_poses[:,:3,3] -= obj_transl\n ext_poses[:,:3,3] /= 100.\n ext_poses = ext_poses[:, :, [1,2,0,3]]\n return ext_poses\n\n\ndef save_model_card(\n repo_id: str,\n videos=None,\n base_model: str = None,\n validation_prompt=None,\n repo_folder=None,\n fps=8,\n):\n widget_dict = []\n if videos is not None:\n for i, video in enumerate(videos):\n export_to_video(video, os.path.join(repo_folder, f\"final_video_{i}.mp4\", fps=fps))\n widget_dict.append(\n {\"text\": validation_prompt if validation_prompt else \" \", \"output\": {\"url\": f\"video_{i}.mp4\"}}\n )\n\n model_description = f\"\"\"\n# CogVideoX LoRA - {repo_id}\n\n\n\n## Model description\n\nThese are {repo_id} LoRA weights for {base_model}.\n\nThe weights were trained using the [CogVideoX Diffusers trainer](https://github.com/huggingface/diffusers/blob/main/examples/cogvideo/train_cogvideox_lora.py).\n\nWas LoRA for the text encoder enabled? No.\n\n## Download model\n\n[Download the *.safetensors LoRA]({repo_id}/tree/main) in the Files & versions tab.\n\n## Use it with the [🧨 diffusers library](https://github.com/huggingface/diffusers)\n\n```py\nfrom diffusers import CogVideoXPipeline\nimport torch\n\npipe = CogVideoXPipeline.from_pretrained(\"THUDM/CogVideoX-5b\", torch_dtype=torch.bfloat16).to(\"cuda\")\npipe.load_lora_weights(\"{repo_id}\", weight_name=\"pytorch_lora_weights.safetensors\", adapter_name=[\"cogvideox-lora\"])\n\n# The LoRA adapter weights are determined by what was used for training.\n# In this case, we assume `--lora_alpha` is 32 and `--rank` is 64.\n# It can be made lower or higher from what was used in training to decrease or amplify the effect\n# of the LoRA upto a tolerance, beyond which one might notice no effect at all or overflows.\npipe.set_adapters([\"cogvideox-lora\"], [32 / 64])\n\nvideo = pipe(\"{validation_prompt}\", guidance_scale=6, use_dynamic_cfg=True).frames[0]\n```\n\nFor more details, including weighting, merging and fusing LoRAs, check the [documentation on loading LoRAs in diffusers](https://huggingface.co/docs/diffusers/main/en/using-diffusers/loading_adapters)\n\n## License\n\nPlease adhere to the licensing terms as described [here](https://huggingface.co/THUDM/CogVideoX-5b/blob/main/LICENSE) and [here](https://huggingface.co/THUDM/CogVideoX-2b/blob/main/LICENSE).\n\"\"\"\n model_card = load_or_create_model_card(\n repo_id_or_path=repo_id,\n from_training=True,\n license=\"other\",\n base_model=base_model,\n prompt=validation_prompt,\n model_description=model_description,\n widget=widget_dict,\n )\n tags = [\n \"text-to-video\",\n \"diffusers-training\",\n \"diffusers\",\n \"lora\",\n \"cogvideox\",\n \"cogvideox-diffusers\",\n \"template:sd-lora\",\n ]\n\n model_card = populate_model_card(model_card, tags=tags)\n model_card.save(os.path.join(repo_folder, \"README.md\"))\n\n\ndef log_validation(\n pipe,\n args,\n accelerator,\n pipeline_args,\n epoch,\n is_final_validation: bool = False,\n):\n logger.info(\n f\"Running validation... \\n Generating {args.num_validation_videos} videos with prompt: {pipeline_args['prompt']}.\"\n )\n # We train on the simplified learning objective. If we were previously predicting a variance, we need the scheduler to ignore it\n scheduler_args = {}\n\n if \"variance_type\" in pipe.scheduler.config:\n variance_type = pipe.scheduler.config.variance_type\n\n if variance_type in [\"learned\", \"learned_range\"]:\n variance_type = \"fixed_small\"\n\n scheduler_args[\"variance_type\"] = variance_type\n\n pipe.scheduler = CogVideoXDPMScheduler.from_config(pipe.scheduler.config, **scheduler_args)\n pipe = pipe.to(accelerator.device)\n # pipe.set_progress_bar_config(disable=True)\n\n # run inference\n generator = torch.Generator(device=accelerator.device).manual_seed(args.seed) if args.seed else None\n\n videos = []\n for _ in range(args.num_validation_videos):\n video = pipe(**pipeline_args, generator=generator, output_type=\"np\").frames[0]\n videos.append(video)\n\n for tracker in accelerator.trackers:\n phase_name = \"test\" if is_final_validation else \"validation\"\n if tracker.name == \"wandb\":\n video_filenames = []\n for i, video in enumerate(videos):\n prompt = (\n pipeline_args[\"prompt\"][:25]\n .replace(\" \", \"_\")\n .replace(\" \", \"_\")\n .replace(\"'\", \"_\")\n .replace('\"', \"_\")\n .replace(\"/\", \"_\")\n )\n filename = os.path.join(args.output_dir, f\"{phase_name}_video_{i}_{prompt}.mp4\")\n export_to_video(video, filename, fps=8)\n video_filenames.append(filename)\n\n tracker.log(\n {\n phase_name: [\n wandb.Video(filename, caption=f\"{i}: {pipeline_args['prompt']}\")\n for i, filename in enumerate(video_filenames)\n ]\n }\n )\n\n free_memory()\n\n return videos\n\n\ndef _get_t5_prompt_embeds(\n tokenizer: T5Tokenizer,\n text_encoder: T5EncoderModel,\n prompt: Union[str, List[str]],\n num_videos_per_prompt: int = 1,\n max_sequence_length: int = 226,\n device: Optional[torch.device] = None,\n dtype: Optional[torch.dtype] = None,\n text_input_ids=None,\n):\n prompt = [prompt] if isinstance(prompt, str) else prompt\n batch_size = len(prompt)\n\n if tokenizer is not None:\n text_inputs = tokenizer(\n prompt,\n padding=\"max_length\",\n max_length=max_sequence_length,\n truncation=True,\n add_special_tokens=True,\n return_tensors=\"pt\",\n )\n text_input_ids = text_inputs.input_ids\n else:\n if text_input_ids is None:\n raise ValueError(\"`text_input_ids` must be provided when the tokenizer is not specified.\")\n\n prompt_embeds = text_encoder(text_input_ids.to(device))[0]\n prompt_attention_mask = text_inputs.attention_mask\n prompt_embeds = prompt_embeds.to(dtype=dtype, device=device)\n prompt_attention_mask = prompt_attention_mask.to(dtype=dtype, device=device)\n\n # duplicate text embeddings for each generation per prompt, using mps friendly method\n _, seq_len, _ = prompt_embeds.shape\n prompt_embeds = prompt_embeds.repeat(1, num_videos_per_prompt, 1)\n prompt_embeds = prompt_embeds.view(batch_size * num_videos_per_prompt, seq_len, -1)\n\n return prompt_embeds, prompt_attention_mask\n\n\ndef encode_prompt(\n tokenizer: T5Tokenizer,\n text_encoder: T5EncoderModel,\n prompt: Union[str, List[str]],\n num_videos_per_prompt: int = 1,\n max_sequence_length: int = 226,\n device: Optional[torch.device] = None,\n dtype: Optional[torch.dtype] = None,\n text_input_ids=None,\n):\n prompt = [prompt] if isinstance(prompt, str) else prompt\n prompt_embeds, prompt_attention_mask = _get_t5_prompt_embeds(\n tokenizer,\n text_encoder,\n prompt=prompt,\n num_videos_per_prompt=num_videos_per_prompt,\n max_sequence_length=max_sequence_length,\n device=device,\n dtype=dtype,\n text_input_ids=text_input_ids,\n )\n return prompt_embeds, prompt_attention_mask\n\n\ndef compute_prompt_embeddings(\n tokenizer, text_encoder, prompt, max_sequence_length, device, dtype, requires_grad: bool = False\n):\n if requires_grad:\n prompt_embeds, prompt_attention_mask = encode_prompt(\n tokenizer,\n text_encoder,\n prompt,\n num_videos_per_prompt=1,\n max_sequence_length=max_sequence_length,\n device=device,\n dtype=dtype,\n )\n else:\n with torch.no_grad():\n prompt_embeds, prompt_attention_mask = encode_prompt(\n tokenizer,\n text_encoder,\n prompt,\n num_videos_per_prompt=1,\n max_sequence_length=max_sequence_length,\n device=device,\n dtype=dtype,\n )\n return prompt_embeds, prompt_attention_mask\n\n\ndef prepare_rotary_positional_embeddings(\n height: int,\n width: int,\n num_frames: int,\n vae_scale_factor_spatial: int = 8,\n patch_size: int = 2,\n attention_head_dim: int = 64,\n device: Optional[torch.device] = None,\n base_height: int = 480,\n base_width: int = 720,\n) -> Tuple[torch.Tensor, torch.Tensor]:\n grid_height = height // (vae_scale_factor_spatial * patch_size)\n grid_width = width // (vae_scale_factor_spatial * patch_size)\n base_size_width = base_width // (vae_scale_factor_spatial * patch_size)\n base_size_height = base_height // (vae_scale_factor_spatial * patch_size)\n\n grid_crops_coords = get_resize_crop_region_for_grid((grid_height, grid_width), base_size_width, base_size_height)\n freqs_cos, freqs_sin = get_3d_rotary_pos_embed(\n embed_dim=attention_head_dim,\n crops_coords=grid_crops_coords,\n grid_size=(grid_height, grid_width),\n temporal_size=num_frames,\n )\n\n freqs_cos = freqs_cos.to(device=device)\n freqs_sin = freqs_sin.to(device=device)\n return freqs_cos, freqs_sin\n\n\ndef get_optimizer(args, params_to_optimize, use_deepspeed: bool = False):\n # Use DeepSpeed optimzer\n if use_deepspeed:\n from accelerate.utils import DummyOptim\n\n return DummyOptim(\n params_to_optimize,\n lr=args.learning_rate,\n betas=(args.adam_beta1, args.adam_beta2),\n eps=args.adam_epsilon,\n weight_decay=args.adam_weight_decay,\n )\n\n # Optimizer creation\n supported_optimizers = [\"adam\", \"adamw\", \"prodigy\"]\n if args.optimizer not in supported_optimizers:\n logger.warning(\n f\"Unsupported choice of optimizer: {args.optimizer}. Supported optimizers include {supported_optimizers}. Defaulting to AdamW\"\n )\n args.optimizer = \"adamw\"\n\n if args.use_8bit_adam and not (args.optimizer.lower() not in [\"adam\", \"adamw\"]):\n logger.warning(\n f\"use_8bit_adam is ignored when optimizer is not set to 'Adam' or 'AdamW'. Optimizer was \"\n f\"set to {args.optimizer.lower()}\"\n )\n\n if args.use_8bit_adam:\n try:\n import bitsandbytes as bnb\n except ImportError:\n raise ImportError(\n \"To use 8-bit Adam, please install the bitsandbytes library: `pip install bitsandbytes`.\"\n )\n\n if args.optimizer.lower() == \"adamw\":\n optimizer_class = bnb.optim.AdamW8bit if args.use_8bit_adam else torch.optim.AdamW\n\n optimizer = optimizer_class(\n params_to_optimize,\n betas=(args.adam_beta1, args.adam_beta2),\n eps=args.adam_epsilon,\n weight_decay=args.adam_weight_decay,\n )\n elif args.optimizer.lower() == \"adam\":\n optimizer_class = bnb.optim.Adam8bit if args.use_8bit_adam else torch.optim.Adam\n\n optimizer = optimizer_class(\n params_to_optimize,\n betas=(args.adam_beta1, args.adam_beta2),\n eps=args.adam_epsilon,\n weight_decay=args.adam_weight_decay,\n )\n elif args.optimizer.lower() == \"prodigy\":\n try:\n import prodigyopt\n except ImportError:\n raise ImportError(\"To use Prodigy, please install the prodigyopt library: `pip install prodigyopt`\")\n\n optimizer_class = prodigyopt.Prodigy\n\n if args.learning_rate <= 0.1:\n logger.warning(\n \"Learning rate is too low. When using prodigy, it's generally better to set learning rate around 1.0\"\n )\n\n optimizer = optimizer_class(\n params_to_optimize,\n lr=args.learning_rate,\n betas=(args.adam_beta1, args.adam_beta2),\n beta3=args.prodigy_beta3,\n weight_decay=args.adam_weight_decay,\n eps=args.adam_epsilon,\n decouple=args.prodigy_decouple,\n use_bias_correction=args.prodigy_use_bias_correction,\n safeguard_warmup=args.prodigy_safeguard_warmup,\n )\n\n return optimizer\n\n\ndef main(args):\n if args.report_to == \"wandb\" and args.hub_token is not None:\n raise ValueError(\n \"You cannot use both --report_to=wandb and --hub_token due to a security risk of exposing your token.\"\n \" Please use `huggingface-cli login` to authenticate with the Hub.\"\n )\n\n if torch.backends.mps.is_available() and args.mixed_precision == \"bf16\":\n # due to pytorch#99272, MPS does not yet support bfloat16.\n raise ValueError(\n \"Mixed precision training with bfloat16 is not supported on MPS. Please use fp16 (recommended) or fp32 instead.\"\n )\n\n logging_dir = Path(args.output_dir, args.logging_dir)\n\n accelerator_project_config = ProjectConfiguration(project_dir=args.output_dir, logging_dir=logging_dir)\n kwargs = DistributedDataParallelKwargs(find_unused_parameters=True)\n accelerator = Accelerator(\n gradient_accumulation_steps=args.gradient_accumulation_steps,\n mixed_precision=args.mixed_precision,\n log_with=args.report_to,\n project_config=accelerator_project_config,\n kwargs_handlers=[kwargs],\n )\n\n # Disable AMP for MPS.\n if torch.backends.mps.is_available():\n accelerator.native_amp = False\n\n if args.report_to == \"wandb\":\n if not is_wandb_available():\n raise ImportError(\"Make sure to install wandb if you want to use it for logging during training.\")\n\n # Make one log on every process with the configuration for debugging.\n logging.basicConfig(\n format=\"%(asctime)s - %(levelname)s - %(name)s - %(message)s\",\n datefmt=\"%m/%d/%Y %H:%M:%S\",\n level=logging.INFO,\n )\n logger.info(accelerator.state, main_process_only=False)\n if accelerator.is_local_main_process:\n transformers.utils.logging.set_verbosity_warning()\n diffusers.utils.logging.set_verbosity_info()\n else:\n transformers.utils.logging.set_verbosity_error()\n diffusers.utils.logging.set_verbosity_error()\n\n # If passed along, set the training seed now.\n if args.seed is not None:\n set_seed(args.seed)\n\n # Handle the repository creation\n if accelerator.is_main_process:\n if args.output_dir is not None:\n os.makedirs(args.output_dir, exist_ok=True)\n\n if args.push_to_hub:\n repo_id = create_repo(\n repo_id=args.hub_model_id or Path(args.output_dir).name,\n exist_ok=True,\n ).repo_id\n\n # Prepare models and scheduler\n tokenizer = AutoTokenizer.from_pretrained(\n args.pretrained_model_name_or_path, subfolder=\"tokenizer\", revision=args.revision\n )\n\n text_encoder = T5EncoderModel.from_pretrained(\n args.pretrained_model_name_or_path, subfolder=\"text_encoder\", revision=args.revision\n )\n\n # CogVideoX-2b weights are stored in float16\n # CogVideoX-5b and CogVideoX-5b-I2V weights are stored in bfloat16\n load_dtype = torch.bfloat16 if \"5b\" in args.pretrained_model_name_or_path.lower() else torch.float16\n\n pipe = CogVideoXPipeline.from_pretrained(\n args.pretrained_model_name_or_path,\n torch_dtype=torch.bfloat16,\n finetune_init=args.finetune_init,\n )\n\n if args.resume_from_checkpoint:\n transformer = CogVideoXTransformer3DModel.from_pretrained(\n args.resume_from_checkpoint,\n torch_dtype=load_dtype,\n finetune_init=args.finetune_init,\n )\n else:\n transformer = CogVideoXTransformer3DModel.from_pretrained(\n args.pretrained_model_name_or_path,\n subfolder=\"transformer\",\n torch_dtype=load_dtype,\n revision=args.revision,\n variant=args.variant,\n finetune_init=args.finetune_init,\n )\n dim = args.dim\n for idx in range(len(transformer.transformer_blocks)):\n if idx%args.block_interval == 0:\n transformer.transformer_blocks[idx].attn_injector = copy.deepcopy(transformer.transformer_blocks[idx].attn1)\n transformer.transformer_blocks[idx].pose_fuse_layer = nn.Linear(12, dim)\n transformer.transformer_blocks[idx].attn_null_feature = nn.Parameter(torch.zeros([dim]))\n \n pipe.transformer = transformer\n\n if args.lora_path:\n pipe.load_lora_weights(args.lora_path, weight_name=\"pytorch_lora_weights.safetensors\", adapter_name=\"default\")\n pipe.fuse_lora(components=['transformer'] ,lora_scale=args.lora_scale)\n\n transformer = pipe.transformer\n\n vae = AutoencoderKLCogVideoX.from_pretrained(\n args.pretrained_model_name_or_path, subfolder=\"vae\", revision=args.revision, variant=args.variant\n )\n\n scheduler = CogVideoXDPMScheduler.from_pretrained(args.pretrained_model_name_or_path, subfolder=\"scheduler\")\n\n if args.enable_slicing:\n vae.enable_slicing()\n if args.enable_tiling:\n vae.enable_tiling()\n\n # We only train the additional injector layers\n text_encoder.requires_grad_(False)\n transformer.requires_grad_(False)\n vae.requires_grad_(False)\n\n for idx in range(len(transformer.transformer_blocks)):\n if hasattr(transformer.transformer_blocks[idx], \"attn_injector\"):\n for name, param_module in transformer.transformer_blocks[idx].attn_injector.named_modules():\n for params in param_module.parameters():\n params.requires_grad_(True)\n if hasattr(transformer.transformer_blocks[idx], \"pose_fuse_layer\"):\n for name, param_module in transformer.transformer_blocks[idx].pose_fuse_layer.named_modules():\n for params in param_module.parameters():\n params.requires_grad_(True)\n if hasattr(transformer.transformer_blocks[idx], \"attn_null_feature\"):\n transformer.transformer_blocks[idx].attn_null_feature.requires_grad_(True)\n\n # For mixed precision training we cast all non-trainable weights (vae, text_encoder and transformer) to half-precision\n # as these weights are only used for inference, keeping weights in full precision is not required.\n weight_dtype = torch.float32\n if accelerator.state.deepspeed_plugin:\n # DeepSpeed is handling precision, use what's in the DeepSpeed config\n if (\n \"fp16\" in accelerator.state.deepspeed_plugin.deepspeed_config\n and accelerator.state.deepspeed_plugin.deepspeed_config[\"fp16\"][\"enabled\"]\n ):\n weight_dtype = torch.float16\n if (\n \"bf16\" in accelerator.state.deepspeed_plugin.deepspeed_config\n and accelerator.state.deepspeed_plugin.deepspeed_config[\"bf16\"][\"enabled\"]\n ):\n weight_dtype = torch.float16\n else:\n if accelerator.mixed_precision == \"fp16\":\n weight_dtype = torch.float16\n elif accelerator.mixed_precision == \"bf16\":\n weight_dtype = torch.bfloat16\n\n if torch.backends.mps.is_available() and weight_dtype == torch.bfloat16:\n # due to pytorch#99272, MPS does not yet support bfloat16.\n raise ValueError(\n \"Mixed precision training with bfloat16 is not supported on MPS. Please use fp16 (recommended) or fp32 instead.\"\n )\n\n text_encoder.to(accelerator.device, dtype=weight_dtype)\n transformer.to(accelerator.device, dtype=weight_dtype)\n vae.to(accelerator.device, dtype=weight_dtype)\n\n if args.gradient_checkpointing:\n transformer.enable_gradient_checkpointing()\n\n def unwrap_model(model):\n model = accelerator.unwrap_model(model)\n model = model._orig_mod if is_compiled_module(model) else model\n return model\n\n # create custom saving & loading hooks so that `accelerator.save_state(...)` serializes in a nice format\n def save_model_hook(models, weights, output_dir):\n if accelerator.is_main_process:\n \n for model in models:\n \n if isinstance(model, type(unwrap_model(transformer))):\n \n from diffusers import CogVideoXTransformer3DModel\n model_base = CogVideoXTransformer3DModel.from_pretrained(os.path.join(args.pretrained_model_name_or_path, 'transformer'), torch_dtype=torch.bfloat16)\n model_save = copy.deepcopy(model)\n\n for idx in range(len(model_save.transformer_blocks)):\n if hasattr(model_save.transformer_blocks[idx], 'attn_injector'):\n model_base.transformer_blocks[idx].attn_injector = model_save.transformer_blocks[idx].attn_injector\n model_base.transformer_blocks[idx].pose_fuse_layer = model_save.transformer_blocks[idx].pose_fuse_layer\n model_base.transformer_blocks[idx].attn_null_feature = model_save.transformer_blocks[idx].attn_null_feature\n \n model_base.save_pretrained(os.path.join(output_dir))\n \n else:\n raise ValueError(f\"unexpected save model: {model.__class__}\")\n\n # make sure to pop weight so that corresponding model is not saved again\n weights.pop()\n\n def load_model_hook(models, input_dir):\n transformer_ = None\n\n while len(models) > 0:\n model = models.pop()\n\n if isinstance(model, type(unwrap_model(transformer))):\n transformer_ = model\n else:\n raise ValueError(f\"Unexpected save model: {model.__class__}\")\n\n lora_state_dict = CogVideoXPipeline.lora_state_dict(input_dir)\n\n transformer_state_dict = {\n f'{k.replace(\"transformer.\", \"\")}': v for k, v in lora_state_dict.items() if k.startswith(\"transformer.\")\n }\n transformer_state_dict = convert_unet_state_dict_to_peft(transformer_state_dict)\n incompatible_keys = set_peft_model_state_dict(transformer_, transformer_state_dict, adapter_name=\"default\")\n if incompatible_keys is not None:\n # check only for unexpected keys\n unexpected_keys = getattr(incompatible_keys, \"unexpected_keys\", None)\n if unexpected_keys:\n logger.warning(\n f\"Loading adapter weights from state_dict led to unexpected keys not found in the model: \"\n f\" {unexpected_keys}. \"\n )\n \n # Make sure the trainable params are in float32. This is again needed since the base models\n # are in `weight_dtype`. More details:\n # https://github.com/huggingface/diffusers/pull/6514#discussion_r1449796804\n if args.mixed_precision == \"fp16\":\n # only upcast trainable parameters (LoRA) into fp32\n cast_training_params([transformer_])\n\n accelerator.register_save_state_pre_hook(save_model_hook)\n accelerator.register_load_state_pre_hook(load_model_hook)\n\n # Enable TF32 for faster training on Ampere GPUs,\n # cf https://pytorch.org/docs/stable/notes/cuda.html#tensorfloat-32-tf32-on-ampere-devices\n if args.allow_tf32 and torch.cuda.is_available():\n torch.backends.cuda.matmul.allow_tf32 = True\n\n if args.scale_lr:\n args.learning_rate = (\n args.learning_rate * args.gradient_accumulation_steps * args.train_batch_size * accelerator.num_processes\n )\n\n # Make sure the trainable params are in float32.\n if args.mixed_precision == \"fp16\":\n # only upcast trainable parameters (LoRA) into fp32\n cast_training_params([transformer], dtype=torch.float32)\n\n transformer_trainable_parameters = list(filter(lambda p: p.requires_grad, transformer.parameters()))\n\n # Optimization parameters\n transformer_parameters_with_lr = {\"params\": transformer_trainable_parameters, \"lr\": args.learning_rate}\n params_to_optimize = [transformer_parameters_with_lr]\n\n use_deepspeed_optimizer = (\n accelerator.state.deepspeed_plugin is not None\n and \"optimizer\" in accelerator.state.deepspeed_plugin.deepspeed_config\n )\n use_deepspeed_scheduler = (\n accelerator.state.deepspeed_plugin is not None\n and \"scheduler\" not in accelerator.state.deepspeed_plugin.deepspeed_config\n )\n\n optimizer = get_optimizer(args, params_to_optimize, use_deepspeed=use_deepspeed_optimizer)\n\n # Dataset and DataLoader\n train_dataset = VideoDataset(\n instance_data_root=args.instance_data_root,\n dataset_name=args.dataset_name,\n dataset_config_name=args.dataset_config_name,\n height=args.height,\n width=args.width,\n fps=args.fps,\n max_num_frames=args.max_num_frames,\n skip_frames_start=args.skip_frames_start,\n skip_frames_end=args.skip_frames_end,\n cache_dir=args.cache_dir,\n id_token=args.id_token,\n )\n \n def encode_video(video):\n video = video.to(accelerator.device, dtype=vae.dtype)\n video = video.permute(0, 2, 1, 3, 4) # [B, C, F, H, W]\n latent_dist = vae.encode(video).latent_dist.sample()\n latent_dist = latent_dist * vae.config.scaling_factor\n return latent_dist\n \n train_dataloader = DataLoader(\n train_dataset,\n batch_size=args.train_batch_size,\n shuffle=True,\n num_workers=args.dataloader_num_workers,\n )\n\n # Scheduler and math around the number of training steps.\n overrode_max_train_steps = False\n num_update_steps_per_epoch = math.ceil(len(train_dataloader) / args.gradient_accumulation_steps)\n if args.max_train_steps is None:\n args.max_train_steps = args.num_train_epochs * num_update_steps_per_epoch\n overrode_max_train_steps = True\n\n if use_deepspeed_scheduler:\n from accelerate.utils import DummyScheduler\n\n lr_scheduler = DummyScheduler(\n name=args.lr_scheduler,\n optimizer=optimizer,\n total_num_steps=args.max_train_steps * accelerator.num_processes,\n num_warmup_steps=args.lr_warmup_steps * accelerator.num_processes,\n )\n else:\n lr_scheduler = get_scheduler(\n args.lr_scheduler,\n optimizer=optimizer,\n num_warmup_steps=args.lr_warmup_steps * accelerator.num_processes,\n num_training_steps=args.max_train_steps * accelerator.num_processes,\n num_cycles=args.lr_num_cycles,\n power=args.lr_power,\n )\n\n # Prepare everything with our `accelerator`.\n transformer, optimizer, train_dataloader, lr_scheduler = accelerator.prepare(\n transformer, optimizer, train_dataloader, lr_scheduler\n )\n \n # We need to recalculate our total training steps as the size of the training dataloader may have changed.\n num_update_steps_per_epoch = math.ceil(len(train_dataloader) / args.gradient_accumulation_steps)\n if overrode_max_train_steps:\n args.max_train_steps = args.num_train_epochs * num_update_steps_per_epoch\n # Afterwards we recalculate our number of training epochs\n args.num_train_epochs = math.ceil(args.max_train_steps / num_update_steps_per_epoch)\n\n # We need to initialize the trackers we use, and also store our configuration.\n # The trackers initializes automatically on the main process.\n if accelerator.is_main_process:\n tracker_name = args.tracker_name or \"cogvideox-injector\"\n accelerator.init_trackers(tracker_name, config=vars(args))\n\n # Train!\n total_batch_size = args.train_batch_size * accelerator.num_processes * args.gradient_accumulation_steps\n num_trainable_parameters = sum(param.numel() for model in params_to_optimize for param in model[\"params\"])\n\n logger.info(\"***** Running training *****\")\n logger.info(f\" Num trainable parameters = {num_trainable_parameters}\")\n logger.info(f\" Num examples = {len(train_dataset)}\")\n logger.info(f\" Num batches each epoch = {len(train_dataloader)}\")\n logger.info(f\" Num epochs = {args.num_train_epochs}\")\n logger.info(f\" Instantaneous batch size per device = {args.train_batch_size}\")\n logger.info(f\" Total train batch size (w. parallel, distributed & accumulation) = {total_batch_size}\")\n logger.info(f\" Gradient accumulation steps = {args.gradient_accumulation_steps}\")\n logger.info(f\" Total optimization steps = {args.max_train_steps}\")\n global_step = 0\n first_epoch = 0\n\n # Potentially load in the weights and states from a previous save\n if not args.resume_from_checkpoint:\n initial_global_step = 0\n else:\n if args.resume_from_checkpoint != \"latest\":\n path = os.path.basename(args.resume_from_checkpoint)\n else:\n # Get the mos recent checkpoint\n dirs = os.listdir(args.output_dir)\n dirs = [d for d in dirs if d.startswith(\"checkpoint\")]\n dirs = sorted(dirs, key=lambda x: int(x.split(\"-\")[1]))\n path = dirs[-1] if len(dirs) > 0 else None\n\n if path is None:\n accelerator.print(\n f\"Checkpoint '{args.resume_from_checkpoint}' does not exist. Starting a new training run.\"\n )\n args.resume_from_checkpoint = None\n initial_global_step = 0\n else:\n accelerator.print(f\"Resuming from checkpoint {path}\")\n # accelerator.load_state(os.path.join(args.output_dir, path))\n global_step = int(path.split(\"-\")[1])\n initial_global_step = global_step\n first_epoch = global_step // num_update_steps_per_epoch\n\n progress_bar = tqdm(\n range(0, args.max_train_steps),\n initial=initial_global_step,\n desc=\"Steps\",\n # Only show the progress bar once on each machine.\n disable=not accelerator.is_local_main_process,\n )\n vae_scale_factor_spatial = 2 ** (len(vae.config.block_out_channels) - 1)\n\n # For DeepSpeed training\n model_config = transformer.module.config if hasattr(transformer, \"module\") else transformer.config\n\n empty_prompt_embeds, _ = compute_prompt_embeddings(\n tokenizer,\n text_encoder,\n \"\",\n model_config.max_text_seq_length,\n accelerator.device,\n weight_dtype,\n requires_grad=False,\n )\n\n for epoch in range(first_epoch, args.num_train_epochs):\n transformer.train()\n\n for step, batch in enumerate(train_dataloader):\n models_to_accumulate = [transformer]\n\n with accelerator.accumulate(models_to_accumulate):\n\n model_input = encode_video(batch[\"video\"]).permute(0, 2, 1, 3, 4).to(dtype=weight_dtype) # [B, F, C, H, W]\n batch_size, num_frames, num_channels, height, width = model_input.shape\n prompts = batch[\"prompt\"]\n prompts_list = batch[\"prompts_list\"]\n pose_embeds = batch[\"pose_embeds\"].to(dtype=weight_dtype)\n\n # encode prompts\n prompt_embeds, _ = compute_prompt_embeddings(\n tokenizer,\n text_encoder,\n prompts,\n model_config.max_text_seq_length,\n accelerator.device,\n weight_dtype,\n requires_grad=False,\n )\n \n # entity prompts\n entity_num = len(prompts_list)\n prompt_entities_embeds = []\n prompt_entities_attention_mask = []\n\n for idx in range(entity_num):\n prompt_entity_embeds, prompt_entity_attention_mask = compute_prompt_embeddings(\n tokenizer,\n text_encoder,\n prompts_list[idx],\n model_config.max_text_seq_length,\n accelerator.device,\n weight_dtype,\n requires_grad=False,\n )\n prompt_entities_embeds.append(prompt_entity_embeds.to(weight_dtype))\n prompt_entities_attention_mask.append(prompt_entity_attention_mask.to(weight_dtype))\n \n prompt_entities_embeds = rearrange(torch.stack(prompt_entities_embeds), \"n b t d -> b n t d\")\n prompt_entities_attention_mask = rearrange(torch.stack(prompt_entities_attention_mask), \"n b t -> b n t\")\n\n # empty prompt\n empty_prompt_embeds = empty_prompt_embeds.repeat(batch_size, 1, 1)\n\n # Sample noise that will be added to the latents\n noise = torch.randn_like(model_input)\n\n # Sample a random timestep for each image\n timesteps = torch.randint(\n 0, scheduler.config.num_train_timesteps, (batch_size,), device=model_input.device\n )\n timesteps = timesteps.long()\n\n # Prepare rotary embeds\n image_rotary_emb = (\n prepare_rotary_positional_embeddings(\n height=args.height,\n width=args.width,\n num_frames=num_frames,\n vae_scale_factor_spatial=vae_scale_factor_spatial,\n patch_size=model_config.patch_size,\n attention_head_dim=model_config.attention_head_dim,\n device=accelerator.device,\n )\n if model_config.use_rotary_positional_embeddings\n else None\n )\n\n # Add noise to the model input according to the noise magnitude at each timestep\n # (this is the forward diffusion process)\n noisy_model_input = scheduler.add_noise(model_input, noise, timesteps)\n \n # Predict the noise residual\n model_output = transformer(\n hidden_states=noisy_model_input,\n encoder_hidden_states=prompt_embeds,\n empty_encoder_hidden_states=empty_prompt_embeds,\n prompt_entities_embeds=prompt_entities_embeds,\n prompt_entities_attention_mask=prompt_entities_attention_mask,\n pose_embeds=pose_embeds,\n timestep=timesteps,\n image_rotary_emb=image_rotary_emb,\n return_dict=False,\n )[0]\n model_pred = scheduler.get_velocity(model_output, noisy_model_input, timesteps)\n\n alphas_cumprod = scheduler.alphas_cumprod[timesteps]\n weights = 1 / (1 - alphas_cumprod)\n while len(weights.shape) < len(model_pred.shape):\n weights = weights.unsqueeze(-1)\n\n target = model_input\n\n loss = torch.mean((weights * (model_pred - target) ** 2).reshape(batch_size, -1), dim=1)\n loss = loss.mean()\n accelerator.backward(loss)\n\n if accelerator.sync_gradients:\n params_to_clip = transformer.parameters()\n accelerator.clip_grad_norm_(params_to_clip, args.max_grad_norm)\n\n if accelerator.state.deepspeed_plugin is None:\n optimizer.step()\n optimizer.zero_grad()\n\n lr_scheduler.step()\n\n # Checks if the accelerator has performed an optimization step behind the scenes\n if accelerator.sync_gradients:\n progress_bar.update(1)\n global_step += 1\n\n if accelerator.is_main_process:\n if global_step % args.checkpointing_steps == 0:\n # _before_ saving state, check if this save would set us over the `checkpoints_total_limit`\n if args.checkpoints_total_limit is not None:\n checkpoints = os.listdir(args.output_dir)\n checkpoints = [d for d in checkpoints if d.startswith(\"checkpoint\")]\n checkpoints = sorted(checkpoints, key=lambda x: int(x.split(\"-\")[1]))\n\n # before we save the new checkpoint, we need to have at _most_ `checkpoints_total_limit - 1` checkpoints\n if len(checkpoints) >= args.checkpoints_total_limit:\n num_to_remove = len(checkpoints) - args.checkpoints_total_limit + 1\n removing_checkpoints = checkpoints[0:num_to_remove]\n\n logger.info(\n f\"{len(checkpoints)} checkpoints already exist, removing {len(removing_checkpoints)} checkpoints\"\n )\n logger.info(f\"Removing checkpoints: {', '.join(removing_checkpoints)}\")\n\n for removing_checkpoint in removing_checkpoints:\n removing_checkpoint = os.path.join(args.output_dir, removing_checkpoint)\n shutil.rmtree(removing_checkpoint)\n\n save_path = os.path.join(args.output_dir, f\"checkpoint-{global_step}\")\n accelerator.save_state(save_path)\n logger.info(f\"Saved state to {save_path}\")\n\n logs = {\"loss\": loss.detach().item(), \"lr\": lr_scheduler.get_last_lr()[0]}\n progress_bar.set_postfix(**logs)\n accelerator.log(logs, step=global_step)\n\n if global_step >= args.max_train_steps:\n break\n \n # Save the transformer layers\n accelerator.wait_for_everyone()\n accelerator.end_training()\n\n\nif __name__ == \"__main__\":\n args = get_args()\n main(args)" }, { "path": "CogVideo/finetune/train_cogvideox_lora.py", "content": "\"\"\"\nAdapted from CogVideoX-5B: https://github.com/THUDM/CogVideo by Xiao Fu (CUHK)\n\"\"\"\n\nimport argparse\nimport logging\nimport math\nimport os\nimport shutil\nfrom pathlib import Path\nfrom typing import List, Optional, Tuple, Union\n\nimport torch\nimport json\nimport numpy as np\nimport random\nimport cv2\nimport decord\nfrom einops import rearrange\nimport transformers\nfrom accelerate import Accelerator\nfrom accelerate.logging import get_logger\nfrom accelerate.utils import DistributedDataParallelKwargs, ProjectConfiguration, set_seed\nfrom huggingface_hub import create_repo, upload_folder\nfrom peft import LoraConfig, get_peft_model_state_dict, set_peft_model_state_dict\nfrom torch.utils.data import DataLoader, Dataset\nfrom torchvision import transforms\nfrom tqdm.auto import tqdm\nfrom transformers import AutoTokenizer, T5EncoderModel, T5Tokenizer\nfrom io import BytesIO\nimport imageio.v2 as imageio\n\nimport diffusers\nfrom diffusers import AutoencoderKLCogVideoX, CogVideoXDPMScheduler, CogVideoXPipeline, CogVideoXTransformer3DModel\nfrom diffusers.models.embeddings import get_3d_rotary_pos_embed\nfrom diffusers.optimization import get_scheduler\nfrom diffusers.pipelines.cogvideo.pipeline_cogvideox import get_resize_crop_region_for_grid\nfrom diffusers.training_utils import (\n cast_training_params,\n free_memory,\n)\nfrom diffusers.utils import check_min_version, convert_unet_state_dict_to_peft, export_to_video, is_wandb_available\nfrom diffusers.utils.hub_utils import load_or_create_model_card, populate_model_card\nfrom diffusers.utils.torch_utils import is_compiled_module\n\n\nif is_wandb_available():\n import wandb\n\n# Will error if the minimal version of diffusers is not installed. Remove at your own risks.\ncheck_min_version(\"0.31.0.dev0\")\n\nlogger = get_logger(__name__)\n\n\ndef get_args():\n parser = argparse.ArgumentParser(description=\"Simple example of a training script for CogVideoX.\")\n\n # Model information\n parser.add_argument(\n \"--pretrained_model_name_or_path\",\n type=str,\n default=None,\n required=True,\n help=\"Path to pretrained model or model identifier from huggingface.co/models.\",\n )\n parser.add_argument(\n \"--revision\",\n type=str,\n default=None,\n required=False,\n help=\"Revision of pretrained model identifier from huggingface.co/models.\",\n )\n parser.add_argument(\n \"--variant\",\n type=str,\n default=None,\n help=\"Variant of the model files of the pretrained model identifier from huggingface.co/models, 'e.g.' fp16\",\n )\n parser.add_argument(\n \"--cache_dir\",\n type=str,\n default=None,\n help=\"The directory where the downloaded models and datasets will be stored.\",\n )\n\n # Dataset information\n parser.add_argument(\n \"--dataset_name\",\n type=str,\n default=None,\n help=(\n \"The name of the Dataset (from the HuggingFace hub) containing the training data of instance images (could be your own, possibly private,\"\n \" dataset). It can also be a path pointing to a local copy of a dataset in your filesystem,\"\n \" or to a folder containing files that 🤗 Datasets can understand.\"\n ),\n )\n parser.add_argument(\n \"--dataset_config_name\",\n type=str,\n default=None,\n help=\"The config of the Dataset, leave as None if there's only one config.\",\n )\n parser.add_argument(\n \"--instance_data_root\",\n type=str,\n default=None,\n help=(\"A folder containing the training data.\"),\n )\n parser.add_argument(\n \"--id_token\", type=str, default=None, help=\"Identifier token appended to the start of each prompt if provided.\"\n )\n parser.add_argument(\n \"--dataloader_num_workers\",\n type=int,\n default=0,\n help=(\n \"Number of subprocesses to use for data loading. 0 means that the data will be loaded in the main process.\"\n ),\n )\n\n # Validation\n parser.add_argument(\n \"--validation_prompt\",\n type=str,\n default=None,\n help=\"One or more prompt(s) that is used during validation to verify that the model is learning. Multiple validation prompts should be separated by the '--validation_prompt_seperator' string.\",\n )\n parser.add_argument(\n \"--validation_prompt_separator\",\n type=str,\n default=\":::\",\n help=\"String that separates multiple validation prompts\",\n )\n parser.add_argument(\n \"--num_validation_videos\",\n type=int,\n default=1,\n help=\"Number of videos that should be generated during validation per `validation_prompt`.\",\n )\n parser.add_argument(\n \"--validation_epochs\",\n type=int,\n default=50,\n help=(\n \"Run validation every X epochs. Validation consists of running the prompt `args.validation_prompt` multiple times: `args.num_validation_videos`.\"\n ),\n )\n parser.add_argument(\n \"--guidance_scale\",\n type=float,\n default=6,\n help=\"The guidance scale to use while sampling validation videos.\",\n )\n parser.add_argument(\n \"--use_dynamic_cfg\",\n action=\"store_true\",\n default=False,\n help=\"Whether or not to use the default cosine dynamic guidance schedule when sampling validation videos.\",\n )\n\n # Training information\n parser.add_argument(\"--seed\", type=int, default=None, help=\"A seed for reproducible training.\")\n parser.add_argument(\n \"--rank\",\n type=int,\n default=128,\n help=(\"The dimension of the LoRA update matrices.\"),\n )\n parser.add_argument(\n \"--lora_alpha\",\n type=float,\n default=128,\n help=(\"The scaling factor to scale LoRA weight update. The actual scaling factor is `lora_alpha / rank`\"),\n )\n parser.add_argument(\n \"--mixed_precision\",\n type=str,\n default=None,\n choices=[\"no\", \"fp16\", \"bf16\"],\n help=(\n \"Whether to use mixed precision. Choose between fp16 and bf16 (bfloat16). Bf16 requires PyTorch >=\"\n \" 1.10.and an Nvidia Ampere GPU. Default to the value of accelerate config of the current system or the\"\n \" flag passed with the `accelerate.launch` command. Use this argument to override the accelerate config.\"\n ),\n )\n parser.add_argument(\n \"--output_dir\",\n type=str,\n default=\"cogvideox-lora\",\n help=\"The output directory where the model predictions and checkpoints will be written.\",\n )\n parser.add_argument(\n \"--height\",\n type=int,\n default=480,\n help=\"All input videos are resized to this height.\",\n )\n parser.add_argument(\n \"--width\",\n type=int,\n default=720,\n help=\"All input videos are resized to this width.\",\n )\n parser.add_argument(\"--fps\", type=int, default=8, help=\"All input videos will be used at this FPS.\")\n parser.add_argument(\n \"--max_num_frames\", type=int, default=49, help=\"All input videos will be truncated to these many frames.\"\n )\n parser.add_argument(\n \"--skip_frames_start\",\n type=int,\n default=0,\n help=\"Number of frames to skip from the beginning of each input video. Useful if training data contains intro sequences.\",\n )\n parser.add_argument(\n \"--skip_frames_end\",\n type=int,\n default=0,\n help=\"Number of frames to skip from the end of each input video. Useful if training data contains outro sequences.\",\n )\n parser.add_argument(\n \"--random_flip\",\n action=\"store_true\",\n help=\"whether to randomly flip videos horizontally\",\n )\n parser.add_argument(\n \"--train_batch_size\", type=int, default=4, help=\"Batch size (per device) for the training dataloader.\"\n )\n parser.add_argument(\"--num_train_epochs\", type=int, default=1)\n parser.add_argument(\n \"--max_train_steps\",\n type=int,\n default=None,\n help=\"Total number of training steps to perform. If provided, overrides `--num_train_epochs`.\",\n )\n parser.add_argument(\n \"--checkpointing_steps\",\n type=int,\n default=500,\n help=(\n \"Save a checkpoint of the training state every X updates. These checkpoints can be used both as final\"\n \" checkpoints in case they are better than the last checkpoint, and are also suitable for resuming\"\n \" training using `--resume_from_checkpoint`.\"\n ),\n )\n parser.add_argument(\n \"--checkpoints_total_limit\",\n type=int,\n default=None,\n help=(\"Max number of checkpoints to store.\"),\n )\n parser.add_argument(\n \"--resume_from_checkpoint\",\n type=str,\n default=None,\n help=(\n \"Whether training should be resumed from a previous checkpoint. Use a path saved by\"\n ' `--checkpointing_steps`, or `\"latest\"` to automatically select the last available checkpoint.'\n ),\n )\n parser.add_argument(\n \"--gradient_accumulation_steps\",\n type=int,\n default=1,\n help=\"Number of updates steps to accumulate before performing a backward/update pass.\",\n )\n parser.add_argument(\n \"--gradient_checkpointing\",\n action=\"store_true\",\n help=\"Whether or not to use gradient checkpointing to save memory at the expense of slower backward pass.\",\n )\n parser.add_argument(\n \"--learning_rate\",\n type=float,\n default=1e-4,\n help=\"Initial learning rate (after the potential warmup period) to use.\",\n )\n parser.add_argument(\n \"--scale_lr\",\n action=\"store_true\",\n default=False,\n help=\"Scale the learning rate by the number of GPUs, gradient accumulation steps, and batch size.\",\n )\n parser.add_argument(\n \"--lr_scheduler\",\n type=str,\n default=\"constant\",\n help=(\n 'The scheduler type to use. Choose between [\"linear\", \"cosine\", \"cosine_with_restarts\", \"polynomial\",'\n ' \"constant\", \"constant_with_warmup\"]'\n ),\n )\n parser.add_argument(\n \"--lr_warmup_steps\", type=int, default=500, help=\"Number of steps for the warmup in the lr scheduler.\"\n )\n parser.add_argument(\n \"--lr_num_cycles\",\n type=int,\n default=1,\n help=\"Number of hard resets of the lr in cosine_with_restarts scheduler.\",\n )\n parser.add_argument(\"--lr_power\", type=float, default=1.0, help=\"Power factor of the polynomial scheduler.\")\n parser.add_argument(\n \"--enable_slicing\",\n action=\"store_true\",\n default=False,\n help=\"Whether or not to use VAE slicing for saving memory.\",\n )\n parser.add_argument(\n \"--enable_tiling\",\n action=\"store_true\",\n default=False,\n help=\"Whether or not to use VAE tiling for saving memory.\",\n )\n\n # Optimizer\n parser.add_argument(\n \"--optimizer\",\n type=lambda s: s.lower(),\n default=\"adam\",\n choices=[\"adam\", \"adamw\", \"prodigy\"],\n help=(\"The optimizer type to use.\"),\n )\n parser.add_argument(\n \"--use_8bit_adam\",\n action=\"store_true\",\n help=\"Whether or not to use 8-bit Adam from bitsandbytes. Ignored if optimizer is not set to AdamW\",\n )\n parser.add_argument(\n \"--adam_beta1\", type=float, default=0.9, help=\"The beta1 parameter for the Adam and Prodigy optimizers.\"\n )\n parser.add_argument(\n \"--adam_beta2\", type=float, default=0.95, help=\"The beta2 parameter for the Adam and Prodigy optimizers.\"\n )\n parser.add_argument(\n \"--prodigy_beta3\",\n type=float,\n default=None,\n help=\"Coefficients for computing the Prodigy optimizer's stepsize using running averages. If set to None, uses the value of square root of beta2.\",\n )\n parser.add_argument(\"--prodigy_decouple\", action=\"store_true\", help=\"Use AdamW style decoupled weight decay\")\n parser.add_argument(\"--adam_weight_decay\", type=float, default=1e-04, help=\"Weight decay to use for unet params\")\n parser.add_argument(\n \"--adam_epsilon\",\n type=float,\n default=1e-08,\n help=\"Epsilon value for the Adam optimizer and Prodigy optimizers.\",\n )\n parser.add_argument(\"--max_grad_norm\", default=1.0, type=float, help=\"Max gradient norm.\")\n parser.add_argument(\"--prodigy_use_bias_correction\", action=\"store_true\", help=\"Turn on Adam's bias correction.\")\n parser.add_argument(\n \"--prodigy_safeguard_warmup\",\n action=\"store_true\",\n help=\"Remove lr from the denominator of D estimate to avoid issues during warm-up stage.\",\n )\n\n # Other information\n parser.add_argument(\"--tracker_name\", type=str, default=None, help=\"Project tracker name\")\n parser.add_argument(\"--push_to_hub\", action=\"store_true\", help=\"Whether or not to push the model to the Hub.\")\n parser.add_argument(\"--hub_token\", type=str, default=None, help=\"The token to use to push to the Model Hub.\")\n parser.add_argument(\n \"--hub_model_id\",\n type=str,\n default=None,\n help=\"The name of the repository to keep in sync with the local `output_dir`.\",\n )\n parser.add_argument(\n \"--logging_dir\",\n type=str,\n default=\"logs\",\n help=\"Directory where logs are stored.\",\n )\n parser.add_argument(\n \"--allow_tf32\",\n action=\"store_true\",\n help=(\n \"Whether or not to allow TF32 on Ampere GPUs. Can be used to speed up training. For more information, see\"\n \" https://pytorch.org/docs/stable/notes/cuda.html#tensorfloat-32-tf32-on-ampere-devices\"\n ),\n )\n parser.add_argument(\n \"--report_to\",\n type=str,\n default=None,\n help=(\n 'The integration to report the results and logs to. Supported platforms are `\"tensorboard\"`'\n ' (default), `\"wandb\"` and `\"comet_ml\"`. Use `\"all\"` to report to all integrations.'\n ),\n )\n\n return parser.parse_args()\n\ndef parse_matrix(matrix_str):\n rows = matrix_str.strip().split('] [')\n matrix = []\n for row in rows:\n row = row.replace('[', '').replace(']', '')\n matrix.append(list(map(float, row.split())))\n return np.array(matrix)\n\nclass VideoDataset(Dataset):\n def __init__(\n self,\n instance_data_root: Optional[str] = None,\n dataset_name: Optional[str] = None,\n dataset_config_name: Optional[str] = None,\n height: int = 480,\n width: int = 720,\n fps: int = 8,\n max_num_frames: int = 49,\n skip_frames_start: int = 0,\n skip_frames_end: int = 0,\n cache_dir: Optional[str] = None,\n id_token: Optional[str] = None,\n ) -> None:\n super().__init__()\n\n self.instance_data_root = Path(instance_data_root) if instance_data_root is not None else None\n self.dataset_name = dataset_name\n self.dataset_config_name = dataset_config_name\n self.height = height\n self.width = width\n self.sample_size = (self.height, self.width)\n self.fps = fps\n self.max_num_frames = max_num_frames\n self.skip_frames_start = skip_frames_start\n self.skip_frames_end = skip_frames_end\n self.cache_dir = cache_dir\n self.id_token = id_token or \"\"\n self.pixel_transforms = [\n transforms.Resize(self.sample_size),\n transforms.Normalize(\n mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5], inplace=True\n ),\n ]\n \n self.video_names = []\n\n # ----------------------------------- Released Dataset -----------------------------------\n scenes = ['Desert', 'HDRI']\n scene_location_pair = {\n 'Desert' : 'desert',\n 'HDRI' : \n {\n 'loc1' : 'snowy street',\n 'loc2' : 'park',\n 'loc3' : 'indoor open space',\n 'loc11' : 'gymnastics room',\n 'loc13' : 'autumn forest',\n }\n }\n\n for scene in scenes:\n video_path = os.path.join(self.instance_data_root, '480_720', scene)\n video_names = os.listdir(video_path)\n locations_path = os.path.join(video_path, \"location_data.json\")\n with open(locations_path, 'r') as f: locations = json.load(f)\n locations_info = {locations[idx]['name']:locations[idx] for idx in range(len(locations))}\n for video_name in video_names:\n if video_name.endswith('Hemi12_1') == True:\n if scene != 'HDRI':\n location = scene_location_pair[scene]\n else:\n location = scene_location_pair['HDRI'][video_name.split('_')[1]]\n self.video_names.append((scene, video_name, location, locations_info))\n\n # ----------------------------------- Internal Dataset -----------------------------------\n # scenes = ['AsianTown_480_720', 'Desert_480_720', 'HDRI_480_720', 'Forest_480_720']\n # scene_location_pair = {\n # 'AsianTown_480_720' : 'asian town',\n # 'Desert_480_720' : 'desert',\n # 'Forest_480_720' : 'crossland',\n # 'MatrixCity' : 'city',\n # 'HDRI_480_720' : \n # {\n # 'loc1' : 'snowy street',\n # 'loc2' : 'park',\n # 'loc3' : 'indoor open space',\n # 'loc11' : 'gymnastics room',\n # 'loc13' : 'autumn forest',\n # }\n # }\n \n # for scene in scenes:\n # video_path = os.path.join(self.instance_data_root, scene)\n # video_names = os.listdir(video_path)\n # locations_path = os.path.join(video_path, \"location_data.json\")\n # with open(locations_path, 'r') as f: locations = json.load(f)\n # locations_info = {locations[idx]['name']:locations[idx] for idx in range(len(locations))}\n # for video_name in video_names:\n # if video_name.endswith('Hemi12_1') == True:\n # if scene != 'HDRI_480_720':\n # location = scene_location_pair[scene]\n # else:\n # location = scene_location_pair['HDRI_480_720'][video_name.split('_')[1]]\n # self.video_names.append((scene, video_name, location, locations_info))\n\n self.cam_num = 12\n self.max_objs_num = 3\n self.length = len(self.video_names)\n self.captions_path = os.path.join(self.instance_data_root, \"CharacterInfo.json\")\n with open(self.captions_path, 'r') as f: captions = json.load(f)['CharacterInfo']\n self.captions_info = {int(captions[idx]['index']):captions[idx]['eng'] for idx in range(len(captions))}\n\n self.cams_path = os.path.join(self.instance_data_root, \"Hemi12_transforms.json\")\n with open(self.cams_path, 'r') as f: self.cams_info = json.load(f)\n cam_poses = []\n for i, key in enumerate(self.cams_info.keys()):\n if \"C_\" in key:\n cam_poses.append(parse_matrix(self.cams_info[key]))\n cam_poses = np.stack(cam_poses)\n cam_poses = np.transpose(cam_poses, (0,2,1))\n cam_poses = cam_poses[:,:,[1,2,0,3]]\n cam_poses[:,:3,3] /= 100.\n self.cam_poses = cam_poses\n self.sample_n_frames = 49\n\n def __len__(self):\n return self.length\n \n def save_images2video(self, images, video_name):\n fps = 8\n format = \"mp4\"\n codec = \"libx264\" \n ffmpeg_params = [\"-crf\", str(12)]\n pixelformat = \"yuv420p\" \n video_stream = BytesIO()\n \n with imageio.get_writer(\n video_stream,\n fps=fps,\n format=format,\n codec=codec,\n ffmpeg_params=ffmpeg_params,\n pixelformat=pixelformat,\n ) as writer:\n for idx in range(len(images)):\n writer.append_data(images[idx])\n \n video_data = video_stream.getvalue()\n output_path = os.path.join(video_name + \".mp4\")\n with open(output_path, \"wb\") as f:\n f.write(video_data)\n\n def __getitem__(self, idx):\n \n while True:\n try:\n (scene, video_name, location, locations_info) = self.video_names[idx]\n\n with open(os.path.join(self.instance_data_root, '480_720', scene, video_name, video_name+'.json'), 'r') as f: objs_file = json.load(f)\n objs_num = len(objs_file['0'])\n video_index = random.randint(1, self.cam_num-1)\n\n location_name = video_name.split('_')[1]\n location_info = locations_info[location_name]\n cam_pose = self.cam_poses[video_index-1]\n obj_transl = location_info['coordinates']['CameraTarget']['position']\n\n video_caption_concat = ''\n video_caption_list = []\n obj_poses_list = []\n\n for obj_idx in range(objs_num):\n \n obj_name_index = objs_file['0'][obj_idx]['index'] \n video_caption = self.captions_info[obj_name_index]\n\n if video_caption.startswith(\" \"):\n video_caption = video_caption[1:]\n if video_caption.endswith(\".\"):\n video_caption = video_caption[:-1]\n video_caption = video_caption.lower()\n video_caption_list.append(video_caption)\n \n obj_poses = self.load_sceneposes(objs_file, obj_idx, obj_transl)\n obj_poses = np.linalg.inv(cam_pose) @ obj_poses\n obj_poses_list.append(obj_poses)\n\n for obj_idx in range(objs_num):\n video_caption = video_caption_list[obj_idx]\n if obj_idx == objs_num - 1:\n if objs_num == 1:\n video_caption_concat += video_caption + ' is moving in the ' + location\n else:\n video_caption_concat += video_caption + ' are moving in the ' + location\n else:\n video_caption_concat += video_caption + ' and '\n\n obj_poses_all = torch.from_numpy(np.array(obj_poses_list))\n\n total_frames = 99\n current_sample_stride = 1.75\n cropped_length = int(self.sample_n_frames * current_sample_stride)\n start_frame_ind = random.randint(10, max(10, total_frames - cropped_length - 1))\n end_frame_ind = min(start_frame_ind + cropped_length, total_frames)\n frame_indices = np.linspace(start_frame_ind, end_frame_ind - 1, self.sample_n_frames, dtype=int)\n \n video_frames_path = os.path.join(self.instance_data_root, '480_720', scene, video_name, 'videos', video_name+ f'_C_{video_index:02d}_35mm.mp4')\n cap = cv2.VideoCapture(video_frames_path)\n height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))\n width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))\n \n # get local rank\n ctx = decord.cpu(0)\n reader = decord.VideoReader(video_frames_path, ctx=ctx, height=height, width=width)\n assert len(reader) == total_frames or len(reader) == total_frames+1\n frame_indexes = [frame_idx for frame_idx in range(total_frames)]\n try:\n video_chunk = reader.get_batch(frame_indexes).asnumpy() \n except:\n video_chunk = reader.get_batch(frame_indexes).numpy()\n\n frame_inverse = torch.rand(1).item() > 0.5\n if frame_inverse:\n frame_indices = frame_indices[::-1]\n video_name += '_inv'\n\n pixel_values = np.array([video_chunk[indice] for indice in frame_indices])\n pixel_values = rearrange(torch.from_numpy(pixel_values) / 255.0, \"f h w c -> f c h w\")\n pixel_values = self.pixel_transforms[0](pixel_values)\n pixel_values = self.pixel_transforms[1](pixel_values)\n\n # interpolation\n trunc_frame_indices = np.zeros_like(frame_indices[::4])\n trunc_frame_indices[0] = frame_indices[0]\n trunc_frame_indices[1:] = ((frame_indices[1:][::4] + frame_indices[4:][::4]) / 2).astype(np.int64)\n\n obj_poses_all = obj_poses_all[:, trunc_frame_indices]\n pose_embeds = rearrange(obj_poses_all[:, :, :3, :], \"b f p q -> b f (q p)\").contiguous()\n\n break\n \n except Exception as e:\n \n (scene, video_name, location, locations_info) = self.video_names[idx]\n\n with open(f'invalid_scene.txt', 'a+') as f:\n f.write(f'{scene} {video_name} {location}')\n f.write('\\n')\n\n idx = random.randint(0, self.length - 1)\n\n return {\n \"prompt\": video_caption_concat, \n \"video\": pixel_values,\n \"video_name\": video_name,\n }\n\n def load_sceneposes(self, objs_file, obj_idx, obj_transl):\n ext_poses = []\n for i, key in enumerate(objs_file.keys()):\n ext_poses.append(parse_matrix(objs_file[key][obj_idx]['matrix']))\n ext_poses = np.stack(ext_poses)\n ext_poses = np.transpose(ext_poses, (0,2,1))\n ext_poses[:,:3,3] -= obj_transl\n ext_poses[:,:3,3] /= 100.\n ext_poses = ext_poses[:, :, [1,2,0,3]]\n return ext_poses\n\n\ndef save_model_card(\n repo_id: str,\n videos=None,\n base_model: str = None,\n validation_prompt=None,\n repo_folder=None,\n fps=8,\n):\n widget_dict = []\n if videos is not None:\n for i, video in enumerate(videos):\n export_to_video(video, os.path.join(repo_folder, f\"final_video_{i}.mp4\", fps=fps))\n widget_dict.append(\n {\"text\": validation_prompt if validation_prompt else \" \", \"output\": {\"url\": f\"video_{i}.mp4\"}}\n )\n\n model_description = f\"\"\"\n# CogVideoX LoRA - {repo_id}\n\n\n\n## Model description\n\nThese are {repo_id} LoRA weights for {base_model}.\n\nThe weights were trained using the [CogVideoX Diffusers trainer](https://github.com/huggingface/diffusers/blob/main/examples/cogvideo/train_cogvideox_lora.py).\n\nWas LoRA for the text encoder enabled? No.\n\n## Download model\n\n[Download the *.safetensors LoRA]({repo_id}/tree/main) in the Files & versions tab.\n\n## Use it with the [🧨 diffusers library](https://github.com/huggingface/diffusers)\n\n```py\nfrom diffusers import CogVideoXPipeline\nimport torch\n\npipe = CogVideoXPipeline.from_pretrained(\"THUDM/CogVideoX-5b\", torch_dtype=torch.bfloat16).to(\"cuda\")\npipe.load_lora_weights(\"{repo_id}\", weight_name=\"pytorch_lora_weights.safetensors\", adapter_name=[\"cogvideox-lora\"])\n\n# The LoRA adapter weights are determined by what was used for training.\n# In this case, we assume `--lora_alpha` is 32 and `--rank` is 64.\n# It can be made lower or higher from what was used in training to decrease or amplify the effect\n# of the LoRA upto a tolerance, beyond which one might notice no effect at all or overflows.\npipe.set_adapters([\"cogvideox-lora\"], [32 / 64])\n\nvideo = pipe(\"{validation_prompt}\", guidance_scale=6, use_dynamic_cfg=True).frames[0]\n```\n\nFor more details, including weighting, merging and fusing LoRAs, check the [documentation on loading LoRAs in diffusers](https://huggingface.co/docs/diffusers/main/en/using-diffusers/loading_adapters)\n\n## License\n\nPlease adhere to the licensing terms as described [here](https://huggingface.co/THUDM/CogVideoX-5b/blob/main/LICENSE) and [here](https://huggingface.co/THUDM/CogVideoX-2b/blob/main/LICENSE).\n\"\"\"\n model_card = load_or_create_model_card(\n repo_id_or_path=repo_id,\n from_training=True,\n license=\"other\",\n base_model=base_model,\n prompt=validation_prompt,\n model_description=model_description,\n widget=widget_dict,\n )\n tags = [\n \"text-to-video\",\n \"diffusers-training\",\n \"diffusers\",\n \"lora\",\n \"cogvideox\",\n \"cogvideox-diffusers\",\n \"template:sd-lora\",\n ]\n\n model_card = populate_model_card(model_card, tags=tags)\n model_card.save(os.path.join(repo_folder, \"README.md\"))\n\n\ndef log_validation(\n pipe,\n args,\n accelerator,\n pipeline_args,\n epoch,\n is_final_validation: bool = False,\n):\n logger.info(\n f\"Running validation... \\n Generating {args.num_validation_videos} videos with prompt: {pipeline_args['prompt']}.\"\n )\n # We train on the simplified learning objective. If we were previously predicting a variance, we need the scheduler to ignore it\n scheduler_args = {}\n\n if \"variance_type\" in pipe.scheduler.config:\n variance_type = pipe.scheduler.config.variance_type\n\n if variance_type in [\"learned\", \"learned_range\"]:\n variance_type = \"fixed_small\"\n\n scheduler_args[\"variance_type\"] = variance_type\n\n pipe.scheduler = CogVideoXDPMScheduler.from_config(pipe.scheduler.config, **scheduler_args)\n pipe = pipe.to(accelerator.device)\n # pipe.set_progress_bar_config(disable=True)\n\n # run inference\n generator = torch.Generator(device=accelerator.device).manual_seed(args.seed) if args.seed else None\n\n videos = []\n for _ in range(args.num_validation_videos):\n video = pipe(**pipeline_args, generator=generator, output_type=\"np\").frames[0]\n videos.append(video)\n\n for tracker in accelerator.trackers:\n phase_name = \"test\" if is_final_validation else \"validation\"\n if tracker.name == \"wandb\":\n video_filenames = []\n for i, video in enumerate(videos):\n prompt = (\n pipeline_args[\"prompt\"][:25]\n .replace(\" \", \"_\")\n .replace(\" \", \"_\")\n .replace(\"'\", \"_\")\n .replace('\"', \"_\")\n .replace(\"/\", \"_\")\n )\n filename = os.path.join(args.output_dir, f\"{phase_name}_video_{i}_{prompt}.mp4\")\n export_to_video(video, filename, fps=8)\n video_filenames.append(filename)\n\n tracker.log(\n {\n phase_name: [\n wandb.Video(filename, caption=f\"{i}: {pipeline_args['prompt']}\")\n for i, filename in enumerate(video_filenames)\n ]\n }\n )\n\n free_memory()\n\n return videos\n\n\ndef _get_t5_prompt_embeds(\n tokenizer: T5Tokenizer,\n text_encoder: T5EncoderModel,\n prompt: Union[str, List[str]],\n num_videos_per_prompt: int = 1,\n max_sequence_length: int = 226,\n device: Optional[torch.device] = None,\n dtype: Optional[torch.dtype] = None,\n text_input_ids=None,\n):\n prompt = [prompt] if isinstance(prompt, str) else prompt\n batch_size = len(prompt)\n\n if tokenizer is not None:\n text_inputs = tokenizer(\n prompt,\n padding=\"max_length\",\n max_length=max_sequence_length,\n truncation=True,\n add_special_tokens=True,\n return_tensors=\"pt\",\n )\n text_input_ids = text_inputs.input_ids\n else:\n if text_input_ids is None:\n raise ValueError(\"`text_input_ids` must be provided when the tokenizer is not specified.\")\n\n prompt_embeds = text_encoder(text_input_ids.to(device))[0]\n prompt_embeds = prompt_embeds.to(dtype=dtype, device=device)\n\n # duplicate text embeddings for each generation per prompt, using mps friendly method\n _, seq_len, _ = prompt_embeds.shape\n prompt_embeds = prompt_embeds.repeat(1, num_videos_per_prompt, 1)\n prompt_embeds = prompt_embeds.view(batch_size * num_videos_per_prompt, seq_len, -1)\n\n return prompt_embeds\n\n\ndef encode_prompt(\n tokenizer: T5Tokenizer,\n text_encoder: T5EncoderModel,\n prompt: Union[str, List[str]],\n num_videos_per_prompt: int = 1,\n max_sequence_length: int = 226,\n device: Optional[torch.device] = None,\n dtype: Optional[torch.dtype] = None,\n text_input_ids=None,\n):\n prompt = [prompt] if isinstance(prompt, str) else prompt\n prompt_embeds = _get_t5_prompt_embeds(\n tokenizer,\n text_encoder,\n prompt=prompt,\n num_videos_per_prompt=num_videos_per_prompt,\n max_sequence_length=max_sequence_length,\n device=device,\n dtype=dtype,\n text_input_ids=text_input_ids,\n )\n return prompt_embeds\n\n\ndef compute_prompt_embeddings(\n tokenizer, text_encoder, prompt, max_sequence_length, device, dtype, requires_grad: bool = False\n):\n if requires_grad:\n prompt_embeds = encode_prompt(\n tokenizer,\n text_encoder,\n prompt,\n num_videos_per_prompt=1,\n max_sequence_length=max_sequence_length,\n device=device,\n dtype=dtype,\n )\n else:\n with torch.no_grad():\n prompt_embeds = encode_prompt(\n tokenizer,\n text_encoder,\n prompt,\n num_videos_per_prompt=1,\n max_sequence_length=max_sequence_length,\n device=device,\n dtype=dtype,\n )\n return prompt_embeds\n\n\ndef prepare_rotary_positional_embeddings(\n height: int,\n width: int,\n num_frames: int,\n vae_scale_factor_spatial: int = 8,\n patch_size: int = 2,\n attention_head_dim: int = 64,\n device: Optional[torch.device] = None,\n base_height: int = 480,\n base_width: int = 720,\n) -> Tuple[torch.Tensor, torch.Tensor]:\n grid_height = height // (vae_scale_factor_spatial * patch_size)\n grid_width = width // (vae_scale_factor_spatial * patch_size)\n base_size_width = base_width // (vae_scale_factor_spatial * patch_size)\n base_size_height = base_height // (vae_scale_factor_spatial * patch_size)\n\n grid_crops_coords = get_resize_crop_region_for_grid((grid_height, grid_width), base_size_width, base_size_height)\n freqs_cos, freqs_sin = get_3d_rotary_pos_embed(\n embed_dim=attention_head_dim,\n crops_coords=grid_crops_coords,\n grid_size=(grid_height, grid_width),\n temporal_size=num_frames,\n )\n\n freqs_cos = freqs_cos.to(device=device)\n freqs_sin = freqs_sin.to(device=device)\n return freqs_cos, freqs_sin\n\n\ndef get_optimizer(args, params_to_optimize, use_deepspeed: bool = False):\n # Use DeepSpeed optimzer\n if use_deepspeed:\n from accelerate.utils import DummyOptim\n\n return DummyOptim(\n params_to_optimize,\n lr=args.learning_rate,\n betas=(args.adam_beta1, args.adam_beta2),\n eps=args.adam_epsilon,\n weight_decay=args.adam_weight_decay,\n )\n\n\n # Optimizer creation\n supported_optimizers = [\"adam\", \"adamw\", \"prodigy\"]\n if args.optimizer not in supported_optimizers:\n logger.warning(\n f\"Unsupported choice of optimizer: {args.optimizer}. Supported optimizers include {supported_optimizers}. Defaulting to AdamW\"\n )\n args.optimizer = \"adamw\"\n\n if args.use_8bit_adam and not (args.optimizer.lower() not in [\"adam\", \"adamw\"]):\n logger.warning(\n f\"use_8bit_adam is ignored when optimizer is not set to 'Adam' or 'AdamW'. Optimizer was \"\n f\"set to {args.optimizer.lower()}\"\n )\n\n if args.use_8bit_adam:\n try:\n import bitsandbytes as bnb\n except ImportError:\n raise ImportError(\n \"To use 8-bit Adam, please install the bitsandbytes library: `pip install bitsandbytes`.\"\n )\n\n if args.optimizer.lower() == \"adamw\":\n optimizer_class = bnb.optim.AdamW8bit if args.use_8bit_adam else torch.optim.AdamW\n\n optimizer = optimizer_class(\n params_to_optimize,\n betas=(args.adam_beta1, args.adam_beta2),\n eps=args.adam_epsilon,\n weight_decay=args.adam_weight_decay,\n )\n elif args.optimizer.lower() == \"adam\":\n optimizer_class = bnb.optim.Adam8bit if args.use_8bit_adam else torch.optim.Adam\n\n optimizer = optimizer_class(\n params_to_optimize,\n betas=(args.adam_beta1, args.adam_beta2),\n eps=args.adam_epsilon,\n weight_decay=args.adam_weight_decay,\n )\n elif args.optimizer.lower() == \"prodigy\":\n try:\n import prodigyopt\n except ImportError:\n raise ImportError(\"To use Prodigy, please install the prodigyopt library: `pip install prodigyopt`\")\n\n optimizer_class = prodigyopt.Prodigy\n\n if args.learning_rate <= 0.1:\n logger.warning(\n \"Learning rate is too low. When using prodigy, it's generally better to set learning rate around 1.0\"\n )\n\n optimizer = optimizer_class(\n params_to_optimize,\n lr=args.learning_rate,\n betas=(args.adam_beta1, args.adam_beta2),\n beta3=args.prodigy_beta3,\n weight_decay=args.adam_weight_decay,\n eps=args.adam_epsilon,\n decouple=args.prodigy_decouple,\n use_bias_correction=args.prodigy_use_bias_correction,\n safeguard_warmup=args.prodigy_safeguard_warmup,\n )\n\n return optimizer\n\n\ndef main(args):\n if args.report_to == \"wandb\" and args.hub_token is not None:\n raise ValueError(\n \"You cannot use both --report_to=wandb and --hub_token due to a security risk of exposing your token.\"\n \" Please use `huggingface-cli login` to authenticate with the Hub.\"\n )\n\n if torch.backends.mps.is_available() and args.mixed_precision == \"bf16\":\n # due to pytorch#99272, MPS does not yet support bfloat16.\n raise ValueError(\n \"Mixed precision training with bfloat16 is not supported on MPS. Please use fp16 (recommended) or fp32 instead.\"\n )\n\n logging_dir = Path(args.output_dir, args.logging_dir)\n\n accelerator_project_config = ProjectConfiguration(project_dir=args.output_dir, logging_dir=logging_dir)\n kwargs = DistributedDataParallelKwargs(find_unused_parameters=True)\n accelerator = Accelerator(\n gradient_accumulation_steps=args.gradient_accumulation_steps,\n mixed_precision=args.mixed_precision,\n log_with=args.report_to,\n project_config=accelerator_project_config,\n kwargs_handlers=[kwargs],\n )\n\n # Disable AMP for MPS.\n if torch.backends.mps.is_available():\n accelerator.native_amp = False\n\n if args.report_to == \"wandb\":\n if not is_wandb_available():\n raise ImportError(\"Make sure to install wandb if you want to use it for logging during training.\")\n\n # Make one log on every process with the configuration for debugging.\n logging.basicConfig(\n format=\"%(asctime)s - %(levelname)s - %(name)s - %(message)s\",\n datefmt=\"%m/%d/%Y %H:%M:%S\",\n level=logging.INFO,\n )\n logger.info(accelerator.state, main_process_only=False)\n if accelerator.is_local_main_process:\n transformers.utils.logging.set_verbosity_warning()\n diffusers.utils.logging.set_verbosity_info()\n else:\n transformers.utils.logging.set_verbosity_error()\n diffusers.utils.logging.set_verbosity_error()\n\n # If passed along, set the training seed now.\n if args.seed is not None:\n set_seed(args.seed)\n\n # Handle the repository creation\n if accelerator.is_main_process:\n if args.output_dir is not None:\n os.makedirs(args.output_dir, exist_ok=True)\n\n if args.push_to_hub:\n repo_id = create_repo(\n repo_id=args.hub_model_id or Path(args.output_dir).name,\n exist_ok=True,\n ).repo_id\n\n # Prepare models and scheduler\n tokenizer = AutoTokenizer.from_pretrained(\n args.pretrained_model_name_or_path, subfolder=\"tokenizer\", revision=args.revision\n )\n\n text_encoder = T5EncoderModel.from_pretrained(\n args.pretrained_model_name_or_path, subfolder=\"text_encoder\", revision=args.revision\n )\n\n # CogVideoX-2b weights are stored in float16\n # CogVideoX-5b and CogVideoX-5b-I2V weights are stored in bfloat16\n load_dtype = torch.bfloat16 if \"5b\" in args.pretrained_model_name_or_path.lower() else torch.float16\n transformer = CogVideoXTransformer3DModel.from_pretrained(\n args.pretrained_model_name_or_path,\n subfolder=\"transformer\",\n torch_dtype=load_dtype,\n revision=args.revision,\n variant=args.variant,\n )\n\n vae = AutoencoderKLCogVideoX.from_pretrained(\n args.pretrained_model_name_or_path, subfolder=\"vae\", revision=args.revision, variant=args.variant\n )\n\n scheduler = CogVideoXDPMScheduler.from_pretrained(args.pretrained_model_name_or_path, subfolder=\"scheduler\")\n\n if args.enable_slicing:\n vae.enable_slicing()\n if args.enable_tiling:\n vae.enable_tiling()\n\n # We only train the additional adapter LoRA layers\n text_encoder.requires_grad_(False)\n transformer.requires_grad_(False)\n vae.requires_grad_(False)\n\n # For mixed precision training we cast all non-trainable weights (vae, text_encoder and transformer) to half-precision\n # as these weights are only used for inference, keeping weights in full precision is not required.\n weight_dtype = torch.float32\n if accelerator.state.deepspeed_plugin:\n # DeepSpeed is handling precision, use what's in the DeepSpeed config\n if (\n \"fp16\" in accelerator.state.deepspeed_plugin.deepspeed_config\n and accelerator.state.deepspeed_plugin.deepspeed_config[\"fp16\"][\"enabled\"]\n ):\n weight_dtype = torch.float16\n if (\n \"bf16\" in accelerator.state.deepspeed_plugin.deepspeed_config\n and accelerator.state.deepspeed_plugin.deepspeed_config[\"bf16\"][\"enabled\"]\n ):\n weight_dtype = torch.float16\n else:\n if accelerator.mixed_precision == \"fp16\":\n weight_dtype = torch.float16\n elif accelerator.mixed_precision == \"bf16\":\n weight_dtype = torch.bfloat16\n\n if torch.backends.mps.is_available() and weight_dtype == torch.bfloat16:\n # due to pytorch#99272, MPS does not yet support bfloat16.\n raise ValueError(\n \"Mixed precision training with bfloat16 is not supported on MPS. Please use fp16 (recommended) or fp32 instead.\"\n )\n\n text_encoder.to(accelerator.device, dtype=weight_dtype)\n transformer.to(accelerator.device, dtype=weight_dtype)\n vae.to(accelerator.device, dtype=weight_dtype)\n\n if args.gradient_checkpointing:\n transformer.enable_gradient_checkpointing()\n\n # now we will add new LoRA weights to the attention layers\n transformer_lora_config = LoraConfig(\n r=args.rank,\n lora_alpha=args.lora_alpha,\n init_lora_weights=True,\n target_modules=[\"to_k\", \"to_q\", \"to_v\", \"to_out.0\"],\n )\n transformer.add_adapter(transformer_lora_config)\n\n def unwrap_model(model):\n model = accelerator.unwrap_model(model)\n model = model._orig_mod if is_compiled_module(model) else model\n return model\n\n # create custom saving & loading hooks so that `accelerator.save_state(...)` serializes in a nice format\n def save_model_hook(models, weights, output_dir):\n if accelerator.is_main_process:\n transformer_lora_layers_to_save = None\n\n for model in models:\n if isinstance(model, type(unwrap_model(transformer))):\n transformer_lora_layers_to_save = get_peft_model_state_dict(model)\n else:\n raise ValueError(f\"unexpected save model: {model.__class__}\")\n\n # make sure to pop weight so that corresponding model is not saved again\n weights.pop()\n\n CogVideoXPipeline.save_lora_weights(\n output_dir,\n transformer_lora_layers=transformer_lora_layers_to_save,\n )\n\n def load_model_hook(models, input_dir):\n transformer_ = None\n\n while len(models) > 0:\n model = models.pop()\n\n if isinstance(model, type(unwrap_model(transformer))):\n transformer_ = model\n else:\n raise ValueError(f\"Unexpected save model: {model.__class__}\")\n\n lora_state_dict = CogVideoXPipeline.lora_state_dict(input_dir)\n\n transformer_state_dict = {\n f'{k.replace(\"transformer.\", \"\")}': v for k, v in lora_state_dict.items() if k.startswith(\"transformer.\")\n }\n transformer_state_dict = convert_unet_state_dict_to_peft(transformer_state_dict)\n incompatible_keys = set_peft_model_state_dict(transformer_, transformer_state_dict, adapter_name=\"default\")\n if incompatible_keys is not None:\n # check only for unexpected keys\n unexpected_keys = getattr(incompatible_keys, \"unexpected_keys\", None)\n if unexpected_keys:\n logger.warning(\n f\"Loading adapter weights from state_dict led to unexpected keys not found in the model: \"\n f\" {unexpected_keys}. \"\n )\n\n # Make sure the trainable params are in float32. This is again needed since the base models\n # are in `weight_dtype`. More details:\n # https://github.com/huggingface/diffusers/pull/6514#discussion_r1449796804\n if args.mixed_precision == \"fp16\":\n # only upcast trainable parameters (LoRA) into fp32\n cast_training_params([transformer_])\n\n accelerator.register_save_state_pre_hook(save_model_hook)\n accelerator.register_load_state_pre_hook(load_model_hook)\n\n # Enable TF32 for faster training on Ampere GPUs,\n # cf https://pytorch.org/docs/stable/notes/cuda.html#tensorfloat-32-tf32-on-ampere-devices\n if args.allow_tf32 and torch.cuda.is_available():\n torch.backends.cuda.matmul.allow_tf32 = True\n\n if args.scale_lr:\n args.learning_rate = (\n args.learning_rate * args.gradient_accumulation_steps * args.train_batch_size * accelerator.num_processes\n )\n\n # Make sure the trainable params are in float32.\n if args.mixed_precision == \"fp16\":\n # only upcast trainable parameters (LoRA) into fp32\n cast_training_params([transformer], dtype=torch.float32)\n\n transformer_lora_parameters = list(filter(lambda p: p.requires_grad, transformer.parameters()))\n\n # Optimization parameters\n transformer_parameters_with_lr = {\"params\": transformer_lora_parameters, \"lr\": args.learning_rate}\n params_to_optimize = [transformer_parameters_with_lr]\n\n use_deepspeed_optimizer = (\n accelerator.state.deepspeed_plugin is not None\n and \"optimizer\" in accelerator.state.deepspeed_plugin.deepspeed_config\n )\n use_deepspeed_scheduler = (\n accelerator.state.deepspeed_plugin is not None\n and \"scheduler\" not in accelerator.state.deepspeed_plugin.deepspeed_config\n )\n\n optimizer = get_optimizer(args, params_to_optimize, use_deepspeed=use_deepspeed_optimizer)\n\n # Dataset and DataLoader\n train_dataset = VideoDataset(\n instance_data_root=args.instance_data_root,\n dataset_name=args.dataset_name,\n dataset_config_name=args.dataset_config_name,\n height=args.height,\n width=args.width,\n fps=args.fps,\n max_num_frames=args.max_num_frames,\n skip_frames_start=args.skip_frames_start,\n skip_frames_end=args.skip_frames_end,\n cache_dir=args.cache_dir,\n id_token=args.id_token,\n )\n\n def encode_video(video):\n video = video.to(accelerator.device, dtype=vae.dtype)\n video = video.permute(0, 2, 1, 3, 4) # [B, C, F, H, W]\n latent_dist = vae.encode(video).latent_dist.sample()\n latent_dist = latent_dist * vae.config.scaling_factor\n return latent_dist\n\n train_dataloader = DataLoader(\n train_dataset,\n batch_size=args.train_batch_size,\n shuffle=True,\n # collate_fn=collate_fn,\n num_workers=args.dataloader_num_workers,\n )\n\n # Scheduler and math around the number of training steps.\n overrode_max_train_steps = False\n num_update_steps_per_epoch = math.ceil(len(train_dataloader) / args.gradient_accumulation_steps)\n if args.max_train_steps is None:\n args.max_train_steps = args.num_train_epochs * num_update_steps_per_epoch\n overrode_max_train_steps = True\n\n if use_deepspeed_scheduler:\n from accelerate.utils import DummyScheduler\n\n lr_scheduler = DummyScheduler(\n name=args.lr_scheduler,\n optimizer=optimizer,\n total_num_steps=args.max_train_steps * accelerator.num_processes,\n num_warmup_steps=args.lr_warmup_steps * accelerator.num_processes,\n )\n else:\n lr_scheduler = get_scheduler(\n args.lr_scheduler,\n optimizer=optimizer,\n num_warmup_steps=args.lr_warmup_steps * accelerator.num_processes,\n num_training_steps=args.max_train_steps * accelerator.num_processes,\n num_cycles=args.lr_num_cycles,\n power=args.lr_power,\n )\n\n # Prepare everything with our `accelerator`.\n transformer, optimizer, train_dataloader, lr_scheduler = accelerator.prepare(\n transformer, optimizer, train_dataloader, lr_scheduler\n )\n\n # We need to recalculate our total training steps as the size of the training dataloader may have changed.\n num_update_steps_per_epoch = math.ceil(len(train_dataloader) / args.gradient_accumulation_steps)\n if overrode_max_train_steps:\n args.max_train_steps = args.num_train_epochs * num_update_steps_per_epoch\n # Afterwards we recalculate our number of training epochs\n args.num_train_epochs = math.ceil(args.max_train_steps / num_update_steps_per_epoch)\n\n # We need to initialize the trackers we use, and also store our configuration.\n # The trackers initializes automatically on the main process.\n if accelerator.is_main_process:\n tracker_name = args.tracker_name or \"cogvideox-lora\"\n accelerator.init_trackers(tracker_name, config=vars(args))\n\n # Train!\n total_batch_size = args.train_batch_size * accelerator.num_processes * args.gradient_accumulation_steps\n num_trainable_parameters = sum(param.numel() for model in params_to_optimize for param in model[\"params\"])\n\n logger.info(\"***** Running training *****\")\n logger.info(f\" Num trainable parameters = {num_trainable_parameters}\")\n logger.info(f\" Num examples = {len(train_dataset)}\")\n logger.info(f\" Num batches each epoch = {len(train_dataloader)}\")\n logger.info(f\" Num epochs = {args.num_train_epochs}\")\n logger.info(f\" Instantaneous batch size per device = {args.train_batch_size}\")\n logger.info(f\" Total train batch size (w. parallel, distributed & accumulation) = {total_batch_size}\")\n logger.info(f\" Gradient accumulation steps = {args.gradient_accumulation_steps}\")\n logger.info(f\" Total optimization steps = {args.max_train_steps}\")\n global_step = 0\n first_epoch = 0\n\n # Potentially load in the weights and states from a previous save\n if not args.resume_from_checkpoint:\n initial_global_step = 0\n else:\n if args.resume_from_checkpoint != \"latest\":\n path = os.path.basename(args.resume_from_checkpoint)\n else:\n # Get the mos recent checkpoint\n dirs = os.listdir(args.output_dir)\n dirs = [d for d in dirs if d.startswith(\"checkpoint\")]\n dirs = sorted(dirs, key=lambda x: int(x.split(\"-\")[1]))\n path = dirs[-1] if len(dirs) > 0 else None\n\n if path is None:\n accelerator.print(\n f\"Checkpoint '{args.resume_from_checkpoint}' does not exist. Starting a new training run.\"\n )\n args.resume_from_checkpoint = None\n initial_global_step = 0\n else:\n accelerator.print(f\"Resuming from checkpoint {path}\")\n accelerator.load_state(os.path.join(args.output_dir, path))\n global_step = int(path.split(\"-\")[1])\n\n initial_global_step = global_step\n first_epoch = global_step // num_update_steps_per_epoch\n\n progress_bar = tqdm(\n range(0, args.max_train_steps),\n initial=initial_global_step,\n desc=\"Steps\",\n # Only show the progress bar once on each machine.\n disable=not accelerator.is_local_main_process,\n )\n vae_scale_factor_spatial = 2 ** (len(vae.config.block_out_channels) - 1)\n\n # For DeepSpeed training\n model_config = transformer.module.config if hasattr(transformer, \"module\") else transformer.config\n\n for epoch in range(first_epoch, args.num_train_epochs):\n transformer.train()\n\n for step, batch in enumerate(train_dataloader):\n models_to_accumulate = [transformer]\n\n with accelerator.accumulate(models_to_accumulate):\n model_input = encode_video(batch[\"video\"]).permute(0, 2, 1, 3, 4).to(dtype=weight_dtype) # [B, F, C, H, W]\n prompts = batch[\"prompt\"]\n\n # encode prompts\n prompt_embeds = compute_prompt_embeddings(\n tokenizer,\n text_encoder,\n prompts,\n model_config.max_text_seq_length,\n accelerator.device,\n weight_dtype,\n requires_grad=False,\n )\n\n # Sample noise that will be added to the latents\n noise = torch.randn_like(model_input)\n batch_size, num_frames, num_channels, height, width = model_input.shape\n\n # Sample a random timestep for each image\n timesteps = torch.randint(\n 0, scheduler.config.num_train_timesteps, (batch_size,), device=model_input.device\n )\n timesteps = timesteps.long()\n\n # Prepare rotary embeds\n image_rotary_emb = (\n prepare_rotary_positional_embeddings(\n height=args.height,\n width=args.width,\n num_frames=num_frames,\n vae_scale_factor_spatial=vae_scale_factor_spatial,\n patch_size=model_config.patch_size,\n attention_head_dim=model_config.attention_head_dim,\n device=accelerator.device,\n )\n if model_config.use_rotary_positional_embeddings\n else None\n )\n\n # Add noise to the model input according to the noise magnitude at each timestep\n # (this is the forward diffusion process)\n noisy_model_input = scheduler.add_noise(model_input, noise, timesteps)\n\n # Predict the noise residual\n model_output = transformer(\n hidden_states=noisy_model_input,\n encoder_hidden_states=prompt_embeds,\n timestep=timesteps,\n image_rotary_emb=image_rotary_emb,\n return_dict=False,\n )[0]\n model_pred = scheduler.get_velocity(model_output, noisy_model_input, timesteps)\n\n alphas_cumprod = scheduler.alphas_cumprod[timesteps]\n weights = 1 / (1 - alphas_cumprod)\n while len(weights.shape) < len(model_pred.shape):\n weights = weights.unsqueeze(-1)\n\n target = model_input\n\n loss = torch.mean((weights * (model_pred - target) ** 2).reshape(batch_size, -1), dim=1)\n loss = loss.mean()\n accelerator.backward(loss)\n\n if accelerator.sync_gradients:\n params_to_clip = transformer.parameters()\n accelerator.clip_grad_norm_(params_to_clip, args.max_grad_norm)\n\n if accelerator.state.deepspeed_plugin is None:\n optimizer.step()\n optimizer.zero_grad()\n\n lr_scheduler.step()\n\n # Checks if the accelerator has performed an optimization step behind the scenes\n if accelerator.sync_gradients:\n progress_bar.update(1)\n global_step += 1\n\n if accelerator.is_main_process:\n if global_step % args.checkpointing_steps == 0:\n # _before_ saving state, check if this save would set us over the `checkpoints_total_limit`\n if args.checkpoints_total_limit is not None:\n checkpoints = os.listdir(args.output_dir)\n checkpoints = [d for d in checkpoints if d.startswith(\"checkpoint\")]\n checkpoints = sorted(checkpoints, key=lambda x: int(x.split(\"-\")[1]))\n\n # before we save the new checkpoint, we need to have at _most_ `checkpoints_total_limit - 1` checkpoints\n if len(checkpoints) >= args.checkpoints_total_limit:\n num_to_remove = len(checkpoints) - args.checkpoints_total_limit + 1\n removing_checkpoints = checkpoints[0:num_to_remove]\n\n logger.info(\n f\"{len(checkpoints)} checkpoints already exist, removing {len(removing_checkpoints)} checkpoints\"\n )\n logger.info(f\"Removing checkpoints: {', '.join(removing_checkpoints)}\")\n\n for removing_checkpoint in removing_checkpoints:\n removing_checkpoint = os.path.join(args.output_dir, removing_checkpoint)\n shutil.rmtree(removing_checkpoint)\n\n save_path = os.path.join(args.output_dir, f\"checkpoint-{global_step}\")\n accelerator.save_state(save_path)\n logger.info(f\"Saved state to {save_path}\")\n\n logs = {\"loss\": loss.detach().item(), \"lr\": lr_scheduler.get_last_lr()[0]}\n progress_bar.set_postfix(**logs)\n accelerator.log(logs, step=global_step)\n\n if global_step >= args.max_train_steps:\n break\n \n if accelerator.is_main_process:\n if args.validation_prompt is not None and (epoch + 1) % args.validation_epochs == 0:\n # Create pipeline\n pipe = CogVideoXPipeline.from_pretrained(\n args.pretrained_model_name_or_path,\n transformer=unwrap_model(transformer),\n text_encoder=unwrap_model(text_encoder),\n vae=unwrap_model(vae),\n scheduler=scheduler,\n revision=args.revision,\n variant=args.variant,\n torch_dtype=weight_dtype,\n )\n\n validation_prompts = args.validation_prompt.split(args.validation_prompt_separator)\n for validation_prompt in validation_prompts:\n pipeline_args = {\n \"prompt\": validation_prompt,\n \"guidance_scale\": args.guidance_scale,\n \"use_dynamic_cfg\": args.use_dynamic_cfg,\n \"height\": args.height,\n \"width\": args.width,\n }\n\n validation_outputs = log_validation(\n pipe=pipe,\n args=args,\n accelerator=accelerator,\n pipeline_args=pipeline_args,\n epoch=epoch,\n )\n\n # Save the lora layers\n accelerator.wait_for_everyone()\n if accelerator.is_main_process:\n transformer = unwrap_model(transformer)\n dtype = (\n torch.float16\n if args.mixed_precision == \"fp16\"\n else torch.bfloat16\n if args.mixed_precision == \"bf16\"\n else torch.float32\n )\n transformer = transformer.to(dtype)\n transformer_lora_layers = get_peft_model_state_dict(transformer)\n\n CogVideoXPipeline.save_lora_weights(\n save_directory=args.output_dir,\n transformer_lora_layers=transformer_lora_layers,\n )\n\n # Final test inference\n pipe = CogVideoXPipeline.from_pretrained(\n args.pretrained_model_name_or_path,\n revision=args.revision,\n variant=args.variant,\n torch_dtype=weight_dtype,\n )\n pipe.scheduler = CogVideoXDPMScheduler.from_config(pipe.scheduler.config)\n\n if args.enable_slicing:\n pipe.vae.enable_slicing()\n if args.enable_tiling:\n pipe.vae.enable_tiling()\n\n # Load LoRA weights\n lora_scaling = args.lora_alpha / args.rank\n pipe.load_lora_weights(args.output_dir, adapter_name=\"cogvideox-lora\")\n pipe.set_adapters([\"cogvideox-lora\"], [lora_scaling])\n\n # Run inference\n validation_outputs = []\n if args.validation_prompt and args.num_validation_videos > 0:\n validation_prompts = args.validation_prompt.split(args.validation_prompt_separator)\n for validation_prompt in validation_prompts:\n pipeline_args = {\n \"prompt\": validation_prompt,\n \"guidance_scale\": args.guidance_scale,\n \"use_dynamic_cfg\": args.use_dynamic_cfg,\n \"height\": args.height,\n \"width\": args.width,\n }\n\n video = log_validation(\n pipe=pipe,\n args=args,\n accelerator=accelerator,\n pipeline_args=pipeline_args,\n epoch=epoch,\n is_final_validation=True,\n )\n validation_outputs.extend(video)\n\n if args.push_to_hub:\n save_model_card(\n repo_id,\n videos=validation_outputs,\n base_model=args.pretrained_model_name_or_path,\n validation_prompt=args.validation_prompt,\n repo_folder=args.output_dir,\n fps=args.fps,\n )\n upload_folder(\n repo_id=repo_id,\n folder_path=args.output_dir,\n commit_message=\"End of training\",\n ignore_patterns=[\"step_*\", \"epoch_*\"],\n )\n\n accelerator.end_training()\n\n\nif __name__ == \"__main__\":\n args = get_args()\n main(args)" }, { "path": "CogVideo/inference/3dtrajmaster_inference.py", "content": "\"\"\"\nAdapted from CogVideoX-5B: https://github.com/THUDM/CogVideo by Xiao Fu (CUHK)\n\"\"\"\n\nimport argparse\nfrom typing import Literal\nimport copy\n\nimport torch\nfrom diffusers import (\n CogVideoXDDIMScheduler,\n CogVideoXDPMScheduler,\n CogVideoXImageToVideoPipeline,\n CogVideoXVideoToVideoPipeline,\n)\nimport os\nimport json\n\nimport datetime\nimport sys\nsys.path.append('../finetune')\n\nfrom models.pipeline_cogvideox import CogVideoXPipeline\nfrom models.cogvideox_transformer_3d import CogVideoXTransformer3DModel\n\nfrom diffusers.utils import export_to_video, load_image, load_video\nimport json\nimport numpy as np\nimport random\nfrom einops import rearrange\n\ndef parse_matrix(matrix_str):\n rows = matrix_str.strip().split('] [')\n matrix = []\n for row in rows:\n row = row.replace('[', '').replace(']', '')\n matrix.append(list(map(float, row.split())))\n return np.array(matrix)\n\ndef load_sceneposes(objs_file, obj_idx, obj_transl):\n ext_poses = []\n for i, key in enumerate(objs_file.keys()):\n ext_poses.append(parse_matrix(objs_file[key][obj_idx]['matrix']))\n ext_poses = np.stack(ext_poses)\n ext_poses = np.transpose(ext_poses, (0,2,1))\n ext_poses[:,:3,3] -= obj_transl\n ext_poses[:,:3,3] /= 100.\n ext_poses = ext_poses[:, :, [1,2,0,3]]\n return ext_poses\n\ndef get_pose_embeds(scene, video_name, instance_data_root, locations_info, cam_poses):\n\n with open(os.path.join(instance_data_root, \"480_720\", scene, video_name, video_name+'.json'), 'r') as f: objs_file = json.load(f)\n objs_num = len(objs_file['0'])\n video_index = 12\n\n location_name = video_name.split('_')[1]\n location_info = locations_info[location_name]\n cam_pose = cam_poses[video_index-1]\n obj_transl = location_info['coordinates']['CameraTarget']['position']\n\n obj_poses_list = []\n for obj_idx in range(objs_num):\n obj_poses = load_sceneposes(objs_file, obj_idx, obj_transl)\n obj_poses = np.linalg.inv(cam_pose) @ obj_poses\n obj_poses_list.append(obj_poses)\n\n obj_poses_all = torch.from_numpy(np.array(obj_poses_list))\n \n total_frames = 99\n sample_n_frames = 49\n current_sample_stride = 1.75\n start_frame_ind = 10\n\n cropped_length = int(sample_n_frames * current_sample_stride)\n end_frame_ind = min(start_frame_ind + cropped_length, total_frames)\n frame_indices = np.linspace(start_frame_ind, end_frame_ind - 1, sample_n_frames, dtype=int)\n\n # interpolation\n trunc_frame_indices = np.zeros_like(frame_indices[::4])\n trunc_frame_indices[0] = frame_indices[0]\n trunc_frame_indices[1:] = ((frame_indices[1:][::4] + frame_indices[4:][::4])/2).astype(np.int64)\n\n obj_poses_all = obj_poses_all[:, trunc_frame_indices]\n pose_embeds = rearrange(obj_poses_all[:, :, :3, :], \"n f p q -> n f (q p)\").contiguous().to(torch.bfloat16)\n\n return pose_embeds\n \n\ndef init_cam_poses(instance_data_root):\n cam_num = 12\n cams_path = os.path.join(instance_data_root, \"Hemi12_transforms.json\")\n with open(cams_path, 'r') as f: cams_info = json.load(f)\n cam_poses = []\n for i, key in enumerate(cams_info.keys()):\n if \"C_\" in key:\n cam_poses.append(parse_matrix(cams_info[key]))\n cam_poses = np.stack(cam_poses)\n cam_poses = np.transpose(cam_poses, (0,2,1))\n cam_poses = cam_poses[:,:,[1,2,0,3]]\n cam_poses[:,:3,3] /= 100.\n return cam_poses\n\ndef generate_video(\n model_path: str,\n ckpt_path: str,\n lora_path: str = None,\n lora_scale: float = 1.0,\n output_path: str = \"./output.mp4\",\n image_or_video_path: str = \"\",\n annealed_sample_step: int = 15,\n num_inference_steps: int = 50,\n guidance_scale: float = 6.0,\n num_videos_per_prompt: int = 1,\n dtype: torch.dtype = torch.bfloat16,\n generate_type: str = Literal[\"t2v\", \"i2v\", \"v2v\"], # i2v: image to video, v2v: video to video\n seed: int = 42,\n):\n \"\"\"\n Generates a video based on the given prompt and saves it to the specified path.\n\n Parameters:\n - prompt (str): The description of the video to be generated.\n - model_path (str): The path of the pre-trained model to be used.\n - lora_path (str): The path of the LoRA weights to be used.\n - lora_scale (float): \n - output_path (str): The path where the generated video will be saved.\n - num_inference_steps (int): Number of steps for the inference process. More steps can result in better quality.\n - guidance_scale (float): The scale for classifier-free guidance. Higher values can lead to better alignment with the prompt.\n - num_videos_per_prompt (int): Number of videos to generate per prompt.\n - dtype (torch.dtype): The data type for computation (default is torch.bfloat16).\n - generate_type (str): The type of video generation (e.g., 't2v', 'i2v', 'v2v').·\n - seed (int): The seed for reproducibility.\n \"\"\"\n \n # 1. Load the pre-trained CogVideoX pipeline with the specified precision (bfloat16).\n # add device_map=\"balanced\" in the from_pretrained function and remove the enable_model_cpu_offload()\n # function to use Multi GPUs.\n\n transformer = CogVideoXTransformer3DModel.from_pretrained(ckpt_path, torch_dtype=dtype)\n pipe = CogVideoXPipeline.from_pretrained(model_path, \n transformer=transformer,\n torch_dtype=dtype\n )\n pipe.transformer_ori = copy.deepcopy(pipe.transformer).to(\"cuda\")\n\n # If you're using with lora, add this code\n if lora_path:\n pipe.load_lora_weights(lora_path, weight_name=\"pytorch_lora_weights.safetensors\", adapter_name=\"default\")\n pipe.fuse_lora(components=['transformer'] ,lora_scale=lora_scale)\n \n # 2. Set Scheduler.\n pipe.scheduler = CogVideoXDPMScheduler.from_config(pipe.scheduler.config, timestep_spacing=\"trailing\")\n\n # 3. Enable CPU offload for the model.\n pipe.to(\"cuda\")\n pipe.enable_sequential_cpu_offload()\n pipe.vae.enable_slicing()\n pipe.vae.enable_tiling()\n \n # 4. Load object poses, scene and object prompts\n instance_data_root = \"/m2v_intern/fuxiao/360Motion-Dataset\"\n scene = \"Desert\"\n locations_path = os.path.join(instance_data_root, \"480_720\", scene, \"location_data.json\")\n with open(locations_path, 'r') as f: locations = json.load(f)\n locations_info = {locations[idx]['name']:locations[idx] for idx in range(len(locations))}\n cam_poses = init_cam_poses(instance_data_root)\n\n video_names = os.listdir(os.path.join(instance_data_root, \"480_720\", scene))\n video_names.remove('location_data.json')\n \n with open('./test_sets.json', 'r') as f: test_sets = json.load(f)\n\n for idx in range(len(test_sets)):\n \n eval_set = test_sets[str(idx)]\n video_caption_list = eval_set['entity_prompts']\n objs_num = len(video_caption_list) \n location = eval_set['loc_prompt']\n video_name = eval_set['video_name']\n\n pose_embeds = get_pose_embeds(scene, video_name, instance_data_root, locations_info, cam_poses)\n\n prompt = \"\"\n for obj_idx in range(objs_num):\n video_caption = video_caption_list[obj_idx]\n if obj_idx == objs_num - 1:\n if objs_num == 1:\n prompt += video_caption + ' is moving in the ' + location\n else:\n prompt += video_caption + ' are moving in the ' + location\n else:\n prompt += video_caption + ' and '\n\n # 5. Generate the video frames based on the prompt.\n video_generate = pipe(\n prompt=prompt,\n prompts_list=video_caption_list,\n pose_embeds=pose_embeds[None],\n num_videos_per_prompt=num_videos_per_prompt,\n annealed_sample_step=annealed_sample_step,\n num_inference_steps=num_inference_steps,\n num_frames=49,\n use_dynamic_cfg=True,\n guidance_scale=guidance_scale,\n generator=torch.Generator().manual_seed(seed),\n ).frames[0]\n\n # 6. Export the generated frames to a video file. fps must be 8 for original video.\n save_video_name = ''\n save_video_name += str(objs_num) + '_' + video_name + '_' + location + '_'\n for obj_idx in range(objs_num):\n video_caption = video_caption_list[obj_idx][:30]\n video_caption = video_caption.replace(' ', '_')\n save_video_name += video_caption + '_'\n save_video_name += '.mp4'\n save_video_name = save_video_name.replace('_.mp4', '.mp4')\n save_video_path = os.path.join(output_path, save_video_name)\n export_to_video(video_generate, save_video_path, fps=8)\n\n with open(save_video_path.replace('.mp4', '.txt'), 'a+') as f:\n f.write(video_name)\n f.write('\\n')\n for obj_idx in range(objs_num):\n f.write(video_caption_list[obj_idx])\n f.write('\\n')\n f.write(location)\n f.write('\\n')\n\n\nif __name__ == \"__main__\":\n parser = argparse.ArgumentParser(description=\"Generate a video from a text prompt using CogVideoX\")\n parser.add_argument(\n \"--model_path\", type=str, default=\"THUDM/CogVideoX-5b\", help=\"The path of the pre-trained model to be used\"\n )\n parser.add_argument(\n \"--ckpt_path\", type=str, default=\"THUDM/CogVideoX-5b\", help=\"The path of the pre-trained transformer to be used\"\n )\n parser.add_argument(\"--lora_path\", type=str, default=None, help=\"The path of the LoRA weights to be used\")\n parser.add_argument(\"--lora_scale\", type=float, default=1.0)\n parser.add_argument(\n \"--output_path\", type=str, default=\"./output.mp4\", help=\"The path where the generated video will be saved\"\n )\n parser.add_argument(\"--guidance_scale\", type=float, default=6.0, help=\"The scale for classifier-free guidance\")\n parser.add_argument(\n \"--num_inference_steps\", type=int, default=50, help=\"Number of steps for the inference process\"\n )\n parser.add_argument(\"--annealed_sample_step\", type=int, default=15, help=\"Number of videos to generate per prompt\")\n parser.add_argument(\"--num_videos_per_prompt\", type=int, default=1, help=\"Number of videos to generate per prompt\")\n parser.add_argument(\n \"--generate_type\", type=str, default=\"t2v\", help=\"The type of video generation (e.g., 't2v', 'i2v', 'v2v')\"\n )\n parser.add_argument(\n \"--dtype\", type=str, default=\"bfloat16\", help=\"The data type for computation (e.g., 'float16' or 'bfloat16')\"\n )\n parser.add_argument(\"--seed\", type=int, default=42, help=\"The seed for reproducibility\")\n\n args = parser.parse_args()\n dtype = torch.float16 if args.dtype == \"float16\" else torch.bfloat16\n os.makedirs(args.output_path, exist_ok=True)\n generate_video(\n model_path=args.model_path,\n ckpt_path=args.ckpt_path,\n lora_path=args.lora_path,\n lora_scale=args.lora_scale,\n output_path=args.output_path,\n annealed_sample_step=args.annealed_sample_step,\n num_inference_steps=args.num_inference_steps,\n guidance_scale=args.guidance_scale,\n num_videos_per_prompt=args.num_videos_per_prompt,\n dtype=dtype,\n generate_type=args.generate_type,\n seed=args.seed,\n )\n" }, { "path": "CogVideo/inference/entity_zoo.txt", "content": "[\n \"a fire spirit with long, twisting flames resembling flowing red and orange hair, a bright yellow core\",\n \"a gentle breeze with soft tendrils of pale blue mist resembling flowing fabric, delicate streaks of white vapor\",\n \"an ember haze with flickering orange and red flames giving off a warm glow, smoldering dark red wisps\",\n \"a luminous storm cloud with layers of deep gray and blue churning together, intermittent flashes of lightning\",\n \"a storm entity with dark swirling clouds as a body, streaks of electric blue lightning shooting across it\",\n \"a cloud creature with billowing white and gray plumes forming a soft, rounded body, wisps of darker fog\",\n \"a foggy apparition with pale gray wisps drifting together in a soft, undefined form, tiny white sparkles\",\n \"a man with short spiky brown hair, athletic build, a navy blue jacket, beige cargo pants, and black sneakers\",\n \"a woman with long wavy blonde hair, petite figure, a red floral dress, white sandals, and a yellow shoulder bag\",\n \"a man with a shaved head, broad shoulders, a gray graphic t-shirt, dark jeans, and brown leather boots\",\n \"a woman with shoulder-length straight auburn hair, a slender figure, a green button-up blouse, black leggings, white sneakers\",\n \"a man with messy black hair, tall frame, a plaid red and black shirt, faded blue jeans, and tan hiking boots\",\n \"a man with medium-length straight brown hair, tall and slender, a gray crew-neck t-shirt, beige trousers, dark green sneakers\",\n \"a woman with short curly black hair, slender build, a pink hoodie, light gray joggers, and blue sneakers\",\n \"a man with short black wavy hair, lean figure, a green and yellow plaid shirt, dark brown pants, and black suede shoes\",\n \"a man with curly black hair, muscular build, a dark green hoodie, gray joggers, and white running shoes\",\n \"a woman with short blonde hair, slim athletic build, a red leather jacket, dark blue jeans, and white sneakers\",\n \"a man with medium-length wavy brown hair, lean build, a black bomber jacket, olive green cargo pants, and brown hiking boots\",\n \"a man with buzz-cut blonde hair, stocky build, a gray zip-up sweater, black shorts, and red basketball shoes\",\n \"a woman with long straight black hair, toned build, a blue denim jacket, light gray leggings, and black slip-on shoes\",\n \"a man with short curly red hair, average build, a black leather jacket, dark blue cargo pants, and white sneakers\",\n \"a woman with shoulder-length wavy brown hair, slim build, a green parka, black leggings, and gray hiking boots\",\n \"a man with short straight black hair, tall and lean build, a navy blue sweater, khaki shorts, and brown sandals\",\n \"a woman with pixie-cut blonde hair, athletic build, a red windbreaker, blue ripped jeans, and black combat boots\",\n \"a man with medium-length wavy gray hair, muscular build, a maroon t-shirt, beige chinos, and brown loafers\",\n \"a woman with long curly black hair, average build, a purple hoodie, black athletic shorts, and white running shoes\",\n \"a man with short spiky blonde hair, slim build, a black trench coat, blue jeans, and brown hiking shoes\",\n \"a dog with a fluffy coat, wagging tail, and warm golden-brown fur, exuding a gentle and friendly charm\",\n \"a tiger with vibrant orange and black stripes, piercing yellow eyes, and a powerful stance, exuding strength and grace\",\n \"a giraffe with golden-yellow fur, long legs, a tall slender neck, and patches of brown spots, exuding elegance and calm\",\n \"an alpaca with soft white wool, short legs, a thick neck, and a fluffy head of fur, radiating gentle charm\",\n \"a zebra with black and white stripes, sturdy legs, a short neck, and a sleek mane running down its back\",\n \"a deer with sleek tan fur, long slender legs, a graceful neck, and tiny antlers atop its head\",\n \"a gazelle with light golden fur, long slender legs, a thin neck, and short, sharp horns, embodying elegance and agility\",\n \"a horse with chestnut brown fur, muscular legs, a slim neck, and a flowing mane, exuding strength and grace\",\n \"a sleek black panther with a smooth, glossy coat, emerald green eyes, and a powerful stance\",\n \"a cheetah with golden fur covered in black spots, intense amber eyes, and a slender, agile body\",\n \"a regal lion with a thick, flowing golden mane, sharp brown eyes, and a powerful muscular frame\",\n \"a snow leopard with pale gray fur adorned with dark rosettes, icy blue eyes, and a stealthy, poised posture\",\n \"a jaguar with a golden-yellow coat dotted with intricate black rosettes, deep green eyes, and a muscular build\",\n \"a wolf with thick silver-gray fur, alert golden eyes, and a lean yet strong body, exuding confidence and boldness\",\n \"a tiger with a pristine white coat marked by bold black stripes, bright blue eyes, and a graceful, poised form\",\n \"a lynx with tufted ears, soft reddish-brown fur with faint spots, and intense yellow-green eyes\",\n \"a bear with dark brown fur, small but fierce black eyes, and a broad and muscular build, radiating power\",\n \"a swift fox with reddish-orange fur, a bushy tail tipped with white, and sharp, intelligent amber eyes\",\n \"a falcon with blue-gray feathers, sharp talons, and keen yellow eyes fixed on its prey below\",\n \"a fox with sleek russet fur, a bushy tail tipped with black, and bright green and cunning eyes\",\n \"a kangaroo with brown fur, powerful hind legs, and a muscular tail, showcasing its strength and agility\",\n \"a polar bear with thick white fur, strong paws, and a black nose, embodying the essence of the Arctic\",\n \"a cheetah with a slender build, spotted golden fur, and sharp eyes, epitomizing speed and agility\",\n \"a dolphin with sleek grey skin, a curved dorsal fin, and intelligent, playful eyes, reflecting its nature\",\n \"a wolf with a body covered in thick silver fur, sharp ears, and piercing yellow eyes, showcasing its alertness\",\n \"a leopard with a body covered in golden fur, dark rosettes, and a long muscular tail, emphasizing its strength\",\n \"a penguin with a body covered in smooth black-and-white feathers, short wings, and webbed feet\",\n \"a gazelle with a body covered in sleek tan fur, long legs, and elegant curved horns, showcasing its grace\",\n \"a rabbit with a body covered in soft fur, quick hops, and a playful demeanor, showcasing its energy\",\n \"a koala with a body covered in soft grey fur, large round ears, and a black nose, radiating cuteness\",\n \"a rhinoceros with a body covered in thick grey skin, a massive horn on its snout, and sturdy legs\",\n \"a flamingo with a body covered in pink feathers, long slender legs, and a gracefully curved neck\",\n \"a parrot with bright red, blue, and yellow feathers, a curved beak, and sharp intelligent eyes\",\n \"a hippopotamus with a body covered in thick grey-brown skin, massive jaws, and a large body\",\n \"a crocodile with a body covered in scaly green skin, a powerful tail, and sharp teeth\",\n \"a moose with a body covered in thick brown fur, massive antlers, and a bulky frame\",\n \"a chameleon with a body covered in vibrant green scales, bulging eyes, and a curled tail, showcasing its unique charm\",\n \"a lemur with a body covered in soft grey fur, a ringed tail, and wide yellow eyes, and curious expression\",\n \"a squirrel with a body covered in bushy red fur, large eyes, and a fluffy tail\",\n \"a panda with a body covered in fluffy black-and-white fur, a round face, and gentle eyes, radiating warmth\",\n \"a porcupine with a body covered in spiky brown quills, a small nose, and curious eyes\",\n \"a sedan with a sleek metallic silver body, long wheelbase, a low-profile hood, and a small rear spoiler\",\n \"a private jet with a shiny silver body, elongated wings, a slim nose, and a compact rear stabilizer\",\n \"an SUV with a matte black exterior, elevated suspension, a tall roofline, and a compact rear roof rack\",\n \"a pickup truck with rugged dark green paint, extended cab, raised suspension, and a modest cargo bed cover\",\n \"a vintage convertible with a body covered in shiny red paint, chrome bumpers, and a stylish design\",\n \"a futuristic electric car with a minimalist silver design, slim LED lights, and smooth curves\",\n \"a family minivan with a spacious interior, sliding doors, and a metallic blue exterior\",\n \"a compact electric vehicle with a silver finish, aerodynamic profile, and efficient battery\",\n \"a sporty roadster with a convertible top, silver trim, and a powerful engine\",\n \"a retro coupe with a body covered in teal paint, round headlights, and a shiny chrome grille\",\n \"a firefighting robot with a water cannon arm, heat sensors, and durable red-and-silver exterior\",\n \"a companion robot with a friendly digital face, a smooth white exterior, and social interaction algorithms\",\n \"an industrial welding robot with articulated arms, a laser precision welder, and heat-resistant shields\",\n \"a surveillance drone robot with extendable camera arms, thermal vision, and a stealth black body\",\n \"a disaster rescue robot with reinforced limbs, advanced AI, and a rugged body designed to navigate\",\n \"an exploration rover robot with solar panels, durable wheels, and advanced sensors for planetary exploration\",\n \"a fluttering butterfly with intricate wing patterns, vivid colors, and graceful flight\",\n]" }, { "path": "CogVideo/inference/location_zoo.txt", "content": "[\n 'fjord',\n 'sunset beach',\n 'cave',\n 'snowy tundra',\n 'prairie',\n 'asian town',\n 'rainforest',\n 'canyon',\n 'savanna',\n 'urban rooftop garden',\n 'swamp',\n 'riverbank',\n 'coral reef',\n 'volcanic landscape',\n 'wind farm',\n 'town street',\n 'night city square',\n 'mall lobby',\n 'glacier',\n 'seaside street',\n 'gymnastics room',\n 'abandoned factory',\n 'autumn forest',\n 'mountain village',\n 'coastal harbor',\n 'ancient ruins',\n 'modern metropolis',\n 'desert',\n 'forest',\n 'city',\n 'snowy street',\n 'park',\n]" }, { "path": "CogVideo/pyproject.toml", "content": "[tool.ruff]\nline-length = 119\n\n[tool.ruff.lint]\n# Never enforce `E501` (line length violations).\nignore = [\"C901\", \"E501\", \"E741\", \"F402\", \"F823\"]\nselect = [\"C\", \"E\", \"F\", \"I\", \"W\"]\n\n# Ignore import violations in all `__init__.py` files.\n[tool.ruff.lint.per-file-ignores]\n\"__init__.py\" = [\"E402\", \"F401\", \"F403\", \"F811\"]\n\n[tool.ruff.lint.isort]\nlines-after-imports = 2\n\n[tool.ruff.format]\n# Like Black, use double quotes for strings.\nquote-style = \"double\"\n\n# Like Black, indent with spaces, rather than tabs.\nindent-style = \"space\"\n\n# Like Black, respect magic trailing commas.\nskip-magic-trailing-comma = false\n\n# Like Black, automatically detect the appropriate line ending.\nline-ending = \"auto\"\n" }, { "path": "CogVideo/requirements.txt", "content": "diffusers==0.31.0\naccelerate==1.1.1\ntransformers==4.46.2\nnumpy==1.26.0\n# torch==2.5.0\n# torchvision==0.20.0\nsentencepiece==0.2.0\nSwissArmyTransformer==0.4.12\ngradio==5.5.0\nimageio==2.35.1\nimageio-ffmpeg==0.5.1\nopenai==1.54.0\nmoviepy==1.0.3\nscikit-video==1.1.11\nopencv-python\npeft==0.12.0\ndecord\nwandb" }, { "path": "CogVideo/tools/caption/README.md", "content": "# Video Caption\n\nTypically, most video data does not come with corresponding descriptive text, so it is necessary to convert the video\ndata into textual descriptions to provide the essential training data for text-to-video models.\n\n## Update and News\n- 🔥🔥 **News**: ```2024/9/19```: The caption model used in the CogVideoX training process to convert video data into text\n descriptions, [CogVLM2-Caption](https://huggingface.co/THUDM/cogvlm2-llama3-caption), is now open-source. Feel\n free to download and use it.\n\n\n## Video Caption via CogVLM2-Caption\n\n🤗 [Hugging Face](https://huggingface.co/THUDM/cogvlm2-llama3-caption) | 🤖 [ModelScope](https://modelscope.cn/models/ZhipuAI/cogvlm2-llama3-caption/) \n\nCogVLM2-Caption is a video captioning model used to generate training data for the CogVideoX model.\n\n### Install\n```shell\npip install -r requirements.txt\n```\n\n### Usage\n\n```shell\npython video_caption.py\n```\n\nExample:\n
\n \n
\n\n## Video Caption via CogVLM2-Video\n\n[Code](https://github.com/THUDM/CogVLM2/tree/main/video_demo) | 🤗 [Hugging Face](https://huggingface.co/THUDM/cogvlm2-video-llama3-chat) | 🤖 [ModelScope](https://modelscope.cn/models/ZhipuAI/cogvlm2-video-llama3-chat) | 📑 [Blog](https://cogvlm2-video.github.io/) ｜ [💬 Online Demo](http://cogvlm2-online.cogviewai.cn:7868/)\n\nCogVLM2-Video is a versatile video understanding model equipped with timestamp-based question answering capabilities.\nUsers can input prompts such as `Please describe this video in detail.` to the model to obtain a detailed video caption:\n
\n \n
\n\nUsers can use the provided [code](https://github.com/THUDM/CogVLM2/tree/main/video_demo) to load the model or configure a RESTful API to generate video captions.\n\n## Citation\n\n🌟 If you find our work helpful, please leave us a star and cite our paper.\n\nCogVLM2-Caption:\n```\n@article{yang2024cogvideox,\n title={CogVideoX: Text-to-Video Diffusion Models with An Expert Transformer},\n author={Yang, Zhuoyi and Teng, Jiayan and Zheng, Wendi and Ding, Ming and Huang, Shiyu and Xu, Jiazheng and Yang, Yuanming and Hong, Wenyi and Zhang, Xiaohan and Feng, Guanyu and others},\n journal={arXiv preprint arXiv:2408.06072},\n year={2024}\n}\n```\nCogVLM2-Video:\n```\n@article{hong2024cogvlm2,\n title={CogVLM2: Visual Language Models for Image and Video Understanding},\n author={Hong, Wenyi and Wang, Weihan and Ding, Ming and Yu, Wenmeng and Lv, Qingsong and Wang, Yan and Cheng, Yean and Huang, Shiyu and Ji, Junhui and Xue, Zhao and others},\n journal={arXiv preprint arXiv:2408.16500},\n year={2024}\n}\n```" }, { "path": "CogVideo/tools/caption/README_ja.md", "content": "# ビデオキャプション\n\n通常、ほとんどのビデオデータには対応する説明文が付いていないため、ビデオデータをテキストの説明に変換して、テキストからビデオへのモデルに必要なトレーニングデータを提供する必要があります。\n\n## 更新とニュース\n- 🔥🔥 **ニュース**: ```2024/9/19```：CogVideoX\n のトレーニングプロセスで、ビデオデータをテキストに変換するためのキャプションモデル [CogVLM2-Caption](https://huggingface.co/THUDM/cogvlm2-llama3-caption)\n がオープンソース化されました。ぜひダウンロードしてご利用ください。\n## CogVLM2-Captionによるビデオキャプション\n\n🤗 [Hugging Face](https://huggingface.co/THUDM/cogvlm2-llama3-caption) | 🤖 [ModelScope](https://modelscope.cn/models/ZhipuAI/cogvlm2-llama3-caption/) \n\nCogVLM2-Captionは、CogVideoXモデルのトレーニングデータを生成するために使用されるビデオキャプションモデルです。\n\n### インストール\n```shell\npip install -r requirements.txt\n```\n\n### 使用方法\n```shell\npython video_caption.py\n```\n\n例:\n
\n \n
\n\n\n\n## CogVLM2-Video を使用したビデオキャプション\n\n[Code](https://github.com/THUDM/CogVLM2/tree/main/video_demo) | 🤗 [Hugging Face](https://huggingface.co/THUDM/cogvlm2-video-llama3-chat) | 🤖 [ModelScope](https://modelscope.cn/models/ZhipuAI/cogvlm2-video-llama3-chat) | 📑 [Blog](https://cogvlm2-video.github.io/) ｜ [💬 Online Demo](http://cogvlm2-online.cogviewai.cn:7868/)\n\n\nCogVLM2-Video は、タイムスタンプベースの質問応答機能を備えた多機能なビデオ理解モデルです。ユーザーは `このビデオを詳細に説明してください。` などのプロンプトをモデルに入力して、詳細なビデオキャプションを取得できます：\n
\n \n
\n\nユーザーは提供された[コード](https://github.com/THUDM/CogVLM2/tree/main/video_demo)を使用してモデルをロードするか、RESTful API を構成してビデオキャプションを生成できます。\n\n## Citation\n\n🌟 If you find our work helpful, please leave us a star and cite our paper.\n\nCogVLM2-Caption:\n```\n@article{yang2024cogvideox,\n title={CogVideoX: Text-to-Video Diffusion Models with An Expert Transformer},\n author={Yang, Zhuoyi and Teng, Jiayan and Zheng, Wendi and Ding, Ming and Huang, Shiyu and Xu, Jiazheng and Yang, Yuanming and Hong, Wenyi and Zhang, Xiaohan and Feng, Guanyu and others},\n journal={arXiv preprint arXiv:2408.06072},\n year={2024}\n}\n```\nCogVLM2-Video:\n```\n@article{hong2024cogvlm2,\n title={CogVLM2: Visual Language Models for Image and Video Understanding},\n author={Hong, Wenyi and Wang, Weihan and Ding, Ming and Yu, Wenmeng and Lv, Qingsong and Wang, Yan and Cheng, Yean and Huang, Shiyu and Ji, Junhui and Xue, Zhao and others},\n journal={arXiv preprint arXiv:2408.16500},\n year={2024}\n}\n```\n" }, { "path": "CogVideo/tools/caption/README_zh.md", "content": "# 视频Caption\n\n通常，大多数视频数据不带有相应的描述性文本，因此需要将视频数据转换为文本描述，以提供必要的训练数据用于文本到视频模型。\n\n## 项目更新\n- 🔥🔥 **News**: ```2024/9/19```: CogVideoX 训练过程中用于将视频数据转换为文本描述的 Caption\n 模型 [CogVLM2-Caption](https://huggingface.co/THUDM/cogvlm2-llama3-caption)\n 已经开源。欢迎前往下载并使用。\n\n## 通过 CogVLM2-Caption 模型生成视频Caption\n\n🤗 [Hugging Face](https://huggingface.co/THUDM/cogvlm2-llama3-caption) | 🤖 [ModelScope](https://modelscope.cn/models/ZhipuAI/cogvlm2-llama3-caption/) \n\nCogVLM2-Caption是用于生成CogVideoX模型训练数据的视频caption模型。\n\n### 安装依赖\n```shell\npip install -r requirements.txt\n```\n\n### 运行caption模型\n\n```shell\npython video_caption.py\n```\n\n示例：\n
\n \n
\n\n## 通过 CogVLM2-Video 模型生成视频Caption\n\n[Code](https://github.com/THUDM/CogVLM2/tree/main/video_demo) | 🤗 [Hugging Face](https://huggingface.co/THUDM/cogvlm2-video-llama3-chat) | 🤖 [ModelScope](https://modelscope.cn/models/ZhipuAI/cogvlm2-video-llama3-chat) | 📑 [Blog](https://cogvlm2-video.github.io/) ｜ [💬 Online Demo](http://cogvlm2-online.cogviewai.cn:7868/)\n\nCogVLM2-Video 是一个多功能的视频理解模型，具备基于时间戳的问题回答能力。用户可以输入诸如 `Describe this video in detail.` 的提示语给模型，以获得详细的视频Caption：\n\n\n
\n \n
\n\n用户可以使用提供的[代码](https://github.com/THUDM/CogVLM2/tree/main/video_demo)加载模型或配置 RESTful API 来生成视频Caption。\n\n\n## Citation\n\n🌟 If you find our work helpful, please leave us a star and cite our paper.\n\nCogVLM2-Caption:\n```\n@article{yang2024cogvideox,\n title={CogVideoX: Text-to-Video Diffusion Models with An Expert Transformer},\n author={Yang, Zhuoyi and Teng, Jiayan and Zheng, Wendi and Ding, Ming and Huang, Shiyu and Xu, Jiazheng and Yang, Yuanming and Hong, Wenyi and Zhang, Xiaohan and Feng, Guanyu and others},\n journal={arXiv preprint arXiv:2408.06072},\n year={2024}\n}\n```\nCogVLM2-Video:\n```\n@article{hong2024cogvlm2,\n title={CogVLM2: Visual Language Models for Image and Video Understanding},\n author={Hong, Wenyi and Wang, Weihan and Ding, Ming and Yu, Wenmeng and Lv, Qingsong and Wang, Yan and Cheng, Yean and Huang, Shiyu and Ji, Junhui and Xue, Zhao and others},\n journal={arXiv preprint arXiv:2408.16500},\n year={2024}\n}\n```" }, { "path": "CogVideo/tools/caption/requirements.txt", "content": "decord>=0.6.0\n#根据https://download.pytorch.org/whl/torch/，python版本为[3.8,3.11]\ntorch==2.1.0\ntorchvision== 0.16.0\npytorchvideo==0.1.5\nxformers\ntransformers==4.42.4\n#git+https://github.com/huggingface/transformers.git\nhuggingface-hub>=0.23.0\npillow\nchainlit>=1.0\npydantic>=2.7.1\ntimm>=0.9.16\nopenai>=1.30.1\nloguru>=0.7.2\npydantic>=2.7.1\neinops\nsse-starlette>=2.1.0\nflask\ngunicorn\ngevent\nrequests\ngradio" }, { "path": "CogVideo/tools/caption/video_caption.py", "content": "import io\n\nimport argparse\nimport numpy as np\nimport torch\nfrom decord import cpu, VideoReader, bridge\nfrom transformers import AutoModelForCausalLM, AutoTokenizer\n\nMODEL_PATH = \"THUDM/cogvlm2-llama3-caption\"\n\nDEVICE = 'cuda' if torch.cuda.is_available() else 'cpu'\nTORCH_TYPE = torch.bfloat16 if torch.cuda.is_available() and torch.cuda.get_device_capability()[\n 0] >= 8 else torch.float16\n\nparser = argparse.ArgumentParser(description=\"CogVLM2-Video CLI Demo\")\nparser.add_argument('--quant', type=int, choices=[4, 8], help='Enable 4-bit or 8-bit precision loading', default=0)\nargs = parser.parse_args([])\n\n\ndef load_video(video_data, strategy='chat'):\n bridge.set_bridge('torch')\n mp4_stream = video_data\n num_frames = 24\n decord_vr = VideoReader(io.BytesIO(mp4_stream), ctx=cpu(0))\n\n frame_id_list = None\n total_frames = len(decord_vr)\n if strategy == 'base':\n clip_end_sec = 60\n clip_start_sec = 0\n start_frame = int(clip_start_sec * decord_vr.get_avg_fps())\n end_frame = min(total_frames,\n int(clip_end_sec * decord_vr.get_avg_fps())) if clip_end_sec is not None else total_frames\n frame_id_list = np.linspace(start_frame, end_frame - 1, num_frames, dtype=int)\n elif strategy == 'chat':\n timestamps = decord_vr.get_frame_timestamp(np.arange(total_frames))\n timestamps = [i[0] for i in timestamps]\n max_second = round(max(timestamps)) + 1\n frame_id_list = []\n for second in range(max_second):\n closest_num = min(timestamps, key=lambda x: abs(x - second))\n index = timestamps.index(closest_num)\n frame_id_list.append(index)\n if len(frame_id_list) >= num_frames:\n break\n\n video_data = decord_vr.get_batch(frame_id_list)\n video_data = video_data.permute(3, 0, 1, 2)\n return video_data\n\n\ntokenizer = AutoTokenizer.from_pretrained(\n MODEL_PATH,\n trust_remote_code=True,\n)\n\nmodel = AutoModelForCausalLM.from_pretrained(\n MODEL_PATH,\n torch_dtype=TORCH_TYPE,\n trust_remote_code=True\n).eval().to(DEVICE)\n\n\ndef predict(prompt, video_data, temperature):\n strategy = 'chat'\n\n video = load_video(video_data, strategy=strategy)\n\n history = []\n query = prompt\n inputs = model.build_conversation_input_ids(\n tokenizer=tokenizer,\n query=query,\n images=[video],\n history=history,\n template_version=strategy\n )\n inputs = {\n 'input_ids': inputs['input_ids'].unsqueeze(0).to('cuda'),\n 'token_type_ids': inputs['token_type_ids'].unsqueeze(0).to('cuda'),\n 'attention_mask': inputs['attention_mask'].unsqueeze(0).to('cuda'),\n 'images': [[inputs['images'][0].to('cuda').to(TORCH_TYPE)]],\n }\n gen_kwargs = {\n \"max_new_tokens\": 2048,\n \"pad_token_id\": 128002,\n \"top_k\": 1,\n \"do_sample\": False,\n \"top_p\": 0.1,\n \"temperature\": temperature,\n }\n with torch.no_grad():\n outputs = model.generate(**inputs, **gen_kwargs)\n outputs = outputs[:, inputs['input_ids'].shape[1]:]\n response = tokenizer.decode(outputs[0], skip_special_tokens=True)\n return response\n\n\ndef test():\n prompt = \"Please describe this video in detail.\"\n temperature = 0.1\n video_data = open('test.mp4', 'rb').read()\n response = predict(prompt, video_data, temperature)\n print(response)\n\n\nif __name__ == '__main__':\n test()\n" }, { "path": "CogVideo/tools/convert_weight_sat2hf.py", "content": "\"\"\"\nThis script demonstrates how to convert and generate video from a text prompt\nusing CogVideoX with 🤗Huggingface Diffusers Pipeline.\nThis script requires the `diffusers>=0.30.2` library to be installed.\n\nFunctions:\n - reassign_query_key_value_inplace: Reassigns the query, key, and value weights in-place.\n - reassign_query_key_layernorm_inplace: Reassigns layer normalization for query and key in-place.\n - reassign_adaln_norm_inplace: Reassigns adaptive layer normalization in-place.\n - remove_keys_inplace: Removes specified keys from the state_dict in-place.\n - replace_up_keys_inplace: Replaces keys in the \"up\" block in-place.\n - get_state_dict: Extracts the state_dict from a saved checkpoint.\n - update_state_dict_inplace: Updates the state_dict with new key assignments in-place.\n - convert_transformer: Converts a transformer checkpoint to the CogVideoX format.\n - convert_vae: Converts a VAE checkpoint to the CogVideoX format.\n - get_args: Parses command-line arguments for the script.\n - generate_video: Generates a video from a text prompt using the CogVideoX pipeline.\n\"\"\"\n\nimport argparse\nfrom typing import Any, Dict\n\nimport torch\nfrom transformers import T5EncoderModel, T5Tokenizer\n\nfrom diffusers import (\n AutoencoderKLCogVideoX,\n CogVideoXDDIMScheduler,\n CogVideoXImageToVideoPipeline,\n CogVideoXPipeline,\n CogVideoXTransformer3DModel,\n)\n\n\ndef reassign_query_key_value_inplace(key: str, state_dict: Dict[str, Any]):\n to_q_key = key.replace(\"query_key_value\", \"to_q\")\n to_k_key = key.replace(\"query_key_value\", \"to_k\")\n to_v_key = key.replace(\"query_key_value\", \"to_v\")\n to_q, to_k, to_v = torch.chunk(state_dict[key], chunks=3, dim=0)\n state_dict[to_q_key] = to_q\n state_dict[to_k_key] = to_k\n state_dict[to_v_key] = to_v\n state_dict.pop(key)\n\n\ndef reassign_query_key_layernorm_inplace(key: str, state_dict: Dict[str, Any]):\n layer_id, weight_or_bias = key.split(\".\")[-2:]\n\n if \"query\" in key:\n new_key = f\"transformer_blocks.{layer_id}.attn1.norm_q.{weight_or_bias}\"\n elif \"key\" in key:\n new_key = f\"transformer_blocks.{layer_id}.attn1.norm_k.{weight_or_bias}\"\n\n state_dict[new_key] = state_dict.pop(key)\n\n\ndef reassign_adaln_norm_inplace(key: str, state_dict: Dict[str, Any]):\n layer_id, _, weight_or_bias = key.split(\".\")[-3:]\n\n weights_or_biases = state_dict[key].chunk(12, dim=0)\n norm1_weights_or_biases = torch.cat(weights_or_biases[0:3] + weights_or_biases[6:9])\n norm2_weights_or_biases = torch.cat(weights_or_biases[3:6] + weights_or_biases[9:12])\n\n norm1_key = f\"transformer_blocks.{layer_id}.norm1.linear.{weight_or_bias}\"\n state_dict[norm1_key] = norm1_weights_or_biases\n\n norm2_key = f\"transformer_blocks.{layer_id}.norm2.linear.{weight_or_bias}\"\n state_dict[norm2_key] = norm2_weights_or_biases\n\n state_dict.pop(key)\n\n\ndef remove_keys_inplace(key: str, state_dict: Dict[str, Any]):\n state_dict.pop(key)\n\n\ndef replace_up_keys_inplace(key: str, state_dict: Dict[str, Any]):\n key_split = key.split(\".\")\n layer_index = int(key_split[2])\n replace_layer_index = 4 - 1 - layer_index\n\n key_split[1] = \"up_blocks\"\n key_split[2] = str(replace_layer_index)\n new_key = \".\".join(key_split)\n\n state_dict[new_key] = state_dict.pop(key)\n\n\nTRANSFORMER_KEYS_RENAME_DICT = {\n \"transformer.final_layernorm\": \"norm_final\",\n \"transformer\": \"transformer_blocks\",\n \"attention\": \"attn1\",\n \"mlp\": \"ff.net\",\n \"dense_h_to_4h\": \"0.proj\",\n \"dense_4h_to_h\": \"2\",\n \".layers\": \"\",\n \"dense\": \"to_out.0\",\n \"input_layernorm\": \"norm1.norm\",\n \"post_attn1_layernorm\": \"norm2.norm\",\n \"time_embed.0\": \"time_embedding.linear_1\",\n \"time_embed.2\": \"time_embedding.linear_2\",\n \"mixins.patch_embed\": \"patch_embed\",\n \"mixins.final_layer.norm_final\": \"norm_out.norm\",\n \"mixins.final_layer.linear\": \"proj_out\",\n \"mixins.final_layer.adaLN_modulation.1\": \"norm_out.linear\",\n \"mixins.pos_embed.pos_embedding\": \"patch_embed.pos_embedding\", # Specific to CogVideoX-5b-I2V\n}\n\nTRANSFORMER_SPECIAL_KEYS_REMAP = {\n \"query_key_value\": reassign_query_key_value_inplace,\n \"query_layernorm_list\": reassign_query_key_layernorm_inplace,\n \"key_layernorm_list\": reassign_query_key_layernorm_inplace,\n \"adaln_layer.adaLN_modulations\": reassign_adaln_norm_inplace,\n \"embed_tokens\": remove_keys_inplace,\n \"freqs_sin\": remove_keys_inplace,\n \"freqs_cos\": remove_keys_inplace,\n \"position_embedding\": remove_keys_inplace,\n}\n\nVAE_KEYS_RENAME_DICT = {\n \"block.\": \"resnets.\",\n \"down.\": \"down_blocks.\",\n \"downsample\": \"downsamplers.0\",\n \"upsample\": \"upsamplers.0\",\n \"nin_shortcut\": \"conv_shortcut\",\n \"encoder.mid.block_1\": \"encoder.mid_block.resnets.0\",\n \"encoder.mid.block_2\": \"encoder.mid_block.resnets.1\",\n \"decoder.mid.block_1\": \"decoder.mid_block.resnets.0\",\n \"decoder.mid.block_2\": \"decoder.mid_block.resnets.1\",\n}\n\nVAE_SPECIAL_KEYS_REMAP = {\n \"loss\": remove_keys_inplace,\n \"up.\": replace_up_keys_inplace,\n}\n\nTOKENIZER_MAX_LENGTH = 226\n\n\ndef get_state_dict(saved_dict: Dict[str, Any]) -> Dict[str, Any]:\n state_dict = saved_dict\n if \"model\" in saved_dict.keys():\n state_dict = state_dict[\"model\"]\n if \"module\" in saved_dict.keys():\n state_dict = state_dict[\"module\"]\n if \"state_dict\" in saved_dict.keys():\n state_dict = state_dict[\"state_dict\"]\n return state_dict\n\n\ndef update_state_dict_inplace(state_dict: Dict[str, Any], old_key: str, new_key: str) -> Dict[str, Any]:\n state_dict[new_key] = state_dict.pop(old_key)\n\n\ndef convert_transformer(\n ckpt_path: str,\n num_layers: int,\n num_attention_heads: int,\n use_rotary_positional_embeddings: bool,\n i2v: bool,\n dtype: torch.dtype,\n):\n PREFIX_KEY = \"model.diffusion_model.\"\n\n original_state_dict = get_state_dict(torch.load(ckpt_path, map_location=\"cpu\", mmap=True))\n transformer = CogVideoXTransformer3DModel(\n in_channels=32 if i2v else 16,\n num_layers=num_layers,\n num_attention_heads=num_attention_heads,\n use_rotary_positional_embeddings=use_rotary_positional_embeddings,\n use_learned_positional_embeddings=i2v,\n ).to(dtype=dtype)\n\n for key in list(original_state_dict.keys()):\n new_key = key[len(PREFIX_KEY) :]\n for replace_key, rename_key in TRANSFORMER_KEYS_RENAME_DICT.items():\n new_key = new_key.replace(replace_key, rename_key)\n update_state_dict_inplace(original_state_dict, key, new_key)\n\n for key in list(original_state_dict.keys()):\n for special_key, handler_fn_inplace in TRANSFORMER_SPECIAL_KEYS_REMAP.items():\n if special_key not in key:\n continue\n handler_fn_inplace(key, original_state_dict)\n transformer.load_state_dict(original_state_dict, strict=True)\n return transformer\n\n\ndef convert_vae(ckpt_path: str, scaling_factor: float, dtype: torch.dtype):\n original_state_dict = get_state_dict(torch.load(ckpt_path, map_location=\"cpu\", mmap=True))\n vae = AutoencoderKLCogVideoX(scaling_factor=scaling_factor).to(dtype=dtype)\n\n for key in list(original_state_dict.keys()):\n new_key = key[:]\n for replace_key, rename_key in VAE_KEYS_RENAME_DICT.items():\n new_key = new_key.replace(replace_key, rename_key)\n update_state_dict_inplace(original_state_dict, key, new_key)\n\n for key in list(original_state_dict.keys()):\n for special_key, handler_fn_inplace in VAE_SPECIAL_KEYS_REMAP.items():\n if special_key not in key:\n continue\n handler_fn_inplace(key, original_state_dict)\n\n vae.load_state_dict(original_state_dict, strict=True)\n return vae\n\n\ndef get_args():\n parser = argparse.ArgumentParser()\n parser.add_argument(\n \"--transformer_ckpt_path\", type=str, default=None, help=\"Path to original transformer checkpoint\")\n parser.add_argument(\"--vae_ckpt_path\", type=str, default=None, help=\"Path to original vae checkpoint\")\n parser.add_argument(\"--output_path\", type=str, required=True, help=\"Path where converted model should be saved\")\n parser.add_argument(\"--fp16\", action=\"store_true\", default=False, help=\"Whether to save the model weights in fp16\")\n parser.add_argument(\"--bf16\", action=\"store_true\", default=False, help=\"Whether to save the model weights in bf16\")\n parser.add_argument(\n \"--push_to_hub\", action=\"store_true\", default=False, help=\"Whether to push to HF Hub after saving\"\n )\n parser.add_argument(\n \"--text_encoder_cache_dir\", type=str, default=None, help=\"Path to text encoder cache directory\"\n )\n # For CogVideoX-2B, num_layers is 30. For 5B, it is 42\n parser.add_argument(\"--num_layers\", type=int, default=30, help=\"Number of transformer blocks\")\n # For CogVideoX-2B, num_attention_heads is 30. For 5B, it is 48\n parser.add_argument(\"--num_attention_heads\", type=int, default=30, help=\"Number of attention heads\")\n # For CogVideoX-2B, use_rotary_positional_embeddings is False. For 5B, it is True\n parser.add_argument(\n \"--use_rotary_positional_embeddings\", action=\"store_true\", default=False, help=\"Whether to use RoPE or not\"\n )\n # For CogVideoX-2B, scaling_factor is 1.15258426. For 5B, it is 0.7\n parser.add_argument(\"--scaling_factor\", type=float, default=1.15258426, help=\"Scaling factor in the VAE\")\n # For CogVideoX-2B, snr_shift_scale is 3.0. For 5B, it is 1.0\n parser.add_argument(\"--snr_shift_scale\", type=float, default=3.0, help=\"Scaling factor in the VAE\")\n parser.add_argument(\"--i2v\", action=\"store_true\", default=False, help=\"Whether to save the model weights in fp16\")\n return parser.parse_args()\n\n\nif __name__ == \"__main__\":\n args = get_args()\n\n transformer = None\n vae = None\n\n if args.fp16 and args.bf16:\n raise ValueError(\"You cannot pass both --fp16 and --bf16 at the same time.\")\n\n dtype = torch.float16 if args.fp16 else torch.bfloat16 if args.bf16 else torch.float32\n\n if args.transformer_ckpt_path is not None:\n transformer = convert_transformer(\n args.transformer_ckpt_path,\n args.num_layers,\n args.num_attention_heads,\n args.use_rotary_positional_embeddings,\n args.i2v,\n dtype,\n )\n if args.vae_ckpt_path is not None:\n vae = convert_vae(args.vae_ckpt_path, args.scaling_factor, dtype)\n\n text_encoder_id = \"google/t5-v1_1-xxl\"\n tokenizer = T5Tokenizer.from_pretrained(text_encoder_id, model_max_length=TOKENIZER_MAX_LENGTH)\n text_encoder = T5EncoderModel.from_pretrained(text_encoder_id, cache_dir=args.text_encoder_cache_dir)\n # Apparently, the conversion does not work anymore without this :shrug:\n for param in text_encoder.parameters():\n param.data = param.data.contiguous()\n\n scheduler = CogVideoXDDIMScheduler.from_config(\n {\n \"snr_shift_scale\": args.snr_shift_scale,\n \"beta_end\": 0.012,\n \"beta_schedule\": \"scaled_linear\",\n \"beta_start\": 0.00085,\n \"clip_sample\": False,\n \"num_train_timesteps\": 1000,\n \"prediction_type\": \"v_prediction\",\n \"rescale_betas_zero_snr\": True,\n \"set_alpha_to_one\": True,\n \"timestep_spacing\": \"trailing\",\n }\n )\n if args.i2v:\n pipeline_cls = CogVideoXImageToVideoPipeline\n else:\n pipeline_cls = CogVideoXPipeline\n\n pipe = pipeline_cls(\n tokenizer=tokenizer,\n text_encoder=text_encoder,\n vae=vae,\n transformer=transformer,\n scheduler=scheduler,\n )\n\n if args.fp16:\n pipe = pipe.to(dtype=torch.float16)\n if args.bf16:\n pipe = pipe.to(dtype=torch.bfloat16)\n\n # We don't use variant here because the model must be run in fp16 (2B) or bf16 (5B). It would be weird\n # for users to specify variant when the default is not fp32 and they want to run with the correct default (which\n # is either fp16/bf16 here).\n pipe.save_pretrained(args.output_path, safe_serialization=True, push_to_hub=args.push_to_hub)\n" }, { "path": "CogVideo/tools/export_sat_lora_weight.py", "content": "from typing import Any, Dict\nimport torch \nimport argparse \nfrom diffusers.loaders.lora_base import LoraBaseMixin\nfrom diffusers.models.modeling_utils import load_state_dict\n\n\ndef get_state_dict(saved_dict: Dict[str, Any]) -> Dict[str, Any]:\n state_dict = saved_dict\n if \"model\" in saved_dict.keys():\n state_dict = state_dict[\"model\"]\n if \"module\" in saved_dict.keys():\n state_dict = state_dict[\"module\"]\n if \"state_dict\" in saved_dict.keys():\n state_dict = state_dict[\"state_dict\"]\n return state_dict\n\nLORA_KEYS_RENAME = {\n\n 'attention.query_key_value.matrix_A.0': 'attn1.to_q.lora_A.weight',\n 'attention.query_key_value.matrix_A.1': 'attn1.to_k.lora_A.weight',\n 'attention.query_key_value.matrix_A.2': 'attn1.to_v.lora_A.weight',\n 'attention.query_key_value.matrix_B.0': 'attn1.to_q.lora_B.weight',\n 'attention.query_key_value.matrix_B.1': 'attn1.to_k.lora_B.weight',\n 'attention.query_key_value.matrix_B.2': 'attn1.to_v.lora_B.weight',\n 'attention.dense.matrix_A.0': 'attn1.to_out.0.lora_A.weight',\n 'attention.dense.matrix_B.0': 'attn1.to_out.0.lora_B.weight'\n}\n\n\n\nPREFIX_KEY = \"model.diffusion_model.\"\nSAT_UNIT_KEY = \"layers\"\nLORA_PREFIX_KEY = \"transformer_blocks\"\n\n\n\ndef export_lora_weight(ckpt_path,lora_save_directory):\n\n merge_original_state_dict = get_state_dict(torch.load(ckpt_path, map_location=\"cpu\", mmap=True))\n\n\n lora_state_dict = {}\n for key in list(merge_original_state_dict.keys()):\n new_key = key[len(PREFIX_KEY) :]\n for special_key, lora_keys in LORA_KEYS_RENAME.items():\n if new_key.endswith(special_key):\n new_key = new_key.replace(special_key, lora_keys)\n new_key = new_key.replace(SAT_UNIT_KEY, LORA_PREFIX_KEY)\n\n lora_state_dict[new_key] = merge_original_state_dict[key]\n\n\n\n # final length should be 240 \n if len(lora_state_dict) != 240:\n raise ValueError(\"lora_state_dict length is not 240\")\n\n lora_state_dict.keys()\n\n LoraBaseMixin.write_lora_layers(\n state_dict=lora_state_dict,\n save_directory=lora_save_directory,\n is_main_process=True,\n weight_name=None,\n save_function=None,\n safe_serialization=True\n )\n\n\ndef get_args():\n parser = argparse.ArgumentParser()\n parser.add_argument(\n \"--sat_pt_path\", type=str, required=True, help=\"Path to original sat transformer checkpoint\"\n )\n parser.add_argument(\"--lora_save_directory\", type=str, required=True, help=\"Path where converted lora should be saved\") \n return parser.parse_args()\n\n\nif __name__ == \"__main__\":\n args = get_args()\n\n export_lora_weight(args.sat_pt_path, args.lora_save_directory)\n" }, { "path": "CogVideo/tools/llm_flux_cogvideox/generate.sh", "content": "#!/bin/bash\n\nNUM_VIDEOS=10\nINFERENCE_STEPS=50\nGUIDANCE_SCALE=7.0\nOUTPUT_DIR_PREFIX=\"outputs/gpu_\"\nLOG_DIR_PREFIX=\"logs/gpu_\"\n\nVIDEO_MODEL_PATH=\"/share/official_pretrains/hf_home/CogVideoX-5b-I2V\"\nLLM_MODEL_PATH=\"/share/home/zyx/Models/Meta-Llama-3.1-8B-Instruct\"\nIMAGE_MODEL_PATH = \"share/home/zyx/Models/FLUX.1-dev\"\n\n#VIDEO_MODEL_PATH=\"THUDM/CogVideoX-5B-I2V\"\n#LLM_MODEL_PATH=\"THUDM/glm-4-9b-chat\"\n#IMAGE_MODEL_PATH = \"black-forest-labs/FLUX.1-dev\"\n\nCUDA_DEVICES=${CUDA_VISIBLE_DEVICES:-\"0\"}\n\nIFS=',' read -r -a GPU_ARRAY <<< \"$CUDA_DEVICES\"\n\nfor i in \"${!GPU_ARRAY[@]}\"\ndo\n GPU=${GPU_ARRAY[$i]}\n echo \"Starting task on GPU $GPU...\"\n CUDA_VISIBLE_DEVICES=$GPU nohup python3 llm_flux_cogvideox.py \\\n --caption_generator_model_id $LLM_MODEL_PATH \\\n --image_generator_model_id $IMAGE_MODEL_PATH \\\n --model_path $VIDEO_MODEL_PATH \\\n --num_videos $NUM_VIDEOS \\\n --image_generator_num_inference_steps $INFERENCE_STEPS \\\n --guidance_scale $GUIDANCE_SCALE \\\n --use_dynamic_cfg \\\n --output_dir ${OUTPUT_DIR_PREFIX}${GPU} \\\n > ${LOG_DIR_PREFIX}${GPU}.log 2>&1 &\ndone" }, { "path": "CogVideo/tools/llm_flux_cogvideox/gradio_page.py", "content": "import os\nimport gradio as gr\nimport gc\nimport random\nimport torch\nimport numpy as np\nfrom PIL import Image\nimport transformers\nfrom diffusers import CogVideoXImageToVideoPipeline, CogVideoXDPMScheduler, DiffusionPipeline\nfrom diffusers.utils import export_to_video\nfrom transformers import AutoTokenizer\nfrom datetime import datetime, timedelta\nimport threading\nimport time\nimport moviepy.editor as mp\n\ntorch.set_float32_matmul_precision(\"high\")\n\n# Set default values\ncaption_generator_model_id = \"/share/home/zyx/Models/Meta-Llama-3.1-8B-Instruct\"\nimage_generator_model_id = \"/share/home/zyx/Models/FLUX.1-dev\"\nvideo_generator_model_id = \"/share/official_pretrains/hf_home/CogVideoX-5b-I2V\"\nseed = 1337\n\nos.makedirs(\"./output\", exist_ok=True)\nos.makedirs(\"./gradio_tmp\", exist_ok=True)\n\ntokenizer = AutoTokenizer.from_pretrained(caption_generator_model_id, trust_remote_code=True)\ncaption_generator = transformers.pipeline(\n \"text-generation\",\n model=caption_generator_model_id,\n device_map=\"balanced\",\n model_kwargs={\n \"local_files_only\": True,\n \"torch_dtype\": torch.bfloat16,\n },\n trust_remote_code=True,\n tokenizer=tokenizer\n)\n\nimage_generator = DiffusionPipeline.from_pretrained(\n image_generator_model_id,\n torch_dtype=torch.bfloat16,\n device_map=\"balanced\"\n)\n# image_generator.to(\"cuda\")\n\nvideo_generator = CogVideoXImageToVideoPipeline.from_pretrained(\n video_generator_model_id,\n torch_dtype=torch.bfloat16,\n device_map=\"balanced\"\n)\n\nvideo_generator.vae.enable_slicing()\nvideo_generator.vae.enable_tiling()\n\nvideo_generator.scheduler = CogVideoXDPMScheduler.from_config(\n video_generator.scheduler.config, timestep_spacing=\"trailing\"\n)\n\n# Define prompts\nSYSTEM_PROMPT = \"\"\"\nYou are part of a team of people that create videos using generative models. You use a video-generation model that can generate a video about anything you describe.\n\nFor example, if you respond with \"A beautiful morning in the woods with the sun peaking through the trees\", the video generation model will create a video of exactly as described. Your task is to summarize the descriptions of videos provided by users and create detailed prompts to feed into the generative model.\n\nThere are a few rules to follow:\n- You will only ever output a single video description per request.\n- If the user mentions to summarize the prompt in [X] words, make sure not to exceed the limit.\n\nYour responses should just be the video generation prompt. Here are examples:\n- \"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.\"\n- \"A street artist, clad in a worn-out denim jacket and a colorful bandana, stands before a vast concrete wall in the heart of the city, holding a can of spray paint, spray-painting a colorful bird on a mottled wall.\"\n\"\"\".strip()\n\nUSER_PROMPT = \"\"\"\nCould you generate a prompt for a video generation model? Please limit the prompt to [{0}] words.\n\"\"\".strip()\n\n\ndef generate_caption(prompt):\n num_words = random.choice([25, 50, 75, 100])\n user_prompt = USER_PROMPT.format(num_words)\n\n messages = [\n {\"role\": \"system\", \"content\": SYSTEM_PROMPT},\n {\"role\": \"user\", \"content\": prompt + \"\\n\" + user_prompt},\n ]\n\n response = caption_generator(\n messages,\n max_new_tokens=226,\n return_full_text=False\n )\n caption = response[0][\"generated_text\"]\n if caption.startswith(\"\\\"\") and caption.endswith(\"\\\"\"):\n caption = caption[1:-1]\n return caption\n\n\ndef generate_image(caption, progress=gr.Progress(track_tqdm=True)):\n image = image_generator(\n prompt=caption,\n height=480,\n width=720,\n num_inference_steps=30,\n guidance_scale=3.5,\n ).images[0]\n return image, image # One for output One for State\n\n\ndef generate_video(\n caption,\n image,\n progress=gr.Progress(track_tqdm=True)\n):\n generator = torch.Generator().manual_seed(seed)\n video_frames = video_generator(\n image=image,\n prompt=caption,\n height=480,\n width=720,\n num_frames=49,\n num_inference_steps=50,\n guidance_scale=6,\n use_dynamic_cfg=True,\n generator=generator,\n ).frames[0]\n video_path = save_video(video_frames)\n gif_path = convert_to_gif(video_path)\n return video_path, gif_path\n\n\ndef save_video(tensor):\n timestamp = datetime.now().strftime(\"%Y%m%d_%H%M%S\")\n video_path = f\"./output/{timestamp}.mp4\"\n os.makedirs(os.path.dirname(video_path), exist_ok=True)\n export_to_video(tensor, video_path, fps=8)\n return video_path\n\n\ndef convert_to_gif(video_path):\n clip = mp.VideoFileClip(video_path)\n clip = clip.set_fps(8)\n clip = clip.resize(height=240)\n gif_path = video_path.replace(\".mp4\", \".gif\")\n clip.write_gif(gif_path, fps=8)\n return gif_path\n\n\ndef delete_old_files():\n while True:\n now = datetime.now()\n cutoff = now - timedelta(minutes=10)\n directories = [\"./output\", \"./gradio_tmp\"]\n\n for directory in directories:\n for filename in os.listdir(directory):\n file_path = os.path.join(directory, filename)\n if os.path.isfile(file_path):\n file_mtime = datetime.fromtimestamp(os.path.getmtime(file_path))\n if file_mtime < cutoff:\n os.remove(file_path)\n time.sleep(600)\n\n\nthreading.Thread(target=delete_old_files, daemon=True).start()\n\nwith gr.Blocks() as demo:\n gr.Markdown(\"\"\"\n
\n LLM + FLUX + CogVideoX-I2V Space 🤗\n
\n \"\"\")\n with gr.Row():\n with gr.Column():\n prompt = gr.Textbox(label=\"Prompt\", placeholder=\"Enter your prompt here\", lines=5)\n generate_caption_button = gr.Button(\"Generate Caption\")\n caption = gr.Textbox(label=\"Caption\", placeholder=\"Caption will appear here\", lines=5)\n generate_image_button = gr.Button(\"Generate Image\")\n image_output = gr.Image(label=\"Generated Image\")\n state_image = gr.State()\n generate_caption_button.click(fn=generate_caption, inputs=prompt, outputs=caption)\n generate_image_button.click(fn=generate_image, inputs=caption, outputs=[image_output, state_image])\n with gr.Column():\n video_output = gr.Video(label=\"Generated Video\", width=720, height=480)\n download_video_button = gr.File(label=\"📥 Download Video\", visible=False)\n download_gif_button = gr.File(label=\"📥 Download GIF\", visible=False)\n generate_video_button = gr.Button(\"Generate Video from Image\")\n generate_video_button.click(fn=generate_video, inputs=[caption, state_image],\n outputs=[video_output, download_gif_button])\n\nif __name__ == \"__main__\":\n demo.launch()\n" }, { "path": "CogVideo/tools/llm_flux_cogvideox/llm_flux_cogvideox.py", "content": "\"\"\"\nThe original experimental code for this project can be found at:\n\nhttps://gist.github.com/a-r-r-o-w/d070cce059ab4ceab3a9f289ff83c69c\n\nBy using this code, description prompts will be generated through a local large language model, and images will be\ngenerated using the black-forest-labs/FLUX.1-dev model, followed by video generation via CogVideoX.\nThe entire process utilizes open-source solutions, without the need for any API keys.\n\nYou can use the generate.sh file in the same folder to automate running this code\nfor batch generation of videos and images.\n\nbash generate.sh\n\n\"\"\"\n\nimport argparse\nimport gc\nimport json\nimport os\nimport pathlib\nimport random\nfrom typing import Any, Dict\n\nfrom transformers import AutoTokenizer\n\nos.environ[\"TORCH_LOGS\"] = \"+dynamo,recompiles,graph_breaks\"\nos.environ[\"TORCHDYNAMO_VERBOSE\"] = \"1\"\n\nimport numpy as np\nimport torch\nimport transformers\nfrom diffusers import CogVideoXImageToVideoPipeline, CogVideoXDPMScheduler, DiffusionPipeline\nfrom diffusers.utils.logging import get_logger\nfrom diffusers.utils import export_to_video\n\ntorch.set_float32_matmul_precision(\"high\")\n\nlogger = get_logger(__name__)\n\nSYSTEM_PROMPT = \"\"\"\nYou are part of a team of people that create videos using generative models. You use a video-generation model that can generate a video about anything you describe.\n\nFor example, if you respond with \"A beautiful morning in the woods with the sun peaking through the trees\", the video generation model will create a video of exactly as described. You task is to summarize the descriptions of videos provided to by users, and create details prompts to feed into the generative model.\n\nThere are a few rules to follow:\n- You will only ever output a single video description per request.\n- If the user mentions to summarize the prompt in [X] words, make sure to not exceed the limit.\n\nYou responses should just be the video generation prompt. Here are examples:\n- “A lone figure stands on a city rooftop at night, gazing up at the full moon. The moon glows brightly, casting a gentle light over the quiet cityscape. Below, the windows of countless homes shine with warm lights, creating a contrast between the bustling life below and the peaceful solitude above. The scene captures the essence of the Mid-Autumn Festival, where despite the distance, the figure feels connected to loved ones through the shared beauty of the moonlit sky.”\n- \"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.\"\n- \"A street artist, clad in a worn-out denim jacket and a colorful banana, stands before a vast concrete wall in the heart, holding a can of spray paint, spray-painting a colorful bird on a mottled wall\"\n\"\"\".strip()\n\nUSER_PROMPT = \"\"\"\nCould you generate a prompt for a video generation model? \nPlease limit the prompt to [{0}] words.\n\"\"\".strip()\n\n\ndef get_args():\n parser = argparse.ArgumentParser()\n parser.add_argument(\n \"--num_videos\",\n type=int,\n default=5,\n help=\"Number of unique videos you would like to generate.\"\n )\n parser.add_argument(\n \"--model_path\",\n type=str,\n default=\"THUDM/CogVideoX-5B\",\n help=\"The path of Image2Video CogVideoX-5B\",\n )\n parser.add_argument(\n \"--caption_generator_model_id\",\n type=str,\n default=\"THUDM/glm-4-9b-chat\",\n help=\"Caption generation model. default GLM-4-9B\",\n )\n parser.add_argument(\n \"--caption_generator_cache_dir\",\n type=str,\n default=None,\n help=\"Cache directory for caption generation model.\"\n )\n parser.add_argument(\n \"--image_generator_model_id\",\n type=str,\n default=\"black-forest-labs/FLUX.1-dev\",\n help=\"Image generation model.\"\n )\n parser.add_argument(\n \"--image_generator_cache_dir\",\n type=str,\n default=None,\n help=\"Cache directory for image generation model.\"\n )\n parser.add_argument(\n \"--image_generator_num_inference_steps\",\n type=int,\n default=50,\n help=\"Caption generation model.\"\n )\n parser.add_argument(\n \"--guidance_scale\",\n type=float,\n default=7,\n help=\"Guidance scale to be use for generation.\"\n )\n parser.add_argument(\n \"--use_dynamic_cfg\",\n action=\"store_true\",\n help=\"Whether or not to use cosine dynamic guidance for generation [Recommended].\",\n )\n parser.add_argument(\n \"--output_dir\",\n type=str,\n default=\"outputs/\",\n help=\"Location where generated images and videos should be stored.\",\n )\n parser.add_argument(\n \"--compile\",\n action=\"store_true\",\n help=\"Whether or not to compile the transformer of image and video generators.\"\n )\n parser.add_argument(\n \"--enable_vae_tiling\",\n action=\"store_true\",\n help=\"Whether or not to use VAE tiling when encoding/decoding.\"\n )\n parser.add_argument(\n \"--seed\",\n type=int,\n default=42,\n help=\"Seed for reproducibility.\"\n )\n return parser.parse_args()\n\n\ndef reset_memory():\n gc.collect()\n torch.cuda.empty_cache()\n torch.cuda.reset_peak_memory_stats()\n torch.cuda.reset_accumulated_memory_stats()\n\n\n@torch.no_grad()\ndef main(args: Dict[str, Any]) -> None:\n output_dir = pathlib.Path(args.output_dir)\n os.makedirs(output_dir.as_posix(), exist_ok=True)\n\n random.seed(args.seed)\n np.random.seed(args.seed)\n torch.manual_seed(args.seed)\n torch.cuda.manual_seed_all(args.seed)\n\n reset_memory()\n tokenizer = AutoTokenizer.from_pretrained(args.caption_generator_model_id, trust_remote_code=True)\n caption_generator = transformers.pipeline(\n \"text-generation\",\n model=args.caption_generator_model_id,\n device_map=\"auto\",\n model_kwargs={\n \"local_files_only\": True,\n \"cache_dir\": args.caption_generator_cache_dir,\n \"torch_dtype\": torch.bfloat16,\n },\n trust_remote_code=True,\n tokenizer=tokenizer\n )\n\n captions = []\n for i in range(args.num_videos):\n num_words = random.choice([50, 75, 100])\n user_prompt = USER_PROMPT.format(num_words)\n\n messages = [\n {\"role\": \"system\", \"content\": SYSTEM_PROMPT},\n {\"role\": \"user\", \"content\": user_prompt},\n ]\n\n outputs = caption_generator(messages, max_new_tokens=226)\n caption = outputs[0][\"generated_text\"][-1][\"content\"]\n if caption.startswith(\"\\\"\") and caption.endswith(\"\\\"\"):\n caption = caption[1:-1]\n captions.append(caption)\n logger.info(f\"Generated caption: {caption}\")\n\n with open(output_dir / \"captions.json\", \"w\") as file:\n json.dump(captions, file)\n\n del caption_generator\n reset_memory()\n\n image_generator = DiffusionPipeline.from_pretrained(\n args.image_generator_model_id,\n cache_dir=args.image_generator_cache_dir,\n torch_dtype=torch.bfloat16\n )\n image_generator.to(\"cuda\")\n\n if args.compile:\n image_generator.transformer = torch.compile(image_generator.transformer, mode=\"max-autotune\", fullgraph=True)\n\n if args.enable_vae_tiling:\n image_generator.vae.enable_tiling()\n\n images = []\n for index, caption in enumerate(captions):\n image = image_generator(\n prompt=caption,\n height=480,\n width=720,\n num_inference_steps=args.image_generator_num_inference_steps,\n guidance_scale=3.5,\n ).images[0]\n filename = caption[:25].replace(\".\", \"_\").replace(\"'\", \"_\").replace('\"', \"_\").replace(\",\", \"_\")\n image.save(output_dir / f\"{index}_{filename}.png\")\n images.append(image)\n\n del image_generator\n reset_memory()\n\n video_generator = CogVideoXImageToVideoPipeline.from_pretrained(\n args.model_path, torch_dtype=torch.bfloat16).to(\"cuda\")\n video_generator.scheduler = CogVideoXDPMScheduler.from_config(\n video_generator.scheduler.config,\n timestep_spacing=\"trailing\")\n\n if args.compile:\n video_generator.transformer = torch.compile(video_generator.transformer, mode=\"max-autotune\", fullgraph=True)\n\n if args.enable_vae_tiling:\n video_generator.vae.enable_tiling()\n\n generator = torch.Generator().manual_seed(args.seed)\n for index, (caption, image) in enumerate(zip(captions, images)):\n video = video_generator(\n image=image,\n prompt=caption,\n height=480,\n width=720,\n num_frames=49,\n num_inference_steps=50,\n guidance_scale=args.guidance_scale,\n use_dynamic_cfg=args.use_dynamic_cfg,\n generator=generator,\n ).frames[0]\n filename = caption[:25].replace(\".\", \"_\").replace(\"'\", \"_\").replace('\"', \"_\").replace(\",\", \"_\")\n export_to_video(video, output_dir / f\"{index}_{filename}.mp4\", fps=8)\n\n\nif __name__ == \"__main__\":\n args = get_args()\n main(args)\n" }, { "path": "CogVideo/tools/load_cogvideox_lora.py", "content": "# Copyright 2024 The HuggingFace Team.\n# All rights reserved.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n# http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\n\nimport math \nimport random \nimport time\nfrom diffusers.utils import export_to_video\nfrom diffusers.image_processor import VaeImageProcessor\nfrom datetime import datetime, timedelta\nfrom diffusers import CogVideoXPipeline, CogVideoXDDIMScheduler, CogVideoXDPMScheduler\nimport os\nimport torch\nimport argparse\n\n\ndevice = \"cuda\" if torch.cuda.is_available() else \"cpu\"\n\n\ndef get_args():\n parser = argparse.ArgumentParser()\n parser.add_argument(\n \"--pretrained_model_name_or_path\",\n type=str,\n default=None,\n required=True,\n help=\"Path to pretrained model or model identifier from huggingface.co/models.\",\n )\n parser.add_argument(\n \"--lora_weights_path\",\n type=str,\n default=None,\n required=True,\n help=\"Path to lora weights.\",\n )\n parser.add_argument(\n \"--lora_r\",\n type=int,\n default=128,\n help=\"\"\"LoRA weights have a rank parameter, with the default for 2B trans set at 128 and 5B trans set at 256. \n This part is used to calculate the value for lora_scale, which is by default divided by the alpha value, \n used for stable learning and to prevent underflow. In the SAT training framework,\n alpha is set to 1 by default. The higher the rank, the better the expressive capability,\n but it requires more memory and training time. Increasing this number blindly isn't always better.\n The formula for lora_scale is: lora_r / alpha.\n \"\"\",\n )\n parser.add_argument(\n \"--lora_alpha\",\n type=int,\n default=1,\n help=\"\"\"LoRA weights have a rank parameter, with the default for 2B trans set at 128 and 5B trans set at 256. \n This part is used to calculate the value for lora_scale, which is by default divided by the alpha value, \n used for stable learning and to prevent underflow. In the SAT training framework,\n alpha is set to 1 by default. The higher the rank, the better the expressive capability,\n but it requires more memory and training time. Increasing this number blindly isn't always better.\n The formula for lora_scale is: lora_r / alpha.\n \"\"\",\n )\n parser.add_argument(\n \"--prompt\",\n type=str,\n help=\"prompt\",\n )\n parser.add_argument(\n \"--output_dir\",\n type=str,\n default=\"output\",\n help=\"The output directory where the model predictions and checkpoints will be written.\",\n )\n return parser.parse_args()\n\n\nif __name__ == \"__main__\":\n args = get_args()\n pipe = CogVideoXPipeline.from_pretrained(args.pretrained_model_name_or_path, torch_dtype=torch.bfloat16).to(device)\n pipe.load_lora_weights(args.lora_weights_path, weight_name=\"pytorch_lora_weights.safetensors\", adapter_name=\"cogvideox-lora\")\n # pipe.fuse_lora(lora_scale=args.lora_alpha/args.lora_r, ['transformer'])\n lora_scaling=args.lora_alpha/args.lora_r\n pipe.set_adapters([\"cogvideox-lora\"], [lora_scaling])\n\n\n pipe.scheduler = CogVideoXDPMScheduler.from_config(pipe.scheduler.config, timestep_spacing=\"trailing\")\n\n os.makedirs(args.output_dir, exist_ok=True)\n \n latents = pipe(\n prompt=args.prompt,\n num_videos_per_prompt=1,\n num_inference_steps=50,\n num_frames=49,\n use_dynamic_cfg=True,\n output_type=\"pt\",\n guidance_scale=3.0,\n generator=torch.Generator(device=\"cpu\").manual_seed(42),\n ).frames\n batch_size = latents.shape[0]\n batch_video_frames = []\n for batch_idx in range(batch_size):\n pt_image = latents[batch_idx]\n pt_image = torch.stack([pt_image[i] for i in range(pt_image.shape[0])])\n\n image_np = VaeImageProcessor.pt_to_numpy(pt_image)\n image_pil = VaeImageProcessor.numpy_to_pil(image_np)\n batch_video_frames.append(image_pil)\n\n timestamp = datetime.now().strftime(\"%Y%m%d_%H%M%S\")\n video_path = f\"{args.output_dir}/{timestamp}.mp4\"\n os.makedirs(os.path.dirname(video_path), exist_ok=True)\n tensor = batch_video_frames[0]\n fps=math.ceil((len(batch_video_frames[0]) - 1) / 6)\n\n export_to_video(tensor, video_path, fps=fps)" }, { "path": "CogVideo/tools/parallel_inference/parallel_inference_xdit.py", "content": "\"\"\"\nThis is a parallel inference script for CogVideo. The original script\ncan be found from the xDiT project at\n\nhttps://github.com/xdit-project/xDiT/blob/main/examples/cogvideox_example.py\n\nBy using this code, the inference process is parallelized on multiple GPUs,\nand thus speeded up.\n\nUsage:\n1. pip install xfuser\n2. mkdir results\n3. run the following command to generate video\ntorchrun --nproc_per_node=4 parallel_inference_xdit.py \\\n --model --ulysses_degree 1 --ring_degree 2 \\\n --use_cfg_parallel --height 480 --width 720 --num_frames 9 \\\n --prompt 'A small dog.'\n\nYou can also use the run.sh file in the same folder to automate running this\ncode for batch generation of videos, by running:\n\nsh ./run.sh\n\n\"\"\"\n\nimport time\nimport torch\nimport torch.distributed\nfrom diffusers import AutoencoderKLTemporalDecoder\nfrom xfuser import xFuserCogVideoXPipeline, xFuserArgs\nfrom xfuser.config import FlexibleArgumentParser\nfrom xfuser.core.distributed import (\n get_world_group,\n get_data_parallel_rank,\n get_data_parallel_world_size,\n get_runtime_state,\n is_dp_last_group,\n)\nfrom diffusers.utils import export_to_video\n\n\ndef main():\n parser = FlexibleArgumentParser(description=\"xFuser Arguments\")\n args = xFuserArgs.add_cli_args(parser).parse_args()\n engine_args = xFuserArgs.from_cli_args(args)\n\n # Check if ulysses_degree is valid\n num_heads = 30\n if engine_args.ulysses_degree > 0 and num_heads % engine_args.ulysses_degree != 0:\n raise ValueError(\n f\"ulysses_degree ({engine_args.ulysses_degree}) must be a divisor of the number of heads ({num_heads})\"\n )\n\n engine_config, input_config = engine_args.create_config()\n local_rank = get_world_group().local_rank\n\n pipe = xFuserCogVideoXPipeline.from_pretrained(\n pretrained_model_name_or_path=engine_config.model_config.model,\n engine_config=engine_config,\n torch_dtype=torch.bfloat16,\n )\n if args.enable_sequential_cpu_offload:\n pipe.enable_model_cpu_offload(gpu_id=local_rank)\n pipe.vae.enable_tiling()\n else:\n device = torch.device(f\"cuda:{local_rank}\")\n pipe = pipe.to(device)\n\n torch.cuda.reset_peak_memory_stats()\n start_time = time.time()\n\n output = pipe(\n height=input_config.height,\n width=input_config.width,\n num_frames=input_config.num_frames,\n prompt=input_config.prompt,\n num_inference_steps=input_config.num_inference_steps,\n generator=torch.Generator(device=\"cuda\").manual_seed(input_config.seed),\n guidance_scale=6,\n ).frames[0]\n\n end_time = time.time()\n elapsed_time = end_time - start_time\n peak_memory = torch.cuda.max_memory_allocated(device=f\"cuda:{local_rank}\")\n\n parallel_info = (\n f\"dp{engine_args.data_parallel_degree}_cfg{engine_config.parallel_config.cfg_degree}_\"\n f\"ulysses{engine_args.ulysses_degree}_ring{engine_args.ring_degree}_\"\n f\"tp{engine_args.tensor_parallel_degree}_\"\n f\"pp{engine_args.pipefusion_parallel_degree}_patch{engine_args.num_pipeline_patch}\"\n )\n if is_dp_last_group():\n world_size = get_data_parallel_world_size()\n resolution = f\"{input_config.width}x{input_config.height}\"\n output_filename = f\"results/cogvideox_{parallel_info}_{resolution}.mp4\"\n export_to_video(output, output_filename, fps=8)\n print(f\"output saved to {output_filename}\")\n\n if get_world_group().rank == get_world_group().world_size - 1:\n print(f\"epoch time: {elapsed_time:.2f} sec, memory: {peak_memory/1e9} GB\")\n get_runtime_state().destory_distributed_env()\n\n\nif __name__ == \"__main__\":\n main()\n" }, { "path": "CogVideo/tools/parallel_inference/run.sh", "content": "set -x\n\nexport PYTHONPATH=$PWD:$PYTHONPATH\n\n# Select the model type\n# The model is downloaded to a specified location on disk, \n# or you can simply use the model's ID on Hugging Face, \n# which will then be downloaded to the default cache path on Hugging Face.\n\nexport MODEL_TYPE=\"CogVideoX\"\n# Configuration for different model types\n# script, model_id, inference_step\ndeclare -A MODEL_CONFIGS=(\n [\"CogVideoX\"]=\"parallel_inference_xdit.py /cfs/dit/CogVideoX-2b 20\"\n)\n\nif [[ -v MODEL_CONFIGS[$MODEL_TYPE] ]]; then\n IFS=' ' read -r SCRIPT MODEL_ID INFERENCE_STEP <<< \"${MODEL_CONFIGS[$MODEL_TYPE]}\"\n export SCRIPT MODEL_ID INFERENCE_STEP\nelse\n echo \"Invalid MODEL_TYPE: $MODEL_TYPE\"\n exit 1\nfi\n\nmkdir -p ./results\n\n# task args\nif [ \"$MODEL_TYPE\" = \"CogVideoX\" ]; then\n TASK_ARGS=\"--height 480 --width 720 --num_frames 9\"\nfi\n\n# CogVideoX asserts sp_degree == ulysses_degree*ring_degree <= 2. Also, do not set the pipefusion degree.\nif [ \"$MODEL_TYPE\" = \"CogVideoX\" ]; then\nN_GPUS=4\nPARALLEL_ARGS=\"--ulysses_degree 2 --ring_degree 1\"\nCFG_ARGS=\"--use_cfg_parallel\"\nfi\n\n\ntorchrun --nproc_per_node=$N_GPUS ./$SCRIPT \\\n--model $MODEL_ID \\\n$PARALLEL_ARGS \\\n$TASK_ARGS \\\n$PIPEFUSION_ARGS \\\n$OUTPUT_ARGS \\\n--num_inference_steps $INFERENCE_STEP \\\n--warmup_steps 0 \\\n--prompt \"A small dog.\" \\\n$CFG_ARGS \\\n$PARALLLEL_VAE \\\n$COMPILE_FLAG\n" }, { "path": "CogVideo/tools/replicate/cog.yaml", "content": "# Configuration for Cog ⚙️\n# Reference: https://cog.run/yaml\n\nbuild:\n # set to true if your model requires a GPU\n gpu: true\n\n # a list of ubuntu apt packages to install\n system_packages:\n - \"libgl1-mesa-glx\"\n - \"libglib2.0-0\"\n\n # python version in the form '3.11' or '3.11.4'\n python_version: \"3.11\"\n\n # a list of packages in the format ==\n python_packages:\n - diffusers>=0.30.3\n - accelerate>=0.34.2\n - transformers>=4.44.2\n - numpy==1.26.0\n - torch>=2.4.0\n - torchvision>=0.19.0\n - sentencepiece>=0.2.0\n - SwissArmyTransformer>=0.4.12\n - imageio>=2.35.1\n - imageio-ffmpeg>=0.5.1\n - openai>=1.45.0\n - moviepy>=1.0.3\n - pillow==9.5.0\n - pydantic==1.10.7\n run:\n - curl -o /usr/local/bin/pget -L \"https://github.com/replicate/pget/releases/download/v0.8.2/pget_linux_x86_64\" && chmod +x /usr/local/bin/pget\n\n# predict.py defines how predictions are run on your model\npredict: \"predict_t2v.py:Predictor\"\n# predict: \"predict_i2v.py:Predictor\"\n" }, { "path": "CogVideo/tools/replicate/predict_i2v.py", "content": "# Prediction interface for Cog ⚙️\n# https://cog.run/python\n\nimport os\nimport subprocess\nimport time\nimport torch\nfrom diffusers import CogVideoXImageToVideoPipeline\nfrom diffusers.utils import export_to_video, load_image\nfrom cog import BasePredictor, Input, Path\n\n\nMODEL_CACHE = \"model_cache_i2v\"\nMODEL_URL = (\n f\"https://weights.replicate.delivery/default/THUDM/CogVideo/{MODEL_CACHE}.tar\"\n)\nos.environ[\"HF_DATASETS_OFFLINE\"] = \"1\"\nos.environ[\"TRANSFORMERS_OFFLINE\"] = \"1\"\nos.environ[\"HF_HOME\"] = MODEL_CACHE\nos.environ[\"TORCH_HOME\"] = MODEL_CACHE\nos.environ[\"HF_DATASETS_CACHE\"] = MODEL_CACHE\nos.environ[\"TRANSFORMERS_CACHE\"] = MODEL_CACHE\nos.environ[\"HUGGINGFACE_HUB_CACHE\"] = MODEL_CACHE\n\n\ndef download_weights(url, dest):\n start = time.time()\n print(\"downloading url: \", url)\n print(\"downloading to: \", dest)\n subprocess.check_call([\"pget\", \"-x\", url, dest], close_fds=False)\n print(\"downloading took: \", time.time() - start)\n\n\nclass Predictor(BasePredictor):\n def setup(self) -> None:\n \"\"\"Load the model into memory to make running multiple predictions efficient\"\"\"\n\n if not os.path.exists(MODEL_CACHE):\n download_weights(MODEL_URL, MODEL_CACHE)\n\n # model_id: THUDM/CogVideoX-5b-I2V\n self.pipe = CogVideoXImageToVideoPipeline.from_pretrained(\n MODEL_CACHE, torch_dtype=torch.bfloat16\n ).to(\"cuda\")\n\n self.pipe.enable_model_cpu_offload()\n self.pipe.vae.enable_tiling()\n\n def predict(\n self,\n prompt: str = Input(\n description=\"Input prompt\", default=\"Starry sky slowly rotating.\"\n ),\n image: Path = Input(description=\"Input image\"),\n num_inference_steps: int = Input(\n description=\"Number of denoising steps\", ge=1, le=500, default=50\n ),\n guidance_scale: float = Input(\n description=\"Scale for classifier-free guidance\", ge=1, le=20, default=6\n ),\n num_frames: int = Input(\n description=\"Number of frames for the output video\", default=49\n ),\n seed: int = Input(\n description=\"Random seed. Leave blank to randomize the seed\", default=None\n ),\n ) -> Path:\n \"\"\"Run a single prediction on the model\"\"\"\n\n if seed is None:\n seed = int.from_bytes(os.urandom(2), \"big\")\n print(f\"Using seed: {seed}\")\n\n img = load_image(image=str(image))\n\n video = self.pipe(\n prompt=prompt,\n image=img,\n num_videos_per_prompt=1,\n num_inference_steps=num_inference_steps,\n num_frames=num_frames,\n guidance_scale=guidance_scale,\n generator=torch.Generator(device=\"cuda\").manual_seed(seed),\n ).frames[0]\n\n out_path = \"/tmp/out.mp4\"\n\n export_to_video(video, out_path, fps=8)\n return Path(out_path)\n" }, { "path": "CogVideo/tools/replicate/predict_t2v.py", "content": "# Prediction interface for Cog ⚙️\n# https://cog.run/python\n\nimport os\nimport subprocess\nimport time\nimport torch\nfrom diffusers import CogVideoXPipeline\nfrom diffusers.utils import export_to_video\nfrom cog import BasePredictor, Input, Path\n\n\nMODEL_CACHE = \"model_cache\"\nMODEL_URL = (\n f\"https://weights.replicate.delivery/default/THUDM/CogVideo/{MODEL_CACHE}.tar\"\n)\nos.environ[\"HF_DATASETS_OFFLINE\"] = \"1\"\nos.environ[\"TRANSFORMERS_OFFLINE\"] = \"1\"\nos.environ[\"HF_HOME\"] = MODEL_CACHE\nos.environ[\"TORCH_HOME\"] = MODEL_CACHE\nos.environ[\"HF_DATASETS_CACHE\"] = MODEL_CACHE\nos.environ[\"TRANSFORMERS_CACHE\"] = MODEL_CACHE\nos.environ[\"HUGGINGFACE_HUB_CACHE\"] = MODEL_CACHE\n\n\ndef download_weights(url, dest):\n start = time.time()\n print(\"downloading url: \", url)\n print(\"downloading to: \", dest)\n subprocess.check_call([\"pget\", \"-x\", url, dest], close_fds=False)\n print(\"downloading took: \", time.time() - start)\n\n\nclass Predictor(BasePredictor):\n def setup(self) -> None:\n \"\"\"Load the model into memory to make running multiple predictions efficient\"\"\"\n\n if not os.path.exists(MODEL_CACHE):\n download_weights(MODEL_URL, MODEL_CACHE)\n\n # model_id: THUDM/CogVideoX-5b\n self.pipe = CogVideoXPipeline.from_pretrained(\n MODEL_CACHE,\n torch_dtype=torch.bfloat16,\n ).to(\"cuda\")\n\n self.pipe.enable_model_cpu_offload()\n self.pipe.vae.enable_tiling()\n\n def predict(\n self,\n prompt: str = Input(\n description=\"Input prompt\",\n default=\"A panda, dressed in a small, red jacket and a tiny hat, sits on a wooden stool in a serene bamboo forest. The panda's fluffy paws strum a miniature acoustic guitar, producing soft, melodic tunes. Nearby, a few other pandas gather, watching curiously and some clapping in rhythm. Sunlight filters through the tall bamboo, casting a gentle glow on the scene. The panda's face is expressive, showing concentration and joy as it plays. The background includes a small, flowing stream and vibrant green foliage, enhancing the peaceful and magical atmosphere of this unique musical performance.\",\n ),\n num_inference_steps: int = Input(\n description=\"Number of denoising steps\", ge=1, le=500, default=50\n ),\n guidance_scale: float = Input(\n description=\"Scale for classifier-free guidance\", ge=1, le=20, default=6\n ),\n num_frames: int = Input(\n description=\"Number of frames for the output video\", default=49\n ),\n seed: int = Input(\n description=\"Random seed. Leave blank to randomize the seed\", default=None\n ),\n ) -> Path:\n \"\"\"Run a single prediction on the model\"\"\"\n\n if seed is None:\n seed = int.from_bytes(os.urandom(2), \"big\")\n print(f\"Using seed: {seed}\")\n\n video = self.pipe(\n prompt=prompt,\n num_videos_per_prompt=1,\n num_inference_steps=num_inference_steps,\n num_frames=num_frames,\n guidance_scale=guidance_scale,\n generator=torch.Generator(device=\"cuda\").manual_seed(seed),\n ).frames[0]\n\n out_path = \"/tmp/out.mp4\"\n\n export_to_video(video, out_path, fps=8)\n return Path(out_path)\n" }, { "path": "CogVideo/tools/venhancer/README.md", "content": "# Enhance CogVideoX Generated Videos with VEnhancer\n\nThis tutorial will guide you through using the VEnhancer tool to enhance videos generated by CogVideoX, including\nachieving higher frame rates and higher resolutions.\n\n## Model Introduction\n\nVEnhancer implements spatial super-resolution, temporal super-resolution (frame interpolation), and video refinement in\na unified framework. It can flexibly adapt to different upsampling factors (e.g., 1x~8x) for spatial or temporal\nsuper-resolution. Additionally, it provides flexible control to modify the refinement strength, enabling it to handle\ndiverse video artifacts.\n\nVEnhancer follows the design of ControlNet, copying the architecture and weights of the multi-frame encoder and middle\nblock from a pre-trained video diffusion model to build a trainable conditional network. This video ControlNet accepts\nlow-resolution keyframes and noisy full-frame latents as inputs. In addition to the time step t and prompt, our proposed\nvideo-aware conditioning also includes noise augmentation level σ and downscaling factor s as additional network\nconditioning inputs.\n\n## Hardware Requirements\n\n+ Operating System: Linux (requires xformers dependency)\n+ Hardware: NVIDIA GPU with at least 60GB of VRAM per card. Machines such as H100, A100 are recommended.\n\n## Quick Start\n\n1. Clone the repository and install dependencies as per the official instructions:\n\n```shell\ngit clone https://github.com/Vchitect/VEnhancer.git\ncd VEnhancer\n## Torch and other dependencies can use those from CogVideoX. If you need to create a new environment, use the following commands:\npip install torch==2.0.1 torchvision==0.15.2 torchaudio==2.0.2\n\n## Install required dependencies\npip install -r requirements.txt\n```\n\nWhere:\n\n- `input_path` is the path to the input video\n- `prompt` is the description of the video content. The prompt used by this tool should be shorter, not exceeding 77\n words. You may need to simplify the prompt used for generating the CogVideoX video.\n- `target_fps` is the target frame rate for the video. Typically, 16 fps is already smooth, with 24 fps as the default\n value.\n- `up_scale` is recommend to be set to 2,3,4. The target resolution is limited to be around 2k and below.\n- `noise_aug` value depends on the input video quality. Lower quality needs higher noise levels, which corresponds to\n stronger refinement. 250~300 is for very low-quality videos. good videos: <= 200.\n- `steps` if you want fewer steps, please change solver_mode to \"normal\" first, then decline the number of steps. \"\n fast\" solver_mode has fixed steps (15).\n The code will automatically download the required models from Hugging Face during execution.\n\nTypical runtime logs are as follows:\n\n```shell\n/share/home/zyx/.conda/envs/cogvideox/lib/python3.10/site-packages/xformers/ops/fmha/flash.py:211: FutureWarning: `torch.library.impl_abstract` was renamed to `torch.library.register_fake`. Please use that instead; we will remove `torch.library.impl_abstract` in a future version of PyTorch.\n @torch.library.impl_abstract(\"xformers_flash::flash_fwd\")\n/share/home/zyx/.conda/envs/cogvideox/lib/python3.10/site-packages/xformers/ops/fmha/flash.py:344: FutureWarning: `torch.library.impl_abstract` was renamed to `torch.library.register_fake`. Please use that instead; we will remove `torch.library.impl_abstract` in a future version of PyTorch.\n @torch.library.impl_abstract(\"xformers_flash::flash_bwd\")\n2024-08-20 13:25:17,553 - video_to_video - INFO - checkpoint_path: ./ckpts/venhancer_paper.pt\n/share/home/zyx/.conda/envs/cogvideox/lib/python3.10/site-packages/open_clip/factory.py:88: FutureWarning: You are using `torch.load` with `weights_only=False` (the current default value), which uses the default pickle module implicitly. It is possible to construct malicious pickle data which will execute arbitrary code during unpickling (See https://github.com/pytorch/pytorch/blob/main/SECURITY.md#untrusted-models for more details). In a future release, the default value for `weights_only` will be flipped to `True`. This limits the functions that could be executed during unpickling. Arbitrary objects will no longer be allowed to be loaded via this mode unless they are explicitly allowlisted by the user via `torch.serialization.add_safe_globals`. We recommend you start setting `weights_only=True` for any use case where you don't have full control of the loaded file. Please open an issue on GitHub for any issues related to this experimental feature.\n checkpoint = torch.load(checkpoint_path, map_location=map_location)\n2024-08-20 13:25:37,486 - video_to_video - INFO - Build encoder with FrozenOpenCLIPEmbedder\n/share/home/zyx/Code/VEnhancer/video_to_video/video_to_video_model.py:35: FutureWarning: You are using `torch.load` with `weights_only=False` (the current default value), which uses the default pickle module implicitly. It is possible to construct malicious pickle data which will execute arbitrary code during unpickling (See https://github.com/pytorch/pytorch/blob/main/SECURITY.md#untrusted-models for more details). In a future release, the default value for `weights_only` will be flipped to `True`. This limits the functions that could be executed during unpickling. Arbitrary objects will no longer be allowed to be loaded via this mode unless they are explicitly allowlisted by the user via `torch.serialization.add_safe_globals`. We recommend you start setting `weights_only=True` for any use case where you don't have full control of the loaded file. Please open an issue on GitHub for any issues related to this experimental feature.\n load_dict = torch.load(cfg.model_path, map_location='cpu')\n2024-08-20 13:25:55,391 - video_to_video - INFO - Load model path ./ckpts/venhancer_paper.pt, with local status \n2024-08-20 13:25:55,392 - video_to_video - INFO - Build diffusion with GaussianDiffusion\n2024-08-20 13:26:16,092 - video_to_video - INFO - input video path: inputs/000000.mp4\n2024-08-20 13:26:16,093 - video_to_video - INFO - text: Wide-angle aerial shot at dawn,soft morning light casting long shadows,an elderly man walking his dog through a quiet,foggy park,trees and benches in the background,peaceful and serene atmosphere\n2024-08-20 13:26:16,156 - video_to_video - INFO - input frames length: 49\n2024-08-20 13:26:16,156 - video_to_video - INFO - input fps: 8.0\n2024-08-20 13:26:16,156 - video_to_video - INFO - target_fps: 24.0\n2024-08-20 13:26:16,311 - video_to_video - INFO - input resolution: (480, 720)\n2024-08-20 13:26:16,312 - video_to_video - INFO - target resolution: (1320, 1982)\n2024-08-20 13:26:16,312 - video_to_video - INFO - noise augmentation: 250\n2024-08-20 13:26:16,312 - video_to_video - INFO - scale s is set to: 8\n2024-08-20 13:26:16,399 - video_to_video - INFO - video_data shape: torch.Size([145, 3, 1320, 1982])\n/share/home/zyx/Code/VEnhancer/video_to_video/video_to_video_model.py:108: FutureWarning: `torch.cuda.amp.autocast(args...)` is deprecated. Please use `torch.amp.autocast('cuda', args...)` instead.\n with amp.autocast(enabled=True):\n2024-08-20 13:27:19,605 - video_to_video - INFO - step: 0\n2024-08-20 13:30:12,020 - video_to_video - INFO - step: 1\n2024-08-20 13:33:04,956 - video_to_video - INFO - step: 2\n2024-08-20 13:35:58,691 - video_to_video - INFO - step: 3\n2024-08-20 13:38:51,254 - video_to_video - INFO - step: 4\n2024-08-20 13:41:44,150 - video_to_video - INFO - step: 5\n2024-08-20 13:44:37,017 - video_to_video - INFO - step: 6\n2024-08-20 13:47:30,037 - video_to_video - INFO - step: 7\n2024-08-20 13:50:22,838 - video_to_video - INFO - step: 8\n2024-08-20 13:53:15,844 - video_to_video - INFO - step: 9\n2024-08-20 13:56:08,657 - video_to_video - INFO - step: 10\n2024-08-20 13:59:01,648 - video_to_video - INFO - step: 11\n2024-08-20 14:01:54,541 - video_to_video - INFO - step: 12\n2024-08-20 14:04:47,488 - video_to_video - INFO - step: 13\n2024-08-20 14:10:13,637 - video_to_video - INFO - sampling, finished.\n\n```\n\nRunning on a single A100 GPU, enhancing each 6-second CogVideoX generated video with default settings will consume 60GB\nof VRAM and take 40-50 minutes." }, { "path": "CogVideo/tools/venhancer/README_ja.md", "content": "\n# VEnhancer で CogVideoX によって生成されたビデオを強化する\n\nこのチュートリアルでは、VEnhancer ツールを使用して、CogVideoX で生成されたビデオを強化し、より高いフレームレートと高い解像度を実現する方法を説明します。\n\n## モデルの紹介\n\nVEnhancer は、空間超解像、時間超解像（フレーム補間）、およびビデオのリファインメントを統一されたフレームワークで実現します。空間または時間の超解像のために、さまざまなアップサンプリング係数（例：1x〜8x）に柔軟に対応できます。さらに、多様なビデオアーティファクトを処理するために、リファインメント強度を変更する柔軟な制御を提供します。\n\nVEnhancer は ControlNet の設計に従い、事前訓練されたビデオ拡散モデルのマルチフレームエンコーダーとミドルブロックのアーキテクチャとウェイトをコピーして、トレーニング可能な条件ネットワークを構築します。このビデオ ControlNet は、低解像度のキーフレームとノイズを含む完全なフレームを入力として受け取ります。さらに、タイムステップ t とプロンプトに加えて、提案されたビデオ対応条件により、ノイズ増幅レベル σ およびダウンスケーリングファクター s が追加のネットワーク条件として使用されます。\n\n## ハードウェア要件\n\n+ オペレーティングシステム: Linux (xformers 依存関係が必要)\n+ ハードウェア: 単一カードあたり少なくとも 60GB の VRAM を持つ NVIDIA GPU。H100、A100 などのマシンを推奨します。\n\n## クイックスタート\n\n1. 公式の指示に従ってリポジトリをクローンし、依存関係をインストールします。\n\n```shell\ngit clone https://github.com/Vchitect/VEnhancer.git\ncd VEnhancer\n## Torch などの依存関係は CogVideoX の依存関係を使用できます。新しい環境を作成する必要がある場合は、以下のコマンドを使用してください。\npip install torch==2.0.1 torchvision==0.15.2 torchaudio==2.0.2\n\n## 必須の依存関係をインストールします。\npip install -r requirements.txt\n```\n\n2. コードを実行します。\n\n```shell\npython enhance_a_video.py --up_scale 4 --target_fps 24 --noise_aug 250 --solver_mode 'fast' --steps 15 --input_path inputs/000000.mp4 --prompt 'Wide-angle aerial shot at dawn, soft morning light casting long shadows, an elderly man walking his dog through a quiet, foggy park, trees and benches in the background, peaceful and serene atmosphere' --save_dir 'results/'\n```\n\n次の設定を行います：\n\n- `input_path` 是输入视频的路径\n- `prompt` 是视频内容的描述。此工具使用的提示词应更短，不超过77个字。您可能需要简化用于生成CogVideoX视频的提示词。\n- `target_fps` 是视频的目标帧率。通常，16 fps已经很流畅，默认值为24 fps。\n- `up_scale` 推荐设置为2、3或4。目标分辨率限制在2k左右及以下。\n- `noise_aug` 的值取决于输入视频的质量。质量较低的视频需要更高的噪声级别，这对应于更强的优化。250~300适用于非常低质量的视频。对于高质量视频，设置为≤200。\n- `steps` 如果想减少步数，请先将solver_mode改为“normal”，然后减少步数。“fast”模式的步数是固定的（15步）。\n 代码在执行过程中会自动从Hugging Face下载所需的模型。\n\nコードの実行中に、必要なモデルは Hugging Face から自動的にダウンロードされます。\n\n```shell\n/share/home/zyx/.conda/envs/cogvideox/lib/python3.10/site-packages/xformers/ops/fmha/flash.py:211: FutureWarning: `torch.library.impl_abstract` was renamed to `torch.library.register_fake`. Please use that instead; we will remove `torch.library.impl_abstract` in a future version of PyTorch.\n @torch.library.impl_abstract(\"xformers_flash::flash_fwd\")\n/share/home/zyx/.conda/envs/cogvideox/lib/python3.10/site-packages/xformers/ops/fmha/flash.py:344: FutureWarning: `torch.library.impl_abstract` was renamed to `torch.library.register_fake`. Please use that instead; we will remove `torch.library.impl_abstract` in a future version of PyTorch.\n @torch.library.impl_abstract(\"xformers_flash::flash_bwd\")\n2024-08-20 13:25:17,553 - video_to_video - INFO - checkpoint_path: ./ckpts/venhancer_paper.pt\n/share/home/zyx/.conda/envs/cogvideox/lib/python3.10/site-packages/open_clip/factory.py:88: FutureWarning: You are using `torch.load` with `weights_only=False` (the current default value), which uses the default pickle module implicitly. It is possible to construct malicious pickle data which will execute arbitrary code during unpickling (See https://github.com/pytorch/pytorch/blob/main/SECURITY.md#untrusted-models for more details). In a future release, the default value for `weights_only` will be flipped to `True`. This limits the functions that could be executed during unpickling. Arbitrary objects will no longer be allowed to be loaded via this mode unless they are explicitly allowlisted by the user via `torch.serialization.add_safe_globals`. We recommend you start setting `weights_only=True` for any use case where you don't have full control of the loaded file. Please open an issue on GitHub for any issues related to this experimental feature.\n checkpoint = torch.load(checkpoint_path, map_location=map_location)\n2024-08-20 13:25:37,486 - video_to_video - INFO - Build encoder with FrozenOpenCLIPEmbedder\n/share/home/zyx/Code/VEnhancer/video_to_video/video_to_video_model.py:35: FutureWarning: You are using `torch.load` with `weights_only=False` (the current default value), which uses the default pickle module implicitly. It is possible to construct malicious pickle data which will execute arbitrary code during unpickling (See https://github.com/pytorch/pytorch/blob/main/SECURITY.md#untrusted-models for more details). In a future release, the default value for `weights_only` will be flipped to `True`. This limits the functions that could be executed during unpickling. Arbitrary objects will no longer be allowed to be loaded via this mode unless they are explicitly allowlisted by the user via `torch.serialization.add_safe_globals`. We recommend you start setting `weights_only=True` for any use case where you don't have full control of the loaded file. Please open an issue on GitHub for any issues related to this experimental feature.\n load_dict = torch.load(cfg.model_path, map_location='cpu')\n2024-08-20 13:25:55,391 - video_to_video - INFO - Load model path ./ckpts/venhancer_paper.pt, with local status \n2024-08-20 13:25:55,392 - video_to_video - INFO - Build diffusion with GaussianDiffusion\n2024-08-20 13:26:16,092 - video_to_video - INFO - input video path: inputs/000000.mp4\n2024-08-20 13:26:16,093 - video_to_video - INFO - text: Wide-angle aerial shot at dawn,soft morning light casting long shadows,an elderly man walking his dog through a quiet,foggy park,trees and benches in the background,peaceful and serene atmosphere\n2024-08-20 13:26:16,156 - video_to_video - INFO - input frames length: 49\n2024-08-20 13:26:16,156 - video_to_video - INFO - input fps: 8.0\n2024-08-20 13:26:16,156 - video_to_video - INFO - target_fps: 24.0\n2024-08-20 13:26:16,311 - video_to_video - INFO - input resolution: (480, 720)\n2024-08-20 13:26:16,312 - video_to_video - INFO - target resolution: (1320, 1982)\n2024-08-20 13:26:16,312 - video_to_video - INFO - noise augmentation: 250\n2024-08-20 13:26:16,312 - video_to_video - INFO - scale s is set to: 8\n2024-08-20 13:26:16,399 - video_to_video - INFO - video_data shape: torch.Size([145, 3, 1320, 1982])\n/share/home/zyx/Code/VEnhancer/video_to_video/video_to_video_model.py:108: FutureWarning: `torch.cuda.amp.autocast(args...)` is deprecated. Please use `torch.amp.autocast('cuda', args...)` instead.\n with amp.autocast(enabled=True):\n2024-08-20 13:27:19,605 - video_to_video - INFO - step: 0\n2024-08-20 13:30:12,020 - video_to_video - INFO - step: 1\n2024-08-20 13:33:04,956 - video_to_video - INFO - step: 2\n2024-08-20 13:35:58,691 - video_to_video - INFO - step: 3\n2024-08-20 13:38:51,254 - video_to_video - INFO - step: 4\n2024-08-20 13:41:44,150 - video_to_video - INFO - step: 5\n2024-08-20 13:44:37,017 - video_to_video - INFO - step: 6\n2024-08-20 13:47:30,037 - video_to_video - INFO - step: 7\n2024-08-20 13:50:22,838 - video_to_video - INFO - step: 8\n2024-08-20 13:53:15,844 - video_to_video - INFO - step: 9\n2024-08-20 13:56:08,657 - video_to_video - INFO - step: 10\n2024-08-20 13:59:01,648 - video_to_video - INFO - step: 11\n2024-08-20 14:01:54,541 - video_to_video - INFO - step: 12\n2024-08-20 14:04:47,488 - video_to_video - INFO - step: 13\n2024-08-20 14:10:13,637 - video_to_video - INFO - sampling, finished.\n\n```\n\nA100 GPU を単一で使用している場合、CogVideoX によって生成された 6 秒間のビデオを強化するには、デフォルト設定で 60GB の VRAM を消費し、40〜50 分かかります。\n" }, { "path": "CogVideo/tools/venhancer/README_zh.md", "content": "# 使用 VEnhancer 对 CogVdieoX 生成视频进行增强\n\n本教程将要使用 VEnhancer 工具 对 CogVdieoX 生成视频进行增强, 包括更高的帧率和更高的分辨率\n\n## 模型介绍\n\nVEnhancer 在一个统一的框架中实现了空间超分辨率、时间超分辨率（帧插值）和视频优化。它可以灵活地适应不同的上采样因子（例如，1x~\n8x）用于空间或时间超分辨率。此外，它提供了灵活的控制，以修改优化强度，从而处理多样化的视频伪影。\n\nVEnhancer 遵循 ControlNet 的设计，复制了预训练的视频扩散模型的多帧编码器和中间块的架构和权重，构建了一个可训练的条件网络。这个视频\nControlNet 接受低分辨率关键帧和包含噪声的完整帧作为输入。此外，除了时间步 t 和提示词外，我们提出的视频感知条件还将噪声增强的噪声级别\nσ 和降尺度因子 s 作为附加的网络条件输入。\n\n## 硬件需求\n\n+ 操作系统: Linux (需要依赖xformers)\n+ 硬件: NVIDIA GPU 并至少保证单卡显存超过60G，推荐使用 H100，A100等机器。\n\n## 快速上手\n\n1. 按照官方指引克隆仓库并安装依赖\n\n```shell\ngit clone https://github.com/Vchitect/VEnhancer.git\ncd VEnhancer\n## torch等依赖可以使用CogVideoX的依赖，如果你需要创建一个新的环境，可以使用以下命令\npip install torch==2.0.1 torchvision==0.15.2 torchaudio==2.0.2\n\n## 安装必须的依赖\npip install -r requirements.txt\n```\n\n2. 运行代码\n\n```shell\npython enhance_a_video.py \\\n--up_scale 4 --target_fps 24 --noise_aug 250 \\\n--solver_mode 'fast' --steps 15 \\\n--input_path inputs/000000.mp4 \\\n--prompt 'Wide-angle aerial shot at dawn,soft morning light casting long shadows,an elderly man walking his dog through a quiet,foggy park,trees and benches in the background,peaceful and serene atmosphere' \\\n--save_dir 'results/' \n```\n\n其中:\n\n- `input_path` 是输入视频的路径\n- `prompt` 是视频内容的描述。此工具使用的提示词应更短，不超过77个字。您可能需要简化用于生成CogVideoX视频的提示词。\n- `target_fps` 是视频的目标帧率。通常，16 fps已经很流畅，默认值为24 fps。\n- `up_scale` 推荐设置为2、3或4。目标分辨率限制在2k左右及以下。\n- `noise_aug` 的值取决于输入视频的质量。质量较低的视频需要更高的噪声级别，这对应于更强的优化。250~300适用于非常低质量的视频。对于高质量视频，设置为≤200。\n- `steps` 如果想减少步数，请先将solver_mode改为“normal”，然后减少步数。“fast”模式的步数是固定的（15步）。\n 代码在执行过程中会自动从Hugging Face下载所需的模型。\n\n代码运行过程中，会自动从Huggingface拉取需要的模型\n\n运行日志通常如下:\n\n```shell\n/share/home/zyx/.conda/envs/cogvideox/lib/python3.10/site-packages/xformers/ops/fmha/flash.py:211: FutureWarning: `torch.library.impl_abstract` was renamed to `torch.library.register_fake`. Please use that instead; we will remove `torch.library.impl_abstract` in a future version of PyTorch.\n @torch.library.impl_abstract(\"xformers_flash::flash_fwd\")\n/share/home/zyx/.conda/envs/cogvideox/lib/python3.10/site-packages/xformers/ops/fmha/flash.py:344: FutureWarning: `torch.library.impl_abstract` was renamed to `torch.library.register_fake`. Please use that instead; we will remove `torch.library.impl_abstract` in a future version of PyTorch.\n @torch.library.impl_abstract(\"xformers_flash::flash_bwd\")\n2024-08-20 13:25:17,553 - video_to_video - INFO - checkpoint_path: ./ckpts/venhancer_paper.pt\n/share/home/zyx/.conda/envs/cogvideox/lib/python3.10/site-packages/open_clip/factory.py:88: FutureWarning: You are using `torch.load` with `weights_only=False` (the current default value), which uses the default pickle module implicitly. It is possible to construct malicious pickle data which will execute arbitrary code during unpickling (See https://github.com/pytorch/pytorch/blob/main/SECURITY.md#untrusted-models for more details). In a future release, the default value for `weights_only` will be flipped to `True`. This limits the functions that could be executed during unpickling. Arbitrary objects will no longer be allowed to be loaded via this mode unless they are explicitly allowlisted by the user via `torch.serialization.add_safe_globals`. We recommend you start setting `weights_only=True` for any use case where you don't have full control of the loaded file. Please open an issue on GitHub for any issues related to this experimental feature.\n checkpoint = torch.load(checkpoint_path, map_location=map_location)\n2024-08-20 13:25:37,486 - video_to_video - INFO - Build encoder with FrozenOpenCLIPEmbedder\n/share/home/zyx/Code/VEnhancer/video_to_video/video_to_video_model.py:35: FutureWarning: You are using `torch.load` with `weights_only=False` (the current default value), which uses the default pickle module implicitly. It is possible to construct malicious pickle data which will execute arbitrary code during unpickling (See https://github.com/pytorch/pytorch/blob/main/SECURITY.md#untrusted-models for more details). In a future release, the default value for `weights_only` will be flipped to `True`. This limits the functions that could be executed during unpickling. Arbitrary objects will no longer be allowed to be loaded via this mode unless they are explicitly allowlisted by the user via `torch.serialization.add_safe_globals`. We recommend you start setting `weights_only=True` for any use case where you don't have full control of the loaded file. Please open an issue on GitHub for any issues related to this experimental feature.\n load_dict = torch.load(cfg.model_path, map_location='cpu')\n2024-08-20 13:25:55,391 - video_to_video - INFO - Load model path ./ckpts/venhancer_paper.pt, with local status \n2024-08-20 13:25:55,392 - video_to_video - INFO - Build diffusion with GaussianDiffusion\n2024-08-20 13:26:16,092 - video_to_video - INFO - input video path: inputs/000000.mp4\n2024-08-20 13:26:16,093 - video_to_video - INFO - text: Wide-angle aerial shot at dawn,soft morning light casting long shadows,an elderly man walking his dog through a quiet,foggy park,trees and benches in the background,peaceful and serene atmosphere\n2024-08-20 13:26:16,156 - video_to_video - INFO - input frames length: 49\n2024-08-20 13:26:16,156 - video_to_video - INFO - input fps: 8.0\n2024-08-20 13:26:16,156 - video_to_video - INFO - target_fps: 24.0\n2024-08-20 13:26:16,311 - video_to_video - INFO - input resolution: (480, 720)\n2024-08-20 13:26:16,312 - video_to_video - INFO - target resolution: (1320, 1982)\n2024-08-20 13:26:16,312 - video_to_video - INFO - noise augmentation: 250\n2024-08-20 13:26:16,312 - video_to_video - INFO - scale s is set to: 8\n2024-08-20 13:26:16,399 - video_to_video - INFO - video_data shape: torch.Size([145, 3, 1320, 1982])\n/share/home/zyx/Code/VEnhancer/video_to_video/video_to_video_model.py:108: FutureWarning: `torch.cuda.amp.autocast(args...)` is deprecated. Please use `torch.amp.autocast('cuda', args...)` instead.\n with amp.autocast(enabled=True):\n2024-08-20 13:27:19,605 - video_to_video - INFO - step: 0\n2024-08-20 13:30:12,020 - video_to_video - INFO - step: 1\n2024-08-20 13:33:04,956 - video_to_video - INFO - step: 2\n2024-08-20 13:35:58,691 - video_to_video - INFO - step: 3\n2024-08-20 13:38:51,254 - video_to_video - INFO - step: 4\n2024-08-20 13:41:44,150 - video_to_video - INFO - step: 5\n2024-08-20 13:44:37,017 - video_to_video - INFO - step: 6\n2024-08-20 13:47:30,037 - video_to_video - INFO - step: 7\n2024-08-20 13:50:22,838 - video_to_video - INFO - step: 8\n2024-08-20 13:53:15,844 - video_to_video - INFO - step: 9\n2024-08-20 13:56:08,657 - video_to_video - INFO - step: 10\n2024-08-20 13:59:01,648 - video_to_video - INFO - step: 11\n2024-08-20 14:01:54,541 - video_to_video - INFO - step: 12\n2024-08-20 14:04:47,488 - video_to_video - INFO - step: 13\n2024-08-20 14:10:13,637 - video_to_video - INFO - sampling, finished.\n\n```\n\n使用A100单卡运行，对于每个CogVideoX生产的6秒视频，按照默认配置，会消耗60G显存，并用时40-50分钟。" }, { "path": "CogVideo/weights/put weights here.txt", "content": "" }, { "path": "README.md", "content": "## ___***3DTrajMaster: Mastering 3D Trajectory for Multi-Entity Motion in Video Generation***___\n
\n\n\n![Version](https://img.shields.io/badge/version-1.0.0-blue)  \n  \n  \n  \n  \n\n**[Xiao Fu1](https://fuxiao0719.github.io/), \n[Xian Liu1](https://alvinliu0.github.io/), \n[Xintao Wang2 ✉](https://xinntao.github.io/), \n[Sida Peng3](https://pengsida.net/), \n[Menghan Xia2](https://menghanxia.github.io/), \n[Xiaoyu Shi2](https://xiaoyushi97.github.io/), \n[Ziyang Yuan2](https://scholar.google.ru/citations?user=fWxWEzsAAAAJ&hl=en),
\n[Pengfei Wan2](https://scholar.google.com/citations?user=P6MraaYAAAAJ&hl=en)\n[Di Zhang2](https://openreview.net/profile?id=~Di_ZHANG3),\n[Dahua Lin1✉](http://dahua.site/)** \n
\n1The Chinese University of Hong Kong\n2Kuaishou Technology\n3Zhejiang University\n
\n✉: Corresponding Authors\n\n**ICLR 2025**\n\n
\n\n## 🌟 Introduction\n\n🔥 3DTrajMaster controls **one or multiple entity motions in 3D space with entity-specific 3D trajectories** for text-to-video (T2V) generation. It has the following features:\n- **6 Domain of Freedom (DoF)**: control 3D entity location and orientation.\n- **Diverse Entities**: human, animal, robot, car, even abstract fire, breeze, etc.\n- **Diverse Background**: city, forest, desert, gym, sunset beach, glacier, hall, night city, etc.\n- **Complex 3D trajectories**: 3D occlusion, rotating in place, 180°/continuous 90° turnings, etc.\n- **Fine-grained Entity Prompt**: change human hair, clothing, gender, figure size, accessory, etc.\n\nhttps://github.com/user-attachments/assets/efe1870f-4168-4aff-98b8-dbd9e3802928\n\n🔥 **Release News**\n- `[2025/01/23]` 3DTrajMaster is accepted to ICLR 2025.\n- `[2025/01/22]` Release inference and training codes based on CogVideoX-5B.\n- `[2024/12/10]` Release [paper](https://arxiv.org/pdf/2412.07759), [project page](http://fuxiao0719.github.io/projects/3dtrajmaster), [dataset](https://huggingface.co/datasets/KwaiVGI/360Motion-Dataset), and [eval code](https://github.com/KwaiVGI/3DTrajMaster).\n\n## ⚙️ Quick Start\n\n> **(1) Access to Our Internal Video Model**\n\nAs per company policy, we may not release the proprietary trained model at this time. However, if you wish to access our internal model, please submit your request via (1) [a shared document](https://docs.google.com/spreadsheets/d/1HL96IS33fyzrDeXTt3hJ80ZsnfRBzDoKh8wparoBAGI/edit?pli=1&gid=0#gid=0) or (2) directly via email (`lemonaddie0909@gmail.com`, recommended); we will respond to requests with the generated video as quickly as possible.\nPlease ensure your request includes the following:\n\n1. Entity prompts (1–3, with a maximum of 42 tokens, approximately 20 words per entity)\n2. Location prompt\n3. Trajectory template (you can refer to the trajectory template in our released 360°-Motion Dataset, or simply describe new ones via text)\n\n> **(2) Access to Publicly Available Codebase**\n\nWe open-source a model based on CogVideoX-5B. Below is a comparison between CogVideoX and our internal video model as of 2025.01.15.\n\nhttps://github.com/user-attachments/assets/a49e46d3-92d0-42ec-a89f-a9d43919f620\n\n\n#### Inference\n1. **[Environment Set Up]** Our environment setup is identical to [CogVideoX](https://github.com/THUDM/CogVideo). You can refer to their configuration to complete the environment setup.\n\n ```bash\n conda create -n 3dtrajmaster python=3.10\n conda activate 3dtrajmaster\n pip install -r requirements.txt\n ```\n\n2. **[Download Weights and Dataset]** Download the pretrained checkpoints (CogVideo-5B, LoRA, and injector) from [here](https://huggingface.co/KwaiVGI/3DTrajMaster) and place them in the `CogVideo/weights` directory. Then, download the dataset from [here](https://huggingface.co/datasets/KwaiVGI/360Motion-Dataset). Please note that in both training stages, we use only 11 camera poses and exclude the last camera pose as the novel pose setting.\n\n3. **[Inference on Generalizable Prompts]** Change root path to `CogVideo/inference`. Note a higher LoRA scale and more annealed steps can improve accuracy in prompt generation but may result in lower visual quality. You can modify `test_sets.json` to add novel entity&location prompts. For entity input, you can use GPT to enhance the description to an appropriate length, such as \"Generate a detailed description of approximately 20 words\".\n\n ```bash\n python 3dtrajmaster_inference.py \\\n --model_path ../weights/cogvideox-5b \\\n --ckpt_path ../weights/injector \\\n --lora_path ../weights/lora \\\n --lora_scale 0.6 \\\n --annealed_sample_step 20 \\\n --seed 24 \\\n --output_path output_example\n ```\n\n | Argument | Description |\n |-------------------------|-------------|\n | `--lora_scale` | LoRA alpha weight. Options: 0-1, float. Default: 0.6. |\n | `--annealed_sample_step` | annealed sampling steps during inference. Options: 0-50, int. Default: 20. |\n | Generalizable Robustness | prompt entity number: 1>2>3 |\n | Entity Length | 15-24 words, ~24-40 tokens after T5 embeddings |\n\n The following code snapshot showcases the core components of 3DTrajMaster, namely the plug-and-play 3D-motion grounded object injector.\n\n ```python\n # 1. norm & modulate\n norm_hidden_states, norm_empty_encoder_hidden_states, gate_msa, enc_gate_msa = self.norm1(hidden_states, empty_encoder_hidden_states, temb)\n bz, N_visual, dim = norm_hidden_states.shape\n max_entity_num = 3\n _, entity_num, num_frames, _ = pose_embeds.shape\n\n # 2. pair-wise fusion of trajectory and entity\n attn_input = self.attn_null_feature.repeat(bz, max_entity_num, 50, num_frames, 1)\n pose_embeds = self.pose_fuse_layer(pose_embeds)\n attn_input[:,:entity_num,:,:,:] = pose_embeds.unsqueeze(-3) + prompt_entities_embeds.unsqueeze(-2)\n attn_input = torch.cat((\n rearrange(norm_hidden_states, \"b (n t) d -> b n t d\",n=num_frames), \n rearrange(attn_input, \"b n t f d -> b f (n t) d\")),\n dim=2\n ).flatten(1,2)\n\n # 3. gated self-attention\n attn_hidden_states, attn_encoder_hidden_states = self.attn1_injector(\n hidden_states=attn_input,\n encoder_hidden_states=norm_empty_encoder_hidden_states,\n image_rotary_emb=image_rotary_emb,\n )\n attn_hidden_states = attn_hidden_states[:,:N_visual,:]\n\n hidden_states = hidden_states + gate_msa * attn_hidden_states\n ```\n\n#### Training\n\n1. Change root path to `CogVideo/finetune`. First, train lora module to fit the synthetic data domain.\n\n ```bash\n bash finetune_single_rank_lora.sh\n ```\n\n2. Then, train injector module to learn the entity motion controller. Here we set `--block_interval` to 2 to insert the injector every 2 transformer blocks. You can increase this value for a lighter model, but note that it will require a longer training time. For the initial fine-tuning stage, use `--finetune_init`. If resuming from a pre-trained checkpoint, omit `--finetune_init` and specify `--resume_from_checkpoint $TRANSFORMER_PATH` instead. Note that in both training stages, we use only 11 camera poses and exclude the last camera pose as the novel pose setting.\n\n ```bash\n bash finetune_single_rank_injector.sh\n ```\n\n## 📦 360°-Motion Dataset ([Download 🤗](https://huggingface.co/datasets/KwaiVGI/360Motion-Dataset))\n ```\n ├── 360Motion-Dataset Video Number Cam-Obj Distance (m)\n ├── 480_720/384_672\n ├── Desert (desert) 18,000 [3.06, 13.39]\n ├── location_data.json\n ├── HDRI \n ├── loc1 (snowy street) 3,600 [3.43, 13.02]\n ├── loc2 (park) 3,600 [4.16, 12.22]\n ├── loc3 (indoor open space) 3,600 [3.62, 12.79]\n ├── loc11 (gymnastics room) 3,600 [4.06, 12.32]\n ├── loc13 (autumn forest) 3,600 [4.49, 11.92]\n ├── location_data.json\n ├── RefPic\n ├── CharacterInfo.json\n ├── Hemi12_transforms.json\n ```\n\n> **(1) Released Dataset Information (V1.0.0)**\n\n| Argument | Description |Argument | Description |\n|-------------------------|-------------|-------------------------|-------------|\n| **Video Resolution** | (1) 480×720 (2) 384×672 | **Frames/Duration/FPS** | 99/3.3s/30 |\n| **UE Scenes** | 6 (1 desert+5 HDRIs) | **Video Samples** | (1) 36,000 (2) 36,000 |\n| **Camera Intrinsics (fx,fy)** | (1) 1060.606 (2) 989.899 | **Sensor Width/Height (mm)** | (1) 23.76/15.84 (2) 23.76/13.365 |\n| **Hemi12_transforms.json** | 12 surrounding cameras | **CharacterInfo.json** | entity prompts |\n| **RefPic** | 50 animals | **1/2/3 Trajectory Templates** | 36/60/35 (121 in total) |\n| **{D/N}_{locX}** | {Day/Night}_{LocationX} | **{C}_ {XX}_{35mm}** | {Close-Up Shot}_{Cam. Index(1-12)} _{Focal Length}|\n\n**Note that** the resolution of 384×672 refers to our internal video diffusion resolution. In fact, we render the video at a resolution of 378×672 (aspect ratio 9:16), with a 3-pixel black border added to both the top and bottom.\n\n> **(2) Difference with the Dataset to Train on Our Internal Video Diffusion Model**\n\nThe release of the full dataset regarding more entities and UE scenes is still under our internal license check.\n\n| Argument | Released Dataset | Our Internal Dataset|\n|-------------------------|-------------|-------------------------|\n| **Video Resolution** | (1) 480×720 (2) 384×672 | 384×672 |\n| **Entities** | 50 (all animals) | 70 (20 humans+50 animals) |\n| **Video Samples** | (1) 36,000 (2) 36,000 | 54,000 |\n| **Scenes** | 6 | 9 (+city, forest, asian town) |\n| **Trajectory Templates** | 121 | 96 |\n\n> **(3) Load Dataset Sample**\n\n1. Change root path to `dataset`. We provide a script to load our dataset (video & entity & pose sequence) as follows. It will generate the sampled video for visualization in the same folder path.\n\n ```bash\n python load_dataset.py\n ```\n\n2. Visualize the 6DoF pose sequence via Open3D as follows:\n\n ```bash\n python vis_trajecotry.py\n ```\n After running the visualization script, you will get an interactive window like this. Note that we have converted the right-handed coordinate system (Open3D) to the left-handed coordinate system in order to better align with the motion trajectory of the video:\n\n \n\n\n## 🚀 Benchmark Evaluation (Reproduce Paper Results)\n ```\n ├── eval\n ├── GVHMR\n ├── common_metrics_on_video_quality\n ```\n\n> **(1) Evaluation on 3D Trajectory**\n\n1. Change root path to `eval/GVHMR`. Then follow [GVHMR](https://github.com/zju3dv/GVHMR/blob/main/docs/INSTALL.md) installation to prepare the setups and (recommend using a different Conda environment to avoid package conflicts). Our evaluation input is available at [here](https://drive.google.com/file/d/1DLWioJtvv9u4snybu5DrteVWma12JXq3/view?usp=drive_link). Please note that the 3D trajectories have been downsampled from 77 frames to 20 frames to match the RGB latent space of the 3D VAE.\n\n2. Download the [inference videos](https://drive.google.com/file/d/1jMH2-ZC0ZBgtqej5Sp-E5ebBIX7mk3Xz/view?usp=drive_link) generated by our internal video diffusion model and corresponding [evalution GT poses](https://drive.google.com/file/d/1iFcPSlcKb_rDNJ85UPoThdl22BqR2Xgh/view?usp=drive_link) by using this command (you can check the 3D evaluated trajectory via our provided visualization script):\n ```bash\n bash download_eval_pose.sh\n ```\n\n3. Estimation of human poses on evaluation sets:\n ```bash\n python tools/demo/demo_folder.py -f eval_sets -d outputs/eval_sets_gvhmr -s\n ```\n\n4. Evaluation of all human samples (note to convert the left and right hand coordinate systems) :\n ```bash\n python tools/eval_pose.py -f outputs/eval_sets_gvhmr\n ```\n\n> **(2) Evaluation on Visual Quality**\n\n1. Change root path to `eval/common_metrics_on_video_quality`. Then download [fvd](https://drive.google.com/file/d/1U2hd6qvwKLfp7c8yGgcTqdqrP_lKJElB/view?usp=drive_link), [inference videos](https://drive.google.com/file/d/1jMH2-ZC0ZBgtqej5Sp-E5ebBIX7mk3Xz/view?usp=drive_link) and [base T2V inference videos](https://drive.google.com/file/d/1kfdCDA5koYh9g3IkCCHb4XPch2CJAwek/view?usp=drive_link) using the download script:\n ```bash\n bash download_eval_visual.sh\n ```\n\n2. Evaluation of FVD, FID, and CLIP-SIM metrics.\n ```bash\n pip install pytorch-fid clip\n bash eval_visual.sh\n ```\n\n## 📚 Related Work\n\n- [MotionCtrl](https://github.com/TencentARC/MotionCtrl): the first to control 3D camera motion and 2D object motion in video generation\n- [TC4D](https://sherwinbahmani.github.io/tc4d/): compositional text-to-4D scene generation with 3D trajectory conditions\n- [Tora](https://ali-videoai.github.io/tora_video/): control 2D motions in trajectory-oriented diffusion transformer for video generation\n- [SynCamMaster](https://jianhongbai.github.io/SynCamMaster/): multi-camera synchronized video generation from diverse viewpoints\n- [StyleMaster](https://zixuan-ye.github.io/stylemaster): enable artistic video generation and translation with reference style image\n\n####\n\n## 🔗 Citation\nIf you find this work helpful, please consider citing:\n```BibTeXw\n@inproceedings{fu20243dtrajmaster,\n title={3DTrajMaster: Mastering 3D Trajectory for Multi-Entity Motion in Video Generation},\n author={Fu, Xiao and Liu, Xian and Wang, Xintao and Peng, Sida and Xia, Menghan and Shi, Xiaoyu and Yuan, Ziyang and Wan, Pengfei and Zhang, Di and Lin, Dahua},\n booktitle={ICLR},\n year={2025}\n}\n```\n" }, { "path": "dataset/load_dataset.py", "content": "# Copyright 2024 Xiao Fu, CUHK, Kuaishou Tech. All rights reserved.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n# http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n# --------------------------------------------------------------------------\n# If you find this code useful, we kindly ask you to cite our paper in your work.\n# More information about the method can be found at http://fuxiao0719.github.io/projects/3dtrajmaster\n# --------------------------------------------------------------------------\n\nimport os\nimport numpy as np\nimport json\nimport torch\nimport random\nimport cv2\nimport decord\nfrom einops import rearrange\nfrom utils import *\n\n# --------------------------------------------------------------------------\n# 1. Load scenes infomation\n# --------------------------------------------------------------------------\ndataset_root = 'root_path/360Motion-Dataset'\nvideo_res = '480_720'\nvideo_names = []\nscenes = ['Desert', 'HDRI']\nscene_location_pair = {\n 'Desert' : 'desert',\n 'HDRI' : \n {\n 'loc1' : 'snowy street',\n 'loc2' : 'park',\n 'loc3' : 'indoor open space',\n 'loc11' : 'gymnastics room',\n 'loc13' : 'autumn forest',\n }\n}\nfor scene in scenes:\n video_path = os.path.join(dataset_root, video_res, scene)\n locations_path = os.path.join(video_path, \"location_data.json\")\n with open(locations_path, 'r') as f: locations = json.load(f)\n locations_info = {locations[idx]['name']:locations[idx] for idx in range(len(locations))}\n for video_name in os.listdir(video_path):\n if video_name.endswith('Hemi12_1') == True:\n if scene != 'HDRI':\n location = scene_location_pair[scene]\n else:\n location = scene_location_pair['HDRI'][video_name.split('_')[1]]\n video_names.append((video_res, scene, video_name, location, locations_info))\n\n# --------------------------------------------------------------------------\n# 2. Load 12 surrounding cameras\n# --------------------------------------------------------------------------\ncam_num = 12\nmax_objs_num = 3\nlength = len(video_names)\ncaptions_path = os.path.join(dataset_root, \"CharacterInfo.json\")\nwith open(captions_path, 'r') as f: captions = json.load(f)['CharacterInfo']\ncaptions_info = {int(captions[idx]['index']):captions[idx]['eng'] for idx in range(len(captions))}\ncams_path = os.path.join(dataset_root, \"Hemi12_transforms.json\")\nwith open(cams_path, 'r') as f: cams_info = json.load(f)\ncam_poses = []\nfor i, key in enumerate(cams_info.keys()):\n if \"C_\" in key:\n cam_poses.append(parse_matrix(cams_info[key]))\ncam_poses = np.stack(cam_poses)\ncam_poses = np.transpose(cam_poses, (0,2,1))\ncam_poses = cam_poses[:,:,[1,2,0,3]]\ncam_poses[:,:3,3] /= 100.\ncam_poses = cam_poses\nsample_n_frames = 49\n\n# --------------------------------------------------------------------------\n# 3. Load a sample of video & object poses\n# --------------------------------------------------------------------------\n(video_res, scene, video_name, location, locations_info) = video_names[20]\n\nwith open(os.path.join(dataset_root, video_res, scene, video_name, video_name+'.json'), 'r') as f: objs_file = json.load(f)\nobjs_num = len(objs_file['0'])\nvideo_index = random.randint(1, cam_num-1)\n\nlocation_name = video_name.split('_')[1]\nlocation_info = locations_info[location_name]\ncam_pose = cam_poses[video_index-1]\nobj_transl = location_info['coordinates']['CameraTarget']['position']\n\nprompt = ''\nvideo_caption_list = []\nobj_poses_list = []\n\nfor obj_idx in range(objs_num):\n\n obj_name_index = objs_file['0'][obj_idx]['index'] \n video_caption = captions_info[obj_name_index]\n\n if video_caption.startswith(\" \"):\n video_caption = video_caption[1:]\n if video_caption.endswith(\".\"):\n video_caption = video_caption[:-1]\n video_caption = video_caption.lower()\n video_caption_list.append(video_caption)\n \n obj_poses = load_sceneposes(objs_file, obj_idx, obj_transl)\n obj_poses = np.linalg.inv(cam_pose) @ obj_poses\n obj_poses_list.append(obj_poses)\n\nfor obj_idx in range(objs_num):\n video_caption = video_caption_list[obj_idx]\n if obj_idx == objs_num - 1:\n if objs_num == 1:\n prompt += video_caption + ' is moving in the ' + location\n else:\n prompt += video_caption + ' are moving in the ' + location\n else:\n prompt += video_caption + ' and '\n\nobj_poses_all = torch.from_numpy(np.array(obj_poses_list))\n\ntotal_frames = 99\ncurrent_sample_stride = 1.75\ncropped_length = int(sample_n_frames * current_sample_stride)\nstart_frame_ind = random.randint(10, max(10, total_frames - cropped_length - 1))\nend_frame_ind = min(start_frame_ind + cropped_length, total_frames)\nframe_indices = np.linspace(start_frame_ind, end_frame_ind - 1, sample_n_frames, dtype=int)\n\nvideo_frames_path = os.path.join(dataset_root, video_res, scene, video_name, 'videos', video_name+ f'_C_{video_index:02d}_35mm.mp4')\ncap = cv2.VideoCapture(video_frames_path)\nheight = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))\nwidth = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))\n\n# get local rank\nctx = decord.cpu(0)\nreader = decord.VideoReader(video_frames_path, ctx=ctx, height=height, width=width)\nassert len(reader) == total_frames or len(reader) == total_frames+1\nframe_indexes = [frame_idx for frame_idx in range(total_frames)]\ntry:\n video_chunk = reader.get_batch(frame_indexes).asnumpy() \nexcept:\n video_chunk = reader.get_batch(frame_indexes).numpy()\n\npixel_values = np.array([video_chunk[indice] for indice in frame_indices])\npixel_values = rearrange(torch.from_numpy(pixel_values) / 255.0, \"f h w c -> f c h w\")\n\nsave_video = True\nif save_video:\n video_data = (pixel_values.cpu().to(torch.float32).numpy() * 255).astype(np.uint8)\n video_data = rearrange(video_data, \"f c h w -> f h w c\")\n save_images2video(video_data, video_name, 12)\n" }, { "path": "dataset/traj_vis/D_loc1_61_t3n13_003d_Hemi12_1.json", "content": "{\"0\":[{\"index\":113,\"matrix\":\"[-0.00531441 0.999986 0 0] [-0.999986 -0.00531441 0 0] [0 -0 1 0] [18999.8 -22935.3 147.103 1] \"},{\"index\":141,\"matrix\":\"[-0.613263 -0.789878 0 0] [0.789878 -0.613263 -0 0] [0 0 1 0] [18849.5 -22481.6 147.103 1] \"},{\"index\":2,\"matrix\":\"[0.951994 0.306116 0 0] [-0.306116 0.951994 0 0] [0 -0 1 0] [19030.7 -22801.6 147.103 1] \"}],\"1\":[{\"index\":113,\"matrix\":\"[-0.00531441 0.999986 0 0] [-0.999986 -0.00531441 0 0] [0 -0 1 0] [18999.8 -22931.3 147.103 1] \"},{\"index\":141,\"matrix\":\"[-0.341759 -0.939788 0 0] [0.939788 -0.341759 -0 0] [0 0 1 0] [18843.7 -22489 147.103 1] \"},{\"index\":2,\"matrix\":\"[0.777737 0.628589 0 0] [-0.628589 0.777737 0 0] [0 -0 1 0] [19037.2 -22799.5 147.103 1] \"}],\"2\":[{\"index\":113,\"matrix\":\"[-0.00531441 0.999986 0 0] [-0.999986 -0.00531441 0 0] [0 -0 1 0] [18999.7 -22927.4 147.103 1] \"},{\"index\":141,\"matrix\":\"[-0.317189 -0.948362 0 0] [0.948362 -0.317189 -0 0] [0 0 1 0] [18840.5 -22497.9 147.103 1] \"},{\"index\":2,\"matrix\":\"[0.734696 0.678396 0 0] [-0.678396 0.734696 0 0] [0 -0 1 0] [19042.5 -22795.2 147.103 1] \"}],\"3\":[{\"index\":113,\"matrix\":\"[-0.00531441 0.999986 0 0] [-0.999986 -0.00531441 0 0] [0 -0 1 0] [18999.7 -22923.5 147.103 1] \"},{\"index\":141,\"matrix\":\"[-0.304602 -0.95248 0 0] [0.95248 -0.304602 -0 0] [0 0 1 0] [18837.5 -22506.8 147.103 1] \"},{\"index\":2,\"matrix\":\"[0.717886 0.696161 0 0] [-0.696161 0.717886 0 0] [0 -0 1 0] [19047.5 -22790.6 147.103 1] \"}],\"4\":[{\"index\":113,\"matrix\":\"[-0.00531441 0.999986 0 0] [-0.999986 -0.00531441 0 0] [0 -0 1 0] [18999.7 -22919.6 147.103 1] \"},{\"index\":141,\"matrix\":\"[-0.291852 -0.956463 0 0] [0.956463 -0.291852 -0 0] [0 0 1 0] [18834.6 -22515.8 147.103 1] \"},{\"index\":2,\"matrix\":\"[0.703315 0.710879 0 0] [-0.710879 0.703315 0 0] [0 -0 1 0] [19052.4 -22785.8 147.103 1] \"}],\"5\":[{\"index\":113,\"matrix\":\"[-0.00531441 0.999986 0 0] [-0.999986 -0.00531441 0 0] [0 -0 1 0] [18999.7 -22915.7 147.103 1] \"},{\"index\":141,\"matrix\":\"[-0.278641 -0.960395 0 0] [0.960395 -0.278641 -0 0] [0 0 1 0] [18831.9 -22524.8 147.103 1] \"},{\"index\":2,\"matrix\":\"[0.688207 0.725515 0 0] [-0.725515 0.688207 0 0] [0 -0 1 0] [19057.2 -22781 147.103 1] \"}],\"6\":[{\"index\":113,\"matrix\":\"[-0.00531441 0.999986 0 0] [-0.999986 -0.00531441 0 0] [0 -0 1 0] [18999.7 -22911.7 147.103 1] \"},{\"index\":141,\"matrix\":\"[-0.264988 -0.964252 0 0] [0.964252 -0.264988 -0 0] [0 0 1 0] [18829.3 -22533.8 147.103 1] \"},{\"index\":2,\"matrix\":\"[0.672214 0.740357 0 0] [-0.740357 0.672214 0 0] [0 -0 1 0] [19061.9 -22776.1 147.103 1] \"}],\"7\":[{\"index\":113,\"matrix\":\"[-0.00531441 0.999986 0 0] [-0.999986 -0.00531441 0 0] [0 -0 1 0] [18999.6 -22907.8 147.103 1] \"},{\"index\":141,\"matrix\":\"[-0.250908 -0.968011 0 0] [0.968011 -0.250908 -0 0] [0 0 1 0] [18826.8 -22542.9 147.103 1] \"},{\"index\":2,\"matrix\":\"[0.655244 0.755417 0 0] [-0.755417 0.655244 0 0] [0 -0 1 0] [19066.4 -22771 147.103 1] \"}],\"8\":[{\"index\":113,\"matrix\":\"[-0.00531441 0.999986 0 0] [-0.999986 -0.00531441 0 0] [0 -0 1 0] [18999.6 -22903.9 147.103 1] \"},{\"index\":141,\"matrix\":\"[-0.23641 -0.971653 0 0] [0.971653 -0.23641 -0 0] [0 0 1 0] [18824.4 -22552 147.103 1] \"},{\"index\":2,\"matrix\":\"[0.637232 0.770672 0 0] [-0.770672 0.637232 0 0] [0 -0 1 0] [19070.9 -22765.9 147.103 1] \"}],\"9\":[{\"index\":113,\"matrix\":\"[-0.00531441 0.999986 0 0] [-0.999986 -0.00531441 0 0] [0 -0 1 0] [18999.6 -22900 147.103 1] \"},{\"index\":141,\"matrix\":\"[-0.221488 -0.975163 0 0] [0.975163 -0.221488 -0 0] [0 0 1 0] [18822.2 -22561.1 147.103 1] \"},{\"index\":2,\"matrix\":\"[0.618071 0.786122 0 0] [-0.786122 0.618071 0 0] [0 -0 1 0] [19075.2 -22760.7 147.103 1] \"}],\"10\":[{\"index\":113,\"matrix\":\"[-0.00531441 0.999986 0 0] [-0.999986 -0.00531441 0 0] [0 -0 1 0] [18999.6 -22896.1 147.103 1] \"},{\"index\":141,\"matrix\":\"[-0.206134 -0.978524 0 0] [0.978524 -0.206134 -0 0] [0 0 1 0] [18820.1 -22570.3 147.103 1] \"},{\"index\":2,\"matrix\":\"[0.597642 0.801763 0 0] [-0.801763 0.597642 0 0] [0 -0 1 0] [19079.4 -22755.3 147.103 1] \"}],\"11\":[{\"index\":113,\"matrix\":\"[-0.00531441 0.999986 0 0] [-0.999986 -0.00531441 0 0] [0 -0 1 0] [18999.5 -22892.1 147.103 1] \"},{\"index\":141,\"matrix\":\"[-0.190335 -0.981719 0 0] [0.981719 -0.190335 -0 0] [0 0 1 0] [18818.2 -22579.5 147.103 1] \"},{\"index\":2,\"matrix\":\"[0.575785 0.817601 0 0] [-0.817601 0.575785 0 0] [0 -0 1 0] [19083.5 -22749.9 147.103 1] \"}],\"12\":[{\"index\":113,\"matrix\":\"[-0.00531441 0.999986 0 0] [-0.999986 -0.00531441 0 0] [0 -0 1 0] [18999.5 -22888.2 147.103 1] \"},{\"index\":141,\"matrix\":\"[-0.174056 -0.984736 0 0] [0.984736 -0.174056 -0 0] [0 0 1 0] [18816.4 -22588.7 147.103 1] \"},{\"index\":2,\"matrix\":\"[0.552297 0.833647 0 0] [-0.833647 0.552297 0 0] [0 -0 1 0] [19087.4 -22744.3 147.103 1] \"}],\"13\":[{\"index\":113,\"matrix\":\"[-0.00531441 0.999986 0 0] [-0.999986 -0.00531441 0 0] [0 -0 1 0] [18999.5 -22884.3 147.103 1] \"},{\"index\":141,\"matrix\":\"[-0.157284 -0.987553 0 0] [0.987553 -0.157284 -0 0] [0 0 1 0] [18814.7 -22598 147.103 1] \"},{\"index\":2,\"matrix\":\"[0.526901 0.849927 0 0] [-0.849927 0.526901 0 0] [0 -0 1 0] [19091.2 -22738.6 147.103 1] \"}],\"14\":[{\"index\":113,\"matrix\":\"[-0.00531441 0.999986 0 0] [-0.999986 -0.00531441 0 0] [0 -0 1 0] [18999.5 -22880.4 147.103 1] \"},{\"index\":141,\"matrix\":\"[-0.139988 -0.990153 0 0] [0.990153 -0.139988 -0 0] [0 0 1 0] [18813.3 -22607.3 147.103 1] \"},{\"index\":2,\"matrix\":\"[0.499271 0.866446 0 0] [-0.866446 0.499271 0 0] [0 -0 1 0] [19094.8 -22732.8 147.103 1] \"}],\"15\":[{\"index\":113,\"matrix\":\"[-0.00531441 0.999986 0 0] [-0.999986 -0.00531441 0 0] [0 -0 1 0] [18999.5 -22876.5 147.103 1] \"},{\"index\":141,\"matrix\":\"[-0.122116 -0.992516 0 0] [0.992516 -0.122116 -0 0] [0 0 1 0] [18811.9 -22616.6 147.103 1] \"},{\"index\":2,\"matrix\":\"[0.468991 0.883203 0 0] [-0.883203 0.468991 0 0] [0 -0 1 0] [19098.2 -22726.9 147.103 1] \"}],\"16\":[{\"index\":113,\"matrix\":\"[-0.00531441 0.999986 0 0] [-0.999986 -0.00531441 0 0] [0 -0 1 0] [18999.4 -22872.5 147.103 1] \"},{\"index\":141,\"matrix\":\"[-0.103621 -0.994617 0 0] [0.994617 -0.103621 -0 0] [0 0 1 0] [18810.8 -22625.9 147.103 1] \"},{\"index\":2,\"matrix\":\"[0.435515 0.900182 0 0] [-0.900182 0.435515 0 0] [0 -0 1 0] [19101.4 -22720.9 147.103 1] \"}],\"17\":[{\"index\":113,\"matrix\":\"[-0.00531441 0.999986 0 0] [-0.999986 -0.00531441 0 0] [0 -0 1 0] [18999.4 -22868.6 147.103 1] \"},{\"index\":141,\"matrix\":\"[-0.0844501 -0.996428 0 0] [0.996428 -0.0844501 -0 0] [0 0 1 0] [18809.8 -22635.3 147.103 1] \"},{\"index\":2,\"matrix\":\"[0.398234 0.917284 0 0] [-0.917284 0.398234 0 0] [0 -0 1 0] [19104.3 -22714.8 147.103 1] \"}],\"18\":[{\"index\":113,\"matrix\":\"[-0.00531441 0.999986 0 0] [-0.999986 -0.00531441 0 0] [0 -0 1 0] [18999.4 -22864.7 147.103 1] \"},{\"index\":141,\"matrix\":\"[-0.0645412 -0.997915 0 0] [0.997915 -0.0645412 -0 0] [0 0 1 0] [18809 -22644.7 147.103 1] \"},{\"index\":2,\"matrix\":\"[0.356564 0.934271 0 0] [-0.934271 0.356564 0 0] [0 -0 1 0] [19107 -22708.6 147.103 1] \"}],\"19\":[{\"index\":113,\"matrix\":\"[-0.00531441 0.999986 0 0] [-0.999986 -0.00531441 0 0] [0 -0 1 0] [18999.4 -22860.8 147.103 1] \"},{\"index\":141,\"matrix\":\"[-0.0438255 -0.999039 0 0] [0.999039 -0.0438255 -0 0] [0 0 1 0] [18808.4 -22654.1 147.103 1] \"},{\"index\":2,\"matrix\":\"[0.310051 0.95072 0 0] [-0.95072 0.310051 0 0] [0 -0 1 0] [19109.5 -22702.2 147.103 1] \"}],\"20\":[{\"index\":113,\"matrix\":\"[-0.00531441 0.999986 0 0] [-0.999986 -0.00531441 0 0] [0 -0 1 0] [18999.4 -22856.9 147.103 1] \"},{\"index\":141,\"matrix\":\"[-0.0222269 -0.999753 0 0] [0.999753 -0.0222269 -0 0] [0 0 1 0] [18808 -22663.4 147.103 1] \"},{\"index\":2,\"matrix\":\"[0.258735 0.965948 0 0] [-0.965948 0.258735 0 0] [0 -0 1 0] [19111.6 -22695.7 147.103 1] \"}],\"21\":[{\"index\":113,\"matrix\":\"[-0.00531441 0.999986 0 0] [-0.999986 -0.00531441 0 0] [0 -0 1 0] [18999.3 -22852.9 147.103 1] \"},{\"index\":141,\"matrix\":\"[0.00032029 -1 0 0] [1 0.00032029 -0 0] [0 0 1 0] [18807.8 -22672.9 147.103 1] \"},{\"index\":2,\"matrix\":\"[0.200412 0.979712 0 0] [-0.979712 0.200412 0 0] [0 -0 1 0] [19113.3 -22689.2 147.103 1] \"}],\"22\":[{\"index\":113,\"matrix\":\"[-0.00531441 0.999986 0 0] [-0.999986 -0.00531441 0 0] [0 -0 1 0] [18999.3 -22849 147.103 1] \"},{\"index\":141,\"matrix\":\"[0.0239348 -0.999714 0 0] [0.999714 0.0239348 -0 0] [0 0 1 0] [18807.8 -22682.3 147.103 1] \"},{\"index\":2,\"matrix\":\"[0.122083 0.99252 0 0] [-0.99252 0.122083 0 0] [0 -0 1 0] [19114.7 -22682.5 147.103 1] \"}],\"23\":[{\"index\":113,\"matrix\":\"[-0.00531441 0.999986 0 0] [-0.999986 -0.00531441 0 0] [0 -0 1 0] [18999.3 -22845.1 147.103 1] \"},{\"index\":141,\"matrix\":\"[0.0487004 -0.998813 0 0] [0.998813 0.0487004 -0 0] [0 0 1 0] [18808 -22691.7 147.103 1] \"},{\"index\":2,\"matrix\":\"[0.0327411 0.999464 0 0] [-0.999464 0.0327411 0 0] [0 -0 1 0] [19115.5 -22675.7 147.103 1] \"}],\"24\":[{\"index\":113,\"matrix\":\"[-0.00531441 0.999986 0 0] [-0.999986 -0.00531441 0 0] [0 -0 1 0] [18999.3 -22841.2 147.103 1] \"},{\"index\":141,\"matrix\":\"[0.0747293 -0.997204 0 0] [0.997204 0.0747293 -0 0] [0 0 1 0] [18808.5 -22701.1 147.103 1] \"},{\"index\":2,\"matrix\":\"[-0.0557011 0.998447 0 0] [-0.998447 -0.0557011 0 0] [0 -0 1 0] [19115.7 -22668.9 147.103 1] \"}],\"25\":[{\"index\":113,\"matrix\":\"[-0.00531441 0.999986 0 0] [-0.999986 -0.00531441 0 0] [0 -0 1 0] [18999.3 -22837.3 147.103 1] \"},{\"index\":141,\"matrix\":\"[0.102095 -0.994775 0 0] [0.994775 0.102095 -0 0] [0 0 1 0] [18809.2 -22710.4 147.103 1] \"},{\"index\":2,\"matrix\":\"[-0.13492 0.990856 0 0] [-0.990856 -0.13492 0 0] [0 -0 1 0] [19115.4 -22662.2 147.103 1] \"}],\"26\":[{\"index\":113,\"matrix\":\"[-0.00531441 0.999986 0 0] [-0.999986 -0.00531441 0 0] [0 -0 1 0] [18999.2 -22833.3 147.103 1] \"},{\"index\":141,\"matrix\":\"[0.13093 -0.991392 0 0] [0.991392 0.13093 -0 0] [0 0 1 0] [18810.2 -22719.8 147.103 1] \"},{\"index\":2,\"matrix\":\"[-0.202816 0.979217 0 0] [-0.979217 -0.202816 0 0] [0 -0 1 0] [19114.4 -22655.4 147.103 1] \"}],\"27\":[{\"index\":113,\"matrix\":\"[-0.00531441 0.999986 0 0] [-0.999986 -0.00531441 0 0] [0 -0 1 0] [18999.2 -22829.4 147.103 1] \"},{\"index\":141,\"matrix\":\"[0.161334 -0.9869 0 0] [0.9869 0.161334 -0 0] [0 0 1 0] [18811.4 -22729.1 147.103 1] \"},{\"index\":2,\"matrix\":\"[-0.261224 0.965278 0 0] [-0.965278 -0.261224 0 0] [0 -0 1 0] [19113.1 -22648.7 147.103 1] \"}],\"28\":[{\"index\":113,\"matrix\":\"[-0.00531441 0.999986 0 0] [-0.999986 -0.00531441 0 0] [0 -0 1 0] [18999.2 -22825.5 147.103 1] \"},{\"index\":141,\"matrix\":\"[0.193408 -0.981118 0 0] [0.981118 0.193408 -0 0] [0 0 1 0] [18812.9 -22738.4 147.103 1] \"},{\"index\":2,\"matrix\":\"[-0.312923 0.949779 0 0] [-0.949779 -0.312923 0 0] [0 -0 1 0] [19111.3 -22642.2 147.103 1] \"}],\"29\":[{\"index\":113,\"matrix\":\"[-0.00531441 0.999986 0 0] [-0.999986 -0.00531441 0 0] [0 -0 1 0] [18999.2 -22821.6 147.103 1] \"},{\"index\":141,\"matrix\":\"[0.227223 -0.973843 0 0] [0.973843 0.227223 -0 0] [0 0 1 0] [18814.7 -22747.6 147.103 1] \"},{\"index\":2,\"matrix\":\"[-0.360193 0.932878 0 0] [-0.932878 -0.360193 0 0] [0 -0 1 0] [19109.2 -22635.7 147.103 1] \"}],\"30\":[{\"index\":113,\"matrix\":\"[-0.00531441 0.999986 0 0] [-0.999986 -0.00531441 0 0] [0 -0 1 0] [18999.2 -22817.7 147.103 1] \"},{\"index\":141,\"matrix\":\"[0.262798 -0.964851 0 0] [0.964851 0.262798 -0 0] [0 0 1 0] [18816.9 -22756.8 147.103 1] \"},{\"index\":2,\"matrix\":\"[-0.404567 0.914508 0 0] [-0.914508 -0.404567 0 0] [0 -0 1 0] [19106.7 -22629.4 147.103 1] \"}],\"31\":[{\"index\":113,\"matrix\":\"[-0.00531441 0.999986 0 0] [-0.999986 -0.00531441 0 0] [0 -0 1 0] [18999.1 -22813.7 147.103 1] \"},{\"index\":141,\"matrix\":\"[0.300213 -0.953872 0 0] [0.953872 0.300213 -0 0] [0 0 1 0] [18819.3 -22765.9 147.103 1] \"},{\"index\":2,\"matrix\":\"[-0.447013 0.894527 0 0] [-0.894527 -0.447013 0 0] [0 -0 1 0] [19104 -22623.2 147.103 1] \"}],\"32\":[{\"index\":113,\"matrix\":\"[-0.00531441 0.999986 0 0] [-0.999986 -0.00531441 0 0] [0 -0 1 0] [18999.1 -22809.8 147.103 1] \"},{\"index\":141,\"matrix\":\"[0.339367 -0.940654 0 0] [0.940654 0.339367 -0 0] [0 0 1 0] [18822.2 -22774.9 147.103 1] \"},{\"index\":2,\"matrix\":\"[-0.488136 0.872768 0 0] [-0.872768 -0.488136 0 0] [0 -0 1 0] [19100.9 -22617.1 147.103 1] \"}],\"33\":[{\"index\":113,\"matrix\":\"[-0.00531441 0.999986 0 0] [-0.999986 -0.00531441 0 0] [0 -0 1 0] [18999.1 -22805.9 147.103 1] \"},{\"index\":141,\"matrix\":\"[0.380199 -0.924905 0 0] [0.924905 0.380199 -0 0] [0 0 1 0] [18825.3 -22783.7 147.103 1] \"},{\"index\":2,\"matrix\":\"[-0.52829 0.849064 0 0] [-0.849064 -0.52829 0 0] [0 -0 1 0] [19097.6 -22611.1 147.103 1] \"}],\"34\":[{\"index\":113,\"matrix\":\"[-0.00531441 0.999986 0 0] [-0.999986 -0.00531441 0 0] [0 -0 1 0] [18999.1 -22802 147.103 1] \"},{\"index\":141,\"matrix\":\"[0.422527 -0.90635 0 0] [0.90635 0.422527 -0 0] [0 0 1 0] [18828.9 -22792.4 147.103 1] \"},{\"index\":2,\"matrix\":\"[-0.56771 0.823229 0 0] [-0.823229 -0.56771 0 0] [0 -0 1 0] [19094 -22605.3 147.103 1] \"}],\"35\":[{\"index\":113,\"matrix\":\"[-0.00531441 0.999986 0 0] [-0.999986 -0.00531441 0 0] [0 -0 1 0] [18999 -22798.1 147.103 1] \"},{\"index\":141,\"matrix\":\"[0.466092 -0.884736 0 0] [0.884736 0.466092 -0 0] [0 0 1 0] [18832.9 -22800.9 147.103 1] \"},{\"index\":2,\"matrix\":\"[-0.606509 0.795076 0 0] [-0.795076 -0.606509 0 0] [0 -0 1 0] [19090.1 -22599.7 147.103 1] \"}],\"36\":[{\"index\":113,\"matrix\":\"[-0.00531441 0.999986 0 0] [-0.999986 -0.00531441 0 0] [0 -0 1 0] [18999 -22794.1 147.103 1] \"},{\"index\":141,\"matrix\":\"[0.510543 -0.859852 0 0] [0.859852 0.510543 -0 0] [0 0 1 0] [18837.3 -22809.3 147.103 1] \"},{\"index\":2,\"matrix\":\"[-0.64471 0.764428 0 0] [-0.764428 -0.64471 0 0] [0 -0 1 0] [19086 -22594.3 147.103 1] \"}],\"37\":[{\"index\":113,\"matrix\":\"[-0.00531441 0.999986 0 0] [-0.999986 -0.00531441 0 0] [0 -0 1 0] [18999 -22790.2 147.103 1] \"},{\"index\":141,\"matrix\":\"[0.55544 -0.831557 0 0] [0.831557 0.55544 -0 0] [0 0 1 0] [18842.1 -22817.3 147.103 1] \"},{\"index\":2,\"matrix\":\"[-0.682253 0.731116 0 0] [-0.731116 -0.682253 0 0] [0 -0 1 0] [19081.6 -22589.1 147.103 1] \"}],\"38\":[{\"index\":113,\"matrix\":\"[-0.00531441 0.999986 0 0] [-0.999986 -0.00531441 0 0] [0 -0 1 0] [18999 -22786.3 147.103 1] \"},{\"index\":141,\"matrix\":\"[0.60026 -0.799805 0 0] [0.799805 0.60026 -0 0] [0 0 1 0] [18847.3 -22825.2 147.103 1] \"},{\"index\":2,\"matrix\":\"[-0.719002 0.695008 0 0] [-0.695008 -0.719002 0 0] [0 -0 1 0] [19077 -22584.2 147.103 1] \"}],\"39\":[{\"index\":113,\"matrix\":\"[-0.00531441 0.999986 0 0] [-0.999986 -0.00531441 0 0] [0 -0 1 0] [18999 -22782.4 147.103 1] \"},{\"index\":141,\"matrix\":\"[0.644427 -0.764665 0 0] [0.764665 0.644427 -0 0] [0 0 1 0] [18853 -22832.7 147.103 1] \"},{\"index\":2,\"matrix\":\"[-0.754745 0.656018 0 0] [-0.656018 -0.754745 0 0] [0 -0 1 0] [19072.1 -22579.4 147.103 1] \"}],\"40\":[{\"index\":113,\"matrix\":\"[-0.00531441 0.999986 0 0] [-0.999986 -0.00531441 0 0] [0 -0 1 0] [18998.9 -22778.5 147.103 1] \"},{\"index\":141,\"matrix\":\"[0.68732 -0.726355 0 0] [0.726355 0.68732 -0 0] [0 0 1 0] [18859 -22839.9 147.103 1] \"},{\"index\":2,\"matrix\":\"[-0.789206 0.614128 0 0] [-0.614128 -0.789206 0 0] [0 -0 1 0] [19066.9 -22575 147.103 1] \"}],\"41\":[{\"index\":113,\"matrix\":\"[-0.00531441 0.999986 0 0] [-0.999986 -0.00531441 0 0] [0 -0 1 0] [18998.9 -22774.5 147.103 1] \"},{\"index\":141,\"matrix\":\"[0.728345 -0.685211 0 0] [0.685211 0.728345 -0 0] [0 0 1 0] [18865.5 -22846.7 147.103 1] \"},{\"index\":2,\"matrix\":\"[-0.822049 0.569417 0 0] [-0.569417 -0.822049 0 0] [0 -0 1 0] [19061.6 -22570.8 147.103 1] \"}],\"42\":[{\"index\":113,\"matrix\":\"[-0.00531441 0.999986 0 0] [-0.999986 -0.00531441 0 0] [0 -0 1 0] [18998.9 -22770.6 147.103 1] \"},{\"index\":141,\"matrix\":\"[0.766948 -0.641709 0 0] [0.641709 0.766948 -0 0] [0 0 1 0] [18872.3 -22853.2 147.103 1] \"},{\"index\":2,\"matrix\":\"[-0.852907 0.522062 0 0] [-0.522062 -0.852907 0 0] [0 -0 1 0] [19056 -22566.9 147.103 1] \"}],\"43\":[{\"index\":113,\"matrix\":\"[-0.00531441 0.999986 0 0] [-0.999986 -0.00531441 0 0] [0 -0 1 0] [18998.9 -22766.7 147.103 1] \"},{\"index\":141,\"matrix\":\"[0.802643 -0.596459 0 0] [0.596459 0.802643 -0 0] [0 0 1 0] [18879.6 -22859.2 147.103 1] \"},{\"index\":2,\"matrix\":\"[-0.8814 0.472371 0 0] [-0.472371 -0.8814 0 0] [0 -0 1 0] [19050.2 -22563.4 147.103 1] \"}],\"44\":[{\"index\":113,\"matrix\":\"[-0.00531441 0.999986 0 0] [-0.999986 -0.00531441 0 0] [0 -0 1 0] [18998.9 -22762.8 147.103 1] \"},{\"index\":141,\"matrix\":\"[0.835163 -0.550002 0 0] [0.550002 0.835163 -0 0] [0 0 1 0] [18887.1 -22864.8 147.103 1] \"},{\"index\":2,\"matrix\":\"[-0.907165 0.420775 0 0] [-0.420775 -0.907165 0 0] [0 -0 1 0] [19044.2 -22560.2 147.103 1] \"}],\"45\":[{\"index\":113,\"matrix\":\"[-0.00531441 0.999986 0 0] [-0.999986 -0.00531441 0 0] [0 -0 1 0] [18998.8 -22758.9 147.103 1] \"},{\"index\":141,\"matrix\":\"[0.864292 -0.502991 0 0] [0.502991 0.864292 -0 0] [0 0 1 0] [18895 -22870 147.103 1] \"},{\"index\":2,\"matrix\":\"[-0.929899 0.367815 0 0] [-0.367815 -0.929899 0 0] [0 -0 1 0] [19038 -22557.3 147.103 1] \"}],\"46\":[{\"index\":113,\"matrix\":\"[-0.00531441 0.999986 0 0] [-0.999986 -0.00531441 0 0] [0 -0 1 0] [18998.8 -22754.9 147.103 1] \"},{\"index\":141,\"matrix\":\"[0.889922 -0.456113 0 0] [0.456113 0.889922 -0 0] [0 0 1 0] [18903.1 -22874.7 147.103 1] \"},{\"index\":2,\"matrix\":\"[-0.949391 0.314096 0 0] [-0.314096 -0.949391 0 0] [0 -0 1 0] [19031.7 -22554.8 147.103 1] \"}],\"47\":[{\"index\":113,\"matrix\":\"[-0.00531441 0.999986 0 0] [-0.999986 -0.00531441 0 0] [0 -0 1 0] [18998.8 -22751 147.103 1] \"},{\"index\":141,\"matrix\":\"[0.912239 -0.409659 0 0] [0.409659 0.912239 -0 0] [0 0 1 0] [18911.5 -22879 147.103 1] \"},{\"index\":2,\"matrix\":\"[-0.965524 0.260313 0 0] [-0.260313 -0.965524 0 0] [0 -0 1 0] [19025.2 -22552.6 147.103 1] \"}],\"48\":[{\"index\":113,\"matrix\":\"[-0.00531441 0.999986 0 0] [-0.999986 -0.00531441 0 0] [0 -0 1 0] [18998.8 -22747.1 147.103 1] \"},{\"index\":141,\"matrix\":\"[0.931546 -0.363624 0 0] [0.363624 0.931546 -0 0] [0 0 1 0] [18920 -22882.9 147.103 1] \"},{\"index\":2,\"matrix\":\"[-0.978323 0.207085 0 0] [-0.207085 -0.978323 0 0] [0 -0 1 0] [19018.7 -22550.9 147.103 1] \"}],\"49\":[{\"index\":113,\"matrix\":\"[-0.00531441 0.999986 0 0] [-0.999986 -0.00531441 0 0] [0 -0 1 0] [18998.8 -22743.2 147.103 1] \"},{\"index\":141,\"matrix\":\"[0.94818 -0.317734 0 0] [0.317734 0.94818 -0 0] [0 0 1 0] [18928.8 -22886.3 147.103 1] \"},{\"index\":2,\"matrix\":\"[-0.987885 0.155185 0 0] [-0.155185 -0.987885 0 0] [0 -0 1 0] [19012 -22549.5 147.103 1] \"}],\"50\":[{\"index\":113,\"matrix\":\"[-0.00531441 0.999986 0 0] [-0.999986 -0.00531441 0 0] [0 -0 1 0] [18998.7 -22739.3 147.103 1] \"},{\"index\":141,\"matrix\":\"[0.962364 -0.271763 0 0] [0.271763 0.962364 -0 0] [0 0 1 0] [18937.7 -22889.3 147.103 1] \"},{\"index\":2,\"matrix\":\"[-0.994528 0.104469 0 0] [-0.104469 -0.994528 0 0] [0 -0 1 0] [19005.3 -22548.4 147.103 1] \"}],\"51\":[{\"index\":113,\"matrix\":\"[-0.00531441 0.999986 0 0] [-0.999986 -0.00531441 0 0] [0 -0 1 0] [18998.7 -22735.3 147.103 1] \"},{\"index\":141,\"matrix\":\"[0.974232 -0.225548 0 0] [0.225548 0.974232 -0 0] [0 0 1 0] [18946.8 -22891.8 147.103 1] \"},{\"index\":2,\"matrix\":\"[-0.998523 0.0543235 0 0] [-0.0543235 -0.998523 0 0] [0 -0 1 0] [18998.5 -22547.7 147.103 1] \"}],\"52\":[{\"index\":113,\"matrix\":\"[-0.00531441 0.999986 0 0] [-0.999986 -0.00531441 0 0] [0 -0 1 0] [18998.7 -22731.4 147.103 1] \"},{\"index\":141,\"matrix\":\"[0.983864 -0.17892 0 0] [0.17892 0.983864 -0 0] [0 0 1 0] [18955.9 -22893.9 147.103 1] \"},{\"index\":2,\"matrix\":\"[-0.999991 0.00428697 0 0] [-0.00428697 -0.999991 0 0] [0 -0 1 0] [18991.7 -22547.3 147.103 1] \"}],\"53\":[{\"index\":113,\"matrix\":\"[-0.00531441 0.999986 0 0] [-0.999986 -0.00531441 0 0] [0 -0 1 0] [18998.7 -22727.5 147.103 1] \"},{\"index\":141,\"matrix\":\"[0.99128 -0.131775 0 0] [0.131775 0.99128 -0 0] [0 0 1 0] [18965.2 -22895.6 147.103 1] \"},{\"index\":2,\"matrix\":\"[-0.998943 -0.0459677 0 0] [0.0459677 -0.998943 -0 0] [0 0 1 0] [18984.9 -22547.3 147.103 1] \"}],\"54\":[{\"index\":113,\"matrix\":\"[-0.00531441 0.999986 0 0] [-0.999986 -0.00531441 0 0] [0 -0 1 0] [18998.7 -22723.6 147.103 1] \"},{\"index\":141,\"matrix\":\"[0.996469 -0.0839633 0 0] [0.0839633 0.996469 -0 0] [0 0 1 0] [18974.5 -22896.9 147.103 1] \"},{\"index\":2,\"matrix\":\"[-0.995314 -0.0966938 0 0] [0.0966938 -0.995314 -0 0] [0 0 1 0] [18978.1 -22547.6 147.103 1] \"}],\"55\":[{\"index\":113,\"matrix\":\"[-0.00531441 0.999986 0 0] [-0.999986 -0.00531441 0 0] [0 -0 1 0] [18998.6 -22719.7 147.103 1] \"},{\"index\":141,\"matrix\":\"[0.999369 -0.0355101 0 0] [0.0355101 0.999369 -0 0] [0 0 1 0] [18983.9 -22897.7 147.103 1] \"},{\"index\":2,\"matrix\":\"[-0.988976 -0.148073 0 0] [0.148073 -0.988976 -0 0] [0 0 1 0] [18971.3 -22548.3 147.103 1] \"}],\"56\":[{\"index\":113,\"matrix\":\"[-0.00531441 0.999986 0 0] [-0.999986 -0.00531441 0 0] [0 -0 1 0] [18998.6 -22715.7 147.103 1] \"},{\"index\":141,\"matrix\":\"[0.999908 0.0135796 0 0] [-0.0135796 0.999908 0 0] [0 -0 1 0] [18993.3 -22898 147.103 1] \"},{\"index\":2,\"matrix\":\"[-0.979751 -0.20022 0 0] [0.20022 -0.979751 -0 0] [0 0 1 0] [18964.6 -22549.3 147.103 1] \"}],\"57\":[{\"index\":113,\"matrix\":\"[-0.00531441 0.999986 0 0] [-0.999986 -0.00531441 0 0] [0 -0 1 0] [18998.6 -22711.8 147.103 1] \"},{\"index\":141,\"matrix\":\"[0.997998 0.063239 0 0] [-0.063239 0.997998 0 0] [0 -0 1 0] [19002.7 -22897.9 147.103 1] \"},{\"index\":2,\"matrix\":\"[-0.967418 -0.253185 0 0] [0.253185 -0.967418 -0 0] [0 0 1 0] [18958 -22550.6 147.103 1] \"}],\"58\":[{\"index\":113,\"matrix\":\"[-0.00531441 0.999986 0 0] [-0.999986 -0.00531441 0 0] [0 -0 1 0] [18998.6 -22707.9 147.103 1] \"},{\"index\":141,\"matrix\":\"[0.993556 0.113345 0 0] [-0.113345 0.993556 0 0] [0 -0 1 0] [19012.1 -22897.3 147.103 1] \"},{\"index\":2,\"matrix\":\"[-0.95174 -0.306905 0 0] [0.306905 -0.95174 -0 0] [0 0 1 0] [18951.4 -22552.4 147.103 1] \"}],\"59\":[{\"index\":113,\"matrix\":\"[-0.00531441 0.999986 0 0] [-0.999986 -0.00531441 0 0] [0 -0 1 0] [18998.5 -22704 147.103 1] \"},{\"index\":141,\"matrix\":\"[0.986506 0.163722 0 0] [-0.163722 0.986506 0 0] [0 -0 1 0] [19021.4 -22896.2 147.103 1] \"},{\"index\":2,\"matrix\":\"[-0.932457 -0.36128 0 0] [0.36128 -0.932457 -0 0] [0 0 1 0] [18944.9 -22554.5 147.103 1] \"}],\"60\":[{\"index\":113,\"matrix\":\"[-0.00531441 0.999986 0 0] [-0.999986 -0.00531441 0 0] [0 -0 1 0] [18998.5 -22700.1 147.103 1] \"},{\"index\":141,\"matrix\":\"[0.976803 0.214142 0 0] [-0.214142 0.976803 0 0] [0 -0 1 0] [19030.7 -22894.7 147.103 1] \"},{\"index\":2,\"matrix\":\"[-0.909323 -0.416092 0 0] [0.416092 -0.909323 -0 0] [0 0 1 0] [18938.6 -22556.9 147.103 1] \"}],\"61\":[{\"index\":113,\"matrix\":\"[-0.00531441 0.999986 0 0] [-0.999986 -0.00531441 0 0] [0 -0 1 0] [18998.5 -22696.1 147.103 1] \"},{\"index\":141,\"matrix\":\"[0.964431 0.264335 0 0] [-0.264335 0.964431 0 0] [0 -0 1 0] [19039.9 -22892.6 147.103 1] \"},{\"index\":2,\"matrix\":\"[-0.882132 -0.471002 0 0] [0.471002 -0.882132 -0 0] [0 0 1 0] [18932.4 -22559.7 147.103 1] \"}],\"62\":[{\"index\":113,\"matrix\":\"[-0.00531441 0.999986 0 0] [-0.999986 -0.00531441 0 0] [0 -0 1 0] [18998.5 -22692.2 147.103 1] \"},{\"index\":141,\"matrix\":\"[0.949422 0.314004 0 0] [-0.314004 0.949422 0 0] [0 -0 1 0] [19049 -22890.2 147.103 1] \"},{\"index\":2,\"matrix\":\"[-0.850733 -0.525599 0 0] [0.525599 -0.850733 -0 0] [0 0 1 0] [18926.4 -22562.9 147.103 1] \"}],\"63\":[{\"index\":113,\"matrix\":\"[-0.00531441 0.999986 0 0] [-0.999986 -0.00531441 0 0] [0 -0 1 0] [18998.5 -22688.3 147.103 1] \"},{\"index\":141,\"matrix\":\"[0.931857 0.362827 0 0] [-0.362827 0.931857 0 0] [0 -0 1 0] [19057.9 -22887.2 147.103 1] \"},{\"index\":2,\"matrix\":\"[-0.815099 -0.579321 0 0] [0.579321 -0.815099 -0 0] [0 0 1 0] [18920.6 -22566.5 147.103 1] \"}],\"64\":[{\"index\":113,\"matrix\":\"[-0.00531441 0.999986 0 0] [-0.999986 -0.00531441 0 0] [0 -0 1 0] [18998.4 -22684.4 147.103 1] \"},{\"index\":141,\"matrix\":\"[0.911861 0.4105 0 0] [-0.4105 0.911861 0 0] [0 -0 1 0] [19066.7 -22883.8 147.103 1] \"},{\"index\":2,\"matrix\":\"[-0.775337 -0.631548 0 0] [0.631548 -0.775337 -0 0] [0 0 1 0] [18915 -22570.5 147.103 1] \"}],\"65\":[{\"index\":113,\"matrix\":\"[-0.00531441 0.999986 0 0] [-0.999986 -0.00531441 0 0] [0 -0 1 0] [18998.4 -22680.5 147.103 1] \"},{\"index\":141,\"matrix\":\"[0.889615 0.456711 0 0] [-0.456711 0.889615 0 0] [0 -0 1 0] [19075.3 -22879.9 147.103 1] \"},{\"index\":2,\"matrix\":\"[-0.731725 -0.6816 0 0] [0.6816 -0.731725 -0 0] [0 0 1 0] [18909.8 -22574.8 147.103 1] \"}],\"66\":[{\"index\":113,\"matrix\":\"[-0.00531441 0.999986 0 0] [-0.999986 -0.00531441 0 0] [0 -0 1 0] [18998.4 -22676.5 147.103 1] \"},{\"index\":141,\"matrix\":\"[0.865334 0.501196 0 0] [-0.501196 0.865334 0 0] [0 -0 1 0] [19083.6 -22875.6 147.103 1] \"},{\"index\":2,\"matrix\":\"[-0.684729 -0.728797 0 0] [0.728797 -0.684729 -0 0] [0 0 1 0] [18904.8 -22579.4 147.103 1] \"}],\"67\":[{\"index\":113,\"matrix\":\"[-0.00531441 0.999986 0 0] [-0.999986 -0.00531441 0 0] [0 -0 1 0] [18998.4 -22672.6 147.103 1] \"},{\"index\":141,\"matrix\":\"[0.839263 0.543725 0 0] [-0.543725 0.839263 0 0] [0 -0 1 0] [19091.8 -22870.9 147.103 1] \"},{\"index\":2,\"matrix\":\"[-0.634996 -0.772516 0 0] [0.772516 -0.634996 -0 0] [0 0 1 0] [18900.1 -22584.4 147.103 1] \"}],\"68\":[{\"index\":113,\"matrix\":\"[-0.00531441 0.999986 0 0] [-0.999986 -0.00531441 0 0] [0 -0 1 0] [18998.4 -22668.7 147.103 1] \"},{\"index\":141,\"matrix\":\"[0.811667 0.584121 0 0] [-0.584121 0.811667 0 0] [0 -0 1 0] [19099.7 -22865.8 147.103 1] \"},{\"index\":2,\"matrix\":\"[-0.583283 -0.812269 0 0] [0.812269 -0.583283 -0 0] [0 0 1 0] [18895.8 -22589.6 147.103 1] \"}],\"69\":[{\"index\":113,\"matrix\":\"[-0.00531441 0.999986 0 0] [-0.999986 -0.00531441 0 0] [0 -0 1 0] [18998.3 -22664.8 147.103 1] \"},{\"index\":141,\"matrix\":\"[0.782812 0.622258 0 0] [-0.622258 0.782812 0 0] [0 -0 1 0] [19107.3 -22860.3 147.103 1] \"},{\"index\":2,\"matrix\":\"[-0.530414 -0.847739 0 0] [0.847739 -0.530414 -0 0] [0 0 1 0] [18891.8 -22595.1 147.103 1] \"}],\"70\":[{\"index\":113,\"matrix\":\"[-0.00531441 0.999986 0 0] [-0.999986 -0.00531441 0 0] [0 -0 1 0] [18998.3 -22660.9 147.103 1] \"},{\"index\":141,\"matrix\":\"[0.752969 0.658056 0 0] [-0.658056 0.752969 0 0] [0 -0 1 0] [19114.7 -22854.5 147.103 1] \"},{\"index\":2,\"matrix\":\"[-0.477171 -0.878811 0 0] [0.878811 -0.477171 -0 0] [0 0 1 0] [18888.2 -22600.9 147.103 1] \"}],\"71\":[{\"index\":113,\"matrix\":\"[-0.00531441 0.999986 0 0] [-0.999986 -0.00531441 0 0] [0 -0 1 0] [18998.3 -22656.9 147.103 1] \"},{\"index\":141,\"matrix\":\"[0.722357 0.691521 0 0] [-0.691521 0.722357 0 0] [0 -0 1 0] [19121.7 -22848.3 147.103 1] \"},{\"index\":2,\"matrix\":\"[-0.424227 -0.905556 0 0] [0.905556 -0.424227 -0 0] [0 0 1 0] [18885 -22606.9 147.103 1] \"}],\"72\":[{\"index\":113,\"matrix\":\"[-0.00531441 0.999986 0 0] [-0.999986 -0.00531441 0 0] [0 -0 1 0] [18998.3 -22653 147.103 1] \"},{\"index\":141,\"matrix\":\"[0.691244 0.722621 0 0] [-0.722621 0.691244 0 0] [0 -0 1 0] [19128.5 -22841.8 147.103 1] \"},{\"index\":2,\"matrix\":\"[-0.372103 -0.928191 0 0] [0.928191 -0.372103 -0 0] [0 0 1 0] [18882.1 -22613.1 147.103 1] \"}],\"73\":[{\"index\":113,\"matrix\":\"[-0.00531441 0.999986 0 0] [-0.999986 -0.00531441 0 0] [0 -0 1 0] [18998.3 -22649.1 147.103 1] \"},{\"index\":141,\"matrix\":\"[0.659808 0.751434 0 0] [-0.751434 0.659808 0 0] [0 -0 1 0] [19135 -22835 147.103 1] \"},{\"index\":2,\"matrix\":\"[-0.321151 -0.947028 0 0] [0.947028 -0.321151 -0 0] [0 0 1 0] [18879.6 -22619.4 147.103 1] \"}],\"74\":[{\"index\":113,\"matrix\":\"[-0.00531441 0.999986 0 0] [-0.999986 -0.00531441 0 0] [0 -0 1 0] [18998.2 -22645.2 147.103 1] \"},{\"index\":141,\"matrix\":\"[0.62823 0.778027 0 0] [-0.778027 0.62823 0 0] [0 -0 1 0] [19141.3 -22827.9 147.103 1] \"},{\"index\":2,\"matrix\":\"[-0.271566 -0.96242 0 0] [0.96242 -0.271566 -0 0] [0 0 1 0] [18877.4 -22625.8 147.103 1] \"}],\"75\":[{\"index\":113,\"matrix\":\"[-0.00531441 0.999986 0 0] [-0.999986 -0.00531441 0 0] [0 -0 1 0] [18998.2 -22641.3 147.103 1] \"},{\"index\":141,\"matrix\":\"[0.596664 0.802491 0 0] [-0.802491 0.596664 0 0] [0 -0 1 0] [19147.2 -22820.6 147.103 1] \"},{\"index\":2,\"matrix\":\"[-0.22342 -0.974722 0 0] [0.974722 -0.22342 -0 0] [0 0 1 0] [18875.5 -22632.4 147.103 1] \"}],\"76\":[{\"index\":113,\"matrix\":\"[-0.00531441 0.999986 0 0] [-0.999986 -0.00531441 0 0] [0 -0 1 0] [18998.2 -22637.3 147.103 1] \"},{\"index\":141,\"matrix\":\"[0.565227 0.824936 0 0] [-0.824936 0.565227 0 0] [0 -0 1 0] [19152.8 -22813 147.103 1] \"},{\"index\":2,\"matrix\":\"[-0.176687 -0.984267 0 0] [0.984267 -0.176687 -0 0] [0 0 1 0] [18874 -22639 147.103 1] \"}],\"77\":[{\"index\":113,\"matrix\":\"[-0.00531441 0.999986 0 0] [-0.999986 -0.00531441 0 0] [0 -0 1 0] [18998.2 -22633.4 147.103 1] \"},{\"index\":141,\"matrix\":\"[0.534027 0.845468 0 0] [-0.845468 0.534027 0 0] [0 -0 1 0] [19158.1 -22805.3 147.103 1] \"},{\"index\":2,\"matrix\":\"[-0.131276 -0.991346 0 0] [0.991346 -0.131276 -0 0] [0 0 1 0] [18872.8 -22645.7 147.103 1] \"}],\"78\":[{\"index\":113,\"matrix\":\"[-0.00531441 0.999986 0 0] [-0.999986 -0.00531441 0 0] [0 -0 1 0] [18998.2 -22629.5 147.103 1] \"},{\"index\":141,\"matrix\":\"[0.503141 0.864204 0 0] [-0.864204 0.503141 0 0] [0 -0 1 0] [19163.1 -22797.3 147.103 1] \"},{\"index\":2,\"matrix\":\"[-0.0870326 -0.996205 0 0] [0.996205 -0.0870326 -0 0] [0 0 1 0] [18871.9 -22652.4 147.103 1] \"}],\"79\":[{\"index\":113,\"matrix\":\"[-0.00531441 0.999986 0 0] [-0.999986 -0.00531441 0 0] [0 -0 1 0] [18998.1 -22625.6 147.103 1] \"},{\"index\":141,\"matrix\":\"[0.472633 0.88126 0 0] [-0.88126 0.472633 0 0] [0 -0 1 0] [19167.8 -22789.2 147.103 1] \"},{\"index\":2,\"matrix\":\"[-0.0437215 -0.999044 0 0] [0.999044 -0.0437215 -0 0] [0 0 1 0] [18871.3 -22659.2 147.103 1] \"}],\"80\":[{\"index\":113,\"matrix\":\"[-0.00531441 0.999986 0 0] [-0.999986 -0.00531441 0 0] [0 -0 1 0] [18998.1 -22621.7 147.103 1] \"},{\"index\":141,\"matrix\":\"[0.442541 0.896748 0 0] [-0.896748 0.442541 0 0] [0 -0 1 0] [19172.3 -22780.9 147.103 1] \"},{\"index\":2,\"matrix\":\"[-0.0214401 -0.99977 0 0] [0.99977 -0.0214401 -0 0] [0 0 1 0] [18871 -22666 147.103 1] \"}],\"81\":[{\"index\":113,\"matrix\":\"[-0.00531441 0.999986 0 0] [-0.999986 -0.00531441 0 0] [0 -0 1 0] [18998.1 -22617.7 147.103 1] \"},{\"index\":141,\"matrix\":\"[0.412894 0.910779 0 0] [-0.910779 0.412894 0 0] [0 -0 1 0] [19176.5 -22772.5 147.103 1] \"},{\"index\":2,\"matrix\":\"[-0.0214401 -0.99977 0 0] [0.99977 -0.0214401 -0 0] [0 0 1 0] [18870.9 -22672.8 147.103 1] \"}],\"82\":[{\"index\":113,\"matrix\":\"[-0.00531441 0.999986 0 0] [-0.999986 -0.00531441 0 0] [0 -0 1 0] [18998.1 -22613.8 147.103 1] \"},{\"index\":141,\"matrix\":\"[0.383706 0.923455 0 0] [-0.923455 0.383706 0 0] [0 -0 1 0] [19180.3 -22763.9 147.103 1] \"},{\"index\":2,\"matrix\":\"[-0.0214401 -0.99977 0 0] [0.99977 -0.0214401 -0 0] [0 0 1 0] [18870.7 -22679.6 147.103 1] \"}],\"83\":[{\"index\":113,\"matrix\":\"[-0.00531441 0.999986 0 0] [-0.999986 -0.00531441 0 0] [0 -0 1 0] [18998 -22609.9 147.103 1] \"},{\"index\":141,\"matrix\":\"[0.354978 0.934875 0 0] [-0.934875 0.354978 0 0] [0 -0 1 0] [19184 -22755.2 147.103 1] \"},{\"index\":2,\"matrix\":\"[-0.0214401 -0.99977 0 0] [0.99977 -0.0214401 -0 0] [0 0 1 0] [18870.6 -22686.4 147.103 1] \"}],\"84\":[{\"index\":113,\"matrix\":\"[-0.00531441 0.999986 0 0] [-0.999986 -0.00531441 0 0] [0 -0 1 0] [18998 -22606 147.103 1] \"},{\"index\":141,\"matrix\":\"[0.326699 0.945128 0 0] [-0.945128 0.326699 0 0] [0 -0 1 0] [19187.3 -22746.4 147.103 1] \"},{\"index\":2,\"matrix\":\"[-0.0214401 -0.99977 0 0] [0.99977 -0.0214401 -0 0] [0 0 1 0] [18870.4 -22693.2 147.103 1] \"}],\"85\":[{\"index\":113,\"matrix\":\"[-0.00531441 0.999986 0 0] [-0.999986 -0.00531441 0 0] [0 -0 1 0] [18998 -22602.1 147.103 1] \"},{\"index\":141,\"matrix\":\"[0.298848 0.954301 0 0] [-0.954301 0.298848 0 0] [0 -0 1 0] [19190.4 -22737.5 147.103 1] \"},{\"index\":2,\"matrix\":\"[-0.0214401 -0.99977 0 0] [0.99977 -0.0214401 -0 0] [0 0 1 0] [18870.3 -22700 147.103 1] \"}],\"86\":[{\"index\":113,\"matrix\":\"[-0.00531441 0.999986 0 0] [-0.999986 -0.00531441 0 0] [0 -0 1 0] [18998 -22598.1 147.103 1] \"},{\"index\":141,\"matrix\":\"[0.271393 0.962469 0 0] [-0.962469 0.271393 0 0] [0 -0 1 0] [19193.2 -22728.5 147.103 1] \"},{\"index\":2,\"matrix\":\"[-0.0214401 -0.99977 0 0] [0.99977 -0.0214401 -0 0] [0 0 1 0] [18870.1 -22706.8 147.103 1] \"}],\"87\":[{\"index\":113,\"matrix\":\"[-0.00531441 0.999986 0 0] [-0.999986 -0.00531441 0 0] [0 -0 1 0] [18998 -22594.2 147.103 1] \"},{\"index\":141,\"matrix\":\"[0.244292 0.969702 0 0] [-0.969702 0.244292 0 0] [0 -0 1 0] [19195.7 -22719.5 147.103 1] \"},{\"index\":2,\"matrix\":\"[-0.0214401 -0.99977 0 0] [0.99977 -0.0214401 -0 0] [0 0 1 0] [18870 -22713.6 147.103 1] \"}],\"88\":[{\"index\":113,\"matrix\":\"[-0.00531441 0.999986 0 0] [-0.999986 -0.00531441 0 0] [0 -0 1 0] [18997.9 -22590.3 147.103 1] \"},{\"index\":141,\"matrix\":\"[0.21749 0.976063 0 0] [-0.976063 0.21749 0 0] [0 -0 1 0] [19198 -22710.4 147.103 1] \"},{\"index\":2,\"matrix\":\"[-0.0214401 -0.99977 0 0] [0.99977 -0.0214401 -0 0] [0 0 1 0] [18869.8 -22720.4 147.103 1] \"}],\"89\":[{\"index\":113,\"matrix\":\"[-0.00531441 0.999986 0 0] [-0.999986 -0.00531441 0 0] [0 -0 1 0] [18997.9 -22586.4 147.103 1] \"},{\"index\":141,\"matrix\":\"[0.190923 0.981605 0 0] [-0.981605 0.190923 0 0] [0 -0 1 0] [19200.1 -22701.2 147.103 1] \"},{\"index\":2,\"matrix\":\"[-0.0214401 -0.99977 0 0] [0.99977 -0.0214401 -0 0] [0 0 1 0] [18869.7 -22727.2 147.103 1] \"}],\"90\":[{\"index\":113,\"matrix\":\"[-0.00531441 0.999986 0 0] [-0.999986 -0.00531441 0 0] [0 -0 1 0] [18997.9 -22582.5 147.103 1] \"},{\"index\":141,\"matrix\":\"[0.164513 0.986375 0 0] [-0.986375 0.164513 0 0] [0 -0 1 0] [19201.9 -22691.9 147.103 1] \"},{\"index\":2,\"matrix\":\"[-0.0214401 -0.99977 0 0] [0.99977 -0.0214401 -0 0] [0 0 1 0] [18869.6 -22734 147.103 1] \"}],\"91\":[{\"index\":113,\"matrix\":\"[-0.00531441 0.999986 0 0] [-0.999986 -0.00531441 0 0] [0 -0 1 0] [18997.9 -22578.5 147.103 1] \"},{\"index\":141,\"matrix\":\"[0.138169 0.990409 0 0] [-0.990409 0.138169 0 0] [0 -0 1 0] [19203.4 -22682.7 147.103 1] \"},{\"index\":2,\"matrix\":\"[-0.0214401 -0.99977 0 0] [0.99977 -0.0214401 -0 0] [0 0 1 0] [18869.4 -22740.8 147.103 1] \"}],\"92\":[{\"index\":113,\"matrix\":\"[-0.00531441 0.999986 0 0] [-0.999986 -0.00531441 0 0] [0 -0 1 0] [18997.9 -22574.6 147.103 1] \"},{\"index\":141,\"matrix\":\"[0.111776 0.993733 0 0] [-0.993733 0.111776 0 0] [0 -0 1 0] [19204.7 -22673.4 147.103 1] \"},{\"index\":2,\"matrix\":\"[-0.0214401 -0.99977 0 0] [0.99977 -0.0214401 -0 0] [0 0 1 0] [18869.3 -22747.6 147.103 1] \"}],\"93\":[{\"index\":113,\"matrix\":\"[-0.00531441 0.999986 0 0] [-0.999986 -0.00531441 0 0] [0 -0 1 0] [18997.8 -22570.7 147.103 1] \"},{\"index\":141,\"matrix\":\"[0.0851806 0.996366 0 0] [-0.996366 0.0851806 0 0] [0 -0 1 0] [19205.8 -22664 147.103 1] \"},{\"index\":2,\"matrix\":\"[-0.0214401 -0.99977 0 0] [0.99977 -0.0214401 -0 0] [0 0 1 0] [18869.1 -22754.4 147.103 1] \"}],\"94\":[{\"index\":113,\"matrix\":\"[-0.00531441 0.999986 0 0] [-0.999986 -0.00531441 0 0] [0 -0 1 0] [18997.8 -22566.8 147.103 1] \"},{\"index\":141,\"matrix\":\"[0.0582103 0.998304 0 0] [-0.998304 0.0582103 0 0] [0 -0 1 0] [19206.6 -22654.6 147.103 1] \"},{\"index\":2,\"matrix\":\"[-0.0214401 -0.99977 0 0] [0.99977 -0.0214401 -0 0] [0 0 1 0] [18869 -22761.2 147.103 1] \"}],\"95\":[{\"index\":113,\"matrix\":\"[-0.00531441 0.999986 0 0] [-0.999986 -0.00531441 0 0] [0 -0 1 0] [18997.8 -22562.9 147.103 1] \"},{\"index\":141,\"matrix\":\"[0.0306518 0.99953 0 0] [-0.99953 0.0306518 0 0] [0 -0 1 0] [19207.1 -22645.2 147.103 1] \"},{\"index\":2,\"matrix\":\"[-0.0214401 -0.99977 0 0] [0.99977 -0.0214401 -0 0] [0 0 1 0] [18868.8 -22768 147.103 1] \"}],\"96\":[{\"index\":113,\"matrix\":\"[-0.00531441 0.999986 0 0] [-0.999986 -0.00531441 0 0] [0 -0 1 0] [18997.8 -22558.9 147.103 1] \"},{\"index\":141,\"matrix\":\"[0.00217772 0.999998 0 0] [-0.999998 0.00217772 0 0] [0 -0 1 0] [19207.4 -22635.8 147.103 1] \"},{\"index\":2,\"matrix\":\"[-0.0214401 -0.99977 0 0] [0.99977 -0.0214401 -0 0] [0 0 1 0] [18868.8 -22768 147.103 1] \"}],\"97\":[{\"index\":113,\"matrix\":\"[-0.00531441 0.999986 0 0] [-0.999986 -0.00531441 0 0] [0 -0 1 0] [18997.8 -22555 147.103 1] \"},{\"index\":141,\"matrix\":\"[-0.0275859 0.999619 0 0] [-0.999619 -0.0275859 0 0] [0 -0 1 0] [19207.4 -22626.4 147.103 1] \"},{\"index\":2,\"matrix\":\"[-0.0214401 -0.99977 0 0] [0.99977 -0.0214401 -0 0] [0 0 1 0] [18868.8 -22768 147.103 1] \"}],\"98\":[{\"index\":113,\"matrix\":\"[-0.00531441 0.999986 0 0] [-0.999986 -0.00531441 0 0] [0 -0 1 0] [18997.7 -22551.1 147.103 1] \"},{\"index\":141,\"matrix\":\"[-0.0591986 0.998246 0 0] [-0.998246 -0.0591986 0 0] [0 -0 1 0] [19207.2 -22617 147.103 1] \"},{\"index\":2,\"matrix\":\"[-0.0214401 -0.99977 0 0] [0.99977 -0.0214401 -0 0] [0 0 1 0] [18868.8 -22768 147.103 1] \"}]}" }, { "path": "dataset/traj_vis/Hemi12_transforms.json", "content": "{\r\n \"C_01_35mm\": \"[-0.8622445326446021 -0.497817113029644 -0.09334070869305826 0] [0.49999999999999994 -0.8660254037844387 0.0 0] [-0.08083542493543144 -0.04667035434652912 0.9956342260592881 0] [692.820323027551 399.99999999999994 0.0 1]\",\r\n \"C_02_35mm\": \"[-0.49781711302964426 -0.862244532644602 -0.09334070869305827 0] [0.8660254037844386 -0.5000000000000002 0.0 0] [-0.04667035434652916 -0.08083542493543144 0.9956342260592881 0] [400.0000000000001 692.8203230275509 0.0 1]\",\r\n \"C_03_35mm\": \"[-1.6011019497192098e-16 -0.9956342260592881 -0.09334070869305827 0] [1.0 -1.6081226496766366e-16 0.0 0] [-1.5010330778617594e-17 -0.09334070869305827 0.9956342260592881 0] [4.898587196589413e-14 800.0 0.0 1]\",\r\n \"C_04_35mm\": \"[0.49781711302964377 -0.8622445326446022 -0.09334070869305827 0] [0.8660254037844388 0.4999999999999997 0.0 0] [0.04667035434652911 -0.08083542493543147 0.9956342260592881 0] [-399.99999999999983 692.820323027551 0.0 1]\",\r\n \"C_05_35mm\": \"[0.8622445326446021 -0.4978171130296439 -0.09334070869305826 0] [0.49999999999999983 0.8660254037844387 0.0 0] [0.08083542493543144 -0.046670354346529115 0.9956342260592881 0] [-692.820323027551 399.99999999999994 0.0 1]\",\r\n \"C_06_35mm\": \"[0.9956342260592881 -1.2193002680650596e-16 -0.09334070869305827 0] [1.2246467991473532e-16 1.0 0.0 0] [0.09334070869305827 -1.1430940013109933e-17 0.9956342260592881 0] [-800.0 9.797174393178826e-14 0.0 1]\",\r\n \"C_07_35mm\": \"[0.862244532644602 0.49781711302964415 -0.09334070869305827 0] [-0.5000000000000001 0.8660254037844386 0.0 0] [0.08083542493543144 0.04667035434652914 0.9956342260592881 0] [-692.8203230275509 -400.0000000000001 0.0 1]\",\r\n \"C_08_35mm\": \"[0.4978171130296444 0.8622445326446019 -0.09334070869305827 0] [-0.8660254037844385 0.5000000000000003 0.0 0] [0.046670354346529164 0.08083542493543144 0.9956342260592881 0] [-400.00000000000034 -692.8203230275508 0.0 1]\",\r\n \"C_09_35mm\": \"[2.820402217784269e-16 0.9956342260592881 -0.09334070869305827 0] [-1.0 2.83276944882399e-16 0.0 0] [2.6441270791727528e-17 0.09334070869305827 0.9956342260592881 0] [-1.4695761589768238e-13 -800.0 0.0 1]\",\r\n \"C_10_35mm\": \"[-0.49781711302964426 0.862244532644602 -0.09334070869305827 0] [-0.8660254037844386 -0.5000000000000002 0.0 0] [-0.04667035434652916 0.08083542493543144 0.9956342260592881 0] [400.0000000000001 -692.8203230275509 0.0 1]\",\r\n \"C_11_35mm\": \"[-0.8622445326446019 0.4978171130296444 -0.09334070869305827 0] [-0.5000000000000003 -0.8660254037844385 0.0 0] [-0.08083542493543144 0.046670354346529164 0.9956342260592881 0] [692.8203230275507 -400.00000000000034 0.0 1]\",\r\n \"C_12_35mm\": \"[-0.9956342260592881 1.2193002680650596e-16 -0.09334070869305827 0] [-1.2246467991473532e-16 -1.0 0.0 0] [-0.09334070869305827 1.1430940013109933e-17 0.9956342260592881 0] [800.0 -1.9594348786357651e-13 0.0 1]\"\r\n}" }, { "path": "dataset/traj_vis/location_data_desert.json", "content": "[\r\n {\r\n \"name\": \"loc1\",\r\n \"coordinates\": {\r\n \"CameraRig_Rail\": {\r\n \"position\": [\r\n 0,\r\n 0,\r\n 0\r\n ],\r\n \"rotation\": [\r\n 0,\r\n 0,\r\n 0\r\n ],\r\n \"scale\": [\r\n 1,\r\n 1,\r\n 1\r\n ]\r\n },\r\n \"CameraTarget\": {\r\n \"position\": [\r\n 19000.0,\r\n -22700.0,\r\n 0.0\r\n ],\r\n \"rotation\": [\r\n 0.0,\r\n 0.0,\r\n 0.0\r\n ],\r\n \"scale\": [\r\n 1.0,\r\n 1.0,\r\n 1.0\r\n ]\r\n },\r\n \"CameraComponent\": {\r\n \"position\": [\r\n 0.0,\r\n 0.0,\r\n 0.0\r\n ],\r\n \"rotation\": [\r\n 0.0,\r\n 0.0,\r\n 0.0\r\n ],\r\n \"scale\": [\r\n 1.0,\r\n 1.0,\r\n 1.0\r\n ]\r\n }\r\n },\r\n \"Height\": \"H10\"\r\n },\r\n {\r\n \"name\": \"loc2\",\r\n \"coordinates\": {\r\n \"CameraRig_Rail\": {\r\n \"position\": [\r\n 0,\r\n 0,\r\n 0\r\n ],\r\n \"rotation\": [\r\n 0,\r\n 0,\r\n 0\r\n ],\r\n \"scale\": [\r\n 1,\r\n 1,\r\n 1\r\n ]\r\n },\r\n \"CameraTarget\": {\r\n \"position\": [\r\n -200.0,\r\n -11500.0,\r\n 130.0\r\n ],\r\n \"rotation\": [\r\n 0.0,\r\n 0.0,\r\n 0.0\r\n ],\r\n \"scale\": [\r\n 1.0,\r\n 1.0,\r\n 1.0\r\n ]\r\n },\r\n \"CameraComponent\": {\r\n \"position\": [\r\n 0.0,\r\n 0.0,\r\n 0.0\r\n ],\r\n \"rotation\": [\r\n 0.0,\r\n 0.0,\r\n 0.0\r\n ],\r\n \"scale\": [\r\n 1.0,\r\n 1.0,\r\n 1.0\r\n ]\r\n }\r\n },\r\n \"Height\": \"H10\"\r\n },\r\n {\r\n \"name\": \"loc3\",\r\n \"coordinates\": {\r\n \"CameraRig_Rail\": {\r\n \"position\": [\r\n 0,\r\n 0,\r\n 0\r\n ],\r\n \"rotation\": [\r\n 0,\r\n 0,\r\n 0\r\n ],\r\n \"scale\": [\r\n 1,\r\n 1,\r\n 1\r\n ]\r\n },\r\n \"CameraTarget\": {\r\n \"position\": [\r\n -22500.0,\r\n -12900.0,\r\n 20.0\r\n ],\r\n \"rotation\": [\r\n 0.0,\r\n 0.0,\r\n 0.0\r\n ],\r\n \"scale\": [\r\n 1.0,\r\n 1.0,\r\n 1.0\r\n ]\r\n },\r\n \"CameraComponent\": {\r\n \"position\": [\r\n 0.0,\r\n 0.0,\r\n 0.0\r\n ],\r\n \"rotation\": [\r\n 0.0,\r\n 0.0,\r\n 0.0\r\n ],\r\n \"scale\": [\r\n 1.0,\r\n 1.0,\r\n 1.0\r\n ]\r\n }\r\n },\r\n \"Height\": \"H10\"\r\n },\r\n {\r\n \"name\": \"loc4\",\r\n \"coordinates\": {\r\n \"CameraRig_Rail\": {\r\n \"position\": [\r\n 0,\r\n 0,\r\n 0\r\n ],\r\n \"rotation\": [\r\n 0,\r\n 0,\r\n 0\r\n ],\r\n \"scale\": [\r\n 1,\r\n 1,\r\n 1\r\n ]\r\n },\r\n \"CameraTarget\": {\r\n \"position\": [\r\n -22000.0,\r\n 6600.0,\r\n 40.0\r\n ],\r\n \"rotation\": [\r\n 0.0,\r\n 0.0,\r\n 0.0\r\n ],\r\n \"scale\": [\r\n 1.0,\r\n 1.0,\r\n 1.0\r\n ]\r\n },\r\n \"CameraComponent\": {\r\n \"position\": [\r\n 0.0,\r\n 0.0,\r\n 0.0\r\n ],\r\n \"rotation\": [\r\n 0.0,\r\n 0.0,\r\n 0.0\r\n ],\r\n \"scale\": [\r\n 1.0,\r\n 1.0,\r\n 1.0\r\n ]\r\n }\r\n },\r\n \"Height\": \"H10\"\r\n },\r\n {\r\n \"name\": \"loc5\",\r\n \"coordinates\": {\r\n \"CameraRig_Rail\": {\r\n \"position\": [\r\n 0,\r\n 0,\r\n 0\r\n ],\r\n \"rotation\": [\r\n 0,\r\n 0,\r\n 0\r\n ],\r\n \"scale\": [\r\n 1,\r\n 1,\r\n 1\r\n ]\r\n },\r\n \"CameraTarget\": {\r\n \"position\": [\r\n 1300.0,\r\n 28000.0,\r\n 20.0\r\n ],\r\n \"rotation\": [\r\n 0.0,\r\n 0.0,\r\n 0.0\r\n ],\r\n \"scale\": [\r\n 1.0,\r\n 1.0,\r\n 1.0\r\n ]\r\n },\r\n \"CameraComponent\": {\r\n \"position\": [\r\n 0.0,\r\n 0.0,\r\n 0.0\r\n ],\r\n \"rotation\": [\r\n 0.0,\r\n 0.0,\r\n 0.0\r\n ],\r\n \"scale\": [\r\n 1.0,\r\n 1.0,\r\n 1.0\r\n ]\r\n }\r\n },\r\n \"Height\": \"H10\"\r\n }\r\n]" }, { "path": "dataset/utils.py", "content": "# Copyright 2024 Xiao Fu, CUHK, Kuaishou Tech. All rights reserved.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n# http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n# --------------------------------------------------------------------------\n# If you find this code useful, we kindly ask you to cite our paper in your work.\n# More information about the method can be found at http://fuxiao0719.github.io/projects/3dtrajmaster\n# --------------------------------------------------------------------------\nimport os\nimport numpy as np\nfrom io import BytesIO\nimport imageio.v2 as imageio\nimport open3d as o3d\nimport math\nimport trimesh\nimport json\n\n\ndef get_camera_frustum(img_size, K, W2C, frustum_length=0.5, color=[0., 1., 0.]):\n W, H = img_size\n hfov = np.rad2deg(np.arctan(W / 2. / K[0, 0]) * 2.)\n vfov = np.rad2deg(np.arctan(H / 2. / K[1, 1]) * 2.)\n half_w = frustum_length * np.tan(np.deg2rad(hfov / 2.))\n half_h = frustum_length * np.tan(np.deg2rad(vfov / 2.))\n\n # build view frustum for camera (I, 0)\n frustum_points = np.array([[0., 0., 0.], # frustum origin\n [-half_w, -half_h, frustum_length], # top-left image corner\n [half_w, -half_h, frustum_length], # top-right image corner\n [half_w, half_h, frustum_length], # bottom-right image corner\n [-half_w, half_h, frustum_length]]) # bottom-left image corner\n frustum_lines = np.array([[0, i] for i in range(1, 5)] + [[i, (i+1)] for i in range(1, 4)] + [[4, 1]])\n frustum_colors = np.tile(np.array(color).reshape((1, 3)), (frustum_lines.shape[0], 1))\n\n # frustum_colors = np.vstack((np.tile(np.array([[1., 0., 0.]]), (4, 1)),\n # np.tile(np.array([[0., 1., 0.]]), (4, 1))))\n\n # transform view frustum from (I, 0) to (R, t)\n C2W = np.linalg.inv(W2C)\n frustum_points = np.dot(np.hstack((frustum_points, np.ones_like(frustum_points[:, 0:1]))), C2W.T)\n frustum_points = frustum_points[:, :3] / frustum_points[:, 3:4]\n\n return frustum_points, frustum_lines, frustum_colors\n\n\ndef frustums2lineset(frustums):\n N = len(frustums)\n merged_points = np.zeros((N*5, 3)) # 5 vertices per frustum\n merged_lines = np.zeros((N*8, 2)) # 8 lines per frustum\n merged_colors = np.zeros((N*8, 3)) # each line gets a color\n\n for i, (frustum_points, frustum_lines, frustum_colors) in enumerate(frustums):\n merged_points[i*5:(i+1)*5, :] = frustum_points\n merged_lines[i*8:(i+1)*8, :] = frustum_lines + i*5\n merged_colors[i*8:(i+1)*8, :] = frustum_colors\n\n lineset = o3d.geometry.LineSet()\n lineset.points = o3d.utility.Vector3dVector(merged_points)\n lineset.lines = o3d.utility.Vector2iVector(merged_lines)\n lineset.colors = o3d.utility.Vector3dVector(merged_colors)\n\n return lineset\n\ndef visualize_cameras(colored_camera_dicts, sphere_radius, camera_size=0.1, geometry_file=None, geometry_type='mesh'):\n sphere = o3d.geometry.TriangleMesh.create_sphere(radius=sphere_radius, resolution=10)\n sphere = o3d.geometry.LineSet.create_from_triangle_mesh(sphere)\n sphere.paint_uniform_color((1, 0, 0))\n\n coord_frame = o3d.geometry.TriangleMesh.create_coordinate_frame(size=0.5, origin=[0., 0., 0.])\n things_to_draw = [sphere, coord_frame]\n\n idx = 0\n for color, camera_dict in colored_camera_dicts:\n idx += 1\n\n cnt = 0\n frustums = []\n for img_name in sorted(camera_dict.keys()):\n K = np.array(camera_dict[img_name]['K']).reshape((4, 4))\n W2C = np.array(camera_dict[img_name]['W2C']).reshape((4, 4))\n C2W = np.linalg.inv(W2C)\n img_size = camera_dict[img_name]['img_size']\n frustums.append(get_camera_frustum(img_size, K, W2C, frustum_length=camera_size, color=color))\n cnt += 1\n cameras = frustums2lineset(frustums)\n things_to_draw.append(cameras)\n\n if geometry_file is not None:\n if geometry_type == 'mesh':\n geometry = o3d.io.read_triangle_mesh(geometry_file)\n geometry.compute_vertex_normals()\n elif geometry_type == 'pointcloud':\n geometry = o3d.io.read_point_cloud(geometry_file)\n else:\n raise Exception('Unknown geometry_type: ', geometry_type)\n\n things_to_draw.append(geometry)\n\n o3d.visualization.draw_geometries(things_to_draw)\n\ndef parse_matrix(matrix_str):\n rows = matrix_str.strip().split('] [')\n matrix = []\n for row in rows:\n row = row.replace('[', '').replace(']', '')\n matrix.append(list(map(float, row.split())))\n return np.array(matrix)\n\ndef load_sceneposes(objs_file, obj_idx, obj_transl):\n ext_poses = []\n for i, key in enumerate(objs_file.keys()):\n ext_poses.append(parse_matrix(objs_file[key][obj_idx]['matrix']))\n ext_poses = np.stack(ext_poses)\n ext_poses = np.transpose(ext_poses, (0,2,1))\n ext_poses[:,:3,3] -= obj_transl\n ext_poses[:,:3,3] /= 100.\n ext_poses = ext_poses[:, :, [1,2,0,3]]\n return ext_poses\n\ndef save_images2video(images, video_name, fps):\n fps = fps\n format = \"mp4\" \n codec = \"libx264\" \n ffmpeg_params = [\"-crf\", str(12)]\n pixelformat = \"yuv420p\" \n video_stream = BytesIO()\n\n with imageio.get_writer(\n video_stream,\n fps=fps,\n format=format,\n codec=codec,\n ffmpeg_params=ffmpeg_params,\n pixelformat=pixelformat,\n ) as writer:\n for idx in range(len(images)):\n writer.append_data(images[idx])\n\n video_data = video_stream.getvalue()\n\n output_path = os.path.join(video_name + \".mp4\")\n with open(output_path, \"wb\") as f:\n f.write(video_data)\n\ndef normalize(x):\n return x / np.linalg.norm(x)\n\ndef viewmatrix(z, up, pos):\n vec2 = normalize(z)\n vec1_avg = up\n vec0 = normalize(np.cross(vec1_avg, vec2))\n vec1 = normalize(np.cross(vec2, vec0))\n m = np.stack([vec0, vec1, vec2, pos], 1)\n return m\n\ndef matrix_to_euler_angles(matrix):\n sy = math.sqrt(matrix[0][0] * matrix[0][0] + matrix[1][0] * matrix[1][0])\n singular = sy < 1e-6\n\n if not singular:\n x = math.atan2(matrix[2][1], matrix[2][2])\n y = math.atan2(-matrix[2][0], sy)\n z = math.atan2(matrix[1][0], matrix[0][0])\n else:\n x = math.atan2(-matrix[1][2], matrix[1][1])\n y = math.atan2(-matrix[2][0], sy)\n z = 0\n\n return math.degrees(x), math.degrees(y), math.degrees(z)\n\ndef eul2rot(theta) :\n\n R = np.array([[np.cos(theta[1])*np.cos(theta[2]), np.sin(theta[0])*np.sin(theta[1])*np.cos(theta[2]) - np.sin(theta[2])*np.cos(theta[0]), np.sin(theta[1])*np.cos(theta[0])*np.cos(theta[2]) + np.sin(theta[0])*np.sin(theta[2])],\n [np.sin(theta[2])*np.cos(theta[1]), np.sin(theta[0])*np.sin(theta[1])*np.sin(theta[2]) + np.cos(theta[0])*np.cos(theta[2]), np.sin(theta[1])*np.sin(theta[2])*np.cos(theta[0]) - np.sin(theta[0])*np.cos(theta[2])],\n [-np.sin(theta[1]), np.sin(theta[0])*np.cos(theta[1]), np.cos(theta[0])*np.cos(theta[1])]])\n\n return R.T\n\ndef extract_location_rotation(data):\n results = {}\n for key, value in data.items():\n matrix = parse_matrix(value)\n location = np.array([matrix[3][0], matrix[3][1], matrix[3][2]])\n rotation = eul2rot(matrix_to_euler_angles(matrix))\n transofmed_matrix = np.identity(4)\n transofmed_matrix[:3,3] = location\n transofmed_matrix[:3,:3] = rotation\n results[key] = transofmed_matrix\n return results\n\ndef get_cam_points_vis(W, H, intrinsics, ext_pose, color,frustum_length):\n cam = get_camera_frustum((W, H), intrinsics, np.linalg.inv(ext_pose), frustum_length=frustum_length, color=[0., 0., 1.])\n cam_points = cam[0]\n for item in cam[1]:\n cam_points = np.concatenate((cam_points, np.linspace(cam[0][item[0]], cam[0][item[1]], num=1000, endpoint=True, retstep=False, dtype=None)))\n cam_points[:,0]*=-1\n cam_points = trimesh.points.PointCloud(vertices = cam_points, colors=[0, 255, 0, 255])\n cam_points_vis = o3d.geometry.PointCloud()\n cam_points_vis.points = o3d.utility.Vector3dVector(cam_points)\n cam_points_vis.paint_uniform_color(color)\n return cam_points_vis\n\ndef batch_axis_angle_to_rotation_matrix(r_batch):\n batch_size = r_batch.shape[0]\n rotation_matrices = []\n \n for i in range(batch_size):\n r = r_batch[i]\n theta = np.linalg.norm(r)\n if theta == 0:\n rotation_matrices.append(np.eye(3))\n else:\n k = r / theta \n kx, ky, kz = k\n \n K = np.array([\n [0, -kz, ky],\n [kz, 0, -kx],\n [-ky, kx, 0]\n ])\n \n R = np.eye(3) + np.sin(theta) * K + (1 - np.cos(theta)) * np.dot(K, K)\n rotation_matrices.append(R)\n \n return np.array(rotation_matrices)" }, { "path": "dataset/vis_trajectory.py", "content": "import trimesh\nimport numpy as np\nimport imageio\nimport copy\nimport cv2\nimport os\nfrom glob import glob\nimport open3d\nfrom multiprocessing import Pool\nimport json\nfrom utils import *\n\nif __name__ == '__main__' :\n\n H = 480\n W = 720\n intrinsics = np.array([[1060.606,0.],\n [0., 1060.606]])\n \n cam_path = \"traj_vis/Hemi12_transforms.json\"\n location_path = \"traj_vis/location_data_desert.json\"\n video_name = \"D_loc1_61_t3n13_003d_Hemi12_1.json\"\n\n with open(location_path, 'r') as f: locations = json.load(f)\n locations_info = {locations[idx]['name']:locations[idx] for idx in range(len(locations))}\n location_name = video_name.split('_')[1]\n location_info = locations_info[location_name]\n translation = location_info['coordinates']['CameraTarget']['position']\n vis_all = []\n \n # vis cam\n with open(cam_path, 'r') as file:\n data = json.load(file)\n cam_poses = []\n for i, key in enumerate(data.keys()):\n if \"C_\" in key:\n cam_poses.append(parse_matrix(data[key]))\n\n cam_poses = np.stack(cam_poses)\n cam_poses = np.transpose(cam_poses, (0,2,1))\n cam_poses[:,:3,3] /= 100.\n cam_poses = cam_poses[:,:,[1,2,0,3]]\n relative_pose = np.linalg.inv(cam_poses[0])\n cam_poses = relative_pose @ cam_poses\n # convert right-hand coord to left-hand coord\n cam_poses[:,:3,3][:,1] *= -1.\n cam_poses[:,:,:2] *= -1.\n\n cam_num = len(cam_poses)\n for idx in range(cam_num):\n cam_pose = cam_poses[idx]\n cam_points_vis = get_cam_points_vis(W, H, intrinsics, cam_pose, [0.4, 0.4, 0.4], frustum_length=1.)\n vis_all.append(cam_points_vis)\n\n # vis gt obj poses\n start_frame_ind = 10\n sample_n_frames = 77\n frame_indices = np.linspace(start_frame_ind, start_frame_ind + sample_n_frames - 1, sample_n_frames, dtype=int)\n \n with open('traj_vis/'+video_name, 'r') as file:\n data = json.load(file)\n obj_poses = []\n for i, key in enumerate(data.keys()):\n obj_poses.append(parse_matrix(data[key][0]['matrix']))\n obj_poses = np.stack(obj_poses)\n obj_poses = np.transpose(obj_poses, (0,2,1))\n obj_poses[:,:3,3] -= translation\n obj_poses[:,:3,3] /= 100.\n obj_poses = obj_poses[:, :, [1,2,0,3]]\n obj_poses = relative_pose @ obj_poses\n obj_poses = obj_poses[frame_indices]\n # convert right-hand coord to left-hand coord\n obj_poses[:,:3,3][:,1] *= -1.\n obj_poses[:,:,:2] *= -1.\n\n obj_num = len(obj_poses)\n for idx in range(obj_num):\n obj_pose = obj_poses[idx]\n if idx % 5 == 0:\n cam_points_vis = get_cam_points_vis(W, H, intrinsics, obj_pose, [0.8, 0., 0.], frustum_length=0.5)\n vis_all.append(cam_points_vis)\n\n if len(data[key])>=2:\n with open('traj_vis/'+video_name, 'r') as file:\n data = json.load(file)\n obj_poses = []\n for i, key in enumerate(data.keys()):\n obj_poses.append(parse_matrix(data[key][1]['matrix']))\n obj_poses = np.stack(obj_poses)\n obj_poses = np.transpose(obj_poses, (0,2,1))\n obj_poses[:,:3,3] -= translation\n obj_poses[:,:3,3] /= 100.\n obj_poses = obj_poses[:, :, [1,2,0,3]]\n obj_poses = relative_pose @ obj_poses\n obj_poses = obj_poses[frame_indices]\n # convert right-hand coord to left-hand coord\n obj_poses[:,:3,3][:,1] *= -1.\n obj_poses[:,:,:2] *= -1.\n obj_num = len(obj_poses)\n for idx in range(obj_num):\n obj_pose = obj_poses[idx]\n if (idx % 5 == 0) :\n cam_points_vis = get_cam_points_vis(W, H, intrinsics, obj_pose, [0., 0.8,0.], frustum_length=0.5)\n vis_all.append(cam_points_vis)\n\n if len(data[key])>=3:\n with open('traj_vis/'+video_name, 'r') as file:\n data = json.load(file)\n obj_poses = []\n for i, key in enumerate(data.keys()):\n obj_poses.append(parse_matrix(data[key][2]['matrix']))\n obj_poses = np.stack(obj_poses)\n obj_poses = np.transpose(obj_poses, (0,2,1))\n obj_poses[:,:3,3] -= translation\n obj_poses[:,:3,3] /= 100.\n obj_poses = obj_poses[:, :, [1,2,0,3]]\n obj_poses = relative_pose @ obj_poses\n obj_poses = obj_poses[frame_indices]\n # convert right-hand coord to left-hand coord\n obj_poses[:,:3,3][:,1] *= -1.\n obj_poses[:,:,:2] *= -1.\n obj_num = len(obj_poses)\n for idx in range(obj_num):\n obj_pose = obj_poses[idx]\n if (idx % 5 == 0):\n cam_points_vis = get_cam_points_vis(W, H, intrinsics, obj_pose, [0., 0., 0.8], frustum_length=0.5)\n vis_all.append(cam_points_vis)\n\n # vis coordinates\n axis = open3d.geometry.TriangleMesh.create_coordinate_frame(size=2, origin=[0,0,0])\n \n open3d.visualization.draw_geometries(vis_all)\n" }, { "path": "eval/GVHMR/.gitignore", "content": ".vscode\n.hydra\ninputs\noutputs\n\n# All file or folders start with tmp will be ignored\ntmp* \n\n# Byte-compiled / optimized / DLL files\n__pycache__/\n*.py[cod]\n*$py.class\n\n# C extensions\n*.so\n\n# Distribution / packaging\n.Python\nbuild/\ndevelop-eggs/\ndist/\ndownloads/\neggs/\n.eggs/\nlib64/\nparts/\nsdist/\nvar/\nwheels/\nshare/python-wheels/\n*.egg-info/\n.installed.cfg\n*.egg\nMANIFEST\n\n# PyInstaller\n# Usually these files are written by a python script from a template\n# before PyInstaller builds the exe, so as to inject date/other infos into it.\n*.manifest\n*.spec\n\n# Installer logs\npip-log.txt\npip-delete-this-directory.txt\n\n# Unit test / coverage reports\nhtmlcov/\n.tox/\n.nox/\n.coverage\n.coverage.*\n.cache\nnosetests.xml\ncoverage.xml\n*.cover\n*.py,cover\n.hypothesis/\n.pytest_cache/\ncover/\n\n# Translations\n*.mo\n*.pot\n\n# Django stuff:\n*.log\nlocal_settings.py\ndb.sqlite3\ndb.sqlite3-journal\n\n# Flask stuff:\ninstance/\n.webassets-cache\n\n# Scrapy stuff:\n.scrapy\n\n# Sphinx documentation\ndocs/_build/\n\n# PyBuilder\n.pybuilder/\ntarget/\n\n# Jupyter Notebook\n.ipynb_checkpoints\n\n# IPython\nprofile_default/\nipython_config.py\n\n# pyenv\n# For a library or package, you might want to ignore these files since the code is\n# intended to run in multiple environments; otherwise, check them in:\n# .python-version\n\n# pipenv\n# According to pypa/pipenv#598, it is recommended to include Pipfile.lock in version control.\n# However, in case of collaboration, if having platform-specific dependencies or dependencies\n# having no cross-platform support, pipenv may install dependencies that don't work, or not\n# install all needed dependencies.\n#Pipfile.lock\n\n# poetry\n# Similar to Pipfile.lock, it is generally recommended to include poetry.lock in version control.\n# This is especially recommended for binary packages to ensure reproducibility, and is more\n# commonly ignored for libraries.\n# https://python-poetry.org/docs/basic-usage/#commit-your-poetrylock-file-to-version-control\n#poetry.lock\n\n# pdm\n# Similar to Pipfile.lock, it is generally recommended to include pdm.lock in version control.\n#pdm.lock\n# pdm stores project-wide configurations in .pdm.toml, but it is recommended to not include it\n# in version control.\n# https://pdm.fming.dev/#use-with-ide\n.pdm.toml\n\n# PEP 582; used by e.g. github.com/David-OConnor/pyflow and github.com/pdm-project/pdm\n__pypackages__/\n\n# Celery stuff\ncelerybeat-schedule\ncelerybeat.pid\n\n# SageMath parsed files\n*.sage.py\n\n# Environments\n.env\n.venv\nenv/\nvenv/\nENV/\nenv.bak/\nvenv.bak/\n\n# Spyder project settings\n.spyderproject\n.spyproject\n\n# Rope project settings\n.ropeproject\n\n# mkdocs documentation\n/site\n\n# mypy\n.mypy_cache/\n.dmypy.json\ndmypy.json\n\n# Pyre type checker\n.pyre/\n\n# pytype static type analyzer\n.pytype/\n\n# Cython debug symbols\ncython_debug/\n\n# \n.DS_Store/\n\n# PyCharm\n# JetBrains specific template is maintained in a separate JetBrains.gitignore that can\n# be found at https://github.com/github/gitignore/blob/main/Global/JetBrains.gitignore\n# and can be added to the global gitignore or merged into this file. For a more nuclear\n# option (not recommended) you can uncomment the following to ignore the entire idea folder.\n#.idea/\n\n# torchsparse\ntorchsparse\n\n# tensorboard\ntensorboard\n\n# glove\nglove" }, { "path": "eval/GVHMR/.gitmodules", "content": "[submodule \"third-party/DPVO\"]\n\tpath = third-party/DPVO\n\turl = https://github.com/princeton-vl/DPVO.git\n" }, { "path": "eval/GVHMR/LICENSE", "content": "Copyright 2022-2023 3D Vision Group at the State Key Lab of CAD&CG, \nZhejiang University. All Rights Reserved. \n\nFor more information see \nIf you use this software, please cite the corresponding publications \nlisted on the above website. \n\nPermission to use, copy, modify and distribute this software and its \ndocumentation for educational, research and non-profit purposes only. \nAny modification based on this work must be open-source and prohibited \nfor commercial use. \nYou must retain, in the source form of any derivative works that you \ndistribute, all copyright, patent, trademark, and attribution notices \nfrom the source form of this work. \n \nFor commercial uses of this software, please send email to xwzhou@zju.edu.cn" }, { "path": "eval/GVHMR/README.md", "content": "# GVHMR: World-Grounded Human Motion Recovery via Gravity-View Coordinates\n### [Project Page](https://zju3dv.github.io/gvhmr) | [Paper](https://arxiv.org/abs/2409.06662)\n\n> World-Grounded Human Motion Recovery via Gravity-View Coordinates \n> [Zehong Shen](https://zehongs.github.io/)\\*,\n[Huaijin Pi](https://phj128.github.io/)\\*,\n[Yan Xia](https://isshikihugh.github.io/scholar),\n[Zhi Cen](https://scholar.google.com/citations?user=Xyy-uFMAAAAJ),\n[Sida Peng](https://pengsida.net/)†,\n[Zechen Hu](https://zju3dv.github.io/gvhmr),\n[Hujun Bao](http://www.cad.zju.edu.cn/home/bao/),\n[Ruizhen Hu](https://csse.szu.edu.cn/staff/ruizhenhu/),\n[Xiaowei Zhou](https://xzhou.me/) \n> SIGGRAPH Asia 2024\n\n

\n \"animated\"\n

\n\n## Setup\n\nPlease see [installation](docs/INSTALL.md) for details.\n\n## Quick Start\n\n### [ Google Colab demo for GVHMR](https://colab.research.google.com/drive/1N9WSchizHv2bfQqkE9Wuiegw_OT7mtGj?usp=sharing)\n\n### [ HuggingFace demo for GVHMR](https://huggingface.co/spaces/LittleFrog/GVHMR)\n\n### Demo\nDemo entries are provided in `tools/demo`. Use `-s` to skip visual odometry if you know the camera is static, otherwise the camera will be estimated by DPVO.\nWe also provide a script `demo_folder.py` to inference a entire folder.\n```shell\npython tools/demo/demo.py --video=docs/example_video/tennis.mp4 -s\npython tools/demo/demo_folder.py -f inputs/demo/folder_in -d outputs/demo/folder_out -s\n```\n\n### Reproduce\n1. **Test**:\nTo reproduce the 3DPW, RICH, and EMDB results in a single run, use the following command:\n ```shell\n python tools/train.py global/task=gvhmr/test_3dpw_emdb_rich exp=gvhmr/mixed/mixed ckpt_path=inputs/checkpoints/gvhmr/gvhmr_siga24_release.ckpt\n ```\n To test individual datasets, change `global/task` to `gvhmr/test_3dpw`, `gvhmr/test_rich`, or `gvhmr/test_emdb`.\n\n2. **Train**:\nTo train the model, use the following command:\n ```shell\n # The gvhmr_siga24_release.ckpt is trained with 2x4090 for 420 epochs, note that different GPU settings may lead to different results.\n python tools/train.py exp=gvhmr/mixed/mixed\n ```\n During training, note that we do not employ post-processing as in the test script, so the global metrics results will differ (but should still be good for comparison with baseline methods).\n\n# Citation\n\nIf you find this code useful for your research, please use the following BibTeX entry.\n\n```\n@inproceedings{shen2024gvhmr,\n title={World-Grounded Human Motion Recovery via Gravity-View Coordinates},\n author={Shen, Zehong and Pi, Huaijin and Xia, Yan and Cen, Zhi and Peng, Sida and Hu, Zechen and Bao, Hujun and Hu, Ruizhen and Zhou, Xiaowei},\n booktitle={SIGGRAPH Asia Conference Proceedings},\n year={2024}\n}\n```\n\n# Acknowledgement\n\nWe thank the authors of\n[WHAM](https://github.com/yohanshin/WHAM),\n[4D-Humans](https://github.com/shubham-goel/4D-Humans),\nand [ViTPose-Pytorch](https://github.com/gpastal24/ViTPose-Pytorch) for their great works, without which our project/code would not be possible.\n" }, { "path": "eval/GVHMR/docs/INSTALL.md", "content": "# Install\n\n## Environment\n\n```bash\ngit clone https://github.com/zju3dv/GVHMR --recursive\ncd GVHMR\n\nconda create -y -n gvhmr python=3.10\nconda activate gvhmr\npip install -r requirements.txt\npip install -e .\n# to install gvhmr in other repo as editable, try adding \"python.analysis.extraPaths\": [\"path/to/your/package\"] to settings.json\n\n# DPVO\ncd third-party/DPVO\nwget https://gitlab.com/libeigen/eigen/-/archive/3.4.0/eigen-3.4.0.zip\nunzip eigen-3.4.0.zip -d thirdparty && rm -rf eigen-3.4.0.zip\npip install torch-scatter -f \"https://data.pyg.org/whl/torch-2.3.0+cu121.html\"\npip install numba pypose\nexport CUDA_HOME=/usr/local/cuda-12.1/\nexport PATH=$PATH:/usr/local/cuda-12.1/bin/\npip install -e .\n```\n\n## Inputs & Outputs\n\n```bash\nmkdir inputs\nmkdir outputs\n```\n\n**Weights**\n\n```bash\nmkdir -p inputs/checkpoints\n\n# 1. You need to sign up for downloading [SMPL](https://smpl.is.tue.mpg.de/) and [SMPLX](https://smpl-x.is.tue.mpg.de/). And the checkpoints should be placed in the following structure:\n\ninputs/checkpoints/\n├── body_models/smplx/\n│ └── SMPLX_{GENDER}.npz # SMPLX (We predict SMPLX params + evaluation)\n└── body_models/smpl/\n └── SMPL_{GENDER}.pkl # SMPL (rendering and evaluation)\n\n# 2. Download other pretrained models from Google-Drive (By downloading, you agree to the corresponding licences): https://drive.google.com/drive/folders/1eebJ13FUEXrKBawHpJroW0sNSxLjh9xD?usp=drive_link\n\ninputs/checkpoints/\n├── dpvo/\n│ └── dpvo.pth\n├── gvhmr/\n│ └── gvhmr_siga24_release.ckpt\n├── hmr2/\n│ └── epoch=10-step=25000.ckpt\n├── vitpose/\n│ └── vitpose-h-multi-coco.pth\n└── yolo/\n └── yolov8x.pt\n```\n\n**Data**\n\nWe provide preprocessed data for training and evaluation.\nNote that we do not intend to distribute the original datasets, and you need to download them (annotation, videos, etc.) from the original websites.\n*We're unable to provide the original data due to the license restrictions.*\nBy downloading the preprocessed data, you agree to the original dataset's terms of use and use the data for research purposes only.\n\nYou can download them from [Google-Drive](https://drive.google.com/drive/folders/10sEef1V_tULzddFxzCmDUpsIqfv7eP-P?usp=drive_link). Please place them in the \"inputs\" folder and execute the following commands:\n\n```bash\ncd inputs\n# Train\ntar -xzvf AMASS_hmr4d_support.tar.gz\ntar -xzvf BEDLAM_hmr4d_support.tar.gz\ntar -xzvf H36M_hmr4d_support.tar.gz\n# Test\ntar -xzvf 3DPW_hmr4d_support.tar.gz\ntar -xzvf EMDB_hmr4d_support.tar.gz\ntar -xzvf RICH_hmr4d_support.tar.gz\n\n# The folder structure should be like this:\ninputs/\n├── AMASS/hmr4d_support/\n├── BEDLAM/hmr4d_support/\n├── H36M/hmr4d_support/\n├── 3DPW/hmr4d_support/\n├── EMDB/hmr4d_support/\n└── RICH/hmr4d_support/\n```\n" }, { "path": "eval/GVHMR/download_eval_pose.sh", "content": "gdown https://drive.google.com/uc\\?id\\=1jMH2-ZC0ZBgtqej5Sp-E5ebBIX7mk3Xz\ngdown https://drive.google.com/uc\\?id\\=1iFcPSlcKb_rDNJ85UPoThdl22BqR2Xgh\n\nunzip eval_sets.zip\nrm -rf eval_sets.zip" }, { "path": "eval/GVHMR/eval.sh", "content": "python tools/demo/demo_folder.py -f eval_sets -d outputs/eval_sets_gvhmr -s\npython tools/eval_pose.py -f outputs/eval_sets_gvhmr_v2\n" }, { "path": "eval/GVHMR/hmr4d/__init__.py", "content": "import os\nfrom pathlib import Path\n\nPROJ_ROOT = Path(__file__).resolve().parents[1]\n\n\ndef os_chdir_to_proj_root():\n \"\"\"useful for running notebooks in different directories.\"\"\"\n os.chdir(PROJ_ROOT)\n" }, { "path": "eval/GVHMR/hmr4d/build_gvhmr.py", "content": "from omegaconf import OmegaConf\nfrom hmr4d import PROJ_ROOT\nfrom hydra.utils import instantiate\nfrom hmr4d.model.gvhmr.gvhmr_pl_demo import DemoPL\n\n\ndef build_gvhmr_demo():\n cfg = OmegaConf.load(PROJ_ROOT / \"hmr4d/configs/demo_gvhmr_model/siga24_release.yaml\")\n gvhmr_demo_pl: DemoPL = instantiate(cfg.model, _recursive_=False)\n gvhmr_demo_pl.load_pretrained_model(PROJ_ROOT / \"inputs/checkpoints/gvhmr/gvhmr_siga24_release.ckpt\")\n return gvhmr_demo_pl.eval()\n" }, { "path": "eval/GVHMR/hmr4d/configs/__init__.py", "content": "from dataclasses import dataclass\nfrom hydra.core.config_store import ConfigStore\nfrom hydra_zen import builds\n\nimport argparse\nfrom hydra import compose, initialize_config_module\nimport os\n\nos.environ[\"HYDRA_FULL_ERROR\"] = \"1\"\n\nMainStore = ConfigStore.instance()\n\n\ndef register_store_gvhmr():\n \"\"\"Register group options to MainStore\"\"\"\n from . import store_gvhmr\n\n\ndef parse_args_to_cfg():\n \"\"\"\n Use minimal Hydra API to parse args and return cfg.\n This function don't do _run_hydra which create log file hierarchy.\n \"\"\"\n parser = argparse.ArgumentParser()\n parser.add_argument(\"--config-name\", \"-cn\", default=\"train\")\n parser.add_argument(\n \"overrides\",\n nargs=\"*\",\n help=\"Any key=value arguments to override config values (use dots for.nested=overrides)\",\n )\n args = parser.parse_args()\n\n # Cfg\n with initialize_config_module(version_base=\"1.3\", config_module=f\"hmr4d.configs\"):\n cfg = compose(config_name=args.config_name, overrides=args.overrides)\n\n return cfg\n" }, { "path": "eval/GVHMR/hmr4d/configs/data/mocap/testY.yaml", "content": "# definition of lightning datamodule (dataset + dataloader)\n_target_: hmr4d.datamodule.mocap_trainX_testY.DataModule\n\ndataset_opts:\n test: ${test_datasets}\n\nloader_opts:\n test:\n batch_size: 1\n num_workers: 0\n" }, { "path": "eval/GVHMR/hmr4d/configs/data/mocap/trainX_testY.yaml", "content": "# definition of lightning datamodule (dataset + dataloader)\n_target_: hmr4d.datamodule.mocap_trainX_testY.DataModule\n\ndataset_opts:\n train: ${train_datasets}\n val: ${test_datasets}\n\nloader_opts:\n train:\n batch_size: 32\n num_workers: 8\n val:\n batch_size: 1\n num_workers: 1\n\nlimit_each_trainset: null" }, { "path": "eval/GVHMR/hmr4d/configs/demo.yaml", "content": "defaults:\n - _self_\n - model: gvhmr/gvhmr_pl_demo\n - network: gvhmr/relative_transformer\n - endecoder: gvhmr/v1_amass_local_bedlam_cam\n\npipeline:\n _target_: hmr4d.model.gvhmr.pipeline.gvhmr_pipeline.Pipeline\n args_denoiser3d: ${network}\n args:\n endecoder_opt: ${endecoder}\n normalize_cam_angvel: True\n weights: null\n static_conf: null\n\nckpt_path: inputs/checkpoints/gvhmr/gvhmr_siga24_release.ckpt\n\n# ================================ #\n# global setting #\n# ================================ #\n\nvideo_name: ???\noutput_root: outputs/demo\noutput_dir: \"${output_root}/${video_name}\"\npreprocess_dir: ${output_dir}/preprocess\nvideo_path: \"${output_dir}/0_input_video.mp4\"\n\n# Options\nstatic_cam: False\nverbose: False\n\npaths:\n bbx: ${preprocess_dir}/bbx.pt\n bbx_xyxy_video_overlay: ${preprocess_dir}/bbx_xyxy_video_overlay.mp4\n vit_features: ${preprocess_dir}/vit_features.pt\n vitpose: ${preprocess_dir}/vitpose.pt\n vitpose_video_overlay: ${preprocess_dir}/vitpose_video_overlay.mp4\n hmr4d_results: ${output_dir}/hmr4d_results.pt\n incam_video: ${output_dir}/1_incam.mp4\n global_video: ${output_dir}/2_global.mp4\n incam_global_horiz_video: ${output_dir}/${video_name}_3_incam_global_horiz.mp4\n slam: ${preprocess_dir}/slam_results.pt\n" }, { "path": "eval/GVHMR/hmr4d/configs/exp/gvhmr/mixed/mixed.yaml", "content": "# @package _global_\ndefaults:\n - override /data: mocap/trainX_testY\n - override /model: gvhmr/gvhmr_pl\n - override /endecoder: gvhmr/v1_amass_local_bedlam_cam\n - override /optimizer: adamw_2e-4\n - override /scheduler_cfg: epoch_half_200_350\n - override /train_datasets:\n - pure_motion_amass/v11\n - imgfeat_bedlam/v2\n - imgfeat_h36m/v1\n - imgfeat_3dpw/v1\n - override /test_datasets:\n - emdb1/v1_fliptest\n - emdb2/v1_fliptest\n - rich/all\n - 3dpw/fliptest\n - override /callbacks:\n - simple_ckpt_saver/every10e_top100\n - prog_bar/prog_reporter_every0.1\n - train_speed_timer/base\n - lr_monitor/pl\n - metric_emdb1\n - metric_emdb2\n - metric_rich\n - metric_3dpw\n - override /network: gvhmr/relative_transformer\n\nexp_name_base: mixed\nexp_name_var: \"\"\nexp_name: ${exp_name_base}${exp_name_var}\ndata_name: mocap_mixed_v1\n\npipeline:\n _target_: hmr4d.model.gvhmr.pipeline.gvhmr_pipeline.Pipeline\n args_denoiser3d: ${network}\n args:\n endecoder_opt: ${endecoder}\n normalize_cam_angvel: True\n weights:\n cr_j3d: 500.\n transl_c: 1.\n cr_verts: 500.\n j2d: 1000.\n verts2d: 1000.\n\n transl_w: 1.\n static_conf_bce: 1.\n\n static_conf:\n vel_thr: 0.15\n\ndata:\n loader_opts:\n train:\n batch_size: 128\n num_workers: 12\n\npl_trainer:\n precision: 16-mixed\n log_every_n_steps: 50\n gradient_clip_val: 0.5\n max_epochs: 500\n check_val_every_n_epoch: 10\n devices: 2\n\nlogger:\n _target_: pytorch_lightning.loggers.TensorBoardLogger\n save_dir: ${output_dir} # /save_dir/name/version/sub_dir\n name: \"\"\n version: \"tb\" # merge name and version\n" }, { "path": "eval/GVHMR/hmr4d/configs/global/debug/debug_train.yaml", "content": "# @package _global_\n\ndata_name: debug\nexp_name: debug\n\n# data:\n# limit_each_trainset: 40\n# loader_opts:\n# train:\n# batch_size: 4\n# num_workers: 0\n# val:\n# batch_size: 1\n# num_workers: 0\n\npl_trainer:\n limit_train_batches: 32\n limit_val_batches: 2\n check_val_every_n_epoch: 3\n enable_checkpointing: False\n devices: 1\n\ncallbacks:\n model_checkpoint: null\n" }, { "path": "eval/GVHMR/hmr4d/configs/global/debug/debug_train_limit_data.yaml", "content": "# @package _global_\n\ndata_name: debug\nexp_name: debug\n\ndata:\n limit_each_trainset: 40\n loader_opts:\n train:\n batch_size: 4\n num_workers: 0\n val:\n batch_size: 1\n num_workers: 0\n\npl_trainer:\n limit_val_batches: 2\n check_val_every_n_epoch: 3\n enable_checkpointing: False\n devices: 1\n\ncallbacks:\n model_checkpoint: null\n" }, { "path": "eval/GVHMR/hmr4d/configs/global/task/gvhmr/test_3dpw.yaml", "content": "# @package _global_\ndefaults:\n - override /data: mocap/testY\n - override /test_datasets:\n - 3dpw/fliptest\n - override /callbacks:\n - metric_3dpw\n - _self_\n\ntask: test\ndata_name: test_mocap\nckpt_path: ??? # will not override previous setting if already set\n\n# lightning utilities\npl_trainer:\n devices: 1\nlogger: null\n" }, { "path": "eval/GVHMR/hmr4d/configs/global/task/gvhmr/test_3dpw_emdb_rich.yaml", "content": "# @package _global_\ndefaults:\n - override /data: mocap/testY\n - override /test_datasets:\n - rich/all\n - emdb1/v1_fliptest\n - emdb2/v1_fliptest\n - 3dpw/fliptest\n - override /callbacks:\n - metric_rich\n - metric_emdb1\n - metric_emdb2\n - metric_3dpw\n - _self_\n\ntask: test\ndata_name: test_mocap\nckpt_path: ??? # will not override previous setting if already set\n\n# lightning utilities\npl_trainer:\n devices: 1\nlogger: null\n" }, { "path": "eval/GVHMR/hmr4d/configs/global/task/gvhmr/test_emdb.yaml", "content": "# @package _global_\ndefaults:\n - override /data: mocap/testY\n - override /test_datasets:\n - emdb1/v1_fliptest\n - emdb2/v1_fliptest\n - override /callbacks:\n - metric_emdb1\n - metric_emdb2\n - _self_\n\ntask: test\ndata_name: test_mocap\nckpt_path: ??? # will not override previous setting if already set\n\n# lightning utilities\npl_trainer:\n devices: 1\nlogger: null\n" }, { "path": "eval/GVHMR/hmr4d/configs/global/task/gvhmr/test_rich.yaml", "content": "# @package _global_\ndefaults:\n - override /data: mocap/testY\n - override /test_datasets:\n - rich/all\n - override /callbacks:\n - metric_rich\n - _self_\n\ntask: test\ndata_name: test_mocap\nckpt_path: ??? # will not override previous setting if already set\n\n# lightning utilities\npl_trainer:\n devices: 1\nlogger: null\n" }, { "path": "eval/GVHMR/hmr4d/configs/hydra/default.yaml", "content": "# enable color logging\ndefaults:\n - override hydra_logging: colorlog\n - override job_logging: colorlog\n\njob_logging:\n formatters:\n simple:\n datefmt: '%m/%d %H:%M:%S'\n format: '[%(asctime)s][%(levelname)s] %(message)s'\n colorlog:\n datefmt: '%m/%d %H:%M:%S'\n format: '[%(cyan)s%(asctime)s%(reset)s][%(log_color)s%(levelname)s%(reset)s] %(message)s'\n handlers:\n file:\n filename: ${output_dir}/${hydra.job.name}.log\n\nrun:\n dir: ${output_dir}" }, { "path": "eval/GVHMR/hmr4d/configs/siga24_release.yaml", "content": "pipeline:\n _target_: hmr4d.model.gvhmr.pipeline.gvhmr_pipeline.Pipeline\n args_denoiser3d: ${network}\n args:\n endecoder_opt: ${endecoder}\n normalize_cam_angvel: true\n weights: null\n static_conf: null\nmodel:\n _target_: hmr4d.model.gvhmr.gvhmr_pl_demo.DemoPL\n pipeline: ${pipeline}\nnetwork:\n _target_: hmr4d.network.gvhmr.relative_transformer.NetworkEncoderRoPEV2\n output_dim: 151\n max_len: 120\n kp2d_mapping: linear_v2\n cliffcam_dim: 3\n cam_angvel_dim: 6\n imgseq_dim: 1024\n f_imgseq_filter: null\n cond_ver: v1\n latent_dim: 512\n num_layers: 12\n num_heads: 8\n mlp_ratio: 4.0\n pred_cam_ver: v2\n pred_cam_dim: 3\n static_conf_dim: 6\n pred_coco17_dim: 0\n dropout: 0.1\n avgbeta: true\nendecoder:\n _target_: hmr4d.model.gvhmr.utils.endecoder.EnDecoder\n stats_name: MM_V1_AMASS_LOCAL_BEDLAM_CAM\n noise_pose_k: 10\n" }, { "path": "eval/GVHMR/hmr4d/configs/store_gvhmr.py", "content": "# Dataset\nimport hmr4d.dataset.pure_motion.amass\nimport hmr4d.dataset.emdb.emdb_motion_test\nimport hmr4d.dataset.rich.rich_motion_test\nimport hmr4d.dataset.threedpw.threedpw_motion_test\nimport hmr4d.dataset.threedpw.threedpw_motion_train\nimport hmr4d.dataset.bedlam.bedlam\nimport hmr4d.dataset.h36m.h36m\n\n# Trainer: Model Optimizer Loss\nimport hmr4d.model.gvhmr.gvhmr_pl\nimport hmr4d.model.gvhmr.utils.endecoder\nimport hmr4d.model.common_utils.optimizer\nimport hmr4d.model.common_utils.scheduler_cfg\n\n# Metric\nimport hmr4d.model.gvhmr.callbacks.metric_emdb\nimport hmr4d.model.gvhmr.callbacks.metric_rich\nimport hmr4d.model.gvhmr.callbacks.metric_3dpw\n\n\n# PL Callbacks\nimport hmr4d.utils.callbacks.simple_ckpt_saver\nimport hmr4d.utils.callbacks.train_speed_timer\nimport hmr4d.utils.callbacks.prog_bar\nimport hmr4d.utils.callbacks.lr_monitor\n\n# Networks\nimport hmr4d.network.gvhmr.relative_transformer\n" }, { "path": "eval/GVHMR/hmr4d/configs/train.yaml", "content": "# ================================ #\n# override #\n# ================================ #\n# specify default configuration; the order determines the override order\ndefaults:\n - _self_\n # pytorch-lightning\n - data: ???\n - model: ???\n - callbacks: null\n\n # system\n - hydra: default\n\n # utility groups that changes a lot\n - pipeline: null\n - network: null\n - optimizer: null\n - scheduler_cfg: default\n - train_datasets: null\n - test_datasets: null\n - endecoder: null # normalize/unnormalize data\n - refiner: null\n\n # global-override\n - exp: ??? # set \"data, model and callbacks\" in yaml\n - global/task: null # dump/test\n - global/hsearch: null # hyper-param search\n - global/debug: null # debug mode\n\n# ================================ #\n# global setting #\n# ================================ #\n# expirement information\ntask: fit # [fit, predict]\nexp_name: ???\ndata_name: ???\n\n# utilities in the entry file\noutput_dir: \"outputs/${data_name}/${exp_name}\"\nckpt_path: null\nresume_mode: null\nseed: 42\n\n# lightning default settings\npl_trainer:\n devices: 1\n num_sanity_val_steps: 0 # disable sanity check\n precision: 32\n inference_mode: False\n\nlogger: null\n" }, { "path": "eval/GVHMR/hmr4d/datamodule/mocap_trainX_testY.py", "content": "import pytorch_lightning as pl\nfrom pytorch_lightning.utilities.combined_loader import CombinedLoader\nfrom hydra.utils import instantiate\nfrom torch.utils.data import DataLoader, ConcatDataset, Subset\nfrom omegaconf import ListConfig, DictConfig\nfrom hmr4d.utils.pylogger import Log\nfrom numpy.random import choice\nfrom torch.utils.data import default_collate\n\n\nimport resource\n\nrlimit = resource.getrlimit(resource.RLIMIT_NOFILE)\nresource.setrlimit(resource.RLIMIT_NOFILE, (4096, rlimit[1]))\n\n\ndef collate_fn(batch):\n \"\"\"Handle meta and Add batch size to the return dict\n Args:\n batch: list of dict, each dict is a data point\n \"\"\"\n # Assume all keys in the batch are the same\n return_dict = {}\n for k in batch[0].keys():\n if k.startswith(\"meta\"): # data information, do not batch\n return_dict[k] = [d[k] for d in batch]\n else:\n return_dict[k] = default_collate([d[k] for d in batch])\n return_dict[\"B\"] = len(batch)\n return return_dict\n\n\nclass DataModule(pl.LightningDataModule):\n def __init__(self, dataset_opts: DictConfig, loader_opts: DictConfig, limit_each_trainset=None):\n \"\"\"This is a general datamodule that can be used for any dataset.\n Train uses ConcatDataset\n Val and Test use CombinedLoader, sequential, completely consumes ecah iterable sequentially, and returns a triplet (data, idx, iterable_idx)\n\n Args:\n dataset_opts: the target of the dataset. e.g. dataset_opts.train = {_target_: ..., limit_size: None}\n loader_opts: the options for the dataset\n limit_each_trainset: limit the size of each dataset, None means no limit, useful for debugging\n \"\"\"\n super().__init__()\n self.loader_opts = loader_opts\n self.limit_each_trainset = limit_each_trainset\n\n # Train uses concat dataset\n if \"train\" in dataset_opts:\n assert \"train\" in self.loader_opts, \"train not in loader_opts\"\n split_opts = dataset_opts.get(\"train\")\n assert isinstance(split_opts, DictConfig), \"split_opts should be a dict for each dataset\"\n dataset = []\n dataset_num = len(split_opts)\n for idx, (k, v) in enumerate(split_opts.items()):\n dataset_i = instantiate(v)\n if self.limit_each_trainset:\n dataset_i = Subset(dataset_i, choice(len(dataset_i), self.limit_each_trainset))\n dataset.append(dataset_i)\n Log.info(f\"[Train Dataset][{idx+1}/{dataset_num}]: name={k}, size={len(dataset[-1])}, {v._target_}\")\n dataset = ConcatDataset(dataset)\n self.trainset = dataset\n Log.info(f\"[Train Dataset][All]: ConcatDataset size={len(dataset)}\")\n Log.info(f\"\")\n\n # Val and Test use sequential dataset\n for split in (\"val\", \"test\"):\n if split not in dataset_opts:\n continue\n assert split in self.loader_opts, f\"split={split} not in loader_opts\"\n split_opts = dataset_opts.get(split)\n assert isinstance(split_opts, DictConfig), \"split_opts should be a dict for each dataset\"\n dataset = []\n dataset_num = len(split_opts)\n for idx, (k, v) in enumerate(split_opts.items()):\n dataset.append(instantiate(v))\n dataset_type = \"Val Dataset\" if split == \"val\" else \"Test Dataset\"\n Log.info(f\"[{dataset_type}][{idx+1}/{dataset_num}]: name={k}, size={len(dataset[-1])}, {v._target_}\")\n setattr(self, f\"{split}sets\", dataset)\n Log.info(f\"\")\n\n def train_dataloader(self):\n if hasattr(self, \"trainset\"):\n return DataLoader(\n self.trainset,\n shuffle=True,\n num_workers=self.loader_opts.train.num_workers,\n persistent_workers=True and self.loader_opts.train.num_workers > 0,\n batch_size=self.loader_opts.train.batch_size,\n drop_last=True,\n collate_fn=collate_fn,\n )\n else:\n return super().train_dataloader()\n\n def val_dataloader(self):\n if hasattr(self, \"valsets\"):\n loaders = []\n for valset in self.valsets:\n loaders.append(\n DataLoader(\n valset,\n shuffle=False,\n num_workers=self.loader_opts.val.num_workers,\n persistent_workers=True and self.loader_opts.val.num_workers > 0,\n batch_size=self.loader_opts.val.batch_size,\n collate_fn=collate_fn,\n )\n )\n return CombinedLoader(loaders, mode=\"sequential\")\n else:\n return None\n\n def test_dataloader(self):\n if hasattr(self, \"testsets\"):\n loaders = []\n for testset in self.testsets:\n loaders.append(\n DataLoader(\n testset,\n shuffle=False,\n num_workers=self.loader_opts.test.num_workers,\n persistent_workers=False,\n batch_size=self.loader_opts.test.batch_size,\n collate_fn=collate_fn,\n )\n )\n return CombinedLoader(loaders, mode=\"sequential\")\n else:\n return super().test_dataloader()\n" }, { "path": "eval/GVHMR/hmr4d/dataset/bedlam/bedlam.py", "content": "from pathlib import Path\nimport numpy as np\nimport torch\nfrom hmr4d.utils.pylogger import Log\nfrom pytorch3d.transforms import axis_angle_to_matrix, matrix_to_axis_angle\nfrom time import time\n\nfrom hmr4d.configs import MainStore, builds\nfrom hmr4d.utils.smplx_utils import make_smplx\nfrom hmr4d.utils.wis3d_utils import make_wis3d, add_motion_as_lines\nfrom hmr4d.utils.vis.renderer_utils import simple_render_mesh_background\nfrom hmr4d.utils.video_io_utils import read_video_np, save_video\n\nimport hmr4d.utils.matrix as matrix\nfrom hmr4d.utils.net_utils import get_valid_mask, repeat_to_max_len, repeat_to_max_len_dict\nfrom hmr4d.dataset.imgfeat_motion.base_dataset import ImgfeatMotionDatasetBase\nfrom hmr4d.dataset.bedlam.utils import mid2featname, mid2vname\nfrom hmr4d.utils.geo_transform import compute_cam_angvel, apply_T_on_points\nfrom hmr4d.utils.geo.hmr_global import get_T_w2c_from_wcparams, get_c_rootparam, get_R_c2gv\n\n\nclass BedlamDatasetV2(ImgfeatMotionDatasetBase):\n \"\"\"mid_to_valid_range and features are newly generated.\"\"\"\n\n MIDINDEX_TO_LOAD = {\n \"all60\": (\"mid_to_valid_range_all60.pt\", \"imgfeats/bedlam_all60\"),\n \"maxspan60\": (\"mid_to_valid_range_maxspan60.pt\", \"imgfeats/bedlam_maxspan60\"),\n }\n\n def __init__(\n self,\n mid_indices=[\"all60\", \"maxspan60\"],\n lazy_load=True, # Load from disk when needed\n random1024=False, # Faster loading for debugging\n ):\n self.root = Path(\"inputs/BEDLAM/hmr4d_support\")\n self.min_motion_frames = 60\n self.max_motion_frames = 120\n self.lazy_load = lazy_load\n self.random1024 = random1024\n\n # speficify mid_index to handle\n if not isinstance(mid_indices, list):\n mid_indices = [mid_indices]\n self.mid_indices = mid_indices\n assert all([m in self.MIDINDEX_TO_LOAD for m in mid_indices])\n\n super().__init__()\n\n def _load_dataset(self):\n Log.info(f\"[BEDLAM] Loading from {self.root}\")\n tic = time()\n # Load mid to valid range\n self.mid_to_valid_range = {}\n self.mid_to_imgfeat_dir = {}\n for m in self.mid_indices:\n fn, feat_dir = self.MIDINDEX_TO_LOAD[m]\n mid_to_valid_range_ = torch.load(self.root / fn)\n self.mid_to_valid_range.update(mid_to_valid_range_)\n self.mid_to_imgfeat_dir.update({mid: self.root / feat_dir for mid in mid_to_valid_range_})\n\n # Load motionfiles\n Log.info(f\"[BEDLAM] Start loading motion files\")\n if self.random1024: # Debug, faster loading\n try:\n Log.info(f\"[BEDLAM] Loading 1024 samples for debugging ...\")\n self.motion_files = torch.load(self.root / \"smplpose_v2_random1024.pth\")\n except:\n Log.info(f\"[BEDLAM] Not found, saving 1024 samples to disk ...\")\n self.motion_files = torch.load(self.root / \"smplpose_v2.pth\")\n keys = list(self.motion_files.keys())\n keys = np.random.choice(keys, 1024, replace=False)\n self.motion_files = {k: self.motion_files[k] for k in keys}\n torch.save(self.motion_files, self.root / \"smplpose_v2_random1024.pth\")\n self.mid_to_valid_range = {k: v for k, v in self.mid_to_valid_range.items() if k in self.motion_files}\n else:\n self.motion_files = torch.load(self.root / \"smplpose_v2.pth\")\n Log.info(f\"[BEDLAM] Motion files loaded. Elapsed: {time() - tic:.2f}s\")\n\n def _get_idx2meta(self):\n # sum_frame = sum([e-s for s, e in self.mid_to_valid_range.values()])\n self.idx2meta = list(self.mid_to_valid_range.keys())\n Log.info(f\"[BEDLAM] {len(self.idx2meta)} sequences. \")\n\n def _load_data(self, idx):\n mid = self.idx2meta[idx]\n # neutral smplx : \"pose\": (F, 63), \"trans\": (F, 3), \"beta\": (10),\n # and : \"skeleton\": (J, 3)\n data = self.motion_files[mid].copy()\n\n # Random select a subset\n range1, range2 = self.mid_to_valid_range[mid] # [range1, range2)\n mlength = range2 - range1\n min_motion_len = self.min_motion_frames\n max_motion_len = self.max_motion_frames\n\n if mlength < min_motion_len: # the minimal mlength is 30 when generating data\n start = range1\n length = mlength\n else:\n effect_max_motion_len = min(max_motion_len, mlength)\n length = np.random.randint(min_motion_len, effect_max_motion_len + 1) # [low, high)\n start = np.random.randint(range1, range2 - length + 1)\n end = start + length\n data[\"start_end\"] = (start, end)\n data[\"length\"] = length\n\n # Update data to a subset\n for k, v in data.items():\n if isinstance(v, torch.Tensor) and len(v.shape) > 1 and k != \"skeleton\":\n data[k] = v[start:end]\n\n # Load img(as feature) : {mid -> 'features', 'bbx_xys', 'img_wh', 'start_end'}\n imgfeat_dir = self.mid_to_imgfeat_dir[mid]\n f_img_dict = torch.load(imgfeat_dir / mid2featname(mid))\n\n # remap (start, end)\n start_mapped = start - f_img_dict[\"start_end\"][0]\n end_mapped = end - f_img_dict[\"start_end\"][0]\n\n data[\"f_imgseq\"] = f_img_dict[\"features\"][start_mapped:end_mapped].float() # (L, 1024)\n data[\"bbx_xys\"] = f_img_dict[\"bbx_xys\"][start_mapped:end_mapped].float() # (L, 4)\n data[\"img_wh\"] = f_img_dict[\"img_wh\"] # (2)\n data[\"kp2d\"] = torch.zeros((end - start), 17, 3) # (L, 17, 3) # do not provide kp2d\n\n return data\n\n def _process_data(self, data, idx):\n length = data[\"length\"]\n\n # SMPL params in cam\n body_pose = data[\"pose\"][:, 3:] # (F, 63)\n betas = data[\"beta\"].repeat(length, 1) # (F, 10)\n global_orient = data[\"global_orient_incam\"] # (F, 3)\n transl = data[\"trans_incam\"] + data[\"cam_ext\"][:, :3, 3] # (F, 3), bedlam convention\n smpl_params_c = {\"body_pose\": body_pose, \"betas\": betas, \"transl\": transl, \"global_orient\": global_orient}\n\n # SMPL params in world\n global_orient_w = data[\"pose\"][:, :3] # (F, 3)\n transl_w = data[\"trans\"] # (F, 3)\n smpl_params_w = {\"body_pose\": body_pose, \"betas\": betas, \"transl\": transl_w, \"global_orient\": global_orient_w}\n\n gravity_vec = torch.tensor([0, -1, 0], dtype=torch.float32) # (3), BEDLAM is ay\n T_w2c = get_T_w2c_from_wcparams(\n global_orient_w=global_orient_w,\n transl_w=transl_w,\n global_orient_c=global_orient,\n transl_c=transl,\n offset=data[\"skeleton\"][0],\n ) # (F, 4, 4)\n R_c2gv = get_R_c2gv(T_w2c[:, :3, :3], gravity_vec) # (F, 3, 3)\n\n # cam_angvel (slightly different from WHAM)\n cam_angvel = compute_cam_angvel(T_w2c[:, :3, :3]) # (F, 6)\n\n # Returns: do not forget to make it batchable! (last lines)\n max_len = self.max_motion_frames\n return_data = {\n \"meta\": {\"data_name\": \"bedlam\", \"idx\": idx},\n \"length\": length,\n \"smpl_params_c\": smpl_params_c,\n \"smpl_params_w\": smpl_params_w,\n \"R_c2gv\": R_c2gv, # (F, 3, 3)\n \"gravity_vec\": gravity_vec, # (3)\n \"bbx_xys\": data[\"bbx_xys\"], # (F, 3)\n \"K_fullimg\": data[\"cam_int\"], # (F, 3, 3)\n \"f_imgseq\": data[\"f_imgseq\"], # (F, D)\n \"kp2d\": data[\"kp2d\"], # (F, 17, 3)\n \"cam_angvel\": cam_angvel, # (F, 6)\n \"mask\": {\n \"valid\": get_valid_mask(max_len, length),\n \"vitpose\": False,\n \"bbx_xys\": True,\n \"f_imgseq\": True,\n \"spv_incam_only\": False,\n },\n }\n\n if False: # check transformation, wis3d: sampled motion (global, incam)\n wis3d = make_wis3d(name=\"debug-data-bedlam\")\n smplx = make_smplx(\"supermotion\")\n\n # global\n smplx_out = smplx(**smpl_params_w)\n w_gt_joints = smplx_out.joints\n add_motion_as_lines(w_gt_joints, wis3d, name=\"w-gt_joints\")\n\n # incam\n smplx_out = smplx(**smpl_params_c)\n c_gt_joints = smplx_out.joints\n add_motion_as_lines(c_gt_joints, wis3d, name=\"c-gt_joints\")\n\n # Check transformation works correctly\n print(\"T_w2c\", (apply_T_on_points(w_gt_joints, T_w2c) - c_gt_joints).abs().max())\n R_c, t_c = get_c_rootparam(\n smpl_params_w[\"global_orient\"], smpl_params_w[\"transl\"], T_w2c, data[\"skeleton\"][0]\n )\n print(\"transl_c\", (t_c - smpl_params_c[\"transl\"]).abs().max())\n R_diff = matrix_to_axis_angle(\n (axis_angle_to_matrix(R_c) @ axis_angle_to_matrix(smpl_params_c[\"global_orient\"]).transpose(-1, -2))\n ).norm(dim=-1)\n print(\"global_orient_c\", R_diff.abs().max()) # < 1e-6\n\n skeleton_beta = smplx.get_skeleton(smpl_params_c[\"betas\"])\n print(\"Skeleton\", (skeleton_beta[0] - data[\"skeleton\"]).abs().max()) # (1.2e-7)\n\n if False: # cam-overlay\n smplx = make_smplx(\"supermotion\")\n\n # *. original bedlam param\n # mid = self.idx2meta[idx]\n # video_path = \"-\".join(mid.replace(\"bedlam_data/\", \"inputs/bedlam/\").split(\"-\")[:-1])\n # npz_file = \"inputs/bedlam/processed_labels/20221024_3-10_100_batch01handhair_static_highSchoolGym.npz\"\n # params = np.load(npz_file, allow_pickle=True)\n # mid2index = {}\n # for j in tqdm(range(len(params[\"video_name\"]))):\n # k = params[\"video_name\"][j] + \"-\" + params[\"sub\"][j]\n # mid2index[k] = j\n # betas = params['shape'][mid2index[mid]][:length]\n # global_orient_incam = torch.from_numpy(params['pose_cam'][121][:, :3])\n # body_pose = torch.from_numpy(params['pose_cam'][121][:, 3:66])\n # transl_incam = torch.from_numpy(params[\"trans_cam\"][121])\n smplx_out = smplx(**smpl_params_c)\n\n # ----- Render Overlay ----- #\n mid = self.idx2meta[idx]\n images = read_video_np(self.root / \"videos\" / mid2vname(mid), data[\"start_end\"][0], data[\"start_end\"][1])\n render_dict = {\n \"K\": data[\"cam_int\"][:1], # only support batch-size 1\n \"faces\": smplx.faces,\n \"verts\": smplx_out.vertices,\n \"background\": images,\n }\n img_overlay = simple_render_mesh_background(render_dict)\n save_video(img_overlay, \"tmp.mp4\", crf=23)\n\n # Batchable\n return_data[\"smpl_params_c\"] = repeat_to_max_len_dict(return_data[\"smpl_params_c\"], max_len)\n return_data[\"smpl_params_w\"] = repeat_to_max_len_dict(return_data[\"smpl_params_w\"], max_len)\n return_data[\"R_c2gv\"] = repeat_to_max_len(return_data[\"R_c2gv\"], max_len)\n return_data[\"bbx_xys\"] = repeat_to_max_len(return_data[\"bbx_xys\"], max_len)\n return_data[\"K_fullimg\"] = repeat_to_max_len(return_data[\"K_fullimg\"], max_len)\n return_data[\"f_imgseq\"] = repeat_to_max_len(return_data[\"f_imgseq\"], max_len)\n return_data[\"kp2d\"] = repeat_to_max_len(return_data[\"kp2d\"], max_len)\n return_data[\"cam_angvel\"] = repeat_to_max_len(return_data[\"cam_angvel\"], max_len)\n return return_data\n\n\ngroup_name = \"train_datasets/imgfeat_bedlam\"\nMainStore.store(name=\"v2\", node=builds(BedlamDatasetV2), group=group_name)\nMainStore.store(name=\"v2_random1024\", node=builds(BedlamDatasetV2, random1024=True), group=group_name)\n" }, { "path": "eval/GVHMR/hmr4d/dataset/bedlam/utils.py", "content": "import torch\nimport numpy as np\nfrom pathlib import Path\n\nresource_dir = Path(__file__).parent / \"resource\"\n\n\ndef mid2vname(mid):\n \"\"\"vname = {scene}/{seq}, Note that it ends with .mp4\"\"\"\n # mid example: \"inputs/bedlam/bedlam_download/20221011_1_250_batch01hand_closeup_suburb_a/mp4/seq_000001.mp4-rp_emma_posed_008\"\n # -> vname: 20221011_1_250_batch01hand_closeup_suburb_a/seq_000001.mp4\n scene = mid.split(\"/\")[-3]\n seq = mid.split(\"/\")[-1].split(\"-\")[0]\n vname = f\"{scene}/{seq}\"\n return vname\n\n\ndef mid2featname(mid):\n \"\"\"featname = {scene}/{seqsubj}, Note that it ends with .pt (extra)\"\"\"\n # mid example: \"inputs/bedlam/bedlam_download/20221011_1_250_batch01hand_closeup_suburb_a/mp4/seq_000001.mp4-rp_emma_posed_008\"\n # -> featname: 20221011_1_250_batch01hand_closeup_suburb_a/seq_000001.mp4-rp_emma_posed_008.pt\n scene = mid.split(\"/\")[-3]\n seqsubj = mid.split(\"/\")[-1]\n featname = f\"{scene}/{seqsubj}.pt\"\n return featname\n\n\ndef featname2mid(featname):\n \"\"\"reverse func of mid2featname, Note that it removes .pt (extra)\"\"\"\n # featname example: 20221011_1_250_batch01hand_closeup_suburb_a/seq_000001.mp4-rp_emma_posed_008.pt\n # -> mid: inputs/bedlam/bedlam_download/20221011_1_250_batch01hand_closeup_suburb_a/mp4/seq_000001.mp4-rp_emma_posed_008\n scene = featname.split(\"/\")[0]\n seqsubj = featname.split(\"/\")[1].strip(\".pt\")\n mid = f\"inputs/bedlam/bedlam_download/{scene}/mp4/{seqsubj}\"\n return mid\n\n\ndef load_vname2lwh():\n return torch.load(resource_dir / \"vname2lwh.pt\")\n" }, { "path": "eval/GVHMR/hmr4d/dataset/emdb/emdb_motion_test.py", "content": "from pathlib import Path\nimport numpy as np\nimport torch\nfrom torch.utils import data\nfrom hmr4d.utils.pylogger import Log\nfrom hmr4d.utils.wis3d_utils import make_wis3d, add_motion_as_lines\n\nfrom hmr4d.utils.geo_transform import compute_cam_angvel\nfrom pytorch3d.transforms import quaternion_to_matrix\nfrom hmr4d.utils.geo.hmr_cam import estimate_K, resize_K\nfrom hmr4d.utils.geo.flip_utils import flip_kp2d_coco17\n\nfrom .utils import EMDB1_NAMES, EMDB2_NAMES\n\nVID_PRESETS = {1: EMDB1_NAMES, 2: EMDB2_NAMES}\n\n\nfrom hmr4d.configs import MainStore, builds\n\n\nclass EmdbSmplFullSeqDataset(data.Dataset):\n def __init__(self, split=1, flip_test=False):\n \"\"\"\n split: 1 for EMDB-1, 2 for EMDB-2\n flip_test: if True, extra flip data will be returned\n \"\"\"\n super().__init__()\n self.dataset_name = \"EMDB\"\n self.split = split\n self.dataset_id = f\"EMDB_{split}\"\n Log.info(f\"[{self.dataset_name}] Full sequence, split={split}\")\n\n # Load evaluation protocol from WHAM labels\n tic = Log.time()\n self.emdb_dir = Path(\"inputs/EMDB/hmr4d_support\")\n # 'name', 'gender', 'smpl_params', 'mask', 'K_fullimg', 'T_w2c', 'bbx_xys', 'kp2d', 'features'\n self.labels = torch.load(self.emdb_dir / \"emdb_vit_v4.pt\")\n self.cam_traj = torch.load(self.emdb_dir / \"emdb_dpvo_traj.pt\") # estimated with DPVO\n\n # Setup dataset index\n self.idx2meta = []\n for vid in VID_PRESETS[split]:\n seq_length = len(self.labels[vid][\"mask\"])\n self.idx2meta.append((vid, 0, seq_length)) # start=0, end=seq_length\n Log.info(f\"[{self.dataset_name}] {len(self.idx2meta)} sequences. Elapsed: {Log.time() - tic:.2f}s\")\n\n # If flip_test is enabled, we will return extra data for flipped test\n self.flip_test = flip_test\n if self.flip_test:\n Log.info(f\"[{self.dataset_name}] Flip test enabled\")\n\n def __len__(self):\n return len(self.idx2meta)\n\n def _load_data(self, idx):\n data = {}\n\n # [vid, start, end]\n vid, start, end = self.idx2meta[idx]\n length = end - start\n meta = {\"dataset_id\": self.dataset_id, \"vid\": vid, \"vid-start-end\": (start, end)}\n data.update({\"meta\": meta, \"length\": length})\n\n label = self.labels[vid]\n\n # smpl_params in world\n gender = label[\"gender\"]\n smpl_params = label[\"smpl_params\"]\n mask = label[\"mask\"]\n data.update({\"smpl_params\": smpl_params, \"gender\": gender, \"mask\": mask})\n\n # camera\n # K_fullimg = label[\"K_fullimg\"] # We use estimated K\n width_height = (1440, 1920) if vid != \"P0_09_outdoor_walk\" else (720, 960)\n K_fullimg = estimate_K(*width_height)\n T_w2c = label[\"T_w2c\"]\n data.update({\"K_fullimg\": K_fullimg, \"T_w2c\": T_w2c})\n\n # R_w2c -> cam_angvel\n use_DPVO = False\n if use_DPVO:\n traj = self.cam_traj[data[\"meta\"][\"vid\"]] # (L, 7)\n R_w2c = quaternion_to_matrix(traj[:, [6, 3, 4, 5]]).mT # (L, 3, 3)\n else: # GT\n R_w2c = data[\"T_w2c\"][:, :3, :3] # (L, 3, 3)\n data[\"cam_angvel\"] = compute_cam_angvel(R_w2c) # (L, 6)\n\n # image bbx, features\n bbx_xys = label[\"bbx_xys\"]\n f_imgseq = label[\"features\"]\n kp2d = label[\"kp2d\"]\n data.update({\"bbx_xys\": bbx_xys, \"f_imgseq\": f_imgseq, \"kp2d\": kp2d})\n\n # to render a video\n video_path = self.emdb_dir / f\"videos/{vid}.mp4\"\n frame_id = torch.where(mask)[0].long()\n resize_factor = 0.5\n width_height_render = torch.tensor(width_height) * resize_factor\n K_render = resize_K(K_fullimg, resize_factor)\n bbx_xys_render = bbx_xys * resize_factor\n data[\"meta_render\"] = {\n \"split\": self.split,\n \"name\": vid,\n \"video_path\": str(video_path),\n \"resize_factor\": resize_factor,\n \"frame_id\": frame_id,\n \"width_height\": width_height_render.int(),\n \"K\": K_render,\n \"bbx_xys\": bbx_xys_render,\n \"R_cam_type\": \"DPVO\" if use_DPVO else \"GtGyro\",\n }\n\n # if enable flip_test\n if self.flip_test:\n imgfeat_dir = self.emdb_dir / \"imgfeats/emdb_flip\"\n f_img_dict = torch.load(imgfeat_dir / f\"{vid}.pt\")\n\n flipped_bbx_xys = f_img_dict[\"bbx_xys\"].float() # (L, 3)\n flipped_features = f_img_dict[\"features\"].float() # (L, 1024)\n width = width_height[0]\n flipped_kp2d = flip_kp2d_coco17(kp2d, width) # (L, 17, 3)\n\n R_flip_x = torch.tensor([[-1, 0, 0], [0, 1, 0], [0, 0, 1]]).float()\n flipped_R_w2c = R_flip_x @ R_w2c.clone()\n\n data_flip = {\n \"bbx_xys\": flipped_bbx_xys,\n \"f_imgseq\": flipped_features,\n \"kp2d\": flipped_kp2d,\n \"cam_angvel\": compute_cam_angvel(flipped_R_w2c),\n }\n data[\"flip_test\"] = data_flip\n\n return data\n\n def _process_data(self, data):\n length = data[\"length\"]\n data[\"K_fullimg\"] = data[\"K_fullimg\"][None].repeat(length, 1, 1)\n return data\n\n def __getitem__(self, idx):\n data = self._load_data(idx)\n data = self._process_data(data)\n return data\n\n\n# EMDB-1 and EMDB-2\nMainStore.store(\n name=\"v1\",\n node=builds(EmdbSmplFullSeqDataset, populate_full_signature=True),\n group=\"test_datasets/emdb1\",\n)\nMainStore.store(\n name=\"v1_fliptest\",\n node=builds(EmdbSmplFullSeqDataset, flip_test=True, populate_full_signature=True),\n group=\"test_datasets/emdb1\",\n)\nMainStore.store(\n name=\"v1\",\n node=builds(EmdbSmplFullSeqDataset, split=2, populate_full_signature=True),\n group=\"test_datasets/emdb2\",\n)\nMainStore.store(\n name=\"v1_fliptest\",\n node=builds(EmdbSmplFullSeqDataset, split=2, flip_test=True, populate_full_signature=True),\n group=\"test_datasets/emdb2\",\n)\n" }, { "path": "eval/GVHMR/hmr4d/dataset/emdb/utils.py", "content": "import torch\nimport pickle\nimport numpy as np\nfrom pathlib import Path\nfrom tqdm import tqdm\nfrom hmr4d.utils.geo_transform import convert_lurb_to_bbx_xys\nfrom hmr4d.utils.video_io_utils import get_video_lwh\n\n\ndef name_to_subfolder(name):\n return f\"{name[:2]}/{name[3:]}\"\n\n\ndef name_to_local_pkl_path(name):\n return f\"{name_to_subfolder(name)}/{name}_data.pkl\"\n\n\ndef load_raw_pkl(fp):\n annot = pickle.load(open(fp, \"rb\"))\n annot[\"subfolder\"] = name_to_subfolder(annot[\"name\"])\n return annot\n\n\ndef load_pkl(fp):\n annot = pickle.load(open(fp, \"rb\"))\n # ['gender', 'name', 'emdb1', 'emdb2', 'n_frames', 'good_frames_mask', 'camera', 'smpl', 'kp2d', 'bboxes', 'subfolder']\n data = {}\n\n F = annot[\"n_frames\"]\n smpl_params = {\n \"body_pose\": annot[\"smpl\"][\"poses_body\"], # (F, 69)\n \"betas\": annot[\"smpl\"][\"betas\"][None].repeat(F, axis=0), # (F, 10)\n \"global_orient\": annot[\"smpl\"][\"poses_root\"], # (F, 3)\n \"transl\": annot[\"smpl\"][\"trans\"], # (F, 3)\n }\n smpl_params = {k: torch.from_numpy(v).float() for k, v in smpl_params.items()}\n\n data[\"name\"] = annot[\"name\"]\n data[\"gender\"] = annot[\"gender\"]\n data[\"smpl_params\"] = smpl_params\n data[\"mask\"] = torch.from_numpy(annot[\"good_frames_mask\"]).bool() # (L,)\n data[\"K_fullimg\"] = torch.from_numpy(annot[\"camera\"][\"intrinsics\"]).float() # (3, 3)\n data[\"T_w2c\"] = torch.from_numpy(annot[\"camera\"][\"extrinsics\"]).float() # (L, 4, 4)\n bbx_lurb = torch.from_numpy(annot[\"bboxes\"][\"bboxes\"]).float()\n data[\"bbx_xys\"] = convert_lurb_to_bbx_xys(bbx_lurb) # (L, 3)\n\n return data\n\n\nEMDB1_LIST = [\n \"P1/14_outdoor_climb/P1_14_outdoor_climb_data.pkl\",\n \"P2/23_outdoor_hug_tree/P2_23_outdoor_hug_tree_data.pkl\",\n \"P3/31_outdoor_workout/P3_31_outdoor_workout_data.pkl\",\n \"P3/32_outdoor_soccer_warmup_a/P3_32_outdoor_soccer_warmup_a_data.pkl\",\n \"P3/33_outdoor_soccer_warmup_b/P3_33_outdoor_soccer_warmup_b_data.pkl\",\n \"P5/42_indoor_dancing/P5_42_indoor_dancing_data.pkl\",\n \"P5/44_indoor_rom/P5_44_indoor_rom_data.pkl\",\n \"P6/49_outdoor_big_stairs_down/P6_49_outdoor_big_stairs_down_data.pkl\", # DUPLICATE\n \"P6/50_outdoor_workout/P6_50_outdoor_workout_data.pkl\",\n \"P6/51_outdoor_dancing/P6_51_outdoor_dancing_data.pkl\",\n \"P7/57_outdoor_rock_chair/P7_57_outdoor_rock_chair_data.pkl\", # DUPLICATE\n \"P7/59_outdoor_rom/P7_59_outdoor_rom_data.pkl\",\n \"P7/60_outdoor_workout/P7_60_outdoor_workout_data.pkl\",\n \"P8/64_outdoor_skateboard/P8_64_outdoor_skateboard_data.pkl\", # DUPLICATE\n \"P8/68_outdoor_handstand/P8_68_outdoor_handstand_data.pkl\",\n \"P8/69_outdoor_cartwheel/P8_69_outdoor_cartwheel_data.pkl\",\n \"P9/76_outdoor_sitting/P9_76_outdoor_sitting_data.pkl\",\n]\nEMDB1_NAMES = [\"_\".join(p.split(\"/\")[:2]) for p in EMDB1_LIST]\n\n\nEMDB2_LIST = [\n \"P0/09_outdoor_walk/P0_09_outdoor_walk_data.pkl\",\n \"P2/19_indoor_walk_off_mvs/P2_19_indoor_walk_off_mvs_data.pkl\",\n \"P2/20_outdoor_walk/P2_20_outdoor_walk_data.pkl\",\n \"P2/24_outdoor_long_walk/P2_24_outdoor_long_walk_data.pkl\",\n \"P3/27_indoor_walk_off_mvs/P3_27_indoor_walk_off_mvs_data.pkl\",\n \"P3/28_outdoor_walk_lunges/P3_28_outdoor_walk_lunges_data.pkl\",\n \"P3/29_outdoor_stairs_up/P3_29_outdoor_stairs_up_data.pkl\",\n \"P3/30_outdoor_stairs_down/P3_30_outdoor_stairs_down_data.pkl\",\n \"P4/35_indoor_walk/P4_35_indoor_walk_data.pkl\",\n \"P4/36_outdoor_long_walk/P4_36_outdoor_long_walk_data.pkl\",\n \"P4/37_outdoor_run_circle/P4_37_outdoor_run_circle_data.pkl\",\n \"P5/40_indoor_walk_big_circle/P5_40_indoor_walk_big_circle_data.pkl\",\n \"P6/48_outdoor_walk_downhill/P6_48_outdoor_walk_downhill_data.pkl\",\n \"P6/49_outdoor_big_stairs_down/P6_49_outdoor_big_stairs_down_data.pkl\", # DUPLICATE\n \"P7/55_outdoor_walk/P7_55_outdoor_walk_data.pkl\",\n \"P7/56_outdoor_stairs_up_down/P7_56_outdoor_stairs_up_down_data.pkl\",\n \"P7/57_outdoor_rock_chair/P7_57_outdoor_rock_chair_data.pkl\", # DUPLICATE\n \"P7/58_outdoor_parcours/P7_58_outdoor_parcours_data.pkl\",\n \"P7/61_outdoor_sit_lie_walk/P7_61_outdoor_sit_lie_walk_data.pkl\",\n \"P8/64_outdoor_skateboard/P8_64_outdoor_skateboard_data.pkl\", # DUPLICATE\n \"P8/65_outdoor_walk_straight/P8_65_outdoor_walk_straight_data.pkl\",\n \"P9/77_outdoor_stairs_up/P9_77_outdoor_stairs_up_data.pkl\",\n \"P9/78_outdoor_stairs_up_down/P9_78_outdoor_stairs_up_down_data.pkl\",\n \"P9/79_outdoor_walk_rectangle/P9_79_outdoor_walk_rectangle_data.pkl\",\n \"P9/80_outdoor_walk_big_circle/P9_80_outdoor_walk_big_circle_data.pkl\",\n]\nEMDB2_NAMES = [\"_\".join(p.split(\"/\")[:2]) for p in EMDB2_LIST]\nEMDB_NAMES = list(sorted(set(EMDB1_NAMES + EMDB2_NAMES)))\n\n\ndef _check_annot(emdb_raw_dir=Path(\"inputs/EMDB/EMDB\")):\n for pkl_local_path in set(EMDB1_LIST + EMDB2_LIST):\n annot = load_raw_pkl(emdb_raw_dir / pkl_local_path)\n if any((annot[\"bboxes\"][\"invalid_idxs\"] != np.where(~annot[\"good_frames_mask\"])[0])):\n print(annot[\"name\"])\n\n\ndef _check_length(emdb_raw_dir=Path(\"inputs/EMDB/EMDB\"), emdb_hmr4d_support_dir=Path(\"inputs/EMDB/hmr4d_support\")):\n lengths = []\n for local_pkl_path in tqdm(set(EMDB1_LIST + EMDB2_LIST)):\n data = load_pkl(emdb_raw_dir / local_pkl_path)\n video_path = emdb_hmr4d_support_dir / \"videos\" / f\"{data['name']}.mp4\"\n length, width, height = get_video_lwh(video_path)\n lengths.append(length)\n print(sorted(lengths))\n\n video_ram = length[-1] * (width / 4) * (height / 4) * 3 / 1e6\n print(f\"Video RAM for {lengths[-1]} x {width} x {height}: {video_ram:.2f} MB\")\n" }, { "path": "eval/GVHMR/hmr4d/dataset/h36m/camera-parameters.json", "content": "{\n \"intrinsics\": {\n \"54138969\": {\n \"calibration_matrix\": [\n [\n 1145.04940458804,\n 0.0,\n 512.541504956548\n ],\n [\n 0.0,\n 1143.78109572365,\n 515.4514869776\n ],\n [\n 0.0,\n 0.0,\n 1.0\n ]\n ],\n \"distortion\": [\n -0.207098910824901,\n 0.247775183068982,\n -0.00142447157470321,\n -0.000975698859470499,\n -0.00307515035078854\n ]\n },\n \"55011271\": {\n \"calibration_matrix\": [\n [\n 1149.67569986785,\n 0.0,\n 508.848621645943\n ],\n [\n 0.0,\n 1147.59161666764,\n 508.064917088557\n ],\n [\n 0.0,\n 0.0,\n 1.0\n ]\n ],\n \"distortion\": [\n -0.194213629607385,\n 0.240408539138292,\n -0.0027408943961907,\n -0.001619026613787,\n 0.00681997559022603\n ]\n },\n \"58860488\": {\n \"calibration_matrix\": [\n [\n 1149.14071676148,\n 0.0,\n 519.815837182153\n ],\n [\n 0.0,\n 1148.7989685676,\n 501.402658888552\n ],\n [\n 0.0,\n 0.0,\n 1.0\n ]\n ],\n \"distortion\": [\n -0.208338188251856,\n 0.255488007488945,\n -0.000759999321030303,\n 0.00148438698385668,\n -0.00246049749891915\n ]\n },\n \"60457274\": {\n \"calibration_matrix\": [\n [\n 1145.51133842318,\n 0.0,\n 514.968197319863\n ],\n [\n 0.0,\n 1144.77392807652,\n 501.882018537695\n ],\n [\n 0.0,\n 0.0,\n 1.0\n ]\n ],\n \"distortion\": [\n -0.198384093827848,\n 0.218323676298049,\n -0.00181336200488089,\n -0.000587205583421232,\n -0.00894780704152122\n ]\n }\n },\n \"extrinsics\": {\n \"S1\": {\n \"54138969\": {\n \"R\": [\n [\n -0.9153617321513369,\n 0.40180836633680234,\n 0.02574754463350265\n ],\n [\n 0.051548117060134555,\n 0.1803735689384521,\n -0.9822464900705729\n ],\n [\n -0.399319034032262,\n -0.8977836111057917,\n -0.185819527201491\n ]\n ],\n \"t\": [\n [\n -346.05078140028075\n ],\n [\n 546.9807793144001\n ],\n [\n 5474.481087434061\n ]\n ]\n },\n \"55011271\": {\n \"R\": [\n [\n 0.9281683400814921,\n 0.3721538354721445,\n 0.002248380248018696\n ],\n [\n 0.08166409428175585,\n -0.1977722953267526,\n -0.976840363061605\n ],\n [\n -0.3630902204349604,\n 0.9068559102440475,\n -0.21395758897485287\n ]\n ],\n \"t\": [\n [\n 251.42516271750836\n ],\n [\n 420.9422103702068\n ],\n [\n 5588.195881837821\n ]\n ]\n },\n \"58860488\": {\n \"R\": [\n [\n -0.9141549520542256,\n -0.4027780222811878,\n -0.045722952682337906\n ],\n [\n -0.04562341383935875,\n 0.21430849526487267,\n -0.9756999400261069\n ],\n [\n 0.40278930937200774,\n -0.889854894701693,\n -0.214287280609606\n ]\n ],\n \"t\": [\n [\n 480.482559565337\n ],\n [\n 253.83237471361554\n ],\n [\n 5704.2076793704555\n ]\n ]\n },\n \"60457274\": {\n \"R\": [\n [\n 0.9141562410494211,\n -0.40060705854636447,\n 0.061905989962380774\n ],\n [\n -0.05641000739510571,\n -0.2769531972942539,\n -0.9592261660183036\n ],\n [\n 0.40141783470104664,\n 0.8733904688919611,\n -0.2757767409202658\n ]\n ],\n \"t\": [\n [\n 51.88347637559197\n ],\n [\n 378.4208425426766\n ],\n [\n 4406.149140878431\n ]\n ]\n }\n },\n \"S2\": {\n \"54138969\": {\n \"R\": [\n [\n -0.9072826056858586,\n 0.4200536513985309,\n 0.019829356183203237\n ],\n [\n 0.06404223092375372,\n 0.18462275321422528,\n -0.9807206695353717\n ],\n [\n -0.4156162485733534,\n -0.8885208882982778,\n -0.1944061855483302\n ]\n ],\n \"t\": [\n [\n -253.9473271477662\n ],\n [\n 543.369692173605\n ],\n [\n 5522.981999493327\n ]\n ]\n },\n \"55011271\": {\n \"R\": [\n [\n 0.9195695689704942,\n 0.3926824530407384,\n 0.013867187794489123\n ],\n [\n 0.09616327770610274,\n -0.190692439252443,\n -0.9769283584955307\n ],\n [\n -0.38097825639298405,\n 0.8996871037676718,\n -0.21311659595137136\n ]\n ],\n \"t\": [\n [\n 123.3506735789221\n ],\n [\n 401.02404156275884\n ],\n [\n 5743.522551411228\n ]\n ]\n },\n \"58860488\": {\n \"R\": [\n [\n -0.9231022562305128,\n -0.3793547679556717,\n -0.06302526930870815\n ],\n [\n -0.023520852900409527,\n 0.21928184512961552,\n -0.9753779994829639\n ],\n [\n 0.3838345920067314,\n -0.898891223911909,\n -0.2113423136836923\n ]\n ],\n \"t\": [\n [\n 498.7689000990772\n ],\n [\n 278.0695777621727\n ],\n [\n 5618.721192968872\n ]\n ]\n },\n \"60457274\": {\n \"R\": [\n [\n 0.9239917699501332,\n -0.37272063182115767,\n 0.08554846392108466\n ],\n [\n -0.01857104155727153,\n -0.2671779087245581,\n -0.9634682566151569\n ],\n [\n 0.38196115703026423,\n 0.8886480156419687,\n -0.2537919991167828\n ]\n ],\n \"t\": [\n [\n -55.1478742462578\n ],\n [\n 424.8747833741909\n ],\n [\n 4452.137526291175\n ]\n ]\n }\n },\n \"S3\": {\n \"54138969\": {\n \"R\": [\n [\n -0.909926063968229,\n 0.4142842734534348,\n 0.020077322541766036\n ],\n [\n 0.06112258570603725,\n 0.18181129378483157,\n -0.9814319553432596\n ],\n [\n -0.41024210855042,\n -0.891803338310328,\n -0.19075696094942407\n ]\n ],\n \"t\": [\n [\n -144.30406670344493\n ],\n [\n 546.2767112872957\n ],\n [\n 5569.530692348755\n ]\n ]\n },\n \"55011271\": {\n \"R\": [\n [\n 0.9248703521034336,\n 0.3800681977315835,\n 0.012767022876799783\n ],\n [\n 0.093795468138089,\n -0.1954524371286302,\n -0.9762175756342618\n ],\n [\n -0.3685339088290622,\n 0.9040721817938792,\n -0.21641683887726407\n ]\n ],\n \"t\": [\n [\n -38.93379836342622\n ],\n [\n 375.57502666735104\n ],\n [\n 5759.402838804998\n ]\n ]\n },\n \"58860488\": {\n \"R\": [\n [\n -0.9218827889823751,\n -0.38260686272952316,\n -0.061189149122614306\n ],\n [\n -0.02577019492115,\n 0.21811470471458455,\n -0.9755828374059251\n ],\n [\n 0.3866109419452632,\n -0.897796170731164,\n -0.2109360457310579\n ]\n ],\n \"t\": [\n [\n 596.8162203909545\n ],\n [\n 282.123966506171\n ],\n [\n 5575.726600786697\n ]\n ]\n },\n \"60457274\": {\n \"R\": [\n [\n 0.9244960445794738,\n -0.37161308683612865,\n 0.08491629554468147\n ],\n [\n -0.018795005038688972,\n -0.26693214570791374,\n -0.963532032354589\n ],\n [\n 0.3807280017840865,\n 0.88918555053481,\n -0.2537621827176058\n ]\n ],\n \"t\": [\n [\n -158.57266932864025\n ],\n [\n 433.1881250816\n ],\n [\n 4413.555688648984\n ]\n ]\n }\n },\n \"S4\": {\n \"54138969\": {\n \"R\": [\n [\n -0.906169211683753,\n 0.422346184383899,\n 0.021933087625945674\n ],\n [\n 0.06180306305120707,\n 0.18355044391174938,\n -0.9810655512947585\n ],\n [\n -0.4183751201899252,\n -0.8876558652294037,\n -0.19243004892662768\n ]\n ],\n \"t\": [\n [\n -201.25197932223173\n ],\n [\n 537.4605027947064\n ],\n [\n 5553.966756732112\n ]\n ]\n },\n \"55011271\": {\n \"R\": [\n [\n 0.9205073288493492,\n 0.39058428754662783,\n 0.010496278213208041\n ],\n [\n 0.0923916650578188,\n -0.1914846595009468,\n -0.9771373523735801\n ],\n [\n -0.3796446203523497,\n 0.900431862773358,\n -0.21234976510469855\n ]\n ],\n \"t\": [\n [\n 63.12322044876507\n ],\n [\n 396.6138950755392\n ],\n [\n 5760.7235858284985\n ]\n ]\n },\n \"58860488\": {\n \"R\": [\n [\n -0.9244800422436603,\n -0.37641653359695837,\n -0.060392422769829215\n ],\n [\n -0.02551533125481826,\n 0.2191523935220463,\n -0.9753569583924211\n ],\n [\n 0.3803756292983513,\n -0.9001571094250204,\n -0.21220640656557746\n ]\n ],\n \"t\": [\n [\n 559.4298619884164\n ],\n [\n 278.041710381495\n ],\n [\n 5601.2846874450925\n ]\n ]\n },\n \"60457274\": {\n \"R\": [\n [\n 0.9241606780958346,\n -0.3729066880542538,\n 0.0828712439019392\n ],\n [\n -0.021270464031387784,\n -0.2668345784720987,\n -0.9635075895349796\n ],\n [\n 0.38141133756265905,\n 0.8886731174824772,\n -0.2545299232755129\n ]\n ],\n \"t\": [\n [\n -98.61477305435534\n ],\n [\n 432.68486951797627\n ],\n [\n 4419.390974448715\n ]\n ]\n }\n },\n \"S5\": {\n \"54138969\": {\n \"R\": [\n [\n -0.9042074184788829,\n 0.42657831374650107,\n 0.020973473936051274\n ],\n [\n 0.06390493744399675,\n 0.18368565260974637,\n -0.9809055713959477\n ],\n [\n -0.4222855708380685,\n -0.8856017859436166,\n -0.1933503902128034\n ]\n ],\n \"t\": [\n [\n -219.3059666108619\n ],\n [\n 544.4787497640639\n ],\n [\n 5518.740477016156\n ]\n ]\n },\n \"55011271\": {\n \"R\": [\n [\n 0.9222116004775194,\n 0.38649075753002626,\n 0.012274293810989732\n ],\n [\n 0.09333184463870337,\n -0.19167233853095322,\n -0.9770111982052265\n ],\n [\n -0.3752531555110883,\n 0.902156643264318,\n -0.21283434941998647\n ]\n ],\n \"t\": [\n [\n 103.90282067751986\n ],\n [\n 395.67169468951965\n ],\n [\n 5767.97265758172\n ]\n ]\n },\n \"58860488\": {\n \"R\": [\n [\n -0.9258288614330635,\n -0.3728674116124112,\n -0.06173178026768599\n ],\n [\n -0.023578112500148365,\n 0.220000562347259,\n -0.9752147584905696\n ],\n [\n 0.3772068291381898,\n -0.9014264506460582,\n -0.21247437993123308\n ]\n ],\n \"t\": [\n [\n 520.3272318446208\n ],\n [\n 283.3690958234795\n ],\n [\n 5591.123958858676\n ]\n ]\n },\n \"60457274\": {\n \"R\": [\n [\n 0.9222815489764817,\n -0.3772688722588351,\n 0.0840532119677073\n ],\n [\n -0.021177649402562934,\n -0.26645871124348197,\n -0.9636136478735888\n ],\n [\n 0.3859381447632816,\n 0.88694303832152,\n -0.25373962085111357\n ]\n ],\n \"t\": [\n [\n -79.116431351199\n ],\n [\n 425.59047114848386\n ],\n [\n 4454.481629705836\n ]\n ]\n }\n },\n \"S6\": {\n \"54138969\": {\n \"R\": [\n [\n -0.9149503344107554,\n 0.4034864343564006,\n 0.008036345687245266\n ],\n [\n 0.07174776353922047,\n 0.1822275975157708,\n -0.9806351824867137\n ],\n [\n -0.3971374371533952,\n -0.896655898321083,\n -0.19567845056940925\n ]\n ],\n \"t\": [\n [\n -239.5182864132218\n ],\n [\n 545.8141831785044\n ],\n [\n 5523.931578633363\n ]\n ]\n },\n \"55011271\": {\n \"R\": [\n [\n 0.9197364689900042,\n 0.39209901596964664,\n 0.018525368698999664\n ],\n [\n 0.101478073351267,\n -0.19191459963948,\n -0.9761511087296542\n ],\n [\n -0.37919260045353465,\n 0.899681692667386,\n -0.21630030892357308\n ]\n ],\n \"t\": [\n [\n 169.02510061389722\n ],\n [\n 409.6671223380997\n ],\n [\n 5714.338002825065\n ]\n ]\n },\n \"58860488\": {\n \"R\": [\n [\n -0.916577698818659,\n -0.39393483656788014,\n -0.06856140726771254\n ],\n [\n -0.01984531630322392,\n 0.21607069980297702,\n -0.9761760169700323\n ],\n [\n 0.3993638509543854,\n -0.8933805444629346,\n -0.20586334624209834\n ]\n ],\n \"t\": [\n [\n 521.9864793089763\n ],\n [\n 286.28272817103516\n ],\n [\n 5643.2724406159\n ]\n ]\n },\n \"60457274\": {\n \"R\": [\n [\n 0.9182950552949388,\n -0.3850769011116475,\n 0.09192372735651859\n ],\n [\n -0.015534985886560007,\n -0.26706146429979655,\n -0.9635542737695438\n ],\n [\n 0.3955917790277871,\n 0.8833990913037544,\n -0.25122338635033875\n ]\n ],\n \"t\": [\n [\n -56.29675276801464\n ],\n [\n 420.29579722027506\n ],\n [\n 4499.322693551688\n ]\n ]\n }\n },\n \"S7\": {\n \"54138969\": {\n \"R\": [\n [\n -0.9055764231419416,\n 0.42392653746206904,\n 0.014752378956221508\n ],\n [\n 0.06862812683752326,\n 0.18074371881263407,\n -0.9811329615890764\n ],\n [\n -0.41859469903024304,\n -0.8874784498483331,\n -0.19277053457045695\n ]\n ],\n \"t\": [\n [\n -323.9118424584857\n ],\n [\n 541.7715234126381\n ],\n [\n 5506.569132699328\n ]\n ]\n },\n \"55011271\": {\n \"R\": [\n [\n 0.9212640765077017,\n 0.3886011826562522,\n 0.01617473877914905\n ],\n [\n 0.09922277503271489,\n -0.1946115441987536,\n -0.9758489574618522\n ],\n [\n -0.3760682680727248,\n 0.9006194910741931,\n -0.21784671226815075\n ]\n ],\n \"t\": [\n [\n 178.6238708832376\n ],\n [\n 403.59193467821774\n ],\n [\n 5694.8801003668095\n ]\n ]\n },\n \"58860488\": {\n \"R\": [\n [\n -0.9245069728829368,\n -0.37555597339631824,\n -0.06515034871105972\n ],\n [\n -0.018955014220249332,\n 0.21601110989507338,\n -0.9762068980691586\n ],\n [\n 0.38069353097569036,\n -0.9012751584550871,\n -0.20682244613440448\n ]\n ],\n \"t\": [\n [\n 441.1064712697594\n ],\n [\n 271.91614362573955\n ],\n [\n 5660.120611352617\n ]\n ]\n },\n \"60457274\": {\n \"R\": [\n [\n 0.9228353966173104,\n -0.3744001545228767,\n 0.09055029013436408\n ],\n [\n -0.014982084363704698,\n -0.269786590656035,\n -0.9628035794752281\n ],\n [\n 0.3849030629889691,\n 0.8871525910436372,\n -0.25457791009093983\n ]\n ],\n \"t\": [\n [\n 25.768533743836343\n ],\n [\n 431.05581759025813\n ],\n [\n 4461.872981411145\n ]\n ]\n }\n },\n \"S8\": {\n \"54138969\": {\n \"R\": [\n [\n -0.9115694669712032,\n 0.4106494283805017,\n 0.020202818036194434\n ],\n [\n 0.060907749548984036,\n 0.1834736632003901,\n -0.9811359034082424\n ],\n [\n -0.40660958293025334,\n -0.8931430243150293,\n -0.19226072190306673\n ]\n ],\n \"t\": [\n [\n -82.70216069652597\n ],\n [\n 552.1896311377282\n ],\n [\n 5557.353609418419\n ]\n ]\n },\n \"55011271\": {\n \"R\": [\n [\n 0.931016282525616,\n 0.3647626932499711,\n 0.01252434769597448\n ],\n [\n 0.08939715221301257,\n -0.19463753190599434,\n -0.9767929055586687\n ],\n [\n -0.35385990285476776,\n 0.9105297407479727,\n -0.2138194574051759\n ]\n ],\n \"t\": [\n [\n -209.06289992510443\n ],\n [\n 375.0691429434037\n ],\n [\n 5818.276676972416\n ]\n ]\n },\n \"58860488\": {\n \"R\": [\n [\n -0.9209075762929309,\n -0.3847355178017309,\n -0.0625125368875214\n ],\n [\n -0.02568138180824641,\n 0.21992027027623712,\n -0.9751797482259595\n ],\n [\n 0.38893405939143305,\n -0.8964450100611084,\n -0.21240678280563546\n ]\n ],\n \"t\": [\n [\n 623.0985110132146\n ],\n [\n 290.9053651845054\n ],\n [\n 5534.379001592981\n ]\n ]\n },\n \"60457274\": {\n \"R\": [\n [\n 0.927667052235436,\n -0.3636062759574404,\n 0.08499597802942535\n ],\n [\n -0.01666268768012713,\n -0.26770413351564454,\n -0.9633570738505596\n ],\n [\n 0.37303645269074087,\n 0.8922583555131325,\n -0.2543989622245125\n ]\n ],\n \"t\": [\n [\n -178.36705625795474\n ],\n [\n 423.4669232560848\n ],\n [\n 4421.6448791590965\n ]\n ]\n }\n },\n \"S9\": {\n \"54138969\": {\n \"R\": [\n [\n -0.9033486204435297,\n 0.4269119782787646,\n 0.04132109321984796\n ],\n [\n 0.04153061098352977,\n 0.182951140059007,\n -0.9822444139329296\n ],\n [\n -0.4268916470184284,\n -0.8855930460167476,\n -0.18299857527497945\n ]\n ],\n \"t\": [\n [\n -321.2078335720134\n ],\n [\n 467.13452033013084\n ],\n [\n 5514.330338522134\n ]\n ]\n },\n \"55011271\": {\n \"R\": [\n [\n 0.9315720471487059,\n 0.36348288012373176,\n -0.007329176497134756\n ],\n [\n 0.06810069482701912,\n -0.19426747906725159,\n -0.9785818524481906\n ],\n [\n -0.35712157080642226,\n 0.911120377575769,\n -0.20572758986325015\n ]\n ],\n \"t\": [\n [\n 19.193095609487138\n ],\n [\n 404.22842728571936\n ],\n [\n 5702.169280033924\n ]\n ]\n },\n \"58860488\": {\n \"R\": [\n [\n -0.9269344193869241,\n -0.3732303525241731,\n -0.03862235247246717\n ],\n [\n -0.04725991098820678,\n 0.218240494552814,\n -0.9747500127472326\n ],\n [\n 0.37223525218497616,\n -0.901704048173249,\n -0.21993345934341726\n ]\n ],\n \"t\": [\n [\n 455.40107288876885\n ],\n [\n 273.3589338272866\n ],\n [\n 5657.814488280711\n ]\n ]\n },\n \"60457274\": {\n \"R\": [\n [\n 0.915460708083783,\n -0.39734606500700814,\n 0.06362229623477154\n ],\n [\n -0.04940628468469528,\n -0.26789167566119776,\n -0.9621814117644814\n ],\n [\n 0.39936288133525055,\n 0.8776959352388969,\n -0.26487569589663096\n ]\n ],\n \"t\": [\n [\n -69.271255294384\n ],\n [\n 422.1843366088847\n ],\n [\n 4457.893374979773\n ]\n ]\n }\n },\n \"S10\": {\n \"54138969\": {\n \"R\": [\n [\n -0.9199955359932982,\n 0.39133749168985454,\n 0.021521648410310328\n ],\n [\n 0.0555185840851712,\n 0.18448351869097226,\n -0.9812662829999691\n ],\n [\n -0.3879766752957989,\n -0.9015657485337887,\n -0.19145051709809383\n ]\n ],\n \"t\": [\n [\n -181.4625993368258\n ],\n [\n 543.5199110634021\n ],\n [\n 5582.194377534298\n ]\n ]\n },\n \"55011271\": {\n \"R\": [\n [\n 0.9152587269115653,\n 0.40266346194010966,\n 0.01279059148853104\n ],\n [\n 0.09843295287698457,\n -0.1927270179143742,\n -0.9763028476624197\n ],\n [\n -0.3906563919867918,\n 0.8948287171209015,\n -0.21603043863220686\n ]\n ],\n \"t\": [\n [\n -22.5707386911355\n ],\n [\n 383.7773845053516\n ],\n [\n 5727.149101385447\n ]\n ]\n },\n \"58860488\": {\n \"R\": [\n [\n -0.9117691356172892,\n -0.4060874893594546,\n -0.06140027948988781\n ],\n [\n -0.03165845257462336,\n 0.21854812554171174,\n -0.9753124931029975\n ],\n [\n 0.4094811176553588,\n -0.887315990956964,\n -0.21212153703897946\n ]\n ],\n \"t\": [\n [\n 579.9870562891809\n ],\n [\n 276.09388439709664\n ],\n [\n 5616.656671116378\n ]\n ]\n },\n \"60457274\": {\n \"R\": [\n [\n 0.9374925472123639,\n -0.3377263929586908,\n 0.08395598501825681\n ],\n [\n -0.009787543064644189,\n -0.2667415901136707,\n -0.9637183863060765\n ],\n [\n 0.3478676873784507,\n 0.9026570819545717,\n -0.25337376437829157\n ]\n ],\n \"t\": [\n [\n -72.59483557976097\n ],\n [\n 445.63607020105314\n ],\n [\n 4402.73689876101\n ]\n ]\n }\n },\n \"S11\": {\n \"54138969\": {\n \"R\": [\n [\n -0.9059013006181885,\n 0.4217144115102914,\n 0.038727105014486805\n ],\n [\n 0.044493184429779696,\n 0.1857199061874203,\n -0.9815948619389944\n ],\n [\n -0.4211450938543295,\n -0.8875049698848251,\n -0.1870073216538954\n ]\n ],\n \"t\": [\n [\n -234.7208032216618\n ],\n [\n 464.34018262882194\n ],\n [\n 5536.652631113797\n ]\n ]\n },\n \"55011271\": {\n \"R\": [\n [\n 0.9216646531492915,\n 0.3879848687925067,\n -0.0014172943441045224\n ],\n [\n 0.07721054863099915,\n -0.18699239961454955,\n -0.979322405373477\n ],\n [\n -0.3802272982247548,\n 0.9024974149959955,\n -0.20230080971229314\n ]\n ],\n \"t\": [\n [\n -11.934348472090557\n ],\n [\n 449.4165893644565\n ],\n [\n 5541.113551868937\n ]\n ]\n },\n \"58860488\": {\n \"R\": [\n [\n -0.9063540572469627,\n -0.42053101768163204,\n -0.04093880896680188\n ],\n [\n -0.0603212197838846,\n 0.22468715090881142,\n -0.9725620980997899\n ],\n [\n 0.4181909532208387,\n -0.8790161246439863,\n -0.2290130547809762\n ]\n ],\n \"t\": [\n [\n 781.127357651581\n ],\n [\n 235.3131620173424\n ],\n [\n 5576.37044019807\n ]\n ]\n },\n \"60457274\": {\n \"R\": [\n [\n 0.91754082476548,\n -0.39226322025776267,\n 0.06517975852741943\n ],\n [\n -0.04531905395586976,\n -0.26600517028098103,\n -0.9629057236990188\n ],\n [\n 0.395050652748768,\n 0.8805514269006645,\n -0.2618476013752581\n ]\n ],\n \"t\": [\n [\n -155.13650339749012\n ],\n [\n 422.16256306729633\n ],\n [\n 4435.416222660868\n ]\n ]\n }\n }\n }\n }" }, { "path": "eval/GVHMR/hmr4d/dataset/h36m/h36m.py", "content": "import torch\nimport numpy as np\nfrom pathlib import Path\nfrom hmr4d.configs import MainStore, builds\n\nfrom hmr4d.utils.pylogger import Log\nfrom hmr4d.dataset.imgfeat_motion.base_dataset import ImgfeatMotionDatasetBase\nfrom pytorch3d.transforms import axis_angle_to_matrix, matrix_to_axis_angle\nfrom hmr4d.utils import matrix\nfrom hmr4d.utils.smplx_utils import make_smplx\nfrom tqdm import tqdm\n\nfrom hmr4d.utils.geo_transform import compute_cam_angvel, apply_T_on_points\nfrom hmr4d.utils.geo.hmr_global import get_tgtcoord_rootparam, get_T_w2c_from_wcparams, get_c_rootparam, get_R_c2gv\n\nfrom hmr4d.utils.wis3d_utils import make_wis3d, add_motion_as_lines\nfrom hmr4d.utils.vis.renderer import Renderer\nimport imageio\nfrom hmr4d.utils.video_io_utils import read_video_np\nfrom hmr4d.utils.net_utils import get_valid_mask, repeat_to_max_len, repeat_to_max_len_dict\n\n\nclass H36mSmplDataset(ImgfeatMotionDatasetBase):\n def __init__(\n self,\n root=\"inputs/H36M/hmr4d_support\",\n original_coord=\"az\",\n motion_frames=120, # H36M's videos are 25fps and very long\n lazy_load=False,\n ):\n # Path\n self.root = Path(root)\n\n # Coord\n self.original_coord = original_coord\n\n # Setting\n self.motion_frames = motion_frames\n self.lazy_load = lazy_load\n\n super().__init__()\n\n def _load_dataset(self):\n # smplpose\n tic = Log.time()\n fn = self.root / \"smplxpose_v1.pt\"\n self.smpl_model = make_smplx(\"supermotion\")\n Log.info(f\"[H36M] Loading from {fn} ...\")\n self.motion_files = torch.load(fn)\n # Dict of {\n # \"smpl_params_glob\": {'body_pose', 'global_orient', 'transl', 'betas'}, FxC\n # \"cam_Rt\": tensor(F, 3),\n # \"cam_K\": tensor(1, 10),\n # }\n self.seqs = list(self.motion_files.keys())\n Log.info(f\"[H36M] {len(self.seqs)} sequences. Elapsed: {Log.time() - tic:.2f}s\")\n\n # img(as feature)\n # vid -> (features, vid, meta {bbx_xys, K_fullimg})\n if not self.lazy_load:\n tic = Log.time()\n fn = self.root / \"vitfeat_h36m.pt\"\n Log.info(f\"[H36M] Fully Loading to RAM ViT-Feat: {fn}\")\n self.f_img_dicts = torch.load(fn)\n Log.info(f\"[H36M] Finished. Elapsed: {Log.time() - tic:.2f}s\")\n else:\n raise NotImplementedError # \"Check BEDLAM-SMPL for lazy_load\"\n\n def _get_idx2meta(self):\n # We expect to see the entire sequence during one epoch,\n # so each sequence will be sampled max(SeqLength // MotionFrames, 1) times\n seq_lengths = []\n self.idx2meta = []\n for vid in self.f_img_dicts:\n seq_length = self.f_img_dicts[vid][\"bbx_xys\"].shape[0]\n num_samples = max(seq_length // self.motion_frames, 1)\n seq_lengths.append(seq_length)\n self.idx2meta.extend([vid] * num_samples)\n hours = sum(seq_lengths) / 25 / 3600\n Log.info(f\"[H36M] has {hours:.1f} hours motion -> Resampled to {len(self.idx2meta)} samples.\")\n\n def _load_data(self, idx):\n sampled_motion = {}\n vid = self.idx2meta[idx]\n motion = self.motion_files[vid]\n seq_length = self.f_img_dicts[vid][\"bbx_xys\"].shape[0] # this is a better choice\n sampled_motion[\"vid\"] = vid\n\n # Random select a subset\n target_length = self.motion_frames\n if target_length > seq_length: # this should not happen\n start = 0\n length = seq_length\n Log.info(f\"[H36M] ({idx}) target length < sequence length: {target_length} <= {seq_length}\")\n else:\n start = np.random.randint(0, seq_length - target_length)\n length = target_length\n end = start + length\n sampled_motion[\"length\"] = length\n sampled_motion[\"start_end\"] = (start, end)\n\n # Select motion subset\n # body_pose, global_orient, transl, betas\n sampled_motion[\"smpl_params_global\"] = {k: v[start:end] for k, v in motion[\"smpl_params_glob\"].items()}\n\n # Image as feature\n f_img_dict = self.f_img_dicts[vid]\n sampled_motion[\"f_imgseq\"] = f_img_dict[\"features\"][start:end].float() # (L, 1024)\n sampled_motion[\"bbx_xys\"] = f_img_dict[\"bbx_xys\"][start:end]\n sampled_motion[\"K_fullimg\"] = f_img_dict[\"K_fullimg\"]\n # sampled_motion[\"kp2d\"] = self.vitpose[vid][start:end].float() # (L, 17, 3)\n sampled_motion[\"kp2d\"] = torch.zeros((end - start), 17, 3) # (L, 17, 3)\n\n # Camera\n sampled_motion[\"T_w2c\"] = motion[\"cam_Rt\"] # (4, 4)\n\n return sampled_motion\n\n def _process_data(self, data, idx):\n length = data[\"length\"]\n\n # SMPL params in world\n smpl_params_w = data[\"smpl_params_global\"].copy() # in az\n\n # SMPL params in cam\n T_w2c = data[\"T_w2c\"] # (4, 4)\n offset = self.smpl_model.get_skeleton(smpl_params_w[\"betas\"][0])[0] # (3)\n global_orient_c, transl_c = get_c_rootparam(\n smpl_params_w[\"global_orient\"],\n smpl_params_w[\"transl\"],\n T_w2c,\n offset,\n )\n smpl_params_c = {\n \"body_pose\": smpl_params_w[\"body_pose\"].clone(), # (F, 63)\n \"betas\": smpl_params_w[\"betas\"].clone(), # (F, 10)\n \"global_orient\": global_orient_c, # (F, 3)\n \"transl\": transl_c, # (F, 3)\n }\n\n # World params\n gravity_vec = torch.tensor([0, 0, -1]).float() # (3), H36M is az\n T_w2c = T_w2c.repeat(length, 1, 1) # (F, 4, 4)\n R_c2gv = get_R_c2gv(T_w2c[..., :3, :3], axis_gravity_in_w=gravity_vec) # (F, 3, 3)\n\n # Image\n bbx_xys = data[\"bbx_xys\"] # (F, 3)\n K_fullimg = data[\"K_fullimg\"].repeat(length, 1, 1) # (F, 3, 3)\n f_imgseq = data[\"f_imgseq\"] # (F, 1024)\n cam_angvel = compute_cam_angvel(T_w2c[:, :3, :3]) # (F, 6) slightly different from WHAM\n\n # Returns: do not forget to make it batchable! (last lines)\n max_len = self.motion_frames\n return_data = {\n \"meta\": {\"data_name\": \"h36m\", \"idx\": idx, \"vid\": data[\"vid\"]},\n \"length\": length,\n \"smpl_params_c\": smpl_params_c,\n \"smpl_params_w\": smpl_params_w,\n \"R_c2gv\": R_c2gv, # (F, 3, 3)\n \"gravity_vec\": gravity_vec, # (3)\n \"bbx_xys\": bbx_xys, # (F, 3)\n \"K_fullimg\": K_fullimg, # (F, 3, 3)\n \"f_imgseq\": f_imgseq, # (F, D)\n \"kp2d\": data[\"kp2d\"], # (F, 17, 3)\n \"cam_angvel\": cam_angvel, # (F, 6)\n \"mask\": {\n \"valid\": get_valid_mask(max_len, length),\n \"vitpose\": False,\n \"bbx_xys\": True,\n \"f_imgseq\": True,\n \"spv_incam_only\": False,\n },\n }\n\n if False: # Render to image to check\n smplx_out = self.smplx(**smpl_params_c)\n # ----- Overlay ----- #\n mid = return_data[\"meta\"][\"mid\"]\n video_path = self.root / f\"videos/{mid}.mp4\"\n images = read_video_np(video_path, data[\"start_end\"][0], data[\"start_end\"][1])\n render_dict = {\n \"K\": K_fullimg[:1], # only support batch size 1\n \"faces\": self.smplx.faces,\n \"verts\": smplx_out.vertices,\n \"background\": images,\n }\n img_overlay = simple_render_mesh_background(render_dict)\n save_video(img_overlay, f\"tmp.mp4\")\n\n # Batchable\n return_data[\"smpl_params_c\"] = repeat_to_max_len_dict(return_data[\"smpl_params_c\"], max_len)\n return_data[\"smpl_params_w\"] = repeat_to_max_len_dict(return_data[\"smpl_params_w\"], max_len)\n return_data[\"R_c2gv\"] = repeat_to_max_len(return_data[\"R_c2gv\"], max_len)\n return_data[\"bbx_xys\"] = repeat_to_max_len(return_data[\"bbx_xys\"], max_len)\n return_data[\"K_fullimg\"] = repeat_to_max_len(return_data[\"K_fullimg\"], max_len)\n return_data[\"f_imgseq\"] = repeat_to_max_len(return_data[\"f_imgseq\"], max_len)\n return_data[\"kp2d\"] = repeat_to_max_len(return_data[\"kp2d\"], max_len)\n return_data[\"cam_angvel\"] = repeat_to_max_len(return_data[\"cam_angvel\"], max_len)\n\n return return_data\n\n\ngroup_name = \"train_datasets/imgfeat_h36m\"\nnode_v1 = builds(H36mSmplDataset)\nMainStore.store(name=\"v1\", node=node_v1, group=group_name)\n" }, { "path": "eval/GVHMR/hmr4d/dataset/h36m/utils.py", "content": "import json\nimport numpy as np\nfrom pathlib import Path\nfrom collections import defaultdict\nimport pickle\nimport torch\n\nRESOURCE_FOLDER = Path(__file__).resolve().parent / \"resource\"\n\ncamera_idx_to_name = {0: \"54138969\", 1: \"55011271\", 2: \"58860488\", 3: \"60457274\"}\n\n\ndef get_vid(pkl_path, cam_id):\n \"\"\".../S6/Posing 1.pkl, 54138969 -> S6@Posing_1@54138969\"\"\"\n sub_id, fn = pkl_path.split(\"/\")[-2:]\n vid = f\"{sub_id}@{fn.split('.')[0].replace(' ', '_')}@{cam_id}\"\n return vid\n\n\ndef get_raw_pkl_paths(h36m_raw_root):\n smpl_param_dir = h36m_raw_root / \"neutrSMPL_H3.6\"\n pkl_paths = []\n for train_sub in [\"S1\", \"S5\", \"S6\", \"S7\", \"S8\"]:\n for pth in (smpl_param_dir / train_sub).glob(\"*.pkl\"):\n if \"aligned\" not in str(pth): # Use world sequence only\n pkl_paths.append(str(pth))\n\n return pkl_paths\n\n\ndef get_cam_KRts():\n \"\"\"\n Returns:\n Ks (torch.Tensor): {cam_id: 3x3}\n Rts (torch.Tensor): {subj_id: {cam_id: 4x4}}\n \"\"\"\n # this file is copied from https://github.com/karfly/human36m-camera-parameters\n cameras_path = RESOURCE_FOLDER / \"camera-parameters.json\"\n with open(cameras_path, \"r\") as f:\n cameras = json.load(f)\n\n # 4 camera ids: '54138969', '55011271', '58860488', '60457274'\n Ks = {}\n for cam in cameras[\"intrinsics\"]:\n Ks[cam] = torch.tensor(cameras[\"intrinsics\"][cam][\"calibration_matrix\"]).float()\n\n # extrinsics\n extrinsics = cameras[\"extrinsics\"]\n Rts = defaultdict(dict)\n for subj in extrinsics:\n for cam in extrinsics[subj]:\n Rt = torch.eye(4)\n Rt[:3, :3] = torch.tensor(extrinsics[subj][cam][\"R\"])\n Rt[:3, [3]] = torch.tensor(extrinsics[subj][cam][\"t\"]) / 1000\n Rts[subj][cam] = Rt.float()\n\n return Ks, Rts\n\n\ndef parse_raw_pkl(pkl_path, to_50hz=True):\n \"\"\"\n raw_pkl @ 200Hz, where video @ 50Hz.\n the frames should be divided by 4, and mannually align with the video.\n \"\"\"\n with open(str(pkl_path), \"rb\") as f:\n data = pickle.load(f, encoding=\"bytes\")\n poses = torch.from_numpy(data[b\"poses\"]).float()\n betas = torch.from_numpy(data[b\"betas\"]).float()\n trans = torch.from_numpy(data[b\"trans\"]).float()\n assert poses.shape[0] == trans.shape[0]\n if to_50hz:\n poses = poses[::4]\n trans = trans[::4]\n\n seq_length = poses.shape[0] # 50FPS\n smpl_params = {\n \"body_pose\": poses[:, 3:],\n \"betas\": betas[None].expand(seq_length, -1),\n \"global_orient\": poses[:, :3],\n \"transl\": trans,\n }\n return smpl_params\n" }, { "path": "eval/GVHMR/hmr4d/dataset/imgfeat_motion/base_dataset.py", "content": "import torch\nfrom torch.utils import data\nimport numpy as np\nfrom pathlib import Path\nfrom hmr4d.utils.pylogger import Log\n\n\nclass ImgfeatMotionDatasetBase(data.Dataset):\n def __init__(self):\n super().__init__()\n self._load_dataset()\n self._get_idx2meta() # -> Set self.idx2meta\n\n def __len__(self):\n return len(self.idx2meta)\n\n def _load_dataset(self):\n raise NotImplemented\n\n def _get_idx2meta(self):\n raise NotImplemented\n\n def _load_data(self, idx):\n raise NotImplemented\n\n def _process_data(self, data, idx):\n raise NotImplemented\n\n def __getitem__(self, idx):\n data = self._load_data(idx)\n data = self._process_data(data, idx)\n return data\n" }, { "path": "eval/GVHMR/hmr4d/dataset/pure_motion/amass.py", "content": "import torch\nimport torch.nn.functional as F\nimport numpy as np\n\nfrom tqdm import tqdm\nfrom pathlib import Path\nfrom hmr4d.utils.pylogger import Log\nfrom hmr4d.configs import MainStore, builds\n\nfrom .base_dataset import BaseDataset\nfrom .utils import *\nfrom hmr4d.utils.geo.hmr_global import get_tgtcoord_rootparam\nfrom hmr4d.utils.wis3d_utils import make_wis3d, add_motion_as_lines, convert_motion_as_line_mesh\n\n\nclass AmassDataset(BaseDataset):\n def __init__(\n self,\n motion_frames=120,\n l_factor=1.5, # speed augmentation\n skip_moyo=True, # not contained in the ICCV19 released version\n cam_augmentation=\"v11\",\n random1024=False, # DEBUG\n limit_size=None,\n ):\n self.root = Path(\"inputs/AMASS/hmr4d_support\")\n self.motion_frames = motion_frames\n self.l_factor = l_factor\n self.random1024 = random1024\n self.skip_moyo = skip_moyo\n self.dataset_name = \"AMASS\"\n super().__init__(cam_augmentation, limit_size)\n\n def _load_dataset(self):\n filename = self.root / \"smplxpose_v2.pth\"\n Log.info(f\"[{self.dataset_name}] Loading from {filename} ...\")\n tic = Log.time()\n if self.random1024: # Debug, faster loading\n try:\n Log.info(f\"[{self.dataset_name}] Loading 1024 samples for debugging ...\")\n self.motion_files = torch.load(self.root / \"smplxpose_v2_random1024.pth\")\n except:\n Log.info(f\"[{self.dataset_name}] Not found! Saving 1024 samples for debugging ...\")\n self.motion_files = torch.load(filename)\n keys = list(self.motion_files.keys())\n keys = np.random.choice(keys, 1024, replace=False)\n self.motion_files = {k: self.motion_files[k] for k in keys}\n torch.save(self.motion_files, self.root / \"smplxpose_v2_random1024.pth\")\n else:\n self.motion_files = torch.load(filename)\n self.seqs = list(self.motion_files.keys())\n Log.info(f\"[{self.dataset_name}] {len(self.seqs)} sequences. Elapsed: {Log.time() - tic:.2f}s\")\n\n def _get_idx2meta(self):\n # We expect to see the entire sequence during one epoch,\n # so each sequence will be sampled max(SeqLength // MotionFrames, 1) times\n seq_lengths = []\n self.idx2meta = []\n\n # Skip too-long idle-prefix\n motion_start_id = {}\n for vid in self.motion_files:\n if self.skip_moyo and \"moyo_smplxn\" in vid:\n continue\n seq_length = self.motion_files[vid][\"pose\"].shape[0]\n start_id = motion_start_id[vid] if vid in motion_start_id else 0\n seq_length = seq_length - start_id\n if seq_length < 25: # Skip clips that are too short\n continue\n num_samples = max(seq_length // self.motion_frames, 1)\n seq_lengths.append(seq_length)\n self.idx2meta.extend([(vid, start_id)] * num_samples)\n hours = sum(seq_lengths) / 30 / 3600\n Log.info(f\"[{self.dataset_name}] has {hours:.1f} hours motion -> Resampled to {len(self.idx2meta)} samples.\")\n\n def _load_data(self, idx):\n \"\"\"\n - Load original data\n - Augmentation: speed-augmentation to L frames\n \"\"\"\n # Load original data\n mid, start_id = self.idx2meta[idx]\n raw_data = self.motion_files[mid]\n raw_len = raw_data[\"pose\"].shape[0] - start_id\n data = {\n \"body_pose\": raw_data[\"pose\"][start_id:, 3:], # (F, 63)\n \"betas\": raw_data[\"beta\"].repeat(raw_len, 1), # (10)\n \"global_orient\": raw_data[\"pose\"][start_id:, :3], # (F, 3)\n \"transl\": raw_data[\"trans\"][start_id:], # (F, 3)\n }\n\n # Get {tgt_len} frames from data\n # Random select a subset with speed augmentation [start, end)\n tgt_len = self.motion_frames\n raw_subset_len = np.random.randint(int(tgt_len / self.l_factor), int(tgt_len * self.l_factor))\n if raw_subset_len <= raw_len:\n start = np.random.randint(0, raw_len - raw_subset_len + 1)\n end = start + raw_subset_len\n else: # interpolation will use all possible frames (results in a slow motion)\n start = 0\n end = raw_len\n data = {k: v[start:end] for k, v in data.items()}\n\n # Interpolation (vec + r6d)\n data_interpolated = interpolate_smpl_params(data, tgt_len)\n\n # AZ -> AY\n data_interpolated[\"global_orient\"], data_interpolated[\"transl\"], _ = get_tgtcoord_rootparam(\n data_interpolated[\"global_orient\"],\n data_interpolated[\"transl\"],\n tsf=\"az->ay\",\n )\n\n data_interpolated[\"data_name\"] = \"amass\"\n return data_interpolated\n\n\ngroup_name = \"train_datasets/pure_motion_amass\"\nMainStore.store(name=\"v11\", node=builds(AmassDataset, cam_augmentation=\"v11\"), group=group_name)\n" }, { "path": "eval/GVHMR/hmr4d/dataset/pure_motion/base_dataset.py", "content": "import torch\nfrom torch.utils.data import Dataset\nfrom pathlib import Path\n\nfrom .utils import *\nfrom .cam_traj_utils import CameraAugmentorV11\nfrom hmr4d.utils.geo.hmr_cam import create_camera_sensor\nfrom hmr4d.utils.geo.hmr_global import get_c_rootparam, get_R_c2gv\nfrom hmr4d.utils.net_utils import get_valid_mask, repeat_to_max_len, repeat_to_max_len_dict\nfrom hmr4d.utils.geo_transform import compute_cam_angvel, apply_T_on_points, project_p2d, cvt_p2d_from_i_to_c\n\nfrom hmr4d.utils.wis3d_utils import make_wis3d, add_motion_as_lines, convert_motion_as_line_mesh\nfrom hmr4d.utils.smplx_utils import make_smplx\n\n\nclass BaseDataset(Dataset):\n def __init__(self, cam_augmentation, limit_size=None):\n super().__init__()\n self.cam_augmentation = cam_augmentation\n self.limit_size = limit_size\n self.smplx = make_smplx(\"supermotion\")\n self.smplx_lite = make_smplx(\"supermotion_smpl24\")\n\n self._load_dataset()\n self._get_idx2meta()\n\n def _load_dataset(self):\n NotImplementedError(\"_load_dataset is not implemented\")\n\n def _get_idx2meta(self):\n self.idx2meta = None\n NotImplementedError(\"_get_idx2meta is not implemented\")\n\n def __len__(self):\n if self.limit_size is not None:\n return min(self.limit_size, len(self.idx2meta))\n return len(self.idx2meta)\n\n def _load_data(self, idx):\n NotImplementedError(\"_load_data is not implemented\")\n\n def _process_data(self, data, idx):\n \"\"\"\n Args:\n data: dict {\n \"body_pose\": (F, 63),\n \"betas\": (F, 10),\n \"global_orient\": (F, 3), in the AY coordinates\n \"transl\": (F, 3), in the AY coordinates\n }\n \"\"\"\n data_name = data[\"data_name\"]\n length = data[\"body_pose\"].shape[0]\n # Augmentation: betas, SMPL (gravity-axis)\n body_pose = data[\"body_pose\"]\n betas = augment_betas(data[\"betas\"], std=0.1)\n global_orient_w, transl_w = rotate_around_axis(data[\"global_orient\"], data[\"transl\"], axis=\"y\")\n del data\n\n # SMPL_params in world\n smpl_params_w = {\n \"body_pose\": body_pose, # (F, 63)\n \"betas\": betas, # (F, 10)\n \"global_orient\": global_orient_w, # (F, 3)\n \"transl\": transl_w, # (F, 3)\n }\n\n # Camera trajectory augmentation\n if self.cam_augmentation == \"v11\":\n # interleave repeat to original length (faster)\n N = 10\n w_j3d = self.smplx_lite(\n smpl_params_w[\"body_pose\"][::N],\n smpl_params_w[\"betas\"][::N],\n smpl_params_w[\"global_orient\"][::N],\n None,\n )\n w_j3d = w_j3d.repeat_interleave(N, dim=0) + smpl_params_w[\"transl\"][:, None] # (F, 24, 3)\n\n if False:\n wis3d = make_wis3d(name=\"debug_amass\")\n add_motion_as_lines(w_j3d, wis3d, \"w_j3d\")\n\n width, height, K_fullimg = create_camera_sensor(1000, 1000, 43.3) # WHAM\n focal_length = K_fullimg[0, 0]\n wham_cam_augmentor = CameraAugmentorV11()\n T_w2c = wham_cam_augmentor(w_j3d, length) # (F, 4, 4)\n\n else:\n raise NotImplementedError\n\n if False: # render\n for idx_render in range(10):\n T_w2c = wham_cam_augmentor(smpl_params_w[\"transl\"])\n\n # targets\n w_j3d = self.smplx(**smpl_params_w).joints[:, :22]\n c_j3d = apply_T_on_points(w_j3d, T_w2c)\n verts, faces, vertex_colors = convert_motion_as_line_mesh(c_j3d)\n vertex_colors = vertex_colors[None] / 255.0\n bg = np.ones((height, width, 3), dtype=np.uint8) * 255\n\n # render\n renderer = Renderer(width, height, device=\"cuda\", faces=faces, K=K_fullimg)\n vname = f\"{idx_render:02d}\"\n out_fn = Path(f\"outputs/dump_render_wham_cam/{vname}.mp4\")\n out_fn.parent.mkdir(exist_ok=True, parents=True)\n writer = imageio.get_writer(out_fn, fps=30, mode=\"I\", format=\"FFMPEG\", macro_block_size=1)\n for i in tqdm(range(len(verts)), desc=f\"Rendering {vname}\"):\n # incam\n # img_overlay_pred = renderer.render_mesh(verts[i].cuda(), bg, [0.8, 0.8, 0.8], VI=1)\n img_overlay_pred = renderer.render_mesh(verts[i].cuda(), bg, vertex_colors, VI=1)\n # if batch[\"meta_render\"][0].get(\"bbx_xys\", None) is not None: # draw bbox lines\n # bbx_xys = batch[\"meta_render\"][0][\"bbx_xys\"][i].cpu().numpy()\n # lu_point = (bbx_xys[:2] - bbx_xys[2:] / 2).astype(int)\n # rd_point = (bbx_xys[:2] + bbx_xys[2:] / 2).astype(int)\n # img_overlay_pred = cv2.rectangle(img_overlay_pred, lu_point, rd_point, (255, 178, 102), 2)\n\n # write\n writer.append_data(img_overlay_pred)\n writer.close()\n pass\n\n # SMPL params in cam\n offset = self.smplx.get_skeleton(smpl_params_w[\"betas\"][0])[0] # (3)\n global_orient_c, transl_c = get_c_rootparam(\n smpl_params_w[\"global_orient\"],\n smpl_params_w[\"transl\"],\n T_w2c,\n offset,\n )\n smpl_params_c = {\n \"body_pose\": smpl_params_w[\"body_pose\"].clone(), # (F, 63)\n \"betas\": smpl_params_w[\"betas\"].clone(), # (F, 10)\n \"global_orient\": global_orient_c, # (F, 3)\n \"transl\": transl_c, # (F, 3)\n }\n\n # World params\n gravity_vec = torch.tensor([0, -1, 0], dtype=torch.float32) # (3), BEDLAM is ay\n R_c2gv = get_R_c2gv(T_w2c[:, :3, :3], gravity_vec) # (F, 3, 3)\n\n # Image\n K_fullimg = K_fullimg.repeat(length, 1, 1) # (F, 3, 3)\n cam_angvel = compute_cam_angvel(T_w2c[:, :3, :3]) # (F, 6)\n\n # Returns: do not forget to make it batchable! (last lines)\n # NOTE: bbx_xys and f_imgseq will be added later\n max_len = length\n return_data = {\n \"meta\": {\"data_name\": data_name, \"idx\": idx, \"T_w2c\": T_w2c},\n \"length\": length,\n \"smpl_params_c\": smpl_params_c,\n \"smpl_params_w\": smpl_params_w,\n \"R_c2gv\": R_c2gv, # (F, 3, 3)\n \"gravity_vec\": gravity_vec, # (3)\n \"bbx_xys\": torch.zeros((length, 3)), # (F, 3) # NOTE: a placeholder\n \"K_fullimg\": K_fullimg, # (F, 3, 3)\n \"f_imgseq\": torch.zeros((length, 1024)), # (F, D) # NOTE: a placeholder\n \"kp2d\": torch.zeros(length, 17, 3), # (F, 17, 3)\n \"cam_angvel\": cam_angvel, # (F, 6)\n \"mask\": {\n \"valid\": get_valid_mask(length, length),\n \"vitpose\": False,\n \"bbx_xys\": False,\n \"f_imgseq\": False,\n \"spv_incam_only\": False,\n },\n }\n\n # Batchable\n return_data[\"smpl_params_c\"] = repeat_to_max_len_dict(return_data[\"smpl_params_c\"], max_len)\n return_data[\"smpl_params_w\"] = repeat_to_max_len_dict(return_data[\"smpl_params_w\"], max_len)\n return_data[\"R_c2gv\"] = repeat_to_max_len(return_data[\"R_c2gv\"], max_len)\n return_data[\"K_fullimg\"] = repeat_to_max_len(return_data[\"K_fullimg\"], max_len)\n return_data[\"cam_angvel\"] = repeat_to_max_len(return_data[\"cam_angvel\"], max_len)\n return return_data\n\n def __getitem__(self, idx):\n data = self._load_data(idx)\n data = self._process_data(data, idx)\n return data\n" }, { "path": "eval/GVHMR/hmr4d/dataset/pure_motion/cam_traj_utils.py", "content": "import torch\nimport torch.nn.functional as F\nimport numpy as np\nfrom numpy.random import rand, randn\nfrom pytorch3d.transforms import (\n axis_angle_to_matrix,\n matrix_to_axis_angle,\n matrix_to_rotation_6d,\n rotation_6d_to_matrix,\n)\nfrom einops import rearrange\nfrom hmr4d.utils.geo.hmr_cam import create_camera_sensor\nfrom hmr4d.utils.geo_transform import transform_mat, apply_T_on_points\nfrom hmr4d.utils.geo.transforms import axis_rotate_to_matrix\nimport hmr4d.utils.matrix as matrix\n\nhalfpi = np.pi / 2\nR_y_upsidedown = torch.tensor([[-1, 0, 0], [0, -1, 0], [0, 0, 1]]).float()\n\n\ndef noisy_interpolation(x, length, step_noise_perc=0.2):\n \"\"\"Non-linear interpolation with noise, although with noise, the jittery is very small\n Args:\n x: (2, C)\n length: scalar\n step_noise_perc: [x0, x1 +-(step_noise_perc * step), x2], where step = x1-x0\n \"\"\"\n assert x.shape[0] == 2 and len(x.shape) == 2\n dim = x.shape[-1]\n output = np.zeros((length, dim))\n\n # Use linsapce(0, 1) +- noise as reference\n linspace = np.repeat(np.linspace(0, 1, length)[None], dim, axis=0) # (D, L)\n noise = (linspace[0, 1] - linspace[0, 0]) * step_noise_perc\n space_noise = np.random.uniform(-noise, noise, (dim, length - 2)) # (D, L-2)\n linspace[:, 1:-1] = linspace[:, 1:-1] + space_noise\n\n # Do 1d interp\n for i in range(dim):\n output[:, i] = np.interp(linspace[i], np.array([0.0, 1.0]), x[:, i])\n return output\n\n\ndef noisy_impluse_interpolation(data1, data2, step_noise_perc=0.2):\n \"\"\"Non-linear interpolation of impluse with noise\"\"\"\n\n dim = data1.shape[-1]\n L = data1.shape[0]\n\n linspace1 = np.stack([np.linspace(0, 1, L // 2) for _ in range(dim)])\n linspace2 = np.stack([np.linspace(0, 1, L // 2)[::-1] for _ in range(dim)])\n linspace = np.concatenate([linspace1, linspace2], axis=-1)\n noise = (linspace[0, 1] - linspace[0, 0]) * step_noise_perc\n space_noise = np.stack([np.random.uniform(-noise, noise, L - 2) for _ in range(dim)])\n\n linspace[:, 1:-1] = linspace[:, 1:-1] + space_noise\n linspace = linspace.T\n output = data1 * (1 - linspace) + data2 * linspace\n return output\n\n\ndef create_camera(w_root, cfg):\n \"\"\"Create static camera pose\n Args:\n w_root: (3,), y-up coordinates\n Returns:\n R_w2c: (3, 3)\n t_w2c: (3)\n \"\"\"\n # Parse\n pitch_std = cfg[\"pitch_std\"]\n pitch_mean = cfg[\"pitch_mean\"]\n roll_std = cfg[\"roll_std\"]\n tz_range1_prob = cfg[\"tz_range1_prob\"]\n tz_range1 = cfg[\"tz_range1\"]\n tz_range2 = cfg[\"tz_range2\"]\n f = cfg[\"f\"]\n w = cfg[\"w\"]\n\n # algo\n yaw = rand() * 2 * np.pi # Look at any direction in xz-plane\n pitch = np.clip(randn() * pitch_std + pitch_mean, -halfpi, halfpi)\n roll = np.clip(randn() * roll_std, -halfpi, halfpi) # Normal-dist\n\n # Note we use OpenCV's camera system by first applying R_y_upsidedown\n yaw_rm = axis_rotate_to_matrix(yaw, axis=\"y\")\n pitch_rm = axis_rotate_to_matrix(pitch, axis=\"x\")\n roll_rm = axis_rotate_to_matrix(roll, axis=\"z\")\n R_w2c = (roll_rm @ pitch_rm @ yaw_rm @ R_y_upsidedown).squeeze(0) # (3, 3)\n\n # Place people in the scene\n if rand() < tz_range1_prob:\n tz = rand() * (tz_range1[1] - tz_range1[0]) + tz_range1[0]\n max_dist_in_fov = (w / 2) / f * tz\n tx = (rand() * 2 - 1) * 0.7 * max_dist_in_fov\n ty = (rand() * 2 - 1) * 0.5 * max_dist_in_fov\n\n else:\n tz = rand() * (tz_range2[1] - tz_range2[0]) + tz_range2[0]\n max_dist_in_fov = (w / 2) / f * tz\n max_dist_in_fov *= 0.9 # add a threshold\n tx = torch.randn(1) * 1.6\n tx = torch.clamp(tx, -max_dist_in_fov, max_dist_in_fov)\n ty = torch.randn(1) * 0.8\n ty = torch.clamp(ty, -max_dist_in_fov, max_dist_in_fov)\n\n dist = torch.tensor([tx, ty, tz], dtype=torch.float)\n t_w2c = dist - torch.matmul(R_w2c, w_root)\n\n return R_w2c, t_w2c\n\n\ndef create_rotation_move(R, length, r_xyz_w_std=[np.pi / 8, np.pi / 4, np.pi / 8]):\n \"\"\"Create rotational move for the camera\n Args:\n R: (3, 3)\n Return:\n R_move: (L, 3, 3)\n \"\"\"\n # Create final camera pose\n assert len(R.size()) == 2\n r_xyz = (2 * rand(3) - 1) * r_xyz_w_std\n Rf = R @ axis_angle_to_matrix(torch.from_numpy(r_xyz).float())\n\n # Inbetweening two poses\n Rs = torch.stack((R, Rf)) # (2, 3, 3)\n rs = matrix_to_rotation_6d(Rs).numpy() # (2, 6)\n rs_move = noisy_interpolation(rs, length) # (L, 6)\n R_move = rotation_6d_to_matrix(torch.from_numpy(rs_move).float())\n\n return R_move\n\n\ndef create_translation_move(R_w2c, t_w2c, length, t_xyz_w_std=[1.0, 0.25, 1.0]):\n \"\"\"Create translational move for the camera\n Args:\n R_w2c: (3, 3),\n t_w2c: (3,),\n \"\"\"\n # Create subject final displacement\n subj_start_final = np.array([[0, 0, 0], randn(3) * t_xyz_w_std])\n subj_move = noisy_interpolation(subj_start_final, length)\n subj_move = torch.from_numpy(subj_move).float() # (L, 3)\n\n # Equal to camera move\n t_move = t_w2c + torch.einsum(\"ij,lj->li\", R_w2c, subj_move)\n\n return t_move\n\n\nclass CameraAugmentorV11:\n cfg_create_camera = {\n \"pitch_mean\": np.pi / 36,\n \"pitch_std\": np.pi / 8,\n \"roll_std\": np.pi / 24,\n \"tz_range1_prob\": 0.4,\n \"tz_range1\": [1.0, 6.0], # uniform sample\n \"tz_range2\": [4.0, 12.0],\n \"tx_scale\": 0.7,\n \"ty_scale\": 0.3,\n }\n\n # r_xyz_w_std = [np.pi / 8, np.pi / 4, np.pi / 8] # in world coords\n r_xyz_w_std = [np.pi / 6, np.pi / 3, np.pi / 6] # in world coords\n t_xyz_w_std = [1.0, 0.25, 1.0] # in world coords\n r_xyz_w_std_half = [x / 2 for x in r_xyz_w_std]\n t_xyz_w_std_half = [x / 2 for x in t_xyz_w_std]\n\n t_factor = 1.0\n tz_bias_factor = 1.0\n\n rotx_impluse_noise = np.pi / 36\n roty_impluse_noise = np.pi / 36\n rotz_impluse_noise = np.pi / 36\n rot_impluse_n = 1\n\n tx_step_noise = 0.0025\n ty_step_noise = 0.0025\n tz_step_noise = 0.0025\n\n tx_impluse_noise = 0.15\n ty_impluse_noise = 0.15\n tz_impluse_noise = 0.15\n t_impluse_n = 1\n\n # === Postprocess === #\n height_max = 4.0\n height_min = -2.0 # -1.5 -> -2.0 allow look upside\n tz_post_min = 0.5\n\n def __init__(self):\n self.w = 1000\n self.f = create_camera_sensor(1000, 1000, 24)[2][0, 0] # use 24mm camera\n self.half_fov_tol = (self.w / 2) / self.f\n\n def create_rotation_track(self, cam_mat, root, rx_factor=1.0, ry_factor=1.0, rz_factor=1.0):\n \"\"\"Create rotational move for the camera with rotating human\"\"\"\n human_mat = matrix.get_TRS(matrix.identity_mat()[None, :3, :3], root)\n cam2human_mat = matrix.get_mat_BtoA(human_mat, cam_mat)\n R = matrix.get_rotation(cam2human_mat)\n\n # Create final camera pose\n yaw = np.random.normal(scale=ry_factor)\n pitch = np.random.normal(scale=rx_factor)\n roll = np.random.normal(scale=rz_factor)\n\n yaw_rm = axis_angle_to_matrix(torch.tensor([0, yaw, 0]).float())\n pitch_rm = axis_angle_to_matrix(torch.tensor([pitch, 0, 0]).float())\n roll_rm = axis_angle_to_matrix(torch.tensor([0, 0, roll]).float())\n Rf = roll_rm @ pitch_rm @ yaw_rm @ R[0]\n\n # Inbetweening two poses\n Rs = torch.stack((R[0], Rf))\n rs = matrix_to_rotation_6d(Rs).numpy()\n rs_move = noisy_interpolation(rs, self.l)\n R_move = rotation_6d_to_matrix(torch.from_numpy(rs_move).float())\n R_move = torch.inverse(R_move)\n return R_move\n\n def create_translation_track(self, cam_mat, root, t_factor=1.0, tz_bias_factor=0.0):\n \"\"\"Create translational move for the camera with tracking human\"\"\"\n delta_T0 = matrix.get_position(cam_mat)[0] - root[0]\n T_new = matrix.get_position(cam_mat)\n\n tz_bias = delta_T0.norm(dim=-1) * tz_bias_factor * np.clip(1 + np.random.normal(scale=0.1), 0.67, 1.5)\n\n T_new[1:] = root[1:] + delta_T0\n cam_mat = matrix.get_TRS(matrix.get_rotation(cam_mat), T_new)\n w2c = torch.inverse(cam_mat)\n T_new = matrix.get_position(w2c)\n\n # Create final camera position\n tx = np.random.normal(scale=t_factor)\n ty = np.random.normal(scale=t_factor)\n tz = np.random.normal(scale=t_factor) + tz_bias\n Ts = np.array([[0, 0, 0], [tx, ty, tz]])\n\n T_move = noisy_interpolation(Ts, self.l)\n T_move = torch.from_numpy(T_move).float()\n return T_move + T_new\n\n def add_stepnoise(self, R, T):\n w2c = matrix.get_TRS(R, T)\n cam_mat = torch.inverse(w2c)\n R_new = matrix.get_rotation(cam_mat)\n T_new = matrix.get_position(cam_mat)\n\n L = R_new.shape[0]\n window = 10\n\n def add_impulse_rot(R_new):\n N = np.random.randint(1, self.rot_impluse_n + 1)\n rx = np.random.normal(scale=self.rotx_impluse_noise, size=N)\n ry = np.random.normal(scale=self.roty_impluse_noise, size=N)\n rz = np.random.normal(scale=self.rotz_impluse_noise, size=N)\n R_impluse_noise = axis_angle_to_matrix(torch.from_numpy(np.array([rx, ry, rz])).float().transpose(0, 1))\n R_noise = R_new.clone()\n last_i = 0\n for i in range(N):\n n_i = np.random.randint(last_i + window, L - (N - i) * window * 2)\n\n # make impluse smooth\n window_R = R_noise[n_i - window : n_i + window].clone()\n window_r = matrix_to_rotation_6d(window_R).numpy()\n impluse_R = R_impluse_noise[i] @ window_R[window]\n window_impluse_R = window_R.clone()\n window_impluse_R[:] = impluse_R[None]\n window_impluse_r = matrix_to_rotation_6d(window_impluse_R).numpy()\n\n window_new_r = noisy_impluse_interpolation(window_r, window_impluse_r)\n window_new_R = rotation_6d_to_matrix(torch.from_numpy(window_new_r).float())\n R_noise[n_i - window : n_i + window] = window_new_R\n last_i = n_i\n R_new = R_noise\n return R_new\n\n def add_impulse_t(T_new):\n N = np.random.randint(1, self.t_impluse_n + 1)\n tx = np.random.normal(scale=self.tx_impluse_noise, size=N)\n ty = np.random.normal(scale=self.ty_impluse_noise, size=N)\n tz = np.random.normal(scale=self.tz_impluse_noise, size=N)\n T_impluse_noise = torch.from_numpy(np.array([tx, ty, tz])).float().transpose(0, 1)\n T_noise = T_new.clone()\n last_i = 0\n for i in range(N):\n n_i = np.random.randint(last_i + window, L - N * window * 2)\n\n # make impluse smooth\n window_T = T_noise[n_i - window : n_i + window].clone()\n window_impluse_T = window_T.clone()\n window_impluse_T += T_impluse_noise[i : i + 1]\n window_impluse_T = window_impluse_T.numpy()\n window_T = window_T.numpy()\n\n window_new_T = noisy_impluse_interpolation(window_T, window_impluse_T)\n window_new_T = torch.from_numpy(window_new_T).float()\n T_noise[n_i - window : n_i + window] = window_new_T\n last_i = n_i\n T_new = T_noise\n return T_new\n\n impulse_type_prob = {\n \"t\": 0.2,\n \"r\": 0.2,\n \"both\": 0.1,\n \"pass\": 0.5,\n }\n impulse_type = np.random.choice(list(impulse_type_prob.keys()), p=list(impulse_type_prob.values()))\n if impulse_type == \"t\":\n # impluse translation only\n T_new = add_impulse_t(T_new)\n elif impulse_type == \"r\":\n # impluse rotation only\n R_new = add_impulse_rot(R_new)\n elif impulse_type == \"both\":\n # impluse rotation and translation\n R_new = add_impulse_rot(R_new)\n T_new = add_impulse_t(T_new)\n else:\n assert impulse_type == \"pass\"\n\n cam_mat_new = matrix.get_TRS(R_new, T_new)\n w2c_new = torch.inverse(cam_mat_new)\n R_new = matrix.get_rotation(w2c_new)\n T_new = matrix.get_position(w2c_new)\n tx = np.random.normal(scale=self.tx_step_noise, size=L)\n ty = np.random.normal(scale=self.ty_step_noise, size=L)\n tz = np.random.normal(scale=self.tz_step_noise, size=L)\n T_new = T_new + torch.from_numpy(np.array([tx, ty, tz])).float().transpose(0, 1)\n\n return R_new, T_new\n\n def __call__(self, w_j3d, length=120):\n \"\"\"\n Args:\n w_j3d: (L, J, 3)\n length: scalar\n \"\"\"\n # Check\n self.l = length\n assert w_j3d.size(0) == self.l, \"currently, only support fixed length\"\n\n # Setup\n w_j3d = w_j3d.clone()\n w_root = w_j3d[:, 0] # (L, 3)\n\n # Simulate a static camera pose\n cfg_camera0 = {**self.cfg_create_camera, \"w\": self.w, \"f\": self.f}\n R0_w2c, t0_w2c = create_camera(w_root[0], cfg_camera0) # (3, 3) and (3,)\n\n # Move camera\n camera_type_prob = {\n \"random\": 0.25,\n \"track\": 0.15,\n \"trackrotate\": 0.10,\n \"trackpush\": 0.05,\n \"trackpull\": 0.05,\n \"static\": 0.4,\n }\n camera_type = np.random.choice(list(camera_type_prob.keys()), p=list(camera_type_prob.values()))\n if camera_type == \"random\": # random move + add noise on cam\n R_w2c = create_rotation_move(R0_w2c, length, self.r_xyz_w_std)\n t_w2c = create_translation_move(R0_w2c, t0_w2c, length, self.t_xyz_w_std)\n R_w2c, t_w2c = self.add_stepnoise(R_w2c, t_w2c)\n\n elif camera_type == \"track\": # track human\n R_w2c = create_rotation_move(R0_w2c, length, self.r_xyz_w_std_half)\n cam_mat = torch.inverse(transform_mat(R0_w2c, t0_w2c)).repeat(length, 1, 1) # (F, 4, 4)\n t_w2c = self.create_translation_track(cam_mat, w_root, 0.5)\n R_w2c, t_w2c = self.add_stepnoise(R_w2c, t_w2c)\n\n elif camera_type == \"trackrotate\": # track human and rotate\n cam_mat = torch.inverse(transform_mat(R0_w2c, t0_w2c)).repeat(length, 1, 1) # (F, 4, 4)\n t_w2c = self.create_translation_track(cam_mat, w_root, 0.5)\n cam_mat = matrix.get_TRS(matrix.get_rotation(cam_mat), t_w2c)\n R_w2c = self.create_rotation_track(cam_mat, w_root, np.pi / 16, np.pi, np.pi / 16)\n R_w2c, t_w2c = self.add_stepnoise(R_w2c, t_w2c)\n\n elif camera_type == \"trackpush\": # track human and push close to human\n R_w2c = create_rotation_move(R0_w2c, length, self.r_xyz_w_std_half)\n # [1/tz_bias_factor, 1] * dist\n cam_mat = torch.inverse(transform_mat(R0_w2c, t0_w2c)).repeat(length, 1, 1) # (F, 4, 4)\n t_w2c = self.create_translation_track(cam_mat, w_root, 0.5, (1.0 / (1 + self.tz_bias_factor) - 1))\n R_w2c, t_w2c = self.add_stepnoise(R_w2c, t_w2c)\n\n elif camera_type == \"trackpull\": # track human and pull far from human\n R_w2c = create_rotation_move(R0_w2c, length, self.r_xyz_w_std_half)\n # [1, (tz_bias_factor + 1)] * dist\n cam_mat = torch.inverse(transform_mat(R0_w2c, t0_w2c)).repeat(length, 1, 1) # (F, 4, 4)\n t_w2c = self.create_translation_track(cam_mat, w_root, 0.5, self.tz_bias_factor)\n R_w2c, t_w2c = self.add_stepnoise(R_w2c, t_w2c)\n\n else:\n assert camera_type == \"static\"\n R_w2c = R0_w2c.repeat(length, 1, 1) # (F, 3, 3)\n t_w2c = t0_w2c.repeat(length, 1) # (F, 3)\n\n # Recompute t_w2c for better camera height\n # cam_w = torch.einsum(\"lji,lj->li\", R_w2c, -t_w2c) # (L, 3), camera center in world: cam_w = - R_w2c^t_w2c @ t\n # height = cam_w[..., 1] - w_root[:, 1]\n # height = torch.clamp(height, self.height_min, self.height_max)\n # new_pos = cam_w.clone()\n # new_pos[:, 1] = w_root[:, 1] + height\n # t_w2c = torch.einsum(\"lij,lj->li\", R_w2c, -new_pos) # (L, 3), new t = -R_w2c @ cam_w\n\n # Recompute t_w2c for better depth and FoV\n c_j3d = torch.einsum(\"lij,lkj->lki\", R_w2c, w_j3d) + t_w2c[:, None] # (L, J, 3)\n delta = torch.zeros_like(t_w2c) # (L, 3) this will be later added to t_w2c\n # - If the person is too close to the camera, push away the person in the z direction\n c_j3d_min = c_j3d[..., 2].min() # scalar\n if c_j3d_min < self.tz_post_min:\n push_away = self.tz_post_min - c_j3d_min\n delta[..., 2] += push_away\n c_j3d[..., 2] += push_away\n # - If the person is not in the FoV, push away the person in the z direction\n c_root = c_j3d[:, 0] # (L, 3)\n half_fov = torch.div(c_root[:, :2], c_root[:, 2:]).abs() # (L, 2), [x/z, y/z]\n if half_fov.max() > self.half_fov_tol:\n max_idx1, max_idx2 = torch.where(torch.max(half_fov) == half_fov)\n max_idx1, max_idx2 = max_idx1[0], max_idx2[0]\n z_trg = c_root[max_idx1, max_idx2].abs() / self.half_fov_tol # extreme fitted z in the fov\n push_away = z_trg - c_root[max_idx1, 2]\n delta[..., 2] += push_away\n t_w2c += delta\n\n T_w2c = transform_mat(R_w2c, t_w2c) # (F, 4, 4)\n return T_w2c\n" }, { "path": "eval/GVHMR/hmr4d/dataset/pure_motion/utils.py", "content": "import torch\nimport torch.nn.functional as F\nfrom pytorch3d.transforms import (\n axis_angle_to_matrix,\n matrix_to_axis_angle,\n matrix_to_rotation_6d,\n rotation_6d_to_matrix,\n)\nfrom einops import rearrange\n\n\ndef aa_to_r6d(x):\n return matrix_to_rotation_6d(axis_angle_to_matrix(x))\n\n\ndef r6d_to_aa(x):\n return matrix_to_axis_angle(rotation_6d_to_matrix(x))\n\n\ndef interpolate_smpl_params(smpl_params, tgt_len):\n \"\"\"\n smpl_params['body_pose'] (L, 63)\n tgt_len: L->L'\n \"\"\"\n betas = smpl_params[\"betas\"]\n body_pose = smpl_params[\"body_pose\"]\n global_orient = smpl_params[\"global_orient\"] # (L, 3)\n transl = smpl_params[\"transl\"] # (L, 3)\n\n # Interpolate\n body_pose = rearrange(aa_to_r6d(body_pose.reshape(-1, 21, 3)), \"l j c -> c j l\")\n body_pose = F.interpolate(body_pose, tgt_len, mode=\"linear\", align_corners=True)\n body_pose = r6d_to_aa(rearrange(body_pose, \"c j l -> l j c\")).reshape(-1, 63)\n\n # although this should be the same as above, we do it for consistency\n betas = rearrange(betas, \"l c -> c 1 l\")\n betas = F.interpolate(betas, tgt_len, mode=\"linear\", align_corners=True)\n betas = rearrange(betas, \"c 1 l -> l c\")\n\n global_orient = rearrange(aa_to_r6d(global_orient.reshape(-1, 1, 3)), \"l j c -> c j l\")\n global_orient = F.interpolate(global_orient, tgt_len, mode=\"linear\", align_corners=True)\n global_orient = r6d_to_aa(rearrange(global_orient, \"c j l -> l j c\")).reshape(-1, 3)\n\n transl = rearrange(transl, \"l c -> c 1 l\")\n transl = F.interpolate(transl, tgt_len, mode=\"linear\", align_corners=True)\n transl = rearrange(transl, \"c 1 l -> l c\")\n\n return {\"body_pose\": body_pose, \"betas\": betas, \"global_orient\": global_orient, \"transl\": transl}\n\n\ndef rotate_around_axis(global_orient, transl, axis=\"y\"):\n \"\"\"Global coordinate augmentation. Random rotation around y-axis\"\"\"\n angle = torch.rand(1) * 2 * torch.pi\n if axis == \"y\":\n aa = torch.tensor([0.0, angle, 0.0]).float().unsqueeze(0)\n rmat = axis_angle_to_matrix(aa)\n\n global_orient = matrix_to_axis_angle(rmat @ axis_angle_to_matrix(global_orient))\n transl = (rmat.squeeze(0) @ transl.T).T\n return global_orient, transl\n\n\ndef augment_betas(betas, std=0.1):\n noise = torch.normal(mean=torch.zeros(10), std=torch.ones(10) * std)\n betas_aug = betas + noise[None]\n return betas_aug\n" }, { "path": "eval/GVHMR/hmr4d/dataset/rich/resource/seqname2imgrange.json", "content": "{\"ParkingLot1_002_burpee3\": [1, 351], \"ParkingLot1_002_overfence1\": [1, 268], \"ParkingLot1_002_overfence2\": [1, 270], \"ParkingLot1_002_stretching1\": [1, 327], \"ParkingLot1_002_pushup1\": [1, 220], \"ParkingLot1_004_pushup2\": [1, 347], \"ParkingLot1_004_burpeejump1\": [1, 296], \"ParkingLot1_004_eating1\": [1, 522], \"ParkingLot1_004_takingphotos1\": [1, 593], \"ParkingLot1_004_phonetalk1\": [1, 724], \"ParkingLot1_005_burpeejump2\": [1, 270], \"ParkingLot1_005_overfence1\": [1, 301], \"ParkingLot1_005_pushup2\": [1, 262], \"ParkingLot1_005_pushup3\": [1, 243], \"ParkingLot1_004_005_greetingchattingeating1\": [275, 849], \"ParkingLot1_007_overfence2\": [1, 263], \"ParkingLot1_007_eating1\": [1, 426], \"ParkingLot1_007_eating2\": [1, 498], \"ParkingLot2_008_phonetalk1\": [171, 1215], \"ParkingLot2_008_burpeejump1\": [78, 505], \"ParkingLot2_008_overfence1\": [161, 459], \"ParkingLot2_008_pushup1\": [165, 459], \"ParkingLot2_008_pushup2\": [107, 719], \"ParkingLot2_008_overfence2\": [138, 632], \"ParkingLot2_008_overfence3\": [100, 661], \"ParkingLot2_008_eating1\": [180, 1332], \"ParkingLot2_014_pushup2\": [80, 420], \"ParkingLot2_014_burpeejump1\": [50, 348], \"ParkingLot2_014_burpeejump2\": [50, 248], \"ParkingLot2_014_phonetalk2\": [121, 1141], \"ParkingLot2_014_takingphotos2\": [91, 906], \"ParkingLot2_014_overfence3\": [40, 502], \"ParkingLot2_015_overfence1\": [170, 692], \"ParkingLot2_015_burpeejump2\": [344, 678], \"ParkingLot2_015_pushup1\": [190, 817], \"ParkingLot2_015_eating2\": [31, 835], \"ParkingLot2_016_burpeejump2\": [100, 793], \"ParkingLot2_016_overfence2\": [100, 720], \"ParkingLot2_016_pushup1\": [61, 680], \"ParkingLot2_016_pushup2\": [100, 570], \"ParkingLot2_016_stretching1\": [100, 691], \"Pavallion_000_yoga2\": [1, 1643], \"Pavallion_000_plankjack\": [1, 900], \"Pavallion_000_phonesiteat\": [1, 1157], \"Pavallion_000_sidebalancerun\": [1, 1091], \"Pavallion_002_plankjack\": [110, 699], \"Pavallion_002_phonesiteat\": [1, 1030], \"Pavallion_003_plankjack\": [1, 764], \"Pavallion_003_phonesiteat\": [75, 838], \"Pavallion_003_sidebalancerun\": [1, 942], \"Pavallion_006_phonesiteat\": [130, 841], \"Pavallion_006_sidebalancerun\": [1, 798], \"Pavallion_006_plankjack\": [1, 615], \"Pavallion_013_phonesiteat\": [1, 1254], \"Pavallion_013_plankjack\": [1, 641], \"Pavallion_013_yoga2\": [1, 884], \"Pavallion_003_018_tossball\": [230, 949], \"LectureHall_018_wipingchairs1\": [1, 1166], \"LectureHall_018_wipingspray1\": [1, 904], \"LectureHall_020_wipingtable1\": [1, 897], \"BBQ_001_juggle\": [0, 297], \"BBQ_001_guitar\": [0, 381], \"ParkingLot1_002_stretching2\": [240, 240], \"ParkingLot1_002_burpee1\": [1, 286], \"ParkingLot1_002_burpee2\": [1, 203], \"ParkingLot1_004_pushup1\": [1, 354], \"ParkingLot1_004_eating2\": [1, 516], \"ParkingLot1_004_phonetalk2\": [1, 960], \"ParkingLot1_004_takingphotos2\": [1, 571], \"ParkingLot1_004_stretching2\": [1, 399], \"ParkingLot1_005_overfence2\": [1, 298], \"ParkingLot1_005_pushup1\": [1, 476], \"ParkingLot1_005_burpeejump1\": [1, 252], \"ParkingLot1_007_burpee2\": [1, 349], \"ParkingLot2_008_eating2\": [160, 1100], \"ParkingLot2_008_burpeejump2\": [129, 492], \"ParkingLot2_014_overfence1\": [95, 547], \"ParkingLot2_014_eating2\": [101, 986], \"ParkingLot2_016_phonetalk5\": [170, 1259], \"Pavallion_002_sidebalancerun\": [1, 655], \"Pavallion_013_sidebalancerun\": [1, 810], \"Pavallion_018_sidebalancerun\": [1, 873], \"LectureHall_018_wipingtable1\": [1, 1280], \"LectureHall_020_wipingchairs1\": [1, 1163], \"LectureHall_003_wipingchairs1\": [1, 724], \"Pavallion_000_yoga1\": [1, 1757], \"Pavallion_002_yoga1\": [1, 613], \"Pavallion_003_yoga1\": [1, 792], \"Pavallion_006_yoga1\": [1, 930], \"Pavallion_018_yoga1\": [1, 880], \"ParkingLot2_017_burpeejump2\": [118, 612], \"ParkingLot2_017_burpeejump1\": [40, 817], \"ParkingLot2_017_overfence1\": [110, 661], \"ParkingLot2_017_overfence2\": [90, 944], \"ParkingLot2_017_eating1\": [97, 895], \"ParkingLot2_017_pushup1\": [191, 719], \"ParkingLot2_017_pushup2\": [74, 811], \"ParkingLot2_009_burpeejump1\": [200, 1085], \"ParkingLot2_009_burpeejump2\": [150, 399], \"ParkingLot2_009_overfence1\": [140, 601], \"ParkingLot2_009_overfence2\": [150, 559], \"LectureHall_009_sidebalancerun1\": [1, 673], \"LectureHall_010_plankjack1\": [1, 532], \"LectureHall_010_sidebalancerun1\": [1, 919], \"LectureHall_021_plankjack1\": [1, 507], \"LectureHall_021_sidebalancerun1\": [1, 855], \"LectureHall_019_wipingchairs1\": [1, 978], \"LectureHall_009_021_reparingprojector1\": [1, 499], \"ParkingLot2_009_spray1\": [145, 1242], \"ParkingLot2_009_impro1\": [100, 990], \"ParkingLot2_009_impro2\": [100, 1140], \"ParkingLot2_009_impro5\": [100, 649], \"Gym_010_pushup1\": [1, 475], \"Gym_010_pushup2\": [1, 407], \"Gym_011_pushup1\": [1, 346], \"Gym_011_pushup2\": [1, 540], \"Gym_011_burpee2\": [1, 479], \"Gym_012_pushup2\": [1, 291], \"Gym_010_mountainclimber1\": [0, 0], \"Gym_010_mountainclimber2\": [1, 471], \"Gym_013_dips1\": [1, 503], \"Gym_013_dips2\": [1, 333], \"Gym_013_dips3\": [1, 502], \"Gym_013_lunge1\": [1, 690], \"Gym_013_lunge2\": [1, 834], \"Gym_013_pushup1\": [1, 861], \"Gym_013_pushup2\": [1, 477], \"Gym_013_burpee4\": [1, 320], \"Gym_010_lunge1\": [1, 337], \"Gym_010_lunge2\": [1, 312], \"Gym_010_dips1\": [1, 572], \"Gym_010_dips2\": [1, 603], \"Gym_010_cooking1\": [1, 779], \"Gym_011_cooking1\": [1, 1141], \"Gym_011_cooking2\": [1, 1145], \"Gym_011_dips1\": [1, 494], \"Gym_011_dips4\": [1, 495], \"Gym_011_dips3\": [1, 320], \"Gym_011_dips2\": [1, 382], \"Gym_012_lunge1\": [1, 225], \"Gym_012_lunge2\": [1, 318], \"Gym_012_cooking2\": [1, 993]}" }, { "path": "eval/GVHMR/hmr4d/dataset/rich/resource/test.txt", "content": "sequence_name\tcapture_name\tscan_name\tid\tmoving_cam\tgender\tscene\taction/scene-interaction\tsubjects\tview_id\nParkingLot2_017_burpeejump2\tParkingLot2\tscan_camcoord\t017\tV\tfemale\tV\tV\tX\t0,2,3\nParkingLot2_017_burpeejump1\tParkingLot2\tscan_camcoord\t017\tV\tfemale\tV\tV\tX\t0,1,5\nParkingLot2_017_overfence1\tParkingLot2\tscan_camcoord\t017\tV\tfemale\tV\tV\tX\t0,3,4\nParkingLot2_017_overfence2\tParkingLot2\tscan_camcoord\t017\tV\tfemale\tV\tV\tX\t0,1,4\nParkingLot2_017_eating1\tParkingLot2\tscan_camcoord\t017\tV\tfemale\tV\tV\tX\t0,2,4\nParkingLot2_017_pushup1\tParkingLot2\tscan_camcoord\t017\tX\tfemale\tV\tV\tX\t0,1,4,5\nParkingLot2_017_pushup2\tParkingLot2\tscan_camcoord\t017\tV\tfemale\tV\tV\tX\t0,4,5\nParkingLot2_009_burpeejump1\tParkingLot2\tscan_camcoord\t009\tX\tfemale\tV\tV\tX\t0,1,2,3\nParkingLot2_009_burpeejump2\tParkingLot2\tscan_camcoord\t009\tX\tfemale\tV\tV\tX\t0,2,3,4\nParkingLot2_009_overfence1\tParkingLot2\tscan_camcoord\t009\tX\tfemale\tV\tV\tX\t0,3,4,5\nParkingLot2_009_overfence2\tParkingLot2\tscan_camcoord\t009\tX\tfemale\tV\tV\tX\t0,1,4,5\nLectureHall_009_sidebalancerun1\tLectureHall\tscan_yoga_scene_camcoord\t009\tX\tfemale\tV\tV\tX\t0,1,4,5\nLectureHall_010_plankjack1\tLectureHall\tscan_yoga_scene_camcoord\t010\tX\tfemale\tV\tV\tX\t0,2,4,6\nLectureHall_010_sidebalancerun1\tLectureHall\tscan_yoga_scene_camcoord\t010\tX\tfemale\tV\tV\tX\t0,1,2,4\nLectureHall_021_plankjack1\tLectureHall\tscan_yoga_scene_camcoord\t021\tX\tfemale\tV\tV\tX\t0,3,5,6\nLectureHall_021_sidebalancerun1\tLectureHall\tscan_yoga_scene_camcoord\t021\tX\tfemale\tV\tV\tX\t0,4,5,6\nLectureHall_019_wipingchairs1\tLectureHall\tscan_chair_scene_camcoord\t019\tX\tfemale\tV\tV\tX\t0,1,2,3\nLectureHall_009_021_reparingprojector1\tLectureHall\tscan_yoga_scene_camcoord\t009\tX\tfemale\tV\tX\tX\t0,3,4,5\nLectureHall_009_021_reparingprojector1\tLectureHall\tscan_yoga_scene_camcoord\t021\tX\tfemale\tV\tX\tX\t0,3,4,5\nParkingLot2_009_spray1\tParkingLot2\tscan_camcoord\t009\tX\tfemale\tV\tX\tX\t0,1,2,3\nParkingLot2_009_impro1\tParkingLot2\tscan_camcoord\t009\tX\tfemale\tV\tX\tX\t0,2,3,4\nParkingLot2_009_impro2\tParkingLot2\tscan_camcoord\t009\tX\tfemale\tV\tX\tX\t0,3,4,5\nParkingLot2_009_impro5\tParkingLot2\tscan_camcoord\t009\tX\tfemale\tV\tX\tX\t0,2,4,5\nGym_010_pushup1\tGym\tscan_camcoord\t010\tX\tfemale\tX\tV\tX\t3,4,5,6\nGym_010_pushup2\tGym\tscan_camcoord\t010\tX\tfemale\tX\tV\tX\t2,3,4,5\nGym_011_pushup1\tGym\tscan_camcoord\t011\tX\tmale\tX\tV\tX\t2,3,4,5\nGym_011_pushup2\tGym\tscan_camcoord\t011\tX\tmale\tX\tV\tX\t2,3,4,5\nGym_011_burpee2\tGym\tscan_camcoord\t011\tX\tmale\tX\tV\tX\t2,3,4,5\nGym_012_pushup2\tGym\tscan_camcoord\t012\tX\tfemale\tX\tV\tX\t3,4,5,6\nGym_010_mountainclimber1\tGym\tscan_camcoord\t010\tX\tfemale\tX\tV\tX\t3,4,5,6\nGym_010_mountainclimber2\tGym\tscan_camcoord\t010\tX\tfemale\tX\tV\tX\t3,4,5,6\nGym_013_dips1\tGym\tscan_camcoord\t013\tX\tfemale\tX\tX\tV\t0,3,4,5\nGym_013_dips2\tGym\tscan_camcoord\t013\tX\tfemale\tX\tX\tV\t1,2,4,5\nGym_013_dips3\tGym\tscan_camcoord\t013\tX\tfemale\tX\tX\tV\t1,2,4,5\nGym_013_lunge1\tGym\tscan_camcoord\t013\tX\tfemale\tX\tX\tV\t1,4,5,6\nGym_013_lunge2\tGym\tscan_camcoord\t013\tX\tfemale\tX\tX\tV\t0,4,5,6\nGym_013_pushup1\tGym\tscan_camcoord\t013\tX\tfemale\tX\tV\tV\t0,3,4,5\nGym_013_pushup2\tGym\tscan_camcoord\t013\tX\tfemale\tX\tV\tV\t1,2,4,5\nGym_013_burpee4\tGym\tscan_camcoord\t013\tX\tfemale\tX\tV\tV\t0,4,5,6\nGym_010_lunge1\tGym\tscan_camcoord\t010\tX\tfemale\tX\tX\tX\t1,4,5,6\nGym_010_lunge2\tGym\tscan_camcoord\t010\tX\tfemale\tX\tX\tX\t0,2,4,5\nGym_010_dips1\tGym\tscan_camcoord\t010\tX\tfemale\tX\tX\tX\t0,4,5,6\nGym_010_dips2\tGym\tscan_camcoord\t010\tX\tfemale\tX\tX\tX\t1,2,4,5\nGym_010_cooking1\tGym\tscan_table_camcoord\t010\tX\tfemale\tX\tX\tX\t1,3,4,5\nGym_011_cooking1\tGym\tscan_table_camcoord\t011\tV\tmale\tX\tX\tX\t4,5,6\nGym_011_cooking2\tGym\tscan_table_camcoord\t011\tV\tmale\tX\tX\tX\t2,4,5\nGym_011_dips1\tGym\tscan_camcoord\t011\tX\tmale\tX\tX\tX\t1,3,4,5\nGym_011_dips4\tGym\tscan_camcoord\t011\tX\tmale\tX\tX\tX\t0,2,4,5\nGym_011_dips3\tGym\tscan_camcoord\t011\tX\tmale\tX\tX\tX\t0,3,4,5\nGym_011_dips2\tGym\tscan_camcoord\t011\tX\tmale\tX\tX\tX\t1,3,4,5\nGym_012_lunge1\tGym\tscan_camcoord\t012\tX\tfemale\tX\tX\tX\t0,3,4,5\nGym_012_lunge2\tGym\tscan_camcoord\t012\tX\tfemale\tX\tX\tX\t0,4,5,6\nGym_012_cooking2\tGym\tscan_table_camcoord\t012\tV\tfemale\tX\tX\tX\t3,4,5" }, { "path": "eval/GVHMR/hmr4d/dataset/rich/resource/train.txt", "content": "sequence_name\tcapture_name\tscan_name\tid\tmoving_cam\tgender\tview_id\nParkingLot1_002_burpee3\tParkingLot1\tscan_camcoord\t002\tX\tmale\t0,1,2,3,4,5,6,7\nParkingLot1_002_overfence1\tParkingLot1\tscan_camcoord\t002\tX\tmale\t0,1,2,3,4,5,6,7\nParkingLot1_002_overfence2\tParkingLot1\tscan_camcoord\t002\tX\tmale\t0,1,2,3,4,5,6,7\nParkingLot1_002_stretching1\tParkingLot1\tscan_camcoord\t002\tX\tmale\t0,1,2,3,4,5,6,7\nParkingLot1_002_pushup1\tParkingLot1\tscan_camcoord\t002\tX\tmale\t0,1,2,3,4,5,6,7\nParkingLot1_004_pushup2\tParkingLot1\tscan_camcoord\t004\tX\tmale\t0,1,2,3,4,5,6,7\nParkingLot1_004_burpeejump1\tParkingLot1\tscan_camcoord\t004\tX\tmale\t0,1,2,3,4,5,6,7\nParkingLot1_004_eating1\tParkingLot1\tscan_camcoord\t004\tX\tmale\t0,1,2,3,4,5,6,7\nParkingLot1_004_takingphotos1\tParkingLot1\tscan_camcoord\t004\tX\tmale\t0,1,2,3,4,5,6,7\nParkingLot1_004_phonetalk1\tParkingLot1\tscan_camcoord\t004\tX\tmale\t0,1,2,3,4,5,6,7\nParkingLot1_005_burpeejump2\tParkingLot1\tscan_camcoord\t005\tX\tmale\t0,1,2,3,4,5,6,7\nParkingLot1_005_overfence1\tParkingLot1\tscan_camcoord\t005\tX\tmale\t0,1,2,3,4,5,6,7\nParkingLot1_005_pushup2\tParkingLot1\tscan_camcoord\t005\tX\tmale\t0,1,2,3,4,5,6,7\nParkingLot1_005_pushup3\tParkingLot1\tscan_camcoord\t005\tX\tmale\t0,1,2,3,4,5,6,7\nParkingLot1_004_005_greetingchattingeating1\tParkingLot1\tscan_camcoord\t004\tX\tmale\t0,1,2,3,4,5,6,7\nParkingLot1_004_005_greetingchattingeating1\tParkingLot1\tscan_camcoord\t005\tX\tmale\t0,1,2,3,4,5,6,7\nParkingLot1_007_overfence2\tParkingLot1\tscan_camcoord\t007\tX\tmale\t0,1,2,3,4,5,6,7\nParkingLot1_007_eating1\tParkingLot1\tscan_camcoord\t007\tX\tmale\t0,1,2,3,4,5,6,7\nParkingLot1_007_eating2\tParkingLot1\tscan_camcoord\t007\tX\tmale\t0,1,2,3,4,5,6,7\nParkingLot2_008_phonetalk1\tParkingLot2\tscan_camcoord\t008\tV\tmale\t0,1,2,3,4,5\nParkingLot2_008_burpeejump1\tParkingLot2\tscan_camcoord\t008\tV\tmale\t0,1,2,3,4,5\nParkingLot2_008_overfence1\tParkingLot2\tscan_camcoord\t008\tV\tmale\t0,1,2,3,4,5\nParkingLot2_008_pushup1\tParkingLot2\tscan_camcoord\t008\tV\tmale\t0,1,2,3,4,5\nParkingLot2_008_pushup2\tParkingLot2\tscan_camcoord\t008\tV\tmale\t0,1,2,3,4,5\nParkingLot2_008_overfence2\tParkingLot2\tscan_camcoord\t008\tV\tmale\t0,1,2,3,4,5\nParkingLot2_008_overfence3\tParkingLot2\tscan_camcoord\t008\tV\tmale\t0,1,2,3,4,5\nParkingLot2_008_eating1\tParkingLot2\tscan_camcoord\t008\tV\tmale\t0,1,2,3,4,5\nParkingLot2_014_pushup2\tParkingLot2\tscan_camcoord\t014\tX\tmale\t0,1,2,3,4,5\nParkingLot2_014_burpeejump1\tParkingLot2\tscan_camcoord\t014\tX\tmale\t0,1,2,3,4,5\nParkingLot2_014_burpeejump2\tParkingLot2\tscan_camcoord\t014\tX\tmale\t0,1,2,3,4,5\nParkingLot2_014_phonetalk2\tParkingLot2\tscan_camcoord\t014\tX\tmale\t0,1,2,3,4,5\nParkingLot2_014_takingphotos2\tParkingLot2\tscan_camcoord\t014\tX\tmale\t0,1,2,3,4,5\nParkingLot2_014_overfence3\tParkingLot2\tscan_camcoord\t014\tX\tmale\t0,1,2,3,4,5\nParkingLot2_015_overfence1\tParkingLot2\tscan_camcoord\t015\tX\tmale\t0,1,2,3,4,5\nParkingLot2_015_burpeejump2\tParkingLot2\tscan_camcoord\t015\tX\tmale\t0,1,2,3,4,5\nParkingLot2_015_pushup1\tParkingLot2\tscan_camcoord\t015\tX\tmale\t0,1,2,3,4,5\nParkingLot2_015_eating2\tParkingLot2\tscan_camcoord\t015\tX\tmale\t0,1,2,3,4,5\nParkingLot2_016_burpeejump2\tParkingLot2\tscan_camcoord\t016\tV\tfemale\t0,1,2,3,4,5\nParkingLot2_016_overfence2\tParkingLot2\tscan_camcoord\t016\tV\tfemale\t0,1,2,3,4,5\nParkingLot2_016_pushup1\tParkingLot2\tscan_camcoord\t016\tV\tfemale\t0,1,2,3,4,5\nParkingLot2_016_pushup2\tParkingLot2\tscan_camcoord\t016\tV\tfemale\t0,1,2,3,4,5\nParkingLot2_016_stretching1\tParkingLot2\tscan_camcoord\t016\tV\tfemale\t0,1,2,3,4,5\nPavallion_000_yoga2\tPavallion\tscan_camcoord\t000\tX\tmale\t0,1,2,3,4,5,6\nPavallion_000_plankjack\tPavallion\tscan_camcoord\t000\tX\tmale\t0,1,2,3,4,5,6\nPavallion_000_phonesiteat\tPavallion\tscan_camcoord\t000\tX\tmale\t0,1,3,4,6\nPavallion_000_sidebalancerun\tPavallion\tscan_camcoord\t000\tX\tmale\t0,1,2,3,4,5,6\nPavallion_002_plankjack\tPavallion\tscan_camcoord\t002\tV\tmale\t0,1,2,3,4,5,6\nPavallion_002_phonesiteat\tPavallion\tscan_camcoord\t002\tV\tmale\t0,1,3,4,6\nPavallion_003_plankjack\tPavallion\tscan_camcoord\t003\tV\tmale\t0,1,2,3,4,5,6\nPavallion_003_phonesiteat\tPavallion\tscan_camcoord\t003\tV\tmale\t0,1,3,4,6\nPavallion_003_sidebalancerun\tPavallion\tscan_camcoord\t003\tV\tmale\t0,1,2,3,4,5,6\nPavallion_006_phonesiteat\tPavallion\tscan_camcoord\t006\tV\tmale\t0,1,3,4,6\nPavallion_006_sidebalancerun\tPavallion\tscan_camcoord\t006\tV\tmale\t0,1,2,3,4,5,6\nPavallion_006_plankjack\tPavallion\tscan_camcoord\t006\tV\tmale\t0,1,2,3,4,5,6\nPavallion_013_phonesiteat\tPavallion\tscan_camcoord\t013\tX\tfemale\t0,1,3,4,6\nPavallion_013_plankjack\tPavallion\tscan_camcoord\t013\tX\tfemale\t0,1,2,3,4,5,6\nPavallion_013_yoga2\tPavallion\tscan_camcoord\t013\tV\tfemale\t0,1,2,3,4,5,6\nPavallion_003_018_tossball\tPavallion\tscan_camcoord\t003\tX\tmale\t0,1,2,3,4,5,6\nPavallion_003_018_tossball\tPavallion\tscan_camcoord\t018\tX\tfemale\t0,1,2,3,4,5,6\nLectureHall_018_wipingchairs1\tLectureHall\tscan_chair_scene_camcoord\t018\tX\tfemale\t0,1,2,3,4,5,6\nLectureHall_018_wipingspray1\tLectureHall\tscan_chair_scene_camcoord\t018\tX\tfemale\t2,3,4\nLectureHall_020_wipingtable1\tLectureHall\tscan_chair_scene_camcoord\t020\tX\tmale\t0,2,4,5,6\nBBQ_001_juggle\tBBQ\tscan_camcoord\t001\tX\tmale\t0,1,2,3,4,5,6,7\nBBQ_001_guitar\tBBQ\tscan_camcoord\t001\tX\tmale\t0,1,2,3,4,5,6,7" }, { "path": "eval/GVHMR/hmr4d/dataset/rich/resource/val.txt", "content": "sequence_name\tcapture_name\tscan_name\tid\tmoving_cam\tgender\tscene\taction/scene-interaction\tsubjects\tview_id\nParkingLot1_002_stretching2\tParkingLot1\tscan_camcoord\t002\tX\tmale\tV\tV\tV\t0,1,2,3,4,5,6,7\nParkingLot1_002_burpee1\tParkingLot1\tscan_camcoord\t002\tX\tmale\tV\tV\tV\t0,1,2,3,4,5,6,7\nParkingLot1_002_burpee2\tParkingLot1\tscan_camcoord\t002\tX\tmale\tV\tV\tV\t0,1,2,3,4,5,6,7\nParkingLot1_004_pushup1\tParkingLot1\tscan_camcoord\t004\tX\tmale\tV\tV\tV\t0,1,2,3,4,5,6,7\nParkingLot1_004_eating2\tParkingLot1\tscan_camcoord\t004\tX\tmale\tV\tV\tV\t0,1,2,3,4,5,6,7\nParkingLot1_004_phonetalk2\tParkingLot1\tscan_camcoord\t004\tX\tmale\tV\tV\tV\t0,1,2,3,4,5,6,7\nParkingLot1_004_takingphotos2\tParkingLot1\tscan_camcoord\t004\tX\tmale\tV\tV\tV\t0,1,2,3,4,5,6,7\nParkingLot1_004_stretching2\tParkingLot1\tscan_camcoord\t004\tX\tmale\tV\tV\tV\t0,1,2,3,4,5,6,7\nParkingLot1_005_overfence2\tParkingLot1\tscan_camcoord\t005\tX\tmale\tV\tV\tV\t0,1,2,3,4,5,6,7\nParkingLot1_005_pushup1\tParkingLot1\tscan_camcoord\t005\tX\tmale\tV\tV\tV\t0,1,2,3,4,5,6,7\nParkingLot1_005_burpeejump1\tParkingLot1\tscan_camcoord\t005\tX\tmale\tV\tV\tV\t0,1,2,3,4,5,6,7\nParkingLot1_007_burpee2\tParkingLot1\tscan_camcoord\t007\tX\tmale\tV\tV\tV\t0,1,2,3,4,5,6,7\nParkingLot2_008_eating2\tParkingLot2\tscan_camcoord\t008\tV\tmale\tV\tV\tV\t0,1,2,3,4,5\nParkingLot2_008_burpeejump2\tParkingLot2\tscan_camcoord\t008\tV\tmale\tV\tV\tV\t0,1,2,3,4,5\nParkingLot2_014_overfence1\tParkingLot2\tscan_camcoord\t014\tX\tmale\tV\tV\tV\t0,1,2,3,4,5\nParkingLot2_014_eating2\tParkingLot2\tscan_camcoord\t014\tX\tmale\tV\tV\tV\t0,1,2,3,4,5\nParkingLot2_016_phonetalk5\tParkingLot2\tscan_camcoord\t016\tV\tfemale\tV\tV\tV\t0,1,2,3,4,5\nPavallion_002_sidebalancerun\tPavallion\tscan_camcoord\t002\tV\tmale\tV\tV\tV\t0,1,2,3,4,5,6\nPavallion_013_sidebalancerun\tPavallion\tscan_camcoord\t013\tX\tfemale\tV\tV\tV\t0,1,2,3,4,5,6\nPavallion_018_sidebalancerun\tPavallion\tscan_camcoord\t018\tV\tfemale\tV\tV\tV\t0,1,2,3,4,5,6\nLectureHall_018_wipingtable1\tLectureHall\tscan_chair_scene_camcoord\t018\tX\tfemale\tV\tV\tV\t0,2,4,5,6\nLectureHall_020_wipingchairs1\tLectureHall\tscan_chair_scene_camcoord\t020\tX\tmale\tV\tV\tV\t0,1,2,3,4,5,6\nLectureHall_003_wipingchairs1\tLectureHall\tscan_chair_scene_camcoord\t003\tX\tmale\tV\tV\tV\t0,1,2,3,4,5,6\nPavallion_000_yoga1\tPavallion\tscan_camcoord\t000\tX\tmale\tV\tX\tV\t0,1,2,3,4,5,6\nPavallion_002_yoga1\tPavallion\tscan_camcoord\t002\tV\tmale\tV\tX\tV\t0,1,2,3,4,5,6\nPavallion_003_yoga1\tPavallion\tscan_camcoord\t003\tV\tmale\tV\tX\tV\t0,1,2,3,4,5,6\nPavallion_006_yoga1\tPavallion\tscan_camcoord\t006\tV\tmale\tV\tX\tV\t0,1,2,3,4,5,6\nPavallion_018_yoga1\tPavallion\tscan_camcoord\t018\tV\tfemale\tV\tX\tV\t0,1,2,3,4,5,6" }, { "path": "eval/GVHMR/hmr4d/dataset/rich/resource/w2az_sahmr.json", "content": "{\"BBQ_scan_camcoord\": [[0.9989829107564298, 0.03367618890797693, -0.029984301180211045, 0.0008183751635392625], [0.03414262169451401, -0.1305975871406019, 0.9908473906797644, -0.005059823133706893], [0.02945208652127451, -0.9908633531086326, -0.13161455111748036, 1.4054905296083466], [0.0, 0.0, 0.0, 1.0]], \"Gym_scan_camcoord\": [[0.9932599733260449, -0.07628732032461205, 0.0872632233306122, -0.047601130084306706], [-0.10233962102690007, -0.22374853741942266, 0.9692590953768503, -0.04091804681182174], [-0.05441716049582774, -0.9716567484252654, -0.23004768176013274, 1.537911791136788], [0.0, 0.0, 0.0, 1.0]], \"Gym_scan_table_camcoord\": [[0.9974451989415423, -0.06250743213795668, 0.03458172980064169, 0.02231858470834599], [-0.04804912583358893, -0.22882402250236075, 0.972281259838159, 0.039081886755815726], [-0.05286167435026744, -0.9714588965331274, -0.2312428501197992, 1.5421821446346522], [0.0, 0.0, 0.0, 1.0]], \"LectureHall_scan_chair_scene_camcoord\": [[0.9992930513998263, 0.030087515976743376, -0.0225419343977731, 0.001998908749589632], [0.030705594681969043, -0.30721111058653017, 0.9511458878570781, -0.025811963513866963], [0.021692484396004613, -0.9511656401040444, -0.307917783192506, 2.060346184503773], [0.0, 0.0, 0.0, 1.0]], \"LectureHall_scan_yoga_scene_camcoord\": [[0.9993358324246812, 0.03030060260429296, -0.020242715082476024, -0.003510046042036605], [0.028600729415016745, -0.3079667078507395, 0.9509671419836329, -0.01748548118379142], [0.022580795137075255, -0.9509144968594153, -0.3086287856852993, 2.0424701474796567], [0.0, 0.0, 0.0, 1.0]], \"ParkingLot1_scan_camcoord\": [[0.9989627324729327, -0.03724260727951709, 0.02620013994738054, 0.0070941466745699025], [-0.03091587075252664, -0.13228243926883107, 0.9907298144280939, -0.0274920377236923], [-0.03343154297742938, -0.9905121627037764, -0.13329661462331338, 1.3859200914120975], [0.0, 0.0, 0.0, 1.0]], \"ParkingLot2_scan_camcoord\": [[0.9989532636786039, -0.04044665659892979, 0.021364572447267097, 0.01646827411554571], [-0.026687287930043047, -0.13600581518076985, 0.9903485279940424, 0.030197722289598695], [-0.03715058073335097, -0.9898820567153364, -0.13694286452455984, 1.4372015171546513], [0.0, 0.0, 0.0, 1.0]], \"Pavallion_scan_camcoord\": [[0.9971864096076799, 0.05693557331723671, -0.048760690979605295, 0.0012478238054067193], [0.05746407703876882, -0.16289761936471214, 0.9849681443861059, -0.006002953831755452], [0.04813672552068054, -0.9849988355812122, -0.16571104235928033, 1.7638454838942128], [0.0, 0.0, 0.0, 1.0]]}" }, { "path": "eval/GVHMR/hmr4d/dataset/rich/rich_motion_test.py", "content": "from pathlib import Path\nimport numpy as np\nimport torch\nfrom torch.utils import data\nfrom hmr4d.utils.pylogger import Log\n\nfrom .rich_utils import (\n get_cam2params,\n get_w2az_sahmr,\n parse_seqname_info,\n get_cam_key_wham_vid,\n)\nfrom hmr4d.utils.geo_transform import apply_T_on_points, transform_mat, compute_cam_angvel\nfrom hmr4d.utils.wis3d_utils import make_wis3d, add_motion_as_lines\nfrom hmr4d.utils.smplx_utils import make_smplx\nfrom pytorch3d.transforms import axis_angle_to_matrix, matrix_to_axis_angle\nfrom hmr4d.utils.geo.hmr_cam import resize_K\n\n\nfrom hmr4d.configs import MainStore, builds\n\n\nVID_PRESETS = {\n \"easytohard\": [\n \"test/Gym_013_burpee4/cam_06\",\n \"test/Gym_011_pushup1/cam_02\",\n \"test/LectureHall_019_wipingchairs1/cam_03\",\n \"test/ParkingLot2_009_overfence1/cam_04\",\n \"test/LectureHall_021_sidebalancerun1/cam_00\",\n \"test/Gym_010_dips2/cam_05\",\n ],\n}\n\n\nclass RichSmplFullSeqDataset(data.Dataset):\n def __init__(self, vid_presets=None):\n \"\"\"\n Args:\n vid_presets is a key in VID_PRESETS\n \"\"\"\n super().__init__()\n self.dataset_name = \"RICH\"\n self.dataset_id = \"RICH\"\n Log.info(f\"[{self.dataset_name}] Full sequence, Test\")\n tic = Log.time()\n\n # Load evaluation protocol from WHAM labels\n self.rich_dir = Path(\"inputs/RICH/hmr4d_support\")\n self.labels = torch.load(self.rich_dir / \"rich_test_labels.pt\")\n self.preproc_data = torch.load(self.rich_dir / \"rich_test_preproc.pt\")\n vids = select_subset(self.labels, vid_presets)\n\n # Setup dataset index\n self.idx2meta = []\n for vid in vids:\n seq_length = len(self.labels[vid][\"frame_id\"])\n self.idx2meta.append((vid, 0, seq_length)) # start=0, end=seq_length\n # print(sum([end - start for _, _, start, end in self.idx2meta]))\n\n # Prepare ground truth motion in ay-coordinate\n self.w2az = get_w2az_sahmr() # scan_name -> T_w2az, w-coordinate refers to cam-1-coordinate\n self.cam2params = get_cam2params() # cam_key -> (T_w2c, K)\n seqname_info = parse_seqname_info(skip_multi_persons=True) # {k: (scan_name, subject_id, gender, cam_ids)}\n self.seqname_to_scanname = {k: v[0] for k, v in seqname_info.items()}\n\n Log.info(f\"[RICH] {len(self.idx2meta)} sequences. Elapsed: {Log.time() - tic:.2f}s\")\n\n def __len__(self):\n return len(self.idx2meta)\n\n def _load_data(self, idx):\n data = {}\n\n # [start, end), when loading data from labels\n vid, start, end = self.idx2meta[idx]\n label = self.labels[vid]\n preproc_data = self.preproc_data[vid]\n\n length = end - start\n meta = {\"dataset_id\": \"RICH\", \"vid\": vid, \"vid-start-end\": (start, end)}\n data.update({\"meta\": meta, \"length\": length})\n\n # SMPLX\n data.update({\"gt_smpl_params\": label[\"gt_smplx_params\"], \"gender\": label[\"gender\"]})\n\n # camera\n cam_key = get_cam_key_wham_vid(vid)\n scan_name = self.seqname_to_scanname[vid.split(\"/\")[1]]\n T_w2c, K = self.cam2params[cam_key] # (4, 4) (3, 3)\n T_w2az = self.w2az[scan_name]\n data.update({\"T_w2c\": T_w2c, \"T_w2az\": T_w2az, \"K\": K})\n\n # image features\n data.update(\n {\n \"f_imgseq\": preproc_data[\"f_imgseq\"],\n \"bbx_xys\": preproc_data[\"bbx_xys\"],\n \"img_wh\": preproc_data[\"img_wh\"],\n \"kp2d\": preproc_data[\"kp2d\"],\n }\n )\n\n # to render a video\n video_path = self.rich_dir / \"video\" / vid / \"video.mp4\"\n frame_id = label[\"frame_id\"] # (F,)\n width, height = data[\"img_wh\"] / 4 # Video saved has been downsampled 1/4\n K_render = resize_K(K, 0.25)\n bbx_xys_render = data[\"bbx_xys\"] / 4\n data[\"meta_render\"] = {\n \"name\": vid.replace(\"/\", \"@\"),\n \"video_path\": str(video_path),\n \"frame_id\": frame_id,\n \"width_height\": (width, height),\n \"K\": K_render,\n \"bbx_xys\": bbx_xys_render,\n }\n\n return data\n\n def _process_data(self, data):\n # T_w2az is pre-computed by using floor clue. az2zy uses a rotation along x-axis.\n R_az2ay = axis_angle_to_matrix(torch.tensor([1.0, 0.0, 0.0]) * -torch.pi / 2) # (3, 3)\n T_w2ay = transform_mat(R_az2ay, R_az2ay.new([0, 0, 0])) @ data[\"T_w2az\"] # (4, 4)\n\n if False: # Visualize groundtruth and observation\n self.rich_smplx = {\n \"male\": make_smplx(\"rich-smplx\", gender=\"male\"),\n \"female\": make_smplx(\"rich-smplx\", gender=\"female\"),\n }\n wis3d = make_wis3d(name=\"debug-rich-smpl_dataset\")\n rich_smplx = make_smplx(\"rich-smplx\", gender=data[\"gender\"])\n smplx_out = rich_smplx(**data[\"gt_smpl_params\"])\n smplx_verts_ay = apply_T_on_points(smplx_out.vertices, T_w2ay)\n for i in range(400):\n wis3d.set_scene_id(i)\n wis3d.add_mesh(smplx_out.vertices[i], rich_smplx.bm.faces, name=f\"gt-smplx\")\n wis3d.add_mesh(smplx_verts_ay[i], rich_smplx.bm.faces, name=f\"gt-smplx-ay\")\n\n # process img feature with xys\n length = data[\"length\"]\n f_imgseq = data[\"f_imgseq\"] # (F, 1024)\n R_w2c = data[\"T_w2c\"][:3, :3].repeat(length, 1, 1) # (L, 4, 4)\n cam_angvel = compute_cam_angvel(R_w2c) # (L, 6)\n\n # Return\n data = {\n # --- not batched\n \"task\": \"CAP-Seq\",\n \"meta\": data[\"meta\"],\n \"meta_render\": data[\"meta_render\"],\n # --- we test on single sequence, so set kv manually\n \"length\": length,\n \"f_imgseq\": f_imgseq,\n \"cam_angvel\": cam_angvel,\n \"bbx_xys\": data[\"bbx_xys\"], # (F, 3)\n \"K_fullimg\": data[\"K\"][None].expand(length, -1, -1), # (F, 3, 3)\n \"kp2d\": data[\"kp2d\"], # (F, 17, 3)\n # --- dataset specific\n \"model\": \"smplx\",\n \"gender\": data[\"gender\"],\n \"gt_smpl_params\": data[\"gt_smpl_params\"],\n \"T_w2ay\": T_w2ay, # (4, 4)\n \"T_w2c\": data[\"T_w2c\"], # (4, 4)\n }\n return data\n\n def __getitem__(self, idx):\n data = self._load_data(idx)\n data = self._process_data(data)\n return data\n\n\ndef select_subset(labels, vid_presets):\n vids = list(labels.keys())\n if vid_presets != None: # Use a subset of the videos\n vids = VID_PRESETS[vid_presets]\n return vids\n\n\n#\ngroup_name = \"test_datasets/rich\"\nbase_node = builds(RichSmplFullSeqDataset, vid_presets=None, populate_full_signature=True)\nMainStore.store(name=\"all\", node=base_node, group=group_name)\nMainStore.store(name=\"easy_to_hard\", node=base_node(vid_presets=\"easytohard\"), group=group_name)\nMainStore.store(name=\"postproc\", node=base_node(vid_presets=\"postproc\"), group=group_name)\n" }, { "path": "eval/GVHMR/hmr4d/dataset/rich/rich_utils.py", "content": "import torch\nimport cv2\nimport numpy as np\nfrom hmr4d.utils.geo_transform import apply_T_on_points, project_p2d\nfrom pathlib import Path\nimport json\nimport time\n\n# ----- Meta sample utils ----- #\n\n\ndef sample_idx2meta(idx2meta, sample_interval):\n \"\"\"\n 1. remove frames that < 45\n 2. sample frames by sample_interval\n 3. sorted\n \"\"\"\n idx2meta = [\n v\n for k, v in idx2meta.items()\n if int(v[\"frame_name\"]) > 45 and (int(v[\"frame_name\"]) + int(v[\"cam_id\"])) % sample_interval == 0\n ]\n idx2meta = sorted(idx2meta, key=lambda meta: meta[\"img_key\"])\n return idx2meta\n\n\ndef remove_bbx_invisible_frame(idx2meta, img2gtbbx):\n raw_img_lu = np.array([0.0, 0.0])\n raw_img_rb_type1 = np.array([4112.0, 3008.0]) - 1 # horizontal\n raw_img_rb_type2 = np.array([3008.0, 4112.0]) - 1 # vertical\n\n idx2meta_new = []\n for meta in idx2meta:\n gtbbx_center = np.array([img2gtbbx[meta[\"img_key\"]][[0, 2]].mean(), img2gtbbx[meta[\"img_key\"]][[1, 3]].mean()])\n if (gtbbx_center < raw_img_lu).any():\n continue\n raw_img_rb = raw_img_rb_type1 if meta[\"cam_key\"] not in [\"Pavallion_3\", \"Pavallion_5\"] else raw_img_rb_type2\n if (gtbbx_center > raw_img_rb).any():\n continue\n idx2meta_new.append(meta)\n return idx2meta_new\n\n\ndef remove_extra_rules(idx2meta):\n multi_person_seqs = [\"LectureHall_009_021_reparingprojector1\"]\n idx2meta = [meta for meta in idx2meta if meta[\"seq_name\"] not in multi_person_seqs]\n return idx2meta\n\n\n# ----- Image utils ----- #\n\n\ndef compute_bbx(dataset, data):\n \"\"\"\n Use gt_smplh_params to compute bbx (w.r.t. original image resolution)\n Args:\n dataset: rich_pose.RichPose\n data: dict\n\n # This function need extra scripts to run\n from hmr4d.utils.smplx_utils import make_smplx\n self.smplh_male = make_smplx(\"rich-smplh\", gender=\"male\")\n self.smplh_female = make_smplx(\"rich-smplh\", gender=\"female\")\n self.smplh = {\n \"male\": self.smplh_male,\n \"female\": self.smplh_female,\n }\n \"\"\"\n gender = data[\"meta\"][\"gender\"]\n smplh_params = {k: v.reshape(1, -1) for k, v in data[\"gt_smplh_params\"].items()}\n smplh_opt = dataset.smplh[gender](**smplh_params)\n verts_3d_w = smplh_opt.vertices\n T_w2c, K = data[\"T_w2c\"], data[\"K\"]\n verts_3d_c = apply_T_on_points(verts_3d_w, T_w2c[None])\n verts_2d = project_p2d(verts_3d_c, K[None])[0]\n min_2d = verts_2d.T.min(-1)[0]\n max_2d = verts_2d.T.max(-1)[0]\n bbx = torch.stack([min_2d, max_2d]).reshape(-1).numpy()\n return bbx\n\n\ndef get_2d(dataset, data):\n gender = data[\"meta\"][\"gender\"]\n smplh_params = {k: v.reshape(1, -1) for k, v in data[\"gt_smplh_params\"].items()}\n smplh_opt = dataset.smplh[gender](**smplh_params)\n joints_3d_w = smplh_opt.joints\n T_w2c, K = data[\"T_w2c\"], data[\"K\"]\n joints_3d_c = apply_T_on_points(joints_3d_w, T_w2c[None])\n joints_2d = project_p2d(joints_3d_c, K[None])[0]\n conf = torch.ones((73, 1))\n keypoints = torch.cat([joints_2d, conf], dim=1)\n return keypoints\n\n\ndef squared_crop_and_resize(dataset, img, bbx_lurb, dst_size=224, state=None):\n if state is not None:\n np.random.set_state(state)\n center_rand = dataset.BBX_CENTER * (np.random.random(2) * 2 - 1)\n center_x = (bbx_lurb[0] + bbx_lurb[2]) / 2 + center_rand[0]\n center_y = (bbx_lurb[1] + bbx_lurb[3]) / 2 + center_rand[1]\n ori_half_size = max(bbx_lurb[2] - bbx_lurb[0], bbx_lurb[3] - bbx_lurb[1]) / 2\n ori_half_size *= 1 + 0.15 + dataset.BBX_ZOOM * np.random.random() # zoom\n\n src = np.array(\n [\n [center_x - ori_half_size, center_y - ori_half_size],\n [center_x + ori_half_size, center_y - ori_half_size],\n [center_x, center_y],\n ],\n dtype=np.float32,\n )\n dst = np.array([[0, 0], [dst_size - 1, 0], [dst_size / 2 - 0.5, dst_size / 2 - 0.5]], dtype=np.float32)\n\n A = cv2.getAffineTransform(src, dst)\n img_crop = cv2.warpAffine(img, A, (dst_size, dst_size), flags=cv2.INTER_LINEAR)\n bbx_new = np.array(\n [center_x - ori_half_size, center_y - ori_half_size, center_x + ori_half_size, center_y + ori_half_size],\n dtype=bbx_lurb.dtype,\n )\n return img_crop, bbx_new, A\n\n\n# Augment bbx\ndef get_augmented_square_bbx(bbx_lurb, per_shift=0.1, per_zoomout=0.2, base_zoomout=0.15, state=None):\n \"\"\"\n Args:\n per_shift: in percent, maximum random shift\n per_zoomout: in percent, maximum random zoom\n \"\"\"\n if state is not None:\n np.random.set_state(state)\n maxsize_bbx = max(bbx_lurb[2] - bbx_lurb[0], bbx_lurb[3] - bbx_lurb[1])\n # shift of center\n shift = maxsize_bbx * per_shift * (np.random.random(2) * 2 - 1)\n center_x = (bbx_lurb[0] + bbx_lurb[2]) / 2 + shift[0]\n center_y = (bbx_lurb[1] + bbx_lurb[3]) / 2 + shift[1]\n # zoomout of half-size\n halfsize_bbx = maxsize_bbx / 2\n halfsize_bbx *= 1 + base_zoomout + per_zoomout * np.random.random()\n\n bbx_lurb = np.array(\n [\n center_x - halfsize_bbx,\n center_y - halfsize_bbx,\n center_x + halfsize_bbx,\n center_y + halfsize_bbx,\n ]\n )\n return bbx_lurb\n\n\ndef get_squared_bbx_region_and_resize(frames, bbx_xys, dst_size=224):\n \"\"\"\n Args:\n frames: (F, H, W, 3)\n bbx_xys: (F, 3), xys\n \"\"\"\n frames_np = frames.numpy() if isinstance(frames, torch.Tensor) else frames\n bbx_xys = bbx_xys if isinstance(bbx_xys, torch.Tensor) else torch.tensor(bbx_xys) # use tensor\n srcs = torch.stack(\n [\n torch.stack([bbx_xys[:, 0] - bbx_xys[:, 2] / 2, bbx_xys[:, 1] - bbx_xys[:, 2] / 2], dim=-1),\n torch.stack([bbx_xys[:, 0] + bbx_xys[:, 2] / 2, bbx_xys[:, 1] - bbx_xys[:, 2] / 2], dim=-1),\n bbx_xys[:, :2],\n ],\n dim=1,\n ) # (F, 3, 2)\n dst = np.array([[0, 0], [dst_size - 1, 0], [dst_size / 2 - 0.5, dst_size / 2 - 0.5]], dtype=np.float32)\n As = np.stack([cv2.getAffineTransform(src, dst) for src in srcs.numpy()])\n\n img_crops = np.stack(\n [cv2.warpAffine(frames_np[i], As[i], (dst_size, dst_size), flags=cv2.INTER_LINEAR) for i in range(len(As))]\n )\n img_crops = torch.from_numpy(img_crops)\n As = torch.from_numpy(As)\n return img_crops, As\n\n\n# ----- Camera utils ----- #\n\n\ndef extract_cam_xml(xml_path=\"\", dtype=torch.float32):\n import xml.etree.ElementTree as ET\n\n tree = ET.parse(xml_path)\n\n extrinsics_mat = [float(s) for s in tree.find(\"./CameraMatrix/data\").text.split()]\n intrinsics_mat = [float(s) for s in tree.find(\"./Intrinsics/data\").text.split()]\n distortion_vec = [float(s) for s in tree.find(\"./Distortion/data\").text.split()]\n\n return {\n \"ext_mat\": torch.tensor(extrinsics_mat).float(),\n \"int_mat\": torch.tensor(intrinsics_mat).float(),\n \"dis_vec\": torch.tensor(distortion_vec).float(),\n }\n\n\ndef get_cam2params(scene_info_root=None):\n \"\"\"\n Args:\n scene_info_root: this could be repalced by path to scan_calibration\n \"\"\"\n if scene_info_root is not None:\n cam_params = {}\n cam_xml_files = Path(scene_info_root).glob(\"*/calibration/*.xml\")\n for cam_xml_file in cam_xml_files:\n cam_param = extract_cam_xml(cam_xml_file)\n T_w2c = cam_param[\"ext_mat\"].reshape(3, 4)\n T_w2c = torch.cat([T_w2c, torch.tensor([[0, 0, 0, 1.0]])], dim=0) # (4, 4)\n K = cam_param[\"int_mat\"].reshape(3, 3)\n cap_name = cam_xml_file.parts[-3]\n cam_id = int(cam_xml_file.stem)\n cam_key = f\"{cap_name}_{cam_id}\"\n cam_params[cam_key] = (T_w2c, K)\n else:\n cam_params = torch.load(Path(__file__).parent / \"resource/cam2params.pt\")\n return cam_params\n\n\n# ----- Parse Raw Resource ----- #\n\n\ndef get_w2az_sahmr():\n \"\"\"\n Returns:\n w2az_sahmr: dict, {scan_name: Tw2az}, Tw2az is a tensor of (4,4)\n \"\"\"\n fn = Path(__file__).parent / \"resource/w2az_sahmr.json\"\n with open(fn, \"r\") as f:\n kvs = json.load(f).items()\n w2az_sahmr = {k: torch.tensor(v) for k, v in kvs}\n return w2az_sahmr\n\n\ndef has_multi_persons(seq_name):\n \"\"\"\n Args:\n seq_name: e.g. LectureHall_009_021_reparingprojector1\n \"\"\"\n return len(seq_name.split(\"_\")) != 3\n\n\ndef parse_seqname_info(skip_multi_persons=True):\n \"\"\"\n This function will skip multi-person sequences.\n Returns:\n sname_to_info: scan_name, subject_id, gender, cam_ids\n \"\"\"\n fns = [Path(__file__).parent / f\"resource/{split}.txt\" for split in [\"train\", \"val\", \"test\"]]\n # Train / Val&Test Header:\n # sequence_name\tcapture_name\tscan_name\tid\tmoving_cam\tgender\tview_id\n # sequence_name\tcapture_name\tscan_name\tid\tmoving_cam\tgender\tscene\taction/scene-interaction\tsubjects\tview_id\n sname_to_info = {}\n for fn in fns:\n with open(fn, \"r\") as f:\n for line in f.readlines()[1:]:\n raw_values = line.strip().split()\n seq_name = raw_values[0]\n if skip_multi_persons and has_multi_persons(seq_name):\n continue\n scan_name = f\"{raw_values[1]}_{raw_values[2]}\"\n subject_id = int(raw_values[3])\n gender = raw_values[5]\n cam_ids = [int(c) for c in raw_values[-1].split(\",\")]\n sname_to_info[seq_name] = (scan_name, subject_id, gender, cam_ids)\n return sname_to_info\n\n\ndef get_seqnames_of_split(splits=[\"train\"], skip_multi_persons=True):\n if not isinstance(splits, list):\n splits = [splits]\n fns = [Path(__file__).parent / f\"resource/{split}.txt\" for split in splits]\n seqnames = []\n for fn in fns:\n with open(fn, \"r\") as f:\n for line in f.readlines()[1:]:\n seq_name = line.strip().split()[0]\n if skip_multi_persons and has_multi_persons(seq_name):\n continue\n seqnames.append(seq_name)\n return seqnames\n\n\ndef get_seqname_to_imgrange():\n \"\"\"Each sequence has a different range of image ids.\"\"\"\n from tqdm import tqdm\n\n split_seqnames = {split: get_seqnames_of_split(split) for split in [\"train\", \"val\", \"test\"]}\n seqname_to_imgrange = {}\n for split in [\"train\", \"val\", \"test\"]:\n for seqname in tqdm(split_seqnames[split]):\n img_root = Path(\"inputs/RICH\") / \"images_ds4\" / split # compressed (not original)\n img_dir = img_root / seqname\n img_names = sorted([n.name for n in img_dir.glob(\"**/*.jpeg\")])\n if len(img_names) == 0:\n img_range = (0, 0)\n else:\n img_range = (int(img_names[0].split(\"_\")[0]), int(img_names[-1].split(\"_\")[0]))\n seqname_to_imgrange[seqname] = img_range\n return seqname_to_imgrange\n\n\n# ----- Compose keys ----- #\n\n\ndef get_img_key(seq_name, cam_id, f_id):\n assert len(seq_name.split(\"_\")) == 3\n subject_id = int(seq_name.split(\"_\")[1])\n return f\"{seq_name}_{int(cam_id)}_{int(f_id):05d}_{subject_id}\"\n\n\ndef get_seq_cam_fn(img_root, seq_name, cam_id):\n \"\"\"\n Args:\n img_root: \"inputs/RICH/images_ds4/train\"\n \"\"\"\n img_root = Path(img_root)\n cam_id = int(cam_id)\n return str(img_root / f\"{seq_name}/cam_{cam_id:02d}\")\n\n\ndef get_img_fn(img_root, seq_name, cam_id, f_id):\n \"\"\"\n Args:\n img_root: \"inputs/RICH/images_ds4/train\"\n \"\"\"\n img_root = Path(img_root)\n cam_id = int(cam_id)\n f_id = int(f_id)\n return str(img_root / f\"{seq_name}/cam_{cam_id:02d}\" / f\"{f_id:05d}_{cam_id:02d}.jpeg\")\n\n\n# ----- WHAM ----- #\n\n\ndef get_cam_key_wham_vid(vid):\n _, sname, cname = vid.split(\"/\")\n scene = sname.split(\"_\")[0]\n cid = int(cname.split(\"_\")[1])\n cam_key = f\"{scene}_{cid}\"\n return cam_key\n\n\ndef get_K_wham_vid(vid):\n cam_key = get_cam_key_wham_vid(vid)\n cam2params = get_cam2params()\n K = cam2params[cam_key][1]\n return K\n\n\nclass RichVid2Tc2az:\n def __init__(self) -> None:\n self.w2az = get_w2az_sahmr() # scan_name: tensor 4,4\n seqname_info = parse_seqname_info(skip_multi_persons=True) # {k: (scan_name, subject_id, gender, cam_ids)}\n self.seqname_to_scanname = {k: v[0] for k, v in seqname_info.items()}\n self.cam2params = get_cam2params() # cam_key -> (T_w2c, K)\n\n def __call__(self, vid):\n cam_key = get_cam_key_wham_vid(vid)\n scan_name = self.seqname_to_scanname[vid.split(\"/\")[1]]\n T_w2c, K = self.cam2params[cam_key] # (4, 4) (3, 3)\n T_w2az = self.w2az[scan_name]\n T_c2az = T_w2az @ T_w2c.inverse()\n return T_c2az\n\n def get_T_w2az(self, vid):\n cam_key = get_cam_key_wham_vid(vid)\n scan_name = self.seqname_to_scanname[vid.split(\"/\")[1]]\n T_w2az = self.w2az[scan_name]\n return T_w2az\n" }, { "path": "eval/GVHMR/hmr4d/dataset/threedpw/threedpw_motion_test.py", "content": "import torch\nfrom torch.utils import data\nfrom pathlib import Path\n\nfrom hmr4d.utils.pylogger import Log\nfrom hmr4d.utils.wis3d_utils import make_wis3d, add_motion_as_lines\nfrom hmr4d.utils.geo_transform import compute_cam_angvel\nfrom hmr4d.utils.geo.hmr_cam import estimate_K, resize_K\nfrom hmr4d.utils.geo.flip_utils import flip_kp2d_coco17\n\nfrom hmr4d.configs import MainStore, builds\n\nVID_HARD = []\n# VID_HARD = [\"downtown_bar_00_1\"]\n\n\nclass ThreedpwSmplFullSeqDataset(data.Dataset):\n def __init__(self, flip_test=False, skip_invalid=False):\n super().__init__()\n self.dataset_name = \"3DPW\"\n self.skip_invalid = skip_invalid\n Log.info(f\"[{self.dataset_name}] Full sequence\")\n\n # Load evaluation protocol from WHAM labels\n self.threedpw_dir = Path(\"inputs/3DPW/hmr4d_support\")\n # ['vname', 'K_fullimg', 'T_w2c', 'smpl_params', 'gender', 'mask_raw', 'mask_wham', 'img_wh']\n self.labels = torch.load(self.threedpw_dir / \"test_3dpw_gt_labels.pt\")\n self.vid2bbx = torch.load(self.threedpw_dir / \"preproc_test_bbx.pt\")\n self.vid2kp2d = torch.load(self.threedpw_dir / \"preproc_test_kp2d_v0.pt\")\n\n # Setup dataset index\n self.idx2meta = list(self.labels)\n if len(VID_HARD) > 0: # Pick subsets for fast testing\n self.idx2meta = VID_HARD\n Log.info(f\"[{self.dataset_name}] {len(self.idx2meta)} sequences.\")\n\n # If flip_test is enabled, we will return extra data for flipped test\n self.flip_test = flip_test\n if self.flip_test:\n Log.info(f\"[{self.dataset_name}] Flip test enabled\")\n\n def __len__(self):\n return len(self.idx2meta)\n\n def _load_data(self, idx):\n data = {}\n vid = self.idx2meta[idx]\n meta = {\"dataset_id\": self.dataset_name, \"vid\": vid}\n data.update({\"meta\": meta})\n\n # Add useful data\n label = self.labels[vid]\n mask = label[\"mask_wham\"]\n width_height = label[\"img_wh\"]\n data.update(\n {\n \"length\": len(mask), # F\n \"smpl_params\": label[\"smpl_params\"], # world\n \"gender\": label[\"gender\"], # str\n \"T_w2c\": label[\"T_w2c\"], # (F, 4, 4)\n \"mask\": mask, # (F)\n }\n )\n K_fullimg = label[\"K_fullimg\"] # (3, 3)\n if False:\n K_fullimg = estimate_K(*width_height)\n data[\"K_fullimg\"] = K_fullimg\n\n # Preprocessed: bbx, kp2d, image as feature\n bbx_xys = self.vid2bbx[vid][\"bbx_xys\"] # (F, 3)\n kp2d = self.vid2kp2d[vid] # (F, 17, 3)\n cam_angvel = compute_cam_angvel(data[\"T_w2c\"][:, :3, :3]) # (L, 6)\n data.update({\"bbx_xys\": bbx_xys, \"kp2d\": kp2d, \"cam_angvel\": cam_angvel})\n\n imgfeat_dir = self.threedpw_dir / \"imgfeats/3dpw_test\"\n f_img_dict = torch.load(imgfeat_dir / f\"{vid}.pt\")\n f_imgseq = f_img_dict[\"features\"].float()\n data[\"f_imgseq\"] = f_imgseq # (F, 1024)\n\n # to render a video\n vname = label[\"vname\"]\n video_path = self.threedpw_dir / f\"videos/{vname}.mp4\"\n frame_id = torch.where(mask)[0].long()\n ds = 0.5\n K_render = resize_K(K_fullimg, ds)\n bbx_xys_render = bbx_xys * ds\n kp2d_render = kp2d.clone()\n kp2d_render[..., :2] *= ds\n data[\"meta_render\"] = {\n \"name\": vid,\n \"video_path\": str(video_path),\n \"ds\": ds,\n \"frame_id\": frame_id,\n \"K\": K_render,\n \"bbx_xys\": bbx_xys_render,\n \"kp2d\": kp2d_render,\n }\n\n if self.flip_test:\n imgfeat_dir = self.threedpw_dir / \"imgfeats/3dpw_test_flip\"\n f_img_dict = torch.load(imgfeat_dir / f\"{vid}.pt\")\n flipped_bbx_xys = f_img_dict[\"bbx_xys\"].float() # (L, 3)\n flipped_features = f_img_dict[\"features\"].float() # (L, 1024)\n flipped_kp2d = flip_kp2d_coco17(kp2d, width_height[0]) # (L, 17, 3)\n\n R_flip_x = torch.tensor([[-1, 0, 0], [0, 1, 0], [0, 0, 1]]).float()\n flipped_R_w2c = R_flip_x @ data[\"T_w2c\"][:, :3, :3].clone()\n\n data_flip = {\n \"bbx_xys\": flipped_bbx_xys,\n \"f_imgseq\": flipped_features,\n \"kp2d\": flipped_kp2d,\n \"cam_angvel\": compute_cam_angvel(flipped_R_w2c),\n }\n data[\"flip_test\"] = data_flip\n return data\n\n def _process_data(self, data):\n length = data[\"length\"]\n data[\"K_fullimg\"] = data[\"K_fullimg\"][None].repeat(length, 1, 1)\n\n if self.skip_invalid: # Drop all invalid frames\n mask = data[\"mask\"].clone()\n data[\"length\"] = sum(mask)\n data[\"smpl_params\"] = {k: v[mask].clone() for k, v in data[\"smpl_params\"].items()}\n data[\"T_w2c\"] = data[\"T_w2c\"][mask].clone()\n data[\"mask\"] = data[\"mask\"][mask].clone()\n data[\"K_fullimg\"] = data[\"K_fullimg\"][mask].clone()\n data[\"bbx_xys\"] = data[\"bbx_xys\"][mask].clone()\n data[\"kp2d\"] = data[\"kp2d\"][mask].clone()\n data[\"cam_angvel\"] = data[\"cam_angvel\"][mask].clone()\n data[\"f_imgseq\"] = data[\"f_imgseq\"][mask].clone()\n data[\"flip_test\"] = {k: v[mask].clone() for k, v in data[\"flip_test\"].items()}\n\n return data\n\n def __getitem__(self, idx):\n data = self._load_data(idx)\n data = self._process_data(data)\n return data\n\n\n# 3DPW\nMainStore.store(\n name=\"fliptest\",\n node=builds(ThreedpwSmplFullSeqDataset, flip_test=True),\n group=\"test_datasets/3dpw\",\n)\nMainStore.store(\n name=\"v1\",\n node=builds(ThreedpwSmplFullSeqDataset, flip_test=False),\n group=\"test_datasets/3dpw\",\n)\n" }, { "path": "eval/GVHMR/hmr4d/dataset/threedpw/threedpw_motion_train.py", "content": "import torch\nfrom torch.utils import data\nfrom pathlib import Path\nimport numpy as np\n\nfrom hmr4d.utils.pylogger import Log\nfrom hmr4d.utils.wis3d_utils import make_wis3d, add_motion_as_lines\nfrom hmr4d.utils.geo_transform import compute_cam_angvel\nfrom hmr4d.utils.geo.hmr_cam import estimate_K, resize_K\nfrom hmr4d.utils.geo.flip_utils import flip_kp2d_coco17\nfrom hmr4d.dataset.imgfeat_motion.base_dataset import ImgfeatMotionDatasetBase\nfrom hmr4d.utils.net_utils import get_valid_mask, repeat_to_max_len, repeat_to_max_len_dict\nfrom hmr4d.utils.smplx_utils import make_smplx\nfrom hmr4d.utils.video_io_utils import get_video_lwh, read_video_np, save_video\nfrom hmr4d.utils.vis.renderer_utils import simple_render_mesh_background\n\nfrom hmr4d.configs import MainStore, builds\n\n\nclass ThreedpwSmplDataset(ImgfeatMotionDatasetBase):\n def __init__(self):\n # Path\n self.hmr4d_support_dir = Path(\"inputs/3DPW/hmr4d_support\")\n self.dataset_name = \"3DPW\"\n\n # Setting\n self.min_motion_frames = 60\n self.max_motion_frames = 120\n super().__init__()\n\n def _load_dataset(self):\n self.train_labels = torch.load(self.hmr4d_support_dir / \"train_3dpw_gt_labels.pt\")\n self.refit_smplx = torch.load(self.hmr4d_support_dir / \"train_refit_smplx.pt\")\n if True: # Remove clips that have obvious error\n update_list = {\n \"courtyard_basketball_00_1\": [(0, 300), (340, 468)],\n \"courtyard_laceShoe_00_0\": [(0, 620), (780, 931)],\n \"courtyard_rangeOfMotions_00_1\": [(0, 370), (410, 601)],\n \"courtyard_shakeHands_00_1\": [(0, 100), (120, 391)],\n }\n for k, v in update_list.items():\n self.refit_smplx[k][\"valid_range_list\"] = v\n\n self.f_img_folder = self.hmr4d_support_dir / \"imgfeats/3dpw_train_smplx_refit\"\n Log.info(f\"[{self.dataset_name}] Train\")\n\n def _get_idx2meta(self):\n # We expect to see the entire sequence during one epoch,\n # so each sequence will be sampled max(SeqLength // MotionFrames, 1) times\n seq_lengths = []\n self.idx2meta = []\n for vid in self.refit_smplx:\n valid_range_list = self.refit_smplx[vid][\"valid_range_list\"]\n for start, end in valid_range_list:\n seq_length = end - start\n num_samples = max(seq_length // self.max_motion_frames, 1)\n seq_lengths.append(seq_length)\n self.idx2meta.extend([(vid, start, end)] * num_samples)\n minutes = sum(seq_lengths) / 25 / 60\n Log.info(\n f\"[{self.dataset_name}] has {minutes:.1f} minutes motion -> Resampled to {len(self.idx2meta)} samples.\"\n )\n\n def _load_data(self, idx):\n data = {}\n vid, range1, range2 = self.idx2meta[idx]\n\n # Random select a subset\n mlength = range2 - range1\n min_motion_len = self.min_motion_frames\n max_motion_len = self.max_motion_frames\n\n if mlength < min_motion_len: # this may happen, the minimal mlength is around 30\n start = range1\n length = mlength\n else:\n effect_max_motion_len = min(max_motion_len, mlength)\n length = np.random.randint(min_motion_len, effect_max_motion_len + 1) # [low, high)\n start = np.random.randint(range1, range2 - length + 1)\n end = start + length\n data[\"length\"] = length\n data[\"meta\"] = {\"data_name\": self.dataset_name, \"idx\": idx, \"vid\": vid, \"start_end\": (start, end)}\n\n # Select motion subset\n data[\"smplx_params_incam\"] = {k: v[start:end] for k, v in self.refit_smplx[vid][\"smplx_params_incam\"].items()}\n data[\"K_fullimg\"] = self.train_labels[vid][\"K_fullimg\"]\n data[\"T_w2c\"] = self.train_labels[vid][\"T_w2c\"][start:end]\n\n # Img (as feature):\n f_img_dict = torch.load(self.f_img_folder / f\"{vid}.pt\")\n data[\"bbx_xys\"] = f_img_dict[\"bbx_xys\"][start:end] # (F, 3)\n data[\"f_imgseq\"] = f_img_dict[\"features\"][start:end].float() # (F, 3)\n data[\"img_wh\"] = f_img_dict[\"img_wh\"] # (2)\n data[\"kp2d\"] = torch.zeros((end - start), 17, 3) # (L, 17, 3) # do not provide kp2d\n\n return data\n\n def _process_data(self, data, idx):\n length = data[\"length\"]\n\n smpl_params_c = data[\"smplx_params_incam\"]\n smpl_params_w_zero = {k: torch.zeros_like(v) for k, v in smpl_params_c.items()}\n K_fullimg = data[\"K_fullimg\"][None].repeat(length, 1, 1)\n cam_angvel = compute_cam_angvel(data[\"T_w2c\"][:, :3, :3])\n\n max_len = self.max_motion_frames\n return_data = {\n \"meta\": data[\"meta\"],\n \"length\": length,\n \"smpl_params_c\": smpl_params_c,\n \"smpl_params_w\": smpl_params_w_zero,\n \"R_c2gv\": torch.zeros(length, 3, 3), # (F, 3, 3)\n \"gravity_vec\": torch.zeros(3), # (3)\n \"bbx_xys\": data[\"bbx_xys\"], # (F, 3)\n \"K_fullimg\": K_fullimg, # (F, 3, 3)\n \"f_imgseq\": data[\"f_imgseq\"], # (F, D)\n \"kp2d\": data[\"kp2d\"], # (F, 17, 3)\n \"cam_angvel\": cam_angvel, # (F, 6)\n \"mask\": {\n \"valid\": get_valid_mask(max_len, length),\n \"vitpose\": False,\n \"bbx_xys\": True,\n \"f_imgseq\": True,\n \"spv_incam_only\": True,\n },\n }\n\n if False: # Debug, render incam\n start, end = data[\"meta\"][\"start_end\"]\n vid = data[\"meta\"][\"vid\"]\n\n ds = 0.5\n faces = smplx.faces\n smplx = make_smplx(\"supermotion\")\n smplx_c_verts = smplx(**return_data[\"smpl_params_c\"]).vertices\n K_render = resize_K(K_fullimg, ds)\n\n video_path = self.hmr4d_support_dir / f\"videos/{vid[:-2]}.mp4\"\n images = read_video_np(video_path, scale=ds, start_frame=start, end_frame=end)\n\n render_dict = {\n \"K\": K_render[:1], # only support batch size 1\n \"faces\": faces,\n \"verts\": smplx_c_verts,\n \"background\": images,\n }\n img_overlay = simple_render_mesh_background(render_dict, VI=10)\n save_video(img_overlay, f\"tmp.mp4\", crf=28)\n\n # Batchable\n return_data[\"smpl_params_c\"] = repeat_to_max_len_dict(return_data[\"smpl_params_c\"], max_len)\n return_data[\"smpl_params_w\"] = repeat_to_max_len_dict(return_data[\"smpl_params_w\"], max_len)\n return_data[\"R_c2gv\"] = repeat_to_max_len(return_data[\"R_c2gv\"], max_len)\n return_data[\"bbx_xys\"] = repeat_to_max_len(return_data[\"bbx_xys\"], max_len)\n return_data[\"K_fullimg\"] = repeat_to_max_len(return_data[\"K_fullimg\"], max_len)\n return_data[\"f_imgseq\"] = repeat_to_max_len(return_data[\"f_imgseq\"], max_len)\n return_data[\"kp2d\"] = repeat_to_max_len(return_data[\"kp2d\"], max_len)\n return_data[\"cam_angvel\"] = repeat_to_max_len(return_data[\"cam_angvel\"], max_len)\n\n return return_data\n\n\n# 3DPW\nMainStore.store(name=\"v1\", node=builds(ThreedpwSmplDataset), group=\"train_datasets/imgfeat_3dpw\")\n" }, { "path": "eval/GVHMR/hmr4d/dataset/threedpw/utils.py", "content": "import json\nimport numpy as np\nfrom pathlib import Path\nfrom collections import defaultdict\nimport pickle\nimport torch\nimport joblib\n\nRESOURCE_FOLDER = Path(__file__).resolve().parent / \"resource\"\n\n\ndef read_raw_pkl(pkl_path):\n with open(pkl_path, \"rb\") as f:\n data = pickle.load(f, encoding=\"bytes\")\n\n num_subjects = len(data[b\"poses\"])\n F = data[b\"poses\"][0].shape[0]\n smpl_params = []\n for i in range(num_subjects):\n smpl_params.append(\n {\n \"body_pose\": torch.from_numpy(data[b\"poses\"][i][:, 3:72]).float(), # (F, 69)\n \"betas\": torch.from_numpy(data[b\"betas\"][i][:10]).repeat(F, 1).float(), # (F, 10)\n \"global_orient\": torch.from_numpy(data[b\"poses\"][i][:, :3]).float(), # (F, 3)\n \"transl\": torch.from_numpy(data[b\"trans\"][i]).float(), # (F, 3)\n }\n )\n genders = [\"male\" if g == \"m\" else \"female\" for g in data[b\"genders\"]]\n campose_valid = [torch.from_numpy(v).bool() for v in data[b\"campose_valid\"]]\n\n seq_name = data[b\"sequence\"]\n K_fullimg = torch.from_numpy(data[b\"cam_intrinsics\"]).float()\n T_w2c = torch.from_numpy(data[b\"cam_poses\"]).float()\n\n return_data = {\n \"sequence\": seq_name, # 'courtyard_bodyScannerMotions_00'\n \"K_fullimg\": K_fullimg, # (3, 3), not 55FoV\n \"T_w2c\": T_w2c, # (F, 4, 4)\n \"smpl_params\": smpl_params, # list of dict\n \"genders\": genders, # list of str\n \"campose_valid\": campose_valid, # list of bool-array\n # \"jointPositions\": data[b'jointPositions'], # SMPL, 24x3\n # \"poses2d\": data[b\"poses2d\"], # COCO, 3x18(?)\n }\n return return_data\n\n\ndef load_and_convert_wham_pth(pth):\n \"\"\"\n Convert to {vid: DataDict} style, Add smpl_params_incam\n \"\"\"\n # load\n wham_labels_raw = joblib.load(pth)\n # convert it to {vid: DataDict} style\n wham_labels = {}\n for i, vid in enumerate(wham_labels_raw[\"vid\"]):\n wham_labels[vid] = {k: wham_labels_raw[k][i] for k in wham_labels_raw}\n\n # convert pose and betas as smpl_params_incam (without transl)\n for vid in wham_labels:\n pose = wham_labels[vid][\"pose\"]\n global_orient = pose[:, :3] # (F, 3)\n body_pose = pose[:, 3:] # (F, 69)\n betas = wham_labels[vid][\"betas\"] # (F, 10), all frames are the same\n wham_labels[vid][\"smpl_params_incam\"] = {\n \"body_pose\": body_pose.float(), # (F, 69)\n \"betas\": betas.float(), # (F, 10)\n \"global_orient\": global_orient.float(), # (F, 3)\n }\n\n return wham_labels\n\n\n# Neural-Annot utils\n\n\ndef na_cam_param_to_K_fullimg(cam_param):\n K = torch.eye(3)\n K[[0, 1], [0, 1]] = torch.tensor(cam_param[\"focal\"])\n K[[0, 1], [2, 2]] = torch.tensor(cam_param[\"princpt\"])\n return K\n" }, { "path": "eval/GVHMR/hmr4d/model/common_utils/optimizer.py", "content": "from torch.optim import AdamW, Adam\nfrom hmr4d.configs import MainStore, builds\n\n\noptimizer_cfgs = {\n \"adam_1e-3\": builds(Adam, lr=1e-3, zen_partial=True),\n \"adam_2e-4\": builds(Adam, lr=2e-4, zen_partial=True),\n \"adamw_2e-4\": builds(AdamW, lr=2e-4, zen_partial=True),\n \"adamw_1e-4\": builds(AdamW, lr=1e-4, zen_partial=True),\n \"adamw_5e-5\": builds(AdamW, lr=5e-5, zen_partial=True),\n \"adamw_1e-5\": builds(AdamW, lr=1e-5, zen_partial=True),\n # zero-shot text-to-image generation\n \"adamw_1e-3_dalle\": builds(AdamW, lr=1e-3, weight_decay=1e-4, zen_partial=True),\n}\n\nfor name, cfg in optimizer_cfgs.items():\n MainStore.store(name=name, node=cfg, group=f\"optimizer\")\n" }, { "path": "eval/GVHMR/hmr4d/model/common_utils/scheduler.py", "content": "import torch\nfrom bisect import bisect_right\n\n\nclass WarmupMultiStepLR(torch.optim.lr_scheduler.LRScheduler):\n def __init__(self, optimizer, milestones, warmup=0, gamma=0.1, last_epoch=-1, verbose=\"deprecated\"):\n \"\"\"Assume optimizer does not change lr; Scheduler is called epoch-based\"\"\"\n self.milestones = milestones\n self.warmup = warmup\n assert warmup < milestones[0]\n self.gamma = gamma\n super().__init__(optimizer, last_epoch, verbose)\n\n def get_lr(self):\n base_lrs = self.base_lrs # base lr for each groups\n n_groups = len(base_lrs)\n comming_epoch = self.last_epoch # the lr will be set for the comming epoch, starts from 0\n\n # add extra warmup\n if comming_epoch < self.warmup:\n # e.g. comming_epoch [0, 1, 2] for warmup == 3\n # lr should be base_lr * (last_epoch+1) / (warmup + 1), e.g. [0.25, 0.5, 0.75] * base_lr\n lr_factor = (self.last_epoch + 1) / (self.warmup + 1)\n return [base_lrs[i] * lr_factor for i in range(n_groups)]\n else:\n # bisect_right([3,5,7], 0) -> 0; bisect_right([3,5,7], 5) -> 2\n p = bisect_right(self.milestones, comming_epoch)\n lr_factor = self.gamma**p\n return [base_lrs[i] * lr_factor for i in range(n_groups)]\n" }, { "path": "eval/GVHMR/hmr4d/model/common_utils/scheduler_cfg.py", "content": "from omegaconf import DictConfig, ListConfig\nfrom hmr4d.configs import MainStore, builds\n\n# do not perform scheduling\ndefault = DictConfig({\"scheduler\": None})\nMainStore.store(name=\"default\", node=default, group=f\"scheduler_cfg\")\n\n\n# epoch-based\ndef epoch_half_by(milestones=[100, 200, 300]):\n return DictConfig(\n {\n \"scheduler\": {\n \"_target_\": \"torch.optim.lr_scheduler.MultiStepLR\",\n \"milestones\": milestones,\n \"gamma\": 0.5,\n },\n \"interval\": \"epoch\",\n \"frequency\": 1,\n }\n )\n\n\nMainStore.store(name=\"epoch_half_100_200_300\", node=epoch_half_by([100, 200, 300]), group=f\"scheduler_cfg\")\nMainStore.store(name=\"epoch_half_100_200\", node=epoch_half_by([100, 200]), group=f\"scheduler_cfg\")\nMainStore.store(name=\"epoch_half_200_350\", node=epoch_half_by([200, 350]), group=f\"scheduler_cfg\")\nMainStore.store(name=\"epoch_half_300\", node=epoch_half_by([300]), group=f\"scheduler_cfg\")\n\n\n# epoch-based\ndef warmup_epoch_half_by(warmup=10, milestones=[100, 200, 300]):\n return DictConfig(\n {\n \"scheduler\": {\n \"_target_\": \"hmr4d.model.common_utils.scheduler.WarmupMultiStepLR\",\n \"milestones\": milestones,\n \"warmup\": warmup,\n \"gamma\": 0.5,\n },\n \"interval\": \"epoch\",\n \"frequency\": 1,\n }\n )\n\n\nMainStore.store(name=\"warmup_5_epoch_half_200_350\", node=warmup_epoch_half_by(5, [200, 350]), group=f\"scheduler_cfg\")\nMainStore.store(name=\"warmup_10_epoch_half_200_350\", node=warmup_epoch_half_by(10, [200, 350]), group=f\"scheduler_cfg\")\n" }, { "path": "eval/GVHMR/hmr4d/model/gvhmr/callbacks/metric_3dpw.py", "content": "import torch\nimport pytorch_lightning as pl\nimport numpy as np\nfrom pathlib import Path\nfrom einops import einsum, rearrange\n\nfrom hmr4d.configs import MainStore, builds\nfrom hmr4d.utils.pylogger import Log\nfrom hmr4d.utils.comm.gather import all_gather\nfrom hmr4d.utils.eval.eval_utils import compute_camcoord_metrics, as_np_array\nfrom hmr4d.utils.smplx_utils import make_smplx\nfrom hmr4d.utils.vis.cv2_utils import cv2, draw_bbx_xys_on_image_batch, draw_coco17_skeleton_batch\nfrom hmr4d.utils.vis.renderer_utils import simple_render_mesh_background\nfrom hmr4d.utils.video_io_utils import read_video_np, get_video_lwh, save_video\nfrom hmr4d.utils.geo_transform import apply_T_on_points\nfrom hmr4d.utils.seq_utils import rearrange_by_mask\n\n\nclass MetricMocap(pl.Callback):\n def __init__(self):\n super().__init__()\n # vid->result\n self.metric_aggregator = {\n \"pa_mpjpe\": {},\n \"mpjpe\": {},\n \"pve\": {},\n \"accel\": {},\n }\n\n # SMPLX and SMPL\n self.smplx = make_smplx(\"supermotion_EVAL3DPW\")\n self.smpl = {\"male\": make_smplx(\"smpl\", gender=\"male\"), \"female\": make_smplx(\"smpl\", gender=\"female\")}\n self.J_regressor = torch.load(\"hmr4d/utils/body_model/smpl_3dpw14_J_regressor_sparse.pt\").to_dense()\n self.J_regressor24 = torch.load(\"hmr4d/utils/body_model/smpl_neutral_J_regressor.pt\")\n self.smplx2smpl = torch.load(\"hmr4d/utils/body_model/smplx2smpl_sparse.pt\")\n self.faces_smplx = self.smplx.faces\n self.faces_smpl = self.smpl[\"male\"].faces\n\n # The metrics are calculated similarly for val/test/predict\n self.on_test_batch_end = self.on_validation_batch_end = self.on_predict_batch_end\n\n # Only validation record the metrics with logger\n self.on_test_epoch_end = self.on_validation_epoch_end = self.on_predict_epoch_end\n\n # ================== Batch-based Computation ================== #\n def on_predict_batch_end(self, trainer, pl_module, outputs, batch, batch_idx, dataloader_idx=0):\n \"\"\"The behaviour is the same for val/test/predict\"\"\"\n assert batch[\"B\"] == 1\n dataset_id = batch[\"meta\"][0][\"dataset_id\"]\n if dataset_id != \"3DPW\":\n return\n\n # Move to cuda if not\n self.smplx = self.smplx.cuda()\n for g in [\"male\", \"female\"]:\n self.smpl[g] = self.smpl[g].cuda()\n self.J_regressor = self.J_regressor.cuda()\n self.J_regressor24 = self.J_regressor24.cuda()\n self.smplx2smpl = self.smplx2smpl.cuda()\n\n vid = batch[\"meta\"][0][\"vid\"]\n seq_length = batch[\"length\"][0].item()\n gender = batch[\"gender\"][0]\n T_w2c = batch[\"T_w2c\"][0]\n mask = batch[\"mask\"][0]\n\n # Groundtruth (cam)\n target_w_params = {k: v[0] for k, v in batch[\"smpl_params\"].items()}\n target_w_output = self.smpl[gender](**target_w_params)\n target_w_verts = target_w_output.vertices\n target_c_verts = apply_T_on_points(target_w_verts, T_w2c)\n target_c_j3d = torch.matmul(self.J_regressor, target_c_verts)\n\n # + Prediction -> Metric\n smpl_out = self.smplx(**outputs[\"pred_smpl_params_incam\"])\n pred_c_verts = torch.stack([torch.matmul(self.smplx2smpl, v_) for v_ in smpl_out.vertices])\n pred_c_j3d = einsum(self.J_regressor, pred_c_verts, \"j v, l v i -> l j i\")\n del smpl_out # Prevent OOM\n\n # Metric of current sequence\n batch_eval = {\n \"pred_j3d\": pred_c_j3d,\n \"target_j3d\": target_c_j3d,\n \"pred_verts\": pred_c_verts,\n \"target_verts\": target_c_verts,\n }\n camcoord_metrics = compute_camcoord_metrics(batch_eval, mask=mask, pelvis_idxs=[2, 3])\n for k in camcoord_metrics:\n self.metric_aggregator[k][vid] = as_np_array(camcoord_metrics[k])\n\n if False: # Render incam (simple)\n meta_render = batch[\"meta_render\"][0]\n images = read_video_np(meta_render[\"video_path\"], scale=meta_render[\"ds\"])\n render_dict = {\n \"K\": meta_render[\"K\"][None], # only support batch size 1\n \"faces\": self.smpl[\"male\"].faces,\n \"verts\": pred_c_verts,\n \"background\": images,\n }\n img_overlay = simple_render_mesh_background(render_dict)\n output_fn = Path(\"outputs/3DPW_render_pred_flip\") / f\"{vid}.mp4\"\n save_video(img_overlay, output_fn, crf=28)\n\n if False: # Render incam (with details)\n meta_render = batch[\"meta_render\"][0]\n images = read_video_np(meta_render[\"video_path\"], scale=meta_render[\"ds\"])\n render_dict = {\n \"K\": meta_render[\"K\"][None], # only support batch size 1\n \"faces\": self.smpl[\"male\"].faces,\n \"verts\": pred_c_verts,\n \"background\": images,\n }\n img_overlay = simple_render_mesh_background(render_dict)\n\n # Add COCO17 and bbx to image\n bbx_xys_render = meta_render[\"bbx_xys\"]\n kp2d_render = meta_render[\"kp2d\"]\n img_overlay = draw_coco17_skeleton_batch(img_overlay, kp2d_render, conf_thr=0.5)\n img_overlay = draw_bbx_xys_on_image_batch(bbx_xys_render, img_overlay, mask)\n\n # Add metric\n metric_all = rearrange_by_mask(torch.tensor(camcoord_metrics[\"pa_mpjpe\"]), mask)\n for i in range(len(img_overlay)):\n m = metric_all[i]\n if m == 0: # a not evaluated frame\n continue\n text = f\"PA-MPJPE: {m:.1f}\"\n color = (244, 10, 20) if m > 45 else (0, 205, 0) # red or green\n cv2.putText(img_overlay[i], text, (10, 30), cv2.FONT_HERSHEY_SIMPLEX, 1, color, 2)\n\n output_dir = Path(\"tmp_pred_details\")\n output_dir.mkdir(exist_ok=True, parents=True)\n save_video(img_overlay, output_dir / f\"{vid}.mp4\", crf=24)\n\n # ================== Epoch Summary ================== #\n def on_predict_epoch_end(self, trainer, pl_module):\n \"\"\"Without logger\"\"\"\n local_rank, world_size = trainer.local_rank, trainer.world_size\n monitor_metric = \"pa_mpjpe\"\n\n # Reduce metric_aggregator across all processes\n metric_keys = list(self.metric_aggregator.keys())\n with torch.inference_mode(False): # allow in-place operation of all_gather\n metric_aggregator_gathered = all_gather(self.metric_aggregator) # list of dict\n for metric_key in metric_keys:\n for d in metric_aggregator_gathered:\n self.metric_aggregator[metric_key].update(d[metric_key])\n\n if False: # debug to make sure the all_gather is correct\n print(f\"[RANK {local_rank}/{world_size}]: {self.metric_aggregator[monitor_metric].keys()}\")\n\n total = len(self.metric_aggregator[monitor_metric])\n Log.info(f\"{total} sequences evaluated in {self.__class__.__name__}\")\n if total == 0:\n return\n\n # print monitored metric per sequence\n mm_per_seq = {k: v.mean() for k, v in self.metric_aggregator[monitor_metric].items()}\n if len(mm_per_seq) > 0:\n sorted_mm_per_seq = sorted(mm_per_seq.items(), key=lambda x: x[1], reverse=True)\n n_worst = 5 if trainer.state.stage == \"validate\" else len(sorted_mm_per_seq)\n if local_rank == 0:\n Log.info(\n f\"monitored metric {monitor_metric} per sequence\\n\"\n + \"\\n\".join([f\"{m:5.1f} : {s}\" for s, m in sorted_mm_per_seq[:n_worst]])\n + \"\\n------\"\n )\n\n # average over all batches\n metrics_avg = {k: np.concatenate(list(v.values())).mean() for k, v in self.metric_aggregator.items()}\n if local_rank == 0:\n Log.info(f\"[Metrics] 3DPW:\\n\" + \"\\n\".join(f\"{k}: {v:.1f}\" for k, v in metrics_avg.items()) + \"\\n------\")\n\n # save to logger if available\n if pl_module.logger is not None:\n cur_epoch = pl_module.current_epoch\n for k, v in metrics_avg.items():\n pl_module.logger.log_metrics({f\"val_metric_3DPW/{k}\": v}, step=cur_epoch)\n\n # reset\n for k in self.metric_aggregator:\n self.metric_aggregator[k] = {}\n\n\nnode_3dpw = builds(MetricMocap)\nMainStore.store(name=\"metric_3dpw\", node=node_3dpw, group=\"callbacks\", package=\"callbacks.metric_3dpw\")\n" }, { "path": "eval/GVHMR/hmr4d/model/gvhmr/callbacks/metric_emdb.py", "content": "import torch\nimport torch.nn.functional as F\nimport pytorch_lightning as pl\nfrom pytorch_lightning.utilities import rank_zero_only\nfrom hmr4d.configs import MainStore, builds\n\nfrom hmr4d.utils.comm.gather import all_gather\nfrom hmr4d.utils.pylogger import Log\n\nfrom hmr4d.utils.eval.eval_utils import (\n compute_camcoord_metrics,\n compute_global_metrics,\n compute_camcoord_perjoint_metrics,\n rearrange_by_mask,\n as_np_array,\n)\nfrom hmr4d.utils.geo_transform import apply_T_on_points, compute_T_ayfz2ay\nfrom hmr4d.utils.smplx_utils import make_smplx\nfrom einops import einsum, rearrange\n\nfrom hmr4d.utils.wis3d_utils import make_wis3d, add_motion_as_lines\nfrom hmr4d.utils.vis.renderer import Renderer, get_global_cameras_static\nfrom hmr4d.utils.geo.hmr_cam import estimate_focal_length\nfrom hmr4d.utils.video_io_utils import read_video_np, save_video\nimport imageio\nfrom tqdm import tqdm\nfrom pathlib import Path\nimport numpy as np\nimport cv2\n\n\nclass MetricMocap(pl.Callback):\n def __init__(self, emdb_split=1):\n \"\"\"\n Args:\n emdb_split: 1 to evaluate incam, 2 to evaluate global\n \"\"\"\n super().__init__()\n # vid->result\n if emdb_split == 1:\n self.target_dataset_id = \"EMDB_1\"\n self.metric_aggregator = {\n \"pa_mpjpe\": {},\n \"mpjpe\": {},\n \"pve\": {},\n \"accel\": {},\n }\n elif emdb_split == 2:\n self.target_dataset_id = \"EMDB_2\"\n self.metric_aggregator = {\n \"wa2_mpjpe\": {},\n \"waa_mpjpe\": {},\n \"rte\": {},\n \"jitter\": {},\n \"fs\": {},\n }\n else:\n raise ValueError(f\"Unknown emdb_split: {emdb_split}\")\n\n # SMPL\n self.smplx = make_smplx(\"supermotion\")\n self.smpl_model = {\"male\": make_smplx(\"smpl\", gender=\"male\"), \"female\": make_smplx(\"smpl\", gender=\"female\")}\n\n self.J_regressor = torch.load(\"hmr4d/utils/body_model/smpl_neutral_J_regressor.pt\")\n self.smplx2smpl = torch.load(\"hmr4d/utils/body_model/smplx2smpl_sparse.pt\")\n self.faces_smpl = self.smpl_model[\"male\"].faces\n self.faces_smplx = self.smplx.faces\n\n # The metrics are calculated similarly for val/test/predict\n self.on_test_batch_end = self.on_validation_batch_end = self.on_predict_batch_end\n\n # Only validation record the metrics with logger\n self.on_test_epoch_end = self.on_validation_epoch_end = self.on_predict_epoch_end\n\n # ================== Batch-based Computation ================== #\n def on_predict_batch_end(self, trainer, pl_module, outputs, batch, batch_idx, dataloader_idx=0):\n \"\"\"The behaviour is the same for val/test/predict\"\"\"\n assert batch[\"B\"] == 1\n dataset_id = batch[\"meta\"][0][\"dataset_id\"]\n if dataset_id != self.target_dataset_id:\n return\n\n # Move to cuda if not\n self.smplx = self.smplx.cuda()\n for g in [\"male\", \"female\"]:\n self.smpl_model[g] = self.smpl_model[g].cuda()\n self.J_regressor = self.J_regressor.cuda()\n self.smplx2smpl = self.smplx2smpl.cuda()\n\n vid = batch[\"meta\"][0][\"vid\"]\n seq_length = batch[\"length\"][0].item()\n gender = batch[\"gender\"][0]\n T_w2c = batch[\"T_w2c\"][0]\n mask = batch[\"mask\"][0]\n\n # Groundtruth (world, cam)\n target_w_params = {k: v[0] for k, v in batch[\"smpl_params\"].items()}\n target_w_output = self.smpl_model[gender](**target_w_params)\n target_w_verts = target_w_output.vertices\n target_w_j3d = torch.matmul(self.J_regressor, target_w_verts)\n target_c_verts = apply_T_on_points(target_w_verts, T_w2c)\n target_c_j3d = apply_T_on_points(target_w_j3d, T_w2c)\n\n # + Prediction -> Metric\n if self.target_dataset_id == \"EMDB_1\": # in camera metrics\n # 1. cam\n pred_smpl_params_incam = outputs[\"pred_smpl_params_incam\"]\n smpl_out = self.smplx(**pred_smpl_params_incam)\n pred_c_verts = torch.stack([torch.matmul(self.smplx2smpl, v_) for v_ in smpl_out.vertices])\n pred_c_j3d = einsum(self.J_regressor, pred_c_verts, \"j v, l v i -> l j i\")\n del smpl_out # Prevent OOM\n\n batch_eval = {\n \"pred_j3d\": pred_c_j3d,\n \"target_j3d\": target_c_j3d,\n \"pred_verts\": pred_c_verts,\n \"target_verts\": target_c_verts,\n }\n camcoord_metrics = compute_camcoord_metrics(batch_eval, mask=mask)\n for k in camcoord_metrics:\n self.metric_aggregator[k][vid] = as_np_array(camcoord_metrics[k])\n\n elif self.target_dataset_id == \"EMDB_2\": # global metrics\n # 2. global (align-y axis)\n pred_smpl_params_global = outputs[\"pred_smpl_params_global\"]\n smpl_out = self.smplx(**pred_smpl_params_global)\n pred_ay_verts = torch.stack([torch.matmul(self.smplx2smpl, v_) for v_ in smpl_out.vertices])\n pred_ay_j3d = einsum(self.J_regressor, pred_ay_verts, \"j v, l v i -> l j i\")\n del smpl_out # Prevent OOM\n\n batch_eval = {\n \"pred_j3d_glob\": pred_ay_j3d,\n \"target_j3d_glob\": target_w_j3d,\n \"pred_verts_glob\": pred_ay_verts,\n \"target_verts_glob\": target_w_verts,\n }\n global_metrics = compute_global_metrics(batch_eval, mask=mask)\n for k in global_metrics:\n self.metric_aggregator[k][vid] = as_np_array(global_metrics[k])\n\n if False: # wis3d debug\n wis3d = make_wis3d(name=\"debug-emdb-incam\")\n pred_cr_j3d = pred_c_j3d - pred_c_j3d[:, [0]] # (L, J, 3)\n target_cr_j3d = target_c_j3d - target_c_j3d[:, [0]] # (L, J, 3)\n add_motion_as_lines(pred_cr_j3d, wis3d, name=\"pred_cr_j3d\", const_color=\"blue\")\n add_motion_as_lines(target_cr_j3d, wis3d, name=\"target_cr_j3d\", const_color=\"green\")\n\n if False: # Dump wis3d\n vid = batch[\"meta\"][0][\"vid\"]\n split = batch[\"meta_render\"][0][\"split\"]\n wis3d = make_wis3d(name=f\"dump_emdb{split}-{vid}\")\n R_cam_type = batch[\"meta_render\"][0][\"R_cam_type\"]\n\n pred_cr_j3d = pred_c_j3d - pred_c_j3d[:, [0]] # (L, J, 3)\n target_cr_j3d = target_c_j3d - target_c_j3d[:, [0]] # (L, J, 3)\n add_motion_as_lines(pred_cr_j3d, wis3d, name=\"pred_cr_j3d\", const_color=\"blue\")\n add_motion_as_lines(target_cr_j3d, wis3d, name=\"target_cr_j3d\", const_color=\"green\")\n add_motion_as_lines(pred_ay_j3d, wis3d, name=f\"pred_ay_j3d@{R_cam_type}\")\n # add_motion_as_lines(target_w_j3d, wis3d, name=\"target_ay_j3d\")\n\n if False: # Render incam\n # -- rendering code -- #\n vname = batch[\"meta_render\"][0][\"name\"]\n video_path = batch[\"meta_render\"][0][\"video_path\"]\n width, height = batch[\"meta_render\"][0][\"width_height\"]\n K = batch[\"meta_render\"][0][\"K\"]\n faces = self.faces_smpl\n split = batch[\"meta_render\"][0][\"split\"]\n\n out_fn = f\"outputs/dump_render_emdb{split}/{vname}.mp4\"\n Path(out_fn).parent.mkdir(exist_ok=True, parents=True)\n\n # renderer\n renderer = Renderer(width, height, device=\"cuda\", faces=faces, K=K)\n # not skipping invalid frames\n resize_factor = 0.25\n images = read_video_np(video_path, scale=resize_factor) # (F, H, W, 3), uint8, numpy\n frame_id = batch[\"meta_render\"][0][\"frame_id\"]\n bbx_xys_render = batch[\"meta_render\"][0][\"bbx_xys\"]\n metric_vis = rearrange_by_mask(torch.from_numpy(self.metric_aggregator[\"mpjpe\"][vid]), mask)\n\n # -- render mesh -- #\n verts_incam = pred_c_verts\n output_images = []\n for i in tqdm(range(len(images)), desc=f\"Rendering {vname}\"):\n img = renderer.render_mesh(verts_incam[i].cuda(), images[i], [0.8, 0.8, 0.8])\n # bbx\n bbx_xys_ = bbx_xys_render[i].cpu().numpy()\n lu_point = (bbx_xys_[:2] - bbx_xys_[2:] / 2).astype(int)\n rd_point = (bbx_xys_[:2] + bbx_xys_[2:] / 2).astype(int)\n img = cv2.rectangle(img, lu_point, rd_point, (255, 178, 102), 2)\n\n if metric_vis[i] > 0:\n text = f\"pred mpjpe: {metric_vis[i]:.1f}\"\n text_color = (244, 10, 20) if metric_vis[i] > 80 else (0, 205, 0) # red or green\n cv2.putText(img, text, (10, 30), cv2.FONT_HERSHEY_SIMPLEX, 0.75, text_color, 2)\n\n output_images.append(img)\n save_video(output_images, out_fn, quality=5)\n\n if False: # Visualize incam + global results\n\n def move_to_start_point_face_z(verts):\n \"XZ to origin, Start from the ground, Face-Z\"\n verts = verts.clone() # (L, V, 3)\n xz_mean = verts[0].mean(0)[[0, 2]]\n y_min = verts[0, :, [1]].min()\n offset = torch.tensor([[[xz_mean[0], y_min, xz_mean[1]]]]).to(verts)\n verts = verts - offset\n\n T_ay2ayfz = compute_T_ayfz2ay(einsum(self.J_regressor, verts[[0]], \"j v, l v i -> l j i\"), inverse=True)\n verts = apply_T_on_points(verts, T_ay2ayfz)\n return verts\n\n verts_incam = pred_c_verts.clone()\n # verts_glob = move_to_start_point_face_z(target_ay_verts) # gt\n verts_glob = move_to_start_point_face_z(pred_ay_verts)\n global_R, global_T, global_lights = get_global_cameras_static(verts_glob.cpu())\n\n # -- rendering code (global version FOV=55) -- #\n vname = batch[\"meta_render\"][0][\"name\"]\n width, height = batch[\"meta_render\"][0][\"width_height\"]\n K = batch[\"meta_render\"][0][\"K\"]\n faces = self.faces_smpl\n out_fn = f\"outputs/dump_render_global/{vname}.mp4\"\n Path(out_fn).parent.mkdir(exist_ok=True, parents=True)\n writer = imageio.get_writer(out_fn, fps=30, mode=\"I\", format=\"FFMPEG\", macro_block_size=1)\n\n # two renderers\n renderer_incam = Renderer(width, height, device=\"cuda\", faces=faces, K=K)\n renderer_glob = Renderer(width, height, estimate_focal_length(width, height), device=\"cuda\", faces=faces)\n\n # imgs\n video_path = batch[\"meta_render\"][0][\"video_path\"]\n frame_id = batch[\"meta_render\"][0][\"frame_id\"].cpu().numpy()\n images = read_video_np(video_path, frame_id=frame_id) # (F, H/4, W/4, 3), uint8, numpy\n\n # Actual rendering\n cx, cz = (verts_glob.mean(1).max(0)[0] + verts_glob.mean(1).min(0)[0])[[0, 2]] / 2.0\n scale = (verts_glob.mean(1).max(0)[0] - verts_glob.mean(1).min(0)[0])[[0, 2]].max() * 1.5\n renderer_glob.set_ground(scale, cx.item(), cz.item())\n color = torch.ones(3).float().cuda() * 0.8\n\n for i in tqdm(range(seq_length), desc=f\"Rendering {vname}\"):\n # incam\n img_overlay_pred = renderer_incam.render_mesh(verts_incam[i].cuda(), images[i], [0.8, 0.8, 0.8])\n if batch[\"meta_render\"][0].get(\"bbx_xys\", None) is not None: # draw bbox lines\n bbx_xys = batch[\"meta_render\"][0][\"bbx_xys\"][i].cpu().numpy()\n lu_point = (bbx_xys[:2] - bbx_xys[2:] / 2).astype(int)\n rd_point = (bbx_xys[:2] + bbx_xys[2:] / 2).astype(int)\n img_overlay_pred = cv2.rectangle(img_overlay_pred, lu_point, rd_point, (255, 178, 102), 2)\n pred_mpjpe_ = self.metric_aggregator[\"mpjpe\"][vid][i]\n text = f\"pred mpjpe: {pred_mpjpe_:.1f}\"\n cv2.putText(img_overlay_pred, text, (10, 30), cv2.FONT_HERSHEY_SIMPLEX, 1, (200, 100, 200), 2)\n\n # glob\n cameras = renderer_glob.create_camera(global_R[i], global_T[i])\n img_glob = renderer_glob.render_with_ground(verts_glob[[i]], color[None], cameras, global_lights)\n\n # write\n img = np.concatenate([img_overlay_pred, img_glob], axis=1)\n writer.append_data(img)\n writer.close()\n pass\n\n # ================== Epoch Summary ================== #\n def on_predict_epoch_end(self, trainer, pl_module):\n \"\"\"Without logger\"\"\"\n local_rank, world_size = trainer.local_rank, trainer.world_size\n if \"mpjpe\" in self.metric_aggregator:\n monitor_metric = \"mpjpe\"\n else:\n monitor_metric = list(self.metric_aggregator.keys())[0]\n\n # Reduce metric_aggregator across all processes\n metric_keys = list(self.metric_aggregator.keys())\n with torch.inference_mode(False): # allow in-place operation of all_gather\n metric_aggregator_gathered = all_gather(self.metric_aggregator) # list of dict\n for metric_key in metric_keys:\n for d in metric_aggregator_gathered:\n self.metric_aggregator[metric_key].update(d[metric_key])\n\n total = len(self.metric_aggregator[monitor_metric])\n Log.info(f\"{total} sequences evaluated in {self.__class__.__name__}\")\n if total == 0:\n return\n\n # print monitored metric per sequence\n mm_per_seq = {k: v.mean() for k, v in self.metric_aggregator[monitor_metric].items()}\n if len(mm_per_seq) > 0:\n sorted_mm_per_seq = sorted(mm_per_seq.items(), key=lambda x: x[1], reverse=True)\n n_worst = 5 if trainer.state.stage == \"validate\" else len(sorted_mm_per_seq)\n if local_rank == 0:\n Log.info(\n f\"monitored metric {monitor_metric} per sequence\\n\"\n + \"\\n\".join([f\"{m:5.1f} : {s}\" for s, m in sorted_mm_per_seq[:n_worst]])\n + \"\\n------\"\n )\n\n # average over all batches\n metrics_avg = {k: np.concatenate(list(v.values())).mean() for k, v in self.metric_aggregator.items()}\n if local_rank == 0:\n Log.info(\n f\"[Metrics] {self.target_dataset_id}:\\n\"\n + \"\\n\".join(f\"{k}: {v:.1f}\" for k, v in metrics_avg.items())\n + \"\\n------\"\n )\n\n # save to logger if available\n if pl_module.logger is not None:\n cur_epoch = pl_module.current_epoch\n for k, v in metrics_avg.items():\n pl_module.logger.log_metrics({f\"val_metric_{self.target_dataset_id}/{k}\": v}, step=cur_epoch)\n\n # reset\n for k in self.metric_aggregator:\n self.metric_aggregator[k] = {}\n\n\nemdb1_node = builds(MetricMocap, emdb_split=1)\nemdb2_node = builds(MetricMocap, emdb_split=2)\nMainStore.store(name=\"metric_emdb1\", node=emdb1_node, group=\"callbacks\", package=\"callbacks.metric_emdb1\")\nMainStore.store(name=\"metric_emdb2\", node=emdb2_node, group=\"callbacks\", package=\"callbacks.metric_emdb2\")\n" }, { "path": "eval/GVHMR/hmr4d/model/gvhmr/callbacks/metric_rich.py", "content": "import torch\nimport torch.nn.functional as F\nimport pytorch_lightning as pl\nfrom pytorch_lightning.utilities import rank_zero_only\nfrom hmr4d.configs import MainStore, builds\n\nfrom hmr4d.utils.comm.gather import all_gather\nfrom hmr4d.utils.pylogger import Log\n\nfrom hmr4d.utils.eval.eval_utils import (\n compute_camcoord_metrics,\n compute_global_metrics,\n compute_camcoord_perjoint_metrics,\n as_np_array,\n)\nfrom hmr4d.utils.geo_transform import apply_T_on_points, compute_T_ayfz2ay\nfrom hmr4d.utils.smplx_utils import make_smplx\nfrom einops import einsum, rearrange\n\nfrom pytorch3d.transforms import axis_angle_to_matrix, matrix_to_axis_angle\nfrom hmr4d.utils.wis3d_utils import make_wis3d, add_motion_as_lines, get_colors_by_conf\nfrom hmr4d.utils.vis.renderer import Renderer, get_global_cameras_static, get_ground_params_from_points\nfrom hmr4d.utils.geo.hmr_cam import estimate_focal_length\nfrom hmr4d.utils.video_io_utils import read_video_np, save_video, get_writer\nimport imageio\nfrom tqdm import tqdm\nfrom pathlib import Path\nimport numpy as np\nimport cv2\n\nfrom smplx.joint_names import JOINT_NAMES\nfrom hmr4d.utils.net_utils import repeat_to_max_len, gaussian_smooth\nfrom hmr4d.utils.geo.hmr_global import rollout_vel, get_static_joint_mask\n\n\nclass MetricMocap(pl.Callback):\n def __init__(self):\n super().__init__()\n # vid->result\n self.metric_aggregator = {\n \"pa_mpjpe\": {},\n \"mpjpe\": {},\n \"pve\": {},\n \"accel\": {},\n \"wa2_mpjpe\": {},\n \"waa_mpjpe\": {},\n \"rte\": {},\n \"jitter\": {},\n \"fs\": {},\n }\n\n self.perjoint_metrics = False\n if self.perjoint_metrics:\n body_joint_names = JOINT_NAMES[:22] + [\"left_hand\", \"right_hand\"]\n self.body_joint_names = body_joint_names\n self.perjoint_metric_aggregator = {\n \"mpjpe\": {k: {} for k in body_joint_names},\n }\n self.perjoint_obs_metric_aggregator = {\n \"mpjpe\": {k: {} for k in body_joint_names},\n }\n\n # SMPL\n self.smplx_model = {\n \"male\": make_smplx(\"rich-smplx\", gender=\"male\"),\n \"female\": make_smplx(\"rich-smplx\", gender=\"female\"),\n \"neutral\": make_smplx(\"rich-smplx\", gender=\"neutral\"),\n }\n self.J_regressor = torch.load(\"hmr4d/utils/body_model/smpl_neutral_J_regressor.pt\")\n self.smplx2smpl = torch.load(\"hmr4d/utils/body_model/smplx2smpl_sparse.pt\")\n self.faces_smpl = make_smplx(\"smpl\").faces\n self.faces_smplx = self.smplx_model[\"neutral\"].faces\n\n # The metrics are calculated similarly for val/test/predict\n self.on_test_batch_end = self.on_validation_batch_end = self.on_predict_batch_end\n\n # Only validation record the metrics with logger\n self.on_test_epoch_end = self.on_validation_epoch_end = self.on_predict_epoch_end\n\n # ================== Batch-based Computation ================== #\n def on_predict_batch_end(self, trainer, pl_module, outputs, batch, batch_idx, dataloader_idx=0):\n \"\"\"The behaviour is the same for val/test/predict\"\"\"\n assert batch[\"B\"] == 1\n dataset_id = batch[\"meta\"][0][\"dataset_id\"]\n if dataset_id != \"RICH\":\n return\n\n # Move to cuda if not\n for g in [\"male\", \"female\", \"neutral\"]:\n self.smplx_model[g] = self.smplx_model[g].cuda()\n self.J_regressor = self.J_regressor.cuda()\n self.smplx2smpl = self.smplx2smpl.cuda()\n\n vid = batch[\"meta\"][0][\"vid\"]\n seq_length = batch[\"length\"][0].item()\n gender = batch[\"gender\"][0]\n T_w2ay = batch[\"T_w2ay\"][0]\n T_w2c = batch[\"T_w2c\"][0]\n\n # Groundtruth (world, cam)\n target_w_params = {k: v[0] for k, v in batch[\"gt_smpl_params\"].items()}\n target_w_output = self.smplx_model[gender](**target_w_params)\n target_w_verts = torch.stack([torch.matmul(self.smplx2smpl, v_) for v_ in target_w_output.vertices])\n target_c_verts = apply_T_on_points(target_w_verts, T_w2c)\n target_c_j3d = torch.matmul(self.J_regressor, target_c_verts)\n offset = target_c_j3d[..., [1, 2], :].mean(-2, keepdim=True) # (L, 1, 3)\n target_cr_j3d = target_c_j3d - offset\n target_cr_verts = target_c_verts - offset\n # optional: ay for visual comparison\n target_ay_verts = apply_T_on_points(target_w_verts, T_w2ay)\n target_ay_j3d = torch.matmul(self.J_regressor, target_ay_verts)\n\n # + Prediction -> Metric\n # 1. cam\n pred_smpl_params_incam = outputs[\"pred_smpl_params_incam\"]\n smpl_out = self.smplx_model[\"neutral\"](**pred_smpl_params_incam)\n pred_c_verts = torch.stack([torch.matmul(self.smplx2smpl, v_) for v_ in smpl_out.vertices])\n pred_c_j3d = einsum(self.J_regressor, pred_c_verts, \"j v, l v i -> l j i\")\n offset = pred_c_j3d[..., [1, 2], :].mean(-2, keepdim=True) # (L, 1, 3)\n\n # 2. ay\n pred_smpl_params_global = outputs[\"pred_smpl_params_global\"]\n smpl_out = self.smplx_model[\"neutral\"](**pred_smpl_params_global)\n pred_ay_verts = torch.stack([torch.matmul(self.smplx2smpl, v_) for v_ in smpl_out.vertices])\n pred_ay_j3d = einsum(self.J_regressor, pred_ay_verts, \"j v, l v i -> l j i\")\n\n # Metric of current sequence\n batch_eval = {\n \"pred_j3d\": pred_c_j3d,\n \"target_j3d\": target_c_j3d,\n \"pred_verts\": pred_c_verts,\n \"target_verts\": target_c_verts,\n }\n camcoord_metrics = compute_camcoord_metrics(batch_eval)\n for k in camcoord_metrics:\n self.metric_aggregator[k][vid] = as_np_array(camcoord_metrics[k])\n\n batch_eval = {\n \"pred_j3d_glob\": pred_ay_j3d,\n \"target_j3d_glob\": target_ay_j3d,\n \"pred_verts_glob\": pred_ay_verts,\n \"target_verts_glob\": target_ay_verts,\n }\n global_metrics = compute_global_metrics(batch_eval)\n for k in global_metrics:\n self.metric_aggregator[k][vid] = as_np_array(global_metrics[k])\n\n if False: # global wi3d debug\n wis3d = make_wis3d(name=\"debug-metric-global\")\n add_motion_as_lines(pred_ay_j3d, wis3d, name=\"pred_ay_j3d\")\n add_motion_as_lines(target_ay_j3d, wis3d, name=\"target_ay_j3d\")\n\n if False: # incam visualize debug\n # Print per-sequence error\n Log.info(\n f\"seq {vid} metrics:\\n\"\n + \"\\n\".join(\n f\"{k}: {self.metric_aggregator[k][vid].mean():.1f} (obs:{camcoord_metrics[k].mean():.1f})\"\n for k in camcoord_metrics.keys()\n )\n + \"\\n------\\n\"\n )\n if self.perjoint_metrics:\n Log.info(\n f\"\\n\".join(\n f\"{k}-{j}: {self.perjoint_metric_aggregator[k][j][vid].mean():.1f} (obs:{self.perjoint_obs_metric_aggregator[k][j][vid].mean():.1f})\"\n for j in self.body_joint_names\n for k in self.perjoint_obs_metric_aggregator.keys()\n )\n + \"\\n------\"\n )\n\n # -- metric -- #\n pred_mpjpe = self.metric_aggregator[\"mpjpe\"][vid].mean()\n obs_mpjpe = camcoord_metrics[\"mpjpe\"].mean()\n\n # -- render mesh -- #\n vertices_gt = target_c_verts\n vertices_cr_gt = target_cr_verts + target_cr_verts.new([0, 0, 3.0]) # move forward +z\n vertices_pred = pred_c_verts\n vertices_cr_obs = obs_cr_verts + obs_cr_verts.new([0, 0, 3.0]) # move forward +z\n vertices_cr_pred = pred_cr_verts + pred_cr_verts.new([0, 0, 3.0]) # move forward +z\n\n # -- rendering code -- #\n vname = batch[\"meta_render\"][0][\"name\"]\n K = batch[\"meta_render\"][0][\"K\"]\n width, height = batch[\"meta_render\"][0][\"width_height\"]\n faces = self.faces_smpl\n\n renderer = Renderer(width, height, device=\"cuda\", faces=faces, K=K)\n out_fn = f\"outputs/dump_render/{vname}.mp4\"\n Path(out_fn).parent.mkdir(exist_ok=True, parents=True)\n writer = imageio.get_writer(out_fn, fps=30, mode=\"I\", format=\"FFMPEG\", macro_block_size=1)\n\n # imgs\n video_path = batch[\"meta_render\"][0][\"video_path\"]\n frame_id = batch[\"meta_render\"][0][\"frame_id\"].cpu().numpy()\n vr = decord.VideoReader(video_path)\n images = vr.get_batch(list(frame_id)).numpy() # (F, H/4, W/4, 3), uint8, numpy\n\n for i in tqdm(range(seq_length), desc=f\"Rendering {vname}\"):\n img_overlay_gt = renderer.render_mesh(vertices_gt[i].cuda(), images[i], [39, 194, 128])\n if batch[\"meta_render\"][0].get(\"bbx_xys\", None) is not None: # draw bbox lines\n bbx_xys = batch[\"meta_render\"][0][\"bbx_xys\"][i].cpu().numpy()\n lu_point = (bbx_xys[:2] - bbx_xys[2:] / 2).astype(int)\n rd_point = (bbx_xys[:2] + bbx_xys[2:] / 2).astype(int)\n img_overlay_gt = cv2.rectangle(img_overlay_gt, lu_point, rd_point, (255, 178, 102), 2)\n\n img_overlay_pred = renderer.render_mesh(vertices_pred[i].cuda(), images[i])\n # img_overlay_pred = renderer.render_mesh(vertices_pred[i].cuda(), np.zeros_like(images[i]))\n img = np.concatenate([img_overlay_gt, img_overlay_pred], axis=0)\n\n ####### overlay gt cr first, then overlay pred cr with error color ########\n # overlay gt cr first with blue color\n black_overlay_obs = renderer.render_mesh(\n vertices_cr_gt[i].cuda(), np.zeros_like(images[i]), colors=[39, 194, 128]\n )\n black_overlay_pred = renderer.render_mesh(\n vertices_cr_gt[i].cuda(), np.zeros_like(images[i]), colors=[39, 194, 128]\n )\n\n # get error color\n obs_error = (vertices_cr_gt[i] - vertices_cr_obs[i]).norm(dim=-1)\n pred_error = (vertices_cr_gt[i] - vertices_cr_pred[i]).norm(dim=-1)\n max_error = max(obs_error.max(), pred_error.max())\n obs_error_color = torch.stack(\n [obs_error / max_error, torch.ones_like(obs_error) * 0.6, torch.ones_like(obs_error) * 0.6],\n dim=-1,\n )\n obs_error_color = torch.clip(obs_error_color, 0, 1)\n pred_error_color = torch.stack(\n [pred_error / max_error, torch.ones_like(pred_error) * 0.6, torch.ones_like(pred_error) * 0.6],\n dim=-1,\n )\n pred_error_color = torch.clip(pred_error_color, 0, 1)\n\n # overlay cr with error color\n black_overlay_obs = renderer.render_mesh(\n vertices_cr_obs[i].cuda(), black_overlay_obs, colors=obs_error_color[None]\n )\n black_overlay_pred = renderer.render_mesh(\n vertices_cr_pred[i].cuda(), black_overlay_pred, colors=pred_error_color[None]\n )\n\n # write mpjpe on the img\n obs_mpjpe_ = camcoord_metrics[\"mpjpe\"][i]\n text = f\"obs mpjpe: {obs_mpjpe_:.1f} ({obs_mpjpe:.1f})\"\n cv2.putText(black_overlay_obs, text, (10, 30), cv2.FONT_HERSHEY_SIMPLEX, 1, (100, 200, 200), 2)\n pred_mpjpe_ = self.metric_aggregator[\"mpjpe\"][vid][i]\n text = f\"pred mpjpe: {pred_mpjpe_:.1f} ({pred_mpjpe:.1f})\"\n if pred_mpjpe_ > obs_mpjpe_:\n # large error -> purple\n cv2.putText(black_overlay_pred, text, (10, 30), cv2.FONT_HERSHEY_SIMPLEX, 1, (200, 100, 200), 2)\n else:\n # small error -> yellow\n cv2.putText(black_overlay_pred, text, (10, 30), cv2.FONT_HERSHEY_SIMPLEX, 1, (200, 200, 100), 2)\n black = np.concatenate([black_overlay_obs, black_overlay_pred], axis=0)\n ###########################################\n\n img = np.concatenate([img, black], axis=1)\n\n writer.append_data(img)\n writer.close()\n\n if False: # Visualize incam + global results\n\n def move_to_start_point_face_z(verts):\n \"XZ to origin, Start from the ground, Face-Z\"\n # position\n verts = verts.clone() # (L, V, 3)\n offset = einsum(self.J_regressor, verts[0], \"j v, v i -> j i\")[0] # (3)\n offset[1] = verts[:, :, [1]].min()\n verts = verts - offset\n # face direction\n T_ay2ayfz = compute_T_ayfz2ay(einsum(self.J_regressor, verts[[0]], \"j v, l v i -> l j i\"), inverse=True)\n verts = apply_T_on_points(verts, T_ay2ayfz)\n return verts\n\n verts_incam = pred_c_verts.clone()\n # verts_glob = move_to_start_point_face_z(target_ay_verts) # gt\n verts_glob = move_to_start_point_face_z(pred_ay_verts)\n joints_glob = einsum(self.J_regressor, verts_glob, \"j v, l v i -> l j i\") # (L, J, 3)\n global_R, global_T, global_lights = get_global_cameras_static(\n verts_glob.cpu(),\n beta=4.0,\n cam_height_degree=20,\n target_center_height=1.0,\n vec_rot=-45,\n )\n\n # -- rendering code (global version FOV=55) -- #\n vname = batch[\"meta_render\"][0][\"name\"]\n width, height = batch[\"meta_render\"][0][\"width_height\"]\n K = batch[\"meta_render\"][0][\"K\"]\n faces = self.faces_smpl\n out_fn = f\"outputs/dump_render_global/{vname}.mp4\"\n Path(out_fn).parent.mkdir(exist_ok=True, parents=True)\n\n # two renderers\n renderer_incam = Renderer(width, height, device=\"cuda\", faces=faces, K=K)\n renderer_glob = Renderer(width, height, estimate_focal_length(width, height), device=\"cuda\", faces=faces)\n\n # imgs\n video_path = batch[\"meta_render\"][0][\"video_path\"]\n frame_id = batch[\"meta_render\"][0][\"frame_id\"].cpu().numpy()\n images = read_video_np(video_path)[frame_id] # (F, H/4, W/4, 3), uint8, numpy\n\n # Actual rendering\n scale, cx, cz = get_ground_params_from_points(joints_glob[:, 0], verts_glob)\n renderer_glob.set_ground(scale * 1.5, cx, cz)\n color = torch.ones(3).float().cuda() * 0.8\n\n writer = get_writer(out_fn, fps=30, crf=23)\n for i in tqdm(range(seq_length), desc=f\"Rendering {vname}\"):\n # incam\n img_overlay_pred = renderer_incam.render_mesh(verts_incam[i].cuda(), images[i], [0.8, 0.8, 0.8])\n # if batch[\"meta_render\"][0].get(\"bbx_xys\", None) is not None: # draw bbox lines\n # bbx_xys = batch[\"meta_render\"][0][\"bbx_xys\"][i].cpu().numpy()\n # lu_point = (bbx_xys[:2] - bbx_xys[2:] / 2).astype(int)\n # rd_point = (bbx_xys[:2] + bbx_xys[2:] / 2).astype(int)\n # img_overlay_pred = cv2.rectangle(img_overlay_pred, lu_point, rd_point, (255, 178, 102), 2)\n # pred_mpjpe_ = self.metric_aggregator[\"mpjpe\"][vid][i]\n # text = f\"pred mpjpe: {pred_mpjpe_:.1f}\"\n # cv2.putText(img_overlay_pred, text, (10, 30), cv2.FONT_HERSHEY_SIMPLEX, 1, (200, 100, 200), 2)\n\n # glob\n cameras = renderer_glob.create_camera(global_R[i], global_T[i])\n # img_glob = renderer_glob.render_with_ground(verts_glob[[i]], color_[None], cameras, global_lights)\n img_glob = renderer_glob.render_with_ground(\n verts_glob[[i]], color.clone()[None], cameras, global_lights\n )\n\n # write\n img = np.concatenate([img_overlay_pred, img_glob], axis=1)\n writer.write_frame(img)\n writer.close()\n\n # ================== Epoch Summary ================== #\n def on_predict_epoch_end(self, trainer, pl_module):\n \"\"\"Without logger\"\"\"\n local_rank, world_size = trainer.local_rank, trainer.world_size\n monitor_metric = \"mpjpe\"\n\n # Reduce metric_aggregator across all processes\n metric_keys = list(self.metric_aggregator.keys())\n with torch.inference_mode(False): # allow in-place operation of all_gather\n metric_aggregator_gathered = all_gather(self.metric_aggregator) # list of dict\n for metric_key in metric_keys:\n for d in metric_aggregator_gathered:\n self.metric_aggregator[metric_key].update(d[metric_key])\n\n if False: # debug to make sure the all_gather is correct\n print(f\"[RANK {local_rank}/{world_size}]: {self.metric_aggregator[monitor_metric].keys()}\")\n\n total = len(self.metric_aggregator[monitor_metric])\n Log.info(f\"{total} sequences evaluated in {self.__class__.__name__}\")\n if total == 0:\n return\n\n # print monitored metric per sequence\n mm_per_seq = {k: v.mean() for k, v in self.metric_aggregator[monitor_metric].items()}\n if len(mm_per_seq) > 0:\n sorted_mm_per_seq = sorted(mm_per_seq.items(), key=lambda x: x[1], reverse=True)\n n_worst = 5 if trainer.state.stage == \"validate\" else len(sorted_mm_per_seq)\n if local_rank == 0:\n Log.info(\n f\"monitored metric {monitor_metric} per sequence\\n\"\n + \"\\n\".join([f\"{m:5.1f} : {s}\" for s, m in sorted_mm_per_seq[:n_worst]])\n + \"\\n------\"\n )\n\n # average over all batches\n metrics_avg = {k: np.concatenate(list(v.values())).mean() for k, v in self.metric_aggregator.items()}\n if local_rank == 0:\n Log.info(f\"[Metrics] RICH:\\n\" + \"\\n\".join(f\"{k}: {v:.1f}\" for k, v in metrics_avg.items()) + \"\\n------\")\n\n # save to logger if available\n if pl_module.logger is not None:\n cur_epoch = pl_module.current_epoch\n for k, v in metrics_avg.items():\n pl_module.logger.log_metrics({f\"val_metric_RICH/{k}\": v}, step=cur_epoch)\n\n # reset\n for k in self.metric_aggregator:\n self.metric_aggregator[k] = {}\n\n\nrich_node = builds(MetricMocap)\nMainStore.store(name=\"metric_rich\", node=rich_node, group=\"callbacks\", package=\"callbacks.metric_rich\")\n" }, { "path": "eval/GVHMR/hmr4d/model/gvhmr/gvhmr_pl.py", "content": "from typing import Any, Dict\nimport numpy as np\nfrom pathlib import Path\nimport torch\nimport pytorch_lightning as pl\nfrom hydra.utils import instantiate\nfrom hmr4d.utils.pylogger import Log\nfrom einops import rearrange, einsum\nfrom hmr4d.configs import MainStore, builds\n\nfrom hmr4d.utils.geo_transform import compute_T_ayfz2ay, apply_T_on_points\nfrom hmr4d.utils.wis3d_utils import make_wis3d, add_motion_as_lines\nfrom hmr4d.utils.smplx_utils import make_smplx\nfrom hmr4d.utils.geo.augment_noisy_pose import (\n get_wham_aug_kp3d,\n get_visible_mask,\n get_invisible_legs_mask,\n randomly_occlude_lower_half,\n randomly_modify_hands_legs,\n)\nfrom hmr4d.utils.geo.hmr_cam import perspective_projection, normalize_kp2d, safely_render_x3d_K, get_bbx_xys\n\nfrom hmr4d.utils.video_io_utils import save_video\nfrom hmr4d.utils.vis.cv2_utils import draw_bbx_xys_on_image_batch\nfrom hmr4d.utils.geo.flip_utils import flip_smplx_params, avg_smplx_aa\nfrom hmr4d.model.gvhmr.utils.postprocess import pp_static_joint, pp_static_joint_cam, process_ik\n\n\nclass GvhmrPL(pl.LightningModule):\n def __init__(\n self,\n pipeline,\n optimizer=None,\n scheduler_cfg=None,\n ignored_weights_prefix=[\"smplx\", \"pipeline.endecoder\"],\n ):\n super().__init__()\n self.pipeline = instantiate(pipeline, _recursive_=False)\n self.optimizer = instantiate(optimizer)\n self.scheduler_cfg = scheduler_cfg\n\n # Options\n self.ignored_weights_prefix = ignored_weights_prefix\n\n # The test step is the same as validation\n self.test_step = self.predict_step = self.validation_step\n\n # SMPLX\n self.smplx = make_smplx(\"supermotion_v437coco17\")\n\n def training_step(self, batch, batch_idx):\n B, F = batch[\"smpl_params_c\"][\"body_pose\"].shape[:2]\n\n # Create augmented noisy-obs : gt_j3d(coco17)\n with torch.no_grad():\n gt_verts437, gt_j3d = self.smplx(**batch[\"smpl_params_c\"])\n root_ = gt_j3d[:, :, [11, 12], :].mean(-2, keepdim=True)\n batch[\"gt_j3d\"] = gt_j3d\n batch[\"gt_cr_coco17\"] = gt_j3d - root_\n batch[\"gt_c_verts437\"] = gt_verts437\n batch[\"gt_cr_verts437\"] = gt_verts437 - root_\n\n # bbx_xys\n i_x2d = safely_render_x3d_K(gt_verts437, batch[\"K_fullimg\"], thr=0.3)\n bbx_xys = get_bbx_xys(i_x2d, do_augment=True)\n if False: # trust image bbx_xys seems better\n batch[\"bbx_xys\"] = bbx_xys\n else:\n mask_bbx_xys = batch[\"mask\"][\"bbx_xys\"]\n batch[\"bbx_xys\"][~mask_bbx_xys] = bbx_xys[~mask_bbx_xys]\n if False: # visualize bbx_xys from an iPhone view\n render_w, render_h = 120, 160 # iphone main-lens 24mm 3:4\n ratio = render_w / 1528\n offset = torch.tensor([764 - 500, 1019 - 500]).to(i_x2d)\n i_x2d_render = (i_x2d + offset).clone()\n i_x2d_render = (i_x2d_render * ratio).long().clone()\n torch.clamp_(i_x2d_render[..., 0], 0, render_w - 1)\n torch.clamp_(i_x2d_render[..., 1], 0, render_h - 1)\n bbx_xys_render = bbx_xys.clone()\n bbx_xys_render[..., :2] += offset\n bbx_xys_render *= ratio\n\n output_dir = Path(\"outputs/simulated_bbx_xys\")\n output_dir.mkdir(parents=True, exist_ok=True)\n video_list = []\n for bid in range(B):\n images = torch.zeros(F, render_h, render_w, 3, device=i_x2d.device)\n for fid in range(F):\n images[fid, i_x2d_render[bid, fid, :, 1], i_x2d_render[bid, fid, :, 0]] = 255\n\n images = draw_bbx_xys_on_image_batch(bbx_xys_render[bid].cpu().numpy(), images.cpu().numpy())\n images = np.stack(images).astype(\"uint8\") # (L, H, W, 3)\n images[:, 0, :] = np.array([255, 255, 255])\n images[:, -1, :] = np.array([255, 255, 255])\n images[:, :, 0] = np.array([255, 255, 255])\n images[:, :, -1] = np.array([255, 255, 255])\n video_list.append(images)\n\n # stack videos\n video_output = []\n for i in range(0, len(video_list), 4):\n if i + 4 <= len(video_list):\n video_output.append(np.concatenate(video_list[i : i + 4], axis=2))\n video_output = np.concatenate(video_output, axis=1)\n save_video(video_output, output_dir / f\"{batch_idx}.mp4\", fps=30, quality=5)\n\n # noisy_j3d -> project to i_j2d -> compute a bbx -> normalized kp2d [-1, 1]\n noisy_j3d = gt_j3d + get_wham_aug_kp3d(gt_j3d.shape[:2])\n if True:\n noisy_j3d = randomly_modify_hands_legs(noisy_j3d)\n obs_i_j2d = perspective_projection(noisy_j3d, batch[\"K_fullimg\"]) # (B, L, J, 2)\n j2d_visible_mask = get_visible_mask(gt_j3d.shape[:2]).cuda() # (B, L, J)\n j2d_visible_mask[noisy_j3d[..., 2] < 0.3] = False # Set close-to-image-plane points as invisible\n if True: # Set both legs as invisible for a period\n legs_invisible_mask = get_invisible_legs_mask(gt_j3d.shape[:2]).cuda() # (B, L, J)\n j2d_visible_mask[legs_invisible_mask] = False\n obs_kp2d = torch.cat([obs_i_j2d, j2d_visible_mask[:, :, :, None].float()], dim=-1) # (B, L, J, 3)\n obs = normalize_kp2d(obs_kp2d, batch[\"bbx_xys\"]) # (B, L, J, 3)\n obs[~j2d_visible_mask] = 0 # if not visible, set to (0,0,0)\n batch[\"obs\"] = obs\n\n if True: # Use some detected vitpose (presave data)\n prob = 0.5\n mask_real_vitpose = (torch.rand(B).to(obs_kp2d) < prob) * batch[\"mask\"][\"vitpose\"]\n batch[\"obs\"][mask_real_vitpose] = normalize_kp2d(batch[\"kp2d\"], batch[\"bbx_xys\"])[mask_real_vitpose]\n\n # Set untrusted frames to False\n batch[\"obs\"][~batch[\"mask\"][\"valid\"]] = 0\n\n if False: # wis3d\n wis3d = make_wis3d(name=\"debug-aug-kp3d\")\n add_motion_as_lines(gt_j3d[0], wis3d, name=\"gt_j3d\", skeleton_type=\"coco17\")\n add_motion_as_lines(noisy_j3d[0], wis3d, name=\"noisy_j3d\", skeleton_type=\"coco17\")\n\n # f_imgseq: apply random aug on offline extracted features\n # f_imgseq = batch[\"f_imgseq\"] + torch.randn_like(batch[\"f_imgseq\"]) * 0.1\n # f_imgseq[~batch[\"mask\"][\"f_imgseq\"]] = 0\n # batch[\"f_imgseq\"] = f_imgseq.clone()\n\n # Forward and get loss\n outputs = self.pipeline.forward(batch, train=True)\n\n # Log\n log_kwargs = {\n \"on_epoch\": True,\n \"prog_bar\": True,\n \"logger\": True,\n \"batch_size\": B,\n \"sync_dist\": True,\n }\n self.log(\"train/loss\", outputs[\"loss\"], **log_kwargs)\n for k, v in outputs.items():\n if \"_loss\" in k:\n self.log(f\"train/{k}\", v, **log_kwargs)\n\n return outputs\n\n def validation_step(self, batch, batch_idx, dataloader_idx=0):\n # Options & Check\n do_postproc = self.trainer.state.stage == \"test\" # Only apply postproc in test\n do_flip_test = \"flip_test\" in batch\n do_postproc_not_flip_test = do_postproc and not do_flip_test # later pp when flip_test\n assert batch[\"B\"] == 1, \"Only support batch size 1 in evalution.\"\n\n # ROPE inference\n obs = normalize_kp2d(batch[\"kp2d\"], batch[\"bbx_xys\"])\n if \"mask\" in batch:\n obs[0, ~batch[\"mask\"][0]] = 0\n\n batch_ = {\n \"length\": batch[\"length\"],\n \"obs\": obs,\n \"bbx_xys\": batch[\"bbx_xys\"],\n \"K_fullimg\": batch[\"K_fullimg\"],\n \"cam_angvel\": batch[\"cam_angvel\"],\n \"f_imgseq\": batch[\"f_imgseq\"],\n }\n outputs = self.pipeline.forward(batch_, train=False, postproc=do_postproc_not_flip_test)\n outputs[\"pred_smpl_params_global\"] = {k: v[0] for k, v in outputs[\"pred_smpl_params_global\"].items()}\n outputs[\"pred_smpl_params_incam\"] = {k: v[0] for k, v in outputs[\"pred_smpl_params_incam\"].items()}\n\n if do_flip_test:\n flip_test = batch[\"flip_test\"]\n obs = normalize_kp2d(flip_test[\"kp2d\"], flip_test[\"bbx_xys\"])\n if \"mask\" in batch:\n obs[0, ~batch[\"mask\"][0]] = 0\n\n batch_ = {\n \"length\": batch[\"length\"],\n \"obs\": obs,\n \"bbx_xys\": flip_test[\"bbx_xys\"],\n \"K_fullimg\": batch[\"K_fullimg\"],\n \"cam_angvel\": flip_test[\"cam_angvel\"],\n \"f_imgseq\": flip_test[\"f_imgseq\"],\n }\n flipped_outputs = self.pipeline.forward(batch_, train=False)\n\n # First update incam results\n flipped_outputs[\"pred_smpl_params_incam\"] = {\n k: v[0] for k, v in flipped_outputs[\"pred_smpl_params_incam\"].items()\n }\n smpl_params1 = outputs[\"pred_smpl_params_incam\"]\n smpl_params2 = flip_smplx_params(flipped_outputs[\"pred_smpl_params_incam\"])\n\n smpl_params_avg = smpl_params1.copy()\n smpl_params_avg[\"betas\"] = (smpl_params1[\"betas\"] + smpl_params2[\"betas\"]) / 2\n smpl_params_avg[\"body_pose\"] = avg_smplx_aa(smpl_params1[\"body_pose\"], smpl_params2[\"body_pose\"])\n smpl_params_avg[\"global_orient\"] = avg_smplx_aa(\n smpl_params1[\"global_orient\"], smpl_params2[\"global_orient\"]\n )\n outputs[\"pred_smpl_params_incam\"] = smpl_params_avg\n\n # Then update global results\n outputs[\"pred_smpl_params_global\"][\"betas\"] = smpl_params_avg[\"betas\"]\n outputs[\"pred_smpl_params_global\"][\"body_pose\"] = smpl_params_avg[\"body_pose\"]\n\n # Finally, apply postprocess\n if do_postproc:\n # temporarily recover the original batch-dim\n outputs[\"pred_smpl_params_global\"] = {k: v[None] for k, v in outputs[\"pred_smpl_params_global\"].items()}\n outputs[\"pred_smpl_params_global\"][\"transl\"] = pp_static_joint(outputs, self.pipeline.endecoder)\n body_pose = process_ik(outputs, self.pipeline.endecoder)\n outputs[\"pred_smpl_params_global\"] = {k: v[0] for k, v in outputs[\"pred_smpl_params_global\"].items()}\n\n outputs[\"pred_smpl_params_global\"][\"body_pose\"] = body_pose[0]\n # outputs[\"pred_smpl_params_incam\"][\"body_pose\"] = body_pose[0]\n\n if False: # wis3d\n wis3d = make_wis3d(name=\"debug-rich-cap\")\n smplx_model = make_smplx(\"rich-smplx\", gender=\"neutral\").cuda()\n gender = batch[\"gender\"][0]\n T_w2ay = batch[\"T_w2ay\"][0]\n\n # Prediction\n # add_motion_as_lines(outputs_window[\"pred_ayfz_motion\"][bid], wis3d, name=\"pred_ayfz_motion\")\n\n smplx_out = smplx_model(**pred_smpl_params_global)\n for i in range(len(smplx_out.vertices)):\n wis3d.set_scene_id(i)\n wis3d.add_mesh(smplx_out.vertices[i], smplx_model.bm.faces, name=f\"pred-smplx-global\")\n\n # GT (w)\n smplx_models = {\n \"male\": make_smplx(\"rich-smplx\", gender=\"male\").cuda(),\n \"female\": make_smplx(\"rich-smplx\", gender=\"female\").cuda(),\n }\n gt_smpl_params = {k: v[0, windows[0]] for k, v in batch[\"gt_smpl_params\"].items()}\n gt_smplx_out = smplx_models[gender](**gt_smpl_params)\n\n # GT (ayfz)\n smplx_verts_ay = apply_T_on_points(gt_smplx_out.vertices, T_w2ay)\n smplx_joints_ay = apply_T_on_points(gt_smplx_out.joints, T_w2ay)\n T_ay2ayfz = compute_T_ayfz2ay(smplx_joints_ay[:1], inverse=True)[0] # (4, 4)\n smplx_verts_ayfz = apply_T_on_points(smplx_verts_ay, T_ay2ayfz) # (F, 22, 3)\n\n for i in range(len(smplx_verts_ayfz)):\n wis3d.set_scene_id(i)\n wis3d.add_mesh(smplx_verts_ayfz[i], smplx_models[gender].bm.faces, name=f\"gt-smplx-ayfz\")\n\n breakpoint()\n\n if False: # o3d\n prog_keys = [\n \"pred_smpl_progress\",\n \"pred_localjoints_progress\",\n \"pred_incam_localjoints_progress\",\n ]\n for k in prog_keys:\n if k in outputs_window:\n seq_out = torch.cat(\n [v[:, :l] for v, l in zip(outputs_window[k], length)], dim=1\n ) # (B, P, L, J, 3) -> (P, L, J, 3) -> (P, CL, J, 3)\n outputs[k] = seq_out[None]\n\n return outputs\n\n def configure_optimizers(self):\n params = []\n for k, v in self.pipeline.named_parameters():\n if v.requires_grad:\n params.append(v)\n optimizer = self.optimizer(params=params)\n\n if self.scheduler_cfg[\"scheduler\"] is None:\n return optimizer\n\n scheduler_cfg = dict(self.scheduler_cfg)\n scheduler_cfg[\"scheduler\"] = instantiate(scheduler_cfg[\"scheduler\"], optimizer=optimizer)\n return [optimizer], [scheduler_cfg]\n\n # ============== Utils ================= #\n def on_save_checkpoint(self, checkpoint) -> None:\n for ig_keys in self.ignored_weights_prefix:\n for k in list(checkpoint[\"state_dict\"].keys()):\n if k.startswith(ig_keys):\n # Log.info(f\"Remove key `{ig_keys}' from checkpoint.\")\n checkpoint[\"state_dict\"].pop(k)\n\n def load_pretrained_model(self, ckpt_path):\n \"\"\"Load pretrained checkpoint, and assign each weight to the corresponding part.\"\"\"\n Log.info(f\"[PL-Trainer] Loading ckpt: {ckpt_path}\")\n\n state_dict = torch.load(ckpt_path, \"cpu\")[\"state_dict\"]\n missing, unexpected = self.load_state_dict(state_dict, strict=False)\n real_missing = []\n for k in missing:\n ignored_when_saving = any(k.startswith(ig_keys) for ig_keys in self.ignored_weights_prefix)\n if not ignored_when_saving:\n real_missing.append(k)\n\n if len(real_missing) > 0:\n Log.warn(f\"Missing keys: {real_missing}\")\n if len(unexpected) > 0:\n Log.warn(f\"Unexpected keys: {unexpected}\")\n\n\ngvhmr_pl = builds(\n GvhmrPL,\n pipeline=\"${pipeline}\",\n optimizer=\"${optimizer}\",\n scheduler_cfg=\"${scheduler_cfg}\",\n populate_full_signature=True, # Adds all the arguments to the signature\n)\nMainStore.store(name=\"gvhmr_pl\", node=gvhmr_pl, group=\"model/gvhmr\")\n" }, { "path": "eval/GVHMR/hmr4d/model/gvhmr/gvhmr_pl_demo.py", "content": "import torch\nimport pytorch_lightning as pl\nfrom hydra.utils import instantiate\nfrom hmr4d.utils.pylogger import Log\nfrom hmr4d.configs import MainStore, builds\n\nfrom hmr4d.utils.geo.hmr_cam import normalize_kp2d\n\n\nclass DemoPL(pl.LightningModule):\n def __init__(self, pipeline):\n super().__init__()\n self.pipeline = instantiate(pipeline, _recursive_=False)\n\n @torch.no_grad()\n def predict(self, data, static_cam=False):\n \"\"\"auto add batch dim\n data: {\n \"length\": int, or Torch.Tensor,\n \"kp2d\": (F, 3)\n \"bbx_xys\": (F, 3)\n \"K_fullimg\": (F, 3, 3)\n \"cam_angvel\": (F, 3)\n \"f_imgseq\": (F, 3, 256, 256)\n }\n\n \"\"\"\n # ROPE inference\n batch = {\n \"length\": data[\"length\"][None],\n \"obs\": normalize_kp2d(data[\"kp2d\"], data[\"bbx_xys\"])[None],\n \"bbx_xys\": data[\"bbx_xys\"][None],\n \"K_fullimg\": data[\"K_fullimg\"][None],\n \"cam_angvel\": data[\"cam_angvel\"][None],\n \"f_imgseq\": data[\"f_imgseq\"][None],\n }\n batch = {k: v.cuda() for k, v in batch.items()}\n outputs = self.pipeline.forward(batch, train=False, postproc=True, static_cam=static_cam)\n\n pred = {\n \"smpl_params_global\": {k: v[0] for k, v in outputs[\"pred_smpl_params_global\"].items()},\n \"smpl_params_incam\": {k: v[0] for k, v in outputs[\"pred_smpl_params_incam\"].items()},\n \"K_fullimg\": data[\"K_fullimg\"],\n \"net_outputs\": outputs, # intermediate outputs\n }\n return pred\n\n def load_pretrained_model(self, ckpt_path):\n \"\"\"Load pretrained checkpoint, and assign each weight to the corresponding part.\"\"\"\n Log.info(f\"[PL-Trainer] Loading ckpt type: {ckpt_path}\")\n\n state_dict = torch.load(ckpt_path, \"cpu\")[\"state_dict\"]\n missing, unexpected = self.load_state_dict(state_dict, strict=False)\n if len(missing) > 0:\n Log.warn(f\"Missing keys: {missing}\")\n if len(unexpected) > 0:\n Log.warn(f\"Unexpected keys: {unexpected}\")\n\n\nMainStore.store(name=\"gvhmr_pl_demo\", node=builds(DemoPL, pipeline=\"${pipeline}\"), group=\"model/gvhmr\")\n" }, { "path": "eval/GVHMR/hmr4d/model/gvhmr/pipeline/gvhmr_pipeline.py", "content": "import torch\nimport torch.nn as nn\nimport torch.nn.functional as F\nfrom torch.cuda.amp import autocast\nimport numpy as np\nfrom einops import einsum, rearrange, repeat\nfrom hydra.utils import instantiate\nfrom hmr4d.utils.pylogger import Log\nfrom hmr4d.utils.net_utils import gaussian_smooth\n\nfrom hmr4d.model.gvhmr.utils.endecoder import EnDecoder\nfrom hmr4d.model.gvhmr.utils.postprocess import (\n pp_static_joint,\n process_ik,\n pp_static_joint_cam,\n)\nfrom hmr4d.model.gvhmr.utils import stats_compose\n\nfrom pytorch3d.transforms import (\n matrix_to_rotation_6d,\n rotation_6d_to_matrix,\n axis_angle_to_matrix,\n matrix_to_axis_angle,\n)\nfrom hmr4d.utils.geo.hmr_cam import compute_bbox_info_bedlam, compute_transl_full_cam, get_a_pred_cam, project_to_bi01\nfrom hmr4d.utils.geo.hmr_global import (\n rollout_local_transl_vel,\n get_static_joint_mask,\n get_tgtcoord_rootparam,\n)\nfrom hmr4d.utils.wis3d_utils import make_wis3d, add_motion_as_lines\nfrom hmr4d.utils.smplx_utils import make_smplx\n\n\nclass Pipeline(nn.Module):\n def __init__(self, args, args_denoiser3d, **kwargs):\n super().__init__()\n self.args = args\n self.weights = args.weights # loss weights\n\n # Networks\n self.denoiser3d = instantiate(args_denoiser3d, _recursive_=False)\n # Log.info(self.denoiser3d)\n\n # Normalizer\n self.endecoder: EnDecoder = instantiate(args.endecoder_opt, _recursive_=False)\n if self.args.normalize_cam_angvel:\n cam_angvel_stats = stats_compose.cam_angvel[\"manual\"]\n self.register_buffer(\"cam_angvel_mean\", torch.tensor(cam_angvel_stats[\"mean\"]), persistent=False)\n self.register_buffer(\"cam_angvel_std\", torch.tensor(cam_angvel_stats[\"std\"]), persistent=False)\n\n # ========== Training ========== #\n\n def forward(self, inputs, train=False, postproc=False, static_cam=False):\n outputs = dict()\n length = inputs[\"length\"] # (B,) effective length of each sample\n\n # *. Conditions\n cliff_cam = compute_bbox_info_bedlam(inputs[\"bbx_xys\"], inputs[\"K_fullimg\"]) # (B, L, 3)\n f_cam_angvel = inputs[\"cam_angvel\"]\n if self.args.normalize_cam_angvel:\n f_cam_angvel = (f_cam_angvel - self.cam_angvel_mean) / self.cam_angvel_std\n f_condition = {\n \"obs\": inputs[\"obs\"], # (B, L, J, 3)\n \"f_cliffcam\": cliff_cam, # (B, L, 3)\n \"f_cam_angvel\": f_cam_angvel, # (B, L, C=6)\n \"f_imgseq\": inputs[\"f_imgseq\"], # (B, L, C=1024)\n }\n if train:\n f_condition = randomly_set_null_condition(f_condition, 0.1)\n\n # Forward & output\n model_output = self.denoiser3d(length=length, **f_condition) # pred_x, pred_cam, static_conf_logits\n decode_dict = self.endecoder.decode(model_output[\"pred_x\"]) # (B, L, C) -> dict\n outputs.update({\"model_output\": model_output, \"decode_dict\": decode_dict})\n\n # Post-processing\n outputs[\"pred_smpl_params_incam\"] = {\n \"body_pose\": decode_dict[\"body_pose\"], # (B, L, 63)\n \"betas\": decode_dict[\"betas\"], # (B, L, 10)\n \"global_orient\": decode_dict[\"global_orient\"], # (B, L, 3)\n \"transl\": compute_transl_full_cam(model_output[\"pred_cam\"], inputs[\"bbx_xys\"], inputs[\"K_fullimg\"]),\n }\n if not train:\n pred_smpl_params_global = get_smpl_params_w_Rt_v2( # This function has for-loop\n global_orient_gv=decode_dict[\"global_orient_gv\"],\n local_transl_vel=decode_dict[\"local_transl_vel\"],\n global_orient_c=decode_dict[\"global_orient\"],\n cam_angvel=inputs[\"cam_angvel\"],\n )\n outputs[\"pred_smpl_params_global\"] = {\n \"body_pose\": decode_dict[\"body_pose\"],\n \"betas\": decode_dict[\"betas\"],\n **pred_smpl_params_global,\n }\n outputs[\"static_conf_logits\"] = model_output[\"static_conf_logits\"]\n\n if postproc: # apply post-processing\n if static_cam: # extra post-processing to utilize static camera prior\n outputs[\"pred_smpl_params_global\"][\"transl\"] = pp_static_joint_cam(outputs, self.endecoder)\n else:\n outputs[\"pred_smpl_params_global\"][\"transl\"] = pp_static_joint(outputs, self.endecoder)\n body_pose = process_ik(outputs, self.endecoder)\n decode_dict[\"body_pose\"] = body_pose\n outputs[\"pred_smpl_params_global\"][\"body_pose\"] = body_pose\n outputs[\"pred_smpl_params_incam\"][\"body_pose\"] = body_pose\n\n return outputs\n\n # ========== Compute Loss ========== #\n total_loss = 0\n mask = inputs[\"mask\"][\"valid\"] # (B, L)\n\n # 1. Simple loss: MSE\n pred_x = model_output[\"pred_x\"] # (B, L, C)\n target_x = self.endecoder.encode(inputs) # (B, L, C)\n simple_loss = F.mse_loss(pred_x, target_x, reduction=\"none\")\n mask_simple = mask[:, :, None].expand(-1, -1, pred_x.size(2)).clone() # (B, L, C)\n mask_simple[inputs[\"mask\"][\"spv_incam_only\"], :, 142:] = False # 3dpw training\n simple_loss = (simple_loss * mask_simple).mean()\n total_loss += simple_loss\n outputs[\"simple_loss\"] = simple_loss\n\n # 2. Extra loss\n extra_funcs = [\n compute_extra_incam_loss,\n compute_extra_global_loss,\n ]\n for extra_func in extra_funcs:\n extra_loss, extra_loss_dict = extra_func(inputs, outputs, self)\n total_loss += extra_loss\n outputs.update(extra_loss_dict)\n\n outputs[\"loss\"] = total_loss\n return outputs\n\n\ndef randomly_set_null_condition(f_condition, uncond_prob=0.1):\n \"\"\"Conditions are in shape (B, L, *)\"\"\"\n keys = list(f_condition.keys())\n for k in keys:\n if f_condition[k] is None:\n continue\n f_condition[k] = f_condition[k].clone()\n mask = torch.rand(f_condition[k].shape[:2]) < uncond_prob\n f_condition[k][mask] = 0.0\n return f_condition\n\n\ndef compute_extra_incam_loss(inputs, outputs, ppl):\n model_output = outputs[\"model_output\"]\n decode_dict = outputs[\"decode_dict\"]\n endecoder = ppl.endecoder\n weights = ppl.weights\n args = ppl.args\n\n extra_loss_dict = {}\n extra_loss = 0\n mask = inputs[\"mask\"][\"valid\"] # effective length mask\n mask_reproj = ~inputs[\"mask\"][\"spv_incam_only\"] # do not supervise reproj for 3DPW\n\n # Incam FK\n # prediction\n pred_c_j3d = endecoder.fk_v2(**outputs[\"pred_smpl_params_incam\"])\n pred_cr_j3d = pred_c_j3d - pred_c_j3d[:, :, :1] # (B, L, J, 3)\n\n # gt\n gt_c_j3d = endecoder.fk_v2(**inputs[\"smpl_params_c\"]) # (B, L, J, 3)\n gt_cr_j3d = gt_c_j3d - gt_c_j3d[:, :, :1] # (B, L, J, 3)\n\n # Root aligned C-MPJPE Loss\n if weights.cr_j3d > 0.0:\n cr_j3d_loss = F.mse_loss(pred_cr_j3d, gt_cr_j3d, reduction=\"none\")\n cr_j3d_loss = (cr_j3d_loss * mask[..., None, None]).mean()\n extra_loss += cr_j3d_loss * weights.cr_j3d\n extra_loss_dict[\"cr_j3d_loss\"] = cr_j3d_loss\n\n # Reprojection (to align with image)\n if weights.transl_c > 0.0:\n # pred_transl = decode_dict[\"transl\"] # (B, L, 3)\n # gt_transl = inputs[\"smpl_params_c\"][\"transl\"]\n # transl_c_loss = F.l1_loss(pred_transl, gt_transl, reduction=\"none\")\n # transl_c_loss = (transl_c_loss * mask[..., None]).mean()\n\n # Instead of supervising transl, we convert gt to pred_cam (prevent divide 0)\n pred_cam = model_output[\"pred_cam\"] # (B, L, 3)\n gt_transl = inputs[\"smpl_params_c\"][\"transl\"] # (B, L, 3)\n gt_pred_cam = get_a_pred_cam(gt_transl, inputs[\"bbx_xys\"], inputs[\"K_fullimg\"]) # (B, L, 3)\n gt_pred_cam[gt_pred_cam.isinf()] = -1 # this will be handled by valid_mask\n # (compute_transl_full_cam(gt_pred_cam, inputs[\"bbx_xys\"], inputs[\"K_fullimg\"]) - gt_transl).abs().max()\n\n # Skip gts that are not good during random construction\n gt_j3d_z_min = inputs[\"gt_j3d\"][..., 2].min(dim=-1)[0]\n valid_mask = (\n (gt_j3d_z_min > 0.3)\n * (gt_pred_cam[..., 0] > 0.3)\n * (gt_pred_cam[..., 0] < 5.0)\n * (gt_pred_cam[..., 1] > -3.0)\n * (gt_pred_cam[..., 1] < 3.0)\n * (gt_pred_cam[..., 2] > -3.0)\n * (gt_pred_cam[..., 2] < 3.0)\n * (inputs[\"bbx_xys\"][..., 2] > 0)\n )[..., None]\n transl_c_loss = F.mse_loss(pred_cam, gt_pred_cam, reduction=\"none\")\n transl_c_loss = (transl_c_loss * mask[..., None] * valid_mask).mean()\n\n extra_loss_dict[\"transl_c_loss\"] = transl_c_loss\n extra_loss += transl_c_loss * weights.transl_c\n\n if weights.j2d > 0.0:\n # prevent divide 0 or small value to overflow(fp16)\n reproj_z_thr = 0.3\n pred_c_j3d_z0_mask = pred_c_j3d[..., 2].abs() <= reproj_z_thr\n pred_c_j3d[pred_c_j3d_z0_mask] = reproj_z_thr\n gt_c_j3d_z0_mask = gt_c_j3d[..., 2].abs() <= reproj_z_thr\n gt_c_j3d[gt_c_j3d_z0_mask] = reproj_z_thr\n\n pred_j2d_01 = project_to_bi01(pred_c_j3d, inputs[\"bbx_xys\"], inputs[\"K_fullimg\"])\n gt_j2d_01 = project_to_bi01(gt_c_j3d, inputs[\"bbx_xys\"], inputs[\"K_fullimg\"]) # (B, L, J, 2)\n\n valid_mask = (\n (gt_c_j3d[..., 2] > reproj_z_thr)\n * (pred_c_j3d[..., 2] > reproj_z_thr) # Be safe\n * (gt_j2d_01[..., 0] > 0.0)\n * (gt_j2d_01[..., 0] < 1.0)\n * (gt_j2d_01[..., 1] > 0.0)\n * (gt_j2d_01[..., 1] < 1.0)\n )[..., None]\n valid_mask[~mask_reproj] = False # Do not supervise on 3dpw\n j2d_loss = F.mse_loss(pred_j2d_01, gt_j2d_01, reduction=\"none\")\n j2d_loss = (j2d_loss * mask[..., None, None] * valid_mask).mean()\n\n extra_loss += j2d_loss * weights.j2d\n extra_loss_dict[\"j2d_loss\"] = j2d_loss\n\n if weights.cr_verts > 0:\n # SMPL forward\n pred_c_verts437, pred_c_j17 = endecoder.smplx_model(**outputs[\"pred_smpl_params_incam\"])\n root_ = pred_c_j17[:, :, [11, 12], :].mean(-2, keepdim=True)\n pred_cr_verts437 = pred_c_verts437 - root_\n\n gt_cr_verts437 = inputs[\"gt_cr_verts437\"] # (B, L, 437, 3)\n cr_vert_loss = F.mse_loss(pred_cr_verts437, gt_cr_verts437, reduction=\"none\")\n cr_vert_loss = (cr_vert_loss * mask[:, :, None, None]).mean()\n extra_loss += cr_vert_loss * weights.cr_verts\n extra_loss_dict[\"cr_vert_loss\"] = cr_vert_loss\n\n if weights.verts2d > 0:\n gt_c_verts437 = inputs[\"gt_c_verts437\"] # (B, L, 437, 3)\n\n # prevent divide 0 or small value to overflow(fp16)\n reproj_z_thr = 0.3\n pred_c_verts437_z0_mask = pred_c_verts437[..., 2].abs() <= reproj_z_thr\n pred_c_verts437[pred_c_verts437_z0_mask] = reproj_z_thr\n gt_c_verts437_z0_mask = gt_c_verts437[..., 2].abs() <= reproj_z_thr\n gt_c_verts437[gt_c_verts437_z0_mask] = reproj_z_thr\n\n pred_verts2d_01 = project_to_bi01(pred_c_verts437, inputs[\"bbx_xys\"], inputs[\"K_fullimg\"])\n gt_verts2d_01 = project_to_bi01(gt_c_verts437, inputs[\"bbx_xys\"], inputs[\"K_fullimg\"]) # (B, L, 437, 2)\n\n valid_mask = (\n (gt_c_verts437[..., 2] > reproj_z_thr)\n * (pred_c_verts437[..., 2] > reproj_z_thr) # Be safe\n * (gt_verts2d_01[..., 0] > 0.0)\n * (gt_verts2d_01[..., 0] < 1.0)\n * (gt_verts2d_01[..., 1] > 0.0)\n * (gt_verts2d_01[..., 1] < 1.0)\n )[..., None]\n valid_mask[~mask_reproj] = False # Do not supervise on 3dpw\n verts2d_loss = F.mse_loss(pred_verts2d_01, gt_verts2d_01, reduction=\"none\")\n verts2d_loss = (verts2d_loss * mask[..., None, None] * valid_mask).mean()\n\n extra_loss += verts2d_loss * weights.verts2d\n extra_loss_dict[\"verts2d_loss\"] = verts2d_loss\n\n return extra_loss, extra_loss_dict\n\n\ndef compute_extra_global_loss(inputs, outputs, ppl):\n decode_dict = outputs[\"decode_dict\"]\n endecoder = ppl.endecoder\n weights = ppl.weights\n args = ppl.args\n\n extra_loss_dict = {}\n extra_loss = 0\n mask = inputs[\"mask\"][\"valid\"].clone() # (B, L)\n mask[inputs[\"mask\"][\"spv_incam_only\"]] = False\n\n if weights.transl_w > 0:\n # compute pred_transl_w by rollout\n gt_transl_w = inputs[\"smpl_params_w\"][\"transl\"]\n gt_global_orient_w = inputs[\"smpl_params_w\"][\"global_orient\"]\n local_transl_vel = decode_dict[\"local_transl_vel\"]\n pred_transl_w = rollout_local_transl_vel(local_transl_vel, gt_global_orient_w, gt_transl_w[:, [0]])\n\n trans_w_loss = F.l1_loss(pred_transl_w, gt_transl_w, reduction=\"none\")\n trans_w_loss = (trans_w_loss * mask[..., None]).mean()\n extra_loss += trans_w_loss * weights.transl_w\n extra_loss_dict[\"transl_w_loss\"] = trans_w_loss\n\n # Static-Conf loss\n if weights.static_conf_bce > 0:\n # Compute gt by thresholding velocity\n vel_thr = args.static_conf.vel_thr\n assert vel_thr > 0\n joint_ids = [7, 10, 8, 11, 20, 21] # [L_Ankle, L_foot, R_Ankle, R_foot, L_wrist, R_wrist]\n gt_w_j3d = endecoder.fk_v2(**inputs[\"smpl_params_w\"]) # (B, L, J=22, 3)\n static_gt = get_static_joint_mask(gt_w_j3d, vel_thr=vel_thr, repeat_last=True) # (B, L, J)\n static_gt = static_gt[:, :, joint_ids].float() # (B, L, J')\n pred_static_conf_logits = outputs[\"model_output\"][\"static_conf_logits\"]\n\n static_conf_loss = F.binary_cross_entropy_with_logits(pred_static_conf_logits, static_gt, reduction=\"none\")\n static_conf_loss = (static_conf_loss * mask[..., None]).mean()\n extra_loss += static_conf_loss * weights.static_conf_bce\n extra_loss_dict[\"static_conf_loss\"] = static_conf_loss\n\n return extra_loss, extra_loss_dict\n\n\n@autocast(enabled=False)\ndef get_smpl_params_w_Rt_v2(\n global_orient_gv,\n local_transl_vel,\n global_orient_c,\n cam_angvel,\n):\n \"\"\"Get global R,t in GV0(ay)\n Args:\n cam_angvel: (B, L, 6), defined as R @ R_{w2c}^{t} = R_{w2c}^{t+1}\n \"\"\"\n\n # Get R_ct_to_c0 from cam_angvel\n def as_identity(R):\n is_I = matrix_to_axis_angle(R).norm(dim=-1) < 1e-5\n R[is_I] = torch.eye(3)[None].expand(is_I.sum(), -1, -1).to(R)\n return R\n\n B = cam_angvel.shape[0]\n R_t_to_tp1 = rotation_6d_to_matrix(cam_angvel) # (B, L, 3, 3)\n R_t_to_tp1 = as_identity(R_t_to_tp1)\n\n # Get R_c2gv\n R_gv = axis_angle_to_matrix(global_orient_gv) # (B, L, 3, 3)\n R_c = axis_angle_to_matrix(global_orient_c) # (B, L, 3, 3)\n\n # Camera view direction in GV coordinate: Rc2gv @ [0,0,1]\n R_c2gv = R_gv @ R_c.mT\n view_axis_gv = R_c2gv[:, :, :, 2] # (B, L, 3) Rc2gv is estimated, so the x-axis is not accurate, i.e. != 0\n\n # Rotate axis use camera relative rotation\n R_cnext2gv = R_c2gv @ R_t_to_tp1.mT\n view_axis_gv_next = R_cnext2gv[..., 2]\n\n vec1_xyz = view_axis_gv.clone()\n vec1_xyz[..., 1] = 0\n vec1_xyz = F.normalize(vec1_xyz, dim=-1)\n vec2_xyz = view_axis_gv_next.clone()\n vec2_xyz[..., 1] = 0\n vec2_xyz = F.normalize(vec2_xyz, dim=-1)\n\n aa_tp1_to_t = vec2_xyz.cross(vec1_xyz, dim=-1)\n aa_tp1_to_t_angle = torch.acos(torch.clamp((vec1_xyz * vec2_xyz).sum(dim=-1, keepdim=True), -1.0, 1.0))\n aa_tp1_to_t = F.normalize(aa_tp1_to_t, dim=-1) * aa_tp1_to_t_angle\n\n aa_tp1_to_t = gaussian_smooth(aa_tp1_to_t, dim=-2) # Smooth\n R_tp1_to_t = axis_angle_to_matrix(aa_tp1_to_t).mT # (B, L, 3)\n\n # Get R_t_to_0\n R_t_to_0 = [torch.eye(3)[None].expand(B, -1, -1).to(R_t_to_tp1)]\n for i in range(1, R_t_to_tp1.shape[1]):\n R_t_to_0.append(R_t_to_0[-1] @ R_tp1_to_t[:, i])\n R_t_to_0 = torch.stack(R_t_to_0, dim=1) # (B, L, 3, 3)\n R_t_to_0 = as_identity(R_t_to_0)\n\n global_orient = matrix_to_axis_angle(R_t_to_0 @ R_gv)\n\n # Rollout to global transl\n # Start from transl0, in gv0 -> flip y-axis of gv0\n transl = rollout_local_transl_vel(local_transl_vel, global_orient)\n global_orient, transl, _ = get_tgtcoord_rootparam(global_orient, transl, tsf=\"any->ay\")\n\n smpl_params_w_Rt = {\"global_orient\": global_orient, \"transl\": transl}\n return smpl_params_w_Rt\n" }, { "path": "eval/GVHMR/hmr4d/model/gvhmr/utils/endecoder.py", "content": "import torch\nimport torch.nn as nn\nfrom pytorch3d.transforms import (\n rotation_6d_to_matrix,\n matrix_to_axis_angle,\n axis_angle_to_matrix,\n matrix_to_rotation_6d,\n matrix_to_quaternion,\n quaternion_to_matrix,\n)\nfrom hmr4d.configs import MainStore, builds\nfrom hmr4d.utils.geo.augment_noisy_pose import gaussian_augment\nimport hmr4d.utils.matrix as matrix\nfrom hmr4d.utils.pylogger import Log\nfrom hmr4d.utils.geo.hmr_global import get_local_transl_vel, rollout_local_transl_vel\nfrom hmr4d.utils.smplx_utils import make_smplx\nfrom . import stats_compose\n\n\nclass EnDecoder(nn.Module):\n def __init__(self, stats_name=\"DEFAULT_01\", noise_pose_k=10):\n super().__init__()\n # Load mean, std\n stats = getattr(stats_compose, stats_name)\n Log.info(f\"[EnDecoder] Use {stats_name} for statistics!\")\n self.register_buffer(\"mean\", torch.tensor(stats[\"mean\"]).float(), False)\n self.register_buffer(\"std\", torch.tensor(stats[\"std\"]).float(), False)\n\n # option\n self.noise_pose_k = noise_pose_k\n\n # smpl\n self.smplx_model = make_smplx(\"supermotion_v437coco17\")\n parents = self.smplx_model.parents[:22]\n self.register_buffer(\"parents_tensor\", parents, False)\n self.parents = parents.tolist()\n\n def get_noisyobs(self, data, return_type=\"r6d\"):\n \"\"\"\n Noisy observation contains local pose with noise\n Args:\n data (dict):\n body_pose: (B, L, J*3) or (B, L, J, 3)\n Returns:\n noisy_bosy_pose: (B, L, J, 6) or (B, L, J, 3) or (B, L, 3, 3) depends on return_type\n \"\"\"\n body_pose = data[\"body_pose\"] # (B, L, 63)\n B, L, _ = body_pose.shape\n body_pose = body_pose.reshape(B, L, -1, 3)\n\n # (B, L, J, C)\n return_mapping = {\"R\": 0, \"r6d\": 1, \"aa\": 2}\n return_id = return_mapping[return_type]\n noisy_bosy_pose = gaussian_augment(body_pose, self.noise_pose_k, to_R=True)[return_id]\n return noisy_bosy_pose\n\n def normalize_body_pose_r6d(self, body_pose_r6d):\n \"\"\"body_pose_r6d: (B, L, {J*6}/{J, 6}) -> (B, L, J*6)\"\"\"\n B, L = body_pose_r6d.shape[:2]\n body_pose_r6d = body_pose_r6d.reshape(B, L, -1)\n if self.mean.shape[-1] == 1: # no mean, std provided\n return body_pose_r6d\n body_pose_r6d = (body_pose_r6d - self.mean[:126]) / self.std[:126] # (B, L, C)\n return body_pose_r6d\n\n def fk_v2(self, body_pose, betas, global_orient=None, transl=None, get_intermediate=False):\n \"\"\"\n Args:\n body_pose: (B, L, 63)\n betas: (B, L, 10)\n global_orient: (B, L, 3)\n Returns:\n joints: (B, L, 22, 3)\n \"\"\"\n B, L = body_pose.shape[:2]\n if global_orient is None:\n global_orient = torch.zeros((B, L, 3), device=body_pose.device)\n aa = torch.cat([global_orient, body_pose], dim=-1).reshape(B, L, -1, 3)\n rotmat = axis_angle_to_matrix(aa) # (B, L, 22, 3, 3)\n\n skeleton = self.smplx_model.get_skeleton(betas)[..., :22, :] # (B, L, 22, 3)\n local_skeleton = skeleton - skeleton[:, :, self.parents_tensor]\n local_skeleton = torch.cat([skeleton[:, :, :1], local_skeleton[:, :, 1:]], dim=2)\n\n if transl is not None:\n local_skeleton[..., 0, :] += transl # B, L, 22, 3\n\n mat = matrix.get_TRS(rotmat, local_skeleton) # B, L, 22, 4, 4\n fk_mat = matrix.forward_kinematics(mat, self.parents) # B, L, 22, 4, 4\n joints = matrix.get_position(fk_mat) # B, L, 22, 3\n if not get_intermediate:\n return joints\n else:\n return joints, mat, fk_mat\n\n def get_local_pos(self, betas):\n skeleton = self.smplx_model.get_skeleton(betas)[..., :22, :] # (B, L, 22, 3)\n local_skeleton = skeleton - skeleton[:, :, self.parents_tensor]\n local_skeleton = torch.cat([skeleton[:, :, :1], local_skeleton[:, :, 1:]], dim=2)\n return local_skeleton\n\n def encode(self, inputs):\n \"\"\"\n definition: {\n body_pose_r6d, # (B, L, (J-1)*6) -> 0:126\n betas, # (B, L, 10) -> 126:136\n global_orient_r6d, # (B, L, 6) -> 136:142 incam\n global_orient_gv_r6d: # (B, L, 6) -> 142:148 gv\n local_transl_vel, # (B, L, 3) -> 148:151, smpl-coord\n }\n \"\"\"\n B, L = inputs[\"smpl_params_c\"][\"body_pose\"].shape[:2]\n # cam\n smpl_params_c = inputs[\"smpl_params_c\"]\n body_pose = smpl_params_c[\"body_pose\"].reshape(B, L, 21, 3)\n body_pose_r6d = matrix_to_rotation_6d(axis_angle_to_matrix(body_pose)).flatten(-2)\n betas = smpl_params_c[\"betas\"]\n global_orient_R = axis_angle_to_matrix(smpl_params_c[\"global_orient\"])\n global_orient_r6d = matrix_to_rotation_6d(global_orient_R)\n\n # global\n R_c2gv = inputs[\"R_c2gv\"] # (B, L, 3, 3)\n global_orient_gv_r6d = matrix_to_rotation_6d(R_c2gv @ global_orient_R)\n\n # local_transl_vel\n smpl_params_w = inputs[\"smpl_params_w\"]\n local_transl_vel = get_local_transl_vel(smpl_params_w[\"transl\"], smpl_params_w[\"global_orient\"])\n if False: # debug\n transl_recover = rollout_local_transl_vel(\n local_transl_vel, smpl_params_w[\"global_orient\"], smpl_params_w[\"transl\"][:, [0]]\n )\n print((transl_recover - smpl_params_w[\"transl\"]).abs().max())\n\n # returns\n x = torch.cat([body_pose_r6d, betas, global_orient_r6d, global_orient_gv_r6d, local_transl_vel], dim=-1)\n x_norm = (x - self.mean) / self.std\n return x_norm\n\n def encode_translw(self, inputs):\n \"\"\"\n definition: {\n body_pose_r6d, # (B, L, (J-1)*6) -> 0:126\n betas, # (B, L, 10) -> 126:136\n global_orient_r6d, # (B, L, 6) -> 136:142 incam\n global_orient_gv_r6d: # (B, L, 6) -> 142:148 gv\n local_transl_vel, # (B, L, 3) -> 148:151, smpl-coord\n }\n \"\"\"\n # local_transl_vel\n smpl_params_w = inputs[\"smpl_params_w\"]\n local_transl_vel = get_local_transl_vel(smpl_params_w[\"transl\"], smpl_params_w[\"global_orient\"])\n\n # returns\n x = local_transl_vel\n x_norm = (x - self.mean[-3:]) / self.std[-3:]\n return x_norm\n\n def decode_translw(self, x_norm):\n return x_norm * self.std[-3:] + self.mean[-3:]\n\n def decode(self, x_norm):\n \"\"\"x_norm: (B, L, C)\"\"\"\n B, L, C = x_norm.shape\n x = (x_norm * self.std) + self.mean\n\n body_pose_r6d = x[:, :, :126]\n betas = x[:, :, 126:136]\n global_orient_r6d = x[:, :, 136:142]\n global_orient_gv_r6d = x[:, :, 142:148]\n local_transl_vel = x[:, :, 148:151]\n\n body_pose = matrix_to_axis_angle(rotation_6d_to_matrix(body_pose_r6d.reshape(B, L, -1, 6)))\n body_pose = body_pose.flatten(-2)\n global_orient_c = matrix_to_axis_angle(rotation_6d_to_matrix(global_orient_r6d))\n global_orient_gv = matrix_to_axis_angle(rotation_6d_to_matrix(global_orient_gv_r6d))\n\n output = {\n \"body_pose\": body_pose,\n \"betas\": betas,\n \"global_orient\": global_orient_c,\n \"global_orient_gv\": global_orient_gv,\n \"local_transl_vel\": local_transl_vel,\n }\n\n return output\n\n\ngroup_name = \"endecoder/gvhmr\"\ncfg_base = builds(EnDecoder, populate_full_signature=True)\nMainStore.store(name=\"v1_no_stdmean\", node=cfg_base, group=group_name)\nMainStore.store(name=\"v1\", node=cfg_base(stats_name=\"MM_V1\"), group=group_name)\nMainStore.store(\n name=\"v1_amass_local_bedlam_cam\",\n node=cfg_base(stats_name=\"MM_V1_AMASS_LOCAL_BEDLAM_CAM\"),\n group=group_name,\n)\n\nMainStore.store(name=\"v2\", node=cfg_base(stats_name=\"MM_V2\"), group=group_name)\nMainStore.store(name=\"v2_1\", node=cfg_base(stats_name=\"MM_V2_1\"), group=group_name)\n" }, { "path": "eval/GVHMR/hmr4d/model/gvhmr/utils/postprocess.py", "content": "import torch\nfrom torch.cuda.amp import autocast\nfrom pytorch3d.transforms import (\n matrix_to_rotation_6d,\n rotation_6d_to_matrix,\n axis_angle_to_matrix,\n matrix_to_axis_angle,\n)\n\nimport hmr4d.utils.matrix as matrix\nfrom hmr4d.utils.ik.ccd_ik import CCD_IK\nfrom hmr4d.utils.geo_transform import get_sequence_cammat, transform_mat, apply_T_on_points\nfrom hmr4d.utils.net_utils import gaussian_smooth\nfrom hmr4d.model.gvhmr.utils.endecoder import EnDecoder\n\nfrom hmr4d.utils.wis3d_utils import make_wis3d, add_motion_as_lines\n\n\n@autocast(enabled=False)\ndef pp_static_joint(outputs, endecoder: EnDecoder):\n # Global FK\n pred_w_j3d = endecoder.fk_v2(**outputs[\"pred_smpl_params_global\"])\n L = pred_w_j3d.shape[1]\n joint_ids = [7, 10, 8, 11, 20, 21] # [L_Ankle, L_foot, R_Ankle, R_foot, L_wrist, R_wrist]\n pred_j3d_static = pred_w_j3d.clone()[:, :, joint_ids] # (B, L, J, 3)\n\n ######## update overall movement with static info, and make displacement ~[0,0,0]\n pred_j_disp = pred_j3d_static[:, 1:] - pred_j3d_static[:, :-1] # (B, L-1, J, 3)\n\n static_conf_logits = outputs[\"static_conf_logits\"][:, :-1].clone()\n static_label_ = static_conf_logits > 0 # (B, L-1, J) # avoid non-contact frame\n static_conf_logits = static_conf_logits.float() - (~static_label_ * 1e6) # fp16 cannot go through softmax\n is_static = static_label_.sum(dim=-1) > 0 # (B, L-1)\n\n pred_disp = pred_j_disp * static_conf_logits[..., None].softmax(dim=-2) # (B, L-1, J, 3)\n pred_disp = pred_disp * is_static[..., None, None] # (B, L-1, J, 3)\n pred_disp = pred_disp.sum(-2) # (B, L-1, 3)\n ####################\n\n # Overwrite results:\n if False: # for-loop\n post_w_transl = outputs[\"pred_smpl_params_global\"][\"transl\"].clone() # (B, L, 3)\n for i in range(1, L):\n post_w_transl[:, i:] -= pred_disp[:, i - 1 : i]\n else: # vectorized\n pred_w_transl = outputs[\"pred_smpl_params_global\"][\"transl\"].clone() # (B, L, 3)\n pred_w_disp = pred_w_transl[:, 1:] - pred_w_transl[:, :-1] # (B, L-1, 3)\n pred_w_disp_new = pred_w_disp - pred_disp\n post_w_transl = torch.cumsum(torch.cat([pred_w_transl[:, :1], pred_w_disp_new], dim=1), dim=1)\n post_w_transl[..., 0] = gaussian_smooth(post_w_transl[..., 0], dim=-1)\n post_w_transl[..., 2] = gaussian_smooth(post_w_transl[..., 2], dim=-1)\n\n # Put the sequence on the ground by -min(y), this does not consider foot height, for o3d vis\n post_w_j3d = pred_w_j3d - pred_w_transl.unsqueeze(-2) + post_w_transl.unsqueeze(-2)\n ground_y = post_w_j3d[..., 1].flatten(-2).min(dim=-1)[0] # (B,) Minimum y value\n post_w_transl[..., 1] -= ground_y\n\n return post_w_transl\n\n\n@autocast(enabled=False)\ndef pp_static_joint_cam(outputs, endecoder: EnDecoder):\n \"\"\"Use static joint and static camera assumption to postprocess the global transl\"\"\"\n # input\n pred_smpl_params_incam = outputs[\"pred_smpl_params_incam\"].copy()\n pred_smpl_params_global = outputs[\"pred_smpl_params_global\"]\n static_conf_logits = outputs[\"static_conf_logits\"].clone()[:, :-1] # (B, L-1, J)\n joint_ids = [7, 10, 8, 11, 20, 21] # [L_Ankle, L_foot, R_Ankle, R_foot, L_wrist, R_wrist]\n B, L = pred_smpl_params_incam[\"transl\"].shape[:2]\n assert B == 1\n\n # FK\n pred_w_j3d = endecoder.fk_v2(**pred_smpl_params_global) # (B, L, J, 3)\n # smooth incam results, as this could be noisy\n pred_smpl_params_incam[\"transl\"] = gaussian_smooth(pred_smpl_params_incam[\"transl\"], sigma=5, dim=-2)\n pred_c_j3d = endecoder.fk_v2(**pred_smpl_params_incam) # (B, L, J, 3)\n\n # compute T_c2w (static) from first frame\n R_gv = axis_angle_to_matrix(pred_smpl_params_global[\"global_orient\"][:, 0]) # (B, 3, 3)\n R_c = axis_angle_to_matrix(pred_smpl_params_incam[\"global_orient\"][:, 0]) # (B, 3, 3)\n R_c2w = R_gv @ R_c.mT # (B, 3, 3)\n t_c2w = pred_w_j3d[:, 0, 0] - torch.einsum(\"bij,bj->bi\", R_c2w, pred_c_j3d[:, 0, 0]) # (B, 3)\n T_c2w = transform_mat(R_c2w, t_c2w) # (B, 4, 4)\n pred_c_j3d_in_w = apply_T_on_points(pred_c_j3d, T_c2w[:, None])\n\n # 1. Make transl similar to incam\n post_w_transl = pred_smpl_params_global[\"transl\"].clone() # (B, L, 3)\n post_w_j3d = pred_w_j3d.clone() # (B, L, J, 3)\n cp_thr = torch.tensor([0.25, 0.25, 0.25]).to(post_w_j3d) # Only update very bad pred\n for i in range(1, L):\n cp_diff = post_w_j3d[:, i, 0] - pred_c_j3d_in_w[:, i, 0] # (B, 3)\n cp_diff = cp_diff * ~((cp_diff > -cp_thr) * (cp_diff < cp_thr))\n cp_diff = torch.clamp(cp_diff, -0.02, 0.02)\n post_w_transl[:, i:] -= cp_diff\n post_w_j3d[:, i:] -= (cp_diff)[:, None, None]\n\n # 1. Make stationary joint stay stationary\n # pred_j3d_static = pred_w_j3d.clone()[:, :, joint_ids] # (B, L, J, 3)\n pred_j3d_static = post_w_j3d[:, :, joint_ids] # (B, L, J, 3)\n pred_j_disp = pred_j3d_static[:, 1:] - pred_j3d_static[:, :-1] # (B, L-1, J, 3)\n\n static_label = static_conf_logits.sigmoid() > 0.8 # (B, L-1, J)\n static_label_sumJ = static_label.sum(-1, keepdim=True) # (B, L-1, 1)\n static_label_sumJ = torch.clamp_min(static_label_sumJ, 1) # replace 0 with 1\n pred_disp_sumJ = (pred_j_disp * static_label[..., None]).sum(-2) # (B, L-1, 3)\n pred_disp = pred_disp_sumJ / static_label_sumJ # (B, L-1, 3)\n pred_disp[:, :, 1] = 0 # do not modify y\n\n # Overwrite results (for-loop)\n for i in range(1, L):\n post_w_transl[:, i:] -= pred_disp[:, [i - 1]]\n post_w_j3d[:, i:] -= pred_disp[:, [i - 1], None]\n\n # Put the sequence on the ground by -min(y), this does not consider foot height, for o3d vis\n ground_y = post_w_j3d[..., 1].flatten(-2).min(dim=-1)[0] # (B,) Minimum y value\n post_w_transl[..., 1] -= ground_y\n\n return post_w_transl\n\n\n@autocast(enabled=False)\ndef process_ik(outputs, endecoder):\n static_conf = outputs[\"static_conf_logits\"].sigmoid() # (B, L, J)\n post_w_j3d, local_mat, post_w_mat = endecoder.fk_v2(**outputs[\"pred_smpl_params_global\"], get_intermediate=True)\n\n # sebas rollout merge\n joint_ids = [7, 10, 8, 11, 20, 21] # [L_Ankle, L_foot, R_Ankle, R_foot, L_wrist, R_wrist]\n post_target_j3d = post_w_j3d.clone()\n for i in range(1, post_w_j3d.size(1)):\n prev = post_target_j3d[:, i - 1, joint_ids]\n this = post_w_j3d[:, i, joint_ids]\n c_prev = static_conf[:, i - 1, :, None]\n post_target_j3d[:, i, joint_ids] = prev * c_prev + this * (1 - c_prev)\n\n # ik\n global_rot = matrix.get_rotation(post_w_mat)\n parents = [-1, 0, 0, 0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 9, 9, 12, 13, 14, 16, 17, 18, 19]\n left_leg_chain = [0, 1, 4, 7, 10]\n right_leg_chain = [0, 2, 5, 8, 11]\n left_hand_chain = [9, 13, 16, 18, 20]\n right_hand_chain = [9, 14, 17, 19, 21]\n\n def ik(local_mat, target_pos, target_rot, target_ind, chain):\n local_mat = local_mat.clone()\n IK_solver = CCD_IK(\n local_mat,\n parents,\n target_ind,\n target_pos,\n target_rot,\n kinematic_chain=chain,\n max_iter=2,\n )\n\n chain_local_mat = IK_solver.solve()\n chain_rotmat = matrix.get_rotation(chain_local_mat)\n local_mat[:, :, chain[1:], :-1, :-1] = chain_rotmat[:, :, 1:] # (B, L, J, 3, 3)\n return local_mat\n\n local_mat = ik(local_mat, post_target_j3d[:, :, [7, 10]], global_rot[:, :, [7, 10]], [3, 4], left_leg_chain)\n local_mat = ik(local_mat, post_target_j3d[:, :, [8, 11]], global_rot[:, :, [8, 11]], [3, 4], right_leg_chain)\n local_mat = ik(local_mat, post_target_j3d[:, :, [20]], global_rot[:, :, [20]], [4], left_hand_chain)\n local_mat = ik(local_mat, post_target_j3d[:, :, [21]], global_rot[:, :, [21]], [4], right_hand_chain)\n\n body_pose = matrix_to_axis_angle(matrix.get_rotation(local_mat[:, :, 1:])) # (B, L, J-1, 3, 3)\n body_pose = body_pose.flatten(2) # (B, L, (J-1)*3)\n\n return body_pose\n" }, { "path": "eval/GVHMR/hmr4d/model/gvhmr/utils/stats_compose.py", "content": "# fmt:off\nbody_pose_r6d = {\n \"bedlam\": {\n \"count\": 5417929,\n \"mean\": [ 0.9772, -0.0925, 0.0028, 0.1058, 0.9111, 0.1373, 0.9796, 0.0711,\n -0.0193, -0.0816, 0.8910, 0.1953, 0.9935, 0.0072, 0.0270, -0.0046,\n 0.9200, -0.2511, 0.9752, 0.0477, -0.0990, -0.0613, 0.8242, -0.2730,\n 0.9836, -0.0400, 0.0067, 0.0148, 0.7836, -0.3471, 0.9931, -0.0300,\n -0.0469, 0.0244, 0.9825, -0.0513, 0.9777, 0.0206, 0.1444, -0.0470,\n 0.9603, 0.1521, 0.9804, -0.0362, -0.0902, 0.0500, 0.9546, 0.1337,\n 0.9969, -0.0105, 0.0076, 0.0090, 0.9914, 0.0150, 0.9953, -0.0607,\n 0.0089, 0.0602, 0.9942, 0.0146, 0.9934, -0.0682, -0.0171, 0.0680,\n 0.9932, -0.0017, 0.9790, 0.0294, 0.0065, -0.0338, 0.9706, -0.0456,\n 0.9056, 0.2457, -0.1029, -0.2279, 0.9262, 0.0145, 0.9233, -0.1301,\n 0.1550, 0.1140, 0.9476, 0.0534, 0.9769, -0.0572, -0.0095, 0.0569,\n 0.9690, 0.0472, 0.6782, 0.5746, -0.2378, -0.5546, 0.7212, 0.0917,\n 0.6489, -0.5955, 0.2424, 0.5821, 0.6797, 0.0563, 0.5562, -0.1252,\n -0.5860, 0.0937, 0.9176, -0.1287, 0.4453, 0.1421, 0.6119, -0.1427,\n 0.8996, -0.1136, 0.9186, -0.0881, -0.1463, 0.1087, 0.8692, 0.0845,\n 0.9175, 0.0257, 0.0663, -0.0385, 0.8603, 0.1020],\n \"std\": [0.0429, 0.1392, 0.1236, 0.1323, 0.1645, 0.3086, 0.0375, 0.1406, 0.1172,\n 0.1275, 0.1934, 0.3280, 0.0119, 0.0835, 0.0716, 0.0741, 0.1528, 0.2484,\n 0.0349, 0.0947, 0.1633, 0.0924, 0.3469, 0.3370, 0.0273, 0.1009, 0.1411,\n 0.0680, 0.3876, 0.3323, 0.0103, 0.0735, 0.0712, 0.0690, 0.0246, 0.1617,\n 0.0216, 0.1097, 0.1016, 0.0924, 0.0509, 0.2035, 0.0245, 0.1188, 0.1212,\n 0.1056, 0.0634, 0.2308, 0.0054, 0.0579, 0.0517, 0.0575, 0.0124, 0.1158,\n 0.0076, 0.0654, 0.0367, 0.0644, 0.0118, 0.0592, 0.0116, 0.0829, 0.0361,\n 0.0832, 0.0124, 0.0422, 0.0343, 0.1060, 0.1680, 0.1075, 0.0473, 0.2023,\n 0.0701, 0.2344, 0.2213, 0.2632, 0.0589, 0.1318, 0.0767, 0.2456, 0.2009,\n 0.2666, 0.0542, 0.1106, 0.0347, 0.1080, 0.1718, 0.1117, 0.0459, 0.2025,\n 0.1882, 0.2769, 0.2032, 0.3072, 0.1447, 0.2204, 0.2018, 0.2820, 0.2126,\n 0.3213, 0.1760, 0.2486, 0.4749, 0.1677, 0.2791, 0.2239, 0.0963, 0.2705,\n 0.5540, 0.1846, 0.2572, 0.2411, 0.1287, 0.2878, 0.1151, 0.2993, 0.1557,\n 0.2812, 0.1880, 0.3334, 0.1286, 0.3355, 0.1553, 0.3216, 0.1880, 0.3306]\n },\n \"amass\": {\n \"count\": 7114038,\n \"mean\": [ 9.6969e-01, -5.9719e-02, -3.7700e-02, 5.8256e-02, 9.0800e-01,\n 1.0972e-01, 9.7636e-01, 4.3401e-02, 4.3110e-03, -4.3032e-02,\n 9.0261e-01, 1.4478e-01, 9.9288e-01, 3.5673e-03, 1.6264e-02,\n -2.2260e-03, 9.3470e-01, -2.3495e-01, 9.7147e-01, 5.2553e-02,\n -9.3666e-02, -5.4550e-02, 8.3321e-01, -2.4246e-01, 9.7971e-01,\n -3.8429e-02, 5.3575e-03, 1.5537e-02, 8.1449e-01, -3.0926e-01,\n 9.9532e-01, -9.4398e-03, -3.8328e-02, 8.5141e-03, 9.8880e-01,\n 1.9976e-04, 9.5602e-01, -3.9528e-02, 2.0017e-01, 1.0363e-02,\n 9.5965e-01, 1.3770e-01, 9.6223e-01, -4.6278e-02, -1.5177e-01,\n 6.6705e-02, 9.5545e-01, 1.2519e-01, 9.9767e-01, -1.2616e-02,\n -2.5442e-04, 1.1661e-02, 9.9376e-01, -3.6222e-02, 9.9511e-01,\n -1.0583e-02, 1.2130e-02, 7.6461e-03, 9.9137e-01, 2.0029e-02,\n 9.9295e-01, 7.2917e-03, 4.9454e-03, -8.0286e-03, 9.9137e-01,\n 2.3707e-03, 9.7698e-01, 1.9943e-02, 1.3808e-03, -2.2006e-02,\n 9.7375e-01, -6.7936e-02, 9.2804e-01, 2.5005e-01, -5.7167e-02,\n -2.4047e-01, 9.4246e-01, 2.5863e-02, 9.2957e-01, -2.1329e-01,\n 1.1112e-01, 2.0741e-01, 9.4876e-01, 2.9901e-02, 9.7683e-01,\n -4.1210e-02, 2.3248e-03, 4.0967e-02, 9.7365e-01, 5.7309e-03,\n 6.4513e-01, 6.1999e-01, -2.5469e-01, -6.2342e-01, 6.8177e-01,\n 3.5524e-02, 6.6192e-01, -5.9341e-01, 2.7136e-01, 5.9269e-01,\n 6.8966e-01, 3.1309e-02, 6.8946e-01, -1.1676e-01, -4.9859e-01,\n 4.0969e-02, 9.3656e-01, -1.4875e-01, 6.2787e-01, 1.3793e-01,\n 5.4289e-01, -9.1946e-02, 9.2868e-01, -1.1927e-01, 9.3012e-01,\n -8.3810e-02, -1.1951e-01, 9.7211e-02, 8.9118e-01, 5.9887e-02,\n 9.3033e-01, 7.1047e-02, 7.5264e-02, -8.0679e-02, 8.8562e-01,\n 4.8960e-02],\n \"std\": [0.0612, 0.1390, 0.1779, 0.1415, 0.1826, 0.3268, 0.0440, 0.1382, 0.1542,\n 0.1348, 0.1930, 0.3272, 0.0132, 0.0801, 0.0855, 0.0729, 0.1255, 0.2238,\n 0.0554, 0.1088, 0.1727, 0.0939, 0.3294, 0.3559, 0.0532, 0.1082, 0.1554,\n 0.0768, 0.3446, 0.3407, 0.0120, 0.0650, 0.0584, 0.0632, 0.0198, 0.1335,\n 0.0631, 0.1250, 0.1574, 0.1047, 0.0730, 0.2091, 0.0759, 0.1241, 0.1667,\n 0.1112, 0.0831, 0.2185, 0.0060, 0.0441, 0.0502, 0.0441, 0.0102, 0.0946,\n 0.0237, 0.0722, 0.0610, 0.0738, 0.0479, 0.0949, 0.0369, 0.0943, 0.0610,\n 0.0966, 0.0498, 0.0729, 0.0425, 0.1001, 0.1824, 0.0972, 0.0408, 0.1887,\n 0.0594, 0.1842, 0.1884, 0.2020, 0.0457, 0.1018, 0.0640, 0.1990, 0.1854,\n 0.2133, 0.0467, 0.0910, 0.0392, 0.1049, 0.1776, 0.1037, 0.0413, 0.1945,\n 0.1733, 0.2612, 0.1905, 0.2963, 0.1512, 0.1861, 0.1710, 0.2663, 0.1896,\n 0.3135, 0.1568, 0.2219, 0.3976, 0.1594, 0.2810, 0.1855, 0.0845, 0.2398,\n 0.4398, 0.1629, 0.2685, 0.1990, 0.0998, 0.2556, 0.1137, 0.2837, 0.1419,\n 0.2761, 0.1678, 0.2973, 0.1172, 0.3010, 0.1394, 0.2910, 0.1724, 0.3039]\n }\n}\n\nbetas = {\n \"bedlam\": {\n \"count\": 37855, # so many subjects? \n \"mean\": [ 0.0378, -0.3562, 0.1185, 0.2245, 0.0204, 0.0929, 0.0537, 0.1006,\n -0.1180, 0.0936],\n \"std\":[0.8070, 1.3480, 0.8964, 0.7390, 0.6433, 0.6089, 0.5374, 0.6984, 0.7263,\n 0.5395],\n },\n \"amass\": {\n \"count\": 18086,\n \"mean\": [ 0.2310, 0.1750, 0.2931, -0.1859, -1.1163, -1.1028, -0.2573, 0.3555,\n 0.3732, 0.2852],\n \"std\": [0.8831, 0.7965, 1.0899, 1.1788, 1.2128, 1.1081, 0.9780, 1.1434, 0.8498,\n 1.1462],\n }\n}\n\nglobal_orient_c_r6d = {\n \"bedlam\": {\n \"count\": 5417929,\n \"mean\": [-4.9862e-03, -8.7136e-04, -1.4187e-03, 1.4825e-02, -9.4419e-01,\n -5.1653e-02],\n \"std\": [0.7048, 0.1713, 0.6884, 0.1548, 0.1546, 0.2403],\n },\n}\n\nglobal_orient_gv_r6d = {\n \"bedlam\": {\n \"count\": 5134187,\n \"mean\": [ 3.6018e-04, -2.2327e-04, 2.2316e-03, -4.4879e-02, -9.7435e-01,\n 1.0021e-01],\n \"std\": [0.6070, 0.5355, 0.5873, 0.6285, 0.2336, 0.7675],\n },\n}\n\nlocal_transl_vel = {\n \"none\":{\n \"mean\": [0., 0., 0.],\n \"std\": [1., 1., 1.]\n },\n \"1e-2\":{\n \"mean\": [0., 0., 0.],\n \"std\": [1e-2, 1e-2, 1e-2]\n },\n \"bedlam\": {\n \"count\": 5417929,\n \"mean\": [7.3057e-05, -2.2142e-04, 3.2444e-03],\n \"std\": [0.0065, 0.0091, 0.0114],\n },\n \"amass\": {\n \"count\": 7113068,\n \"mean\": [-0.0002, -0.0006, 0.0069],\n \"std\": [0.0064, 0.0070, 0.0138],\n },\n \"alignhead\":{\n \"count\": 7113068,\n \"mean\":[-2.0822e-04, -1.7966e-06, 6.9816e-03],\n \"std\":[0.0065, 0.0066, 0.0139],\n },\n \"alignhead_absy\":{\n \"count\": 7113068,\n \"mean\":[-0.0002, -0.0316, 0.0070],\n \"std\":[0.0065, 0.1351, 0.0139],\n },\n \"alignhead_absgy\":{\n \"count\": 7113068,\n \"mean\":[[-2.0822e-04, 1.2627e+00, 6.9816e-03]],\n \"std\":[0.0065, 0.1516, 0.0139],\n }\n\n}\n\npred_cam = {\n \"bedlam\": {\n \"count\": 5096332,\n \"mean\": [1.0606, -0.0027, 0.2702],\n \"std\": [0.1784, 0.0956, 0.0764],\n }\n}\n\nvitfeat = {\n \"bedlam\": {\n \"count\": 5546332,\n \"mean\": [-1.3772, 0.2490, 0.0602, -0.1834, 0.2458, 0.5372, 0.3343, -0.3476, -0.1017, -0.0362, -0.0678, 0.2150, -0.2534, 0.1029, 0.8199, -0.4676, 0.6259, -0.3350, 0.0549, -0.4469, 0.2751, -0.1763, 0.1114, -0.2115, -0.0264, 0.5294, 0.8212, -0.4562, 0.4147, -0.0256, -0.1019, 0.2798, 0.9284, 0.4652, 0.6365, 0.6785, -0.0765, 0.0337, -0.2566, -0.0335, -0.1799, 0.7426, 0.2810, -0.7121, -0.0893, 0.1608, -0.2483, 1.5094, -1.4395, -0.3682, -0.4157, -0.0032, -0.0376, -0.0043, 0.2092, 0.3038, -0.2077, -0.4868, -0.1534, 0.2668, 1.2773, 0.2838, -0.4863, -1.2300, 0.0581, -0.3041, 0.1518, 0.7955, -0.4293, 1.4666, 0.3077, 0.3918, 0.1418, 0.1590, 0.8671, -0.3527, 0.5629, 0.1414, 0.0964, -0.1094, -0.0211, -0.0937, 0.1606, -0.7900, 0.0397, 0.0570, 0.7083, -0.5732, 0.1430, -0.2571, 0.5275, 0.6603, 0.3265, 0.4574, -0.3361, -0.1267, 0.3841, 0.1758, -0.6207, -0.3673, 0.8914, 0.4297, -0.8118, 0.2229, -0.2876, 0.2460, 0.4856, -0.1446, -0.2416, 0.1229, 0.2865, 0.7023, -0.2883, 0.3940, -1.5496, 0.4456, 0.6445, 0.2058, -0.4265, 0.3724, 0.1557, -1.4208, -0.1246, 0.1237, -0.3965, 0.0105, -0.0780, 0.6448, -0.1132, 0.8500, -0.2828, 0.4447, 0.6257, -0.2664, -0.8384, -1.8091, -0.2769, 0.1866, 0.6051, -0.2548, 0.9823, -0.2985, -0.2773, -0.4383, 0.1886, 0.2411, 0.2546, 0.2195, -0.0041, 0.1038, -0.6804, 1.2364, 0.5393, 0.0351, 0.4537, -0.8044, -0.1993, -2.1097, -0.8458, 0.1497, 1.6042, 0.6458, -0.5455, 0.0778, 0.0504, -0.5242, -0.3215, -0.0199, 1.1461, -0.3355, -0.3421, -0.3951, 0.0184, -0.0261, 0.2048, 0.0080, 0.6553, -1.3221, 0.5140, 0.5958, -0.2523, 0.9434, -0.0727, 0.1978, 1.1105, -0.4992, 0.3990, 0.2074, 0.3843, -0.0444, 0.0624, -0.8442, -0.0724, -0.5328, 1.1723, 0.8043, 0.6674, 1.5283, 4.2502, 0.0935, 0.3733, 0.1569, 0.0154, 0.0674, 0.0862, -0.2744, -0.4537, 0.1588, -1.9156, 0.0149, -1.0498, -0.0790, 0.0851, -0.5007, 0.3323, -0.1065, 0.0782, 0.0725, -0.5921, -0.1876, 0.0094, -0.3631, 0.0951, 0.1318, 0.0936, 0.5668, -0.0875, -0.4576, -0.4306, 0.5458, 1.0761, 1.1740, -0.0337, 1.3718, -0.2913, -0.3433, 0.5338, -0.4577, -0.4966, 0.2704, 0.3236, 0.4053, 0.0360, 1.1616, -0.2012, 0.7373, 0.0779, -0.0280, -0.4426, 0.0450, 0.2923, 0.0161, -0.4788, 0.1924, -0.3012, 0.0298, -0.7776, -0.2215, 0.4494, -0.1677, 0.2214, 0.0762, -0.3088, 0.4230, 0.0673, -1.0233, 0.0748, -0.4358, -0.2497, -0.0066, 0.1679, -0.1077, -0.4290, 2.5254, -0.8819, -0.8073, 0.2535, 2.0680, -0.4715, 0.3614, -2.9281, 3.1536, 0.3118, -0.0239, 0.7064, -0.6935, -1.1070, -0.1715, -0.0920, -0.2133, -1.0173, 0.0084, -0.1721, 0.2605, -0.6607, -0.0788, -0.3479, -0.2187, 1.0605, 0.2857, 0.7464, 0.9612, -1.1332, 1.5708, -1.0264, 0.6070, 0.4103, -0.1950, -0.0629, -0.0958, -0.2199, -0.2198, -0.4019, 0.2478, -0.3576, 0.0191, -5.8435, 0.0145, -0.2312, 0.9872, 1.1159, 0.3775, 0.1960, -0.5968, -0.2611, -0.0634, -0.1003, 0.7411, -0.8298, -0.1743, 1.8418, 0.3692, -0.4321, 0.0613, -1.9046, 0.5812, 0.2805, 0.1703, -0.2212, -0.0740, -0.2737, -0.3084, 2.9787, -0.1392, 0.3347, 0.0866, -0.8654, -0.4564, -0.7839, 0.1033, -0.0204, 0.1558, -0.1469, 0.2850, -0.1139, 0.8253, 0.7352, -0.6132, 0.0566, 0.3087, -0.1189, 0.1640, 0.2511, 0.5230, -0.0972, -0.5621, -2.5404, 0.3529, -0.2543, -0.6757, 0.2045, -0.0511, -0.2204, 0.1023, 0.0143, 0.4191, -0.3946, -1.0912, 0.8555, 1.0751, -0.0184, -0.3162, 0.1910, 0.6522, -0.5801, 0.2091, -0.8254, -0.3425, 0.3368, -0.0384, -0.4570, 2.5288, -0.3513, -0.1630, 0.1096, -0.5936, 1.5303, -0.4135, -0.2418, -0.0564, -2.6344, -0.1054, 0.8866, -0.2946, -0.4564, -0.6220, 0.2672, -0.9012, 0.3535, 0.2344, -0.0718, 0.0782, 0.0133, 0.2032, -1.2768, 0.1271, -0.5114, -0.0584, -0.8219, -0.1069, 1.5577, -0.1432, -0.6794, 0.9101, 0.6390, 0.3547, -0.6126, -0.1885, 0.2462, -1.1864, 0.0653, -0.7940, 0.5204, 0.5372, 0.5353, -0.4268, -0.2003, -0.2496, -0.0405, 0.3615, -0.1635, 0.1908, -0.0467, 0.7167, 0.1465, 0.4621, 0.1190, -1.6899, 0.6512, 1.3150, -0.1273, 0.0507, 0.2058, -0.1855, 0.1316, 0.1280, 0.5049, 0.0262, -0.0329, 2.0327, -0.6410, 0.4536, 0.0609, 0.1883, -0.5454, -0.5247, 0.1856, 0.7238, 1.4886, -0.1068, 1.7239, -0.8228, -0.2155, 0.5159, 0.2941, -0.0782, -0.0159, 0.1844, -0.1808, -0.1132, 0.4861, 4.0106, 0.0130, 0.2455, -0.1101, 0.0792, 0.4720, -0.1022, 2.0154, -0.4013, 0.5604, 1.3600, -0.5614, 0.3793, -0.1245, 0.2444, 0.1657, 1.7616, 0.6198, 0.1761, -0.6036, -0.1931, 0.4449, 0.2574, -0.2360, 1.1118, 0.0804, 1.1533, 0.2549, 0.3386, 0.2463, 0.0930, -0.6093, -0.1464, 0.2889, 0.2294, -0.5943, 0.1323, 0.5119, 0.1093, -1.0178, 0.4735, 0.3068, 0.3213, -0.0585, -0.3682, -0.6105, -0.7776, 0.1999, 0.9439, -0.4209, 0.1488, 1.3119, -0.4679, -0.3882, 0.2677, -0.1673, -0.5921, -1.2811, -1.0972, 0.3873, 0.0798, -0.0538, 0.0659, -0.1439, -1.3106, -0.5175, 0.4538, -1.0376, -0.9015, 0.7454, -0.0714, -0.4641, 0.2083, 0.0596, -2.9637, 0.3057, 0.2121, -0.2399, 0.6963, 0.1400, 1.7446, 0.9707, -0.3118, -0.3371, 0.0130, 1.0006, -0.2740, 0.1100, -0.9666, 0.7636, 1.2002, -0.0018, -0.3380, 0.1262, 0.5829, -0.0374, 0.0689, 0.2022, -2.0056, -0.2051, -0.4549, 0.0519, 0.4217, -0.7413, 0.0601, 0.4385, 2.8503, -2.7656, 1.2281, -0.1280, 0.6028, 0.4995, 0.0638, -0.3376, 0.2527, -0.1572, -0.4385, -0.6372, 0.2569, 0.4115, 0.4507, 0.6063, -0.1051, 1.2529, 0.2453, -0.7905, -0.3797, -0.2674, 0.2662, 1.5347, -0.3908, 0.8839, -0.6054, -0.4827, -0.3495, 1.2107, -0.4419, -0.6177, 0.1054, 1.0132, -0.3246, -0.1776, 1.1740, -0.0252, 0.0368, -0.7937, -0.9988, -0.0228, 0.0742, -2.4925, 0.5785, 2.3900, 1.2726, -0.3682, -0.8625, -0.3299, 0.3934, 1.4045, -0.6200, -0.0024, 0.2348, -0.1827, -0.5913, -0.6982, 0.2648, 0.2601, 0.9986, 0.1636, 0.8982, -0.4269, 1.7454, -1.9136, -0.9865, -0.0451, 0.2851, -0.5938, -0.3066, 0.0910, -0.3150, -0.4002, 0.4789, 0.0337, -0.6997, -0.2555, -0.6602, -3.0103, 0.2491, -1.0346, 0.3651, 0.2319, 1.0224, -0.2613, 1.6970, 0.7515, 2.1477, 0.1310, 0.2060, 0.1372, 1.0049, -0.8758, -0.3804, -2.1513, 0.8010, -0.2271, -0.2108, 0.3728, -1.7321, -1.0250, -0.2584, -0.2513, 0.2418, -0.7641, 0.2084, -1.3560, 0.5803, 0.1556, -0.3612, 1.3099, -0.2673, 0.4371, -0.8022, 0.1776, -0.5019, 0.1880, -0.2093, 0.0750, -0.7228, -1.3950, 0.1944, -1.5994, -0.2832, 0.0507, 0.1917, 1.2954, 0.0471, 0.3115, -2.2382, -0.3891, -0.0704, 0.3897, 0.0347, 0.9186, -0.8407, 0.9456, 0.5629, 0.3474, -0.4869, 0.4696, -0.4438, 0.0860, -0.8313, -0.0383, 0.2055, 0.4822, -0.1455, -0.1719, -0.2346, -0.4606, 0.8018, 0.3767, -0.0613, 1.9429, -0.6558, -0.0772, -0.1592, -0.1413, 0.4759, -0.0686, 0.9243, -0.2413, -0.1084, -0.2248, -0.0776, 1.4193, -0.0605, 0.1305, -0.2055, 0.0917, 0.6884, -0.0152, 0.1215, 0.2920, -0.0781, -0.0256, 0.3789, -0.1933, 0.1759, 2.3899, 1.0915, -0.7082, -0.4519, -0.2648, -1.2404, -0.2485, 1.0713, 0.1662, -0.1268, 0.3338, -0.0319, 0.1692, -0.5161, 0.9351, 0.1996, -0.2743, 0.0492, -0.0171, 0.1546, 0.2533, -0.0102, 0.6147, 0.0035, -0.2468, -0.2116, -1.7912, 0.2735, 0.4147, 0.4458, 0.6123, 0.0860, 0.2098, -0.3691, -0.2297, -0.6086, -1.0407, -0.7736, -0.3087, -0.0900, -0.1007, -0.3801, -0.3408, -0.4853, -0.3101, -0.8812, 0.0187, -0.9697, -0.2393, 0.1129, -0.5682, 0.4349, 0.1017, 0.2173, -0.0644, -0.9307, 0.9754, 0.2189, 0.2966, -0.4089, -0.2471, -0.7549, 0.3300, 0.7856, 0.1262, 0.2097, -0.5872, 0.9896, 0.5100, 1.0608, -0.7974, 0.1549, -0.1020, 0.4286, 0.0603, -0.6836, -0.4662, -1.2350, -0.0858, -0.5552, 0.0383, 0.2145, -0.4324, -0.5896, 0.9709, -0.0827, -0.2574, 0.2436, -0.1460, 0.5862, 0.4329, -1.2421, 0.0497, -0.0034, 0.2385, -0.1346, 2.0652, 0.8790, -0.2033, -2.6427, 0.3654, -0.1929, -0.0753, -0.9107, 0.9437, 0.3717, -0.7058, -0.2487, -1.0937, -0.7612, 0.9516, -0.7426, -0.0736, 1.2167, 0.6336, 0.2707, -0.7666, -0.1272, -0.8960, 0.3748, 0.7344, 0.7257, 0.3686, -0.5036, -0.2829, 0.0548, 0.3034, -0.2335, -0.3215, 0.0566, -0.2733, -0.3644, 0.0467, -0.0924, -0.5145, -1.7089, 0.4896, 0.0074, 0.2840, 0.1140, -0.0409, -0.3251, 1.0805, 3.0856, -0.3409, 1.2684, -0.0245, -0.0636, -0.0090, 0.1293, -0.3410, -0.0482, 0.1482, 0.2027, 0.5623, 0.0566, 0.6453, -0.0126, 0.0720, -0.0277, 0.0531, 0.1860, -0.1044, -0.6973, 0.3026, 0.4733, -0.1590, 0.4727, 0.8486, 0.4478, 0.1814, 1.0862, 0.0478, 0.2437, -0.5269, -0.0796, -0.4291, 0.4937, -0.0407, -0.6961, -0.0412, 0.6865, 0.0457, 0.1085, -0.4717, -0.1339, 0.8600, 0.6718, -0.3542, -0.5655, 1.3711, 0.0034, 0.3077, 0.0903, 0.3618, 0.3287, -0.1007, 0.0332, -0.3841, -0.3981, 0.1079, -0.4399, 0.1836, 0.0939, -0.1425, -0.2531, -1.2103, 0.0234, -1.3023, -0.0570, -0.0587, 1.1733, 0.0079, 1.0809, 0.4697, -0.1427, 3.3793, -0.1503, 0.4354, 0.0274, 0.3112, -0.3816, 0.0187, -0.1282, -0.4136, 0.3684, 0.6930, 1.3605, 0.4949, 0.4162, -2.2398, 0.4104, 0.6839, 0.4519, 0.0546, -0.0816, 0.0357, 0.1977, -0.8450, 0.1481, 0.1588, -0.1392, -0.3304, -0.3499, -0.8669, 0.1510, 0.1127, 0.9853, -0.3019, -0.3493, -0.0783, -0.8491, 0.0696, 0.7295, -1.0612, 0.1232],\n \"std\": [0.9277, 0.7470, 0.6154, 0.8520, 0.8682, 0.7121, 0.7048, 0.6865, 0.7543, 0.6952, 0.6186, 0.4204, 0.4614, 0.4731, 0.4421, 0.4068, 0.6927, 0.6540, 0.4717, 0.4993, 0.5945, 0.5480, 0.4898, 0.6438, 0.5551, 0.5686, 0.7287, 0.6033, 0.5590, 0.3768, 0.5304, 0.6748, 0.5559, 0.5265, 0.6214, 0.6490, 0.4639, 0.6465, 0.5575, 0.6202, 0.5369, 1.2466, 0.7340, 0.5462, 0.6508, 0.5766, 0.5405, 0.5581, 0.5687, 0.7549, 0.5743, 0.4748, 0.6308, 0.6292, 0.6391, 0.6284, 0.4202, 0.5970, 0.5587, 0.5364, 0.4655, 0.5201, 0.7140, 0.6220, 0.4978, 0.4479, 0.5452, 0.7489, 0.5866, 0.4592, 0.7493, 0.6548, 0.5497, 0.4658, 0.8663, 0.4574, 0.5351, 0.5595, 0.4579, 0.5141, 0.4824, 0.5504, 0.5468, 0.5726, 0.5155, 0.6679, 0.8433, 0.5278, 0.5666, 0.7699, 0.5682, 0.9431, 0.5344, 0.6562, 0.4749, 0.5241, 0.6869, 0.4117, 0.5839, 0.5115, 0.8811, 0.5335, 0.6476, 0.4883, 0.6034, 0.5778, 0.4764, 0.8787, 0.8589, 0.5168, 0.4548, 0.8146, 0.5860, 0.6087, 0.6758, 0.7049, 0.8292, 0.6547, 0.6043, 0.7242, 0.6158, 0.6435, 0.5219, 0.6148, 0.7738, 0.4871, 0.7944, 0.7605, 0.6120, 0.5482, 0.6107, 0.6106, 0.4295, 0.4549, 0.4167, 0.6142, 0.6368, 0.5432, 0.5412, 0.6568, 0.9641, 0.6413, 0.6634, 0.4222, 0.6917, 0.5664, 0.5554, 0.4098, 0.6949, 0.5890, 0.4995, 0.5475, 0.6446, 0.5599, 0.6439, 0.6220, 0.5761, 0.5862, 0.5126, 0.6037, 0.5377, 0.5817, 0.6216, 0.5986, 0.4834, 0.6929, 0.5819, 0.6781, 0.6088, 0.5425, 0.7211, 0.6253, 0.5408, 0.6826, 0.5454, 0.7614, 0.9767, 0.8721, 0.7527, 0.4022, 0.5061, 0.5921, 0.5945, 0.6048, 0.7206, 0.5533, 0.5506, 0.6816, 0.6116, 0.6424, 0.7484, 0.6350, 0.5953, 0.4941, 0.7675, 0.8244, 0.6885, 0.5751, 0.9304, 0.5252, 0.5741, 0.4537, 0.5610, 0.9873, 0.5155, 0.7180, 0.4421, 0.5171, 0.5343, 0.5225, 0.7952, 0.6149, 0.6401, 0.5667, 0.6946, 0.8172, 0.5188, 0.5082, 0.6298, 0.6904, 0.4820, 0.5600, 0.5584, 0.5600, 0.4776, 0.5008, 0.7215, 0.6071, 0.5571, 0.6174, 0.4049, 0.7368, 0.5996, 0.7888, 0.7609, 0.5913, 0.8778, 0.4462, 0.7460, 0.7240, 0.5705, 0.6267, 0.5684, 0.5707, 0.6560, 0.5310, 0.5278, 0.6833, 0.6420, 0.6696, 0.8815, 0.4767, 0.7171, 0.4826, 0.6736, 0.5483, 0.4913, 0.5840, 0.5242, 0.4310, 0.5846, 0.4389, 0.5164, 0.6203, 0.5625, 0.8495, 0.5091, 0.6904, 0.5490, 0.5467, 0.4746, 0.8446, 0.6030, 0.6563, 1.0108, 0.5633, 0.6324, 0.6339, 0.6269, 1.2128, 0.6877, 0.5998, 0.4763, 0.4979, 0.7968, 0.6549, 1.0234, 0.5385, 0.6164, 0.5485, 0.8526, 0.5776, 0.5292, 0.5716, 0.5458, 0.5332, 0.5264, 0.6239, 0.6668, 0.7481, 0.3929, 0.5932, 0.5741, 0.4433, 0.7519, 0.4940, 0.7438, 0.5315, 0.3895, 0.5528, 0.6656, 0.6665, 0.9897, 0.8098, 0.6000, 0.5226, 1.2953, 0.5624, 0.6416, 0.5880, 0.5828, 0.4779, 0.6721, 0.6273, 0.7918, 0.5498, 0.5262, 0.6396, 0.6185, 0.6117, 0.8871, 0.5688, 0.5335, 0.6402, 0.5994, 0.9472, 0.5072, 0.7688, 0.6257, 0.6548, 0.6070, 0.7646, 0.5362, 0.5151, 0.6852, 0.4533, 0.6976, 0.6170, 0.5700, 0.5819, 0.4350, 0.5755, 0.4902, 0.9396, 0.5110, 0.5461, 0.6380, 1.0192, 0.5009, 0.8211, 0.6223, 0.5970, 0.5465, 0.8314, 0.4997, 0.5066, 0.5824, 0.6241, 0.4910, 0.4849, 0.5292, 0.5357, 0.4856, 0.6120, 0.4212, 0.6712, 0.4599, 0.4625, 0.7568, 0.8765, 0.8095, 0.7385, 0.5748, 0.7405, 0.6474, 0.6466, 0.6481, 0.5660, 0.6876, 0.9852, 0.5923, 0.6319, 0.6818, 0.4716, 0.6599, 0.5343, 0.5384, 0.9786, 0.4421, 0.5543, 1.0386, 0.5640, 0.5990, 0.5060, 0.6141, 0.3880, 0.6767, 0.5753, 0.4797, 0.4623, 0.5802, 0.6813, 0.5792, 0.4790, 0.6855, 0.5186, 0.4890, 0.5740, 0.6117, 0.5177, 0.5032, 0.6367, 0.4555, 0.6749, 0.6680, 0.6878, 0.7425, 0.8106, 0.5460, 1.0575, 0.5022, 0.7639, 0.5132, 0.5433, 0.7702, 0.4572, 0.4274, 0.6779, 0.5277, 0.5634, 0.4814, 0.5491, 0.5790, 0.5750, 0.5573, 0.4652, 0.5240, 0.6244, 0.6247, 0.7397, 0.7107, 0.5964, 0.4891, 0.7089, 0.6531, 0.6979, 0.4630, 0.5348, 0.4308, 0.8983, 0.5416, 0.4521, 0.6261, 0.4931, 0.7247, 0.5689, 0.5254, 0.4913, 0.6307, 0.5586, 0.5804, 0.5692, 0.5211, 0.6549, 0.6069, 0.5216, 0.4617, 0.7538, 0.4234, 0.4868, 0.7661, 1.1726, 0.8879, 0.4984, 0.6142, 0.4203, 0.5944, 0.6758, 0.5682, 0.6554, 0.7316, 0.5552, 0.7454, 0.3907, 0.7559, 0.4752, 0.5638, 0.7824, 0.7995, 0.5728, 0.8546, 0.5663, 0.5545, 0.4785, 1.0497, 0.7177, 0.5461, 0.5134, 0.5432, 0.5964, 0.5879, 0.7046, 0.7501, 0.5707, 0.9907, 0.9337, 0.5682, 0.4887, 0.5970, 0.6229, 0.6501, 0.7529, 0.7062, 0.6775, 0.7286, 0.6250, 0.4521, 0.5357, 0.5479, 0.7957, 0.4596, 0.6440, 0.8665, 0.6024, 0.7485, 0.6478, 0.6483, 0.5785, 0.5500, 0.4802, 0.4465, 0.6829, 0.6890, 0.6180, 0.8767, 0.7419, 0.6193, 0.3918, 0.5888, 0.5440, 0.5146, 0.4297, 0.4410, 0.4894, 0.4422, 0.9614, 0.6290, 0.6717, 0.5415, 0.5442, 0.5862, 0.4967, 0.7102, 1.1356, 0.4818, 0.4557, 0.6403, 0.4971, 0.7491, 0.8534, 0.8754, 0.5308, 0.5591, 0.6415, 0.7715, 0.8137, 0.4898, 0.5460, 0.5476, 0.9199, 0.6195, 0.5949, 0.7990, 0.4444, 0.6199, 0.5166, 0.4646, 0.9060, 0.6261, 0.5149, 0.6533, 0.7420, 0.4830, 0.5314, 0.5503, 0.5777, 0.6284, 0.7288, 0.5743, 0.6041, 0.5674, 0.4661, 0.6211, 0.6172, 0.4094, 0.5787, 0.8089, 0.6061, 0.5882, 0.5498, 0.7239, 0.6387, 0.7910, 0.5267, 0.5569, 0.6382, 0.5492, 0.5444, 0.6476, 0.8666, 0.9807, 0.5594, 0.6814, 0.5467, 0.8900, 0.5321, 0.5516, 1.0188, 0.7193, 0.5044, 0.5717, 0.9741, 0.7856, 0.6849, 0.5604, 1.0236, 0.8399, 0.5065, 0.6475, 0.4055, 0.7975, 0.4454, 0.5726, 0.4489, 0.6851, 0.6504, 0.4737, 0.5995, 0.6226, 0.5917, 0.5394, 0.5240, 0.7863, 0.6008, 0.5330, 0.4760, 0.6163, 0.4679, 0.5712, 0.7180, 0.4908, 1.0175, 0.5942, 0.5170, 0.7534, 0.5569, 0.8764, 0.7314, 0.5474, 0.9083, 0.6677, 0.6286, 0.6759, 0.5397, 0.5748, 0.6215, 0.4800, 0.5206, 0.5591, 0.5884, 0.6291, 0.6633, 0.7693, 0.5104, 0.6564, 0.5489, 0.6270, 0.5935, 0.6236, 0.6108, 0.4794, 0.5974, 0.7061, 0.6686, 0.6512, 0.4998, 0.5933, 0.4956, 0.6610, 0.7542, 0.5869, 0.8418, 0.9938, 0.9021, 0.6323, 0.5777, 0.4343, 0.6098, 0.5338, 0.5906, 0.7783, 0.7423, 0.6426, 0.6236, 0.9643, 0.5780, 1.0100, 1.1266, 0.7556, 0.5229, 0.8272, 0.6900, 0.5175, 0.4124, 0.5741, 0.4516, 0.6266, 0.5630, 0.5275, 0.5692, 0.5075, 0.7549, 0.6359, 0.5804, 0.6680, 0.7558, 0.6250, 0.4314, 0.6496, 0.5479, 0.7524, 0.7088, 0.6644, 0.7214, 0.6450, 0.4467, 0.7789, 0.5168, 0.6297, 0.6242, 0.4410, 0.8372, 0.5758, 0.4997, 0.8915, 0.6473, 0.5974, 0.5293, 0.7941, 0.4605, 0.9110, 0.5919, 0.5139, 0.5003, 0.4500, 0.6182, 0.5807, 0.4562, 0.5618, 0.6794, 0.7201, 0.6143, 0.8797, 0.8171, 0.6225, 0.7453, 0.7611, 0.4696, 1.0906, 0.8825, 0.7207, 0.5523, 0.7120, 0.5194, 0.5321, 1.0233, 0.5618, 0.5410, 0.4300, 0.7191, 0.5373, 0.4795, 0.4450, 0.6546, 0.7965, 0.7454, 0.6264, 0.5576, 0.7710, 0.5527, 0.6586, 0.5177, 0.4858, 0.5005, 0.5372, 0.5766, 0.4508, 0.5238, 0.8275, 0.4104, 0.5535, 0.8077, 0.4460, 0.7125, 0.7166, 0.6107, 0.4561, 0.6620, 0.4635, 0.6397, 0.4391, 0.6880, 0.6801, 0.5627, 0.8076, 0.7918, 1.0309, 0.5832, 0.6152, 0.7971, 0.4539, 0.5846, 0.7248, 0.4455, 0.6318, 0.6118, 0.4552, 0.6757, 0.5354, 0.6566, 0.6728, 0.4383, 0.6899, 1.0565, 0.6028, 0.6937, 0.5518, 0.8039, 0.4296, 0.6068, 0.5736, 0.4923, 0.7643, 0.7391, 0.4975, 0.5006, 0.5674, 0.5170, 0.4835, 0.4286, 0.5667, 0.6109, 0.6465, 0.6281, 0.7791, 0.5174, 0.5058, 0.6196, 0.6593, 0.5999, 0.5012, 0.5414, 0.7151, 0.6546, 0.6790, 0.5412, 0.4801, 0.6561, 1.0082, 0.5567, 0.6362, 0.4540, 0.8812, 0.6893, 0.6420, 0.6078, 0.5117, 0.7079, 0.8240, 0.7587, 0.6344, 0.6848, 0.4633, 0.5352, 0.6077, 0.5436, 0.7223, 0.5001, 0.9734, 0.5155, 0.5549, 0.4711, 0.9038, 0.5415, 1.0173, 0.5001, 0.5290, 0.5228, 0.5619, 0.9670, 0.7854, 0.5350, 0.5183, 0.9770, 0.5547, 0.9710, 0.5050, 0.4584, 0.6438, 0.4854, 0.5949, 0.6611, 0.4676, 0.4815, 0.8837, 0.6425, 0.6257, 0.6896, 0.4465, 0.7492, 0.6293, 0.7096, 0.5578, 0.5117, 0.4909, 0.5773, 0.4800, 0.5488, 0.6336, 0.6863, 0.5035, 0.6682, 0.7245, 0.5524, 0.4594, 0.5816, 0.5698, 0.6140, 0.5816, 0.5242, 0.4088, 0.4358, 0.6426, 0.4777, 0.6115, 0.4383, 0.5957, 0.8423, 0.5353, 0.5407, 0.8497, 0.6962, 0.7542, 0.5981, 0.5121, 0.6232, 0.5306, 0.5416, 0.5217, 0.5437, 0.5349, 0.5111, 0.8627, 0.6092, 0.5850, 0.5851, 0.7203, 0.3688, 0.5063, 0.5650, 0.5444, 0.5657, 0.7461, 0.4447, 0.7153, 0.4738, 0.5730, 0.4605, 0.4905, 0.6253, 0.8114, 0.8273, 0.5052, 0.6180, 0.6496, 0.4037, 0.5635, 0.5212, 0.7652, 0.4872, 0.5764, 0.7834, 0.6888, 0.5313, 0.5379, 0.5710, 0.7474, 0.6535, 0.9660, 0.5257, 0.7157, 0.7150, 0.5430, 0.5331, 0.6820, 0.6872, 0.4904, 0.6592, 0.6256, 0.6107, 0.4939, 0.5986, 0.5172, 0.4583],\n },\n \"emdb\": {\n \"count\": 62707,\n \"mean\": [-1.1869, 0.1485, 0.1933, -0.6247, 0.0793, 0.5762, 0.1835, -0.2564, 0.1285, 0.3221, 0.0577, 0.1154, -0.0818, -0.2512, 0.9673, -0.5680, 0.5968, -0.2124, -0.0112, -0.5576, 0.5339, -0.1490, 0.3102, -0.4012, -0.0570, 0.6416, 0.9359, -0.2932, 0.8544, 0.1719, -0.4534, 0.1316, 0.8625, 0.3806, 0.4884, 1.0853, -0.3872, -0.2403, -0.4274, 0.1319, -0.3334, 0.6352, 0.5748, -0.8850, -0.4331, 0.3662, -0.3324, 1.3993, -1.5142, -0.3082, -0.5491, -0.1847, 0.0145, -0.0726, 0.0015, -0.0358, -0.2815, -0.4356, -0.3842, 0.1150, 1.1513, 0.6343, -0.7336, -1.1613, 0.1020, -0.1291, 0.1560, 0.4854, -0.4191, 1.6794, 0.4274, 0.4792, 0.3570, 0.0811, 1.0886, 0.0670, 0.5227, 0.1891, 0.1121, 0.1495, -0.2090, -0.2156, -0.2512, -0.9291, 0.1287, -0.0481, 0.6701, -0.4579, 0.2352, -0.1056, 0.5551, 0.4357, 0.8168, 0.6344, -0.6445, -0.1965, 0.5587, 0.3860, -0.2466, -0.1542, 0.6825, 0.5875, -0.5208, 0.1500, -0.3980, 0.2157, 0.8368, -0.1356, -0.3387, 0.1747, 0.1467, 0.2282, -0.1412, 0.6216, -1.8406, 0.0150, 0.2891, 0.0280, 0.0461, 0.8558, 0.2929, -1.3753, -0.5792, 0.2089, -0.3524, -0.1849, -0.0157, 0.4454, -0.5306, 0.8238, -0.3160, 0.3760, 0.8978, -0.1943, -0.9474, -1.7321, -0.0149, 0.2338, 0.6087, -0.4851, 0.5210, -0.4042, -0.5368, -0.6220, 0.1245, 0.3112, 0.6360, -0.1522, 0.0540, -0.2380, -0.8354, 1.7591, 0.5687, 0.1732, 0.7923, -0.5383, -0.3271, -2.0050, -0.5563, 0.2979, 1.6609, 0.7108, -1.0155, 0.3591, 0.0136, -0.4743, -0.5401, -0.0176, 1.3333, -0.2973, -0.1114, -0.1616, 0.1160, 0.1152, 0.0057, 0.2067, 0.3876, -1.5311, 0.0636, 0.4566, -0.2653, 1.0534, -0.4638, 0.2166, 0.8686, -0.1447, 0.5605, -0.3841, 0.7015, 0.0418, 0.0811, -0.6406, -0.2929, -0.6821, 1.3678, 0.7574, 0.8315, 2.0377, 4.9034, -0.0097, 0.0165, 0.3248, 0.2994, 0.0210, 0.2276, -0.6580, -0.6899, 0.1981, -2.3205, 0.0059, -0.9412, -0.3191, 0.0389, -0.4170, 0.3391, -0.1346, 0.1567, 0.1838, -0.4176, -0.2758, 0.1495, -0.2977, 0.0929, 0.7186, 0.1230, 0.8780, -0.1240, -0.7370, -0.7551, 0.3830, 1.0824, 1.4500, -0.1040, 1.4225, 0.0929, 0.4612, 0.5167, -0.7093, -0.4729, 0.2321, 0.4156, -0.0696, -0.0626, 1.3341, -0.2398, 0.8453, 0.4048, 0.1690, 0.0074, -0.0474, 0.4134, 0.2043, -0.5962, 0.1643, -0.3821, 0.3012, -0.5690, 0.0133, 0.1876, -0.0727, 0.2896, 0.3253, 0.0313, 0.5141, -0.0055, -1.2889, -0.0983, -0.3212, -0.4173, -0.0804, 0.2591, -0.4160, -0.4815, 2.2822, -1.0033, -0.9814, 0.5290, 1.7943, -0.4217, -0.0373, -3.3970, 3.3067, 0.1174, -0.1369, 0.3847, -0.6960, -0.8867, -0.3825, -0.0134, -0.4367, -1.0273, -0.0623, 0.1520, 0.3816, -0.6543, -0.0118, -0.3019, -0.1190, 1.0490, 0.6255, 0.8503, 0.9500, -1.1942, 1.6886, -1.3958, 0.9389, 0.2318, -0.0460, 0.1140, -0.2352, -0.5648, 0.0363, -0.5636, 0.0661, -0.8680, -0.1223, -6.5336, 0.2139, -0.2734, 1.1739, 0.6003, 0.2183, 0.2154, -0.5902, -0.2916, -0.2748, 0.0787, 0.9065, -0.9764, -0.2278, 1.6248, 0.7941, -0.5014, 0.2422, -2.1474, 0.7818, 0.4370, 0.1361, -0.3936, -0.7724, 0.0941, -0.5762, 3.2182, -0.1101, 0.2677, -0.0101, -1.1798, -0.0122, -0.8163, 0.1115, -0.1697, -0.1466, -0.3549, 0.5360, -0.5183, 0.7519, 0.7093, -0.5946, 0.2787, 0.4822, -0.2680, 0.0934, 0.1483, 0.6706, -0.1150, -0.1945, -2.6643, 0.2194, -0.5014, -0.5869, 0.1022, 0.1988, -0.2558, 0.3732, -0.0644, 0.6440, -0.7403, -1.0228, 0.8158, 0.9543, -0.1226, -0.0929, 0.2716, 0.7962, -0.5293, 0.1538, -1.2074, -0.5093, 0.2037, 0.2156, -0.4407, 2.6976, -0.3653, 0.0458, -0.0899, -0.7584, 1.8329, -0.5082, -0.4776, -0.0265, -2.9437, -0.1675, 1.2358, 0.1571, -0.5022, -0.6370, 0.4087, -0.9664, 0.3533, 0.0928, -0.5308, 0.4462, 0.2476, 0.0976, -1.8347, 0.0468, -0.9309, -0.3712, -0.8578, -0.0568, 1.7377, -0.1299, -0.7187, 0.9764, 0.6858, 0.4272, -0.9588, 0.1038, 0.2520, -1.3775, 0.1491, -0.8507, 0.7052, 0.6483, 0.2818, -0.3305, -0.5913, -0.0907, -0.2438, -0.1932, -0.0564, -0.0777, -0.0748, 0.6530, 0.2393, 0.4476, 0.3941, -1.7061, 0.8876, 1.1888, 0.1423, 0.1737, 0.1330, 0.1115, 0.1525, -0.3715, 0.4657, -0.4010, -0.3089, 2.0455, -0.9555, 0.5093, 0.1502, -0.0865, -0.7851, -0.5175, 0.1613, 0.8113, 1.1943, 0.0612, 1.7087, -1.1616, -0.3204, 0.4428, 0.6120, -0.2282, 0.0174, -0.3141, -0.0045, 0.2204, 0.3966, 4.1174, -0.1531, 0.4325, -0.0245, -0.0310, 0.6541, 0.2904, 1.9309, -0.5405, 0.8576, 1.0352, -0.3592, -0.1056, -0.0047, 0.7218, 0.2350, 1.8817, 0.7558, -0.1575, -0.0544, 0.0234, 0.5841, 0.0996, -0.0503, 1.4150, 0.2260, 0.9152, 0.0688, 0.5286, 0.5885, 0.4606, -0.9186, 0.0441, 0.5233, 0.5305, -0.9086, 0.3728, 0.6752, 0.5453, -1.1360, 0.0613, -0.2365, 0.8856, -0.0512, -0.2589, -0.7055, -0.8111, 0.1787, 1.0393, -0.2469, -0.0922, 1.1790, -0.3284, 0.0402, 0.0746, -0.1033, -0.7248, -1.3859, -1.0511, 0.2797, 0.2777, -0.0877, 0.0271, 0.0740, -1.5863, -0.7014, 0.3677, -1.6786, -1.0769, 0.5594, 0.2428, -0.2664, 0.3454, -0.0490, -3.3762, 0.2004, 0.1913, -0.6461, 0.7643, -0.1239, 1.6487, 0.4942, -0.3305, -0.5069, -0.2183, 1.1533, -0.4380, 0.0219, -0.6319, 0.6743, 1.0648, 0.0587, -0.0989, -0.0995, 0.3757, 0.1813, 0.2854, 0.4345, -2.2154, 0.3601, -0.6406, -0.1099, 0.3583, -0.3726, 0.2892, 0.5897, 3.4282, -2.8781, 0.8985, 0.1550, 0.1102, 0.8008, -0.0811, -0.4199, 0.3145, -0.3236, -0.2425, -0.4502, 0.2431, 0.8504, 0.4597, 0.6396, 0.0902, 1.3885, 0.1297, -1.1721, -0.3227, -0.4472, 0.2575, 1.6201, -0.5444, 0.8665, -0.9622, 0.0035, -0.5908, 1.6270, 0.0351, -0.3419, 0.0039, 1.1001, -0.3767, -0.2270, 1.3332, 0.3555, 0.0667, -0.5392, -1.3500, -0.0842, 0.2591, -2.8862, 0.3166, 2.3757, 1.1254, -0.5208, -0.7074, -0.8110, 0.3715, 1.3720, -0.7236, -0.0665, 0.2772, -0.2840, -0.3515, -0.4777, 0.3030, 0.5417, 0.7752, -0.0182, 1.1569, -0.1614, 1.6521, -2.2844, -0.9332, -0.1472, 0.6151, -0.5020, -0.0719, 0.3361, -0.2722, -0.1500, 0.5092, -0.0348, -0.6530, -0.4159, -0.6603, -3.6738, 0.1421, -1.1267, 0.4267, 0.0699, 1.6415, 0.1451, 1.3309, 0.7792, 2.1801, -0.0886, 0.4233, 0.2828, 1.3708, -1.2021, -0.2627, -2.1505, 0.7701, -0.0167, -0.0247, 0.4665, -1.5951, -0.9997, -0.1568, -0.1108, 0.1543, -1.0055, 0.0001, -1.0355, 0.8421, -0.0485, -0.3064, 1.2358, -0.0448, 0.4038, -0.7671, 0.3624, -0.6197, 0.7966, -0.2266, 0.1130, -0.5302, -1.5468, 0.0700, -1.1711, -0.3307, 0.0086, -0.0416, 1.2763, -0.0574, 0.0121, -2.6334, -0.3180, -0.1954, 0.3944, 0.0076, 1.2025, -0.5634, 0.9271, 0.4198, 0.3251, -0.0041, 0.5236, -0.5314, 0.0639, -0.8840, -0.2680, 0.4958, 0.7804, 0.2942, -0.1935, -0.1405, -0.5670, 0.9489, 0.5726, -0.2529, 1.8878, -0.7204, -0.0050, -0.2448, 0.1725, 0.4253, 0.0058, 1.0247, -0.2908, -0.3978, -0.0963, 0.2107, 1.3576, 0.3074, 0.5527, -0.0927, 0.1521, 0.6300, -0.1377, -0.0497, 0.0425, -0.2248, -0.1534, 0.5778, 0.0033, 0.1789, 2.4935, 1.3225, -0.8038, -0.8864, 0.1176, -1.0532, -0.2375, 1.4582, -0.1168, 0.0548, 0.4221, -0.3585, 0.4043, -0.4371, 1.3289, -0.3674, -0.4286, -0.1730, 0.0535, 0.1441, 0.2703, 0.3826, 0.5123, -0.0401, -0.1230, -0.3143, -1.7583, 0.2582, 0.3484, 0.5722, 0.8621, 0.4420, 0.4442, -0.2445, 0.0532, -0.8102, -1.4058, -0.6382, -0.5799, -0.2456, -0.0906, -0.3191, -0.3395, -0.4364, -0.5810, -0.7970, 0.0831, -1.1570, -0.2573, -0.0644, -0.7106, 0.1313, 0.1944, -0.2329, 0.1409, -1.2096, 1.0822, 0.5523, 0.2151, -0.1106, -0.1034, -0.4873, 0.6932, 1.0196, -0.0521, 0.0569, -0.8759, 1.0084, 0.6800, 1.0768, -1.2878, -0.1161, 0.0447, 0.1888, -0.2371, -1.0470, -0.4027, -1.4363, 0.1606, -0.8026, -0.0244, -0.2893, -0.4938, -0.6921, 1.0140, -0.4158, -0.5957, 0.3313, -0.2462, 0.7703, 0.3403, -1.5113, -0.1231, -0.3776, 0.3326, 0.1634, 2.1520, 0.7302, -0.0300, -2.8234, 0.4553, -0.4652, -0.3331, -1.0286, 1.2882, -0.2797, -0.4759, 0.1470, -1.0253, -0.8175, 0.6936, -0.3728, -0.4594, 1.0876, 0.6229, -0.0461, -0.4342, -0.1686, -1.3960, 0.5283, 0.4002, 0.8179, 0.4787, -0.7147, -0.5052, -0.2552, 0.2817, -0.4022, -0.5289, 0.0815, -0.4814, -0.5451, -0.1384, -0.4303, -0.4506, -1.9036, 0.6884, 0.1361, 0.2678, -0.0052, 0.0119, -0.1882, 1.0507, 3.1094, -0.5746, 1.3087, -0.1831, -0.1917, 0.0633, 0.5083, -0.1448, -0.0134, 0.5002, 0.2579, 0.7755, 0.1579, 0.4157, -0.2610, -0.4953, 0.1709, 0.4063, 0.2068, 0.2666, -0.7872, 0.5325, 0.4910, -0.1599, 0.4387, 0.9262, 0.9245, 0.5763, 0.9292, -0.4531, -0.5367, -0.4911, 0.2302, -0.4182, 0.7188, 0.0342, -0.2079, 0.1310, 0.5718, -0.0331, 0.1861, -0.1287, -0.0427, 0.8478, 0.7278, -0.5664, -0.5335, 1.3976, 0.1697, 0.6063, -0.0220, 0.4921, -0.1349, -0.0531, -0.2408, -0.3858, -0.2741, 0.2285, -0.5532, 0.2704, -0.2687, -0.2161, -0.1179, -1.5228, -0.3683, -1.3004, 0.2431, -0.3305, 1.6118, -0.0328, 1.1503, 0.5712, -0.0423, 3.4830, -0.2760, 0.6307, -0.0419, 0.1553, -0.5602, 0.2106, -0.2213, -0.4543, 0.3034, 0.9189, 1.5738, 0.5071, 0.2238, -2.2069, 0.4104, 0.6224, 0.2836, -0.1620, -0.3043, -0.4012, 0.2410, -0.6261, -0.2435, 0.0211, -0.2227, -0.2392, -0.3634, -0.9207, 0.2260, 0.0929, 0.8206, -0.3214, -0.2296, 0.1274, -0.8615, 0.2329, 1.1085, -1.0565, 0.2258],\n \"std\": [0.9963, 0.6391, 0.4956, 0.6280, 0.7591, 0.5610, 0.8236, 0.7139, 0.7494, 0.5686, 0.5042, 0.3464, 0.4228, 0.4171, 0.3526, 0.3710, 0.6288, 0.4674, 0.4413, 0.4741, 0.6553, 0.4882, 0.3697, 0.5507, 0.4961, 0.3683, 0.5604, 0.5302, 0.6027, 0.3023, 0.4882, 0.5746, 0.5314, 0.5031, 0.6145, 0.5994, 0.4285, 0.6399, 0.5362, 0.5403, 0.4677, 1.2902, 0.6126, 0.4145, 0.5068, 0.4667, 0.4825, 0.4275, 0.4381, 0.6758, 0.4866, 0.4136, 0.5262, 0.5698, 0.6550, 0.6492, 0.3450, 0.5948, 0.4219, 0.4973, 0.4483, 0.4336, 0.7440, 0.4595, 0.4366, 0.3634, 0.4430, 0.6587, 0.5073, 0.3533, 0.7036, 0.7039, 0.5312, 0.4701, 0.7512, 0.4102, 0.4227, 0.4488, 0.4158, 0.4676, 0.4521, 0.4560, 0.3917, 0.4757, 0.4348, 0.6013, 0.6715, 0.5179, 0.4834, 0.7451, 0.4845, 0.8893, 0.4188, 0.5963, 0.4306, 0.4551, 0.6417, 0.2886, 0.5378, 0.4316, 0.7568, 0.4818, 0.5494, 0.4736, 0.5841, 0.5043, 0.4265, 0.6994, 0.7652, 0.4344, 0.3931, 0.7198, 0.4169, 0.5794, 0.6720, 0.5694, 0.8603, 0.5307, 0.5893, 0.5763, 0.5292, 0.5228, 0.4156, 0.4901, 0.8334, 0.4574, 0.7241, 0.5346, 0.4063, 0.4147, 0.4979, 0.6599, 0.4173, 0.3715, 0.3828, 0.4492, 0.5576, 0.4060, 0.4353, 0.5315, 0.9834, 0.5548, 0.5679, 0.3506, 0.5419, 0.4256, 0.4187, 0.3570, 0.6316, 0.5870, 0.4832, 0.4862, 0.6072, 0.6781, 0.6152, 0.6708, 0.5008, 0.4435, 0.4229, 0.4973, 0.4301, 0.5363, 0.5478, 0.5388, 0.3952, 0.5961, 0.4721, 0.6389, 0.4450, 0.4841, 0.5594, 0.5234, 0.5224, 0.6326, 0.4469, 0.7397, 0.9551, 0.8426, 0.7576, 0.3893, 0.4382, 0.5222, 0.5234, 0.6035, 0.5764, 0.4043, 0.4741, 0.5471, 0.4229, 0.5962, 0.7127, 0.6205, 0.5671, 0.3766, 0.7455, 0.7315, 0.5891, 0.5372, 0.5957, 0.5342, 0.4010, 0.4453, 0.4609, 0.8789, 0.4353, 0.6297, 0.4126, 0.4149, 0.4597, 0.4859, 0.6733, 0.6096, 0.5719, 0.4494, 0.6353, 0.7537, 0.4643, 0.4577, 0.6485, 0.6069, 0.3603, 0.5821, 0.4807, 0.5192, 0.5329, 0.4153, 0.7329, 0.5444, 0.5742, 0.4593, 0.4003, 0.6770, 0.5428, 0.6781, 0.7920, 0.5037, 0.7615, 0.4537, 0.5931, 0.7333, 0.4880, 0.5469, 0.4698, 0.4917, 0.6256, 0.4947, 0.3974, 0.7559, 0.5916, 0.6547, 0.7502, 0.4682, 0.4517, 0.4888, 0.6472, 0.4755, 0.3927, 0.5845, 0.4135, 0.4091, 0.5860, 0.4544, 0.4051, 0.5547, 0.5322, 0.7200, 0.4595, 0.5484, 0.4758, 0.5259, 0.4137, 0.7149, 0.5638, 0.6221, 0.9309, 0.5637, 0.5657, 0.5711, 0.5651, 1.0484, 0.4435, 0.4587, 0.3716, 0.4108, 0.8114, 0.5531, 1.0675, 0.5825, 0.3841, 0.4500, 0.7335, 0.4767, 0.4162, 0.5679, 0.4880, 0.4614, 0.5118, 0.5198, 0.5619, 0.6869, 0.3536, 0.5128, 0.4722, 0.3722, 0.7705, 0.4556, 0.5365, 0.4999, 0.3254, 0.5268, 0.7580, 0.5932, 0.9908, 0.6171, 0.4912, 0.4439, 0.9135, 0.4658, 0.6566, 0.5500, 0.5423, 0.4725, 0.5415, 0.5550, 0.7519, 0.4220, 0.6024, 0.4821, 0.5268, 0.4583, 0.7421, 0.5200, 0.4541, 0.5197, 0.4562, 0.8381, 0.4423, 0.7400, 0.6578, 0.6459, 0.5316, 0.6877, 0.5362, 0.4215, 0.6455, 0.4363, 0.6716, 0.5795, 0.5587, 0.5234, 0.4456, 0.4991, 0.4244, 0.8959, 0.4744, 0.4440, 0.4437, 0.8485, 0.4237, 0.6907, 0.5582, 0.4315, 0.5458, 0.7341, 0.4731, 0.5065, 0.6181, 0.5643, 0.4407, 0.4353, 0.4732, 0.3769, 0.4162, 0.5028, 0.3689, 0.6656, 0.4598, 0.3735, 0.6801, 0.7902, 0.7101, 0.6292, 0.5732, 0.7452, 0.6803, 0.5065, 0.5261, 0.4644, 0.5021, 0.6714, 0.5226, 0.4455, 0.7599, 0.4380, 0.5468, 0.4595, 0.5308, 0.8445, 0.4413, 0.5196, 0.9241, 0.5414, 0.5018, 0.3832, 0.4950, 0.3185, 0.5330, 0.4844, 0.4481, 0.4517, 0.5104, 0.6092, 0.5712, 0.4164, 0.6590, 0.4888, 0.3930, 0.5419, 0.5486, 0.5165, 0.4390, 0.5542, 0.3883, 0.4074, 0.6213, 0.6185, 0.7711, 0.6565, 0.4925, 1.0624, 0.4690, 0.7498, 0.5333, 0.5290, 0.6258, 0.4473, 0.3862, 0.6571, 0.4873, 0.5240, 0.4127, 0.4445, 0.5094, 0.4754, 0.5769, 0.4786, 0.4510, 0.5130, 0.4897, 0.7568, 0.7398, 0.5718, 0.4229, 0.4929, 0.7470, 0.5901, 0.3772, 0.4914, 0.4074, 0.9471, 0.4967, 0.4323, 0.5259, 0.3591, 0.7202, 0.6012, 0.4573, 0.4296, 0.5578, 0.5218, 0.4640, 0.4522, 0.4029, 0.8071, 0.6086, 0.4832, 0.4202, 0.6781, 0.3862, 0.3920, 0.7543, 1.0257, 0.8849, 0.4181, 0.4722, 0.4069, 0.4854, 0.5405, 0.4676, 0.5547, 0.6282, 0.4275, 0.8011, 0.3308, 0.7135, 0.4315, 0.4915, 0.6616, 0.7376, 0.5742, 0.7461, 0.5443, 0.4749, 0.4906, 1.0020, 0.6306, 0.4435, 0.4559, 0.4360, 0.4047, 0.5802, 0.6109, 0.7836, 0.5163, 0.9777, 0.9272, 0.4618, 0.3534, 0.5218, 0.4479, 0.6498, 0.7145, 0.6224, 0.5671, 0.5042, 0.3885, 0.4079, 0.4481, 0.5406, 0.6944, 0.3744, 0.5942, 0.6770, 0.5934, 0.7417, 0.5662, 0.4753, 0.5063, 0.5003, 0.4510, 0.4358, 0.6455, 0.7740, 0.4780, 0.8687, 0.5533, 0.5700, 0.3518, 0.4868, 0.4154, 0.4798, 0.3266, 0.3536, 0.3789, 0.3805, 0.7909, 0.5760, 0.5784, 0.4993, 0.5787, 0.5324, 0.4496, 0.8483, 1.0794, 0.4820, 0.4135, 0.6231, 0.4668, 0.6684, 0.7052, 0.7616, 0.4881, 0.4150, 0.5793, 0.8068, 0.7793, 0.4721, 0.5230, 0.4810, 0.9577, 0.5537, 0.5583, 0.6645, 0.4334, 0.6398, 0.5011, 0.4081, 0.6255, 0.5372, 0.4846, 0.6125, 0.6509, 0.4413, 0.4762, 0.4917, 0.5940, 0.4950, 0.6753, 0.6653, 0.5210, 0.5599, 0.4678, 0.4868, 0.5985, 0.4160, 0.4874, 0.8380, 0.5382, 0.5701, 0.5448, 0.6131, 0.5674, 0.7120, 0.4070, 0.4434, 0.5725, 0.4919, 0.4805, 0.5997, 0.7108, 0.9824, 0.4765, 0.7575, 0.4452, 0.8892, 0.4639, 0.4962, 1.0346, 0.7584, 0.4312, 0.4835, 0.8968, 0.4799, 0.6864, 0.5641, 1.0694, 0.6750, 0.4288, 0.5159, 0.3649, 0.7699, 0.4386, 0.4449, 0.3923, 0.6499, 0.5612, 0.4541, 0.6261, 0.5444, 0.4369, 0.4124, 0.4174, 0.6129, 0.5005, 0.4779, 0.3929, 0.4865, 0.4338, 0.4114, 0.6266, 0.3669, 1.0147, 0.4856, 0.4867, 0.6250, 0.5368, 0.6699, 0.6411, 0.5296, 0.7614, 0.5643, 0.5843, 0.6846, 0.3923, 0.3928, 0.4964, 0.4490, 0.4755, 0.4104, 0.5468, 0.6040, 0.5808, 0.6283, 0.4316, 0.6127, 0.4635, 0.5303, 0.4261, 0.4668, 0.6121, 0.4063, 0.5571, 0.6130, 0.5874, 0.4987, 0.4113, 0.5401, 0.4028, 0.6598, 0.7740, 0.5384, 0.7890, 0.9379, 0.8801, 0.6222, 0.5356, 0.3990, 0.4802, 0.4107, 0.5475, 0.6936, 0.6865, 0.4776, 0.5211, 0.8844, 0.6517, 1.0729, 0.9252, 0.6953, 0.4177, 0.7587, 0.6628, 0.3629, 0.3685, 0.3758, 0.4439, 0.5236, 0.4905, 0.5290, 0.4184, 0.3940, 0.6498, 0.5411, 0.5662, 0.5519, 0.6107, 0.6385, 0.4127, 0.6277, 0.5255, 0.5926, 0.5653, 0.6570, 0.6034, 0.5312, 0.4128, 0.7292, 0.3620, 0.5067, 0.5314, 0.3908, 0.7561, 0.4494, 0.4501, 0.7682, 0.4939, 0.4198, 0.5256, 0.6339, 0.5123, 0.9018, 0.5054, 0.4879, 0.4567, 0.4145, 0.6046, 0.3835, 0.4289, 0.5254, 0.6191, 0.6610, 0.5933, 0.7890, 0.7817, 0.6299, 0.5977, 0.7094, 0.3737, 1.0318, 0.7045, 0.7785, 0.5376, 0.5861, 0.4233, 0.5538, 1.0604, 0.5690, 0.5249, 0.3747, 0.6036, 0.4707, 0.3617, 0.3665, 0.6184, 0.4878, 0.6193, 0.5311, 0.6187, 0.6748, 0.4493, 0.6137, 0.4601, 0.3855, 0.4183, 0.4986, 0.4832, 0.4192, 0.4416, 0.7202, 0.3724, 0.4899, 0.6939, 0.4272, 0.7122, 0.6950, 0.5565, 0.4417, 0.6186, 0.4753, 0.5919, 0.3763, 0.5643, 0.5347, 0.5454, 0.9336, 0.6594, 0.9747, 0.4970, 0.4725, 0.7820, 0.4113, 0.4942, 0.6699, 0.4159, 0.6766, 0.6564, 0.3947, 0.5381, 0.3874, 0.6686, 0.5628, 0.3904, 0.6647, 0.9821, 0.4343, 0.5455, 0.4879, 0.8165, 0.4153, 0.5544, 0.5179, 0.3821, 0.6678, 0.7883, 0.3372, 0.4702, 0.5044, 0.4584, 0.4769, 0.3787, 0.4377, 0.5435, 0.5899, 0.5378, 0.5986, 0.4887, 0.5390, 0.5464, 0.6330, 0.5010, 0.4244, 0.5249, 0.6770, 0.6314, 0.6404, 0.4605, 0.3649, 0.6489, 1.0657, 0.5497, 0.5357, 0.3651, 0.8484, 0.8126, 0.4873, 0.6711, 0.4401, 0.6181, 0.8585, 0.6000, 0.5654, 0.5416, 0.3504, 0.4671, 0.5499, 0.4409, 0.7650, 0.4980, 0.9734, 0.3568, 0.6037, 0.4361, 0.7880, 0.4726, 0.9902, 0.5020, 0.5178, 0.5065, 0.4543, 0.9039, 0.8296, 0.4451, 0.4436, 0.8518, 0.5201, 0.8668, 0.5122, 0.3412, 0.5849, 0.4815, 0.5795, 0.5664, 0.4384, 0.4593, 0.7974, 0.6570, 0.6522, 0.5490, 0.4195, 0.6821, 0.6133, 0.5692, 0.4780, 0.4574, 0.5090, 0.4488, 0.4269, 0.4153, 0.5143, 0.6560, 0.4480, 0.5482, 0.6997, 0.4377, 0.4166, 0.6103, 0.4671, 0.4449, 0.5672, 0.3296, 0.3898, 0.3778, 0.6572, 0.5555, 0.4047, 0.3720, 0.5728, 0.6867, 0.5435, 0.5001, 0.6808, 0.6373, 0.6849, 0.4826, 0.4767, 0.3736, 0.5070, 0.4442, 0.4302, 0.4339, 0.4614, 0.4735, 0.7977, 0.5657, 0.4047, 0.5261, 0.6204, 0.3413, 0.3996, 0.4236, 0.3303, 0.4193, 0.6074, 0.3941, 0.4802, 0.4114, 0.3880, 0.3460, 0.3767, 0.6491, 0.6893, 0.8560, 0.4244, 0.4307, 0.5702, 0.3635, 0.5170, 0.3975, 0.6187, 0.5012, 0.4976, 0.7149, 0.7001, 0.4834, 0.3844, 0.5179, 0.6909, 0.5862, 1.0062, 0.5099, 0.6410, 0.7432, 0.4219, 0.4655, 0.6067, 0.6674, 0.4618, 0.7115, 0.5300, 0.5284, 0.4208, 0.4955, 0.4561, 0.3723],\n }\n}\n\ncam_angvel = {\n \"emdb_none_test\": {\n \"count\": 42622,\n \"mean\": [1., 0., 0., 0., 1., 0.],\n \"std\": [5.5702e-05, 3.2200e-03, 5.6530e-03, 3.2191e-03, 2.4738e-05, 3.3406e-03],\n },\n \"manual\": {\n \"mean\": [1., 0., 0., 0., 1., 0.],\n \"std\": [0.001, 0.1, 0.1, 0.1, 0.001, 0.1], # manually\n }\n}\n# fmt:on\n\n# ====== Compose ====== #\n\n\ndef compose(targets, sources):\n if len(sources) == 1:\n sources = sources * len(targets)\n mean = []\n std = []\n for t, s in zip(targets, sources):\n mean.extend(t[s][\"mean\"])\n std.extend(t[s][\"std\"])\n return {\"mean\": mean, \"std\": std}\n\n\nDEFAULT_01 = {\"mean\": [0.0], \"std\": [1.0]}\n\nMM_V1 = compose(\n [body_pose_r6d, betas, global_orient_c_r6d, global_orient_gv_r6d, local_transl_vel],\n [\"bedlam\"] * 5,\n)\nMM_V1_AMASS_LOCAL_BEDLAM_CAM = compose(\n [body_pose_r6d, betas, global_orient_c_r6d, global_orient_gv_r6d, local_transl_vel],\n [\"amass\", \"amass\", \"bedlam\", \"bedlam\", \"amass\"],\n)\n\nMM_V2 = compose(\n [body_pose_r6d, betas, global_orient_c_r6d, global_orient_gv_r6d, local_transl_vel],\n [\"bedlam\", \"bedlam\", \"bedlam\", \"bedlam\", \"none\"],\n)\n\nMM_V2_1 = compose(\n [body_pose_r6d, betas, global_orient_c_r6d, global_orient_gv_r6d, local_transl_vel],\n [\"bedlam\", \"bedlam\", \"bedlam\", \"bedlam\", \"1e-2\"],\n)\n" }, { "path": "eval/GVHMR/hmr4d/network/base_arch/embeddings/rotary_embedding.py", "content": "import torch\nimport torch.nn as nn\nfrom einops import repeat, rearrange\nfrom torch.cuda.amp import autocast\n\n\ndef rotate_half(x):\n x = rearrange(x, \"... (d r) -> ... d r\", r=2)\n x1, x2 = x.unbind(dim=-1)\n x = torch.stack((-x2, x1), dim=-1)\n return rearrange(x, \"... d r -> ... (d r)\")\n\n\n@autocast(enabled=False)\ndef apply_rotary_emb(freqs, t, start_index=0, scale=1.0, seq_dim=-2):\n if t.ndim == 3:\n seq_len = t.shape[seq_dim]\n freqs = freqs[-seq_len:].to(t)\n\n rot_dim = freqs.shape[-1]\n end_index = start_index + rot_dim\n\n assert (\n rot_dim <= t.shape[-1]\n ), f\"feature dimension {t.shape[-1]} is not of sufficient size to rotate in all the positions {rot_dim}\"\n\n t_left, t, t_right = t[..., :start_index], t[..., start_index:end_index], t[..., end_index:]\n t = (t * freqs.cos() * scale) + (rotate_half(t) * freqs.sin() * scale)\n return torch.cat((t_left, t, t_right), dim=-1)\n\n\ndef get_encoding(d_model, max_seq_len=4096):\n \"\"\"Return: (L, D)\"\"\"\n t = torch.arange(max_seq_len).float()\n freqs = 1.0 / (10000 ** (torch.arange(0, d_model, 2).float() / d_model))\n freqs = torch.einsum(\"i, j -> i j\", t, freqs)\n freqs = repeat(freqs, \"i j -> i (j r)\", r=2)\n return freqs\n\n\nclass ROPE(nn.Module):\n \"\"\"Minimal impl of a lang-style positional encoding.\"\"\"\n\n def __init__(self, d_model, max_seq_len=4096):\n super().__init__()\n self.d_model = d_model\n self.max_seq_len = max_seq_len\n\n # Pre-cache a freqs tensor\n encoding = get_encoding(d_model, max_seq_len)\n self.register_buffer(\"encoding\", encoding, False)\n\n def rotate_queries_or_keys(self, x):\n \"\"\"\n Args:\n x : (B, H, L, D)\n Returns:\n rotated_x: (B, H, L, D)\n \"\"\"\n\n seq_len, d_model = x.shape[-2:]\n assert d_model == self.d_model\n\n # encoding: (L, D)s\n if seq_len > self.max_seq_len:\n encoding = get_encoding(d_model, seq_len).to(x)\n else:\n encoding = self.encoding[:seq_len]\n\n # encoding: (L, D)\n # x: (B, H, L, D)\n rotated_x = apply_rotary_emb(encoding, x, seq_dim=-2)\n\n return rotated_x\n" }, { "path": "eval/GVHMR/hmr4d/network/base_arch/transformer/encoder_rope.py", "content": "import torch\nimport torch.nn as nn\nimport torch.nn.functional as F\nimport math\nfrom timm.models.vision_transformer import Mlp\nfrom typing import Optional, Tuple\nfrom einops import einsum, rearrange, repeat\nfrom hmr4d.network.base_arch.embeddings.rotary_embedding import ROPE\n\n\nclass RoPEAttention(nn.Module):\n def __init__(self, embed_dim, num_heads, dropout=0.1):\n super().__init__()\n self.embed_dim = embed_dim\n self.num_heads = num_heads\n self.head_dim = embed_dim // num_heads\n\n self.rope = ROPE(self.head_dim, max_seq_len=4096)\n\n self.query = nn.Linear(embed_dim, embed_dim)\n self.key = nn.Linear(embed_dim, embed_dim)\n self.value = nn.Linear(embed_dim, embed_dim)\n self.dropout = nn.Dropout(dropout)\n self.proj = nn.Linear(embed_dim, embed_dim)\n\n def forward(self, x, attn_mask=None, key_padding_mask=None):\n # x: (B, L, C)\n # attn_mask: (L, L)\n # key_padding_mask: (B, L)\n B, L, _ = x.shape\n xq, xk, xv = self.query(x), self.key(x), self.value(x)\n\n xq = xq.reshape(B, L, self.num_heads, -1).transpose(1, 2)\n xk = xk.reshape(B, L, self.num_heads, -1).transpose(1, 2)\n xv = xv.reshape(B, L, self.num_heads, -1).transpose(1, 2)\n\n xq = self.rope.rotate_queries_or_keys(xq) # B, N, L, C\n xk = self.rope.rotate_queries_or_keys(xk) # B, N, L, C\n\n attn_score = einsum(xq, xk, \"b n i c, b n j c -> b n i j\") / math.sqrt(self.head_dim)\n if attn_mask is not None:\n attn_mask = attn_mask.reshape(1, 1, L, L).expand(B, self.num_heads, -1, -1)\n attn_score = attn_score.masked_fill(attn_mask, float(\"-inf\"))\n if key_padding_mask is not None:\n key_padding_mask = key_padding_mask.reshape(B, 1, 1, L).expand(-1, self.num_heads, L, -1)\n attn_score = attn_score.masked_fill(key_padding_mask, float(\"-inf\"))\n\n attn_score = torch.softmax(attn_score, dim=-1)\n attn_score = self.dropout(attn_score)\n output = einsum(attn_score, xv, \"b n i j, b n j c -> b n i c\") # B, N, L, C\n output = output.transpose(1, 2).reshape(B, L, -1) # B, L, C\n output = self.proj(output) # B, L, C\n return output\n\n\nclass EncoderRoPEBlock(nn.Module):\n def __init__(self, hidden_size, num_heads, mlp_ratio=4.0, dropout=0.1, **block_kwargs):\n super().__init__()\n self.norm1 = nn.LayerNorm(hidden_size, elementwise_affine=True, eps=1e-6)\n self.attn = RoPEAttention(hidden_size, num_heads, dropout)\n self.norm2 = nn.LayerNorm(hidden_size, elementwise_affine=True, eps=1e-6)\n mlp_hidden_dim = int(hidden_size * mlp_ratio)\n approx_gelu = lambda: nn.GELU(approximate=\"tanh\")\n self.mlp = Mlp(in_features=hidden_size, hidden_features=mlp_hidden_dim, act_layer=approx_gelu, drop=dropout)\n\n self.gate_msa = nn.Parameter(torch.zeros(1, 1, hidden_size))\n self.gate_mlp = nn.Parameter(torch.zeros(1, 1, hidden_size))\n\n # Zero-out adaLN modulation layers\n nn.init.constant_(self.gate_msa, 0)\n nn.init.constant_(self.gate_mlp, 0)\n\n def forward(self, x, attn_mask=None, tgt_key_padding_mask=None):\n x = x + self.gate_msa * self._sa_block(\n self.norm1(x), attn_mask=attn_mask, key_padding_mask=tgt_key_padding_mask\n )\n x = x + self.gate_mlp * self.mlp(self.norm2(x))\n return x\n\n def _sa_block(self, x, attn_mask=None, key_padding_mask=None):\n # x: (B, L, C)\n x = self.attn(x, attn_mask=attn_mask, key_padding_mask=key_padding_mask)\n return x\n" }, { "path": "eval/GVHMR/hmr4d/network/base_arch/transformer/layer.py", "content": "import torch\nimport torch.nn as nn\nimport torch.nn.functional as F\n\n\ndef zero_module(module):\n \"\"\"\n Zero out the parameters of a module and return it.\n \"\"\"\n for p in module.parameters():\n p.detach().zero_()\n return module\n" }, { "path": "eval/GVHMR/hmr4d/network/gvhmr/relative_transformer.py", "content": "import torch\nimport torch.nn as nn\nimport torch.nn.functional as F\nfrom einops import einsum, rearrange, repeat\nfrom hmr4d.configs import MainStore, builds\n\nfrom hmr4d.network.base_arch.transformer.encoder_rope import EncoderRoPEBlock\nfrom hmr4d.network.base_arch.transformer.layer import zero_module\n\nfrom hmr4d.utils.net_utils import length_to_mask\nfrom timm.models.vision_transformer import Mlp\n\n\nclass NetworkEncoderRoPE(nn.Module):\n def __init__(\n self,\n # x\n output_dim=151,\n max_len=120,\n # condition\n cliffcam_dim=3,\n cam_angvel_dim=6,\n imgseq_dim=1024,\n # intermediate\n latent_dim=512,\n num_layers=12,\n num_heads=8,\n mlp_ratio=4.0,\n # output\n pred_cam_dim=3,\n static_conf_dim=6,\n # training\n dropout=0.1,\n # other\n avgbeta=True,\n ):\n super().__init__()\n\n # input\n self.output_dim = output_dim\n self.max_len = max_len\n\n # condition\n self.cliffcam_dim = cliffcam_dim\n self.cam_angvel_dim = cam_angvel_dim\n self.imgseq_dim = imgseq_dim\n\n # intermediate\n self.latent_dim = latent_dim\n self.num_layers = num_layers\n self.num_heads = num_heads\n self.dropout = dropout\n\n # ===== build model ===== #\n # Input (Kp2d)\n # Main token: map d_obs 2 to 32\n self.learned_pos_linear = nn.Linear(2, 32)\n self.learned_pos_params = nn.Parameter(torch.randn(17, 32), requires_grad=True)\n self.embed_noisyobs = Mlp(\n 17 * 32, hidden_features=self.latent_dim * 2, out_features=self.latent_dim, drop=dropout\n )\n\n self._build_condition_embedder()\n\n # Transformer\n self.blocks = nn.ModuleList(\n [\n EncoderRoPEBlock(self.latent_dim, self.num_heads, mlp_ratio=mlp_ratio, dropout=dropout)\n for _ in range(self.num_layers)\n ]\n )\n\n # Output heads\n self.final_layer = Mlp(self.latent_dim, out_features=self.output_dim)\n self.pred_cam_head = pred_cam_dim > 0 # keep extra_output for easy-loading old ckpt\n if self.pred_cam_head:\n self.pred_cam_head = Mlp(self.latent_dim, out_features=pred_cam_dim)\n self.register_buffer(\"pred_cam_mean\", torch.tensor([1.0606, -0.0027, 0.2702]), False)\n self.register_buffer(\"pred_cam_std\", torch.tensor([0.1784, 0.0956, 0.0764]), False)\n\n self.static_conf_head = static_conf_dim > 0\n if self.static_conf_head:\n self.static_conf_head = Mlp(self.latent_dim, out_features=static_conf_dim)\n\n self.avgbeta = avgbeta\n\n def _build_condition_embedder(self):\n latent_dim = self.latent_dim\n dropout = self.dropout\n self.cliffcam_embedder = nn.Sequential(\n nn.Linear(self.cliffcam_dim, latent_dim),\n nn.SiLU(),\n nn.Dropout(dropout),\n zero_module(nn.Linear(latent_dim, latent_dim)),\n )\n if self.cam_angvel_dim > 0:\n self.cam_angvel_embedder = nn.Sequential(\n nn.Linear(self.cam_angvel_dim, latent_dim),\n nn.SiLU(),\n nn.Dropout(dropout),\n zero_module(nn.Linear(latent_dim, latent_dim)),\n )\n if self.imgseq_dim > 0:\n self.imgseq_embedder = nn.Sequential(\n nn.LayerNorm(self.imgseq_dim),\n zero_module(nn.Linear(self.imgseq_dim, latent_dim)),\n )\n\n def forward(self, length, obs=None, f_cliffcam=None, f_cam_angvel=None, f_imgseq=None):\n \"\"\"\n Args:\n x: None we do not use it\n timesteps: (B,)\n length: (B), valid length of x, if None then use x.shape[2]\n f_imgseq: (B, L, C)\n f_cliffcam: (B, L, 3), CLIFF-Cam parameters (bbx-detection in the full-image)\n f_noisyobs: (B, L, C), nosiy pose observation\n f_cam_angvel: (B, L, 6), Camera angular velocity\n \"\"\"\n B, L, J, C = obs.shape\n assert J == 17 and C == 3\n\n # Main token from observation (2D pose)\n obs = obs.clone()\n visible_mask = obs[..., [2]] > 0.5 # (B, L, J, 1)\n obs[~visible_mask[..., 0]] = 0 # set low-conf to all zeros\n f_obs = self.learned_pos_linear(obs[..., :2]) # (B, L, J, 32)\n f_obs = f_obs * visible_mask + self.learned_pos_params.repeat(B, L, 1, 1) * ~visible_mask\n x = self.embed_noisyobs(f_obs.view(B, L, -1)) # (B, L, J*32) -> (B, L, C)\n\n # Condition\n f_to_add = []\n f_to_add.append(self.cliffcam_embedder(f_cliffcam))\n if hasattr(self, \"cam_angvel_embedder\"):\n f_to_add.append(self.cam_angvel_embedder(f_cam_angvel))\n if f_imgseq is not None and hasattr(self, \"imgseq_embedder\"):\n f_to_add.append(self.imgseq_embedder(f_imgseq))\n\n for f_delta in f_to_add:\n x = x + f_delta\n\n # Setup length and make padding mask\n assert B == length.size(0)\n pmask = ~length_to_mask(length, L) # (B, L)\n\n if L > self.max_len:\n attnmask = torch.ones((L, L), device=x.device, dtype=torch.bool)\n for i in range(L):\n min_ind = max(0, i - self.max_len // 2)\n max_ind = min(L, i + self.max_len // 2)\n max_ind = max(self.max_len, max_ind)\n min_ind = min(L - self.max_len, min_ind)\n attnmask[i, min_ind:max_ind] = False\n else:\n attnmask = None\n\n # Transformer\n for block in self.blocks:\n x = block(x, attn_mask=attnmask, tgt_key_padding_mask=pmask)\n\n # Output\n sample = self.final_layer(x) # (B, L, C)\n if self.avgbeta:\n betas = (sample[..., 126:136] * (~pmask[..., None])).sum(1) / length[:, None] # (B, C)\n betas = repeat(betas, \"b c -> b l c\", l=L)\n sample = torch.cat([sample[..., :126], betas, sample[..., 136:]], dim=-1)\n\n # Output (extra)\n pred_cam = None\n if self.pred_cam_head:\n pred_cam = self.pred_cam_head(x)\n pred_cam = pred_cam * self.pred_cam_std + self.pred_cam_mean\n torch.clamp_min_(pred_cam[..., 0], 0.25) # min_clamp s to 0.25 (prevent negative prediction)\n\n static_conf_logits = None\n if self.static_conf_head:\n static_conf_logits = self.static_conf_head(x) # (B, L, C')\n\n output = {\n \"pred_context\": x,\n \"pred_x\": sample,\n \"pred_cam\": pred_cam,\n \"static_conf_logits\": static_conf_logits,\n }\n return output\n\n\n# Add to MainStore\ngroup_name = \"network/gvhmr\"\nMainStore.store(\n name=\"relative_transformer\",\n node=builds(NetworkEncoderRoPE, populate_full_signature=True),\n group=group_name,\n)\n" }, { "path": "eval/GVHMR/hmr4d/network/hmr2/__init__.py", "content": "import torch\nfrom .hmr2 import HMR2\nfrom pathlib import Path\nfrom .configs import get_config\nfrom hmr4d import PROJ_ROOT\n\nHMR2A_CKPT = PROJ_ROOT / f\"inputs/checkpoints/hmr2/epoch=10-step=25000.ckpt\" # this is HMR2.0a, follow WHAM\n\n\ndef load_hmr2(checkpoint_path=HMR2A_CKPT):\n model_cfg = str((Path(__file__).parent / \"configs/model_config.yaml\").resolve())\n model_cfg = get_config(model_cfg)\n\n # Override some config values, to crop bbox correctly\n if (model_cfg.MODEL.BACKBONE.TYPE == \"vit\") and (\"BBOX_SHAPE\" not in model_cfg.MODEL):\n model_cfg.defrost()\n assert (\n model_cfg.MODEL.IMAGE_SIZE == 256\n ), f\"MODEL.IMAGE_SIZE ({model_cfg.MODEL.IMAGE_SIZE}) should be 256 for ViT backbone\"\n model_cfg.MODEL.BBOX_SHAPE = [192, 256] # (W, H)\n model_cfg.freeze()\n\n # Setup model and Load weights.\n # model = HMR2.load_from_checkpoint(checkpoint_path, strict=False, cfg=model_cfg)\n model = HMR2(model_cfg)\n\n state_dict = torch.load(checkpoint_path, map_location=\"cpu\")[\"state_dict\"]\n keys = [k for k in state_dict.keys() if k.split(\".\")[0] in [\"backbone\", \"smpl_head\"]]\n state_dict = {k: v for k, v in state_dict.items() if k in keys}\n model.load_state_dict(state_dict, strict=True)\n\n return model\n" }, { "path": "eval/GVHMR/hmr4d/network/hmr2/components/__init__.py", "content": "" }, { "path": "eval/GVHMR/hmr4d/network/hmr2/components/pose_transformer.py", "content": "from inspect import isfunction\nfrom typing import Callable, Optional\n\nimport torch\nfrom einops import rearrange\nfrom einops.layers.torch import Rearrange\nfrom torch import nn\n\nfrom .t_cond_mlp import (\n AdaptiveLayerNorm1D,\n FrequencyEmbedder,\n normalization_layer,\n)\n# from .vit import Attention, FeedForward\n\n\ndef exists(val):\n return val is not None\n\n\ndef default(val, d):\n if exists(val):\n return val\n return d() if isfunction(d) else d\n\n\nclass PreNorm(nn.Module):\n def __init__(self, dim: int, fn: Callable, norm: str = \"layer\", norm_cond_dim: int = -1):\n super().__init__()\n self.norm = normalization_layer(norm, dim, norm_cond_dim)\n self.fn = fn\n\n def forward(self, x: torch.Tensor, *args, **kwargs):\n if isinstance(self.norm, AdaptiveLayerNorm1D):\n return self.fn(self.norm(x, *args), **kwargs)\n else:\n return self.fn(self.norm(x), **kwargs)\n\n\nclass FeedForward(nn.Module):\n def __init__(self, dim, hidden_dim, dropout=0.0):\n super().__init__()\n self.net = nn.Sequential(\n nn.Linear(dim, hidden_dim),\n nn.GELU(),\n nn.Dropout(dropout),\n nn.Linear(hidden_dim, dim),\n nn.Dropout(dropout),\n )\n\n def forward(self, x):\n return self.net(x)\n\n\nclass Attention(nn.Module):\n def __init__(self, dim, heads=8, dim_head=64, dropout=0.0):\n super().__init__()\n inner_dim = dim_head * heads\n project_out = not (heads == 1 and dim_head == dim)\n\n self.heads = heads\n self.scale = dim_head**-0.5\n\n self.attend = nn.Softmax(dim=-1)\n self.dropout = nn.Dropout(dropout)\n\n self.to_qkv = nn.Linear(dim, inner_dim * 3, bias=False)\n\n self.to_out = (\n nn.Sequential(nn.Linear(inner_dim, dim), nn.Dropout(dropout))\n if project_out\n else nn.Identity()\n )\n\n def forward(self, x):\n qkv = self.to_qkv(x).chunk(3, dim=-1)\n q, k, v = map(lambda t: rearrange(t, \"b n (h d) -> b h n d\", h=self.heads), qkv)\n\n dots = torch.matmul(q, k.transpose(-1, -2)) * self.scale\n\n attn = self.attend(dots)\n attn = self.dropout(attn)\n\n out = torch.matmul(attn, v)\n out = rearrange(out, \"b h n d -> b n (h d)\")\n return self.to_out(out)\n\n\nclass CrossAttention(nn.Module):\n def __init__(self, dim, context_dim=None, heads=8, dim_head=64, dropout=0.0):\n super().__init__()\n inner_dim = dim_head * heads\n project_out = not (heads == 1 and dim_head == dim)\n\n self.heads = heads\n self.scale = dim_head**-0.5\n\n self.attend = nn.Softmax(dim=-1)\n self.dropout = nn.Dropout(dropout)\n\n context_dim = default(context_dim, dim)\n self.to_kv = nn.Linear(context_dim, inner_dim * 2, bias=False)\n self.to_q = nn.Linear(dim, inner_dim, bias=False)\n\n self.to_out = (\n nn.Sequential(nn.Linear(inner_dim, dim), nn.Dropout(dropout))\n if project_out\n else nn.Identity()\n )\n\n def forward(self, x, context=None):\n context = default(context, x)\n k, v = self.to_kv(context).chunk(2, dim=-1)\n q = self.to_q(x)\n q, k, v = map(lambda t: rearrange(t, \"b n (h d) -> b h n d\", h=self.heads), [q, k, v])\n\n dots = torch.matmul(q, k.transpose(-1, -2)) * self.scale\n\n attn = self.attend(dots)\n attn = self.dropout(attn)\n\n out = torch.matmul(attn, v)\n out = rearrange(out, \"b h n d -> b n (h d)\")\n return self.to_out(out)\n\n\nclass Transformer(nn.Module):\n def __init__(\n self,\n dim: int,\n depth: int,\n heads: int,\n dim_head: int,\n mlp_dim: int,\n dropout: float = 0.0,\n norm: str = \"layer\",\n norm_cond_dim: int = -1,\n ):\n super().__init__()\n self.layers = nn.ModuleList([])\n for _ in range(depth):\n sa = Attention(dim, heads=heads, dim_head=dim_head, dropout=dropout)\n ff = FeedForward(dim, mlp_dim, dropout=dropout)\n self.layers.append(\n nn.ModuleList(\n [\n PreNorm(dim, sa, norm=norm, norm_cond_dim=norm_cond_dim),\n PreNorm(dim, ff, norm=norm, norm_cond_dim=norm_cond_dim),\n ]\n )\n )\n\n def forward(self, x: torch.Tensor, *args):\n for attn, ff in self.layers:\n x = attn(x, *args) + x\n x = ff(x, *args) + x\n return x\n\n\nclass TransformerCrossAttn(nn.Module):\n def __init__(\n self,\n dim: int,\n depth: int,\n heads: int,\n dim_head: int,\n mlp_dim: int,\n dropout: float = 0.0,\n norm: str = \"layer\",\n norm_cond_dim: int = -1,\n context_dim: Optional[int] = None,\n ):\n super().__init__()\n self.layers = nn.ModuleList([])\n for _ in range(depth):\n sa = Attention(dim, heads=heads, dim_head=dim_head, dropout=dropout)\n ca = CrossAttention(\n dim, context_dim=context_dim, heads=heads, dim_head=dim_head, dropout=dropout\n )\n ff = FeedForward(dim, mlp_dim, dropout=dropout)\n self.layers.append(\n nn.ModuleList(\n [\n PreNorm(dim, sa, norm=norm, norm_cond_dim=norm_cond_dim),\n PreNorm(dim, ca, norm=norm, norm_cond_dim=norm_cond_dim),\n PreNorm(dim, ff, norm=norm, norm_cond_dim=norm_cond_dim),\n ]\n )\n )\n\n def forward(self, x: torch.Tensor, *args, context=None, context_list=None):\n if context_list is None:\n context_list = [context] * len(self.layers)\n if len(context_list) != len(self.layers):\n raise ValueError(f\"len(context_list) != len(self.layers) ({len(context_list)} != {len(self.layers)})\")\n\n for i, (self_attn, cross_attn, ff) in enumerate(self.layers):\n x = self_attn(x, *args) + x\n x = cross_attn(x, *args, context=context_list[i]) + x\n x = ff(x, *args) + x\n return x\n\n\nclass DropTokenDropout(nn.Module):\n def __init__(self, p: float = 0.1):\n super().__init__()\n if p < 0 or p > 1:\n raise ValueError(\n \"dropout probability has to be between 0 and 1, \" \"but got {}\".format(p)\n )\n self.p = p\n\n def forward(self, x: torch.Tensor):\n # x: (batch_size, seq_len, dim)\n if self.training and self.p > 0:\n zero_mask = torch.full_like(x[0, :, 0], self.p).bernoulli().bool()\n # TODO: permutation idx for each batch using torch.argsort\n if zero_mask.any():\n x = x[:, ~zero_mask, :]\n return x\n\n\nclass ZeroTokenDropout(nn.Module):\n def __init__(self, p: float = 0.1):\n super().__init__()\n if p < 0 or p > 1:\n raise ValueError(\n \"dropout probability has to be between 0 and 1, \" \"but got {}\".format(p)\n )\n self.p = p\n\n def forward(self, x: torch.Tensor):\n # x: (batch_size, seq_len, dim)\n if self.training and self.p > 0:\n zero_mask = torch.full_like(x[:, :, 0], self.p).bernoulli().bool()\n # Zero-out the masked tokens\n x[zero_mask, :] = 0\n return x\n\n\nclass TransformerEncoder(nn.Module):\n def __init__(\n self,\n num_tokens: int,\n token_dim: int,\n dim: int,\n depth: int,\n heads: int,\n mlp_dim: int,\n dim_head: int = 64,\n dropout: float = 0.0,\n emb_dropout: float = 0.0,\n emb_dropout_type: str = \"drop\",\n emb_dropout_loc: str = \"token\",\n norm: str = \"layer\",\n norm_cond_dim: int = -1,\n token_pe_numfreq: int = -1,\n ):\n super().__init__()\n if token_pe_numfreq > 0:\n token_dim_new = token_dim * (2 * token_pe_numfreq + 1)\n self.to_token_embedding = nn.Sequential(\n Rearrange(\"b n d -> (b n) d\", n=num_tokens, d=token_dim),\n FrequencyEmbedder(token_pe_numfreq, token_pe_numfreq - 1),\n Rearrange(\"(b n) d -> b n d\", n=num_tokens, d=token_dim_new),\n nn.Linear(token_dim_new, dim),\n )\n else:\n self.to_token_embedding = nn.Linear(token_dim, dim)\n self.pos_embedding = nn.Parameter(torch.randn(1, num_tokens, dim))\n if emb_dropout_type == \"drop\":\n self.dropout = DropTokenDropout(emb_dropout)\n elif emb_dropout_type == \"zero\":\n self.dropout = ZeroTokenDropout(emb_dropout)\n else:\n raise ValueError(f\"Unknown emb_dropout_type: {emb_dropout_type}\")\n self.emb_dropout_loc = emb_dropout_loc\n\n self.transformer = Transformer(\n dim, depth, heads, dim_head, mlp_dim, dropout, norm=norm, norm_cond_dim=norm_cond_dim\n )\n\n def forward(self, inp: torch.Tensor, *args, **kwargs):\n x = inp\n\n if self.emb_dropout_loc == \"input\":\n x = self.dropout(x)\n x = self.to_token_embedding(x)\n\n if self.emb_dropout_loc == \"token\":\n x = self.dropout(x)\n b, n, _ = x.shape\n x += self.pos_embedding[:, :n]\n\n if self.emb_dropout_loc == \"token_afterpos\":\n x = self.dropout(x)\n x = self.transformer(x, *args)\n return x\n\n\nclass TransformerDecoder(nn.Module):\n def __init__(\n self,\n num_tokens: int,\n token_dim: int,\n dim: int,\n depth: int,\n heads: int,\n mlp_dim: int,\n dim_head: int = 64,\n dropout: float = 0.0,\n emb_dropout: float = 0.0,\n emb_dropout_type: str = 'drop',\n norm: str = \"layer\",\n norm_cond_dim: int = -1,\n context_dim: Optional[int] = None,\n skip_token_embedding: bool = False,\n ):\n super().__init__()\n if not skip_token_embedding:\n self.to_token_embedding = nn.Linear(token_dim, dim)\n else:\n self.to_token_embedding = nn.Identity()\n if token_dim != dim:\n raise ValueError(\n f\"token_dim ({token_dim}) != dim ({dim}) when skip_token_embedding is True\"\n )\n\n self.pos_embedding = nn.Parameter(torch.randn(1, num_tokens, dim))\n if emb_dropout_type == \"drop\":\n self.dropout = DropTokenDropout(emb_dropout)\n elif emb_dropout_type == \"zero\":\n self.dropout = ZeroTokenDropout(emb_dropout)\n elif emb_dropout_type == \"normal\":\n self.dropout = nn.Dropout(emb_dropout)\n\n self.transformer = TransformerCrossAttn(\n dim,\n depth,\n heads,\n dim_head,\n mlp_dim,\n dropout,\n norm=norm,\n norm_cond_dim=norm_cond_dim,\n context_dim=context_dim,\n )\n\n def forward(self, inp: torch.Tensor, *args, context=None, context_list=None):\n x = self.to_token_embedding(inp)\n b, n, _ = x.shape\n\n x = self.dropout(x)\n x += self.pos_embedding[:, :n]\n\n x = self.transformer(x, *args, context=context, context_list=context_list)\n return x\n\n" }, { "path": "eval/GVHMR/hmr4d/network/hmr2/components/t_cond_mlp.py", "content": "import copy\nfrom typing import List, Optional\n\nimport torch\n\n\nclass AdaptiveLayerNorm1D(torch.nn.Module):\n def __init__(self, data_dim: int, norm_cond_dim: int):\n super().__init__()\n if data_dim <= 0:\n raise ValueError(f\"data_dim must be positive, but got {data_dim}\")\n if norm_cond_dim <= 0:\n raise ValueError(f\"norm_cond_dim must be positive, but got {norm_cond_dim}\")\n self.norm = torch.nn.LayerNorm(\n data_dim\n ) # TODO: Check if elementwise_affine=True is correct\n self.linear = torch.nn.Linear(norm_cond_dim, 2 * data_dim)\n torch.nn.init.zeros_(self.linear.weight)\n torch.nn.init.zeros_(self.linear.bias)\n\n def forward(self, x: torch.Tensor, t: torch.Tensor) -> torch.Tensor:\n # x: (batch, ..., data_dim)\n # t: (batch, norm_cond_dim)\n # return: (batch, data_dim)\n x = self.norm(x)\n alpha, beta = self.linear(t).chunk(2, dim=-1)\n\n # Add singleton dimensions to alpha and beta\n if x.dim() > 2:\n alpha = alpha.view(alpha.shape[0], *([1] * (x.dim() - 2)), alpha.shape[1])\n beta = beta.view(beta.shape[0], *([1] * (x.dim() - 2)), beta.shape[1])\n\n return x * (1 + alpha) + beta\n\n\nclass SequentialCond(torch.nn.Sequential):\n def forward(self, input, *args, **kwargs):\n for module in self:\n if isinstance(module, (AdaptiveLayerNorm1D, SequentialCond, ResidualMLPBlock)):\n # print(f'Passing on args to {module}', [a.shape for a in args])\n input = module(input, *args, **kwargs)\n else:\n # print(f'Skipping passing args to {module}', [a.shape for a in args])\n input = module(input)\n return input\n\n\ndef normalization_layer(norm: Optional[str], dim: int, norm_cond_dim: int = -1):\n if norm == \"batch\":\n return torch.nn.BatchNorm1d(dim)\n elif norm == \"layer\":\n return torch.nn.LayerNorm(dim)\n elif norm == \"ada\":\n assert norm_cond_dim > 0, f\"norm_cond_dim must be positive, got {norm_cond_dim}\"\n return AdaptiveLayerNorm1D(dim, norm_cond_dim)\n elif norm is None:\n return torch.nn.Identity()\n else:\n raise ValueError(f\"Unknown norm: {norm}\")\n\n\ndef linear_norm_activ_dropout(\n input_dim: int,\n output_dim: int,\n activation: torch.nn.Module = torch.nn.ReLU(),\n bias: bool = True,\n norm: Optional[str] = \"layer\", # Options: ada/batch/layer\n dropout: float = 0.0,\n norm_cond_dim: int = -1,\n) -> SequentialCond:\n layers = []\n layers.append(torch.nn.Linear(input_dim, output_dim, bias=bias))\n if norm is not None:\n layers.append(normalization_layer(norm, output_dim, norm_cond_dim))\n layers.append(copy.deepcopy(activation))\n if dropout > 0.0:\n layers.append(torch.nn.Dropout(dropout))\n return SequentialCond(*layers)\n\n\ndef create_simple_mlp(\n input_dim: int,\n hidden_dims: List[int],\n output_dim: int,\n activation: torch.nn.Module = torch.nn.ReLU(),\n bias: bool = True,\n norm: Optional[str] = \"layer\", # Options: ada/batch/layer\n dropout: float = 0.0,\n norm_cond_dim: int = -1,\n) -> SequentialCond:\n layers = []\n prev_dim = input_dim\n for hidden_dim in hidden_dims:\n layers.extend(\n linear_norm_activ_dropout(\n prev_dim, hidden_dim, activation, bias, norm, dropout, norm_cond_dim\n )\n )\n prev_dim = hidden_dim\n layers.append(torch.nn.Linear(prev_dim, output_dim, bias=bias))\n return SequentialCond(*layers)\n\n\nclass ResidualMLPBlock(torch.nn.Module):\n def __init__(\n self,\n input_dim: int,\n hidden_dim: int,\n num_hidden_layers: int,\n output_dim: int,\n activation: torch.nn.Module = torch.nn.ReLU(),\n bias: bool = True,\n norm: Optional[str] = \"layer\", # Options: ada/batch/layer\n dropout: float = 0.0,\n norm_cond_dim: int = -1,\n ):\n super().__init__()\n if not (input_dim == output_dim == hidden_dim):\n raise NotImplementedError(\n f\"input_dim {input_dim} != output_dim {output_dim} is not implemented\"\n )\n\n layers = []\n prev_dim = input_dim\n for i in range(num_hidden_layers):\n layers.append(\n linear_norm_activ_dropout(\n prev_dim, hidden_dim, activation, bias, norm, dropout, norm_cond_dim\n )\n )\n prev_dim = hidden_dim\n self.model = SequentialCond(*layers)\n self.skip = torch.nn.Identity()\n\n def forward(self, x: torch.Tensor, *args, **kwargs) -> torch.Tensor:\n return x + self.model(x, *args, **kwargs)\n\n\nclass ResidualMLP(torch.nn.Module):\n def __init__(\n self,\n input_dim: int,\n hidden_dim: int,\n num_hidden_layers: int,\n output_dim: int,\n activation: torch.nn.Module = torch.nn.ReLU(),\n bias: bool = True,\n norm: Optional[str] = \"layer\", # Options: ada/batch/layer\n dropout: float = 0.0,\n num_blocks: int = 1,\n norm_cond_dim: int = -1,\n ):\n super().__init__()\n self.input_dim = input_dim\n self.model = SequentialCond(\n linear_norm_activ_dropout(\n input_dim, hidden_dim, activation, bias, norm, dropout, norm_cond_dim\n ),\n *[\n ResidualMLPBlock(\n hidden_dim,\n hidden_dim,\n num_hidden_layers,\n hidden_dim,\n activation,\n bias,\n norm,\n dropout,\n norm_cond_dim,\n )\n for _ in range(num_blocks)\n ],\n torch.nn.Linear(hidden_dim, output_dim, bias=bias),\n )\n\n def forward(self, x: torch.Tensor, *args, **kwargs) -> torch.Tensor:\n return self.model(x, *args, **kwargs)\n\n\nclass FrequencyEmbedder(torch.nn.Module):\n def __init__(self, num_frequencies, max_freq_log2):\n super().__init__()\n frequencies = 2 ** torch.linspace(0, max_freq_log2, steps=num_frequencies)\n self.register_buffer(\"frequencies\", frequencies)\n\n def forward(self, x):\n # x should be of size (N,) or (N, D)\n N = x.size(0)\n if x.dim() == 1: # (N,)\n x = x.unsqueeze(1) # (N, D) where D=1\n x_unsqueezed = x.unsqueeze(-1) # (N, D, 1)\n scaled = self.frequencies.view(1, 1, -1) * x_unsqueezed # (N, D, num_frequencies)\n s = torch.sin(scaled)\n c = torch.cos(scaled)\n embedded = torch.cat([s, c, x_unsqueezed], dim=-1).view(\n N, -1\n ) # (N, D * 2 * num_frequencies + D)\n return embedded\n\n" }, { "path": "eval/GVHMR/hmr4d/network/hmr2/configs/__init__.py", "content": "import os\nfrom typing import Dict\nfrom yacs.config import CfgNode as CN\nfrom pathlib import Path\n\n# CACHE_DIR = os.path.join(os.environ.get(\"HOME\"), \"Code/4D-Humans/cache\")\n# CACHE_DIR_4DHUMANS = os.path.join(CACHE_DIR, \"4DHumans\")\n\n\ndef to_lower(x: Dict) -> Dict:\n \"\"\"\n Convert all dictionary keys to lowercase\n Args:\n x (dict): Input dictionary\n Returns:\n dict: Output dictionary with all keys converted to lowercase\n \"\"\"\n return {k.lower(): v for k, v in x.items()}\n\n\n_C = CN(new_allowed=True)\n\n_C.GENERAL = CN(new_allowed=True)\n_C.GENERAL.RESUME = True\n_C.GENERAL.TIME_TO_RUN = 3300\n_C.GENERAL.VAL_STEPS = 100\n_C.GENERAL.LOG_STEPS = 100\n_C.GENERAL.CHECKPOINT_STEPS = 20000\n_C.GENERAL.CHECKPOINT_DIR = \"checkpoints\"\n_C.GENERAL.SUMMARY_DIR = \"tensorboard\"\n_C.GENERAL.NUM_GPUS = 1\n_C.GENERAL.NUM_WORKERS = 4\n_C.GENERAL.MIXED_PRECISION = True\n_C.GENERAL.ALLOW_CUDA = True\n_C.GENERAL.PIN_MEMORY = False\n_C.GENERAL.DISTRIBUTED = False\n_C.GENERAL.LOCAL_RANK = 0\n_C.GENERAL.USE_SYNCBN = False\n_C.GENERAL.WORLD_SIZE = 1\n\n_C.TRAIN = CN(new_allowed=True)\n_C.TRAIN.NUM_EPOCHS = 100\n_C.TRAIN.BATCH_SIZE = 32\n_C.TRAIN.SHUFFLE = True\n_C.TRAIN.WARMUP = False\n_C.TRAIN.NORMALIZE_PER_IMAGE = False\n_C.TRAIN.CLIP_GRAD = False\n_C.TRAIN.CLIP_GRAD_VALUE = 1.0\n_C.LOSS_WEIGHTS = CN(new_allowed=True)\n\n_C.DATASETS = CN(new_allowed=True)\n\n_C.MODEL = CN(new_allowed=True)\n_C.MODEL.IMAGE_SIZE = 224\n\n_C.EXTRA = CN(new_allowed=True)\n_C.EXTRA.FOCAL_LENGTH = 5000\n\n_C.DATASETS.CONFIG = CN(new_allowed=True)\n_C.DATASETS.CONFIG.SCALE_FACTOR = 0.3\n_C.DATASETS.CONFIG.ROT_FACTOR = 30\n_C.DATASETS.CONFIG.TRANS_FACTOR = 0.02\n_C.DATASETS.CONFIG.COLOR_SCALE = 0.2\n_C.DATASETS.CONFIG.ROT_AUG_RATE = 0.6\n_C.DATASETS.CONFIG.TRANS_AUG_RATE = 0.5\n_C.DATASETS.CONFIG.DO_FLIP = True\n_C.DATASETS.CONFIG.FLIP_AUG_RATE = 0.5\n_C.DATASETS.CONFIG.EXTREME_CROP_AUG_RATE = 0.10\n\n\ndef default_config() -> CN:\n \"\"\"\n Get a yacs CfgNode object with the default config values.\n \"\"\"\n # Return a clone so that the defaults will not be altered\n # This is for the \"local variable\" use pattern\n return _C.clone()\n\n\ndef dataset_config(name=\"datasets_tar.yaml\") -> CN:\n \"\"\"\n Get dataset config file\n Returns:\n CfgNode: Dataset config as a yacs CfgNode object.\n \"\"\"\n cfg = CN(new_allowed=True)\n config_file = os.path.join(os.path.dirname(os.path.realpath(__file__)), name)\n cfg.merge_from_file(config_file)\n cfg.freeze()\n return cfg\n\n\ndef dataset_eval_config() -> CN:\n return dataset_config(\"datasets_eval.yaml\")\n\n\ndef get_config(config_file: str, merge: bool = True) -> CN:\n \"\"\"\n Read a config file and optionally merge it with the default config file.\n Args:\n config_file (str): Path to config file.\n merge (bool): Whether to merge with the default config or not.\n Returns:\n CfgNode: Config as a yacs CfgNode object.\n \"\"\"\n if merge:\n cfg = default_config()\n else:\n cfg = CN(new_allowed=True)\n cfg.merge_from_file(config_file)\n\n # ---- Update ---- #\n cfg.SMPL.MODEL_PATH = cfg.SMPL.MODEL_PATH # Not used\n cfg.SMPL.JOINT_REGRESSOR_EXTRA = cfg.SMPL.JOINT_REGRESSOR_EXTRA # Not Used\n cfg.SMPL.MEAN_PARAMS = str(Path(__file__).parent / \"smpl_mean_params.npz\")\n # ---------------- #\n\n cfg.freeze()\n return cfg\n" }, { "path": "eval/GVHMR/hmr4d/network/hmr2/configs/model_config.yaml", "content": "task_name: train\ntags:\n- dev\ntrain: true\ntest: false\nckpt_path: null\nseed: null\nDATASETS:\n TRAIN:\n H36M-TRAIN:\n WEIGHT: 0.3\n MPII-TRAIN:\n WEIGHT: 0.1\n COCO-TRAIN-2014:\n WEIGHT: 0.4\n MPI-INF-TRAIN:\n WEIGHT: 0.2\n VAL:\n COCO-VAL:\n WEIGHT: 1.0\n MOCAP: CMU-MOCAP\n CONFIG:\n SCALE_FACTOR: 0.3\n ROT_FACTOR: 30\n TRANS_FACTOR: 0.02\n COLOR_SCALE: 0.2\n ROT_AUG_RATE: 0.6\n TRANS_AUG_RATE: 0.5\n DO_FLIP: true\n FLIP_AUG_RATE: 0.5\n EXTREME_CROP_AUG_RATE: 0.1\ntrainer:\n _target_: pytorch_lightning.Trainer\n default_root_dir: ${paths.output_dir}\n accelerator: gpu\n devices: 8\n deterministic: false\n num_sanity_val_steps: 0\n log_every_n_steps: ${GENERAL.LOG_STEPS}\n val_check_interval: ${GENERAL.VAL_STEPS}\n precision: 16\n max_steps: ${GENERAL.TOTAL_STEPS}\n move_metrics_to_cpu: true\n limit_val_batches: 1\n track_grad_norm: 2\n strategy: ddp\n num_nodes: 1\n sync_batchnorm: true\npaths:\n root_dir: ${oc.env:PROJECT_ROOT}\n data_dir: ${paths.root_dir}/data/\n log_dir: /fsx/shubham/code/hmr2023/logs_hydra/\n output_dir: ${hydra:runtime.output_dir}\n work_dir: ${hydra:runtime.cwd}\nextras:\n ignore_warnings: false\n enforce_tags: true\n print_config: true\nexp_name: 3001d\nSMPL:\n MODEL_PATH: data/smpl\n GENDER: neutral\n NUM_BODY_JOINTS: 23\n JOINT_REGRESSOR_EXTRA: data/SMPL_to_J19.pkl\n MEAN_PARAMS: data/smpl_mean_params.npz\nEXTRA:\n FOCAL_LENGTH: 5000\n NUM_LOG_IMAGES: 4\n NUM_LOG_SAMPLES_PER_IMAGE: 8\n PELVIS_IND: 39\nMODEL:\n IMAGE_SIZE: 256\n IMAGE_MEAN:\n - 0.485\n - 0.456\n - 0.406\n IMAGE_STD:\n - 0.229\n - 0.224\n - 0.225\n BACKBONE:\n TYPE: vit\n FREEZE: true\n NUM_LAYERS: 50\n OUT_CHANNELS: 2048\n ADD_NECK: false\n FLOW:\n DIM: 144\n NUM_LAYERS: 4\n CONTEXT_FEATURES: 2048\n LAYER_HIDDEN_FEATURES: 1024\n LAYER_DEPTH: 2\n FC_HEAD:\n NUM_FEATURES: 1024\n SMPL_HEAD:\n TYPE: transformer_decoder\n IN_CHANNELS: 2048\n TRANSFORMER_DECODER:\n depth: 6\n heads: 8\n mlp_dim: 1024\n dim_head: 64\n dropout: 0.0\n emb_dropout: 0.0\n norm: layer\n context_dim: 1280\nGENERAL:\n TOTAL_STEPS: 100000\n LOG_STEPS: 100\n VAL_STEPS: 100\n CHECKPOINT_STEPS: 1000\n CHECKPOINT_SAVE_TOP_K: -1\n NUM_WORKERS: 6\n PREFETCH_FACTOR: 2\nTRAIN:\n LR: 0.0001\n WEIGHT_DECAY: 0.0001\n BATCH_SIZE: 512\n LOSS_REDUCTION: mean\n NUM_TRAIN_SAMPLES: 2\n NUM_TEST_SAMPLES: 64\n POSE_2D_NOISE_RATIO: 0.01\n SMPL_PARAM_NOISE_RATIO: 0.005\nLOSS_WEIGHTS:\n KEYPOINTS_3D: 0.05\n KEYPOINTS_2D: 0.01\n GLOBAL_ORIENT: 0.001\n BODY_POSE: 0.001\n BETAS: 0.0005\n ADVERSARIAL: 0.0005\nlocal: {}\n" }, { "path": "eval/GVHMR/hmr4d/network/hmr2/hmr2.py", "content": "import torch\nimport pytorch_lightning as pl\nfrom yacs.config import CfgNode\nfrom .vit import ViT\nfrom .smpl_head import SMPLTransformerDecoderHead\n\nfrom pytorch3d.transforms import matrix_to_axis_angle\nfrom hmr4d.utils.geo.hmr_cam import compute_transl_full_cam\n\n\nclass HMR2(pl.LightningModule):\n def __init__(self, cfg: CfgNode):\n super().__init__()\n self.cfg = cfg\n self.backbone = ViT(\n img_size=(256, 192),\n patch_size=16,\n embed_dim=1280,\n depth=32,\n num_heads=16,\n ratio=1,\n use_checkpoint=False,\n mlp_ratio=4,\n qkv_bias=True,\n drop_path_rate=0.55,\n )\n self.smpl_head = SMPLTransformerDecoderHead(cfg)\n\n def forward(self, batch, feat_mode=True):\n \"\"\"this file has been modified\n Args:\n feat_mode: default True, as we only need the feature token output for the HMR4D project;\n when False, the full process of HMR2 will be executed.\n \"\"\"\n # Backbone\n x = batch[\"img\"][:, :, :, 32:-32]\n vit_feats = self.backbone(x)\n\n # Output head\n if feat_mode:\n token_out = self.smpl_head(vit_feats, only_return_token_out=True) # (B, 1024)\n return token_out\n\n # return full process\n pred_smpl_params, pred_cam, _, token_out = self.smpl_head(vit_feats, only_return_token_out=False)\n output = {}\n output[\"token_out\"] = token_out\n output[\"smpl_params\"] = {\n \"body_pose\": matrix_to_axis_angle(pred_smpl_params[\"body_pose\"]).flatten(-2), # (B, 23, 3)\n \"betas\": pred_smpl_params[\"betas\"], # (B, 10)\n \"global_orient\": matrix_to_axis_angle(pred_smpl_params[\"global_orient\"])[:, 0], # (B, 3)\n \"transl\": compute_transl_full_cam(pred_cam, batch[\"bbx_xys\"], batch[\"K_fullimg\"]), # (B, 3)\n }\n\n return output\n" }, { "path": "eval/GVHMR/hmr4d/network/hmr2/smpl_head.py", "content": "import torch\nimport torch.nn as nn\nimport torch.nn.functional as F\nimport numpy as np\nimport einops\n\nfrom .utils.geometry import rot6d_to_rotmat, aa_to_rotmat\nfrom .components.pose_transformer import TransformerDecoder\n\n\nclass SMPLTransformerDecoderHead(nn.Module):\n \"\"\"Cross-attention based SMPL Transformer decoder\"\"\"\n\n def __init__(self, cfg):\n super().__init__()\n self.cfg = cfg\n self.joint_rep_type = cfg.MODEL.SMPL_HEAD.get(\"JOINT_REP\", \"6d\")\n self.joint_rep_dim = {\"6d\": 6, \"aa\": 3}[self.joint_rep_type]\n npose = self.joint_rep_dim * (cfg.SMPL.NUM_BODY_JOINTS + 1)\n self.npose = npose\n self.input_is_mean_shape = cfg.MODEL.SMPL_HEAD.get(\"TRANSFORMER_INPUT\", \"zero\") == \"mean_shape\"\n transformer_args = dict(\n num_tokens=1,\n token_dim=(npose + 10 + 3) if self.input_is_mean_shape else 1,\n dim=1024,\n )\n transformer_args.update(**dict(cfg.MODEL.SMPL_HEAD.TRANSFORMER_DECODER))\n self.transformer = TransformerDecoder(**transformer_args)\n dim = transformer_args[\"dim\"]\n self.decpose = nn.Linear(dim, npose)\n self.decshape = nn.Linear(dim, 10)\n self.deccam = nn.Linear(dim, 3)\n\n if cfg.MODEL.SMPL_HEAD.get(\"INIT_DECODER_XAVIER\", False):\n # True by default in MLP. False by default in Transformer\n nn.init.xavier_uniform_(self.decpose.weight, gain=0.01)\n nn.init.xavier_uniform_(self.decshape.weight, gain=0.01)\n nn.init.xavier_uniform_(self.deccam.weight, gain=0.01)\n\n mean_params = np.load(cfg.SMPL.MEAN_PARAMS)\n init_body_pose = torch.from_numpy(mean_params[\"pose\"].astype(np.float32)).unsqueeze(0)\n init_betas = torch.from_numpy(mean_params[\"shape\"].astype(\"float32\")).unsqueeze(0)\n init_cam = torch.from_numpy(mean_params[\"cam\"].astype(np.float32)).unsqueeze(0)\n self.register_buffer(\"init_body_pose\", init_body_pose)\n self.register_buffer(\"init_betas\", init_betas)\n self.register_buffer(\"init_cam\", init_cam)\n\n def forward(self, x, only_return_token_out=False):\n batch_size = x.shape[0]\n # vit pretrained backbone is channel-first. Change to token-first\n x = einops.rearrange(x, \"b c h w -> b (h w) c\")\n\n init_body_pose = self.init_body_pose.expand(batch_size, -1)\n init_betas = self.init_betas.expand(batch_size, -1)\n init_cam = self.init_cam.expand(batch_size, -1)\n\n # TODO: Convert init_body_pose to aa rep if needed\n if self.joint_rep_type == \"aa\":\n raise NotImplementedError\n\n pred_body_pose = init_body_pose\n pred_betas = init_betas\n pred_cam = init_cam\n pred_body_pose_list = []\n pred_betas_list = []\n pred_cam_list = []\n for i in range(self.cfg.MODEL.SMPL_HEAD.get(\"IEF_ITERS\", 1)):\n assert i == 0, \"Only support 1 iteration for now\"\n\n # Input token to transformer is zero token\n if self.input_is_mean_shape:\n token = torch.cat([pred_body_pose, pred_betas, pred_cam], dim=1)[:, None, :]\n else:\n token = torch.zeros(batch_size, 1, 1).to(x.device)\n\n # Pass through transformer\n token_out = self.transformer(token, context=x)\n token_out = token_out.squeeze(1) # (B, C)\n\n if only_return_token_out:\n return token_out\n else:\n # Readout from token_out\n pred_body_pose = self.decpose(token_out) + pred_body_pose\n pred_betas = self.decshape(token_out) + pred_betas\n pred_cam = self.deccam(token_out) + pred_cam\n pred_body_pose_list.append(pred_body_pose)\n pred_betas_list.append(pred_betas)\n pred_cam_list.append(pred_cam)\n\n # Convert self.joint_rep_type -> rotmat\n joint_conversion_fn = {\"6d\": rot6d_to_rotmat, \"aa\": lambda x: aa_to_rotmat(x.view(-1, 3).contiguous())}[\n self.joint_rep_type\n ]\n\n pred_smpl_params_list = {}\n pred_smpl_params_list[\"body_pose\"] = torch.cat(\n [joint_conversion_fn(pbp).view(batch_size, -1, 3, 3)[:, 1:, :, :] for pbp in pred_body_pose_list], dim=0\n )\n pred_smpl_params_list[\"betas\"] = torch.cat(pred_betas_list, dim=0)\n pred_smpl_params_list[\"cam\"] = torch.cat(pred_cam_list, dim=0)\n pred_body_pose = joint_conversion_fn(pred_body_pose).view(batch_size, self.cfg.SMPL.NUM_BODY_JOINTS + 1, 3, 3)\n\n pred_smpl_params = {\n \"global_orient\": pred_body_pose[:, [0]],\n \"body_pose\": pred_body_pose[:, 1:],\n \"betas\": pred_betas,\n }\n return pred_smpl_params, pred_cam, pred_smpl_params_list, token_out\n" }, { "path": "eval/GVHMR/hmr4d/network/hmr2/utils/geometry.py", "content": "from typing import Optional\nimport torch\nfrom torch.nn import functional as F\n\n\ndef aa_to_rotmat(theta: torch.Tensor):\n \"\"\"\n Convert axis-angle representation to rotation matrix.\n Works by first converting it to a quaternion.\n Args:\n theta (torch.Tensor): Tensor of shape (B, 3) containing axis-angle representations.\n Returns:\n torch.Tensor: Corresponding rotation matrices with shape (B, 3, 3).\n \"\"\"\n norm = torch.norm(theta + 1e-8, p=2, dim=1)\n angle = torch.unsqueeze(norm, -1)\n normalized = torch.div(theta, angle)\n angle = angle * 0.5\n v_cos = torch.cos(angle)\n v_sin = torch.sin(angle)\n quat = torch.cat([v_cos, v_sin * normalized], dim=1)\n return quat_to_rotmat(quat)\n\n\ndef quat_to_rotmat(quat: torch.Tensor) -> torch.Tensor:\n \"\"\"\n Convert quaternion representation to rotation matrix.\n Args:\n quat (torch.Tensor) of shape (B, 4); 4 <===> (w, x, y, z).\n Returns:\n torch.Tensor: Corresponding rotation matrices with shape (B, 3, 3).\n \"\"\"\n norm_quat = quat\n norm_quat = norm_quat / norm_quat.norm(p=2, dim=1, keepdim=True)\n w, x, y, z = norm_quat[:, 0], norm_quat[:, 1], norm_quat[:, 2], norm_quat[:, 3]\n\n B = quat.size(0)\n\n w2, x2, y2, z2 = w.pow(2), x.pow(2), y.pow(2), z.pow(2)\n wx, wy, wz = w * x, w * y, w * z\n xy, xz, yz = x * y, x * z, y * z\n\n rotMat = torch.stack(\n [\n w2 + x2 - y2 - z2,\n 2 * xy - 2 * wz,\n 2 * wy + 2 * xz,\n 2 * wz + 2 * xy,\n w2 - x2 + y2 - z2,\n 2 * yz - 2 * wx,\n 2 * xz - 2 * wy,\n 2 * wx + 2 * yz,\n w2 - x2 - y2 + z2,\n ],\n dim=1,\n ).view(B, 3, 3)\n return rotMat\n\n\ndef rot6d_to_rotmat(x: torch.Tensor) -> torch.Tensor:\n \"\"\"\n Convert 6D rotation representation to 3x3 rotation matrix.\n Based on Zhou et al., \"On the Continuity of Rotation Representations in Neural Networks\", CVPR 2019\n Args:\n x (torch.Tensor): (B,6) Batch of 6-D rotation representations.\n Returns:\n torch.Tensor: Batch of corresponding rotation matrices with shape (B,3,3).\n \"\"\"\n x = x.reshape(-1, 2, 3).permute(0, 2, 1).contiguous()\n a1 = x[:, :, 0]\n a2 = x[:, :, 1]\n b1 = F.normalize(a1)\n b2 = F.normalize(a2 - torch.einsum(\"bi,bi->b\", b1, a2).unsqueeze(-1) * b1)\n b3 = torch.cross(b1, b2)\n return torch.stack((b1, b2, b3), dim=-1)\n\n\ndef perspective_projection(\n points: torch.Tensor,\n translation: torch.Tensor,\n focal_length: torch.Tensor,\n camera_center: Optional[torch.Tensor] = None,\n rotation: Optional[torch.Tensor] = None,\n) -> torch.Tensor:\n \"\"\"\n Computes the perspective projection of a set of 3D points.\n Args:\n points (torch.Tensor): Tensor of shape (B, N, 3) containing the input 3D points.\n translation (torch.Tensor): Tensor of shape (B, 3) containing the 3D camera translation.\n focal_length (torch.Tensor): Tensor of shape (B, 2) containing the focal length in pixels.\n camera_center (torch.Tensor): Tensor of shape (B, 2) containing the camera center in pixels.\n rotation (torch.Tensor): Tensor of shape (B, 3, 3) containing the camera rotation.\n Returns:\n torch.Tensor: Tensor of shape (B, N, 2) containing the projection of the input points.\n \"\"\"\n batch_size = points.shape[0]\n if rotation is None:\n rotation = torch.eye(3, device=points.device, dtype=points.dtype).unsqueeze(0).expand(batch_size, -1, -1)\n if camera_center is None:\n camera_center = torch.zeros(batch_size, 2, device=points.device, dtype=points.dtype)\n # Populate intrinsic camera matrix K.\n K = torch.zeros([batch_size, 3, 3], device=points.device, dtype=points.dtype)\n K[:, 0, 0] = focal_length[:, 0]\n K[:, 1, 1] = focal_length[:, 1]\n K[:, 2, 2] = 1.0\n K[:, :-1, -1] = camera_center\n\n # Transform points\n points = torch.einsum(\"bij,bkj->bki\", rotation, points)\n points = points + translation.unsqueeze(1)\n\n # Apply perspective distortion\n projected_points = points / points[:, :, -1].unsqueeze(-1)\n\n # Apply camera intrinsics\n projected_points = torch.einsum(\"bij,bkj->bki\", K, projected_points)\n\n return projected_points[:, :, :-1]\n" }, { "path": "eval/GVHMR/hmr4d/network/hmr2/utils/preproc.py", "content": "import cv2\nimport numpy as np\nimport torch\nfrom pathlib import Path\n\nIMAGE_MEAN = torch.tensor([0.485, 0.456, 0.406])\nIMAGE_STD = torch.tensor([0.229, 0.224, 0.225])\n\n\ndef expand_to_aspect_ratio(input_shape, target_aspect_ratio=[192, 256]):\n \"\"\"Increase the size of the bounding box to match the target shape.\"\"\"\n if target_aspect_ratio is None:\n return input_shape\n\n try:\n w, h = input_shape\n except (ValueError, TypeError):\n return input_shape\n\n w_t, h_t = target_aspect_ratio\n if h / w < h_t / w_t:\n h_new = max(w * h_t / w_t, h)\n w_new = w\n else:\n h_new = h\n w_new = max(h * w_t / h_t, w)\n if h_new < h or w_new < w:\n breakpoint()\n return np.array([w_new, h_new])\n\n\ndef crop_and_resize(img, bbx_xy, bbx_s, dst_size=256, enlarge_ratio=1.2):\n \"\"\"\n Args:\n img: (H, W, 3)\n bbx_xy: (2,)\n bbx_s: scalar\n \"\"\"\n hs = bbx_s * enlarge_ratio / 2\n src = np.stack(\n [\n bbx_xy - hs, # left-up corner\n bbx_xy + np.array([hs, -hs]), # right-up corner\n bbx_xy, # center\n ]\n ).astype(np.float32)\n dst = np.array([[0, 0], [dst_size - 1, 0], [dst_size / 2 - 0.5, dst_size / 2 - 0.5]], dtype=np.float32)\n A = cv2.getAffineTransform(src, dst)\n\n img_crop = cv2.warpAffine(img, A, (dst_size, dst_size), flags=cv2.INTER_LINEAR)\n bbx_xys_final = np.array([*bbx_xy, bbx_s * enlarge_ratio])\n return img_crop, bbx_xys_final\n" }, { "path": "eval/GVHMR/hmr4d/network/hmr2/utils/smpl_wrapper.py", "content": "import torch\nimport numpy as np\nimport pickle\nfrom typing import Optional\nimport smplx\nfrom smplx.lbs import vertices2joints\nfrom smplx.utils import SMPLOutput\n\n\nclass SMPL(smplx.SMPLLayer):\n def __init__(self, *args, joint_regressor_extra: Optional[str] = None, update_hips: bool = False, **kwargs):\n \"\"\"\n Extension of the official SMPL implementation to support more joints.\n Args:\n Same as SMPLLayer.\n joint_regressor_extra (str): Path to extra joint regressor.\n \"\"\"\n super(SMPL, self).__init__(*args, **kwargs)\n smpl_to_openpose = [24, 12, 17, 19, 21, 16, 18, 20, 0, 2, 5, 8, 1, 4, 7, 25, 26, 27, 28, 29, 30, 31, 32, 33, 34]\n\n if joint_regressor_extra is not None:\n self.register_buffer(\n \"joint_regressor_extra\",\n torch.tensor(pickle.load(open(joint_regressor_extra, \"rb\"), encoding=\"latin1\"), dtype=torch.float32),\n )\n self.register_buffer(\"joint_map\", torch.tensor(smpl_to_openpose, dtype=torch.long))\n self.update_hips = update_hips\n\n def forward(self, *args, **kwargs) -> SMPLOutput:\n \"\"\"\n Run forward pass. Same as SMPL and also append an extra set of joints if joint_regressor_extra is specified.\n \"\"\"\n smpl_output = super(SMPL, self).forward(*args, **kwargs)\n joints = smpl_output.joints[:, self.joint_map, :]\n if self.update_hips:\n joints[:, [9, 12]] = (\n joints[:, [9, 12]]\n + 0.25 * (joints[:, [9, 12]] - joints[:, [12, 9]])\n + 0.5 * (joints[:, [8]] - 0.5 * (joints[:, [9, 12]] + joints[:, [12, 9]]))\n )\n if hasattr(self, \"joint_regressor_extra\"):\n extra_joints = vertices2joints(self.joint_regressor_extra, smpl_output.vertices)\n joints = torch.cat([joints, extra_joints], dim=1)\n smpl_output.joints = joints\n return smpl_output\n" }, { "path": "eval/GVHMR/hmr4d/network/hmr2/vit.py", "content": "# Copyright (c) OpenMMLab. All rights reserved.\nimport math\n\nimport torch\nfrom functools import partial\nimport torch.nn as nn\nimport torch.nn.functional as F\nimport torch.utils.checkpoint as checkpoint\n\nfrom timm.models.layers import drop_path, to_2tuple, trunc_normal_\n\ndef vit(cfg):\n return ViT(\n img_size=(256, 192),\n patch_size=16,\n embed_dim=1280,\n depth=32,\n num_heads=16,\n ratio=1,\n use_checkpoint=False,\n mlp_ratio=4,\n qkv_bias=True,\n drop_path_rate=0.55,\n )\n\ndef get_abs_pos(abs_pos, h, w, ori_h, ori_w, has_cls_token=True):\n \"\"\"\n Calculate absolute positional embeddings. If needed, resize embeddings and remove cls_token\n dimension for the original embeddings.\n Args:\n abs_pos (Tensor): absolute positional embeddings with (1, num_position, C).\n has_cls_token (bool): If true, has 1 embedding in abs_pos for cls token.\n hw (Tuple): size of input image tokens.\n\n Returns:\n Absolute positional embeddings after processing with shape (1, H, W, C)\n \"\"\"\n cls_token = None\n B, L, C = abs_pos.shape\n if has_cls_token:\n cls_token = abs_pos[:, 0:1]\n abs_pos = abs_pos[:, 1:]\n\n if ori_h != h or ori_w != w:\n new_abs_pos = F.interpolate(\n abs_pos.reshape(1, ori_h, ori_w, -1).permute(0, 3, 1, 2),\n size=(h, w),\n mode=\"bicubic\",\n align_corners=False,\n ).permute(0, 2, 3, 1).reshape(B, -1, C)\n\n else:\n new_abs_pos = abs_pos\n \n if cls_token is not None:\n new_abs_pos = torch.cat([cls_token, new_abs_pos], dim=1)\n return new_abs_pos\n\nclass DropPath(nn.Module):\n \"\"\"Drop paths (Stochastic Depth) per sample (when applied in main path of residual blocks).\n \"\"\"\n def __init__(self, drop_prob=None):\n super(DropPath, self).__init__()\n self.drop_prob = drop_prob\n\n def forward(self, x):\n return drop_path(x, self.drop_prob, self.training)\n \n def extra_repr(self):\n return 'p={}'.format(self.drop_prob)\n\nclass Mlp(nn.Module):\n def __init__(self, in_features, hidden_features=None, out_features=None, act_layer=nn.GELU, drop=0.):\n super().__init__()\n out_features = out_features or in_features\n hidden_features = hidden_features or in_features\n self.fc1 = nn.Linear(in_features, hidden_features)\n self.act = act_layer()\n self.fc2 = nn.Linear(hidden_features, out_features)\n self.drop = nn.Dropout(drop)\n\n def forward(self, x):\n x = self.fc1(x)\n x = self.act(x)\n x = self.fc2(x)\n x = self.drop(x)\n return x\n\nclass Attention(nn.Module):\n def __init__(\n self, dim, num_heads=8, qkv_bias=False, qk_scale=None, attn_drop=0.,\n proj_drop=0., attn_head_dim=None,):\n super().__init__()\n self.num_heads = num_heads\n head_dim = dim // num_heads\n self.dim = dim\n\n if attn_head_dim is not None:\n head_dim = attn_head_dim\n all_head_dim = head_dim * self.num_heads\n\n self.scale = qk_scale or head_dim ** -0.5\n\n self.qkv = nn.Linear(dim, all_head_dim * 3, bias=qkv_bias)\n\n self.attn_drop = nn.Dropout(attn_drop)\n self.proj = nn.Linear(all_head_dim, dim)\n self.proj_drop = nn.Dropout(proj_drop)\n\n def forward(self, x):\n B, N, C = x.shape\n qkv = self.qkv(x)\n qkv = qkv.reshape(B, N, 3, self.num_heads, -1).permute(2, 0, 3, 1, 4)\n q, k, v = qkv[0], qkv[1], qkv[2] # make torchscript happy (cannot use tensor as tuple)\n\n q = q * self.scale\n attn = (q @ k.transpose(-2, -1))\n\n attn = attn.softmax(dim=-1)\n attn = self.attn_drop(attn)\n\n x = (attn @ v).transpose(1, 2).reshape(B, N, -1)\n x = self.proj(x)\n x = self.proj_drop(x)\n\n return x\n\nclass Block(nn.Module):\n\n def __init__(self, dim, num_heads, mlp_ratio=4., qkv_bias=False, qk_scale=None, \n drop=0., attn_drop=0., drop_path=0., act_layer=nn.GELU, \n norm_layer=nn.LayerNorm, attn_head_dim=None\n ):\n super().__init__()\n \n self.norm1 = norm_layer(dim)\n self.attn = Attention(\n dim, num_heads=num_heads, qkv_bias=qkv_bias, qk_scale=qk_scale,\n attn_drop=attn_drop, proj_drop=drop, attn_head_dim=attn_head_dim\n )\n\n # NOTE: drop path for stochastic depth, we shall see if this is better than dropout here\n self.drop_path = DropPath(drop_path) if drop_path > 0. else nn.Identity()\n self.norm2 = norm_layer(dim)\n mlp_hidden_dim = int(dim * mlp_ratio)\n self.mlp = Mlp(in_features=dim, hidden_features=mlp_hidden_dim, act_layer=act_layer, drop=drop)\n\n def forward(self, x):\n x = x + self.drop_path(self.attn(self.norm1(x)))\n x = x + self.drop_path(self.mlp(self.norm2(x)))\n return x\n\n\nclass PatchEmbed(nn.Module):\n \"\"\" Image to Patch Embedding\n \"\"\"\n def __init__(self, img_size=224, patch_size=16, in_chans=3, embed_dim=768, ratio=1):\n super().__init__()\n img_size = to_2tuple(img_size)\n patch_size = to_2tuple(patch_size)\n num_patches = (img_size[1] // patch_size[1]) * (img_size[0] // patch_size[0]) * (ratio ** 2)\n self.patch_shape = (int(img_size[0] // patch_size[0] * ratio), int(img_size[1] // patch_size[1] * ratio))\n self.origin_patch_shape = (int(img_size[0] // patch_size[0]), int(img_size[1] // patch_size[1]))\n self.img_size = img_size\n self.patch_size = patch_size\n self.num_patches = num_patches\n\n self.proj = nn.Conv2d(in_chans, embed_dim, kernel_size=patch_size, stride=(patch_size[0] // ratio), padding=4 + 2 * (ratio//2-1))\n\n def forward(self, x, **kwargs):\n B, C, H, W = x.shape\n x = self.proj(x)\n Hp, Wp = x.shape[2], x.shape[3]\n\n x = x.flatten(2).transpose(1, 2)\n return x, (Hp, Wp)\n\n\nclass HybridEmbed(nn.Module):\n \"\"\" CNN Feature Map Embedding\n Extract feature map from CNN, flatten, project to embedding dim.\n \"\"\"\n def __init__(self, backbone, img_size=224, feature_size=None, in_chans=3, embed_dim=768):\n super().__init__()\n assert isinstance(backbone, nn.Module)\n img_size = to_2tuple(img_size)\n self.img_size = img_size\n self.backbone = backbone\n if feature_size is None:\n with torch.no_grad():\n training = backbone.training\n if training:\n backbone.eval()\n o = self.backbone(torch.zeros(1, in_chans, img_size[0], img_size[1]))[-1]\n feature_size = o.shape[-2:]\n feature_dim = o.shape[1]\n backbone.train(training)\n else:\n feature_size = to_2tuple(feature_size)\n feature_dim = self.backbone.feature_info.channels()[-1]\n self.num_patches = feature_size[0] * feature_size[1]\n self.proj = nn.Linear(feature_dim, embed_dim)\n\n def forward(self, x):\n x = self.backbone(x)[-1]\n x = x.flatten(2).transpose(1, 2)\n x = self.proj(x)\n return x\n\n\nclass ViT(nn.Module):\n\n def __init__(self,\n img_size=224, patch_size=16, in_chans=3, num_classes=80, embed_dim=768, depth=12,\n num_heads=12, mlp_ratio=4., qkv_bias=False, qk_scale=None, drop_rate=0., attn_drop_rate=0.,\n drop_path_rate=0., hybrid_backbone=None, norm_layer=None, use_checkpoint=False, \n frozen_stages=-1, ratio=1, last_norm=True,\n patch_padding='pad', freeze_attn=False, freeze_ffn=False,\n ):\n # Protect mutable default arguments\n super(ViT, self).__init__()\n norm_layer = norm_layer or partial(nn.LayerNorm, eps=1e-6)\n self.num_classes = num_classes\n self.num_features = self.embed_dim = embed_dim # num_features for consistency with other models\n self.frozen_stages = frozen_stages\n self.use_checkpoint = use_checkpoint\n self.patch_padding = patch_padding\n self.freeze_attn = freeze_attn\n self.freeze_ffn = freeze_ffn\n self.depth = depth\n\n if hybrid_backbone is not None:\n self.patch_embed = HybridEmbed(\n hybrid_backbone, img_size=img_size, in_chans=in_chans, embed_dim=embed_dim)\n else:\n self.patch_embed = PatchEmbed(\n img_size=img_size, patch_size=patch_size, in_chans=in_chans, embed_dim=embed_dim, ratio=ratio)\n num_patches = self.patch_embed.num_patches\n\n # since the pretraining model has class token\n self.pos_embed = nn.Parameter(torch.zeros(1, num_patches + 1, embed_dim))\n\n dpr = [x.item() for x in torch.linspace(0, drop_path_rate, depth)] # stochastic depth decay rule\n\n self.blocks = nn.ModuleList([\n Block(\n dim=embed_dim, num_heads=num_heads, mlp_ratio=mlp_ratio, qkv_bias=qkv_bias, qk_scale=qk_scale,\n drop=drop_rate, attn_drop=attn_drop_rate, drop_path=dpr[i], norm_layer=norm_layer,\n )\n for i in range(depth)])\n\n self.last_norm = norm_layer(embed_dim) if last_norm else nn.Identity()\n\n if self.pos_embed is not None:\n trunc_normal_(self.pos_embed, std=.02)\n\n self._freeze_stages()\n\n def _freeze_stages(self):\n \"\"\"Freeze parameters.\"\"\"\n if self.frozen_stages >= 0:\n self.patch_embed.eval()\n for param in self.patch_embed.parameters():\n param.requires_grad = False\n\n for i in range(1, self.frozen_stages + 1):\n m = self.blocks[i]\n m.eval()\n for param in m.parameters():\n param.requires_grad = False\n\n if self.freeze_attn:\n for i in range(0, self.depth):\n m = self.blocks[i]\n m.attn.eval()\n m.norm1.eval()\n for param in m.attn.parameters():\n param.requires_grad = False\n for param in m.norm1.parameters():\n param.requires_grad = False\n\n if self.freeze_ffn:\n self.pos_embed.requires_grad = False\n self.patch_embed.eval()\n for param in self.patch_embed.parameters():\n param.requires_grad = False\n for i in range(0, self.depth):\n m = self.blocks[i]\n m.mlp.eval()\n m.norm2.eval()\n for param in m.mlp.parameters():\n param.requires_grad = False\n for param in m.norm2.parameters():\n param.requires_grad = False\n\n def init_weights(self):\n \"\"\"Initialize the weights in backbone.\n Args:\n pretrained (str, optional): Path to pre-trained weights.\n Defaults to None.\n \"\"\"\n def _init_weights(m):\n if isinstance(m, nn.Linear):\n trunc_normal_(m.weight, std=.02)\n if isinstance(m, nn.Linear) and m.bias is not None:\n nn.init.constant_(m.bias, 0)\n elif isinstance(m, nn.LayerNorm):\n nn.init.constant_(m.bias, 0)\n nn.init.constant_(m.weight, 1.0)\n\n self.apply(_init_weights)\n\n def get_num_layers(self):\n return len(self.blocks)\n\n @torch.jit.ignore\n def no_weight_decay(self):\n return {'pos_embed', 'cls_token'}\n\n def forward_features(self, x):\n B, C, H, W = x.shape\n x, (Hp, Wp) = self.patch_embed(x)\n\n if self.pos_embed is not None:\n # fit for multiple GPU training\n # since the first element for pos embed (sin-cos manner) is zero, it will cause no difference\n x = x + self.pos_embed[:, 1:] + self.pos_embed[:, :1]\n\n for blk in self.blocks:\n if self.use_checkpoint:\n x = checkpoint.checkpoint(blk, x)\n else:\n x = blk(x)\n\n x = self.last_norm(x)\n\n xp = x.permute(0, 2, 1).reshape(B, -1, Hp, Wp).contiguous()\n\n return xp\n\n def forward(self, x):\n x = self.forward_features(x)\n return x\n\n def train(self, mode=True):\n \"\"\"Convert the model into training mode.\"\"\"\n super().train(mode)\n self._freeze_stages()\n" }, { "path": "eval/GVHMR/hmr4d/utils/body_model/README.md", "content": "# README\n\nContents of this folder are modified from HuMoR repository." }, { "path": "eval/GVHMR/hmr4d/utils/body_model/__init__.py", "content": "from .body_model import BodyModel\nfrom .body_model_smplh import BodyModelSMPLH\nfrom .body_model_smplx import BodyModelSMPLX\n" }, { "path": "eval/GVHMR/hmr4d/utils/body_model/body_model.py", "content": "from turtle import forward\nimport numpy as np\n\nimport torch\nimport torch.nn as nn\n\nfrom smplx import SMPL, SMPLH, SMPLX\nfrom smplx.vertex_ids import vertex_ids\nfrom smplx.utils import Struct\n\n\nclass BodyModel(nn.Module):\n \"\"\" \n Wrapper around SMPLX body model class. \n modified by Zehong Shen\n \"\"\"\n\n def __init__(self,\n bm_path,\n num_betas=16,\n use_vtx_selector=False,\n model_type='smplh'):\n super().__init__()\n '''\n Creates the body model object at the given path.\n\n :param bm_path: path to the body model pkl file\n :param model_type: one of [smpl, smplh, smplx]\n :param use_vtx_selector: if true, returns additional vertices as joints that correspond to OpenPose joints\n '''\n self.use_vtx_selector = use_vtx_selector\n cur_vertex_ids = None\n if self.use_vtx_selector:\n cur_vertex_ids = vertex_ids[model_type]\n data_struct = None\n if '.npz' in bm_path:\n # smplx does not support .npz by default, so have to load in manually\n smpl_dict = np.load(bm_path, encoding='latin1')\n data_struct = Struct(**smpl_dict)\n # print(smpl_dict.files)\n if model_type == 'smplh':\n data_struct.hands_componentsl = np.zeros((0))\n data_struct.hands_componentsr = np.zeros((0))\n data_struct.hands_meanl = np.zeros((15 * 3))\n data_struct.hands_meanr = np.zeros((15 * 3))\n V, D, B = data_struct.shapedirs.shape\n data_struct.shapedirs = np.concatenate([data_struct.shapedirs, np.zeros(\n (V, D, SMPL.SHAPE_SPACE_DIM-B))], axis=-1) # super hacky way to let smplh use 16-size beta\n kwargs = {\n 'model_type': model_type,\n 'data_struct': data_struct,\n 'num_betas': num_betas,\n 'vertex_ids': cur_vertex_ids,\n 'use_pca': False,\n 'flat_hand_mean': True,\n # - enable variable batchsize, since we don't need module variable - #\n 'create_body_pose': False,\n 'create_betas': False,\n 'create_global_orient': False,\n 'create_transl': False,\n 'create_left_hand_pose': False,\n 'create_right_hand_pose': False,\n }\n assert(model_type in ['smpl', 'smplh', 'smplx'])\n if model_type == 'smpl':\n self.bm = SMPL(bm_path, **kwargs)\n self.num_joints = SMPL.NUM_JOINTS\n elif model_type == 'smplh':\n self.bm = SMPLH(bm_path, **kwargs)\n self.num_joints = SMPLH.NUM_JOINTS\n elif model_type == 'smplx':\n self.bm = SMPLX(bm_path, **kwargs)\n self.num_joints = SMPLX.NUM_JOINTS\n\n self.model_type = model_type\n\n def forward(self, root_orient=None, pose_body=None, pose_hand=None, pose_jaw=None, pose_eye=None, betas=None,\n trans=None, dmpls=None, expression=None, return_dict=False, **kwargs):\n '''\n Note dmpls are not supported.\n '''\n assert(dmpls is None)\n B = pose_body.shape[0]\n if pose_hand is None:\n pose_hand = torch.zeros((B, 2*SMPLH.NUM_HAND_JOINTS*3), device=pose_body.device)\n if len(betas.shape) == 1:\n betas = betas.reshape((1, -1)).expand(B, -1)\n\n out_obj = self.bm(\n betas=betas,\n global_orient=root_orient,\n body_pose=pose_body,\n left_hand_pose=pose_hand[:, :(SMPLH.NUM_HAND_JOINTS*3)],\n right_hand_pose=pose_hand[:, (SMPLH.NUM_HAND_JOINTS*3):],\n transl=trans,\n expression=expression,\n jaw_pose=pose_jaw,\n leye_pose=None if pose_eye is None else pose_eye[:, :3],\n reye_pose=None if pose_eye is None else pose_eye[:, 3:],\n return_full_pose=True,\n **kwargs\n )\n\n out = {\n 'v': out_obj.vertices,\n 'f': self.bm.faces_tensor,\n 'Jtr': out_obj.joints,\n }\n\n if not self.use_vtx_selector:\n # don't need extra joints\n out['Jtr'] = out['Jtr'][:, :self.num_joints+1] # add one for the root\n\n if not return_dict:\n out = Struct(**out)\n\n return out\n\n def forward_motion(self, **kwargs):\n B, W, _ = kwargs['pose_body'].shape\n kwargs = {k: v.reshape(B*W, v.shape[-1]) for k, v in kwargs.items()}\n\n smpl_opt = self.forward(**kwargs)\n smpl_opt.v = smpl_opt.v.reshape(B, W, -1, 3)\n smpl_opt.Jtr = smpl_opt.Jtr.reshape(B, W, -1, 3)\n\n return smpl_opt\n" }, { "path": "eval/GVHMR/hmr4d/utils/body_model/body_model_smplh.py", "content": "import torch\nimport torch.nn as nn\nimport smplx\n\nkwargs_disable_member_var = {\n \"create_body_pose\": False,\n \"create_betas\": False,\n \"create_global_orient\": False,\n \"create_transl\": False,\n \"create_left_hand_pose\": False,\n \"create_right_hand_pose\": False,\n}\n\n\nclass BodyModelSMPLH(nn.Module):\n \"\"\"Support Batch inference\"\"\"\n\n def __init__(self, model_path, **kwargs):\n super().__init__()\n # enable flexible batchsize, handle missing variable at forward()\n kwargs.update(kwargs_disable_member_var)\n self.bm = smplx.create(model_path=model_path, **kwargs)\n self.faces = self.bm.faces\n self.is_smpl = kwargs.get(\"model_type\", \"smpl\") == \"smpl\"\n if not self.is_smpl:\n self.hand_pose_dim = self.bm.num_pca_comps if self.bm.use_pca else 3 * self.bm.NUM_HAND_JOINTS\n\n # For fast computing of skeleton under beta\n shapedirs = self.bm.shapedirs # (V, 3, 10)\n J_regressor = self.bm.J_regressor[:22, :] # (22, V)\n v_template = self.bm.v_template # (V, 3)\n J_template = J_regressor @ v_template # (22, 3)\n J_shapedirs = torch.einsum(\"jv, vcd -> jcd\", J_regressor, shapedirs) # (22, 3, 10)\n self.register_buffer(\"J_template\", J_template, False)\n self.register_buffer(\"J_shapedirs\", J_shapedirs, False)\n\n def forward(\n self,\n betas=None,\n global_orient=None,\n transl=None,\n body_pose=None,\n left_hand_pose=None,\n right_hand_pose=None,\n **kwargs\n ):\n\n device, dtype = self.bm.shapedirs.device, self.bm.shapedirs.dtype\n\n model_vars = [betas, global_orient, body_pose, transl, left_hand_pose, right_hand_pose]\n batch_size = 1\n for var in model_vars:\n if var is None:\n continue\n batch_size = max(batch_size, len(var))\n\n if global_orient is None:\n global_orient = torch.zeros([batch_size, 3], dtype=dtype, device=device)\n if body_pose is None:\n body_pose = (\n torch.zeros(3 * self.bm.NUM_BODY_JOINTS, device=device, dtype=dtype)[None]\n .expand(batch_size, -1)\n .contiguous()\n )\n if not self.is_smpl:\n if left_hand_pose is None:\n left_hand_pose = (\n torch.zeros(self.hand_pose_dim, device=device, dtype=dtype)[None]\n .expand(batch_size, -1)\n .contiguous()\n )\n if right_hand_pose is None:\n right_hand_pose = (\n torch.zeros(self.hand_pose_dim, device=device, dtype=dtype)[None]\n .expand(batch_size, -1)\n .contiguous()\n )\n if betas is None:\n betas = torch.zeros([batch_size, self.bm.num_betas], dtype=dtype, device=device)\n if transl is None:\n transl = torch.zeros([batch_size, 3], dtype=dtype, device=device)\n\n bm_out = self.bm(\n betas=betas,\n global_orient=global_orient,\n body_pose=body_pose,\n left_hand_pose=left_hand_pose,\n right_hand_pose=right_hand_pose,\n transl=transl,\n **kwargs\n )\n\n return bm_out\n\n def get_skeleton(self, betas):\n \"\"\"betas: (*, 10) -> skeleton_beta: (*, 22, 3)\"\"\"\n skeleton_beta = self.J_template + torch.einsum(\"...d, jcd -> ...jc\", betas, self.J_shapedirs) # (22, 3)\n return skeleton_beta\n" }, { "path": "eval/GVHMR/hmr4d/utils/body_model/body_model_smplx.py", "content": "import torch\nimport torch.nn as nn\nimport smplx\n\nkwargs_disable_member_var = {\n \"create_body_pose\": False,\n \"create_betas\": False,\n \"create_global_orient\": False,\n \"create_transl\": False,\n \"create_left_hand_pose\": False,\n \"create_right_hand_pose\": False,\n \"create_expression\": False,\n \"create_jaw_pose\": False,\n \"create_leye_pose\": False,\n \"create_reye_pose\": False,\n}\n\n\nclass BodyModelSMPLX(nn.Module):\n \"\"\"Support Batch inference\"\"\"\n\n def __init__(self, model_path, **kwargs):\n super().__init__()\n # enable flexible batchsize, handle missing variable at forward()\n kwargs.update(kwargs_disable_member_var)\n self.bm = smplx.create(model_path=model_path, **kwargs)\n self.faces = self.bm.faces\n self.hand_pose_dim = self.bm.num_pca_comps if self.bm.use_pca else 3 * self.bm.NUM_HAND_JOINTS\n\n # For fast computing of skeleton under beta\n shapedirs = self.bm.shapedirs # (V, 3, 10)\n J_regressor = self.bm.J_regressor[:22, :] # (22, V)\n v_template = self.bm.v_template # (V, 3)\n J_template = J_regressor @ v_template # (22, 3)\n J_shapedirs = torch.einsum(\"jv, vcd -> jcd\", J_regressor, shapedirs) # (22, 3, 10)\n self.register_buffer(\"J_template\", J_template, False)\n self.register_buffer(\"J_shapedirs\", J_shapedirs, False)\n\n def forward(\n self,\n betas=None,\n global_orient=None,\n transl=None,\n body_pose=None,\n left_hand_pose=None,\n right_hand_pose=None,\n expression=None,\n jaw_pose=None,\n leye_pose=None,\n reye_pose=None,\n **kwargs\n ):\n\n device, dtype = self.bm.shapedirs.device, self.bm.shapedirs.dtype\n\n model_vars = [\n betas,\n global_orient,\n body_pose,\n transl,\n expression,\n left_hand_pose,\n right_hand_pose,\n jaw_pose,\n leye_pose,\n reye_pose,\n ]\n batch_size = 1\n for var in model_vars:\n if var is None:\n continue\n batch_size = max(batch_size, len(var))\n\n if global_orient is None:\n global_orient = torch.zeros([batch_size, 3], dtype=dtype, device=device)\n if body_pose is None:\n body_pose = (\n torch.zeros(3 * self.bm.NUM_BODY_JOINTS, device=device, dtype=dtype)[None]\n .expand(batch_size, -1)\n .contiguous()\n )\n if left_hand_pose is None:\n left_hand_pose = (\n torch.zeros(self.hand_pose_dim, device=device, dtype=dtype)[None].expand(batch_size, -1).contiguous()\n )\n if right_hand_pose is None:\n right_hand_pose = (\n torch.zeros(self.hand_pose_dim, device=device, dtype=dtype)[None].expand(batch_size, -1).contiguous()\n )\n if jaw_pose is None:\n jaw_pose = torch.zeros([batch_size, 3], dtype=dtype, device=device)\n if leye_pose is None:\n leye_pose = torch.zeros([batch_size, 3], dtype=dtype, device=device)\n if reye_pose is None:\n reye_pose = torch.zeros([batch_size, 3], dtype=dtype, device=device)\n if expression is None:\n expression = torch.zeros([batch_size, self.bm.num_expression_coeffs], dtype=dtype, device=device)\n if betas is None:\n betas = torch.zeros([batch_size, self.bm.num_betas], dtype=dtype, device=device)\n if transl is None:\n transl = torch.zeros([batch_size, 3], dtype=dtype, device=device)\n\n bm_out = self.bm(\n betas=betas,\n global_orient=global_orient,\n body_pose=body_pose,\n left_hand_pose=left_hand_pose,\n right_hand_pose=right_hand_pose,\n transl=transl,\n expression=expression,\n jaw_pose=jaw_pose,\n leye_pose=leye_pose,\n reye_pose=reye_pose,\n **kwargs\n )\n\n return bm_out\n\n def get_skeleton(self, betas):\n \"\"\"betas: (*, 10) -> skeleton_beta: (*, 22, 3)\"\"\"\n skeleton_beta = self.J_template + torch.einsum(\"...d, jcd -> ...jc\", betas, self.J_shapedirs) # (22, 3)\n return skeleton_beta\n\n def forward_bfc(self, **kwargs):\n \"\"\"Wrap (B, F, C) to (B*F, C) and unwrap (B*F, C) to (B, F, C)\"\"\"\n for k in kwargs:\n assert len(kwargs[k].shape) == 3\n B, F = kwargs[\"body_pose\"].shape[:2]\n smplx_out = self.forward(**{k: v.reshape(B * F, -1) for k, v in kwargs.items()})\n smplx_out.vertices = smplx_out.vertices.reshape(B, F, -1, 3)\n smplx_out.joints = smplx_out.joints.reshape(B, F, -1, 3)\n return smplx_out\n" }, { "path": "eval/GVHMR/hmr4d/utils/body_model/min_lbs.py", "content": "import numpy as np\nimport torch\nimport torch.nn as nn\nimport torch.nn.functional as F\nfrom pytorch3d.transforms import axis_angle_to_matrix\nfrom smplx.utils import Struct, to_np, to_tensor\nfrom hmr4d.utils.smplx_utils import forward_kinematics_motion\n\n\nclass MinimalLBS(nn.Module):\n def __init__(self, sp_ids, bm_dir='models/smplh', num_betas=16, model_type='smplh', **kwargs):\n super().__init__()\n self.num_betas = num_betas\n self.sensor_point_vid = torch.tensor(sp_ids)\n\n # load struct data on predefined sensor-point\n self.load_struct_on_sp(f'{bm_dir}/male/model.npz', prefix='male')\n self.load_struct_on_sp(f'{bm_dir}/female/model.npz', prefix='female')\n\n def load_struct_on_sp(self, bm_path, prefix='m'):\n \"\"\"\n Load 4 weights from body-model-struct. \n Keep the sensor points only. Use prefix to label different bm.\n \"\"\"\n num_betas = self.num_betas\n sp_vid = self.sensor_point_vid\n # load data\n data_struct = Struct(**np.load(bm_path, encoding='latin1'))\n\n # v-template\n v_template = to_tensor(to_np(data_struct.v_template)) # (V, 3)\n v_template_sp = v_template[sp_vid] # (N, 3)\n self.register_buffer(f'{prefix}_v_template_sp', v_template_sp, False)\n\n # shapedirs\n shapedirs = to_tensor(to_np(data_struct.shapedirs[:, :, :num_betas])) # (V, 3, NB)\n shapedirs_sp = shapedirs[sp_vid]\n self.register_buffer(f'{prefix}_shapedirs_sp', shapedirs_sp, False)\n\n # posedirs\n posedirs = to_tensor(to_np(data_struct.posedirs)) # (V, 3, 51*9)\n posedirs_sp = posedirs[sp_vid]\n posedirs_sp = posedirs_sp.reshape(len(sp_vid)*3, -1).T # (51*9, N*3)\n self.register_buffer(f'{prefix}_posedirs_sp', posedirs_sp, False)\n\n # lbs_weights\n lbs_weights = to_tensor(to_np(data_struct.weights)) # (V, J+1)\n lbs_weights_sp = lbs_weights[sp_vid]\n self.register_buffer(f'{prefix}_lbs_weights_sp', lbs_weights_sp, False)\n\n def forward(self, root_orient=None, pose_body=None, trans=None, betas=None, A=None, recompute_A=False, genders=None,\n joints_zero=None):\n \"\"\"\n Args:\n root_orient, Optional: (B, T, 3)\n pose_body: (B, T, J*3)\n trans: (B, T, 3)\n betas: (B, T, 16)\n A, Optional: (B, T, J+1, 4, 4)\n recompute_A: if True, root_orient should be given, otherwise use A\n genders, List: ['male', 'female', ...]\n joints_zero: (B, J+1, 3), required when recompute_A is True\n Returns:\n sensor_verts: (B, T, N, 3)\n \"\"\"\n B, T = pose_body.shape[:2]\n\n v_template = torch.stack([getattr(self, f'{g}_v_template_sp') for g in genders]) # (B, N, 3)\n shapedirs = torch.stack([getattr(self, f'{g}_shapedirs_sp') for g in genders]) # (B, N, 3, NB)\n posedirs = torch.stack([getattr(self, f'{g}_posedirs_sp') for g in genders]) # (B, 51*9, N*3)\n lbs_weights = torch.stack([getattr(self, f'{g}_lbs_weights_sp') for g in genders]) # (B, N, J+1)\n\n # ===== LBS, handle T ===== #\n # 2. Add shape contribution\n if betas.shape[1] == 1:\n betas = betas.expand(-1, T, -1)\n blend_shape = torch.einsum('btl,bmkl->btmk', [betas, shapedirs])\n v_shaped = v_template[:, None] + blend_shape\n\n # 3. Add pose blend shapes\n ident = torch.eye(3).to(pose_body)\n aa = pose_body.reshape(B, T, -1, 3)\n R = axis_angle_to_matrix(aa)\n pose_feature = (R - ident).view(B, T, -1)\n dim_pf = pose_feature.shape[-1]\n # (B, T, P) @ (B, P, N*3) -> (B, T, N, 3)\n pose_offsets = torch.matmul(pose_feature, posedirs[:, :dim_pf]).view(B, T, -1, 3)\n v_posed = pose_offsets + v_shaped\n\n # 4. Compute A\n if recompute_A:\n _, _, A = forward_kinematics_motion(root_orient, pose_body, trans, joints_zero)\n\n # 5. Skinning\n W = lbs_weights\n # (B, 1, N, J+1)) @ (B, T, J+1, 16)\n num_joints = A.shape[-3] # 22\n Ts = torch.matmul(W[:, None, :, :num_joints], A.view(B, T, num_joints, 16))\n Ts = Ts.view(B, T, -1, 4, 4) # (B, T, N, 4, 4)\n v_posed_homo = F.pad(v_posed, (0, 1), value=1) # (B, T, N, 4)\n v_homo = torch.matmul(Ts, torch.unsqueeze(v_posed_homo, dim=-1))\n\n # 6. translate\n sensor_verts = v_homo[:, :, :, :3, 0] + trans[:, :, None]\n\n return sensor_verts\n" }, { "path": "eval/GVHMR/hmr4d/utils/body_model/smpl_lite.py", "content": "import numpy as np\nimport torch\nimport torch.nn as nn\nimport torch.nn.functional as F\nfrom pathlib import Path\nfrom pytorch3d.transforms import axis_angle_to_matrix\nfrom smplx.utils import Struct, to_np, to_tensor\nfrom einops import einsum, rearrange\nfrom time import time\n\nimport pickle\n\nfrom .smplx_lite import batch_rigid_transform_v2\n\n\nclass SmplLite(nn.Module):\n def __init__(\n self,\n model_path=\"inputs/checkpoints/body_models/smpl\",\n gender=\"neutral\",\n num_betas=10,\n ):\n super().__init__()\n\n # Load the model\n model_path = Path(model_path)\n if model_path.is_dir():\n smpl_path = Path(model_path) / f\"SMPL_{gender.upper()}.pkl\"\n else:\n smpl_path = model_path\n assert smpl_path.exists()\n with open(smpl_path, \"rb\") as smpl_file:\n data_struct = Struct(**pickle.load(smpl_file, encoding=\"latin1\"))\n self.faces = data_struct.f # (F, 3)\n\n self.register_smpl_buffers(data_struct, num_betas)\n self.register_fast_skeleton_computing_buffers()\n\n def register_smpl_buffers(self, data_struct, num_betas):\n # shapedirs, (V, 3, N_betas), V=10475 for SMPLX\n shapedirs = to_tensor(to_np(data_struct.shapedirs[:, :, :num_betas])).float()\n self.register_buffer(\"shapedirs\", shapedirs, False)\n\n # v_template, (V, 3)\n v_template = to_tensor(to_np(data_struct.v_template)).float()\n self.register_buffer(\"v_template\", v_template, False)\n\n # J_regressor, (J, V), J=55 for SMPLX\n J_regressor = to_tensor(to_np(data_struct.J_regressor)).float()\n self.register_buffer(\"J_regressor\", J_regressor, False)\n\n # posedirs, (54*9, V, 3), note that the first global_orient is not included\n posedirs = to_tensor(to_np(data_struct.posedirs)).float() # (V, 3, 54*9)\n posedirs = rearrange(posedirs, \"v c n -> n v c\")\n self.register_buffer(\"posedirs\", posedirs, False)\n\n # lbs_weights, (V, J), J=55\n lbs_weights = to_tensor(to_np(data_struct.weights)).float()\n self.register_buffer(\"lbs_weights\", lbs_weights, False)\n\n # parents, (J), long\n parents = to_tensor(to_np(data_struct.kintree_table[0])).long()\n parents[0] = -1\n self.register_buffer(\"parents\", parents, False)\n\n def register_fast_skeleton_computing_buffers(self):\n # For fast computing of skeleton under beta\n J_template = self.J_regressor @ self.v_template # (J, 3)\n J_shapedirs = torch.einsum(\"jv, vcd -> jcd\", self.J_regressor, self.shapedirs) # (J, 3, 10)\n self.register_buffer(\"J_template\", J_template, False)\n self.register_buffer(\"J_shapedirs\", J_shapedirs, False)\n\n def get_skeleton(self, betas):\n return self.J_template + einsum(betas, self.J_shapedirs, \"... k, j c k -> ... j c\")\n\n def forward(\n self,\n body_pose,\n betas,\n global_orient,\n transl,\n ):\n \"\"\"\n Args:\n body_pose: (B, L, 63)\n betas: (B, L, 10)\n global_orient: (B, L, 3)\n transl: (B, L, 3)\n Returns:\n vertices: (B, L, V, 3)\n \"\"\"\n # 1. Convert [global_orient, body_pose] to rot_mats\n full_pose = torch.cat([global_orient, body_pose], dim=-1)\n rot_mats = axis_angle_to_matrix(full_pose.reshape(*full_pose.shape[:-1], full_pose.shape[-1] // 3, 3))\n\n # 2. Forward Kinematics\n J = self.get_skeleton(betas) # (*, 55, 3)\n A = batch_rigid_transform_v2(rot_mats, J, self.parents)[1]\n\n # 3. Canonical v_posed = v_template + shaped_offsets + pose_offsets\n pose_feature = rot_mats[..., 1:, :, :] - rot_mats.new([[1, 0, 0], [0, 1, 0], [0, 0, 1]])\n pose_feature = pose_feature.view(*pose_feature.shape[:-3], -1) # (*, 55*3*3)\n v_posed = (\n self.v_template\n + einsum(betas, self.shapedirs, \"... k, v c k -> ... v c\")\n + einsum(pose_feature, self.posedirs, \"... k, k v c -> ... v c\")\n )\n del pose_feature, rot_mats, full_pose\n\n # 4. Skinning\n T = einsum(self.lbs_weights, A, \"v j, ... j c d -> ... v c d\")\n verts = einsum(T[..., :3, :3], v_posed, \"... v c d, ... v d -> ... v c\") + T[..., :3, 3]\n\n # 5. Translation\n verts = verts + transl[..., None, :]\n return verts\n\n\nclass SmplxLiteJ24(SmplLite):\n def __init__(self, **kwargs):\n super().__init__(**kwargs)\n\n # Compute mapping\n smpl2j24 = self.J_regressor # (24, 6890)\n\n jids, smplx_vids = torch.where(smpl2j24 != 0)\n interestd = torch.zeros([len(smplx_vids), 24])\n for idx, (jid, smplx_vid) in enumerate(zip(jids, smplx_vids)):\n interestd[idx, jid] = smpl2j24[jid, smplx_vid]\n self.register_buffer(\"interestd\", interestd, False) # (236, 24)\n\n # Update to vertices of interest\n self.v_template = self.v_template[smplx_vids].clone() # (V', 3)\n self.shapedirs = self.shapedirs[smplx_vids].clone() # (V', 3, K)\n self.posedirs = self.posedirs[:, smplx_vids].clone() # (K, V', 3)\n self.lbs_weights = self.lbs_weights[smplx_vids].clone() # (V', J)\n\n def forward(self, body_pose, betas, global_orient, transl):\n \"\"\"Returns: joints (*, J, 3). (B, L) or (B,) are both supported.\"\"\"\n # Use super class's forward to get verts\n verts = super().forward(body_pose, betas, global_orient, transl) # (*, 236, 3)\n joints = einsum(self.interestd, verts, \"v j, ... v c -> ... j c\")\n return joints\n" }, { "path": "eval/GVHMR/hmr4d/utils/body_model/smpl_vert_segmentation.json", "content": "{\n \"rightHand\": [\n 5442, \n 5443, \n 5444, \n 5445, \n 5446, \n 5447, \n 5448, \n 5449, \n 5450, \n 5451, \n 5452, \n 5453, \n 5454, \n 5455, \n 5456, \n 5457, \n 5458, \n 5459, \n 5460, \n 5461, \n 5462, \n 5463, \n 5464, \n 5465, \n 5466, \n 5467, \n 5468, \n 5469, \n 5470, \n 5471, \n 5472, \n 5473, \n 5474, \n 5475, \n 5476, \n 5477, \n 5478, \n 5479, \n 5480, \n 5481, \n 5482, \n 5483, \n 5484, \n 5485, \n 5486, \n 5487, \n 5492, \n 5493, \n 5494, \n 5495, \n 5496, \n 5497, \n 5502, \n 5503, \n 5504, \n 5505, \n 5506, \n 5507, \n 5508, \n 5509, \n 5510, \n 5511, \n 5512, \n 5513, \n 5514, \n 5515, \n 5516, \n 5517, \n 5518, \n 5519, \n 5520, \n 5521, \n 5522, \n 5523, \n 5524, \n 5525, \n 5526, \n 5527, \n 5530, \n 5531, \n 5532, \n 5533, \n 5534, \n 5535, \n 5536, \n 5537, \n 5538, \n 5539, \n 5540, \n 5541, \n 5542, \n 5543, \n 5544, \n 5545, \n 5546, \n 5547, \n 5548, \n 5549, \n 5550, \n 5551, \n 5552, \n 5553, \n 5554, \n 5555, \n 5556, \n 5557, \n 5558, \n 5559, \n 5560, \n 5561, \n 5562, \n 5569, \n 5571, \n 5574, \n 5575, \n 5576, \n 5577, \n 5578, \n 5579, \n 5580, \n 5581, \n 5582, \n 5583, \n 5588, \n 5589, \n 5592, \n 5593, \n 5594, \n 5595, \n 5596, \n 5597, \n 5598, \n 5599, \n 5600, \n 5601, \n 5602, \n 5603, \n 5604, \n 5605, \n 5610, \n 5611, \n 5612, \n 5613, \n 5614, \n 5621, \n 5622, \n 5625, \n 5631, \n 5632, \n 5633, \n 5634, \n 5635, \n 5636, \n 5637, \n 5638, \n 5639, \n 5640, \n 5641, \n 5643, \n 5644, \n 5645, \n 5646, \n 5649, \n 5650, \n 5652, \n 5653, \n 5654, \n 5655, \n 5656, \n 5657, \n 5658, \n 5659, \n 5660, \n 5661, \n 5662, \n 5663, \n 5664, \n 5667, \n 5670, \n 5671, \n 5672, \n 5673, \n 5674, \n 5675, \n 5682, \n 5683, \n 5684, \n 5685, \n 5686, \n 5687, \n 5688, \n 5689, \n 5690, \n 5692, \n 5695, \n 5697, \n 5698, \n 5699, \n 5700, \n 5701, \n 5707, \n 5708, \n 5709, \n 5710, \n 5711, \n 5712, \n 5713, \n 5714, \n 5715, \n 5716, \n 5717, \n 5718, \n 5719, \n 5720, \n 5721, \n 5723, \n 5724, \n 5725, \n 5726, \n 5727, \n 5728, \n 5729, \n 5730, \n 5731, \n 5732, \n 5735, \n 5736, \n 5737, \n 5738, \n 5739, \n 5740, \n 5745, \n 5746, \n 5748, \n 5749, \n 5750, \n 5751, \n 5752, \n 6056, \n 6057, \n 6066, \n 6067, \n 6158, \n 6159, \n 6160, \n 6161, \n 6162, \n 6163, \n 6164, \n 6165, \n 6166, \n 6167, \n 6168, \n 6169, \n 6170, \n 6171, \n 6172, \n 6173, \n 6174, \n 6175, \n 6176, \n 6177, \n 6178, \n 6179, \n 6180, \n 6181, \n 6182, \n 6183, \n 6184, \n 6185, \n 6186, \n 6187, \n 6188, \n 6189, \n 6190, \n 6191, \n 6192, \n 6193, \n 6194, \n 6195, \n 6196, \n 6197, \n 6198, \n 6199, \n 6200, \n 6201, \n 6202, \n 6203, \n 6204, \n 6205, \n 6206, \n 6207, \n 6208, \n 6209, \n 6210, \n 6211, \n 6212, \n 6213, \n 6214, \n 6215, \n 6216, \n 6217, \n 6218, \n 6219, \n 6220, \n 6221, \n 6222, \n 6223, \n 6224, \n 6225, \n 6226, \n 6227, \n 6228, \n 6229, \n 6230, \n 6231, \n 6232, \n 6233, \n 6234, \n 6235, \n 6236, \n 6237, \n 6238, \n 6239\n ], \n \"rightUpLeg\": [\n 4320, \n 4321, \n 4323, \n 4324, \n 4333, \n 4334, \n 4335, \n 4336, \n 4337, \n 4338, \n 4339, \n 4340, \n 4356, \n 4357, \n 4358, \n 4359, \n 4360, \n 4361, \n 4362, \n 4363, \n 4364, \n 4365, \n 4366, \n 4367, \n 4383, \n 4384, \n 4385, \n 4386, \n 4387, \n 4388, \n 4389, \n 4390, \n 4391, \n 4392, \n 4393, \n 4394, \n 4395, \n 4396, \n 4397, \n 4398, \n 4399, \n 4400, \n 4401, \n 4419, \n 4420, \n 4421, \n 4422, \n 4430, \n 4431, \n 4432, \n 4433, \n 4434, \n 4435, \n 4436, \n 4437, \n 4438, \n 4439, \n 4440, \n 4441, \n 4442, \n 4443, \n 4444, \n 4445, \n 4446, \n 4447, \n 4448, \n 4449, \n 4450, \n 4451, \n 4452, \n 4453, \n 4454, \n 4455, \n 4456, \n 4457, \n 4458, \n 4459, \n 4460, \n 4461, \n 4462, \n 4463, \n 4464, \n 4465, \n 4466, \n 4467, \n 4468, \n 4469, \n 4470, \n 4471, \n 4472, \n 4473, \n 4474, \n 4475, \n 4476, \n 4477, \n 4478, \n 4479, \n 4480, \n 4481, \n 4482, \n 4483, \n 4484, \n 4485, \n 4486, \n 4487, \n 4488, \n 4489, \n 4490, \n 4491, \n 4492, \n 4493, \n 4494, \n 4495, \n 4496, \n 4497, \n 4498, \n 4499, \n 4500, \n 4501, \n 4502, \n 4503, \n 4504, \n 4505, \n 4506, \n 4507, \n 4508, \n 4509, \n 4510, \n 4511, \n 4512, \n 4513, \n 4514, \n 4515, \n 4516, \n 4517, \n 4518, \n 4519, \n 4520, \n 4521, \n 4522, \n 4523, \n 4524, \n 4525, \n 4526, \n 4527, \n 4528, \n 4529, \n 4530, \n 4531, \n 4532, \n 4623, \n 4624, \n 4625, \n 4626, \n 4627, \n 4628, \n 4629, \n 4630, \n 4631, \n 4632, \n 4633, \n 4634, \n 4645, \n 4646, \n 4647, \n 4648, \n 4649, \n 4650, \n 4651, \n 4652, \n 4653, \n 4654, \n 4655, \n 4656, \n 4657, \n 4658, \n 4659, \n 4660, \n 4670, \n 4671, \n 4672, \n 4673, \n 4704, \n 4705, \n 4706, \n 4707, \n 4708, \n 4709, \n 4710, \n 4711, \n 4712, \n 4713, \n 4745, \n 4746, \n 4757, \n 4758, \n 4759, \n 4760, \n 4801, \n 4802, \n 4829, \n 4834, \n 4835, \n 4836, \n 4837, \n 4838, \n 4839, \n 4840, \n 4841, \n 4924, \n 4925, \n 4926, \n 4928, \n 4929, \n 4930, \n 4931, \n 4932, \n 4933, \n 4934, \n 4935, \n 4936, \n 4948, \n 4949, \n 4950, \n 4951, \n 4952, \n 4970, \n 4971, \n 4972, \n 4973, \n 4983, \n 4984, \n 4985, \n 4986, \n 4987, \n 4988, \n 4989, \n 4990, \n 4991, \n 4992, \n 4993, \n 5004, \n 5005, \n 6546, \n 6547, \n 6548, \n 6549, \n 6552, \n 6553, \n 6554, \n 6555, \n 6556, \n 6873, \n 6877\n ], \n \"leftArm\": [\n 626, \n 627, \n 628, \n 629, \n 634, \n 635, \n 680, \n 681, \n 716, \n 717, \n 718, \n 719, \n 769, \n 770, \n 771, \n 772, \n 773, \n 774, \n 775, \n 776, \n 777, \n 778, \n 779, \n 780, \n 784, \n 785, \n 786, \n 787, \n 788, \n 789, \n 790, \n 791, \n 792, \n 793, \n 1231, \n 1232, \n 1233, \n 1234, \n 1258, \n 1259, \n 1260, \n 1261, \n 1271, \n 1281, \n 1282, \n 1310, \n 1311, \n 1314, \n 1315, \n 1340, \n 1341, \n 1342, \n 1343, \n 1355, \n 1356, \n 1357, \n 1358, \n 1376, \n 1377, \n 1378, \n 1379, \n 1380, \n 1381, \n 1382, \n 1383, \n 1384, \n 1385, \n 1386, \n 1387, \n 1388, \n 1389, \n 1390, \n 1391, \n 1392, \n 1393, \n 1394, \n 1395, \n 1396, \n 1397, \n 1398, \n 1399, \n 1400, \n 1402, \n 1403, \n 1405, \n 1406, \n 1407, \n 1408, \n 1409, \n 1410, \n 1411, \n 1412, \n 1413, \n 1414, \n 1415, \n 1416, \n 1428, \n 1429, \n 1430, \n 1431, \n 1432, \n 1433, \n 1438, \n 1439, \n 1440, \n 1441, \n 1442, \n 1443, \n 1444, \n 1445, \n 1502, \n 1505, \n 1506, \n 1507, \n 1508, \n 1509, \n 1510, \n 1538, \n 1541, \n 1542, \n 1543, \n 1545, \n 1619, \n 1620, \n 1621, \n 1622, \n 1631, \n 1632, \n 1633, \n 1634, \n 1635, \n 1636, \n 1637, \n 1638, \n 1639, \n 1640, \n 1641, \n 1642, \n 1645, \n 1646, \n 1647, \n 1648, \n 1649, \n 1650, \n 1651, \n 1652, \n 1653, \n 1654, \n 1655, \n 1656, \n 1658, \n 1659, \n 1661, \n 1662, \n 1664, \n 1666, \n 1667, \n 1668, \n 1669, \n 1670, \n 1671, \n 1672, \n 1673, \n 1674, \n 1675, \n 1676, \n 1677, \n 1678, \n 1679, \n 1680, \n 1681, \n 1682, \n 1683, \n 1684, \n 1696, \n 1697, \n 1698, \n 1703, \n 1704, \n 1705, \n 1706, \n 1707, \n 1708, \n 1709, \n 1710, \n 1711, \n 1712, \n 1713, \n 1714, \n 1715, \n 1716, \n 1717, \n 1718, \n 1719, \n 1720, \n 1725, \n 1731, \n 1732, \n 1733, \n 1734, \n 1735, \n 1737, \n 1739, \n 1740, \n 1745, \n 1746, \n 1747, \n 1748, \n 1749, \n 1751, \n 1761, \n 1830, \n 1831, \n 1844, \n 1845, \n 1846, \n 1850, \n 1851, \n 1854, \n 1855, \n 1858, \n 1860, \n 1865, \n 1866, \n 1867, \n 1869, \n 1870, \n 1871, \n 1874, \n 1875, \n 1876, \n 1877, \n 1878, \n 1882, \n 1883, \n 1888, \n 1889, \n 1892, \n 1900, \n 1901, \n 1902, \n 1903, \n 1904, \n 1909, \n 2819, \n 2820, \n 2821, \n 2822, \n 2895, \n 2896, \n 2897, \n 2898, \n 2899, \n 2900, \n 2901, \n 2902, \n 2903, \n 2945, \n 2946, \n 2974, \n 2975, \n 2976, \n 2977, \n 2978, \n 2979, \n 2980, \n 2981, \n 2982, \n 2983, \n 2984, \n 2985, \n 2986, \n 2987, \n 2988, \n 2989, \n 2990, \n 2991, \n 2992, \n 2993, \n 2994, \n 2995, \n 2996, \n 3002, \n 3013\n ], \n \"leftLeg\": [\n 995, \n 998, \n 999, \n 1002, \n 1004, \n 1005, \n 1008, \n 1010, \n 1012, \n 1015, \n 1016, \n 1018, \n 1019, \n 1043, \n 1044, \n 1047, \n 1048, \n 1049, \n 1050, \n 1051, \n 1052, \n 1053, \n 1054, \n 1055, \n 1056, \n 1057, \n 1058, \n 1059, \n 1060, \n 1061, \n 1062, \n 1063, \n 1064, \n 1065, \n 1066, \n 1067, \n 1068, \n 1069, \n 1070, \n 1071, \n 1072, \n 1073, \n 1074, \n 1075, \n 1076, \n 1077, \n 1078, \n 1079, \n 1080, \n 1081, \n 1082, \n 1083, \n 1084, \n 1085, \n 1086, \n 1087, \n 1088, \n 1089, \n 1090, \n 1091, \n 1092, \n 1093, \n 1094, \n 1095, \n 1096, \n 1097, \n 1098, \n 1099, \n 1100, \n 1101, \n 1102, \n 1103, \n 1104, \n 1105, \n 1106, \n 1107, \n 1108, \n 1109, \n 1110, \n 1111, \n 1112, \n 1113, \n 1114, \n 1115, \n 1116, \n 1117, \n 1118, \n 1119, \n 1120, \n 1121, \n 1122, \n 1123, \n 1124, \n 1125, \n 1126, \n 1127, \n 1128, \n 1129, \n 1130, \n 1131, \n 1132, \n 1133, \n 1134, \n 1135, \n 1136, \n 1148, \n 1149, \n 1150, \n 1151, \n 1152, \n 1153, \n 1154, \n 1155, \n 1156, \n 1157, \n 1158, \n 1175, \n 1176, \n 1177, \n 1178, \n 1179, \n 1180, \n 1181, \n 1182, \n 1183, \n 1369, \n 1370, \n 1371, \n 1372, \n 1373, \n 1374, \n 1375, \n 1464, \n 1465, \n 1466, \n 1467, \n 1468, \n 1469, \n 1470, \n 1471, \n 1472, \n 1473, \n 1474, \n 1522, \n 1523, \n 1524, \n 1525, \n 1526, \n 1527, \n 1528, \n 1529, \n 1530, \n 1531, \n 1532, \n 3174, \n 3175, \n 3176, \n 3177, \n 3178, \n 3179, \n 3180, \n 3181, \n 3182, \n 3183, \n 3184, \n 3185, \n 3186, \n 3187, \n 3188, \n 3189, \n 3190, \n 3191, \n 3192, \n 3193, \n 3194, \n 3195, \n 3196, \n 3197, \n 3198, \n 3199, \n 3200, \n 3201, \n 3202, \n 3203, \n 3204, \n 3205, \n 3206, \n 3207, \n 3208, \n 3209, \n 3210, \n 3319, \n 3320, \n 3321, \n 3322, \n 3323, \n 3324, \n 3325, \n 3326, \n 3327, \n 3328, \n 3329, \n 3330, \n 3331, \n 3332, \n 3333, \n 3334, \n 3335, \n 3432, \n 3433, \n 3434, \n 3435, \n 3436, \n 3469, \n 3472, \n 3473, \n 3474\n ], \n \"leftToeBase\": [\n 3211, \n 3212, \n 3213, \n 3214, \n 3215, \n 3216, \n 3217, \n 3218, \n 3219, \n 3220, \n 3221, \n 3222, \n 3223, \n 3224, \n 3225, \n 3226, \n 3227, \n 3228, \n 3229, \n 3230, \n 3231, \n 3232, \n 3233, \n 3234, \n 3235, \n 3236, \n 3237, \n 3238, \n 3239, \n 3240, \n 3241, \n 3242, \n 3243, \n 3244, \n 3245, \n 3246, \n 3247, \n 3248, \n 3249, \n 3250, \n 3251, \n 3252, \n 3253, \n 3254, \n 3255, \n 3256, \n 3257, \n 3258, \n 3259, \n 3260, \n 3261, \n 3262, \n 3263, \n 3264, \n 3265, \n 3266, \n 3267, \n 3268, \n 3269, \n 3270, \n 3271, \n 3272, \n 3273, \n 3274, \n 3275, \n 3276, \n 3277, \n 3278, \n 3279, \n 3280, \n 3281, \n 3282, \n 3283, \n 3284, \n 3285, \n 3286, \n 3287, \n 3288, \n 3289, \n 3290, \n 3291, \n 3292, \n 3293, \n 3294, \n 3295, \n 3296, \n 3297, \n 3298, \n 3299, \n 3300, \n 3301, \n 3302, \n 3303, \n 3304, \n 3305, \n 3306, \n 3307, \n 3308, \n 3309, \n 3310, \n 3311, \n 3312, \n 3313, \n 3314, \n 3315, \n 3316, \n 3317, \n 3318, \n 3336, \n 3337, \n 3340, \n 3342, \n 3344, \n 3346, \n 3348, \n 3350, \n 3352, \n 3354, \n 3357, \n 3358, \n 3360, \n 3362\n ], \n \"leftFoot\": [\n 3327, \n 3328, \n 3329, \n 3330, \n 3331, \n 3332, \n 3333, \n 3334, \n 3335, \n 3336, \n 3337, \n 3338, \n 3339, \n 3340, \n 3341, \n 3342, \n 3343, \n 3344, \n 3345, \n 3346, \n 3347, \n 3348, \n 3349, \n 3350, \n 3351, \n 3352, \n 3353, \n 3354, \n 3355, \n 3356, \n 3357, \n 3358, \n 3359, \n 3360, \n 3361, \n 3362, \n 3363, \n 3364, \n 3365, \n 3366, \n 3367, \n 3368, \n 3369, \n 3370, \n 3371, \n 3372, \n 3373, \n 3374, \n 3375, \n 3376, \n 3377, \n 3378, \n 3379, \n 3380, \n 3381, \n 3382, \n 3383, \n 3384, \n 3385, \n 3386, \n 3387, \n 3388, \n 3389, \n 3390, \n 3391, \n 3392, \n 3393, \n 3394, \n 3395, \n 3396, \n 3397, \n 3398, \n 3399, \n 3400, \n 3401, \n 3402, \n 3403, \n 3404, \n 3405, \n 3406, \n 3407, \n 3408, \n 3409, \n 3410, \n 3411, \n 3412, \n 3413, \n 3414, \n 3415, \n 3416, \n 3417, \n 3418, \n 3419, \n 3420, \n 3421, \n 3422, \n 3423, \n 3424, \n 3425, \n 3426, \n 3427, \n 3428, \n 3429, \n 3430, \n 3431, \n 3432, \n 3433, \n 3434, \n 3435, \n 3436, \n 3437, \n 3438, \n 3439, \n 3440, \n 3441, \n 3442, \n 3443, \n 3444, \n 3445, \n 3446, \n 3447, \n 3448, \n 3449, \n 3450, \n 3451, \n 3452, \n 3453, \n 3454, \n 3455, \n 3456, \n 3457, \n 3458, \n 3459, \n 3460, \n 3461, \n 3462, \n 3463, \n 3464, \n 3465, \n 3466, \n 3467, \n 3468, \n 3469\n ], \n \"spine1\": [\n 598, \n 599, \n 600, \n 601, \n 610, \n 611, \n 612, \n 613, \n 614, \n 615, \n 616, \n 617, \n 618, \n 619, \n 620, \n 621, \n 642, \n 645, \n 646, \n 647, \n 652, \n 653, \n 658, \n 659, \n 660, \n 661, \n 668, \n 669, \n 670, \n 671, \n 684, \n 685, \n 686, \n 687, \n 688, \n 689, \n 690, \n 691, \n 692, \n 722, \n 723, \n 724, \n 725, \n 736, \n 750, \n 751, \n 761, \n 764, \n 766, \n 767, \n 794, \n 795, \n 891, \n 892, \n 893, \n 894, \n 925, \n 926, \n 927, \n 928, \n 929, \n 940, \n 941, \n 942, \n 943, \n 1190, \n 1191, \n 1192, \n 1193, \n 1194, \n 1195, \n 1196, \n 1197, \n 1200, \n 1201, \n 1202, \n 1212, \n 1236, \n 1252, \n 1253, \n 1254, \n 1255, \n 1268, \n 1269, \n 1270, \n 1329, \n 1330, \n 1348, \n 1349, \n 1351, \n 1420, \n 1421, \n 1423, \n 1424, \n 1425, \n 1426, \n 1436, \n 1437, \n 1756, \n 1757, \n 1758, \n 2839, \n 2840, \n 2841, \n 2842, \n 2843, \n 2844, \n 2845, \n 2846, \n 2847, \n 2848, \n 2849, \n 2850, \n 2851, \n 2870, \n 2871, \n 2883, \n 2906, \n 2908, \n 3014, \n 3017, \n 3025, \n 3030, \n 3033, \n 3034, \n 3037, \n 3039, \n 3040, \n 3041, \n 3042, \n 3043, \n 3044, \n 3076, \n 3077, \n 3079, \n 3480, \n 3505, \n 3511, \n 4086, \n 4087, \n 4088, \n 4089, \n 4098, \n 4099, \n 4100, \n 4101, \n 4102, \n 4103, \n 4104, \n 4105, \n 4106, \n 4107, \n 4108, \n 4109, \n 4130, \n 4131, \n 4134, \n 4135, \n 4140, \n 4141, \n 4146, \n 4147, \n 4148, \n 4149, \n 4156, \n 4157, \n 4158, \n 4159, \n 4172, \n 4173, \n 4174, \n 4175, \n 4176, \n 4177, \n 4178, \n 4179, \n 4180, \n 4210, \n 4211, \n 4212, \n 4213, \n 4225, \n 4239, \n 4240, \n 4249, \n 4250, \n 4255, \n 4256, \n 4282, \n 4283, \n 4377, \n 4378, \n 4379, \n 4380, \n 4411, \n 4412, \n 4413, \n 4414, \n 4415, \n 4426, \n 4427, \n 4428, \n 4429, \n 4676, \n 4677, \n 4678, \n 4679, \n 4680, \n 4681, \n 4682, \n 4683, \n 4686, \n 4687, \n 4688, \n 4695, \n 4719, \n 4735, \n 4736, \n 4737, \n 4740, \n 4751, \n 4752, \n 4753, \n 4824, \n 4825, \n 4828, \n 4893, \n 4894, \n 4895, \n 4897, \n 4898, \n 4899, \n 4908, \n 4909, \n 5223, \n 5224, \n 5225, \n 6300, \n 6301, \n 6302, \n 6303, \n 6304, \n 6305, \n 6306, \n 6307, \n 6308, \n 6309, \n 6310, \n 6311, \n 6312, \n 6331, \n 6332, \n 6342, \n 6366, \n 6367, \n 6475, \n 6477, \n 6478, \n 6481, \n 6482, \n 6485, \n 6487, \n 6488, \n 6489, \n 6490, \n 6491, \n 6878\n ], \n \"spine2\": [\n 570, \n 571, \n 572, \n 573, \n 584, \n 585, \n 586, \n 587, \n 588, \n 589, \n 590, \n 591, \n 592, \n 593, \n 594, \n 595, \n 596, \n 597, \n 602, \n 603, \n 604, \n 605, \n 606, \n 607, \n 608, \n 609, \n 622, \n 623, \n 624, \n 625, \n 638, \n 639, \n 640, \n 641, \n 643, \n 644, \n 648, \n 649, \n 650, \n 651, \n 666, \n 667, \n 672, \n 673, \n 674, \n 675, \n 680, \n 681, \n 682, \n 683, \n 693, \n 694, \n 695, \n 696, \n 697, \n 698, \n 699, \n 700, \n 701, \n 702, \n 703, \n 704, \n 713, \n 714, \n 715, \n 716, \n 717, \n 726, \n 727, \n 728, \n 729, \n 730, \n 731, \n 732, \n 733, \n 735, \n 737, \n 738, \n 739, \n 740, \n 741, \n 742, \n 743, \n 744, \n 745, \n 746, \n 747, \n 748, \n 749, \n 752, \n 753, \n 754, \n 755, \n 756, \n 757, \n 758, \n 759, \n 760, \n 762, \n 763, \n 803, \n 804, \n 805, \n 806, \n 811, \n 812, \n 813, \n 814, \n 817, \n 818, \n 819, \n 820, \n 821, \n 824, \n 825, \n 826, \n 827, \n 828, \n 895, \n 896, \n 930, \n 931, \n 1198, \n 1199, \n 1213, \n 1214, \n 1215, \n 1216, \n 1217, \n 1218, \n 1219, \n 1220, \n 1235, \n 1237, \n 1256, \n 1257, \n 1271, \n 1272, \n 1273, \n 1279, \n 1280, \n 1283, \n 1284, \n 1285, \n 1286, \n 1287, \n 1288, \n 1289, \n 1290, \n 1291, \n 1292, \n 1293, \n 1294, \n 1295, \n 1296, \n 1297, \n 1298, \n 1299, \n 1300, \n 1301, \n 1302, \n 1303, \n 1304, \n 1305, \n 1306, \n 1307, \n 1308, \n 1309, \n 1312, \n 1313, \n 1319, \n 1320, \n 1346, \n 1347, \n 1350, \n 1352, \n 1401, \n 1417, \n 1418, \n 1419, \n 1422, \n 1427, \n 1434, \n 1435, \n 1503, \n 1504, \n 1536, \n 1537, \n 1544, \n 1545, \n 1753, \n 1754, \n 1755, \n 1759, \n 1760, \n 1761, \n 1762, \n 1763, \n 1808, \n 1809, \n 1810, \n 1811, \n 1816, \n 1817, \n 1818, \n 1819, \n 1820, \n 1834, \n 1835, \n 1836, \n 1837, \n 1838, \n 1839, \n 1868, \n 1879, \n 1880, \n 2812, \n 2813, \n 2852, \n 2853, \n 2854, \n 2855, \n 2856, \n 2857, \n 2858, \n 2859, \n 2860, \n 2861, \n 2862, \n 2863, \n 2864, \n 2865, \n 2866, \n 2867, \n 2868, \n 2869, \n 2872, \n 2875, \n 2876, \n 2877, \n 2878, \n 2881, \n 2882, \n 2884, \n 2885, \n 2886, \n 2904, \n 2905, \n 2907, \n 2931, \n 2932, \n 2933, \n 2934, \n 2935, \n 2936, \n 2937, \n 2941, \n 2950, \n 2951, \n 2952, \n 2953, \n 2954, \n 2955, \n 2956, \n 2957, \n 2958, \n 2959, \n 2960, \n 2961, \n 2962, \n 2963, \n 2964, \n 2965, \n 2966, \n 2967, \n 2968, \n 2969, \n 2970, \n 2971, \n 2972, \n 2973, \n 2997, \n 2998, \n 3006, \n 3007, \n 3012, \n 3015, \n 3026, \n 3027, \n 3028, \n 3029, \n 3031, \n 3032, \n 3035, \n 3036, \n 3038, \n 3059, \n 3060, \n 3061, \n 3062, \n 3063, \n 3064, \n 3065, \n 3066, \n 3067, \n 3073, \n 3074, \n 3075, \n 3078, \n 3168, \n 3169, \n 3171, \n 3470, \n 3471, \n 3482, \n 3483, \n 3495, \n 3496, \n 3497, \n 3498, \n 3506, \n 3508, \n 4058, \n 4059, \n 4060, \n 4061, \n 4072, \n 4073, \n 4074, \n 4075, \n 4076, \n 4077, \n 4078, \n 4079, \n 4080, \n 4081, \n 4082, \n 4083, \n 4084, \n 4085, \n 4090, \n 4091, \n 4092, \n 4093, \n 4094, \n 4095, \n 4096, \n 4097, \n 4110, \n 4111, \n 4112, \n 4113, \n 4126, \n 4127, \n 4128, \n 4129, \n 4132, \n 4133, \n 4136, \n 4137, \n 4138, \n 4139, \n 4154, \n 4155, \n 4160, \n 4161, \n 4162, \n 4163, \n 4168, \n 4169, \n 4170, \n 4171, \n 4181, \n 4182, \n 4183, \n 4184, \n 4185, \n 4186, \n 4187, \n 4188, \n 4189, \n 4190, \n 4191, \n 4192, \n 4201, \n 4202, \n 4203, \n 4204, \n 4207, \n 4214, \n 4215, \n 4216, \n 4217, \n 4218, \n 4219, \n 4220, \n 4221, \n 4223, \n 4224, \n 4226, \n 4227, \n 4228, \n 4229, \n 4230, \n 4231, \n 4232, \n 4233, \n 4234, \n 4235, \n 4236, \n 4237, \n 4238, \n 4241, \n 4242, \n 4243, \n 4244, \n 4245, \n 4246, \n 4247, \n 4248, \n 4251, \n 4252, \n 4291, \n 4292, \n 4293, \n 4294, \n 4299, \n 4300, \n 4301, \n 4302, \n 4305, \n 4306, \n 4307, \n 4308, \n 4309, \n 4312, \n 4313, \n 4314, \n 4315, \n 4381, \n 4382, \n 4416, \n 4417, \n 4684, \n 4685, \n 4696, \n 4697, \n 4698, \n 4699, \n 4700, \n 4701, \n 4702, \n 4703, \n 4718, \n 4720, \n 4738, \n 4739, \n 4754, \n 4755, \n 4756, \n 4761, \n 4762, \n 4765, \n 4766, \n 4767, \n 4768, \n 4769, \n 4770, \n 4771, \n 4772, \n 4773, \n 4774, \n 4775, \n 4776, \n 4777, \n 4778, \n 4779, \n 4780, \n 4781, \n 4782, \n 4783, \n 4784, \n 4785, \n 4786, \n 4787, \n 4788, \n 4789, \n 4792, \n 4793, \n 4799, \n 4800, \n 4822, \n 4823, \n 4826, \n 4827, \n 4874, \n 4890, \n 4891, \n 4892, \n 4896, \n 4900, \n 4907, \n 4910, \n 4975, \n 4976, \n 5007, \n 5008, \n 5013, \n 5014, \n 5222, \n 5226, \n 5227, \n 5228, \n 5229, \n 5230, \n 5269, \n 5270, \n 5271, \n 5272, \n 5277, \n 5278, \n 5279, \n 5280, \n 5281, \n 5295, \n 5296, \n 5297, \n 5298, \n 5299, \n 5300, \n 5329, \n 5340, \n 5341, \n 6273, \n 6274, \n 6313, \n 6314, \n 6315, \n 6316, \n 6317, \n 6318, \n 6319, \n 6320, \n 6321, \n 6322, \n 6323, \n 6324, \n 6325, \n 6326, \n 6327, \n 6328, \n 6329, \n 6330, \n 6333, \n 6336, \n 6337, \n 6340, \n 6341, \n 6343, \n 6344, \n 6345, \n 6363, \n 6364, \n 6365, \n 6390, \n 6391, \n 6392, \n 6393, \n 6394, \n 6395, \n 6396, \n 6398, \n 6409, \n 6410, \n 6411, \n 6412, \n 6413, \n 6414, \n 6415, \n 6416, \n 6417, \n 6418, \n 6419, \n 6420, \n 6421, \n 6422, \n 6423, \n 6424, \n 6425, \n 6426, \n 6427, \n 6428, \n 6429, \n 6430, \n 6431, \n 6432, \n 6456, \n 6457, \n 6465, \n 6466, \n 6476, \n 6479, \n 6480, \n 6483, \n 6484, \n 6486, \n 6496, \n 6497, \n 6498, \n 6499, \n 6500, \n 6501, \n 6502, \n 6503, \n 6879\n ], \n \"leftShoulder\": [\n 591, \n 604, \n 605, \n 606, \n 609, \n 634, \n 635, \n 636, \n 637, \n 674, \n 706, \n 707, \n 708, \n 709, \n 710, \n 711, \n 712, \n 713, \n 715, \n 717, \n 730, \n 733, \n 734, \n 735, \n 781, \n 782, \n 783, \n 1238, \n 1239, \n 1240, \n 1241, \n 1242, \n 1243, \n 1244, \n 1245, \n 1290, \n 1291, \n 1294, \n 1316, \n 1317, \n 1318, \n 1401, \n 1402, \n 1403, \n 1404, \n 1509, \n 1535, \n 1545, \n 1808, \n 1810, \n 1811, \n 1812, \n 1813, \n 1814, \n 1815, \n 1818, \n 1819, \n 1821, \n 1822, \n 1823, \n 1824, \n 1825, \n 1826, \n 1827, \n 1828, \n 1829, \n 1830, \n 1831, \n 1832, \n 1833, \n 1837, \n 1840, \n 1841, \n 1842, \n 1843, \n 1844, \n 1845, \n 1846, \n 1847, \n 1848, \n 1849, \n 1850, \n 1851, \n 1852, \n 1853, \n 1854, \n 1855, \n 1856, \n 1857, \n 1858, \n 1859, \n 1861, \n 1862, \n 1863, \n 1864, \n 1872, \n 1873, \n 1880, \n 1881, \n 1884, \n 1885, \n 1886, \n 1887, \n 1890, \n 1891, \n 1893, \n 1894, \n 1895, \n 1896, \n 1897, \n 1898, \n 1899, \n 2879, \n 2880, \n 2881, \n 2886, \n 2887, \n 2888, \n 2889, \n 2890, \n 2891, \n 2892, \n 2893, \n 2894, \n 2903, \n 2938, \n 2939, \n 2940, \n 2941, \n 2942, \n 2943, \n 2944, \n 2945, \n 2946, \n 2947, \n 2948, \n 2949, \n 2965, \n 2967, \n 2969, \n 2999, \n 3000, \n 3001, \n 3002, \n 3003, \n 3004, \n 3005, \n 3008, \n 3009, \n 3010, \n 3011\n ], \n \"rightShoulder\": [\n 4077, \n 4091, \n 4092, \n 4094, \n 4095, \n 4122, \n 4123, \n 4124, \n 4125, \n 4162, \n 4194, \n 4195, \n 4196, \n 4197, \n 4198, \n 4199, \n 4200, \n 4201, \n 4203, \n 4207, \n 4218, \n 4219, \n 4222, \n 4223, \n 4269, \n 4270, \n 4271, \n 4721, \n 4722, \n 4723, \n 4724, \n 4725, \n 4726, \n 4727, \n 4728, \n 4773, \n 4774, \n 4778, \n 4796, \n 4797, \n 4798, \n 4874, \n 4875, \n 4876, \n 4877, \n 4982, \n 5006, \n 5014, \n 5269, \n 5271, \n 5272, \n 5273, \n 5274, \n 5275, \n 5276, \n 5279, \n 5281, \n 5282, \n 5283, \n 5284, \n 5285, \n 5286, \n 5287, \n 5288, \n 5289, \n 5290, \n 5291, \n 5292, \n 5293, \n 5294, \n 5298, \n 5301, \n 5302, \n 5303, \n 5304, \n 5305, \n 5306, \n 5307, \n 5308, \n 5309, \n 5310, \n 5311, \n 5312, \n 5313, \n 5314, \n 5315, \n 5316, \n 5317, \n 5318, \n 5319, \n 5320, \n 5322, \n 5323, \n 5324, \n 5325, \n 5333, \n 5334, \n 5341, \n 5342, \n 5345, \n 5346, \n 5347, \n 5348, \n 5351, \n 5352, \n 5354, \n 5355, \n 5356, \n 5357, \n 5358, \n 5359, \n 5360, \n 6338, \n 6339, \n 6340, \n 6345, \n 6346, \n 6347, \n 6348, \n 6349, \n 6350, \n 6351, \n 6352, \n 6353, \n 6362, \n 6397, \n 6398, \n 6399, \n 6400, \n 6401, \n 6402, \n 6403, \n 6404, \n 6405, \n 6406, \n 6407, \n 6408, \n 6424, \n 6425, \n 6428, \n 6458, \n 6459, \n 6460, \n 6461, \n 6462, \n 6463, \n 6464, \n 6467, \n 6468, \n 6469, \n 6470\n ], \n \"rightFoot\": [\n 6727, \n 6728, \n 6729, \n 6730, \n 6731, \n 6732, \n 6733, \n 6734, \n 6735, \n 6736, \n 6737, \n 6738, \n 6739, \n 6740, \n 6741, \n 6742, \n 6743, \n 6744, \n 6745, \n 6746, \n 6747, \n 6748, \n 6749, \n 6750, \n 6751, \n 6752, \n 6753, \n 6754, \n 6755, \n 6756, \n 6757, \n 6758, \n 6759, \n 6760, \n 6761, \n 6762, \n 6763, \n 6764, \n 6765, \n 6766, \n 6767, \n 6768, \n 6769, \n 6770, \n 6771, \n 6772, \n 6773, \n 6774, \n 6775, \n 6776, \n 6777, \n 6778, \n 6779, \n 6780, \n 6781, \n 6782, \n 6783, \n 6784, \n 6785, \n 6786, \n 6787, \n 6788, \n 6789, \n 6790, \n 6791, \n 6792, \n 6793, \n 6794, \n 6795, \n 6796, \n 6797, \n 6798, \n 6799, \n 6800, \n 6801, \n 6802, \n 6803, \n 6804, \n 6805, \n 6806, \n 6807, \n 6808, \n 6809, \n 6810, \n 6811, \n 6812, \n 6813, \n 6814, \n 6815, \n 6816, \n 6817, \n 6818, \n 6819, \n 6820, \n 6821, \n 6822, \n 6823, \n 6824, \n 6825, \n 6826, \n 6827, \n 6828, \n 6829, \n 6830, \n 6831, \n 6832, \n 6833, \n 6834, \n 6835, \n 6836, \n 6837, \n 6838, \n 6839, \n 6840, \n 6841, \n 6842, \n 6843, \n 6844, \n 6845, \n 6846, \n 6847, \n 6848, \n 6849, \n 6850, \n 6851, \n 6852, \n 6853, \n 6854, \n 6855, \n 6856, \n 6857, \n 6858, \n 6859, \n 6860, \n 6861, \n 6862, \n 6863, \n 6864, \n 6865, \n 6866, \n 6867, \n 6868, \n 6869\n ], \n \"head\": [\n 0, \n 1, \n 2, \n 3, \n 4, \n 5, \n 6, \n 7, \n 8, \n 9, \n 10, \n 11, \n 12, \n 13, \n 14, \n 15, \n 16, \n 17, \n 18, \n 19, \n 20, \n 21, \n 22, \n 23, \n 24, \n 25, \n 26, \n 27, \n 28, \n 29, \n 30, \n 31, \n 32, \n 33, \n 34, \n 35, \n 36, \n 37, \n 38, \n 39, \n 40, \n 41, \n 42, \n 43, \n 44, \n 45, \n 46, \n 47, \n 48, \n 49, \n 50, \n 51, \n 52, \n 53, \n 54, \n 55, \n 56, \n 57, \n 58, \n 59, \n 60, \n 61, \n 62, \n 63, \n 64, \n 65, \n 66, \n 67, \n 68, \n 69, \n 70, \n 71, \n 72, \n 73, \n 74, \n 75, \n 76, \n 77, \n 78, \n 79, \n 80, \n 81, \n 82, \n 83, \n 84, \n 85, \n 86, \n 87, \n 88, \n 89, \n 90, \n 91, \n 92, \n 93, \n 94, \n 95, \n 96, \n 97, \n 98, \n 99, \n 100, \n 101, \n 102, \n 103, \n 104, \n 105, \n 106, \n 107, \n 108, \n 109, \n 110, \n 111, \n 112, \n 113, \n 114, \n 115, \n 116, \n 117, \n 118, \n 119, \n 120, \n 121, \n 122, \n 123, \n 124, \n 125, \n 126, \n 127, \n 128, \n 129, \n 130, \n 131, \n 132, \n 133, \n 134, \n 135, \n 136, \n 137, \n 138, \n 139, \n 140, \n 141, \n 142, \n 143, \n 144, \n 145, \n 146, \n 147, \n 148, \n 149, \n 154, \n 155, \n 156, \n 157, \n 158, \n 159, \n 160, \n 161, \n 162, \n 163, \n 164, \n 165, \n 166, \n 167, \n 168, \n 169, \n 170, \n 171, \n 172, \n 173, \n 176, \n 177, \n 178, \n 179, \n 180, \n 181, \n 182, \n 183, \n 184, \n 185, \n 186, \n 187, \n 188, \n 189, \n 190, \n 191, \n 192, \n 193, \n 194, \n 195, \n 196, \n 197, \n 198, \n 199, \n 200, \n 201, \n 202, \n 203, \n 204, \n 205, \n 220, \n 221, \n 225, \n 226, \n 227, \n 228, \n 229, \n 230, \n 231, \n 232, \n 233, \n 234, \n 235, \n 236, \n 237, \n 238, \n 239, \n 240, \n 241, \n 242, \n 243, \n 244, \n 245, \n 246, \n 247, \n 248, \n 249, \n 250, \n 251, \n 252, \n 253, \n 254, \n 255, \n 258, \n 259, \n 260, \n 261, \n 262, \n 263, \n 264, \n 265, \n 266, \n 267, \n 268, \n 269, \n 270, \n 271, \n 272, \n 273, \n 274, \n 275, \n 276, \n 277, \n 278, \n 279, \n 280, \n 281, \n 282, \n 283, \n 286, \n 287, \n 288, \n 289, \n 290, \n 291, \n 292, \n 293, \n 294, \n 295, \n 303, \n 304, \n 306, \n 307, \n 310, \n 311, \n 312, \n 313, \n 314, \n 315, \n 316, \n 317, \n 318, \n 319, \n 320, \n 321, \n 322, \n 323, \n 324, \n 325, \n 326, \n 327, \n 328, \n 329, \n 330, \n 331, \n 332, \n 335, \n 336, \n 337, \n 338, \n 339, \n 340, \n 341, \n 342, \n 343, \n 344, \n 345, \n 346, \n 347, \n 348, \n 349, \n 350, \n 351, \n 352, \n 353, \n 354, \n 355, \n 356, \n 357, \n 358, \n 359, \n 360, \n 361, \n 362, \n 363, \n 364, \n 365, \n 366, \n 367, \n 368, \n 369, \n 370, \n 371, \n 372, \n 373, \n 374, \n 375, \n 376, \n 377, \n 378, \n 379, \n 380, \n 381, \n 382, \n 383, \n 384, \n 385, \n 386, \n 387, \n 388, \n 389, \n 390, \n 391, \n 392, \n 393, \n 394, \n 395, \n 396, \n 397, \n 398, \n 399, \n 400, \n 401, \n 402, \n 403, \n 404, \n 405, \n 406, \n 407, \n 408, \n 409, \n 410, \n 411, \n 412, \n 413, \n 414, \n 415, \n 416, \n 417, \n 418, \n 419, \n 420, \n 421, \n 422, \n 427, \n 428, \n 429, \n 430, \n 431, \n 432, \n 433, \n 434, \n 435, \n 436, \n 437, \n 438, \n 439, \n 442, \n 443, \n 444, \n 445, \n 446, \n 447, \n 448, \n 449, \n 450, \n 454, \n 455, \n 456, \n 457, \n 458, \n 459, \n 461, \n 462, \n 463, \n 464, \n 465, \n 466, \n 467, \n 468, \n 469, \n 470, \n 471, \n 472, \n 473, \n 474, \n 475, \n 476, \n 477, \n 478, \n 479, \n 480, \n 481, \n 482, \n 483, \n 484, \n 485, \n 486, \n 487, \n 488, \n 489, \n 490, \n 491, \n 492, \n 493, \n 494, \n 495, \n 496, \n 497, \n 498, \n 499, \n 500, \n 501, \n 502, \n 503, \n 504, \n 505, \n 506, \n 507, \n 508, \n 509, \n 510, \n 511, \n 512, \n 513, \n 514, \n 515, \n 516, \n 517, \n 518, \n 519, \n 520, \n 521, \n 522, \n 523, \n 524, \n 525, \n 526, \n 527, \n 528, \n 529, \n 530, \n 531, \n 532, \n 533, \n 534, \n 535, \n 536, \n 537, \n 538, \n 539, \n 540, \n 541, \n 542, \n 543, \n 544, \n 545, \n 546, \n 547, \n 548, \n 549, \n 550, \n 551, \n 552, \n 553, \n 554, \n 555, \n 556, \n 557, \n 558, \n 559, \n 560, \n 561, \n 562, \n 563, \n 564, \n 565, \n 566, \n 567, \n 568, \n 569, \n 574, \n 575, \n 576, \n 577, \n 578, \n 579, \n 580, \n 581, \n 582, \n 583, \n 1764, \n 1765, \n 1766, \n 1770, \n 1771, \n 1772, \n 1773, \n 1774, \n 1775, \n 1776, \n 1777, \n 1778, \n 1905, \n 1906, \n 1907, \n 1908, \n 2779, \n 2780, \n 2781, \n 2782, \n 2783, \n 2784, \n 2785, \n 2786, \n 2787, \n 2788, \n 2789, \n 2790, \n 2791, \n 2792, \n 2793, \n 2794, \n 2795, \n 2796, \n 2797, \n 2798, \n 2799, \n 2800, \n 2801, \n 2802, \n 2803, \n 2804, \n 2805, \n 2806, \n 2807, \n 2808, \n 2809, \n 2810, \n 2811, \n 2814, \n 2815, \n 2816, \n 2817, \n 2818, \n 3045, \n 3046, \n 3047, \n 3048, \n 3051, \n 3052, \n 3053, \n 3054, \n 3055, \n 3056, \n 3058, \n 3069, \n 3070, \n 3071, \n 3072, \n 3161, \n 3162, \n 3163, \n 3165, \n 3166, \n 3167, \n 3485, \n 3486, \n 3487, \n 3488, \n 3489, \n 3490, \n 3491, \n 3492, \n 3493, \n 3494, \n 3499, \n 3512, \n 3513, \n 3514, \n 3515, \n 3516, \n 3517, \n 3518, \n 3519, \n 3520, \n 3521, \n 3522, \n 3523, \n 3524, \n 3525, \n 3526, \n 3527, \n 3528, \n 3529, \n 3530, \n 3531, \n 3532, \n 3533, \n 3534, \n 3535, \n 3536, \n 3537, \n 3538, \n 3539, \n 3540, \n 3541, \n 3542, \n 3543, \n 3544, \n 3545, \n 3546, \n 3547, \n 3548, \n 3549, \n 3550, \n 3551, \n 3552, \n 3553, \n 3554, \n 3555, \n 3556, \n 3557, \n 3558, \n 3559, \n 3560, \n 3561, \n 3562, \n 3563, \n 3564, \n 3565, \n 3566, \n 3567, \n 3568, \n 3569, \n 3570, \n 3571, \n 3572, \n 3573, \n 3574, \n 3575, \n 3576, \n 3577, \n 3578, \n 3579, \n 3580, \n 3581, \n 3582, \n 3583, \n 3584, \n 3585, \n 3586, \n 3587, \n 3588, \n 3589, \n 3590, \n 3591, \n 3592, \n 3593, \n 3594, \n 3595, \n 3596, \n 3597, \n 3598, \n 3599, \n 3600, \n 3601, \n 3602, \n 3603, \n 3604, \n 3605, \n 3606, \n 3607, \n 3608, \n 3609, \n 3610, \n 3611, \n 3612, \n 3613, \n 3614, \n 3615, \n 3616, \n 3617, \n 3618, \n 3619, \n 3620, \n 3621, \n 3622, \n 3623, \n 3624, \n 3625, \n 3626, \n 3627, \n 3628, \n 3629, \n 3630, \n 3631, \n 3632, \n 3633, \n 3634, \n 3635, \n 3636, \n 3637, \n 3638, \n 3639, \n 3640, \n 3641, \n 3642, \n 3643, \n 3644, \n 3645, \n 3646, \n 3647, \n 3648, \n 3649, \n 3650, \n 3651, \n 3652, \n 3653, \n 3654, \n 3655, \n 3656, \n 3657, \n 3658, \n 3659, \n 3660, \n 3661, \n 3666, \n 3667, \n 3668, \n 3669, \n 3670, \n 3671, \n 3672, \n 3673, \n 3674, \n 3675, \n 3676, \n 3677, \n 3678, \n 3679, \n 3680, \n 3681, \n 3682, \n 3683, \n 3684, \n 3685, \n 3688, \n 3689, \n 3690, \n 3691, \n 3692, \n 3693, \n 3694, \n 3695, \n 3696, \n 3697, \n 3698, \n 3699, \n 3700, \n 3701, \n 3702, \n 3703, \n 3704, \n 3705, \n 3706, \n 3707, \n 3708, \n 3709, \n 3710, \n 3711, \n 3712, \n 3713, \n 3714, \n 3715, \n 3716, \n 3717, \n 3732, \n 3733, \n 3737, \n 3738, \n 3739, \n 3740, \n 3741, \n 3742, \n 3743, \n 3744, \n 3745, \n 3746, \n 3747, \n 3748, \n 3749, \n 3750, \n 3751, \n 3752, \n 3753, \n 3754, \n 3755, \n 3756, \n 3757, \n 3758, \n 3759, \n 3760, \n 3761, \n 3762, \n 3763, \n 3764, \n 3765, \n 3766, \n 3767, \n 3770, \n 3771, \n 3772, \n 3773, \n 3774, \n 3775, \n 3776, \n 3777, \n 3778, \n 3779, \n 3780, \n 3781, \n 3782, \n 3783, \n 3784, \n 3785, \n 3786, \n 3787, \n 3788, \n 3789, \n 3790, \n 3791, \n 3792, \n 3793, \n 3794, \n 3795, \n 3798, \n 3799, \n 3800, \n 3801, \n 3802, \n 3803, \n 3804, \n 3805, \n 3806, \n 3807, \n 3815, \n 3816, \n 3819, \n 3820, \n 3821, \n 3822, \n 3823, \n 3824, \n 3825, \n 3826, \n 3827, \n 3828, \n 3829, \n 3830, \n 3831, \n 3832, \n 3833, \n 3834, \n 3835, \n 3836, \n 3837, \n 3838, \n 3841, \n 3842, \n 3843, \n 3844, \n 3845, \n 3846, \n 3847, \n 3848, \n 3849, \n 3850, \n 3851, \n 3852, \n 3853, \n 3854, \n 3855, \n 3856, \n 3857, \n 3858, \n 3859, \n 3860, \n 3861, \n 3862, \n 3863, \n 3864, \n 3865, \n 3866, \n 3867, \n 3868, \n 3869, \n 3870, \n 3871, \n 3872, \n 3873, \n 3874, \n 3875, \n 3876, \n 3877, \n 3878, \n 3879, \n 3880, \n 3881, \n 3882, \n 3883, \n 3884, \n 3885, \n 3886, \n 3887, \n 3888, \n 3889, \n 3890, \n 3891, \n 3892, \n 3893, \n 3894, \n 3895, \n 3896, \n 3897, \n 3898, \n 3899, \n 3900, \n 3901, \n 3902, \n 3903, \n 3904, \n 3905, \n 3906, \n 3907, \n 3908, \n 3909, \n 3910, \n 3911, \n 3912, \n 3913, \n 3914, \n 3915, \n 3916, \n 3917, \n 3922, \n 3923, \n 3924, \n 3925, \n 3926, \n 3927, \n 3928, \n 3929, \n 3930, \n 3931, \n 3932, \n 3933, \n 3936, \n 3937, \n 3938, \n 3939, \n 3940, \n 3941, \n 3945, \n 3946, \n 3947, \n 3948, \n 3949, \n 3950, \n 3951, \n 3952, \n 3953, \n 3954, \n 3955, \n 3956, \n 3957, \n 3958, \n 3959, \n 3960, \n 3961, \n 3962, \n 3963, \n 3964, \n 3965, \n 3966, \n 3967, \n 3968, \n 3969, \n 3970, \n 3971, \n 3972, \n 3973, \n 3974, \n 3975, \n 3976, \n 3977, \n 3978, \n 3979, \n 3980, \n 3981, \n 3982, \n 3983, \n 3984, \n 3985, \n 3986, \n 3987, \n 3988, \n 3989, \n 3990, \n 3991, \n 3992, \n 3993, \n 3994, \n 3995, \n 3996, \n 3997, \n 3998, \n 3999, \n 4000, \n 4001, \n 4002, \n 4003, \n 4004, \n 4005, \n 4006, \n 4007, \n 4008, \n 4009, \n 4010, \n 4011, \n 4012, \n 4013, \n 4014, \n 4015, \n 4016, \n 4017, \n 4018, \n 4019, \n 4020, \n 4021, \n 4022, \n 4023, \n 4024, \n 4025, \n 4026, \n 4027, \n 4028, \n 4029, \n 4030, \n 4031, \n 4032, \n 4033, \n 4034, \n 4035, \n 4036, \n 4037, \n 4038, \n 4039, \n 4040, \n 4041, \n 4042, \n 4043, \n 4044, \n 4045, \n 4046, \n 4047, \n 4048, \n 4049, \n 4050, \n 4051, \n 4052, \n 4053, \n 4054, \n 4055, \n 4056, \n 4057, \n 4062, \n 4063, \n 4064, \n 4065, \n 4066, \n 4067, \n 4068, \n 4069, \n 4070, \n 4071, \n 5231, \n 5232, \n 5233, \n 5235, \n 5236, \n 5237, \n 5238, \n 5239, \n 5240, \n 5241, \n 5242, \n 5243, \n 5366, \n 5367, \n 5368, \n 5369, \n 6240, \n 6241, \n 6242, \n 6243, \n 6244, \n 6245, \n 6246, \n 6247, \n 6248, \n 6249, \n 6250, \n 6251, \n 6252, \n 6253, \n 6254, \n 6255, \n 6256, \n 6257, \n 6258, \n 6259, \n 6260, \n 6261, \n 6262, \n 6263, \n 6264, \n 6265, \n 6266, \n 6267, \n 6268, \n 6269, \n 6270, \n 6271, \n 6272, \n 6275, \n 6276, \n 6277, \n 6278, \n 6279, \n 6492, \n 6493, \n 6494, \n 6495, \n 6880, \n 6881, \n 6882, \n 6883, \n 6884, \n 6885, \n 6886, \n 6887, \n 6888, \n 6889\n ], \n \"rightArm\": [\n 4114, \n 4115, \n 4116, \n 4117, \n 4122, \n 4125, \n 4168, \n 4171, \n 4204, \n 4205, \n 4206, \n 4207, \n 4257, \n 4258, \n 4259, \n 4260, \n 4261, \n 4262, \n 4263, \n 4264, \n 4265, \n 4266, \n 4267, \n 4268, \n 4272, \n 4273, \n 4274, \n 4275, \n 4276, \n 4277, \n 4278, \n 4279, \n 4280, \n 4281, \n 4714, \n 4715, \n 4716, \n 4717, \n 4741, \n 4742, \n 4743, \n 4744, \n 4756, \n 4763, \n 4764, \n 4790, \n 4791, \n 4794, \n 4795, \n 4816, \n 4817, \n 4818, \n 4819, \n 4830, \n 4831, \n 4832, \n 4833, \n 4849, \n 4850, \n 4851, \n 4852, \n 4853, \n 4854, \n 4855, \n 4856, \n 4857, \n 4858, \n 4859, \n 4860, \n 4861, \n 4862, \n 4863, \n 4864, \n 4865, \n 4866, \n 4867, \n 4868, \n 4869, \n 4870, \n 4871, \n 4872, \n 4873, \n 4876, \n 4877, \n 4878, \n 4879, \n 4880, \n 4881, \n 4882, \n 4883, \n 4884, \n 4885, \n 4886, \n 4887, \n 4888, \n 4889, \n 4901, \n 4902, \n 4903, \n 4904, \n 4905, \n 4906, \n 4911, \n 4912, \n 4913, \n 4914, \n 4915, \n 4916, \n 4917, \n 4918, \n 4974, \n 4977, \n 4978, \n 4979, \n 4980, \n 4981, \n 4982, \n 5009, \n 5010, \n 5011, \n 5012, \n 5014, \n 5088, \n 5089, \n 5090, \n 5091, \n 5100, \n 5101, \n 5102, \n 5103, \n 5104, \n 5105, \n 5106, \n 5107, \n 5108, \n 5109, \n 5110, \n 5111, \n 5114, \n 5115, \n 5116, \n 5117, \n 5118, \n 5119, \n 5120, \n 5121, \n 5122, \n 5123, \n 5124, \n 5125, \n 5128, \n 5129, \n 5130, \n 5131, \n 5134, \n 5135, \n 5136, \n 5137, \n 5138, \n 5139, \n 5140, \n 5141, \n 5142, \n 5143, \n 5144, \n 5145, \n 5146, \n 5147, \n 5148, \n 5149, \n 5150, \n 5151, \n 5152, \n 5153, \n 5165, \n 5166, \n 5167, \n 5172, \n 5173, \n 5174, \n 5175, \n 5176, \n 5177, \n 5178, \n 5179, \n 5180, \n 5181, \n 5182, \n 5183, \n 5184, \n 5185, \n 5186, \n 5187, \n 5188, \n 5189, \n 5194, \n 5200, \n 5201, \n 5202, \n 5203, \n 5204, \n 5206, \n 5208, \n 5209, \n 5214, \n 5215, \n 5216, \n 5217, \n 5218, \n 5220, \n 5229, \n 5292, \n 5293, \n 5303, \n 5306, \n 5309, \n 5311, \n 5314, \n 5315, \n 5318, \n 5319, \n 5321, \n 5326, \n 5327, \n 5328, \n 5330, \n 5331, \n 5332, \n 5335, \n 5336, \n 5337, \n 5338, \n 5339, \n 5343, \n 5344, \n 5349, \n 5350, \n 5353, \n 5361, \n 5362, \n 5363, \n 5364, \n 5365, \n 5370, \n 6280, \n 6281, \n 6282, \n 6283, \n 6354, \n 6355, \n 6356, \n 6357, \n 6358, \n 6359, \n 6360, \n 6361, \n 6362, \n 6404, \n 6405, \n 6433, \n 6434, \n 6435, \n 6436, \n 6437, \n 6438, \n 6439, \n 6440, \n 6441, \n 6442, \n 6443, \n 6444, \n 6445, \n 6446, \n 6447, \n 6448, \n 6449, \n 6450, \n 6451, \n 6452, \n 6453, \n 6454, \n 6455, \n 6461, \n 6471\n ], \n \"leftHandIndex1\": [\n 2027, \n 2028, \n 2029, \n 2030, \n 2037, \n 2038, \n 2039, \n 2040, \n 2057, \n 2067, \n 2068, \n 2123, \n 2124, \n 2125, \n 2126, \n 2127, \n 2128, \n 2129, \n 2130, \n 2132, \n 2145, \n 2146, \n 2152, \n 2153, \n 2154, \n 2156, \n 2157, \n 2158, \n 2159, \n 2160, \n 2161, \n 2162, \n 2163, \n 2164, \n 2165, \n 2166, \n 2167, \n 2168, \n 2169, \n 2177, \n 2178, \n 2179, \n 2181, \n 2186, \n 2187, \n 2190, \n 2191, \n 2204, \n 2205, \n 2215, \n 2216, \n 2217, \n 2218, \n 2219, \n 2220, \n 2232, \n 2233, \n 2245, \n 2246, \n 2247, \n 2258, \n 2259, \n 2261, \n 2262, \n 2263, \n 2269, \n 2270, \n 2272, \n 2273, \n 2274, \n 2276, \n 2277, \n 2280, \n 2281, \n 2282, \n 2283, \n 2291, \n 2292, \n 2293, \n 2294, \n 2295, \n 2296, \n 2297, \n 2298, \n 2299, \n 2300, \n 2301, \n 2302, \n 2303, \n 2304, \n 2305, \n 2306, \n 2307, \n 2308, \n 2309, \n 2310, \n 2311, \n 2312, \n 2313, \n 2314, \n 2315, \n 2316, \n 2317, \n 2318, \n 2319, \n 2320, \n 2321, \n 2322, \n 2323, \n 2324, \n 2325, \n 2326, \n 2327, \n 2328, \n 2329, \n 2330, \n 2331, \n 2332, \n 2333, \n 2334, \n 2335, \n 2336, \n 2337, \n 2338, \n 2339, \n 2340, \n 2341, \n 2342, \n 2343, \n 2344, \n 2345, \n 2346, \n 2347, \n 2348, \n 2349, \n 2350, \n 2351, \n 2352, \n 2353, \n 2354, \n 2355, \n 2356, \n 2357, \n 2358, \n 2359, \n 2360, \n 2361, \n 2362, \n 2363, \n 2364, \n 2365, \n 2366, \n 2367, \n 2368, \n 2369, \n 2370, \n 2371, \n 2372, \n 2373, \n 2374, \n 2375, \n 2376, \n 2377, \n 2378, \n 2379, \n 2380, \n 2381, \n 2382, \n 2383, \n 2384, \n 2385, \n 2386, \n 2387, \n 2388, \n 2389, \n 2390, \n 2391, \n 2392, \n 2393, \n 2394, \n 2395, \n 2396, \n 2397, \n 2398, \n 2399, \n 2400, \n 2401, \n 2402, \n 2403, \n 2404, \n 2405, \n 2406, \n 2407, \n 2408, \n 2409, \n 2410, \n 2411, \n 2412, \n 2413, \n 2414, \n 2415, \n 2416, \n 2417, \n 2418, \n 2419, \n 2420, \n 2421, \n 2422, \n 2423, \n 2424, \n 2425, \n 2426, \n 2427, \n 2428, \n 2429, \n 2430, \n 2431, \n 2432, \n 2433, \n 2434, \n 2435, \n 2436, \n 2437, \n 2438, \n 2439, \n 2440, \n 2441, \n 2442, \n 2443, \n 2444, \n 2445, \n 2446, \n 2447, \n 2448, \n 2449, \n 2450, \n 2451, \n 2452, \n 2453, \n 2454, \n 2455, \n 2456, \n 2457, \n 2458, \n 2459, \n 2460, \n 2461, \n 2462, \n 2463, \n 2464, \n 2465, \n 2466, \n 2467, \n 2468, \n 2469, \n 2470, \n 2471, \n 2472, \n 2473, \n 2474, \n 2475, \n 2476, \n 2477, \n 2478, \n 2479, \n 2480, \n 2481, \n 2482, \n 2483, \n 2484, \n 2485, \n 2486, \n 2487, \n 2488, \n 2489, \n 2490, \n 2491, \n 2492, \n 2493, \n 2494, \n 2495, \n 2496, \n 2497, \n 2498, \n 2499, \n 2500, \n 2501, \n 2502, \n 2503, \n 2504, \n 2505, \n 2506, \n 2507, \n 2508, \n 2509, \n 2510, \n 2511, \n 2512, \n 2513, \n 2514, \n 2515, \n 2516, \n 2517, \n 2518, \n 2519, \n 2520, \n 2521, \n 2522, \n 2523, \n 2524, \n 2525, \n 2526, \n 2527, \n 2528, \n 2529, \n 2530, \n 2531, \n 2532, \n 2533, \n 2534, \n 2535, \n 2536, \n 2537, \n 2538, \n 2539, \n 2540, \n 2541, \n 2542, \n 2543, \n 2544, \n 2545, \n 2546, \n 2547, \n 2548, \n 2549, \n 2550, \n 2551, \n 2552, \n 2553, \n 2554, \n 2555, \n 2556, \n 2557, \n 2558, \n 2559, \n 2560, \n 2561, \n 2562, \n 2563, \n 2564, \n 2565, \n 2566, \n 2567, \n 2568, \n 2569, \n 2570, \n 2571, \n 2572, \n 2573, \n 2574, \n 2575, \n 2576, \n 2577, \n 2578, \n 2579, \n 2580, \n 2581, \n 2582, \n 2583, \n 2584, \n 2585, \n 2586, \n 2587, \n 2588, \n 2589, \n 2590, \n 2591, \n 2592, \n 2593, \n 2594, \n 2596, \n 2597, \n 2599, \n 2600, \n 2601, \n 2602, \n 2603, \n 2604, \n 2606, \n 2607, \n 2609, \n 2610, \n 2611, \n 2612, \n 2613, \n 2614, \n 2615, \n 2616, \n 2617, \n 2618, \n 2619, \n 2620, \n 2621, \n 2622, \n 2623, \n 2624, \n 2625, \n 2626, \n 2627, \n 2628, \n 2629, \n 2630, \n 2631, \n 2632, \n 2633, \n 2634, \n 2635, \n 2636, \n 2637, \n 2638, \n 2639, \n 2640, \n 2641, \n 2642, \n 2643, \n 2644, \n 2645, \n 2646, \n 2647, \n 2648, \n 2649, \n 2650, \n 2651, \n 2652, \n 2653, \n 2654, \n 2655, \n 2656, \n 2657, \n 2658, \n 2659, \n 2660, \n 2661, \n 2662, \n 2663, \n 2664, \n 2665, \n 2666, \n 2667, \n 2668, \n 2669, \n 2670, \n 2671, \n 2672, \n 2673, \n 2674, \n 2675, \n 2676, \n 2677, \n 2678, \n 2679, \n 2680, \n 2681, \n 2682, \n 2683, \n 2684, \n 2685, \n 2686, \n 2687, \n 2688, \n 2689, \n 2690, \n 2691, \n 2692, \n 2693, \n 2694, \n 2695, \n 2696\n ], \n \"rightLeg\": [\n 4481, \n 4482, \n 4485, \n 4486, \n 4491, \n 4492, \n 4493, \n 4495, \n 4498, \n 4500, \n 4501, \n 4505, \n 4506, \n 4529, \n 4532, \n 4533, \n 4534, \n 4535, \n 4536, \n 4537, \n 4538, \n 4539, \n 4540, \n 4541, \n 4542, \n 4543, \n 4544, \n 4545, \n 4546, \n 4547, \n 4548, \n 4549, \n 4550, \n 4551, \n 4552, \n 4553, \n 4554, \n 4555, \n 4556, \n 4557, \n 4558, \n 4559, \n 4560, \n 4561, \n 4562, \n 4563, \n 4564, \n 4565, \n 4566, \n 4567, \n 4568, \n 4569, \n 4570, \n 4571, \n 4572, \n 4573, \n 4574, \n 4575, \n 4576, \n 4577, \n 4578, \n 4579, \n 4580, \n 4581, \n 4582, \n 4583, \n 4584, \n 4585, \n 4586, \n 4587, \n 4588, \n 4589, \n 4590, \n 4591, \n 4592, \n 4593, \n 4594, \n 4595, \n 4596, \n 4597, \n 4598, \n 4599, \n 4600, \n 4601, \n 4602, \n 4603, \n 4604, \n 4605, \n 4606, \n 4607, \n 4608, \n 4609, \n 4610, \n 4611, \n 4612, \n 4613, \n 4614, \n 4615, \n 4616, \n 4617, \n 4618, \n 4619, \n 4620, \n 4621, \n 4622, \n 4634, \n 4635, \n 4636, \n 4637, \n 4638, \n 4639, \n 4640, \n 4641, \n 4642, \n 4643, \n 4644, \n 4661, \n 4662, \n 4663, \n 4664, \n 4665, \n 4666, \n 4667, \n 4668, \n 4669, \n 4842, \n 4843, \n 4844, \n 4845, \n 4846, \n 4847, \n 4848, \n 4937, \n 4938, \n 4939, \n 4940, \n 4941, \n 4942, \n 4943, \n 4944, \n 4945, \n 4946, \n 4947, \n 4993, \n 4994, \n 4995, \n 4996, \n 4997, \n 4998, \n 4999, \n 5000, \n 5001, \n 5002, \n 5003, \n 6574, \n 6575, \n 6576, \n 6577, \n 6578, \n 6579, \n 6580, \n 6581, \n 6582, \n 6583, \n 6584, \n 6585, \n 6586, \n 6587, \n 6588, \n 6589, \n 6590, \n 6591, \n 6592, \n 6593, \n 6594, \n 6595, \n 6596, \n 6597, \n 6598, \n 6599, \n 6600, \n 6601, \n 6602, \n 6603, \n 6604, \n 6605, \n 6606, \n 6607, \n 6608, \n 6609, \n 6610, \n 6719, \n 6720, \n 6721, \n 6722, \n 6723, \n 6724, \n 6725, \n 6726, \n 6727, \n 6728, \n 6729, \n 6730, \n 6731, \n 6732, \n 6733, \n 6734, \n 6735, \n 6832, \n 6833, \n 6834, \n 6835, \n 6836, \n 6869, \n 6870, \n 6871, \n 6872\n ], \n \"rightHandIndex1\": [\n 5488, \n 5489, \n 5490, \n 5491, \n 5498, \n 5499, \n 5500, \n 5501, \n 5518, \n 5528, \n 5529, \n 5584, \n 5585, \n 5586, \n 5587, \n 5588, \n 5589, \n 5590, \n 5591, \n 5592, \n 5606, \n 5607, \n 5613, \n 5615, \n 5616, \n 5617, \n 5618, \n 5619, \n 5620, \n 5621, \n 5622, \n 5623, \n 5624, \n 5625, \n 5626, \n 5627, \n 5628, \n 5629, \n 5630, \n 5638, \n 5639, \n 5640, \n 5642, \n 5647, \n 5648, \n 5650, \n 5651, \n 5665, \n 5666, \n 5676, \n 5677, \n 5678, \n 5679, \n 5680, \n 5681, \n 5693, \n 5694, \n 5706, \n 5707, \n 5708, \n 5719, \n 5721, \n 5722, \n 5723, \n 5724, \n 5730, \n 5731, \n 5733, \n 5734, \n 5735, \n 5737, \n 5738, \n 5741, \n 5742, \n 5743, \n 5744, \n 5752, \n 5753, \n 5754, \n 5755, \n 5756, \n 5757, \n 5758, \n 5759, \n 5760, \n 5761, \n 5762, \n 5763, \n 5764, \n 5765, \n 5766, \n 5767, \n 5768, \n 5769, \n 5770, \n 5771, \n 5772, \n 5773, \n 5774, \n 5775, \n 5776, \n 5777, \n 5778, \n 5779, \n 5780, \n 5781, \n 5782, \n 5783, \n 5784, \n 5785, \n 5786, \n 5787, \n 5788, \n 5789, \n 5790, \n 5791, \n 5792, \n 5793, \n 5794, \n 5795, \n 5796, \n 5797, \n 5798, \n 5799, \n 5800, \n 5801, \n 5802, \n 5803, \n 5804, \n 5805, \n 5806, \n 5807, \n 5808, \n 5809, \n 5810, \n 5811, \n 5812, \n 5813, \n 5814, \n 5815, \n 5816, \n 5817, \n 5818, \n 5819, \n 5820, \n 5821, \n 5822, \n 5823, \n 5824, \n 5825, \n 5826, \n 5827, \n 5828, \n 5829, \n 5830, \n 5831, \n 5832, \n 5833, \n 5834, \n 5835, \n 5836, \n 5837, \n 5838, \n 5839, \n 5840, \n 5841, \n 5842, \n 5843, \n 5844, \n 5845, \n 5846, \n 5847, \n 5848, \n 5849, \n 5850, \n 5851, \n 5852, \n 5853, \n 5854, \n 5855, \n 5856, \n 5857, \n 5858, \n 5859, \n 5860, \n 5861, \n 5862, \n 5863, \n 5864, \n 5865, \n 5866, \n 5867, \n 5868, \n 5869, \n 5870, \n 5871, \n 5872, \n 5873, \n 5874, \n 5875, \n 5876, \n 5877, \n 5878, \n 5879, \n 5880, \n 5881, \n 5882, \n 5883, \n 5884, \n 5885, \n 5886, \n 5887, \n 5888, \n 5889, \n 5890, \n 5891, \n 5892, \n 5893, \n 5894, \n 5895, \n 5896, \n 5897, \n 5898, \n 5899, \n 5900, \n 5901, \n 5902, \n 5903, \n 5904, \n 5905, \n 5906, \n 5907, \n 5908, \n 5909, \n 5910, \n 5911, \n 5912, \n 5913, \n 5914, \n 5915, \n 5916, \n 5917, \n 5918, \n 5919, \n 5920, \n 5921, \n 5922, \n 5923, \n 5924, \n 5925, \n 5926, \n 5927, \n 5928, \n 5929, \n 5930, \n 5931, \n 5932, \n 5933, \n 5934, \n 5935, \n 5936, \n 5937, \n 5938, \n 5939, \n 5940, \n 5941, \n 5942, \n 5943, \n 5944, \n 5945, \n 5946, \n 5947, \n 5948, \n 5949, \n 5950, \n 5951, \n 5952, \n 5953, \n 5954, \n 5955, \n 5956, \n 5957, \n 5958, \n 5959, \n 5960, \n 5961, \n 5962, \n 5963, \n 5964, \n 5965, \n 5966, \n 5967, \n 5968, \n 5969, \n 5970, \n 5971, \n 5972, \n 5973, \n 5974, \n 5975, \n 5976, \n 5977, \n 5978, \n 5979, \n 5980, \n 5981, \n 5982, \n 5983, \n 5984, \n 5985, \n 5986, \n 5987, \n 5988, \n 5989, \n 5990, \n 5991, \n 5992, \n 5993, \n 5994, \n 5995, \n 5996, \n 5997, \n 5998, \n 5999, \n 6000, \n 6001, \n 6002, \n 6003, \n 6004, \n 6005, \n 6006, \n 6007, \n 6008, \n 6009, \n 6010, \n 6011, \n 6012, \n 6013, \n 6014, \n 6015, \n 6016, \n 6017, \n 6018, \n 6019, \n 6020, \n 6021, \n 6022, \n 6023, \n 6024, \n 6025, \n 6026, \n 6027, \n 6028, \n 6029, \n 6030, \n 6031, \n 6032, \n 6033, \n 6034, \n 6035, \n 6036, \n 6037, \n 6038, \n 6039, \n 6040, \n 6041, \n 6042, \n 6043, \n 6044, \n 6045, \n 6046, \n 6047, \n 6048, \n 6049, \n 6050, \n 6051, \n 6052, \n 6053, \n 6054, \n 6055, \n 6058, \n 6059, \n 6060, \n 6061, \n 6062, \n 6063, \n 6064, \n 6065, \n 6068, \n 6069, \n 6070, \n 6071, \n 6072, \n 6073, \n 6074, \n 6075, \n 6076, \n 6077, \n 6078, \n 6079, \n 6080, \n 6081, \n 6082, \n 6083, \n 6084, \n 6085, \n 6086, \n 6087, \n 6088, \n 6089, \n 6090, \n 6091, \n 6092, \n 6093, \n 6094, \n 6095, \n 6096, \n 6097, \n 6098, \n 6099, \n 6100, \n 6101, \n 6102, \n 6103, \n 6104, \n 6105, \n 6106, \n 6107, \n 6108, \n 6109, \n 6110, \n 6111, \n 6112, \n 6113, \n 6114, \n 6115, \n 6116, \n 6117, \n 6118, \n 6119, \n 6120, \n 6121, \n 6122, \n 6123, \n 6124, \n 6125, \n 6126, \n 6127, \n 6128, \n 6129, \n 6130, \n 6131, \n 6132, \n 6133, \n 6134, \n 6135, \n 6136, \n 6137, \n 6138, \n 6139, \n 6140, \n 6141, \n 6142, \n 6143, \n 6144, \n 6145, \n 6146, \n 6147, \n 6148, \n 6149, \n 6150, \n 6151, \n 6152, \n 6153, \n 6154, \n 6155, \n 6156, \n 6157\n ], \n \"leftForeArm\": [\n 1546, \n 1547, \n 1548, \n 1549, \n 1550, \n 1551, \n 1552, \n 1553, \n 1554, \n 1555, \n 1556, \n 1557, \n 1558, \n 1559, \n 1560, \n 1561, \n 1562, \n 1563, \n 1564, \n 1565, \n 1566, \n 1567, \n 1568, \n 1569, \n 1570, \n 1571, \n 1572, \n 1573, \n 1574, \n 1575, \n 1576, \n 1577, \n 1578, \n 1579, \n 1580, \n 1581, \n 1582, \n 1583, \n 1584, \n 1585, \n 1586, \n 1587, \n 1588, \n 1589, \n 1590, \n 1591, \n 1592, \n 1593, \n 1594, \n 1595, \n 1596, \n 1597, \n 1598, \n 1599, \n 1600, \n 1601, \n 1602, \n 1603, \n 1604, \n 1605, \n 1606, \n 1607, \n 1608, \n 1609, \n 1610, \n 1611, \n 1612, \n 1613, \n 1614, \n 1615, \n 1616, \n 1617, \n 1618, \n 1620, \n 1621, \n 1623, \n 1624, \n 1625, \n 1626, \n 1627, \n 1628, \n 1629, \n 1630, \n 1643, \n 1644, \n 1646, \n 1647, \n 1650, \n 1651, \n 1654, \n 1655, \n 1657, \n 1658, \n 1659, \n 1660, \n 1661, \n 1662, \n 1663, \n 1664, \n 1665, \n 1666, \n 1685, \n 1686, \n 1687, \n 1688, \n 1689, \n 1690, \n 1691, \n 1692, \n 1693, \n 1694, \n 1695, \n 1699, \n 1700, \n 1701, \n 1702, \n 1721, \n 1722, \n 1723, \n 1724, \n 1725, \n 1726, \n 1727, \n 1728, \n 1729, \n 1730, \n 1732, \n 1736, \n 1738, \n 1741, \n 1742, \n 1743, \n 1744, \n 1750, \n 1752, \n 1900, \n 1909, \n 1910, \n 1911, \n 1912, \n 1913, \n 1914, \n 1915, \n 1916, \n 1917, \n 1918, \n 1919, \n 1920, \n 1921, \n 1922, \n 1923, \n 1924, \n 1925, \n 1926, \n 1927, \n 1928, \n 1929, \n 1930, \n 1931, \n 1932, \n 1933, \n 1934, \n 1935, \n 1936, \n 1937, \n 1938, \n 1939, \n 1940, \n 1941, \n 1942, \n 1943, \n 1944, \n 1945, \n 1946, \n 1947, \n 1948, \n 1949, \n 1950, \n 1951, \n 1952, \n 1953, \n 1954, \n 1955, \n 1956, \n 1957, \n 1958, \n 1959, \n 1960, \n 1961, \n 1962, \n 1963, \n 1964, \n 1965, \n 1966, \n 1967, \n 1968, \n 1969, \n 1970, \n 1971, \n 1972, \n 1973, \n 1974, \n 1975, \n 1976, \n 1977, \n 1978, \n 1979, \n 1980, \n 2019, \n 2059, \n 2060, \n 2073, \n 2089, \n 2098, \n 2099, \n 2100, \n 2101, \n 2102, \n 2103, \n 2104, \n 2105, \n 2106, \n 2107, \n 2108, \n 2109, \n 2110, \n 2111, \n 2112, \n 2147, \n 2148, \n 2206, \n 2207, \n 2208, \n 2209, \n 2228, \n 2230, \n 2234, \n 2235, \n 2241, \n 2242, \n 2243, \n 2244, \n 2279, \n 2286, \n 2873, \n 2874\n ], \n \"rightForeArm\": [\n 5015, \n 5016, \n 5017, \n 5018, \n 5019, \n 5020, \n 5021, \n 5022, \n 5023, \n 5024, \n 5025, \n 5026, \n 5027, \n 5028, \n 5029, \n 5030, \n 5031, \n 5032, \n 5033, \n 5034, \n 5035, \n 5036, \n 5037, \n 5038, \n 5039, \n 5040, \n 5041, \n 5042, \n 5043, \n 5044, \n 5045, \n 5046, \n 5047, \n 5048, \n 5049, \n 5050, \n 5051, \n 5052, \n 5053, \n 5054, \n 5055, \n 5056, \n 5057, \n 5058, \n 5059, \n 5060, \n 5061, \n 5062, \n 5063, \n 5064, \n 5065, \n 5066, \n 5067, \n 5068, \n 5069, \n 5070, \n 5071, \n 5072, \n 5073, \n 5074, \n 5075, \n 5076, \n 5077, \n 5078, \n 5079, \n 5080, \n 5081, \n 5082, \n 5083, \n 5084, \n 5085, \n 5086, \n 5087, \n 5090, \n 5091, \n 5092, \n 5093, \n 5094, \n 5095, \n 5096, \n 5097, \n 5098, \n 5099, \n 5112, \n 5113, \n 5116, \n 5117, \n 5120, \n 5121, \n 5124, \n 5125, \n 5126, \n 5127, \n 5128, \n 5129, \n 5130, \n 5131, \n 5132, \n 5133, \n 5134, \n 5135, \n 5154, \n 5155, \n 5156, \n 5157, \n 5158, \n 5159, \n 5160, \n 5161, \n 5162, \n 5163, \n 5164, \n 5168, \n 5169, \n 5170, \n 5171, \n 5190, \n 5191, \n 5192, \n 5193, \n 5194, \n 5195, \n 5196, \n 5197, \n 5198, \n 5199, \n 5202, \n 5205, \n 5207, \n 5210, \n 5211, \n 5212, \n 5213, \n 5219, \n 5221, \n 5361, \n 5370, \n 5371, \n 5372, \n 5373, \n 5374, \n 5375, \n 5376, \n 5377, \n 5378, \n 5379, \n 5380, \n 5381, \n 5382, \n 5383, \n 5384, \n 5385, \n 5386, \n 5387, \n 5388, \n 5389, \n 5390, \n 5391, \n 5392, \n 5393, \n 5394, \n 5395, \n 5396, \n 5397, \n 5398, \n 5399, \n 5400, \n 5401, \n 5402, \n 5403, \n 5404, \n 5405, \n 5406, \n 5407, \n 5408, \n 5409, \n 5410, \n 5411, \n 5412, \n 5413, \n 5414, \n 5415, \n 5416, \n 5417, \n 5418, \n 5419, \n 5420, \n 5421, \n 5422, \n 5423, \n 5424, \n 5425, \n 5426, \n 5427, \n 5428, \n 5429, \n 5430, \n 5431, \n 5432, \n 5433, \n 5434, \n 5435, \n 5436, \n 5437, \n 5438, \n 5439, \n 5440, \n 5441, \n 5480, \n 5520, \n 5521, \n 5534, \n 5550, \n 5559, \n 5560, \n 5561, \n 5562, \n 5563, \n 5564, \n 5565, \n 5566, \n 5567, \n 5568, \n 5569, \n 5570, \n 5571, \n 5572, \n 5573, \n 5608, \n 5609, \n 5667, \n 5668, \n 5669, \n 5670, \n 5689, \n 5691, \n 5695, \n 5696, \n 5702, \n 5703, \n 5704, \n 5705, \n 5740, \n 5747, \n 6334, \n 6335\n ], \n \"neck\": [\n 148, \n 150, \n 151, \n 152, \n 153, \n 172, \n 174, \n 175, \n 201, \n 202, \n 204, \n 205, \n 206, \n 207, \n 208, \n 209, \n 210, \n 211, \n 212, \n 213, \n 214, \n 215, \n 216, \n 217, \n 218, \n 219, \n 222, \n 223, \n 224, \n 225, \n 256, \n 257, \n 284, \n 285, \n 295, \n 296, \n 297, \n 298, \n 299, \n 300, \n 301, \n 302, \n 303, \n 304, \n 305, \n 306, \n 307, \n 308, \n 309, \n 333, \n 334, \n 423, \n 424, \n 425, \n 426, \n 440, \n 441, \n 451, \n 452, \n 453, \n 460, \n 461, \n 571, \n 572, \n 824, \n 825, \n 826, \n 827, \n 828, \n 829, \n 1279, \n 1280, \n 1312, \n 1313, \n 1319, \n 1320, \n 1331, \n 3049, \n 3050, \n 3057, \n 3058, \n 3059, \n 3068, \n 3164, \n 3661, \n 3662, \n 3663, \n 3664, \n 3665, \n 3685, \n 3686, \n 3687, \n 3714, \n 3715, \n 3716, \n 3717, \n 3718, \n 3719, \n 3720, \n 3721, \n 3722, \n 3723, \n 3724, \n 3725, \n 3726, \n 3727, \n 3728, \n 3729, \n 3730, \n 3731, \n 3734, \n 3735, \n 3736, \n 3737, \n 3768, \n 3769, \n 3796, \n 3797, \n 3807, \n 3808, \n 3809, \n 3810, \n 3811, \n 3812, \n 3813, \n 3814, \n 3815, \n 3816, \n 3817, \n 3818, \n 3819, \n 3839, \n 3840, \n 3918, \n 3919, \n 3920, \n 3921, \n 3934, \n 3935, \n 3942, \n 3943, \n 3944, \n 3950, \n 4060, \n 4061, \n 4312, \n 4313, \n 4314, \n 4315, \n 4761, \n 4762, \n 4792, \n 4793, \n 4799, \n 4800, \n 4807\n ], \n \"rightToeBase\": [\n 6611, \n 6612, \n 6613, \n 6614, \n 6615, \n 6616, \n 6617, \n 6618, \n 6619, \n 6620, \n 6621, \n 6622, \n 6623, \n 6624, \n 6625, \n 6626, \n 6627, \n 6628, \n 6629, \n 6630, \n 6631, \n 6632, \n 6633, \n 6634, \n 6635, \n 6636, \n 6637, \n 6638, \n 6639, \n 6640, \n 6641, \n 6642, \n 6643, \n 6644, \n 6645, \n 6646, \n 6647, \n 6648, \n 6649, \n 6650, \n 6651, \n 6652, \n 6653, \n 6654, \n 6655, \n 6656, \n 6657, \n 6658, \n 6659, \n 6660, \n 6661, \n 6662, \n 6663, \n 6664, \n 6665, \n 6666, \n 6667, \n 6668, \n 6669, \n 6670, \n 6671, \n 6672, \n 6673, \n 6674, \n 6675, \n 6676, \n 6677, \n 6678, \n 6679, \n 6680, \n 6681, \n 6682, \n 6683, \n 6684, \n 6685, \n 6686, \n 6687, \n 6688, \n 6689, \n 6690, \n 6691, \n 6692, \n 6693, \n 6694, \n 6695, \n 6696, \n 6697, \n 6698, \n 6699, \n 6700, \n 6701, \n 6702, \n 6703, \n 6704, \n 6705, \n 6706, \n 6707, \n 6708, \n 6709, \n 6710, \n 6711, \n 6712, \n 6713, \n 6714, \n 6715, \n 6716, \n 6717, \n 6718, \n 6736, \n 6739, \n 6741, \n 6743, \n 6745, \n 6747, \n 6749, \n 6750, \n 6752, \n 6754, \n 6757, \n 6758, \n 6760, \n 6762\n ], \n \"spine\": [\n 616, \n 617, \n 630, \n 631, \n 632, \n 633, \n 654, \n 655, \n 656, \n 657, \n 662, \n 663, \n 664, \n 665, \n 720, \n 721, \n 765, \n 766, \n 767, \n 768, \n 796, \n 797, \n 798, \n 799, \n 889, \n 890, \n 916, \n 917, \n 918, \n 919, \n 921, \n 922, \n 923, \n 924, \n 925, \n 926, \n 1188, \n 1189, \n 1211, \n 1212, \n 1248, \n 1249, \n 1250, \n 1251, \n 1264, \n 1265, \n 1266, \n 1267, \n 1323, \n 1324, \n 1325, \n 1326, \n 1327, \n 1328, \n 1332, \n 1333, \n 1334, \n 1335, \n 1336, \n 1344, \n 1345, \n 1481, \n 1482, \n 1483, \n 1484, \n 1485, \n 1486, \n 1487, \n 1488, \n 1489, \n 1490, \n 1491, \n 1492, \n 1493, \n 1494, \n 1495, \n 1496, \n 1767, \n 2823, \n 2824, \n 2825, \n 2826, \n 2827, \n 2828, \n 2829, \n 2830, \n 2831, \n 2832, \n 2833, \n 2834, \n 2835, \n 2836, \n 2837, \n 2838, \n 2839, \n 2840, \n 2841, \n 2842, \n 2843, \n 2844, \n 2845, \n 2847, \n 2848, \n 2851, \n 3016, \n 3017, \n 3018, \n 3019, \n 3020, \n 3023, \n 3024, \n 3124, \n 3173, \n 3476, \n 3477, \n 3478, \n 3480, \n 3500, \n 3501, \n 3502, \n 3504, \n 3509, \n 3511, \n 4103, \n 4104, \n 4118, \n 4119, \n 4120, \n 4121, \n 4142, \n 4143, \n 4144, \n 4145, \n 4150, \n 4151, \n 4152, \n 4153, \n 4208, \n 4209, \n 4253, \n 4254, \n 4255, \n 4256, \n 4284, \n 4285, \n 4286, \n 4287, \n 4375, \n 4376, \n 4402, \n 4403, \n 4405, \n 4406, \n 4407, \n 4408, \n 4409, \n 4410, \n 4411, \n 4412, \n 4674, \n 4675, \n 4694, \n 4695, \n 4731, \n 4732, \n 4733, \n 4734, \n 4747, \n 4748, \n 4749, \n 4750, \n 4803, \n 4804, \n 4805, \n 4806, \n 4808, \n 4809, \n 4810, \n 4811, \n 4812, \n 4820, \n 4821, \n 4953, \n 4954, \n 4955, \n 4956, \n 4957, \n 4958, \n 4959, \n 4960, \n 4961, \n 4962, \n 4963, \n 4964, \n 4965, \n 4966, \n 4967, \n 4968, \n 5234, \n 6284, \n 6285, \n 6286, \n 6287, \n 6288, \n 6289, \n 6290, \n 6291, \n 6292, \n 6293, \n 6294, \n 6295, \n 6296, \n 6297, \n 6298, \n 6299, \n 6300, \n 6301, \n 6302, \n 6303, \n 6304, \n 6305, \n 6306, \n 6308, \n 6309, \n 6312, \n 6472, \n 6473, \n 6474, \n 6545, \n 6874, \n 6875, \n 6876, \n 6878\n ], \n \"leftUpLeg\": [\n 833, \n 834, \n 838, \n 839, \n 847, \n 848, \n 849, \n 850, \n 851, \n 852, \n 853, \n 854, \n 870, \n 871, \n 872, \n 873, \n 874, \n 875, \n 876, \n 877, \n 878, \n 879, \n 880, \n 881, \n 897, \n 898, \n 899, \n 900, \n 901, \n 902, \n 903, \n 904, \n 905, \n 906, \n 907, \n 908, \n 909, \n 910, \n 911, \n 912, \n 913, \n 914, \n 915, \n 933, \n 934, \n 935, \n 936, \n 944, \n 945, \n 946, \n 947, \n 948, \n 949, \n 950, \n 951, \n 952, \n 953, \n 954, \n 955, \n 956, \n 957, \n 958, \n 959, \n 960, \n 961, \n 962, \n 963, \n 964, \n 965, \n 966, \n 967, \n 968, \n 969, \n 970, \n 971, \n 972, \n 973, \n 974, \n 975, \n 976, \n 977, \n 978, \n 979, \n 980, \n 981, \n 982, \n 983, \n 984, \n 985, \n 986, \n 987, \n 988, \n 989, \n 990, \n 991, \n 992, \n 993, \n 994, \n 995, \n 996, \n 997, \n 998, \n 999, \n 1000, \n 1001, \n 1002, \n 1003, \n 1004, \n 1005, \n 1006, \n 1007, \n 1008, \n 1009, \n 1010, \n 1011, \n 1012, \n 1013, \n 1014, \n 1015, \n 1016, \n 1017, \n 1018, \n 1019, \n 1020, \n 1021, \n 1022, \n 1023, \n 1024, \n 1025, \n 1026, \n 1027, \n 1028, \n 1029, \n 1030, \n 1031, \n 1032, \n 1033, \n 1034, \n 1035, \n 1036, \n 1037, \n 1038, \n 1039, \n 1040, \n 1041, \n 1042, \n 1043, \n 1044, \n 1045, \n 1046, \n 1137, \n 1138, \n 1139, \n 1140, \n 1141, \n 1142, \n 1143, \n 1144, \n 1145, \n 1146, \n 1147, \n 1148, \n 1159, \n 1160, \n 1161, \n 1162, \n 1163, \n 1164, \n 1165, \n 1166, \n 1167, \n 1168, \n 1169, \n 1170, \n 1171, \n 1172, \n 1173, \n 1174, \n 1184, \n 1185, \n 1186, \n 1187, \n 1221, \n 1222, \n 1223, \n 1224, \n 1225, \n 1226, \n 1227, \n 1228, \n 1229, \n 1230, \n 1262, \n 1263, \n 1274, \n 1275, \n 1276, \n 1277, \n 1321, \n 1322, \n 1354, \n 1359, \n 1360, \n 1361, \n 1362, \n 1365, \n 1366, \n 1367, \n 1368, \n 1451, \n 1452, \n 1453, \n 1455, \n 1456, \n 1457, \n 1458, \n 1459, \n 1460, \n 1461, \n 1462, \n 1463, \n 1475, \n 1477, \n 1478, \n 1479, \n 1480, \n 1498, \n 1499, \n 1500, \n 1501, \n 1511, \n 1512, \n 1513, \n 1514, \n 1516, \n 1517, \n 1518, \n 1519, \n 1520, \n 1521, \n 1522, \n 1533, \n 1534, \n 3125, \n 3126, \n 3127, \n 3128, \n 3131, \n 3132, \n 3133, \n 3134, \n 3135, \n 3475, \n 3479\n ], \n \"leftHand\": [\n 1981, \n 1982, \n 1983, \n 1984, \n 1985, \n 1986, \n 1987, \n 1988, \n 1989, \n 1990, \n 1991, \n 1992, \n 1993, \n 1994, \n 1995, \n 1996, \n 1997, \n 1998, \n 1999, \n 2000, \n 2001, \n 2002, \n 2003, \n 2004, \n 2005, \n 2006, \n 2007, \n 2008, \n 2009, \n 2010, \n 2011, \n 2012, \n 2013, \n 2014, \n 2015, \n 2016, \n 2017, \n 2018, \n 2019, \n 2020, \n 2021, \n 2022, \n 2023, \n 2024, \n 2025, \n 2026, \n 2031, \n 2032, \n 2033, \n 2034, \n 2035, \n 2036, \n 2041, \n 2042, \n 2043, \n 2044, \n 2045, \n 2046, \n 2047, \n 2048, \n 2049, \n 2050, \n 2051, \n 2052, \n 2053, \n 2054, \n 2055, \n 2056, \n 2057, \n 2058, \n 2059, \n 2060, \n 2061, \n 2062, \n 2063, \n 2064, \n 2065, \n 2066, \n 2069, \n 2070, \n 2071, \n 2072, \n 2073, \n 2074, \n 2075, \n 2076, \n 2077, \n 2078, \n 2079, \n 2080, \n 2081, \n 2082, \n 2083, \n 2084, \n 2085, \n 2086, \n 2087, \n 2088, \n 2089, \n 2090, \n 2091, \n 2092, \n 2093, \n 2094, \n 2095, \n 2096, \n 2097, \n 2098, \n 2099, \n 2100, \n 2101, \n 2107, \n 2111, \n 2113, \n 2114, \n 2115, \n 2116, \n 2117, \n 2118, \n 2119, \n 2120, \n 2121, \n 2122, \n 2127, \n 2130, \n 2131, \n 2132, \n 2133, \n 2134, \n 2135, \n 2136, \n 2137, \n 2138, \n 2139, \n 2140, \n 2141, \n 2142, \n 2143, \n 2144, \n 2149, \n 2150, \n 2151, \n 2152, \n 2155, \n 2160, \n 2163, \n 2164, \n 2170, \n 2171, \n 2172, \n 2173, \n 2174, \n 2175, \n 2176, \n 2177, \n 2178, \n 2179, \n 2180, \n 2182, \n 2183, \n 2184, \n 2185, \n 2188, \n 2189, \n 2191, \n 2192, \n 2193, \n 2194, \n 2195, \n 2196, \n 2197, \n 2198, \n 2199, \n 2200, \n 2201, \n 2202, \n 2203, \n 2207, \n 2209, \n 2210, \n 2211, \n 2212, \n 2213, \n 2214, \n 2221, \n 2222, \n 2223, \n 2224, \n 2225, \n 2226, \n 2227, \n 2228, \n 2229, \n 2231, \n 2234, \n 2236, \n 2237, \n 2238, \n 2239, \n 2240, \n 2246, \n 2247, \n 2248, \n 2249, \n 2250, \n 2251, \n 2252, \n 2253, \n 2254, \n 2255, \n 2256, \n 2257, \n 2258, \n 2259, \n 2260, \n 2262, \n 2263, \n 2264, \n 2265, \n 2266, \n 2267, \n 2268, \n 2269, \n 2270, \n 2271, \n 2274, \n 2275, \n 2276, \n 2277, \n 2278, \n 2279, \n 2284, \n 2285, \n 2287, \n 2288, \n 2289, \n 2290, \n 2293, \n 2595, \n 2598, \n 2605, \n 2608, \n 2697, \n 2698, \n 2699, \n 2700, \n 2701, \n 2702, \n 2703, \n 2704, \n 2705, \n 2706, \n 2707, \n 2708, \n 2709, \n 2710, \n 2711, \n 2712, \n 2713, \n 2714, \n 2715, \n 2716, \n 2717, \n 2718, \n 2719, \n 2720, \n 2721, \n 2722, \n 2723, \n 2724, \n 2725, \n 2726, \n 2727, \n 2728, \n 2729, \n 2730, \n 2731, \n 2732, \n 2733, \n 2734, \n 2735, \n 2736, \n 2737, \n 2738, \n 2739, \n 2740, \n 2741, \n 2742, \n 2743, \n 2744, \n 2745, \n 2746, \n 2747, \n 2748, \n 2749, \n 2750, \n 2751, \n 2752, \n 2753, \n 2754, \n 2755, \n 2756, \n 2757, \n 2758, \n 2759, \n 2760, \n 2761, \n 2762, \n 2763, \n 2764, \n 2765, \n 2766, \n 2767, \n 2768, \n 2769, \n 2770, \n 2771, \n 2772, \n 2773, \n 2774, \n 2775, \n 2776, \n 2777, \n 2778\n ], \n \"hips\": [\n 631, \n 632, \n 654, \n 657, \n 662, \n 665, \n 676, \n 677, \n 678, \n 679, \n 705, \n 720, \n 796, \n 799, \n 800, \n 801, \n 802, \n 807, \n 808, \n 809, \n 810, \n 815, \n 816, \n 822, \n 823, \n 830, \n 831, \n 832, \n 833, \n 834, \n 835, \n 836, \n 837, \n 838, \n 839, \n 840, \n 841, \n 842, \n 843, \n 844, \n 845, \n 846, \n 855, \n 856, \n 857, \n 858, \n 859, \n 860, \n 861, \n 862, \n 863, \n 864, \n 865, \n 866, \n 867, \n 868, \n 869, \n 871, \n 878, \n 881, \n 882, \n 883, \n 884, \n 885, \n 886, \n 887, \n 888, \n 889, \n 890, \n 912, \n 915, \n 916, \n 917, \n 918, \n 919, \n 920, \n 932, \n 937, \n 938, \n 939, \n 1163, \n 1166, \n 1203, \n 1204, \n 1205, \n 1206, \n 1207, \n 1208, \n 1209, \n 1210, \n 1246, \n 1247, \n 1262, \n 1263, \n 1276, \n 1277, \n 1278, \n 1321, \n 1336, \n 1337, \n 1338, \n 1339, \n 1353, \n 1354, \n 1361, \n 1362, \n 1363, \n 1364, \n 1446, \n 1447, \n 1448, \n 1449, \n 1450, \n 1454, \n 1476, \n 1497, \n 1511, \n 1513, \n 1514, \n 1515, \n 1533, \n 1534, \n 1539, \n 1540, \n 1768, \n 1769, \n 1779, \n 1780, \n 1781, \n 1782, \n 1783, \n 1784, \n 1785, \n 1786, \n 1787, \n 1788, \n 1789, \n 1790, \n 1791, \n 1792, \n 1793, \n 1794, \n 1795, \n 1796, \n 1797, \n 1798, \n 1799, \n 1800, \n 1801, \n 1802, \n 1803, \n 1804, \n 1805, \n 1806, \n 1807, \n 2909, \n 2910, \n 2911, \n 2912, \n 2913, \n 2914, \n 2915, \n 2916, \n 2917, \n 2918, \n 2919, \n 2920, \n 2921, \n 2922, \n 2923, \n 2924, \n 2925, \n 2926, \n 2927, \n 2928, \n 2929, \n 2930, \n 3018, \n 3019, \n 3021, \n 3022, \n 3080, \n 3081, \n 3082, \n 3083, \n 3084, \n 3085, \n 3086, \n 3087, \n 3088, \n 3089, \n 3090, \n 3091, \n 3092, \n 3093, \n 3094, \n 3095, \n 3096, \n 3097, \n 3098, \n 3099, \n 3100, \n 3101, \n 3102, \n 3103, \n 3104, \n 3105, \n 3106, \n 3107, \n 3108, \n 3109, \n 3110, \n 3111, \n 3112, \n 3113, \n 3114, \n 3115, \n 3116, \n 3117, \n 3118, \n 3119, \n 3120, \n 3121, \n 3122, \n 3123, \n 3124, \n 3128, \n 3129, \n 3130, \n 3136, \n 3137, \n 3138, \n 3139, \n 3140, \n 3141, \n 3142, \n 3143, \n 3144, \n 3145, \n 3146, \n 3147, \n 3148, \n 3149, \n 3150, \n 3151, \n 3152, \n 3153, \n 3154, \n 3155, \n 3156, \n 3157, \n 3158, \n 3159, \n 3160, \n 3170, \n 3172, \n 3481, \n 3484, \n 3500, \n 3502, \n 3503, \n 3507, \n 3510, \n 4120, \n 4121, \n 4142, \n 4143, \n 4150, \n 4151, \n 4164, \n 4165, \n 4166, \n 4167, \n 4193, \n 4208, \n 4284, \n 4285, \n 4288, \n 4289, \n 4290, \n 4295, \n 4296, \n 4297, \n 4298, \n 4303, \n 4304, \n 4310, \n 4311, \n 4316, \n 4317, \n 4318, \n 4319, \n 4320, \n 4321, \n 4322, \n 4323, \n 4324, \n 4325, \n 4326, \n 4327, \n 4328, \n 4329, \n 4330, \n 4331, \n 4332, \n 4341, \n 4342, \n 4343, \n 4344, \n 4345, \n 4346, \n 4347, \n 4348, \n 4349, \n 4350, \n 4351, \n 4352, \n 4353, \n 4354, \n 4355, \n 4356, \n 4364, \n 4365, \n 4368, \n 4369, \n 4370, \n 4371, \n 4372, \n 4373, \n 4374, \n 4375, \n 4376, \n 4398, \n 4399, \n 4402, \n 4403, \n 4404, \n 4405, \n 4406, \n 4418, \n 4423, \n 4424, \n 4425, \n 4649, \n 4650, \n 4689, \n 4690, \n 4691, \n 4692, \n 4693, \n 4729, \n 4730, \n 4745, \n 4746, \n 4759, \n 4760, \n 4801, \n 4812, \n 4813, \n 4814, \n 4815, \n 4829, \n 4836, \n 4837, \n 4919, \n 4920, \n 4921, \n 4922, \n 4923, \n 4927, \n 4969, \n 4983, \n 4984, \n 4986, \n 5004, \n 5005, \n 5244, \n 5245, \n 5246, \n 5247, \n 5248, \n 5249, \n 5250, \n 5251, \n 5252, \n 5253, \n 5254, \n 5255, \n 5256, \n 5257, \n 5258, \n 5259, \n 5260, \n 5261, \n 5262, \n 5263, \n 5264, \n 5265, \n 5266, \n 5267, \n 5268, \n 6368, \n 6369, \n 6370, \n 6371, \n 6372, \n 6373, \n 6374, \n 6375, \n 6376, \n 6377, \n 6378, \n 6379, \n 6380, \n 6381, \n 6382, \n 6383, \n 6384, \n 6385, \n 6386, \n 6387, \n 6388, \n 6389, \n 6473, \n 6474, \n 6504, \n 6505, \n 6506, \n 6507, \n 6508, \n 6509, \n 6510, \n 6511, \n 6512, \n 6513, \n 6514, \n 6515, \n 6516, \n 6517, \n 6518, \n 6519, \n 6520, \n 6521, \n 6522, \n 6523, \n 6524, \n 6525, \n 6526, \n 6527, \n 6528, \n 6529, \n 6530, \n 6531, \n 6532, \n 6533, \n 6534, \n 6535, \n 6536, \n 6537, \n 6538, \n 6539, \n 6540, \n 6541, \n 6542, \n 6543, \n 6544, \n 6545, \n 6549, \n 6550, \n 6551, \n 6557, \n 6558, \n 6559, \n 6560, \n 6561, \n 6562, \n 6563, \n 6564, \n 6565, \n 6566, \n 6567, \n 6568, \n 6569, \n 6570, \n 6571, \n 6572, \n 6573\n ]\n}" }, { "path": "eval/GVHMR/hmr4d/utils/body_model/smplx_lite.py", "content": "import numpy as np\nimport torch\nimport torch.nn as nn\nimport torch.nn.functional as F\nfrom pathlib import Path\nfrom pytorch3d.transforms import axis_angle_to_matrix, rotation_6d_to_matrix\nfrom smplx.utils import Struct, to_np, to_tensor\nfrom einops import einsum, rearrange\nfrom time import time\n\nfrom hmr4d import PROJ_ROOT\n\n\nclass SmplxLite(nn.Module):\n def __init__(\n self,\n model_path=PROJ_ROOT / \"inputs/checkpoints/body_models/smplx\",\n gender=\"neutral\",\n num_betas=10,\n ):\n super().__init__()\n\n # Load the model\n model_path = Path(model_path)\n if model_path.is_dir():\n smplx_path = Path(model_path) / f\"SMPLX_{gender.upper()}.npz\"\n else:\n smplx_path = model_path\n assert smplx_path.exists()\n model_data = np.load(smplx_path, allow_pickle=True)\n\n data_struct = Struct(**model_data)\n self.faces = data_struct.f # (F, 3)\n\n self.register_smpl_buffers(data_struct, num_betas)\n # self.register_smplh_buffers(data_struct, num_pca_comps, flat_hand_mean)\n # self.register_smplx_buffers(data_struct)\n self.register_fast_skeleton_computing_buffers()\n\n # default_pose (99,) for torch.cat([global_orient, body_pose, default_pose])\n other_default_pose = torch.cat(\n [\n torch.zeros(9),\n to_tensor(data_struct.hands_meanl).float(),\n to_tensor(data_struct.hands_meanr).float(),\n ]\n )\n self.register_buffer(\"other_default_pose\", other_default_pose, False)\n\n def register_smpl_buffers(self, data_struct, num_betas):\n # shapedirs, (V, 3, N_betas), V=10475 for SMPLX\n shapedirs = to_tensor(to_np(data_struct.shapedirs[:, :, :num_betas])).float()\n self.register_buffer(\"shapedirs\", shapedirs, False)\n\n # v_template, (V, 3)\n v_template = to_tensor(to_np(data_struct.v_template)).float()\n self.register_buffer(\"v_template\", v_template, False)\n\n # J_regressor, (J, V), J=55 for SMPLX\n J_regressor = to_tensor(to_np(data_struct.J_regressor)).float()\n self.register_buffer(\"J_regressor\", J_regressor, False)\n\n # posedirs, (54*9, V, 3), note that the first global_orient is not included\n posedirs = to_tensor(to_np(data_struct.posedirs)).float() # (V, 3, 54*9)\n posedirs = rearrange(posedirs, \"v c n -> n v c\")\n self.register_buffer(\"posedirs\", posedirs, False)\n\n # lbs_weights, (V, J), J=55\n lbs_weights = to_tensor(to_np(data_struct.weights)).float()\n self.register_buffer(\"lbs_weights\", lbs_weights, False)\n\n # parents, (J), long\n parents = to_tensor(to_np(data_struct.kintree_table[0])).long()\n parents[0] = -1\n self.register_buffer(\"parents\", parents, False)\n\n def register_smplh_buffers(self, data_struct, num_pca_comps, flat_hand_mean):\n # hand_pca, (N_pca, 45)\n left_hand_components = to_tensor(data_struct.hands_componentsl[:num_pca_comps]).float()\n right_hand_components = to_tensor(data_struct.hands_componentsr[:num_pca_comps]).float()\n self.register_buffer(\"left_hand_components\", left_hand_components, False)\n self.register_buffer(\"right_hand_components\", right_hand_components, False)\n\n # hand_mean, (45,)\n left_hand_mean = to_tensor(data_struct.hands_meanl).float()\n right_hand_mean = to_tensor(data_struct.hands_meanr).float()\n if not flat_hand_mean:\n left_hand_mean = torch.zeros_like(left_hand_mean)\n right_hand_mean = torch.zeros_like(right_hand_mean)\n self.register_buffer(\"left_hand_mean\", left_hand_mean, False)\n self.register_buffer(\"right_hand_mean\", right_hand_mean, False)\n\n def register_smplx_buffers(self, data_struct):\n # expr_dirs, (V, 3, N_expr)\n expr_dirs = to_tensor(to_np(data_struct.shapedirs[:, :, 300:310])).float()\n self.register_buffer(\"expr_dirs\", expr_dirs, False)\n\n def register_fast_skeleton_computing_buffers(self):\n # For fast computing of skeleton under beta\n J_template = self.J_regressor @ self.v_template # (J, 3)\n J_shapedirs = torch.einsum(\"jv, vcd -> jcd\", self.J_regressor, self.shapedirs) # (J, 3, 10)\n self.register_buffer(\"J_template\", J_template, False)\n self.register_buffer(\"J_shapedirs\", J_shapedirs, False)\n\n def get_skeleton(self, betas):\n return self.J_template + einsum(betas, self.J_shapedirs, \"... k, j c k -> ... j c\")\n\n def forward(\n self,\n body_pose,\n betas,\n global_orient,\n transl=None,\n rotation_type=\"aa\",\n ):\n \"\"\"\n Args:\n body_pose: (B, L, 63)\n betas: (B, L, 10)\n global_orient: (B, L, 3)\n transl: (B, L, 3)\n Returns:\n vertices: (B, L, V, 3)\n \"\"\"\n # 1. Convert [global_orient, body_pose, other_default_pose] to rot_mats\n other_default_pose = self.other_default_pose # (99,)\n if rotation_type == \"aa\":\n other_default_pose = other_default_pose.expand(*body_pose.shape[:-1], -1)\n full_pose = torch.cat([global_orient, body_pose, other_default_pose], dim=-1)\n rot_mats = axis_angle_to_matrix(full_pose.reshape(*full_pose.shape[:-1], 55, 3))\n del full_pose, other_default_pose\n else:\n assert rotation_type == \"r6d\" # useful when doing smplify\n other_default_pose = axis_angle_to_matrix(other_default_pose.view(33, 3))\n part_full_pose = torch.cat([global_orient, body_pose], dim=-1)\n rot_mats = rotation_6d_to_matrix(part_full_pose.view(*part_full_pose.shape[:-1], 22, 6))\n other_default_pose = other_default_pose.expand(*rot_mats.shape[:-3], -1, -1, -1)\n rot_mats = torch.cat([rot_mats, other_default_pose], dim=-3)\n del part_full_pose, other_default_pose\n\n # 2. Forward Kinematics\n J = self.get_skeleton(betas) # (*, 55, 3)\n A = batch_rigid_transform_v2(rot_mats, J, self.parents)[1]\n\n # 3. Canonical v_posed = v_template + shaped_offsets + pose_offsets\n pose_feature = rot_mats[..., 1:, :, :] - rot_mats.new([[1, 0, 0], [0, 1, 0], [0, 0, 1]])\n pose_feature = pose_feature.view(*pose_feature.shape[:-3], -1) # (*, 55*3*3)\n v_posed = (\n self.v_template\n + einsum(betas, self.shapedirs, \"... k, v c k -> ... v c\")\n + einsum(pose_feature, self.posedirs, \"... k, k v c -> ... v c\")\n )\n del pose_feature, rot_mats\n\n # 4. Skinning\n T = einsum(self.lbs_weights, A, \"v j, ... j c d -> ... v c d\")\n verts = einsum(T[..., :3, :3], v_posed, \"... v c d, ... v d -> ... v c\") + T[..., :3, 3]\n\n # 5. Translation\n if transl is not None:\n verts = verts + transl[..., None, :]\n return verts\n\n\nclass SmplxLiteCoco17(SmplxLite):\n \"\"\"Output COCO17 joints (Faster, but cannot output vertices)\"\"\"\n\n def __init__(self, **kwargs):\n super().__init__(**kwargs)\n\n # Compute mapping\n smplx2smpl = torch.load(PROJ_ROOT / \"hmr4d/utils/body_model/smplx2smpl_sparse.pt\")\n COCO17_regressor = torch.load(PROJ_ROOT / \"hmr4d/utils/body_model/smpl_coco17_J_regressor.pt\")\n smplx2coco17 = torch.matmul(COCO17_regressor, smplx2smpl.to_dense())\n\n jids, smplx_vids = torch.where(smplx2coco17 != 0)\n smplx2coco17_interestd = torch.zeros([len(smplx_vids), 17])\n for idx, (jid, smplx_vid) in enumerate(zip(jids, smplx_vids)):\n smplx2coco17_interestd[idx, jid] = smplx2coco17[jid, smplx_vid]\n self.register_buffer(\"smplx2coco17_interestd\", smplx2coco17_interestd, False) # (132, 17)\n\n # Update to vertices of interest\n self.v_template = self.v_template[smplx_vids].clone() # (V', 3)\n self.shapedirs = self.shapedirs[smplx_vids].clone() # (V', 3, K)\n self.posedirs = self.posedirs[:, smplx_vids].clone() # (K, V', 3)\n self.lbs_weights = self.lbs_weights[smplx_vids].clone() # (V', J)\n\n def forward(self, body_pose, betas, global_orient, transl):\n \"\"\"Returns: joints (*, 17, 3). (B, L) or (B,) are both supported.\"\"\"\n # Use super class's forward to get verts\n verts = super().forward(body_pose, betas, global_orient, transl) # (*, 132, 3)\n joints = einsum(self.smplx2coco17_interestd, verts, \"v j, ... v c -> ... j c\")\n return joints\n\n\nclass SmplxLiteV437Coco17(SmplxLite):\n def __init__(self, **kwargs):\n super().__init__(**kwargs)\n\n # Compute mapping (COCO17)\n smplx2smpl = torch.load(PROJ_ROOT / \"hmr4d/utils/body_model/smplx2smpl_sparse.pt\")\n COCO17_regressor = torch.load(PROJ_ROOT / \"hmr4d/utils/body_model/smpl_coco17_J_regressor.pt\")\n smplx2coco17 = torch.matmul(COCO17_regressor, smplx2smpl.to_dense())\n\n jids, smplx_vids = torch.where(smplx2coco17 != 0)\n smplx2coco17_interestd = torch.zeros([len(smplx_vids), 17])\n for idx, (jid, smplx_vid) in enumerate(zip(jids, smplx_vids)):\n smplx2coco17_interestd[idx, jid] = smplx2coco17[jid, smplx_vid]\n self.register_buffer(\"smplx2coco17_interestd\", smplx2coco17_interestd, False) # (132, 17)\n assert len(smplx_vids) == 132\n\n # Verts437\n smplx_vids2 = torch.load(PROJ_ROOT / \"hmr4d/utils/body_model/smplx_verts437.pt\")\n smplx_vids = torch.cat([smplx_vids, smplx_vids2])\n\n # Update to vertices of interest\n self.v_template = self.v_template[smplx_vids].clone() # (V', 3)\n self.shapedirs = self.shapedirs[smplx_vids].clone() # (V', 3, K)\n self.posedirs = self.posedirs[:, smplx_vids].clone() # (K, V', 3)\n self.lbs_weights = self.lbs_weights[smplx_vids].clone() # (V', J)\n\n def forward(self, body_pose, betas, global_orient, transl):\n \"\"\"\n Returns:\n verts_437: (*, 437, 3)\n joints (*, 17, 3). (B, L) or (B,) are both supported.\n \"\"\"\n # Use super class's forward to get verts\n verts = super().forward(body_pose, betas, global_orient, transl) # (*, 132+437, 3)\n\n verts_437 = verts[..., 132:, :].clone()\n joints = einsum(self.smplx2coco17_interestd, verts[..., :132, :], \"v j, ... v c -> ... j c\")\n return verts_437, joints\n\n\nclass SmplxLiteSmplN24(SmplxLite):\n \"\"\"Output SMPL(not smplx)-Neutral 24 joints (Faster, but cannot output vertices)\"\"\"\n\n def __init__(self, **kwargs):\n super().__init__(**kwargs)\n\n # Compute mapping\n smplx2smpl = torch.load(PROJ_ROOT / \"hmr4d/utils/body_model/smplx2smpl_sparse.pt\")\n smpl2joints = torch.load(PROJ_ROOT / \"hmr4d/utils/body_model/smpl_neutral_J_regressor.pt\")\n smplx2joints = torch.matmul(smpl2joints, smplx2smpl.to_dense())\n\n jids, smplx_vids = torch.where(smplx2joints != 0)\n smplx2joints_interested = torch.zeros([len(smplx_vids), smplx2joints.size(0)])\n for idx, (jid, smplx_vid) in enumerate(zip(jids, smplx_vids)):\n smplx2joints_interested[idx, jid] = smplx2joints[jid, smplx_vid]\n self.register_buffer(\"smplx2joints_interested\", smplx2joints_interested, False) # (V', J)\n\n # Update to vertices of interest\n self.v_template = self.v_template[smplx_vids].clone() # (V', 3)\n self.shapedirs = self.shapedirs[smplx_vids].clone() # (V', 3, K)\n self.posedirs = self.posedirs[:, smplx_vids].clone() # (K, V', 3)\n self.lbs_weights = self.lbs_weights[smplx_vids].clone() # (V', J)\n\n def forward(self, body_pose, betas, global_orient, transl):\n \"\"\"Returns: joints (*, J, 3). (B, L) or (B,) are both supported.\"\"\"\n # Use super class's forward to get verts\n verts = super().forward(body_pose, betas, global_orient, transl) # (*, V', 3)\n joints = einsum(self.smplx2joints_interested, verts, \"v j, ... v c -> ... j c\")\n return joints\n\n\ndef batch_rigid_transform_v2(rot_mats, joints, parents):\n \"\"\"\n Args:\n rot_mats: (*, J, 3, 3)\n joints: (*, J, 3)\n \"\"\"\n # check shape, since sometimes beta has shape=1\n rot_mats_shape_prefix = rot_mats.shape[:-3]\n if rot_mats_shape_prefix != joints.shape[:-2]:\n joints = joints.expand(*rot_mats_shape_prefix, -1, -1)\n\n rel_joints = joints.clone()\n rel_joints[..., 1:, :] -= joints[..., parents[1:], :]\n transforms_mat = torch.cat([rot_mats, rel_joints[..., :, None]], dim=-1) # (*, J, 3, 4)\n transforms_mat = F.pad(transforms_mat, [0, 0, 0, 1], value=0.0)\n transforms_mat[..., 3, 3] = 1.0 # (*, J, 4, 4)\n\n transform_chain = [transforms_mat[..., 0, :, :]]\n for i in range(1, parents.shape[0]):\n # Subtract the joint location at the rest pose\n # No need for rotation, since it's identity when at rest\n curr_res = torch.matmul(transform_chain[parents[i]], transforms_mat[..., i, :, :])\n transform_chain.append(curr_res)\n\n transforms = torch.stack(transform_chain, dim=-3) # (*, J, 4, 4)\n\n # The last column of the transformations contains the posed joints\n posed_joints = transforms[..., :3, 3].clone()\n rel_transforms = transforms.clone()\n rel_transforms[..., :3, 3] -= einsum(transforms[..., :3, :3], joints, \"... j c d, ... j d -> ... j c\")\n return posed_joints, rel_transforms\n\n\ndef sync_time():\n torch.cuda.synchronize()\n return time()\n" }, { "path": "eval/GVHMR/hmr4d/utils/body_model/utils.py", "content": "import os\nimport numpy as np\nimport torch\n\nSMPLH_JOINT_NAMES = [\n 'pelvis',\n 'left_hip',\n 'right_hip',\n 'spine1',\n 'left_knee',\n 'right_knee',\n 'spine2',\n 'left_ankle',\n 'right_ankle',\n 'spine3',\n 'left_foot',\n 'right_foot',\n 'neck',\n 'left_collar',\n 'right_collar',\n 'head',\n 'left_shoulder',\n 'right_shoulder',\n 'left_elbow',\n 'right_elbow',\n 'left_wrist',\n 'right_wrist',\n 'left_index1',\n 'left_index2',\n 'left_index3',\n 'left_middle1',\n 'left_middle2',\n 'left_middle3',\n 'left_pinky1',\n 'left_pinky2',\n 'left_pinky3',\n 'left_ring1',\n 'left_ring2',\n 'left_ring3',\n 'left_thumb1',\n 'left_thumb2',\n 'left_thumb3',\n 'right_index1',\n 'right_index2',\n 'right_index3',\n 'right_middle1',\n 'right_middle2',\n 'right_middle3',\n 'right_pinky1',\n 'right_pinky2',\n 'right_pinky3',\n 'right_ring1',\n 'right_ring2',\n 'right_ring3',\n 'right_thumb1',\n 'right_thumb2',\n 'right_thumb3',\n 'nose',\n 'right_eye',\n 'left_eye',\n 'right_ear',\n 'left_ear',\n 'left_big_toe',\n 'left_small_toe',\n 'left_heel',\n 'right_big_toe',\n 'right_small_toe',\n 'right_heel',\n 'left_thumb',\n 'left_index',\n 'left_middle',\n 'left_ring',\n 'left_pinky',\n 'right_thumb',\n 'right_index',\n 'right_middle',\n 'right_ring',\n 'right_pinky',\n]\n\nSMPLH_LEFT_LEG = ['left_hip', 'left_knee', 'left_ankle', 'left_foot']\nSMPLH_RIGHT_LEG = ['right_hip', 'right_knee', 'right_ankle', 'right_foot']\nSMPLH_LEFT_ARM = ['left_collar', 'left_shoulder', 'left_elbow', 'left_wrist']\nSMPLH_RIGHT_ARM = ['right_collar', 'right_shoulder', 'right_elbow', 'right_wrist']\nSMPLH_HEAD = ['neck', 'head']\nSMPLH_SPINE = ['spine1', 'spine2', 'spine3']\n\n# name to 21 index (without pelvis, hand, and extra)\n_name_2_idx = {j: i for i, j in enumerate(SMPLH_JOINT_NAMES[1:22])}\nSMPLH_PART_IDX = {\n 'left_leg': [_name_2_idx[x] for x in SMPLH_LEFT_LEG],\n 'right_leg': [_name_2_idx[x] for x in SMPLH_RIGHT_LEG],\n 'left_arm': [_name_2_idx[x] for x in SMPLH_LEFT_ARM],\n 'right_arm': [_name_2_idx[x] for x in SMPLH_RIGHT_ARM],\n 'two_legs': [_name_2_idx[x] for x in SMPLH_LEFT_LEG + SMPLH_RIGHT_LEG],\n 'left_arm_and_leg': [_name_2_idx[x] for x in SMPLH_LEFT_ARM + SMPLH_LEFT_LEG],\n 'right_arm_and_leg': [_name_2_idx[x] for x in SMPLH_RIGHT_ARM + SMPLH_RIGHT_LEG],\n}\n\n# name to full index\n_name_2_idx_full = {j: i for i, j in enumerate(SMPLH_JOINT_NAMES)}\nSMPLH_PART_IDX_FULL = {\n 'lower_body': [_name_2_idx_full[x] for x in ['pelvis'] + SMPLH_LEFT_LEG + SMPLH_RIGHT_LEG]\n}\n\n# ===== ⬇️ Fitting optimizer ⬇️ ===== #\nSMPL_JOINTS = {'hips': 0, 'leftUpLeg': 1, 'rightUpLeg': 2, 'spine': 3, 'leftLeg': 4, 'rightLeg': 5,\n 'spine1': 6, 'leftFoot': 7, 'rightFoot': 8, 'spine2': 9, 'leftToeBase': 10, 'rightToeBase': 11,\n 'neck': 12, 'leftShoulder': 13, 'rightShoulder': 14, 'head': 15, 'leftArm': 16, 'rightArm': 17,\n 'leftForeArm': 18, 'rightForeArm': 19, 'leftHand': 20, 'rightHand': 21}\n\n# chosen virtual mocap markers that are \"keypoints\" to work with\nKEYPT_VERTS = [4404, 920, 3076, 3169, 823, 4310, 1010, 1085, 4495, 4569, 6615, 3217, 3313, 6713,\n 6785, 3383, 6607, 3207, 1241, 1508, 4797, 4122, 1618, 1569, 5135, 5040, 5691, 5636,\n 5404, 2230, 2173, 2108, 134, 3645, 6543, 3123, 3024, 4194, 1306, 182, 3694, 4294, 744]\n\n\n# From https://github.com/vchoutas/smplify-x/blob/master/smplifyx/utils.py\n# Please see license for usage restrictions.\ndef smpl_to_openpose(model_type='smplx', use_hands=True, use_face=True,\n use_face_contour=False, openpose_format='coco25'):\n ''' Returns the indices of the permutation that maps SMPL to OpenPose\n\n Parameters\n ----------\n model_type: str, optional\n The type of SMPL-like model that is used. The default mapping\n returned is for the SMPLX model\n use_hands: bool, optional\n Flag for adding to the returned permutation the mapping for the\n hand keypoints. Defaults to True\n use_face: bool, optional\n Flag for adding to the returned permutation the mapping for the\n face keypoints. Defaults to True\n use_face_contour: bool, optional\n Flag for appending the facial contour keypoints. Defaults to False\n openpose_format: bool, optional\n The output format of OpenPose. For now only COCO-25 and COCO-19 is\n supported. Defaults to 'coco25'\n\n '''\n if openpose_format.lower() == 'coco25':\n if model_type == 'smpl':\n return np.array([24, 12, 17, 19, 21, 16, 18, 20, 0, 2, 5, 8, 1, 4,\n 7, 25, 26, 27, 28, 29, 30, 31, 32, 33, 34],\n dtype=np.int32)\n elif model_type == 'smplh':\n body_mapping = np.array([52, 12, 17, 19, 21, 16, 18, 20, 0, 2, 5,\n 8, 1, 4, 7, 53, 54, 55, 56, 57, 58, 59,\n 60, 61, 62], dtype=np.int32)\n mapping = [body_mapping]\n if use_hands:\n lhand_mapping = np.array([20, 34, 35, 36, 63, 22, 23, 24, 64,\n 25, 26, 27, 65, 31, 32, 33, 66, 28,\n 29, 30, 67], dtype=np.int32)\n rhand_mapping = np.array([21, 49, 50, 51, 68, 37, 38, 39, 69,\n 40, 41, 42, 70, 46, 47, 48, 71, 43,\n 44, 45, 72], dtype=np.int32)\n mapping += [lhand_mapping, rhand_mapping]\n return np.concatenate(mapping)\n # SMPLX\n elif model_type == 'smplx':\n body_mapping = np.array([55, 12, 17, 19, 21, 16, 18, 20, 0, 2, 5,\n 8, 1, 4, 7, 56, 57, 58, 59, 60, 61, 62,\n 63, 64, 65], dtype=np.int32)\n mapping = [body_mapping]\n if use_hands:\n lhand_mapping = np.array([20, 37, 38, 39, 66, 25, 26, 27,\n 67, 28, 29, 30, 68, 34, 35, 36, 69,\n 31, 32, 33, 70], dtype=np.int32)\n rhand_mapping = np.array([21, 52, 53, 54, 71, 40, 41, 42, 72,\n 43, 44, 45, 73, 49, 50, 51, 74, 46,\n 47, 48, 75], dtype=np.int32)\n\n mapping += [lhand_mapping, rhand_mapping]\n if use_face:\n # end_idx = 127 + 17 * use_face_contour\n face_mapping = np.arange(76, 127 + 17 * use_face_contour,\n dtype=np.int32)\n mapping += [face_mapping]\n\n return np.concatenate(mapping)\n else:\n raise ValueError('Unknown model type: {}'.format(model_type))\n elif openpose_format == 'coco19':\n if model_type == 'smpl':\n return np.array([24, 12, 17, 19, 21, 16, 18, 20, 0, 2, 5, 8,\n 1, 4, 7, 25, 26, 27, 28],\n dtype=np.int32)\n elif model_type == 'smplh':\n body_mapping = np.array([52, 12, 17, 19, 21, 16, 18, 20, 0, 2, 5,\n 8, 1, 4, 7, 53, 54, 55, 56],\n dtype=np.int32)\n mapping = [body_mapping]\n if use_hands:\n lhand_mapping = np.array([20, 34, 35, 36, 57, 22, 23, 24, 58,\n 25, 26, 27, 59, 31, 32, 33, 60, 28,\n 29, 30, 61], dtype=np.int32)\n rhand_mapping = np.array([21, 49, 50, 51, 62, 37, 38, 39, 63,\n 40, 41, 42, 64, 46, 47, 48, 65, 43,\n 44, 45, 66], dtype=np.int32)\n mapping += [lhand_mapping, rhand_mapping]\n return np.concatenate(mapping)\n # SMPLX\n elif model_type == 'smplx':\n body_mapping = np.array([55, 12, 17, 19, 21, 16, 18, 20, 0, 2, 5,\n 8, 1, 4, 7, 56, 57, 58, 59],\n dtype=np.int32)\n mapping = [body_mapping]\n if use_hands:\n lhand_mapping = np.array([20, 37, 38, 39, 60, 25, 26, 27,\n 61, 28, 29, 30, 62, 34, 35, 36, 63,\n 31, 32, 33, 64], dtype=np.int32)\n rhand_mapping = np.array([21, 52, 53, 54, 65, 40, 41, 42, 66,\n 43, 44, 45, 67, 49, 50, 51, 68, 46,\n 47, 48, 69], dtype=np.int32)\n\n mapping += [lhand_mapping, rhand_mapping]\n if use_face:\n face_mapping = np.arange(70, 70 + 51 +\n 17 * use_face_contour,\n dtype=np.int32)\n mapping += [face_mapping]\n\n return np.concatenate(mapping)\n else:\n raise ValueError('Unknown model type: {}'.format(model_type))\n else:\n raise ValueError('Unknown joint format: {}'.format(openpose_format))\n" }, { "path": "eval/GVHMR/hmr4d/utils/callbacks/lr_monitor.py", "content": "from pytorch_lightning.callbacks import LearningRateMonitor\nfrom hmr4d.configs import builds, MainStore\n\n\nMainStore.store(name=\"pl\", node=builds(LearningRateMonitor), group=\"callbacks/lr_monitor\")\n" }, { "path": "eval/GVHMR/hmr4d/utils/callbacks/prog_bar.py", "content": "from collections import OrderedDict\nfrom numbers import Number\nfrom datetime import datetime, timedelta\nfrom typing import Any, Dict, Union\nfrom pytorch_lightning.utilities.types import STEP_OUTPUT\nimport torch\nfrom pytorch_lightning.callbacks.progress.tqdm_progress import TQDMProgressBar, Tqdm, convert_inf\nfrom pytorch_lightning.callbacks.progress import ProgressBar\nfrom pytorch_lightning.utilities import rank_zero_only\nimport pytorch_lightning as pl\n\nfrom hmr4d.utils.pylogger import Log\nfrom time import time\nfrom collections import deque\nimport sys\nfrom hmr4d.configs import MainStore, builds\n\n# ========== Helper functions ========== #\n\n\ndef format_num(n):\n f = \"{0:.3g}\".format(n).replace(\"+0\", \"+\").replace(\"-0\", \"-\")\n n = str(n)\n return f if len(f) < len(n) else n\n\n\ndef convert_kwargs_to_str(**kwargs):\n # Sort in alphabetical order to be more deterministic\n postfix = OrderedDict([])\n for key in sorted(kwargs.keys()):\n new_key = key.split(\"/\")[-1]\n postfix[new_key] = kwargs[key]\n # Preprocess stats according to datatype\n for key in postfix.keys():\n # Number: limit the length of the string\n if isinstance(postfix[key], Number):\n postfix[key] = format_num(postfix[key])\n # Else for any other type, try to get the string conversion\n elif not isinstance(postfix[key], str):\n postfix[key] = str(postfix[key])\n # Else if it's a string, don't need to preprocess anything\n # Stitch together to get the final postfix\n postfix = \", \".join(key + \"=\" + postfix[key].strip() for key in postfix.keys())\n return postfix\n\n\ndef convert_t_to_str(t):\n \"\"\"Convert time in second to string in format hour:minute:second.\n If hour is 0, don't show it. Always show minute and second.\n \"\"\"\n t_str = timedelta(seconds=t) # e.g. 0:00:00.704186\n t_str = str(t_str).split(\".\")[0] # e.g. 0:00:00\n if t_str[:2] == \"0:\":\n t_str = t_str[2:]\n return t_str\n\n\nclass MyTQDMProgressBar(TQDMProgressBar, pl.Callback):\n def init_train_tqdm(self):\n bar = Tqdm(\n desc=\"Training\", # this will be overwritten anyway\n bar_format=\"{desc}{percentage:3.0f}%[{bar:10}][{n_fmt}/{total_fmt}, {elapsed}→{remaining},{rate_fmt}]{postfix}\",\n position=(2 * self.process_position),\n disable=self.is_disabled,\n leave=False,\n smoothing=0,\n dynamic_ncols=False,\n )\n return bar\n\n @rank_zero_only\n def on_train_batch_end(self, trainer, pl_module, outputs, batch, batch_idx):\n # this function also updates the main progress bar\n super().on_train_batch_end(trainer, pl_module, outputs, batch, batch_idx)\n # in this function, we only set the postfix of the main progress bar\n n = batch_idx + 1\n if self._should_update(n, self.train_progress_bar.total):\n # Set post-fix string\n # 1. maximum GPU usage\n max_mem = torch.cuda.max_memory_allocated() / 1024.0 / 1024.0 / 1024.0\n post_fix_str = f\"maxGPU={max_mem:.1f}GB\"\n\n # 2. training metrics\n training_metrics = self.get_metrics(trainer, pl_module)\n training_metrics.pop(\"v_num\", None)\n post_fix_str += \", \" + convert_kwargs_to_str(**training_metrics)\n\n # extra message if applicable\n if \"message\" in outputs:\n post_fix_str += \", \" + outputs[\"message\"]\n\n self.train_progress_bar.set_postfix_str(post_fix_str)\n\n\nclass ProgressReporter(ProgressBar, pl.Callback):\n def __init__(\n self,\n log_every_percent: float = 0.1, # report interval\n exp_name=None, # if None, use pl_module.exp_name or \"Unnamed Experiment\"\n data_name=None, # if None, use pl_module.exp_name or \"Unknown Data\"\n **kwargs,\n ):\n super().__init__()\n self.enable = True\n # 1. Store experiment meta data.\n self.log_every_percent = log_every_percent\n self.exp_name = exp_name\n self.data_name = data_name\n self.batch_time_queue = deque(maxlen=5)\n self.start_prompt = \"🚀\"\n self.finish_prompt = \"✅\"\n # 2. Utils for evaluation\n self.n_finished = 0\n\n def disable(self):\n self.enable = False\n\n def setup(self, trainer: pl.Trainer, pl_module: pl.LightningModule, stage: str) -> None:\n # Connect to the trainer object.\n super().setup(trainer, pl_module, stage)\n self.stage = stage\n self.time_exp_start = time()\n self.epoch_exp_start = trainer.current_epoch\n\n if self.exp_name is None:\n if hasattr(pl_module, \"exp_name\"):\n self.exp_name = pl_module.exp_name\n else:\n self.exp_name = \"Unnamed Experiment\"\n if self.data_name is None:\n if hasattr(pl_module, \"data_name\"):\n self.data_name = pl_module.data_name\n else:\n self.data_name = \"Unknown Data\"\n\n def print(self, *args: Any, **kwargs: Any) -> None:\n print(*args)\n\n def get_metrics(self, trainer: pl.Trainer, pl_module: pl.LightningModule) -> Dict[str, Union[str, float]]:\n \"\"\"Get metrics from trainer for progress bar.\"\"\"\n items = super().get_metrics(trainer, pl_module)\n items.pop(\"v_num\", None)\n return items\n\n def _should_update(self, n_finished: int, total: int) -> bool:\n \"\"\"\n Rule: Log every `log_every_percent` percent, or the last batch.\n \"\"\"\n log_interval = max(int(total * self.log_every_percent), 1)\n able = n_finished % log_interval == 0 or n_finished == total\n if log_interval > 10:\n able = able or n_finished in [5, 10] # always log\n able = able and self.enable\n return able\n\n @rank_zero_only\n def on_train_epoch_start(self, trainer: \"pl.Trainer\", *_: Any) -> None:\n self.print(\"=\" * 80)\n Log.info(\n f\"{self.start_prompt}[FIT][Epoch {trainer.current_epoch}] Data: {self.data_name} Experiment: {self.exp_name}\"\n )\n self.time_train_epoch_start = time()\n\n @rank_zero_only\n def on_train_batch_end(self, trainer, pl_module, outputs, batch, batch_idx):\n super().on_train_batch_end(trainer, pl_module, outputs, batch, batch_idx) # don't forget this :)\n total = self.total_train_batches\n\n # Speed\n n_finished = batch_idx + 1\n percent = 100 * n_finished / total\n time_current = time()\n self.batch_time_queue.append(time_current)\n time_elapsed = time_current - self.time_train_epoch_start # second\n time_remaining = time_elapsed * (total - n_finished) / n_finished # second\n if len(self.batch_time_queue) == 1: # cannot compute speed\n speed = 1 / time_elapsed\n else:\n speed = (len(self.batch_time_queue) - 1) / (self.batch_time_queue[-1] - self.batch_time_queue[0])\n\n # Skip if not update\n if not self._should_update(n_finished, total):\n return\n\n # ===== Set Prefix string ===== #\n # General\n desc = f\"[Train]\"\n\n # Speed: Get elapsed time and estimated remaining time\n time_elapsed_str = convert_t_to_str(time_elapsed)\n time_remaining_str = convert_t_to_str(time_remaining)\n speed_str = f\"{speed:.2f}it/s\" if speed > 1 else f\"{1/speed:.1f}s/it\"\n n_digit = len(str(total))\n desc_speed = (\n f\"[{n_finished:{n_digit}d}/{total}={percent:3.0f}%, {time_elapsed_str} → {time_remaining_str}, {speed_str}]\"\n )\n\n # ===== Set postfix string ===== #\n # 1. maximum GPU usage\n max_mem = torch.cuda.max_memory_allocated() / 1024.0 / 1024.0 / 1024.0\n post_fix_str = f\"maxGPU={max_mem:.1f}GB\"\n\n # 2. training step metrics\n train_metrics = self.get_metrics(trainer, pl_module)\n train_metrics = {k: v for k, v in train_metrics.items() if (\"train\" in k and \"epoch\" not in k)}\n post_fix_str += \", \" + convert_kwargs_to_str(**train_metrics)\n\n # extra message if applicable\n if \"message\" in outputs:\n post_fix_str += \", \" + outputs[\"message\"]\n post_fix_str = f\"[{post_fix_str}]\"\n\n # ===== Output ===== #\n bar_output = f\"{desc}{desc_speed}{post_fix_str}\"\n self.print(bar_output)\n\n @rank_zero_only\n def on_train_epoch_end(self, trainer: pl.Trainer, pl_module: pl.LightningModule) -> None:\n super().on_train_epoch_end(trainer, pl_module)\n\n # Clear\n self.batch_time_queue.clear()\n\n # Estimate Epoch time\n n_finished = trainer.current_epoch + 1 - self.epoch_exp_start\n n_to_finish = trainer.max_epochs - trainer.current_epoch - 1\n time_current = time()\n time_elapsed = time_current - self.time_exp_start\n time_remaining = time_elapsed * n_to_finish / n_finished\n time_elapsed_str = convert_t_to_str(time_elapsed)\n time_remaining_str = convert_t_to_str(time_remaining)\n\n # Metrics\n # training epoch metrics\n train_metrics = self.get_metrics(trainer, pl_module)\n train_metrics = {k: v for k, v in train_metrics.items() if (\"train\" in k and \"epoch\" in k)}\n train_metrics_str = convert_kwargs_to_str(**train_metrics)\n\n Log.info(\n f\"{self.finish_prompt}[FIT][Epoch {trainer.current_epoch}] finished! {time_elapsed_str}→{time_remaining_str} | {train_metrics_str}\"\n )\n\n # ===== Validation/Test/Prediction ===== #\n @rank_zero_only\n def on_validation_epoch_start(self, trainer, pl_module):\n self.time_val_epoch_start = time()\n\n @rank_zero_only\n def on_validation_batch_end(self, trainer, pl_module, outputs, batch, batch_idx, dataloader_idx=0):\n self.n_finished += 1\n n_finished = self.n_finished\n total = self.total_val_batches\n if not self._should_update(n_finished, total):\n return\n\n # General\n desc = f\"[Val]\"\n\n # Speed\n percent = 100 * n_finished / total\n time_current = time()\n time_elapsed = time_current - self.time_val_epoch_start # second\n time_remaining = time_elapsed * (total - n_finished) / n_finished # second\n time_elapsed_str = convert_t_to_str(time_elapsed)\n time_remaining_str = convert_t_to_str(time_remaining)\n desc_speed = f\"[{n_finished}/{total} ={percent:3.0f}%, {time_elapsed_str}→{time_remaining_str}]\"\n\n # Output\n bar_output = f\"{desc} {desc_speed}\"\n self.print(bar_output)\n\n def on_validation_epoch_end(self, trainer: pl.Trainer, pl_module: pl.LightningModule) -> None:\n # Reset\n self.n_finished = 0\n\n\nclass EmojiProgressReporter(ProgressBar, pl.Callback):\n def __init__(\n self,\n refresh_rate_batch: Union[int, None] = 1, # report interval of batch, set None to disable it\n refresh_rate_epoch: int = 1, # report interval of epoch\n **kwargs,\n ):\n super().__init__()\n self.enable = True\n # Store experiment meta data.\n self.refresh_rate_batch = refresh_rate_batch\n self.refresh_rate_epoch = refresh_rate_epoch\n\n # Style of the progress bar.\n self.title_prompt = \"📝\"\n self.prog_prompt = \"🚀\"\n self.timer_prompt = \"⌛️\"\n self.metric_prompt = \"📌\"\n self.finish_prompt = \"✅\"\n\n def disable(self):\n self.enable = False\n\n def setup(self, trainer: pl.Trainer, pl_module: pl.LightningModule, stage: str):\n # Connect to the trainer object.\n super().setup(trainer, pl_module, stage)\n self.stage = stage\n self.time_start_batch = None\n self.time_start_epoch = None\n if hasattr(pl_module, \"exp_name\"):\n self.exp_name = pl_module.exp_name\n else:\n self.exp_name = \"Unnamed Experiment\"\n Log.warn(\"Experiment name not found, please set it to `pl_module.exp_name`!\")\n\n def print(self, *args: Any, **kwargs: Any):\n print(*args)\n\n def get_metrics(self, trainer: pl.Trainer, pl_module: pl.LightningModule) -> Dict[str, Union[str, float]]:\n \"\"\"Get metrics from trainer for progress bar.\"\"\"\n items = super().get_metrics(trainer, pl_module)\n items.pop(\"v_num\", None)\n return dict(sorted(items.items()))\n\n def _should_log_batch(self, n: int) -> bool:\n # Disable batch log.\n if self.refresh_rate_batch is None:\n return False\n # Log at the first & last batch, and every `self.refresh_rate_batch` batches.\n able = n % self.refresh_rate_batch == 0 or n == self.total_train_batches - 1\n able = able and self.enable\n return able\n\n def _should_log_epoch(self, n: int) -> bool:\n # Log at the first & last epoch, and every `self.refresh_rate_epoch` epochs.\n able = n % self.refresh_rate_epoch == 0 or n == self.trainer.max_epochs - 1\n able = able and self.enable\n return able\n\n def timestamp_delta_to_str(self, timestamp_delta: float):\n \"\"\"Convert delta timestamp to string.\"\"\"\n time_rest = timedelta(seconds=timestamp_delta)\n hours, remainder = divmod(time_rest.seconds, 3600)\n minutes, seconds = divmod(remainder, 60)\n time_str = \"\"\n\n # Check if the time is valid. Note that, if `hours` is visible, then `minutes` must be visible.\n if hours <= 0:\n hours = None\n if minutes <= 0:\n minutes = None\n if seconds <= 0:\n seconds = None\n\n time_str += f\"{hours}h \" if hours is not None else \"\"\n time_str += f\"{minutes}m \" if minutes is not None else \"\"\n time_str += f\"{seconds}s\" if seconds is not None else \"\"\n return time_str\n\n @rank_zero_only\n def on_train_batch_start(self, trainer: pl.Trainer, pl_module: pl.LightningModule, batch: Any, batch_idx: int):\n super().on_train_batch_start(trainer, pl_module, batch, batch_idx)\n # Initialize some meta data.\n if self.time_start_batch is None:\n self.time_start_batch = datetime.now().timestamp()\n\n @rank_zero_only\n def on_train_batch_end(self, trainer, pl_module, outputs, batch, batch_idx):\n super().on_train_batch_end(trainer, pl_module, outputs, batch, batch_idx) # don't forget this :)\n # Get some meta data.\n epoch_idx = trainer.current_epoch\n percent = 100 * (batch_idx + 1) / (self.total_train_batches + 1)\n metrics = self.get_metrics(trainer, pl_module)\n\n # Current time.\n time_cur_stamp = datetime.now().timestamp()\n time_cur_str = datetime.fromtimestamp(time_cur_stamp).strftime(\"%m-%d %H:%M:%S\")\n # Rest time.\n time_rest_stamp = (time_cur_stamp - self.time_start_batch) * (100 - percent) / percent\n time_rest_str = self.timestamp_delta_to_str(time_rest_stamp)\n\n if not self._should_log_batch(batch_idx):\n return\n\n # Print the logs.\n self.print(f\"{self.title_prompt} [{self.stage.upper()}] Exp: {self.exp_name}...\")\n self.print(\n f\"{self.prog_prompt} Ep {epoch_idx}: {int(percent):02d}% <= [{batch_idx}/{self.total_train_batches}]\"\n )\n self.print(f\"{self.timer_prompt} Time: {time_cur_str} | Ep Rest: {time_rest_str}\")\n for k, v in metrics.items():\n self.print(f\"{self.metric_prompt} {k}: {v}\")\n self.print(\"\") # Add a blank line.\n\n def on_train_epoch_start(self, trainer: pl.Trainer, pl_module: pl.LightningModule):\n super().on_train_epoch_start(trainer, pl_module)\n # Initialize some meta data.\n self.time_start_batch = None\n if self.time_start_epoch is None:\n self.time_start_epoch = datetime.now().timestamp()\n\n @rank_zero_only\n def on_train_epoch_end(self, trainer: pl.Trainer, pl_module: pl.LightningModule):\n super().on_train_epoch_end(trainer, pl_module)\n # Get some meta data.\n epoch_idx = trainer.current_epoch\n percent = 100 * (epoch_idx + 1) / (self.trainer.max_epochs + 1)\n metrics = self.get_metrics(trainer, pl_module)\n\n # Current time.\n time_cur = datetime.now().timestamp()\n time_str = datetime.fromtimestamp(time_cur).strftime(\"%m-%d %H: %M:%S\")\n # Rest time.\n time_rest_stamp = (time_cur - self.time_start_epoch) * (100 - percent) / percent\n time_rest_str = self.timestamp_delta_to_str(time_rest_stamp)\n\n if not self._should_log_batch(epoch_idx):\n return\n\n # Print the logs.\n self.print(f\">> >> >> >>\")\n self.print(f\"{self.title_prompt} [{self.stage.upper()}] Exp: {self.exp_name}\")\n self.print(f\"{self.finish_prompt} Ep {epoch_idx} finished!\")\n self.print(f\"{self.timer_prompt} Time: {time_str} | Rest: {time_rest_str}\")\n for k, v in metrics.items():\n self.print(f\"{self.metric_prompt} {k}: {v}\")\n self.print(f\"<< << << <<\")\n self.print(\"\") # Add a blank line.\n\n\ngroup_name = \"callbacks/prog_bar\"\nprog_reporter_base = builds(\n ProgressReporter,\n log_every_percent=0.1,\n exp_name=\"${exp_name}\",\n data_name=\"${data_name}\",\n populate_full_signature=True,\n)\nMainStore.store(name=\"prog_reporter_every0.1\", node=prog_reporter_base, group=group_name)\nMainStore.store(name=\"prog_reporter_every0.2\", node=prog_reporter_base(log_every_percent=0.2), group=group_name)\n" }, { "path": "eval/GVHMR/hmr4d/utils/callbacks/simple_ckpt_saver.py", "content": "from pathlib import Path\nimport torch\nimport pytorch_lightning as pl\nfrom pytorch_lightning.callbacks.checkpoint import Checkpoint\nfrom pytorch_lightning.utilities import rank_zero_only\n\nfrom hmr4d.utils.pylogger import Log\nfrom hmr4d.configs import MainStore, builds\n\n\nclass SimpleCkptSaver(Checkpoint):\n \"\"\"\n This callback runs at the end of each training epoch.\n Check {every_n_epochs} and save at most {save_top_k} model if it is time.\n \"\"\"\n\n def __init__(\n self,\n output_dir,\n filename=\"e{epoch:03d}-s{step:06d}.ckpt\",\n save_top_k=1,\n every_n_epochs=1,\n save_last=None,\n save_weights_only=True,\n ):\n super().__init__()\n self.output_dir = Path(output_dir)\n self.filename = filename\n self.save_top_k = save_top_k\n self.every_n_epochs = every_n_epochs\n self.save_last = save_last\n self.save_weights_only = save_weights_only\n\n # Setup output dir\n if rank_zero_only.rank == 0:\n self.output_dir.mkdir(parents=True, exist_ok=True)\n Log.info(f\"[Simple Ckpt Saver]: Save to `{self.output_dir}'\")\n\n @rank_zero_only\n def on_train_epoch_end(self, trainer, pl_module):\n \"\"\"Save a checkpoint at the end of the training epoch.\"\"\"\n if self.every_n_epochs >= 1 and (trainer.current_epoch + 1) % self.every_n_epochs == 0:\n if self.save_top_k == 0:\n return\n\n # Current saved ckpts in the output_dir\n model_paths = []\n for p in sorted(list(self.output_dir.glob(\"*.ckpt\"))):\n model_paths.append(p)\n model_to_remove = model_paths[0] if len(model_paths) >= self.save_top_k else None\n\n # Save cureent checkpoint\n filepath = self.output_dir / self.filename.format(epoch=trainer.current_epoch, step=trainer.global_step)\n checkpoint = {\n \"epoch\": trainer.current_epoch,\n \"global_step\": trainer.global_step,\n \"pytorch-lightning_version\": pl.__version__,\n \"state_dict\": pl_module.state_dict(),\n }\n pl_module.on_save_checkpoint(checkpoint)\n\n if not self.save_weights_only:\n # optimizer\n optimizer_states = []\n for i, optimizer in enumerate(trainer.optimizers):\n # Rely on accelerator to dump optimizer state\n optimizer_state = trainer.strategy.optimizer_state(optimizer)\n optimizer_states.append(optimizer_state)\n checkpoint[\"optimizer_states\"] = optimizer_states\n\n # lr_scheduler\n lr_schedulers = []\n for config in trainer.lr_scheduler_configs:\n lr_schedulers.append(config.scheduler.state_dict())\n checkpoint[\"lr_schedulers\"] = lr_schedulers\n\n # trainer.strategy.checkpoint_io.save_checkpoint(checkpoint, filepath)\n torch.save(checkpoint, filepath)\n\n # Remove the earliest checkpoint\n if model_to_remove:\n trainer.strategy.remove_checkpoint(model_paths[0])\n\n\ngroup_name = \"callbacks/simple_ckpt_saver\"\nbase = builds(SimpleCkptSaver, output_dir=\"${output_dir}/checkpoints/\", populate_full_signature=True)\nMainStore.store(name=\"base\", node=base, group=group_name)\nMainStore.store(name=\"every1e\", node=base, group=group_name)\nMainStore.store(name=\"every2e\", node=base(every_n_epochs=2), group=group_name)\nMainStore.store(name=\"every5e\", node=base(every_n_epochs=5), group=group_name)\nMainStore.store(name=\"every5e_top100\", node=base(every_n_epochs=5, save_top_k=100), group=group_name)\nMainStore.store(name=\"every10e\", node=base(every_n_epochs=10), group=group_name)\nMainStore.store(name=\"every10e_top100\", node=base(every_n_epochs=10, save_top_k=100), group=group_name)\n" }, { "path": "eval/GVHMR/hmr4d/utils/callbacks/train_speed_timer.py", "content": "import pytorch_lightning as pl\nfrom pytorch_lightning.utilities import rank_zero_only\nfrom time import time\nfrom collections import deque\n\nfrom hmr4d.configs import MainStore, builds\n\n\nclass TrainSpeedTimer(pl.Callback):\n def __init__(self, N_avg=5):\n \"\"\"\n This callback times the training speed (averge over recent 5 iterations)\n 1. Data waiting time: this should be small, otherwise the data loading should be improved\n 2. Single batch time: this is the time for one batch of training (excluding data waiting)\n \"\"\"\n super().__init__()\n self.last_batch_end = None\n self.this_batch_start = None\n\n # time queues for averaging\n self.data_waiting_time_queue = deque(maxlen=N_avg)\n self.single_batch_time_queue = deque(maxlen=N_avg)\n\n @rank_zero_only\n def on_train_batch_start(self, trainer, pl_module, batch, batch_idx):\n \"\"\"Count the time of data waiting\"\"\"\n if self.last_batch_end is not None:\n # This should be small, otherwise the data loading should be improved\n data_waiting = time() - self.last_batch_end\n\n # Average the time\n self.data_waiting_time_queue.append(data_waiting)\n average_time = sum(self.data_waiting_time_queue) / len(self.data_waiting_time_queue)\n\n # Log to prog-bar\n pl_module.log(\n \"train_timer/data_waiting\", average_time, on_step=True, on_epoch=False, prog_bar=True, logger=True\n )\n\n self.this_batch_start = time()\n\n @rank_zero_only\n def on_train_batch_end(self, trainer, pl_module, outputs, batch, batch_idx):\n # Effective training time elapsed (excluding data waiting)\n single_batch = time() - self.this_batch_start\n\n # Average the time\n self.single_batch_time_queue.append(single_batch)\n average_time = sum(self.single_batch_time_queue) / len(self.single_batch_time_queue)\n\n # Log iter time\n pl_module.log(\n \"train_timer/single_batch\", average_time, on_step=True, on_epoch=False, prog_bar=False, logger=True\n )\n\n # Set timer for counting data waiting\n self.last_batch_end = time()\n\n @rank_zero_only\n def on_train_epoch_end(self, trainer, pl_module):\n # Reset the timer\n self.last_batch_end = None\n self.this_batch_start = None\n # Clear the queue\n self.data_waiting_time_queue.clear()\n self.single_batch_time_queue.clear()\n\n\ngroup_name = \"callbacks/train_speed_timer\"\nbase = builds(TrainSpeedTimer, populate_full_signature=True)\nMainStore.store(name=\"base\", node=base, group=group_name)\n" }, { "path": "eval/GVHMR/hmr4d/utils/comm/gather.py", "content": "# Copyright (c) Facebook, Inc. and its affiliates. All Rights Reserved\n\"\"\"\n[Copied from detectron2]\nThis file contains primitives for multi-gpu communication.\nThis is useful when doing distributed training.\n\"\"\"\n\nimport functools\nimport logging\nimport numpy as np\nimport pickle\nimport torch\nimport torch.distributed as dist\n\n_LOCAL_PROCESS_GROUP = None\n\"\"\"\nA torch process group which only includes processes that on the same machine as the current process.\nThis variable is set when processes are spawned by `launch()` in \"engine/launch.py\".\n\"\"\"\n\n\ndef get_world_size() -> int:\n if not dist.is_available():\n return 1\n if not dist.is_initialized():\n return 1\n return dist.get_world_size()\n\n\ndef get_rank() -> int:\n if not dist.is_available():\n return 0\n if not dist.is_initialized():\n return 0\n return dist.get_rank()\n\n\ndef get_local_rank() -> int:\n \"\"\"\n Returns:\n The rank of the current process within the local (per-machine) process group.\n \"\"\"\n if not dist.is_available():\n return 0\n if not dist.is_initialized():\n return 0\n assert _LOCAL_PROCESS_GROUP is not None\n return dist.get_rank(group=_LOCAL_PROCESS_GROUP)\n\n\ndef get_local_size() -> int:\n \"\"\"\n Returns:\n The size of the per-machine process group,\n i.e. the number of processes per machine.\n \"\"\"\n if not dist.is_available():\n return 1\n if not dist.is_initialized():\n return 1\n return dist.get_world_size(group=_LOCAL_PROCESS_GROUP)\n\n\ndef is_main_process() -> bool:\n return get_rank() == 0\n\n\ndef synchronize():\n \"\"\"\n Helper function to synchronize (barrier) among all processes when\n using distributed training\n \"\"\"\n if not dist.is_available():\n return\n if not dist.is_initialized():\n return\n world_size = dist.get_world_size()\n if world_size == 1:\n return\n dist.barrier()\n\n\n@functools.lru_cache()\ndef _get_global_gloo_group():\n \"\"\"\n Return a process group based on gloo backend, containing all the ranks\n The result is cached.\n \"\"\"\n if dist.get_backend() == \"nccl\":\n return dist.new_group(backend=\"gloo\")\n else:\n return dist.group.WORLD\n\n\ndef _serialize_to_tensor(data, group):\n backend = dist.get_backend(group)\n assert backend in [\"gloo\", \"nccl\"]\n device = torch.device(\"cpu\" if backend == \"gloo\" else \"cuda\")\n\n buffer = pickle.dumps(data)\n if len(buffer) > 1024**3:\n logger = logging.getLogger(__name__)\n logger.warning(\n \"Rank {} trying to all-gather {:.2f} GB of data on device {}\".format(\n get_rank(), len(buffer) / (1024**3), device\n )\n )\n storage = torch.ByteStorage.from_buffer(buffer)\n tensor = torch.ByteTensor(storage).to(device=device)\n return tensor\n\n\ndef _pad_to_largest_tensor(tensor, group):\n \"\"\"\n Returns:\n list[int]: size of the tensor, on each rank\n Tensor: padded tensor that has the max size\n \"\"\"\n world_size = dist.get_world_size(group=group)\n assert world_size >= 1, \"comm.gather/all_gather must be called from ranks within the given group!\"\n local_size = torch.tensor([tensor.numel()], dtype=torch.int64, device=tensor.device)\n size_list = [torch.zeros([1], dtype=torch.int64, device=tensor.device) for _ in range(world_size)]\n dist.all_gather(size_list, local_size, group=group)\n\n size_list = [int(size.item()) for size in size_list]\n\n max_size = max(size_list)\n\n # we pad the tensor because torch all_gather does not support\n # gathering tensors of different shapes\n if local_size != max_size:\n padding = torch.zeros((max_size - local_size,), dtype=torch.uint8, device=tensor.device)\n tensor = torch.cat((tensor, padding), dim=0)\n return size_list, tensor\n\n\ndef all_gather(data, group=None):\n \"\"\"\n Run all_gather on arbitrary picklable data (not necessarily tensors).\n\n Args:\n data: any picklable object\n group: a torch process group. By default, will use a group which\n contains all ranks on gloo backend.\n\n Returns:\n list[data]: list of data gathered from each rank\n \"\"\"\n if get_world_size() == 1:\n return [data]\n if group is None:\n group = _get_global_gloo_group()\n if dist.get_world_size(group) == 1:\n return [data]\n\n tensor = _serialize_to_tensor(data, group)\n\n size_list, tensor = _pad_to_largest_tensor(tensor, group)\n max_size = max(size_list)\n\n # receiving Tensor from all ranks\n tensor_list = [torch.empty((max_size,), dtype=torch.uint8, device=tensor.device) for _ in size_list]\n dist.all_gather(tensor_list, tensor, group=group)\n\n data_list = []\n for size, tensor in zip(size_list, tensor_list):\n buffer = tensor.cpu().numpy().tobytes()[:size]\n data_list.append(pickle.loads(buffer))\n\n return data_list\n\n\ndef gather(data, dst=0, group=None):\n \"\"\"\n Run gather on arbitrary picklable data (not necessarily tensors).\n\n Args:\n data: any picklable object\n dst (int): destination rank\n group: a torch process group. By default, will use a group which\n contains all ranks on gloo backend.\n\n Returns:\n list[data]: on dst, a list of data gathered from each rank. Otherwise,\n an empty list.\n \"\"\"\n if get_world_size() == 1:\n return [data]\n if group is None:\n group = _get_global_gloo_group()\n if dist.get_world_size(group=group) == 1:\n return [data]\n rank = dist.get_rank(group=group)\n\n tensor = _serialize_to_tensor(data, group)\n size_list, tensor = _pad_to_largest_tensor(tensor, group)\n\n # receiving Tensor from all ranks\n if rank == dst:\n max_size = max(size_list)\n tensor_list = [torch.empty((max_size,), dtype=torch.uint8, device=tensor.device) for _ in size_list]\n dist.gather(tensor, tensor_list, dst=dst, group=group)\n\n data_list = []\n for size, tensor in zip(size_list, tensor_list):\n buffer = tensor.cpu().numpy().tobytes()[:size]\n data_list.append(pickle.loads(buffer))\n return data_list\n else:\n dist.gather(tensor, [], dst=dst, group=group)\n return []\n\n\ndef shared_random_seed():\n \"\"\"\n Returns:\n int: a random number that is the same across all workers.\n If workers need a shared RNG, they can use this shared seed to\n create one.\n\n All workers must call this function, otherwise it will deadlock.\n \"\"\"\n ints = np.random.randint(2**31)\n all_ints = all_gather(ints)\n return all_ints[0]\n\n\ndef reduce_dict(input_dict, average=True):\n \"\"\"\n Reduce the values in the dictionary from all processes so that process with rank\n 0 has the reduced results.\n\n Args:\n input_dict (dict): inputs to be reduced. All the values must be scalar CUDA Tensor.\n average (bool): whether to do average or sum\n\n Returns:\n a dict with the same keys as input_dict, after reduction.\n \"\"\"\n world_size = get_world_size()\n if world_size < 2:\n return input_dict\n with torch.no_grad():\n names = []\n values = []\n # sort the keys so that they are consistent across processes\n for k in sorted(input_dict.keys()):\n names.append(k)\n values.append(input_dict[k])\n values = torch.stack(values, dim=0)\n dist.reduce(values, dst=0)\n if dist.get_rank() == 0 and average:\n # only main process gets accumulated, so only divide by\n # world_size in this case\n values /= world_size\n reduced_dict = {k: v for k, v in zip(names, values)}\n return reduced_dict\n" }, { "path": "eval/GVHMR/hmr4d/utils/eval/eval_utils.py", "content": "import torch\nimport numpy as np\n\n\n@torch.no_grad()\ndef compute_camcoord_metrics(batch, pelvis_idxs=[1, 2], fps=30, mask=None):\n \"\"\"\n Args:\n batch (dict): {\n \"pred_j3d\": (..., J, 3) tensor\n \"target_j3d\":\n \"pred_verts\":\n \"target_verts\":\n }\n Returns:\n cam_coord_metrics (dict): {\n \"pa_mpjpe\": (..., ) numpy array\n \"mpjpe\":\n \"pve\":\n \"accel\":\n }\n \"\"\"\n # All data is in camera coordinates\n pred_j3d = batch[\"pred_j3d\"].cpu() # (..., J, 3)\n target_j3d = batch[\"target_j3d\"].cpu()\n pred_verts = batch[\"pred_verts\"].cpu()\n target_verts = batch[\"target_verts\"].cpu()\n\n if mask is not None:\n mask = mask.cpu()\n pred_j3d = pred_j3d[mask].clone()\n target_j3d = target_j3d[mask].clone()\n pred_verts = pred_verts[mask].clone()\n target_verts = target_verts[mask].clone()\n assert \"mask\" not in batch\n\n # Align by pelvis\n pred_j3d, target_j3d, pred_verts, target_verts = batch_align_by_pelvis(\n [pred_j3d, target_j3d, pred_verts, target_verts], pelvis_idxs=pelvis_idxs\n )\n\n # Metrics\n m2mm = 1000\n S1_hat = batch_compute_similarity_transform_torch(pred_j3d, target_j3d)\n pa_mpjpe = compute_jpe(S1_hat, target_j3d) * m2mm\n mpjpe = compute_jpe(pred_j3d, target_j3d) * m2mm\n pve = compute_jpe(pred_verts, target_verts) * m2mm\n accel = compute_error_accel(joints_pred=pred_j3d, joints_gt=target_j3d, fps=fps)\n\n camcoord_metrics = {\n \"pa_mpjpe\": pa_mpjpe,\n \"mpjpe\": mpjpe,\n \"pve\": pve,\n \"accel\": accel,\n }\n return camcoord_metrics\n\n\n@torch.no_grad()\ndef compute_global_metrics(batch, mask=None):\n \"\"\"Follow WHAM, the input has skipped invalid frames\n Args:\n batch (dict): {\n \"pred_j3d_glob\": (F, J, 3) tensor\n \"target_j3d_glob\":\n \"pred_verts_glob\":\n \"target_verts_glob\":\n }\n Returns:\n global_metrics (dict): {\n \"wa2_mpjpe\": (F, ) numpy array\n \"waa_mpjpe\":\n \"rte\":\n \"jitter\":\n \"fs\":\n }\n \"\"\"\n # All data is in global coordinates\n pred_j3d_glob = batch[\"pred_j3d_glob\"].cpu() # (..., J, 3)\n target_j3d_glob = batch[\"target_j3d_glob\"].cpu()\n pred_verts_glob = batch[\"pred_verts_glob\"].cpu()\n target_verts_glob = batch[\"target_verts_glob\"].cpu()\n if mask is not None:\n mask = mask.cpu()\n pred_j3d_glob = pred_j3d_glob[mask].clone()\n target_j3d_glob = target_j3d_glob[mask].clone()\n pred_verts_glob = pred_verts_glob[mask].clone()\n target_verts_glob = target_verts_glob[mask].clone()\n assert \"mask\" not in batch\n\n seq_length = pred_j3d_glob.shape[0]\n\n # Use chunk to compare\n chunk_length = 100\n wa2_mpjpe, waa_mpjpe = [], []\n for start in range(0, seq_length, chunk_length):\n end = min(seq_length, start + chunk_length)\n\n target_j3d = target_j3d_glob[start:end].clone().cpu()\n pred_j3d = pred_j3d_glob[start:end].clone().cpu()\n\n w_j3d = first_align_joints(target_j3d, pred_j3d)\n wa_j3d = global_align_joints(target_j3d, pred_j3d)\n\n if False:\n from hmr4d.utils.wis3d_utils import make_wis3d, add_motion_as_lines\n\n wis3d = make_wis3d(name=\"debug-metric_utils\")\n add_motion_as_lines(target_j3d, wis3d, name=\"target_j3d\")\n add_motion_as_lines(pred_j3d, wis3d, name=\"pred_j3d\")\n add_motion_as_lines(w_j3d, wis3d, name=\"pred_w2_j3d\")\n add_motion_as_lines(wa_j3d, wis3d, name=\"pred_wa_j3d\")\n\n wa2_mpjpe.append(compute_jpe(target_j3d, w_j3d))\n waa_mpjpe.append(compute_jpe(target_j3d, wa_j3d))\n\n # Metrics\n m2mm = 1000\n wa2_mpjpe = np.concatenate(wa2_mpjpe) * m2mm\n waa_mpjpe = np.concatenate(waa_mpjpe) * m2mm\n\n # Additional Metrics\n rte = compute_rte(target_j3d_glob[:, 0].cpu(), pred_j3d_glob[:, 0].cpu()) * 1e2\n jitter = compute_jitter(pred_j3d_glob, fps=30)\n foot_sliding = compute_foot_sliding(target_verts_glob, pred_verts_glob) * m2mm\n\n global_metrics = {\n \"wa2_mpjpe\": wa2_mpjpe,\n \"waa_mpjpe\": waa_mpjpe,\n \"rte\": rte,\n \"jitter\": jitter,\n \"fs\": foot_sliding,\n }\n return global_metrics\n\n\n@torch.no_grad()\ndef compute_camcoord_perjoint_metrics(batch, pelvis_idxs=[1, 2]):\n \"\"\"\n Args:\n batch (dict): {\n \"pred_j3d\": (..., J, 3) tensor\n \"target_j3d\":\n }\n Returns:\n cam_coord_metrics (dict): {\n \"pa_mpjpe\": (..., ) numpy array\n \"mpjpe\":\n \"pve\":\n \"accel\":\n }\n \"\"\"\n # All data is in camera coordinates\n pred_j3d = batch[\"pred_j3d\"].cpu() # (..., J, 3)\n target_j3d = batch[\"target_j3d\"].cpu()\n pred_verts = batch[\"pred_verts\"].cpu()\n target_verts = batch[\"target_verts\"].cpu()\n\n # Align by pelvis\n pred_j3d, target_j3d, pred_verts, target_verts = batch_align_by_pelvis(\n [pred_j3d, target_j3d, pred_verts, target_verts], pelvis_idxs=pelvis_idxs\n )\n # Metrics\n m2mm = 1000\n perjoint_mpjpe = compute_perjoint_jpe(pred_j3d, target_j3d) * m2mm\n\n camcoord_perjoint_metrics = {\n \"mpjpe\": perjoint_mpjpe,\n }\n return camcoord_perjoint_metrics\n\n\n# ===== Utilities =====\n\n\ndef compute_jpe(S1, S2):\n return torch.sqrt(((S1 - S2) ** 2).sum(dim=-1)).mean(dim=-1).numpy()\n\n\ndef compute_perjoint_jpe(S1, S2):\n return torch.sqrt(((S1 - S2) ** 2).sum(dim=-1)).numpy()\n\n\ndef batch_align_by_pelvis(data_list, pelvis_idxs=[1, 2]):\n \"\"\"\n Assumes data is given as [pred_j3d, target_j3d, pred_verts, target_verts].\n Each data is in shape of (frames, num_points, 3)\n Pelvis is notated as one / two joints indices.\n Align all data to the corresponding pelvis location.\n \"\"\"\n\n pred_j3d, target_j3d, pred_verts, target_verts = data_list\n\n pred_pelvis = pred_j3d[:, pelvis_idxs].mean(dim=1, keepdims=True).clone()\n target_pelvis = target_j3d[:, pelvis_idxs].mean(dim=1, keepdims=True).clone()\n\n # Align to the pelvis\n pred_j3d = pred_j3d - pred_pelvis\n target_j3d = target_j3d - target_pelvis\n pred_verts = pred_verts - pred_pelvis\n target_verts = target_verts - target_pelvis\n\n return (pred_j3d, target_j3d, pred_verts, target_verts)\n\n\ndef batch_compute_similarity_transform_torch(S1, S2):\n \"\"\"\n Computes a similarity transform (sR, t) that takes\n a set of 3D points S1 (3 x N) closest to a set of 3D points S2,\n where R is an 3x3 rotation matrix, t 3x1 translation, s scale.\n i.e. solves the orthogonal Procrutes problem.\n \"\"\"\n transposed = False\n if S1.shape[0] != 3 and S1.shape[0] != 2:\n S1 = S1.permute(0, 2, 1)\n S2 = S2.permute(0, 2, 1)\n transposed = True\n assert S2.shape[1] == S1.shape[1]\n\n # 1. Remove mean.\n mu1 = S1.mean(axis=-1, keepdims=True)\n mu2 = S2.mean(axis=-1, keepdims=True)\n\n X1 = S1 - mu1\n X2 = S2 - mu2\n\n # 2. Compute variance of X1 used for scale.\n var1 = torch.sum(X1**2, dim=1).sum(dim=1)\n\n # 3. The outer product of X1 and X2.\n K = X1.bmm(X2.permute(0, 2, 1))\n\n # 4. Solution that Maximizes trace(R'K) is R=U*V', where U, V are\n # singular vectors of K.\n U, s, V = torch.svd(K)\n\n # Construct Z that fixes the orientation of R to get det(R)=1.\n Z = torch.eye(U.shape[1], device=S1.device).unsqueeze(0)\n Z = Z.repeat(U.shape[0], 1, 1)\n Z[:, -1, -1] *= torch.sign(torch.det(U.bmm(V.permute(0, 2, 1))))\n\n # Construct R.\n R = V.bmm(Z.bmm(U.permute(0, 2, 1)))\n\n # 5. Recover scale.\n scale = torch.cat([torch.trace(x).unsqueeze(0) for x in R.bmm(K)]) / var1\n\n # 6. Recover translation.\n t = mu2 - (scale.unsqueeze(-1).unsqueeze(-1) * (R.bmm(mu1)))\n\n # 7. Error:\n S1_hat = scale.unsqueeze(-1).unsqueeze(-1) * R.bmm(S1) + t\n\n if transposed:\n S1_hat = S1_hat.permute(0, 2, 1)\n\n return S1_hat\n\n\ndef compute_error_accel(joints_gt, joints_pred, valid_mask=None, fps=None):\n \"\"\"\n Use [i-1, i, i+1] to compute acc at frame_i. The acceleration error:\n 1/(n-2) \\sum_{i=1}^{n-1} X_{i-1} - 2X_i + X_{i+1}\n Note that for each frame that is not visible, three entries(-1, 0, +1) in the\n acceleration error will be zero'd out.\n Args:\n joints_gt : (F, J, 3)\n joints_pred : (F, J, 3)\n valid_mask : (F)\n Returns:\n error_accel (F-2) when valid_mask is None, else (F'), F' <= F-2\n \"\"\"\n # (F, J, 3) -> (F-2) per-joint\n accel_gt = joints_gt[:-2] - 2 * joints_gt[1:-1] + joints_gt[2:]\n accel_pred = joints_pred[:-2] - 2 * joints_pred[1:-1] + joints_pred[2:]\n normed = np.linalg.norm(accel_pred - accel_gt, axis=-1).mean(axis=-1)\n if fps is not None:\n normed = normed * fps**2\n\n if valid_mask is None:\n new_vis = np.ones(len(normed), dtype=bool)\n else:\n invis = np.logical_not(valid_mask)\n invis1 = np.roll(invis, -1)\n invis2 = np.roll(invis, -2)\n new_invis = np.logical_or(invis, np.logical_or(invis1, invis2))[:-2]\n new_vis = np.logical_not(new_invis)\n if new_vis.sum() == 0:\n print(\"Warning!!! no valid acceleration error to compute.\")\n\n return normed[new_vis]\n\n\ndef compute_rte(target_trans, pred_trans):\n # Compute the global alignment\n _, rot, trans = align_pcl(target_trans[None, :], pred_trans[None, :], fixed_scale=True)\n pred_trans_hat = (torch.einsum(\"tij,tnj->tni\", rot, pred_trans[None, :]) + trans[None, :])[0]\n\n # Compute the entire displacement of ground truth trajectory\n disps, disp = [], 0\n for p1, p2 in zip(target_trans, target_trans[1:]):\n delta = (p2 - p1).norm(2, dim=-1)\n disp += delta\n disps.append(disp)\n\n # Compute absolute root-translation-error (RTE)\n rte = torch.norm(target_trans - pred_trans_hat, 2, dim=-1)\n\n # Normalize it to the displacement\n return (rte / disp).numpy()\n\n\ndef compute_jitter(joints, fps=30):\n \"\"\"compute jitter of the motion\n Args:\n joints (N, J, 3).\n fps (float).\n Returns:\n jitter (N-3).\n \"\"\"\n pred_jitter = torch.norm(\n (joints[3:] - 3 * joints[2:-1] + 3 * joints[1:-2] - joints[:-3]) * (fps**3),\n dim=2,\n ).mean(dim=-1)\n\n return pred_jitter.cpu().numpy() / 10.0\n\n\ndef compute_foot_sliding(target_verts, pred_verts, thr=1e-2):\n \"\"\"compute foot sliding error\n The foot ground contact label is computed by the threshold of 1 cm/frame\n Args:\n target_verts (N, 6890, 3).\n pred_verts (N, 6890, 3).\n Returns:\n error (N frames in contact).\n \"\"\"\n assert target_verts.shape == pred_verts.shape\n assert target_verts.shape[-2] == 6890\n\n # Foot vertices idxs\n foot_idxs = [3216, 3387, 6617, 6787]\n\n # Compute contact label\n foot_loc = target_verts[:, foot_idxs]\n foot_disp = (foot_loc[1:] - foot_loc[:-1]).norm(2, dim=-1)\n contact = foot_disp[:] < thr\n\n pred_feet_loc = pred_verts[:, foot_idxs]\n pred_disp = (pred_feet_loc[1:] - pred_feet_loc[:-1]).norm(2, dim=-1)\n\n error = pred_disp[contact]\n\n return error.cpu().numpy()\n\n\ndef convert_joints22_to_24(joints22, ratio2220=0.3438, ratio2321=0.3345):\n joints24 = torch.zeros(*joints22.shape[:-2], 24, 3).to(joints22.device)\n joints24[..., :22, :] = joints22\n joints24[..., 22, :] = joints22[..., 20, :] + ratio2220 * (joints22[..., 20, :] - joints22[..., 18, :])\n joints24[..., 23, :] = joints22[..., 21, :] + ratio2321 * (joints22[..., 21, :] - joints22[..., 19, :])\n return joints24\n\n\ndef align_pcl(Y, X, weight=None, fixed_scale=False):\n \"\"\"align similarity transform to align X with Y using umeyama method\n X' = s * R * X + t is aligned with Y\n :param Y (*, N, 3) first trajectory\n :param X (*, N, 3) second trajectory\n :param weight (*, N, 1) optional weight of valid correspondences\n :returns s (*, 1), R (*, 3, 3), t (*, 3)\n \"\"\"\n *dims, N, _ = Y.shape\n N = torch.ones(*dims, 1, 1) * N\n\n if weight is not None:\n Y = Y * weight\n X = X * weight\n N = weight.sum(dim=-2, keepdim=True) # (*, 1, 1)\n\n # subtract mean\n my = Y.sum(dim=-2) / N[..., 0] # (*, 3)\n mx = X.sum(dim=-2) / N[..., 0]\n y0 = Y - my[..., None, :] # (*, N, 3)\n x0 = X - mx[..., None, :]\n\n if weight is not None:\n y0 = y0 * weight\n x0 = x0 * weight\n\n # correlation\n C = torch.matmul(y0.transpose(-1, -2), x0) / N # (*, 3, 3)\n U, D, Vh = torch.linalg.svd(C) # (*, 3, 3), (*, 3), (*, 3, 3)\n\n S = torch.eye(3).reshape(*(1,) * (len(dims)), 3, 3).repeat(*dims, 1, 1)\n neg = torch.det(U) * torch.det(Vh.transpose(-1, -2)) < 0\n S[neg, 2, 2] = -1\n\n R = torch.matmul(U, torch.matmul(S, Vh)) # (*, 3, 3)\n\n D = torch.diag_embed(D) # (*, 3, 3)\n if fixed_scale:\n s = torch.ones(*dims, 1, device=Y.device, dtype=torch.float32)\n else:\n var = torch.sum(torch.square(x0), dim=(-1, -2), keepdim=True) / N # (*, 1, 1)\n s = torch.diagonal(torch.matmul(D, S), dim1=-2, dim2=-1).sum(dim=-1, keepdim=True) / var[..., 0] # (*, 1)\n\n t = my - s * torch.matmul(R, mx[..., None])[..., 0] # (*, 3)\n\n return s, R, t\n\n\ndef global_align_joints(gt_joints, pred_joints):\n \"\"\"\n :param gt_joints (T, J, 3)\n :param pred_joints (T, J, 3)\n \"\"\"\n s_glob, R_glob, t_glob = align_pcl(gt_joints.reshape(-1, 3), pred_joints.reshape(-1, 3))\n pred_glob = s_glob * torch.einsum(\"ij,tnj->tni\", R_glob, pred_joints) + t_glob[None, None]\n return pred_glob\n\n\ndef first_align_joints(gt_joints, pred_joints):\n \"\"\"\n align the first two frames\n :param gt_joints (T, J, 3)\n :param pred_joints (T, J, 3)\n \"\"\"\n # (1, 1), (1, 3, 3), (1, 3)\n s_first, R_first, t_first = align_pcl(gt_joints[:2].reshape(1, -1, 3), pred_joints[:2].reshape(1, -1, 3))\n pred_first = s_first * torch.einsum(\"tij,tnj->tni\", R_first, pred_joints) + t_first[:, None]\n return pred_first\n\n\ndef rearrange_by_mask(x, mask):\n \"\"\"\n x (L, *)\n mask (M,), M >= L\n \"\"\"\n M = mask.size(0)\n L = x.size(0)\n if M == L:\n return x\n assert M > L\n assert mask.sum() == L\n x_rearranged = torch.zeros((M, *x.size()[1:]), dtype=x.dtype, device=x.device)\n x_rearranged[mask] = x\n return x_rearranged\n\n\ndef as_np_array(d):\n if isinstance(d, torch.Tensor):\n return d.cpu().numpy()\n elif isinstance(d, np.ndarray):\n return d\n else:\n return np.array(d)\n" }, { "path": "eval/GVHMR/hmr4d/utils/geo/augment_noisy_pose.py", "content": "import torch\nfrom pytorch3d.transforms import axis_angle_to_matrix, matrix_to_axis_angle, matrix_to_rotation_6d\nimport hmr4d.utils.matrix as matrix\nfrom hmr4d import PROJ_ROOT\n\nCOCO17_AUG = {k: v.flatten() for k, v in torch.load(PROJ_ROOT / \"hmr4d/utils/body_model/coco_aug_dict.pth\").items()}\nCOCO17_AUG_CUDA = {}\nCOCO17_TREE = [[5, 6], 0, 0, 1, 2, -1, -1, 5, 6, 7, 8, -1, -1, 11, 12, 13, 14, 15, 15, 15, 16, 16, 16]\n\n\ndef gaussian_augment(body_pose, std_angle=10.0, to_R=True):\n \"\"\"\n Args:\n body_pose torch.Tensor: (..., J, 3) axis-angle if to_R is True, else rotmat (..., J, 3, 3)\n std_angle: scalar or list, in degree\n \"\"\"\n\n body_pose = body_pose.clone()\n\n if to_R:\n body_pose_R = axis_angle_to_matrix(body_pose) # (B, L, J, 3, 3)\n else:\n body_pose_R = body_pose\n shape = body_pose_R.shape[:-2]\n device = body_pose.device\n\n # 1. Simulate noise\n # angle:\n std_angle = torch.tensor(std_angle).to(device).reshape(-1) # allow scalar or list\n noise_angle = torch.randn(shape, device=device) * std_angle * torch.pi / 180\n\n # axis: avoid zero vector\n noise_axis = torch.rand((*shape, 3), device=device)\n mask_ = torch.norm(noise_axis, dim=-1) < 1e-6\n noise_axis[mask_] = 1\n\n noise_axis = noise_axis / torch.norm(noise_axis, dim=-1, keepdim=True)\n noise_aa = noise_angle[..., None] * noise_axis # (B, L, J, 3)\n noise_R = axis_angle_to_matrix(noise_aa) # (B, L, J, 3, 3)\n\n # 2. Add noise to body pose\n new_body_pose_R = matrix.get_mat_BfromA(body_pose_R, noise_R) # (B, L, J, 3, 3)\n # new_body_pose_R = torch.matmul(noise_R, body_pose_R)\n new_body_pose_r6d = matrix_to_rotation_6d(new_body_pose_R) # (B, L, J, 6)\n new_body_pose_aa = matrix_to_axis_angle(new_body_pose_R) # (B, L, J, 3)\n\n return new_body_pose_R, new_body_pose_r6d, new_body_pose_aa\n\n\n# ========= Augment Joint 3D ======== #\n\n\ndef get_jitter(shape=(8, 120), s_jittering=5e-2):\n \"\"\"Guassian jitter modeling.\"\"\"\n jittering_noise = (\n torch.normal(\n mean=torch.zeros((*shape, 17, 3)),\n std=COCO17_AUG[\"jittering\"].reshape(1, 1, 17, 1).expand(*shape, -1, 3),\n )\n * s_jittering\n )\n return jittering_noise\n\n\ndef get_jitter_cuda(shape=(8, 120), s_jittering=5e-2):\n if \"jittering\" not in COCO17_AUG_CUDA:\n COCO17_AUG_CUDA[\"jittering\"] = COCO17_AUG[\"jittering\"].cuda().reshape(1, 1, 17, 1)\n jittering = COCO17_AUG_CUDA[\"jittering\"]\n jittering_noise = torch.randn((*shape, 17, 3), device=\"cuda\") * jittering * s_jittering\n return jittering_noise\n\n\ndef get_lfhp(shape=(8, 120), s_peak=3e-1, s_peak_mask=5e-3):\n \"\"\"Low-frequency high-peak noise modeling.\"\"\"\n\n def get_peak_noise_mask():\n peak_noise_mask = torch.rand(*shape, 17) * COCO17_AUG[\"pmask\"]\n peak_noise_mask = peak_noise_mask < s_peak_mask\n return peak_noise_mask\n\n peak_noise_mask = get_peak_noise_mask() # (B, L, 17)\n peak_noise = peak_noise_mask.float().unsqueeze(-1).repeat(1, 1, 1, 3)\n peak_noise = peak_noise * torch.randn(3) * COCO17_AUG[\"peak\"].reshape(17, 1) * s_peak\n return peak_noise\n\n\ndef get_lfhp_cuda(shape=(8, 120), s_peak=3e-1, s_peak_mask=5e-3):\n if \"peak\" not in COCO17_AUG_CUDA:\n COCO17_AUG_CUDA[\"pmask\"] = COCO17_AUG[\"pmask\"].cuda()\n COCO17_AUG_CUDA[\"peak\"] = COCO17_AUG[\"peak\"].cuda().reshape(17, 1)\n\n pmask = COCO17_AUG_CUDA[\"pmask\"]\n peak = COCO17_AUG_CUDA[\"peak\"]\n peak_noise_mask = torch.rand(*shape, 17, device=\"cuda\") * pmask < s_peak_mask\n peak_noise = (\n peak_noise_mask.float().unsqueeze(-1).expand(-1, -1, -1, 3) * torch.randn(3, device=\"cuda\") * peak * s_peak\n )\n return peak_noise\n\n\ndef get_bias(shape=(8, 120), s_bias=1e-1):\n \"\"\"Bias noise modeling.\"\"\"\n b, l = shape\n bias_noise = torch.normal(mean=torch.zeros((b, 17, 3)), std=COCO17_AUG[\"bias\"].reshape(1, 17, 1)) * s_bias\n bias_noise = bias_noise[:, None].expand(-1, l, -1, -1) # (B, L, J, 3), the whole sequence is moved by the same bias\n return bias_noise\n\n\ndef get_bias_cuda(shape=(8, 120), s_bias=1e-1):\n if \"bias\" not in COCO17_AUG_CUDA:\n COCO17_AUG_CUDA[\"bias\"] = COCO17_AUG[\"bias\"].cuda().reshape(1, 17, 1)\n\n bias = COCO17_AUG_CUDA[\"bias\"]\n bias_noise = torch.randn((shape[0], 17, 3), device=\"cuda\") * bias * s_bias\n bias_noise = bias_noise[:, None].expand(-1, shape[1], -1, -1)\n return bias_noise\n\n\ndef get_wham_aug_kp3d(shape=(8, 120)):\n # aug = get_bias(shape).cuda() + get_lfhp(shape).cuda() + get_jitter(shape).cuda()\n aug = get_bias_cuda(shape) + get_lfhp_cuda(shape) + get_jitter_cuda(shape)\n return aug\n\n\ndef get_visible_mask(shape=(8, 120), s_mask=0.03):\n \"\"\"Mask modeling.\"\"\"\n # Per-frame and joint\n mask = torch.rand(*shape, 17) < s_mask\n visible = (~mask).clone() # (B, L, 17)\n\n visible = visible.reshape(-1, 17) # (BL, 17)\n for child in range(17):\n parent = COCO17_TREE[child]\n if parent == -1:\n continue\n if isinstance(parent, list):\n visible[:, child] *= visible[:, parent[0]] * visible[:, parent[1]]\n else:\n visible[:, child] *= visible[:, parent]\n visible = visible.reshape(*shape, 17).clone() # (B, L, J)\n return visible\n\n\ndef get_invisible_legs_mask(shape, s_mask=0.03):\n \"\"\"\n Both legs are invisible for a random duration.\n \"\"\"\n B, L = shape\n starts = torch.randint(0, L - 90, (B,))\n ends = starts + torch.randint(30, 90, (B,))\n mask_range = torch.arange(L).unsqueeze(0).expand(B, -1)\n mask_to_apply = (mask_range >= starts.unsqueeze(1)) & (mask_range < ends.unsqueeze(1))\n mask_to_apply = mask_to_apply.unsqueeze(2).expand(-1, -1, 17).clone()\n mask_to_apply[:, :, :11] = False # only both legs are invisible\n mask_to_apply = mask_to_apply & (torch.rand(B, 1, 1) < s_mask)\n return mask_to_apply\n\n\ndef randomly_occlude_lower_half(i_x2d, s_mask=0.03):\n \"\"\"\n Randomly occlude the lower half of the image.\n \"\"\"\n raise NotImplementedError\n B, L, N, _ = i_x2d.shape\n i_x2d = i_x2d.clone()\n\n # a period of time when the lower half of the image is invisible\n starts = torch.randint(0, L - 90, (B,))\n ends = starts + torch.randint(30, 90, (B,))\n mask_range = torch.arange(L).unsqueeze(0).expand(B, -1)\n mask_to_apply = (mask_range >= starts.unsqueeze(1)) & (mask_range < ends.unsqueeze(1))\n mask_to_apply = mask_to_apply.unsqueeze(2).expand(-1, -1, N) # (B, L, N)\n\n # only the lower half of the image is invisible\n i_x2d\n i_x2d[..., 1] / 2\n\n mask_to_apply = mask_to_apply & (torch.rand(B, 1, 1) < s_mask)\n return mask_to_apply\n\n\ndef randomly_modify_hands_legs(j3d):\n hands = [9, 10]\n legs = [15, 16]\n\n B, L, J, _ = j3d.shape\n p_switch_hand = 0.001\n p_switch_leg = 0.001\n p_wrong_hand0 = 0.001\n p_wrong_hand1 = 0.001\n p_wrong_leg0 = 0.001\n p_wrong_leg1 = 0.001\n\n mask = torch.rand(B, L) < p_switch_hand\n j3d[mask][:, hands] = j3d[mask][:, hands[::-1]]\n mask = torch.rand(B, L) < p_switch_leg\n j3d[mask][:, legs] = j3d[mask][:, legs[::-1]]\n mask = torch.rand(B, L) < p_wrong_hand0\n j3d[mask][:, 9] = j3d[mask][:, 10]\n mask = torch.rand(B, L) < p_wrong_hand1\n j3d[mask][:, 10] = j3d[mask][:, 9]\n mask = torch.rand(B, L) < p_wrong_leg0\n j3d[mask][:, 15] = j3d[mask][:, 16]\n mask = torch.rand(B, L) < p_wrong_leg1\n j3d[mask][:, 16] = j3d[mask][:, 15]\n\n return j3d\n" }, { "path": "eval/GVHMR/hmr4d/utils/geo/flip_utils.py", "content": "import torch\nfrom pytorch3d.transforms import axis_angle_to_matrix, matrix_to_axis_angle\n\n\ndef flip_heatmap_coco17(output_flipped):\n assert output_flipped.ndim == 4, \"output_flipped should be [B, J, H, W]\"\n shape_ori = output_flipped.shape\n channels = 1\n output_flipped = output_flipped.reshape(shape_ori[0], -1, channels, shape_ori[2], shape_ori[3])\n output_flipped_back = output_flipped.clone()\n\n # Swap left-right parts\n for left, right in [[1, 2], [3, 4], [5, 6], [7, 8], [9, 10], [11, 12], [13, 14], [15, 16]]:\n output_flipped_back[:, left, ...] = output_flipped[:, right, ...]\n output_flipped_back[:, right, ...] = output_flipped[:, left, ...]\n output_flipped_back = output_flipped_back.reshape(shape_ori)\n # Flip horizontally\n output_flipped_back = output_flipped_back.flip(3)\n return output_flipped_back\n\n\ndef flip_bbx_xys(bbx_xys, w):\n \"\"\"\n bbx_xys: (F, 3)\n \"\"\"\n bbx_xys_flip = bbx_xys.clone()\n bbx_xys_flip[:, 0] = w - bbx_xys_flip[:, 0]\n return bbx_xys_flip\n\n\ndef flip_kp2d_coco17(kp2d, w):\n \"\"\"Flip keypoints.\"\"\"\n kp2d = kp2d.clone()\n flipped_parts = [0, 2, 1, 4, 3, 6, 5, 8, 7, 10, 9, 12, 11, 14, 13, 16, 15]\n kp2d = kp2d[..., flipped_parts, :]\n kp2d[..., 0] = w - kp2d[..., 0]\n return kp2d\n\n\ndef flip_smplx_params(smplx_params):\n \"\"\"Flip pose.\n The flipping is based on SMPLX parameters.\n \"\"\"\n rotation = torch.cat([smplx_params[\"global_orient\"], smplx_params[\"body_pose\"]], dim=1)\n\n BN = rotation.shape[0]\n pose = rotation.reshape(BN, -1).transpose(0, 1)\n\n SMPL_JOINTS_FLIP_PERM = [0, 2, 1, 3, 5, 4, 6, 8, 7, 9, 11, 10, 12, 14, 13, 15, 17, 16, 19, 18, 21, 20] # , 23, 22]\n SMPL_POSE_FLIP_PERM = []\n for i in SMPL_JOINTS_FLIP_PERM:\n SMPL_POSE_FLIP_PERM.append(3 * i)\n SMPL_POSE_FLIP_PERM.append(3 * i + 1)\n SMPL_POSE_FLIP_PERM.append(3 * i + 2)\n\n pose = pose[SMPL_POSE_FLIP_PERM]\n\n # we also negate the second and the third dimension of the axis-angle\n pose[1::3] = -pose[1::3]\n pose[2::3] = -pose[2::3]\n pose = pose.transpose(0, 1).reshape(BN, -1, 3)\n\n smplx_params_flipped = smplx_params.copy()\n smplx_params_flipped[\"global_orient\"] = pose[:, :1]\n smplx_params_flipped[\"body_pose\"] = pose[:, 1:]\n return smplx_params_flipped\n\n\ndef avg_smplx_aa(aa1, aa2):\n def avg_rot(rot):\n # input [B,...,3,3] --> output [...,3,3]\n rot = rot.mean(dim=0)\n U, _, V = torch.svd(rot)\n rot = U @ V.transpose(-1, -2)\n return rot\n\n B, J3 = aa1.shape\n aa1 = aa1.reshape(B, -1, 3)\n aa2 = aa2.reshape(B, -1, 3)\n\n R1 = axis_angle_to_matrix(aa1)\n R2 = axis_angle_to_matrix(aa2)\n R_avg = avg_rot(torch.stack([R1, R2]))\n aa_avg = matrix_to_axis_angle(R_avg).reshape(B, -1)\n\n return aa_avg\n" }, { "path": "eval/GVHMR/hmr4d/utils/geo/hmr_cam.py", "content": "import torch\nimport numpy as np\nfrom hmr4d.utils.geo_transform import project_p2d, convert_bbx_xys_to_lurb, cvt_to_bi01_p2d\n\n\ndef estimate_focal_length(img_w, img_h):\n return (img_w**2 + img_h**2) ** 0.5 # Diagonal FOV = 2*arctan(0.5) * 180/pi = 53\n\n\ndef estimate_K(img_w, img_h):\n focal_length = estimate_focal_length(img_w, img_h)\n K = torch.eye(3).float()\n K[0, 0] = focal_length\n K[1, 1] = focal_length\n K[0, 2] = img_w / 2.0\n K[1, 2] = img_h / 2.0\n return K\n\n\ndef convert_K_to_K4(K):\n K4 = torch.stack([K[0, 0], K[1, 1], K[0, 2], K[1, 2]]).float()\n return K4\n\n\ndef convert_f_to_K(focal_length, img_w, img_h):\n K = torch.eye(3).float()\n K[0, 0] = focal_length\n K[1, 1] = focal_length\n K[0, 2] = img_w / 2.0\n K[1, 2] = img_h / 2.0\n return K\n\n\ndef resize_K(K, f=0.5):\n K = K.clone() * f\n K[..., 2, 2] = 1.0\n return K\n\n\ndef create_camera_sensor(width=None, height=None, f_fullframe=None):\n if width is None or height is None:\n # The 4:3 aspect ratio is widely adopted by image sensors in mobile phones.\n if np.random.rand() < 0.5:\n width, height = 1200, 1600\n else:\n width, height = 1600, 1200\n\n # Sample FOV from common options:\n # 1. wide-angle lenses are common in mobile phones,\n # 2. telephoto lenses has less perspective effect, which should makes it easy to learn\n if f_fullframe is None:\n f_fullframe_options = [24, 26, 28, 30, 35, 40, 50, 60, 70]\n f_fullframe = np.random.choice(f_fullframe_options)\n\n # We use diag to map focal-length: https://www.nikonians.org/reviews/fov-tables\n diag_fullframe = (24**2 + 36**2) ** 0.5\n diag_img = (width**2 + height**2) ** 0.5\n focal_length = diag_img / diag_fullframe * f_fullframe\n\n K_fullimg = torch.eye(3)\n K_fullimg[0, 0] = focal_length\n K_fullimg[1, 1] = focal_length\n K_fullimg[0, 2] = width / 2\n K_fullimg[1, 2] = height / 2\n\n return width, height, K_fullimg\n\n\n# ====== Compute cliffcam ===== #\n\n\ndef convert_xys_to_cliff_cam_wham(xys, res):\n \"\"\"\n Args:\n xys: (N, 3) in pixel. Note s should not be touched by 200\n res: (2), e.g. [4112., 3008.] (w,h)\n Returns:\n cliff_cam: (N, 3), normalized representation\n \"\"\"\n\n def normalize_keypoints_to_image(x, res):\n \"\"\"\n Args:\n x: (N, 2), centers\n res: (2), e.g. [4112., 3008.]\n Returns:\n x_normalized: (N, 2)\n \"\"\"\n res = res.to(x.device)\n scale = res.max(-1)[0].reshape(-1)\n mean = torch.stack([res[..., 0] / scale, res[..., 1] / scale], dim=-1).to(x.device)\n x = 2 * x / scale.reshape(*[1 for i in range(len(x.shape[1:]))]) - mean.reshape(\n *[1 for i in range(len(x.shape[1:-1]))], -1\n )\n return x\n\n centers = normalize_keypoints_to_image(xys[:, :2], res) # (N, 2)\n scale = xys[:, 2:] / res.max()\n location = torch.cat((centers, scale), dim=-1)\n return location\n\n\ndef compute_bbox_info_bedlam(bbx_xys, K_fullimg):\n \"\"\"impl as in BEDLAM\n Args:\n bbx_xys: ((B), N, 3), in pixel space described by K_fullimg\n K_fullimg: ((B), (N), 3, 3)\n Returns:\n bbox_info: ((B), N, 3)\n \"\"\"\n fl = K_fullimg[..., 0, 0].unsqueeze(-1)\n icx = K_fullimg[..., 0, 2]\n icy = K_fullimg[..., 1, 2]\n\n cx, cy, b = bbx_xys[..., 0], bbx_xys[..., 1], bbx_xys[..., 2]\n bbox_info = torch.stack([cx - icx, cy - icy, b], dim=-1)\n bbox_info = bbox_info / fl\n return bbox_info\n\n\n# ====== Convert Prediction to Cam-t ===== #\n\n\ndef compute_transl_full_cam(pred_cam, bbx_xys, K_fullimg):\n s, tx, ty = pred_cam[..., 0], pred_cam[..., 1], pred_cam[..., 2]\n focal_length = K_fullimg[..., 0, 0]\n\n icx = K_fullimg[..., 0, 2]\n icy = K_fullimg[..., 1, 2]\n sb = s * bbx_xys[..., 2]\n cx = 2 * (bbx_xys[..., 0] - icx) / (sb + 1e-9)\n cy = 2 * (bbx_xys[..., 1] - icy) / (sb + 1e-9)\n tz = 2 * focal_length / (sb + 1e-9)\n\n cam_t = torch.stack([tx + cx, ty + cy, tz], dim=-1)\n return cam_t\n\n\ndef get_a_pred_cam(transl, bbx_xys, K_fullimg):\n \"\"\"Inverse operation of compute_transl_full_cam\"\"\"\n assert transl.ndim == bbx_xys.ndim # (*, L, 3)\n assert K_fullimg.ndim == (bbx_xys.ndim + 1) # (*, L, 3, 3)\n f = K_fullimg[..., 0, 0]\n cx = K_fullimg[..., 0, 2]\n cy = K_fullimg[..., 1, 2]\n gt_s = 2 * f / (transl[..., 2] * bbx_xys[..., 2]) # (B, L)\n gt_x = transl[..., 0] - transl[..., 2] / f * (bbx_xys[..., 0] - cx)\n gt_y = transl[..., 1] - transl[..., 2] / f * (bbx_xys[..., 1] - cy)\n gt_pred_cam = torch.stack([gt_s, gt_x, gt_y], dim=-1)\n return gt_pred_cam\n\n\n# ====== 3D to 2D ===== #\n\n\ndef project_to_bi01(points, bbx_xys, K_fullimg):\n \"\"\"\n points: (B, L, J, 3)\n bbx_xys: (B, L, 3)\n K_fullimg: (B, L, 3, 3)\n \"\"\"\n # p2d = project_p2d(points, K_fullimg)\n p2d = perspective_projection(points, K_fullimg)\n bbx_lurb = convert_bbx_xys_to_lurb(bbx_xys)\n p2d_bi01 = cvt_to_bi01_p2d(p2d, bbx_lurb)\n return p2d_bi01\n\n\ndef perspective_projection(points, K):\n # points: (B, L, J, 3)\n # K: (B, L, 3, 3)\n projected_points = points / points[..., -1].unsqueeze(-1)\n projected_points = torch.einsum(\"...ij,...kj->...ki\", K, projected_points.float())\n return projected_points[..., :-1]\n\n\n# ====== 2D (bbx from j2d) ===== #\n\n\ndef normalize_kp2d(obs_kp2d, bbx_xys, clamp_scale_min=False):\n \"\"\"\n Args:\n obs_kp2d: (B, L, J, 3) [x, y, c]\n bbx_xys: (B, L, 3)\n Returns:\n obs: (B, L, J, 3) [x, y, c]\n \"\"\"\n obs_xy = obs_kp2d[..., :2] # (B, L, J, 2)\n obs_conf = obs_kp2d[..., 2] # (B, L, J)\n center = bbx_xys[..., :2]\n scale = bbx_xys[..., [2]]\n\n # Mark keypoints outside the bounding box as invisible\n xy_max = center + scale / 2\n xy_min = center - scale / 2\n invisible_mask = (\n (obs_xy[..., 0] < xy_min[..., None, 0])\n + (obs_xy[..., 0] > xy_max[..., None, 0])\n + (obs_xy[..., 1] < xy_min[..., None, 1])\n + (obs_xy[..., 1] > xy_max[..., None, 1])\n )\n obs_conf = obs_conf * ~invisible_mask\n if clamp_scale_min:\n scale = scale.clamp(min=1e-5)\n normalized_obs_xy = 2 * (obs_xy - center.unsqueeze(-2)) / scale.unsqueeze(-2)\n\n return torch.cat([normalized_obs_xy, obs_conf[..., None]], dim=-1)\n\n\ndef get_bbx_xys(i_j2d, bbx_ratio=[192, 256], do_augment=False, base_enlarge=1.2):\n \"\"\"Args: (B, L, J, 3) [x,y,c] -> Returns: (B, L, 3)\"\"\"\n # Center\n min_x = i_j2d[..., 0].min(-1)[0]\n max_x = i_j2d[..., 0].max(-1)[0]\n min_y = i_j2d[..., 1].min(-1)[0]\n max_y = i_j2d[..., 1].max(-1)[0]\n center_x = (min_x + max_x) / 2\n center_y = (min_y + max_y) / 2\n\n # Size\n h = max_y - min_y # (B, L)\n w = max_x - min_x # (B, L)\n\n if True: # fit w and h into aspect-ratio\n aspect_ratio = bbx_ratio[0] / bbx_ratio[1]\n mask1 = w > aspect_ratio * h\n h[mask1] = w[mask1] / aspect_ratio\n mask2 = w < aspect_ratio * h\n w[mask2] = h[mask2] * aspect_ratio\n\n # apply a common factor to enlarge the bounding box\n bbx_size = torch.max(h, w) * base_enlarge\n\n if do_augment:\n B, L = bbx_size.shape[:2]\n device = bbx_size.device\n if True:\n scaleFactor = torch.rand((B, L), device=device) * 0.3 + 1.05 # 1.05~1.35\n txFactor = torch.rand((B, L), device=device) * 1.6 - 0.8 # -0.8~0.8\n tyFactor = torch.rand((B, L), device=device) * 1.6 - 0.8 # -0.8~0.8\n else:\n scaleFactor = torch.rand((B, 1), device=device) * 0.3 + 1.05 # 1.05~1.35\n txFactor = torch.rand((B, 1), device=device) * 1.6 - 0.8 # -0.8~0.8\n tyFactor = torch.rand((B, 1), device=device) * 1.6 - 0.8 # -0.8~0.8\n\n raw_bbx_size = bbx_size / base_enlarge\n bbx_size = raw_bbx_size * scaleFactor\n center_x += raw_bbx_size / 2 * ((scaleFactor - 1) * txFactor)\n center_y += raw_bbx_size / 2 * ((scaleFactor - 1) * tyFactor)\n\n return torch.stack([center_x, center_y, bbx_size], dim=-1)\n\n\ndef safely_render_x3d_K(x3d, K_fullimg, thr):\n \"\"\"\n Args:\n x3d: (B, L, V, 3), should as least have a safe points (not examined here)\n K_fullimg: (B, L, 3, 3)\n Returns:\n bbx_xys: (B, L, 3)\n i_x2d: (B, L, V, 2)\n \"\"\"\n # For each frame, update unsafe z ( 0:\n x3d[..., 2][x3d_unsafe_mask] = thr\n if False:\n from hmr4d.utils.wis3d_utils import make_wis3d\n\n wis3d = make_wis3d(name=\"debug-update-z\")\n bs, ls, vs = torch.where(x3d_unsafe_mask)\n bs = torch.unique(bs)\n for b in bs:\n for f in range(x3d.size(1)):\n wis3d.set_scene_id(f)\n wis3d.add_point_cloud(x3d[b, f], name=\"unsafe\")\n pass\n\n # renfer\n i_x2d = perspective_projection(x3d, K_fullimg) # (B, L, V, 2)\n return i_x2d\n\n\ndef get_bbx_xys_from_xyxy(bbx_xyxy, base_enlarge=1.2):\n \"\"\"\n Args:\n bbx_xyxy: (N, 4) [x1, y1, x2, y2]\n Returns:\n bbx_xys: (N, 3) [center_x, center_y, size]\n \"\"\"\n\n i_p2d = torch.stack([bbx_xyxy[:, [0, 1]], bbx_xyxy[:, [2, 3]]], dim=1) # (L, 2, 2)\n bbx_xys = get_bbx_xys(i_p2d[None], base_enlarge=base_enlarge)[0]\n return bbx_xys\n\n\ndef bbx_xyxy_from_x(p2d):\n \"\"\"\n Args:\n p2d: (*, V, 2) - Tensor containing 2D points.\n\n Returns:\n bbx_xyxy: (*, 4) - Bounding box coordinates in the format (xmin, ymin, xmax, ymax).\n \"\"\"\n # Compute the minimum and maximum coordinates for the bounding box\n xy_min = p2d.min(dim=-2).values # (*, 2)\n xy_max = p2d.max(dim=-2).values # (*, 2)\n\n # Concatenate min and max coordinates to form the bounding box\n bbx_xyxy = torch.cat([xy_min, xy_max], dim=-1) # (*, 4)\n\n return bbx_xyxy\n\n\ndef bbx_xyxy_from_masked_x(p2d, mask):\n \"\"\"\n Args:\n p2d: (*, V, 2) - Tensor containing 2D points.\n mask: (*, V) - Boolean tensor indicating valid points.\n\n Returns:\n bbx_xyxy: (*, 4) - Bounding box coordinates in the format (xmin, ymin, xmax, ymax).\n \"\"\"\n # Ensure the shapes of p2d and mask are compatible\n assert p2d.shape[:-1] == mask.shape, \"The shape of p2d and mask are not compatible.\"\n\n # Flatten the input tensors for batch processing\n p2d_flat = p2d.view(-1, p2d.shape[-2], p2d.shape[-1])\n mask_flat = mask.view(-1, mask.shape[-1])\n\n # Set masked out values to a large positive and negative value respectively\n p2d_min = torch.where(mask_flat.unsqueeze(-1), p2d_flat, torch.tensor(float(\"inf\")).to(p2d_flat))\n p2d_max = torch.where(mask_flat.unsqueeze(-1), p2d_flat, torch.tensor(float(\"-inf\")).to(p2d_flat))\n\n # Compute the minimum and maximum coordinates for the bounding box\n xy_min = p2d_min.min(dim=1).values # (BL, 2)\n xy_max = p2d_max.max(dim=1).values # (BL, 2)\n\n # Concatenate min and max coordinates to form the bounding box\n bbx_xyxy = torch.cat([xy_min, xy_max], dim=-1) # (BL, 4)\n\n # Reshape back to the original shape prefix\n bbx_xyxy = bbx_xyxy.view(*p2d.shape[:-2], 4)\n\n return bbx_xyxy\n\n\ndef bbx_xyxy_ratio(xyxy1, xyxy2):\n \"\"\"Designed for fov/unbounded\n Args:\n xyxy1: (*, 4)\n xyxy2: (*, 4)\n Return:\n ratio: (*), squared_area(xyxy1) / squared_area(xyxy2)\n \"\"\"\n area1 = (xyxy1[..., 2] - xyxy1[..., 0]) * (xyxy1[..., 3] - xyxy1[..., 1])\n area2 = (xyxy2[..., 2] - xyxy2[..., 0]) * (xyxy2[..., 3] - xyxy2[..., 1])\n # Check\n area1[~torch.isfinite(area1)] = 0 # replace inf in area1 with 0\n assert (area2 > 0).all(), \"area2 should be positive\"\n return area1 / area2\n\n\ndef get_mesh_in_fov_category(mask):\n \"\"\"mask: (L, V)\n The definition:\n 1. FullyVisible: The mesh in every frame is entirely within the field of view (FOV).\n 2. PartiallyVisible: In some frames, parts of the mesh are outside the FOV, while other parts are within the FOV.\n 3. PartiallyOut: In some frames, the mesh is completely outside the FOV, while in others, it is visible.\n 4. FullyOut: The mesh is completely outside the FOV in every frame.\n \"\"\"\n mask = mask.clone().cpu()\n is_class1 = mask.all() # FullyVisible\n is_class2 = mask.any(1).all() * ~is_class1 # PartiallyVisible\n is_class4 = ~(mask.any()) # PartiallyOut\n is_class3 = ~is_class1 * ~is_class2 * ~is_class4 # FullyOut\n\n mask_frame_any_verts = mask.any(1)\n assert is_class1.int() + is_class2.int() + is_class3.int() + is_class4.int() == 1\n class_type = is_class1.int() + 2 * is_class2.int() + 3 * is_class3.int() + 4 * is_class4.int()\n return class_type.item(), mask_frame_any_verts\n\n\ndef get_infov_mask(p2d, w_real, h_real):\n \"\"\"\n Args:\n p2d: (B, L, V, 2)\n w_real, h_real: (B, L) or int\n Returns:\n mask: (B, L, V)\n \"\"\"\n x, y = p2d[..., 0], p2d[..., 1]\n if isinstance(w_real, int):\n mask = (x >= 0) * (x < w_real) * (y >= 0) * (y < h_real)\n else:\n mask = (x >= 0) * (x < w_real[..., None]) * (y >= 0) * (y < h_real[..., None])\n return mask\n" }, { "path": "eval/GVHMR/hmr4d/utils/geo/hmr_global.py", "content": "import torch\nfrom pytorch3d.transforms import axis_angle_to_matrix, matrix_to_axis_angle, matrix_to_quaternion, quaternion_to_matrix\nimport hmr4d.utils.matrix as matrix\nfrom hmr4d.utils.net_utils import gaussian_smooth\n\n\ndef get_R_c2gv(R_w2c, axis_gravity_in_w=[0, 0, -1]):\n \"\"\"\n Args:\n R_w2c: (*, 3, 3)\n Returns:\n R_c2gv: (*, 3, 3)\n \"\"\"\n if isinstance(axis_gravity_in_w, list):\n axis_gravity_in_w = torch.tensor(axis_gravity_in_w).float() # gravity direction in world coord\n axis_z_in_c = torch.tensor([0, 0, 1]).float()\n\n # get gv-coord axes in in c-coord\n axis_y_of_gv = R_w2c @ axis_gravity_in_w # (*, 3)\n axis_x_of_gv = axis_y_of_gv.cross(axis_z_in_c.expand_as(axis_y_of_gv), dim=-1)\n # normalize\n axis_x_of_gv_norm = axis_x_of_gv.norm(dim=-1, keepdim=True)\n axis_x_of_gv = axis_x_of_gv / (axis_x_of_gv_norm + 1e-5)\n axis_x_of_gv[axis_x_of_gv_norm.squeeze(-1) < 1e-5] = torch.tensor([1.0, 0.0, 0.0]) # use cam x-axis as axis_x_of_gv\n axis_z_of_gv = axis_x_of_gv.cross(axis_y_of_gv, dim=-1)\n\n R_gv2c = torch.stack([axis_x_of_gv, axis_y_of_gv, axis_z_of_gv], dim=-1) # (*, 3, 3)\n R_c2gv = R_gv2c.transpose(-1, -2) # (*, 3, 3)\n return R_c2gv\n\n\ntsf_axisangle = {\n \"ay->ay\": [0, 0, 0],\n \"any->ay\": [0, 0, torch.pi],\n \"az->ay\": [-torch.pi / 2, 0, 0],\n \"ay->any\": [0, 0, torch.pi],\n}\n\n\ndef get_tgtcoord_rootparam(global_orient, transl, gravity_vec=None, tgt_gravity_vec=None, tsf=\"ay->ay\"):\n \"\"\"Rotate around the origin center, to match the new gravity direction\n Args:\n global_orient: torch.tensor, (*, 3)\n transl: torch.tensor, (*, 3)\n gravity_vec: torch.tensor, (3,)\n tgt_gravity_vec: torch.tensor, (3,)\n Returns:\n tgt_global_orient: torch.tensor, (*, 3)\n tgt_transl: torch.tensor, (*, 3)\n R_g2tg: (3, 3)\n \"\"\"\n # get rotation matrix\n device = global_orient.device\n if gravity_vec is None and tgt_gravity_vec is None:\n aa = torch.tensor(tsf_axisangle[tsf]).to(device)\n R_g2tg = axis_angle_to_matrix(aa) # (3, 3)\n else:\n raise NotImplementedError\n # TODO: Impl this function\n gravity_vec = torch.tensor(gravity_vec).float().to(device)\n gravity_vec = gravity_vec / gravity_vec.norm()\n tgt_gravity_vec = torch.tensor(tgt_gravity_vec).float().to(device)\n tgt_gravity_vec = tgt_gravity_vec / tgt_gravity_vec.norm()\n # pick one identity axis\n axis_identity = torch.tensor([0, 0, 0]).float().to(device)\n for i in (gravity_vec == 0) & (tgt_global_orient == 0):\n if i:\n axis_identity[i] = 1\n break\n\n # rotate\n global_orient_R = axis_angle_to_matrix(global_orient) # (*, 3, 3)\n tgt_global_orient = matrix_to_axis_angle(R_g2tg @ global_orient_R) # (*, 3, 3)\n tgt_transl = torch.einsum(\"...ij,...j->...i\", R_g2tg, transl)\n\n return tgt_global_orient, tgt_transl, R_g2tg\n\n\ndef get_c_rootparam(global_orient, transl, T_w2c, offset):\n \"\"\"\n Args:\n global_orient: torch.tensor, (F, 3)\n transl: torch.tensor, (F, 3)\n T_w2c: torch.tensor, (*, 4, 4)\n offset: torch.tensor, (3,)\n Returns:\n R_c: torch.tensor, (F, 3)\n t_c: torch.tensor, (F, 3)\n \"\"\"\n assert global_orient.shape == transl.shape and len(global_orient.shape) == 2\n R_w = axis_angle_to_matrix(global_orient) # (F, 3, 3)\n t_w = transl # (F, 3)\n\n R_w2c = T_w2c[..., :3, :3] # (*, 3, 3)\n t_w2c = T_w2c[..., :3, 3] # (*, 3)\n if len(R_w2c.shape) == 2:\n R_w2c = R_w2c[None].expand(R_w.size(0), -1, -1) # (F, 3, 3)\n t_w2c = t_w2c[None].expand(t_w.size(0), -1)\n\n R_c = matrix_to_axis_angle(R_w2c @ R_w) # (F, 3)\n t_c = torch.einsum(\"fij,fj->fi\", R_w2c, t_w + offset) + t_w2c - offset # (F, 3)\n return R_c, t_c\n\n\ndef get_T_w2c_from_wcparams(global_orient_w, transl_w, global_orient_c, transl_c, offset):\n \"\"\"\n Args:\n global_orient_w: torch.tensor, (F, 3)\n transl_w: torch.tensor, (F, 3)\n global_orient_c: torch.tensor, (F, 3)\n transl_c: torch.tensor, (F, 3)\n offset: torch.tensor, (*, 3)\n Returns:\n T_w2c: torch.tensor, (F, 4, 4)\n \"\"\"\n assert global_orient_w.shape == transl_w.shape and len(global_orient_w.shape) == 2\n assert global_orient_c.shape == transl_c.shape and len(global_orient_c.shape) == 2\n\n R_w = axis_angle_to_matrix(global_orient_w) # (F, 3, 3)\n t_w = transl_w # (F, 3)\n R_c = axis_angle_to_matrix(global_orient_c) # (F, 3, 3)\n t_c = transl_c # (F, 3)\n\n R_w2c = R_c @ R_w.transpose(-1, -2) # (F, 3, 3)\n t_w2c = t_c + offset - torch.einsum(\"fij,fj->fi\", R_w2c, t_w + offset) # (F, 3)\n T_w2c = torch.eye(4, device=global_orient_w.device).repeat(R_w.size(0), 1, 1) # (F, 4, 4)\n T_w2c[..., :3, :3] = R_w2c # (F, 3, 3)\n T_w2c[..., :3, 3] = t_w2c # (F, 3)\n return T_w2c\n\n\ndef get_local_transl_vel(transl, global_orient):\n \"\"\"\n transl velocity is in local coordinate (or, SMPL-coord)\n Args:\n transl: (*, L, 3)\n global_orient: (*, L, 3)\n Returns:\n transl_vel: (*, L, 3)\n \"\"\"\n assert len(transl.shape) == len(global_orient.shape)\n global_orient_R = axis_angle_to_matrix(global_orient) # (B, L, 3, 3)\n transl_vel = transl[..., 1:, :] - transl[..., :-1, :] # (B, L-1, 3)\n transl_vel = torch.cat([transl_vel, transl_vel[..., [-1], :]], dim=-2) # (B, L, 3) last-padding\n\n # v_local = R^T @ v_global\n local_transl_vel = torch.einsum(\"...lij,...li->...lj\", global_orient_R, transl_vel)\n return local_transl_vel\n\n\ndef rollout_local_transl_vel(local_transl_vel, global_orient, transl_0=None):\n \"\"\"\n transl velocity is in local coordinate (or, SMPL-coord)\n Args:\n local_transl_vel: (*, L, 3)\n global_orient: (*, L, 3)\n transl_0: (*, 1, 3), if not provided, the start point is 0\n Returns:\n transl: (*, L, 3)\n \"\"\"\n global_orient_R = axis_angle_to_matrix(global_orient)\n transl_vel = torch.einsum(\"...lij,...lj->...li\", global_orient_R, local_transl_vel)\n\n # set start point\n if transl_0 is None:\n transl_0 = transl_vel[..., :1, :].clone().detach().zero_()\n transl_ = torch.cat([transl_0, transl_vel[..., :-1, :]], dim=-2)\n\n # rollout from start point\n transl = torch.cumsum(transl_, dim=-2)\n return transl\n\n\ndef get_local_transl_vel_alignhead(transl, global_orient):\n # assume global_orient is ay\n global_orient_rot = axis_angle_to_matrix(global_orient) # (*, 3, 3)\n global_orient_quat = matrix_to_quaternion(global_orient_rot) # (*, 4)\n\n global_orient_quat_xyzw = matrix.quat_wxyz2xyzw(global_orient_quat) # (*, 4)\n head_quat_xyzw = matrix.calc_heading_quat(global_orient_quat_xyzw, head_ind=2, gravity_axis=\"y\") # (*, 4)\n head_quat = matrix.quat_xyzw2wxyz(head_quat_xyzw) # (*, 4)\n head_rot = quaternion_to_matrix(head_quat)\n head_aa = matrix_to_axis_angle(head_rot)\n\n local_transl_vel_alignhead = get_local_transl_vel(transl, head_aa)\n return local_transl_vel_alignhead\n\n\ndef rollout_local_transl_vel_alignhead(local_transl_vel_alignhead, global_orient, transl_0=None):\n # assume global_orient is ay\n global_orient_rot = axis_angle_to_matrix(global_orient) # (*, 3, 3)\n global_orient_quat = matrix_to_quaternion(global_orient_rot) # (*, 4)\n\n global_orient_quat_xyzw = matrix.quat_wxyz2xyzw(global_orient_quat) # (*, 4)\n head_quat_xyzw = matrix.calc_heading_quat(global_orient_quat_xyzw, head_ind=2, gravity_axis=\"y\") # (*, 4)\n head_quat = matrix.quat_xyzw2wxyz(head_quat_xyzw) # (*, 4)\n head_rot = quaternion_to_matrix(head_quat)\n head_aa = matrix_to_axis_angle(head_rot)\n\n transl = rollout_local_transl_vel(local_transl_vel_alignhead, head_aa, transl_0)\n return transl\n\n\ndef get_local_transl_vel_alignhead_absy(transl, global_orient):\n # assume global_orient is ay\n global_orient_rot = axis_angle_to_matrix(global_orient) # (*, 3, 3)\n global_orient_quat = matrix_to_quaternion(global_orient_rot) # (*, 4)\n\n global_orient_quat_xyzw = matrix.quat_wxyz2xyzw(global_orient_quat) # (*, 4)\n head_quat_xyzw = matrix.calc_heading_quat(global_orient_quat_xyzw, head_ind=2, gravity_axis=\"y\") # (*, 4)\n head_quat = matrix.quat_xyzw2wxyz(head_quat_xyzw) # (*, 4)\n head_rot = quaternion_to_matrix(head_quat)\n head_aa = matrix_to_axis_angle(head_rot)\n\n local_transl_vel_alignhead = get_local_transl_vel(transl, head_aa)\n abs_y = torch.cumsum(local_transl_vel_alignhead[..., [1]], dim=-2) # (*, L, 1)\n local_transl_vel_alignhead_absy = torch.cat(\n [local_transl_vel_alignhead[..., [0]], abs_y, local_transl_vel_alignhead[..., [2]]], dim=-1\n )\n\n return local_transl_vel_alignhead_absy\n\n\ndef rollout_local_transl_vel_alignhead_absy(local_transl_vel_alignhead_absy, global_orient, transl_0=None):\n # assume global_orient is ay\n global_orient_rot = axis_angle_to_matrix(global_orient) # (*, 3, 3)\n global_orient_quat = matrix_to_quaternion(global_orient_rot) # (*, 4)\n\n global_orient_quat_xyzw = matrix.quat_wxyz2xyzw(global_orient_quat) # (*, 4)\n head_quat_xyzw = matrix.calc_heading_quat(global_orient_quat_xyzw, head_ind=2, gravity_axis=\"y\") # (*, 4)\n head_quat = matrix.quat_xyzw2wxyz(head_quat_xyzw) # (*, 4)\n head_rot = quaternion_to_matrix(head_quat)\n head_aa = matrix_to_axis_angle(head_rot)\n\n local_transl_vel_alignhead_y = (\n local_transl_vel_alignhead_absy[..., 1:, [1]] - local_transl_vel_alignhead_absy[..., :-1, [1]]\n )\n local_transl_vel_alignhead_y = torch.cat(\n [local_transl_vel_alignhead_absy[..., :1, [1]], local_transl_vel_alignhead_y], dim=-2\n )\n local_transl_vel_alignhead = torch.cat(\n [\n local_transl_vel_alignhead_absy[..., [0]],\n local_transl_vel_alignhead_y,\n local_transl_vel_alignhead_absy[..., [2]],\n ],\n dim=-1,\n )\n\n transl = rollout_local_transl_vel(local_transl_vel_alignhead, head_aa, transl_0)\n return transl\n\n\ndef get_local_transl_vel_alignhead_absgy(transl, global_orient):\n # assume global_orient is ay\n global_orient_rot = axis_angle_to_matrix(global_orient) # (*, 3, 3)\n global_orient_quat = matrix_to_quaternion(global_orient_rot) # (*, 4)\n\n global_orient_quat_xyzw = matrix.quat_wxyz2xyzw(global_orient_quat) # (*, 4)\n head_quat_xyzw = matrix.calc_heading_quat(global_orient_quat_xyzw, head_ind=2, gravity_axis=\"y\") # (*, 4)\n head_quat = matrix.quat_xyzw2wxyz(head_quat_xyzw) # (*, 4)\n head_rot = quaternion_to_matrix(head_quat)\n head_aa = matrix_to_axis_angle(head_rot)\n\n local_transl_vel_alignhead = get_local_transl_vel(transl, head_aa)\n abs_y = transl[..., [1]] # (*, L, 1)\n local_transl_vel_alignhead_absy = torch.cat(\n [local_transl_vel_alignhead[..., [0]], abs_y, local_transl_vel_alignhead[..., [2]]], dim=-1\n )\n\n return local_transl_vel_alignhead_absy\n\n\ndef rollout_local_transl_vel_alignhead_absgy(local_transl_vel_alignhead_absgy, global_orient, transl_0=None):\n # assume global_orient is ay\n global_orient_rot = axis_angle_to_matrix(global_orient) # (*, 3, 3)\n global_orient_quat = matrix_to_quaternion(global_orient_rot) # (*, 4)\n\n global_orient_quat_xyzw = matrix.quat_wxyz2xyzw(global_orient_quat) # (*, 4)\n head_quat_xyzw = matrix.calc_heading_quat(global_orient_quat_xyzw, head_ind=2, gravity_axis=\"y\") # (*, 4)\n head_quat = matrix.quat_xyzw2wxyz(head_quat_xyzw) # (*, 4)\n head_rot = quaternion_to_matrix(head_quat)\n head_aa = matrix_to_axis_angle(head_rot)\n\n local_transl_vel_alignhead_y = (\n local_transl_vel_alignhead_absgy[..., 1:, [1]] - local_transl_vel_alignhead_absgy[..., :-1, [1]]\n )\n local_transl_vel_alignhead_y = torch.cat(\n [local_transl_vel_alignhead_y, local_transl_vel_alignhead_y[..., -1:, :]], dim=-2\n )\n if transl_0 is not None:\n transl_0 = transl_0.clone()\n transl_0[..., 1] = local_transl_vel_alignhead_absgy[..., :1, 1]\n else:\n transl_0 = local_transl_vel_alignhead_absgy.clone()[..., :1, :] # (*, 1, 3)\n transl_0[..., :1, 0] = 0.0\n transl_0[..., :1, 2] = 0.0\n\n local_transl_vel_alignhead = torch.cat(\n [\n local_transl_vel_alignhead_absgy[..., [0]],\n local_transl_vel_alignhead_y,\n local_transl_vel_alignhead_absgy[..., [2]],\n ],\n dim=-1,\n )\n\n transl = rollout_local_transl_vel(local_transl_vel_alignhead, head_aa, transl_0)\n return transl\n\n\ndef rollout_vel(vel, transl_0=None):\n \"\"\"\n Args:\n vel: (*, L, 3)\n transl_0: (*, 1, 3), if not provided, the start point is 0\n Returns:\n transl: (*, L, 3)\n \"\"\"\n # set start point\n if transl_0 is None:\n assert len(vel.shape) == len(transl_0.shape)\n transl_0 = vel[..., :1, :].clone().detach().zero_()\n transl_ = torch.cat([transl_0, vel[..., :-1, :]], dim=-2)\n\n # rollout from start point\n transl = torch.cumsum(transl_, dim=-2)\n return transl\n\n\ndef get_static_joint_mask(w_j3d, vel_thr=0.25, smooth=False, repeat_last=False):\n \"\"\"\n w_j3d: (*, L, J, 3)\n vel_thr: HuMoR uses 0.15m/s\n \"\"\"\n joint_v_ = (w_j3d[..., 1:, :, :] - w_j3d[..., :-1, :, :]).pow(2).sum(-1).sqrt() / 0.033 # (*, L-1, J)\n if smooth:\n joint_v_ = gaussian_smooth(joint_v_, 3, -2)\n\n static_joint_mask = joint_v_ < vel_thr # 1 as stable, 0 as moving\n\n if repeat_last: # repeat the last frame, this makes the shape same as w_j3d\n static_joint_mask = torch.cat([static_joint_mask, static_joint_mask[..., [-1], :]], dim=-2)\n\n return static_joint_mask\n" }, { "path": "eval/GVHMR/hmr4d/utils/geo/quaternion.py", "content": "# Copyright (c) 2018-present, Facebook, Inc.\n# All rights reserved.\n#\n# This source code is licensed under the license found in the\n# LICENSE file in the root directory of this source tree.\n#\n\nimport torch\nimport numpy as np\n\n_EPS4 = np.finfo(float).eps * 4.0\n\ntry:\n _FLOAT_EPS = np.finfo(np.float).eps\nexcept:\n _FLOAT_EPS = np.finfo(float).eps\n\n\n# PyTorch-backed implementations\ndef qinv(q):\n assert q.shape[-1] == 4, \"q must be a tensor of shape (*, 4)\"\n mask = torch.ones_like(q)\n mask[..., 1:] = -mask[..., 1:]\n return q * mask\n\n\ndef qinv_np(q):\n assert q.shape[-1] == 4, \"q must be a tensor of shape (*, 4)\"\n return qinv(torch.from_numpy(q).float()).numpy()\n\n\ndef qnormalize(q):\n assert q.shape[-1] == 4, \"q must be a tensor of shape (*, 4)\"\n return q / torch.clamp(torch.norm(q, dim=-1, keepdim=True), min=1e-8)\n\n\ndef qmul(q, r):\n \"\"\"\n Multiply quaternion(s) q with quaternion(s) r.\n Expects two equally-sized tensors of shape (*, 4), where * denotes any number of dimensions.\n Returns q*r as a tensor of shape (*, 4).\n \"\"\"\n assert q.shape[-1] == 4\n assert r.shape[-1] == 4\n\n original_shape = q.shape\n\n # Compute outer product\n terms = torch.bmm(r.reshape(-1, 4, 1), q.reshape(-1, 1, 4))\n\n w = terms[:, 0, 0] - terms[:, 1, 1] - terms[:, 2, 2] - terms[:, 3, 3]\n x = terms[:, 0, 1] + terms[:, 1, 0] - terms[:, 2, 3] + terms[:, 3, 2]\n y = terms[:, 0, 2] + terms[:, 1, 3] + terms[:, 2, 0] - terms[:, 3, 1]\n z = terms[:, 0, 3] - terms[:, 1, 2] + terms[:, 2, 1] + terms[:, 3, 0]\n return torch.stack((w, x, y, z), dim=1).view(original_shape)\n\n\ndef qrot(q, v):\n \"\"\"\n Rotate vector(s) v about the rotation described by quaternion(s) q.\n Expects a tensor of shape (*, 4) for q and a tensor of shape (*, 3) for v,\n where * denotes any number of dimensions.\n Returns a tensor of shape (*, 3).\n \"\"\"\n assert q.shape[-1] == 4\n assert v.shape[-1] == 3\n assert q.shape[:-1] == v.shape[:-1]\n\n original_shape = list(v.shape)\n # print(q.shape)\n q = q.contiguous().view(-1, 4)\n v = v.contiguous().view(-1, 3)\n\n qvec = q[:, 1:]\n uv = torch.cross(qvec, v, dim=1)\n uuv = torch.cross(qvec, uv, dim=1)\n return (v + 2 * (q[:, :1] * uv + uuv)).view(original_shape)\n\n\ndef qeuler(q, order, epsilon=0, deg=True):\n \"\"\"\n Convert quaternion(s) q to Euler angles.\n Expects a tensor of shape (*, 4), where * denotes any number of dimensions.\n Returns a tensor of shape (*, 3).\n \"\"\"\n assert q.shape[-1] == 4\n\n original_shape = list(q.shape)\n original_shape[-1] = 3\n q = q.view(-1, 4)\n\n q0 = q[:, 0]\n q1 = q[:, 1]\n q2 = q[:, 2]\n q3 = q[:, 3]\n\n if order == \"xyz\":\n x = torch.atan2(2 * (q0 * q1 - q2 * q3), 1 - 2 * (q1 * q1 + q2 * q2))\n y = torch.asin(torch.clamp(2 * (q1 * q3 + q0 * q2), -1 + epsilon, 1 - epsilon))\n z = torch.atan2(2 * (q0 * q3 - q1 * q2), 1 - 2 * (q2 * q2 + q3 * q3))\n elif order == \"yzx\":\n x = torch.atan2(2 * (q0 * q1 - q2 * q3), 1 - 2 * (q1 * q1 + q3 * q3))\n y = torch.atan2(2 * (q0 * q2 - q1 * q3), 1 - 2 * (q2 * q2 + q3 * q3))\n z = torch.asin(torch.clamp(2 * (q1 * q2 + q0 * q3), -1 + epsilon, 1 - epsilon))\n elif order == \"zxy\":\n x = torch.asin(torch.clamp(2 * (q0 * q1 + q2 * q3), -1 + epsilon, 1 - epsilon))\n y = torch.atan2(2 * (q0 * q2 - q1 * q3), 1 - 2 * (q1 * q1 + q2 * q2))\n z = torch.atan2(2 * (q0 * q3 - q1 * q2), 1 - 2 * (q1 * q1 + q3 * q3))\n elif order == \"xzy\":\n x = torch.atan2(2 * (q0 * q1 + q2 * q3), 1 - 2 * (q1 * q1 + q3 * q3))\n y = torch.atan2(2 * (q0 * q2 + q1 * q3), 1 - 2 * (q2 * q2 + q3 * q3))\n z = torch.asin(torch.clamp(2 * (q0 * q3 - q1 * q2), -1 + epsilon, 1 - epsilon))\n elif order == \"yxz\":\n x = torch.asin(torch.clamp(2 * (q0 * q1 - q2 * q3), -1 + epsilon, 1 - epsilon))\n y = torch.atan2(2 * (q1 * q3 + q0 * q2), 1 - 2 * (q1 * q1 + q2 * q2))\n z = torch.atan2(2 * (q1 * q2 + q0 * q3), 1 - 2 * (q1 * q1 + q3 * q3))\n elif order == \"zyx\":\n x = torch.atan2(2 * (q0 * q1 + q2 * q3), 1 - 2 * (q1 * q1 + q2 * q2))\n y = torch.asin(torch.clamp(2 * (q0 * q2 - q1 * q3), -1 + epsilon, 1 - epsilon))\n z = torch.atan2(2 * (q0 * q3 + q1 * q2), 1 - 2 * (q2 * q2 + q3 * q3))\n else:\n raise\n\n if deg:\n return torch.stack((x, y, z), dim=1).view(original_shape) * 180 / np.pi\n else:\n return torch.stack((x, y, z), dim=1).view(original_shape)\n\n\n# Numpy-backed implementations\n\n\ndef qmul_np(q, r):\n q = torch.from_numpy(q).contiguous().float()\n r = torch.from_numpy(r).contiguous().float()\n return qmul(q, r).numpy()\n\n\ndef qrot_np(q, v):\n q = torch.from_numpy(q).contiguous().float()\n v = torch.from_numpy(v).contiguous().float()\n return qrot(q, v).numpy()\n\n\ndef qeuler_np(q, order, epsilon=0, use_gpu=False):\n if use_gpu:\n q = torch.from_numpy(q).cuda().float()\n return qeuler(q, order, epsilon).cpu().numpy()\n else:\n q = torch.from_numpy(q).contiguous().float()\n return qeuler(q, order, epsilon).numpy()\n\n\ndef qfix(q):\n \"\"\"\n Enforce quaternion continuity across the time dimension by selecting\n the representation (q or -q) with minimal distance (or, equivalently, maximal dot product)\n between two consecutive frames.\n\n Expects a tensor of shape (L, J, 4), where L is the sequence length and J is the number of joints.\n Returns a tensor of the same shape.\n \"\"\"\n assert len(q.shape) == 3\n assert q.shape[-1] == 4\n\n result = q.copy()\n dot_products = np.sum(q[1:] * q[:-1], axis=2)\n mask = dot_products < 0\n mask = (np.cumsum(mask, axis=0) % 2).astype(bool)\n result[1:][mask] *= -1\n return result\n\n\ndef euler2quat(e, order, deg=True):\n \"\"\"\n Convert Euler angles to quaternions.\n \"\"\"\n assert e.shape[-1] == 3\n\n original_shape = list(e.shape)\n original_shape[-1] = 4\n\n e = e.view(-1, 3)\n\n ## if euler angles in degrees\n if deg:\n e = e * np.pi / 180.0\n\n x = e[:, 0]\n y = e[:, 1]\n z = e[:, 2]\n\n rx = torch.stack((torch.cos(x / 2), torch.sin(x / 2), torch.zeros_like(x), torch.zeros_like(x)), dim=1)\n ry = torch.stack((torch.cos(y / 2), torch.zeros_like(y), torch.sin(y / 2), torch.zeros_like(y)), dim=1)\n rz = torch.stack((torch.cos(z / 2), torch.zeros_like(z), torch.zeros_like(z), torch.sin(z / 2)), dim=1)\n\n result = None\n for coord in order:\n if coord == \"x\":\n r = rx\n elif coord == \"y\":\n r = ry\n elif coord == \"z\":\n r = rz\n else:\n raise\n if result is None:\n result = r\n else:\n result = qmul(result, r)\n\n # Reverse antipodal representation to have a non-negative \"w\"\n if order in [\"xyz\", \"yzx\", \"zxy\"]:\n result *= -1\n\n return result.view(original_shape)\n\n\ndef expmap_to_quaternion(e):\n \"\"\"\n Convert axis-angle rotations (aka exponential maps) to quaternions.\n Stable formula from \"Practical Parameterization of Rotations Using the Exponential Map\".\n Expects a tensor of shape (*, 3), where * denotes any number of dimensions.\n Returns a tensor of shape (*, 4).\n \"\"\"\n assert e.shape[-1] == 3\n\n original_shape = list(e.shape)\n original_shape[-1] = 4\n e = e.reshape(-1, 3)\n\n theta = np.linalg.norm(e, axis=1).reshape(-1, 1)\n w = np.cos(0.5 * theta).reshape(-1, 1)\n xyz = 0.5 * np.sinc(0.5 * theta / np.pi) * e\n return np.concatenate((w, xyz), axis=1).reshape(original_shape)\n\n\ndef euler_to_quaternion(e, order):\n \"\"\"\n Convert Euler angles to quaternions.\n \"\"\"\n assert e.shape[-1] == 3\n\n original_shape = list(e.shape)\n original_shape[-1] = 4\n\n e = e.reshape(-1, 3)\n\n x = e[:, 0]\n y = e[:, 1]\n z = e[:, 2]\n\n rx = np.stack((np.cos(x / 2), np.sin(x / 2), np.zeros_like(x), np.zeros_like(x)), axis=1)\n ry = np.stack((np.cos(y / 2), np.zeros_like(y), np.sin(y / 2), np.zeros_like(y)), axis=1)\n rz = np.stack((np.cos(z / 2), np.zeros_like(z), np.zeros_like(z), np.sin(z / 2)), axis=1)\n\n result = None\n for coord in order:\n if coord == \"x\":\n r = rx\n elif coord == \"y\":\n r = ry\n elif coord == \"z\":\n r = rz\n else:\n raise\n if result is None:\n result = r\n else:\n result = qmul_np(result, r)\n\n # Reverse antipodal representation to have a non-negative \"w\"\n if order in [\"xyz\", \"yzx\", \"zxy\"]:\n result *= -1\n\n return result.reshape(original_shape)\n\n\ndef quaternion_to_matrix(quaternions):\n \"\"\"\n Convert rotations given as quaternions to rotation matrices.\n Args:\n quaternions: quaternions with real part first,\n as tensor of shape (..., 4).\n Returns:\n Rotation matrices as tensor of shape (..., 3, 3).\n \"\"\"\n r, i, j, k = torch.unbind(quaternions, -1)\n two_s = 2.0 / (quaternions * quaternions).sum(-1)\n\n o = torch.stack(\n (\n 1 - two_s * (j * j + k * k),\n two_s * (i * j - k * r),\n two_s * (i * k + j * r),\n two_s * (i * j + k * r),\n 1 - two_s * (i * i + k * k),\n two_s * (j * k - i * r),\n two_s * (i * k - j * r),\n two_s * (j * k + i * r),\n 1 - two_s * (i * i + j * j),\n ),\n -1,\n )\n return o.reshape(quaternions.shape[:-1] + (3, 3))\n\n\ndef quaternion_to_matrix_np(quaternions):\n q = torch.from_numpy(quaternions).contiguous().float()\n return quaternion_to_matrix(q).numpy()\n\n\ndef quaternion_to_cont6d_np(quaternions):\n rotation_mat = quaternion_to_matrix_np(quaternions)\n cont_6d = np.concatenate([rotation_mat[..., 0], rotation_mat[..., 1]], axis=-1)\n return cont_6d\n\n\ndef quaternion_to_cont6d(quaternions):\n rotation_mat = quaternion_to_matrix(quaternions)\n cont_6d = torch.cat([rotation_mat[..., 0], rotation_mat[..., 1]], dim=-1)\n return cont_6d\n\n\ndef cont6d_to_matrix(cont6d):\n assert cont6d.shape[-1] == 6, \"The last dimension must be 6\"\n x_raw = cont6d[..., 0:3]\n y_raw = cont6d[..., 3:6]\n\n x = x_raw / torch.norm(x_raw, dim=-1, keepdim=True)\n z = torch.cross(x, y_raw, dim=-1)\n z = z / torch.norm(z, dim=-1, keepdim=True)\n\n y = torch.cross(z, x, dim=-1)\n\n x = x[..., None]\n y = y[..., None]\n z = z[..., None]\n\n mat = torch.cat([x, y, z], dim=-1)\n return mat\n\n\ndef cont6d_to_matrix_np(cont6d):\n q = torch.from_numpy(cont6d).contiguous().float()\n return cont6d_to_matrix(q).numpy()\n\n\ndef qpow(q0, t, dtype=torch.float):\n \"\"\"q0 : tensor of quaternions\n t: tensor of powers\n \"\"\"\n q0 = qnormalize(q0)\n theta0 = torch.acos(q0[..., :1])\n\n ## if theta0 is close to zero, add epsilon to avoid NaNs\n mask = (theta0 <= 10e-10) * (theta0 >= -10e-10)\n mask = mask.float()\n theta0 = (1 - mask) * theta0 + mask * 10e-10\n v0 = q0[..., 1:] / torch.sin(theta0)\n\n if isinstance(t, torch.Tensor):\n # Do not check here\n q = torch.zeros(t.shape + q0.shape, device=q0.device)\n theta = t.view(-1, 1) * theta0.view(1, -1)\n else: ## if t is a number\n q = torch.zeros(q0.shape, device=q0.device)\n theta = t * theta0\n\n q[..., :1] = torch.cos(theta)\n q[..., 1:] = v0 * torch.sin(theta)\n\n return q.to(dtype)\n\n\ndef qslerp(q0, q1, t):\n \"\"\"\n q0: starting quaternion\n q1: ending quaternion\n t: array of points along the way\n\n Returns:\n Tensor of Slerps: t.shape + q0.shape\n \"\"\"\n\n q0 = qnormalize(q0)\n q1 = qnormalize(q1)\n q_ = qpow(qmul(q1, qinv(q0)), t)\n\n return qmul(q_, q0)\n\n\ndef qbetween(v0, v1):\n \"\"\"\n find the quaternion used to rotate v0 to v1\n \"\"\"\n assert v0.shape[-1] == 3, \"v0 must be of the shape (*, 3)\"\n assert v1.shape[-1] == 3, \"v1 must be of the shape (*, 3)\"\n\n v = torch.cross(v0, v1, dim=-1)\n\n w = torch.sqrt((v0**2).sum(dim=-1, keepdim=True) * (v1**2).sum(dim=-1, keepdim=True)) + (v0 * v1).sum(\n dim=-1, keepdim=True\n )\n y_vec = torch.zeros_like(v)\n y_vec[..., 1] = 1.0\n # if v0 is (0, 0, -1), v1 is (0, 0, 1), v will be 0 and w will also be 0 -> this makes below situation comes v=1 w = 2\n mask = v.norm(dim=-1) == 0\n # if v0 is (0, 0, 1), v1 is (0, 0, 1), v will be 0 and w will be 2 -> do nothing\n mask2 = w.sum(dim=-1).abs() <= 1e-4\n mask = torch.logical_and(mask, mask2)\n v[mask] = y_vec[mask]\n\n return qnormalize(torch.cat([w, v], dim=-1))\n\n\ndef qbetween_np(v0, v1):\n \"\"\"\n find the quaternion used to rotate v0 to v1\n \"\"\"\n assert v0.shape[-1] == 3, \"v0 must be of the shape (*, 3)\"\n assert v1.shape[-1] == 3, \"v1 must be of the shape (*, 3)\"\n\n v0 = torch.from_numpy(v0).float()\n v1 = torch.from_numpy(v1).float()\n return qbetween(v0, v1).numpy()\n\n\ndef lerp(p0, p1, t):\n if not isinstance(t, torch.Tensor):\n t = torch.Tensor([t])\n\n new_shape = t.shape + p0.shape\n new_view_t = t.shape + torch.Size([1] * len(p0.shape))\n new_view_p = torch.Size([1] * len(t.shape)) + p0.shape\n p0 = p0.view(new_view_p).expand(new_shape)\n p1 = p1.view(new_view_p).expand(new_shape)\n t = t.view(new_view_t).expand(new_shape)\n\n return p0 + t * (p1 - p0)\n" }, { "path": "eval/GVHMR/hmr4d/utils/geo/transforms.py", "content": "import torch\n\n\ndef axis_rotate_to_matrix(angle, axis=\"x\"):\n \"\"\"Get rotation matrix for rotating around one axis\n Args:\n angle: (N, 1)\n Returns:\n R: (N, 3, 3)\n \"\"\"\n if isinstance(angle, float):\n angle = torch.tensor([angle], dtype=torch.float)\n\n c = torch.cos(angle)\n s = torch.sin(angle)\n z = torch.zeros_like(angle)\n o = torch.ones_like(angle)\n if axis == \"x\":\n R = torch.stack([o, z, z, z, c, -s, z, s, c], dim=1).view(-1, 3, 3)\n elif axis == \"y\":\n R = torch.stack([c, z, s, z, o, z, -s, z, c], dim=1).view(-1, 3, 3)\n else:\n assert axis == \"z\"\n R = torch.stack([c, -s, z, s, c, z, z, z, o], dim=1).view(-1, 3, 3)\n return R\n" }, { "path": "eval/GVHMR/hmr4d/utils/geo_transform.py", "content": "import numpy as np\nimport cv2\nimport torch\nimport torch.nn.functional as F\nfrom pytorch3d.transforms import so3_exp_map, so3_log_map\nfrom pytorch3d.transforms import matrix_to_quaternion, quaternion_to_axis_angle, matrix_to_rotation_6d\nimport pytorch3d.ops.knn as knn\nfrom hmr4d.utils.pylogger import Log\nfrom pytorch3d.transforms import euler_angles_to_matrix\nimport hmr4d.utils.matrix as matrix\nfrom einops import einsum, rearrange, repeat\nfrom hmr4d.utils.geo.quaternion import qbetween\n\n\ndef homo_points(points):\n \"\"\"\n Args:\n points: (..., C)\n Returns: (..., C+1), with 1 padded\n \"\"\"\n return F.pad(points, [0, 1], value=1.0)\n\n\ndef apply_Ts_on_seq_points(points, Ts):\n \"\"\"\n perform translation matrix on related point\n Args:\n points: (..., N, 3)\n Ts: (..., N, 4, 4)\n Returns: (..., N, 3)\n \"\"\"\n points = torch.torch.einsum(\"...ki,...i->...k\", Ts[..., :3, :3], points) + Ts[..., :3, 3]\n return points\n\n\ndef apply_T_on_points(points, T):\n \"\"\"\n Args:\n points: (..., N, 3)\n T: (..., 4, 4)\n Returns: (..., N, 3)\n \"\"\"\n points_T = torch.einsum(\"...ki,...ji->...jk\", T[..., :3, :3], points) + T[..., None, :3, 3]\n return points_T\n\n\ndef T_transforms_points(T, points, pattern):\n \"\"\"manual mode of apply_T_on_points\n T: (..., 4, 4)\n points: (..., 3)\n pattern: \"... c d, ... d -> ... c\"\n \"\"\"\n return einsum(T, homo_points(points), pattern)[..., :3]\n\n\ndef project_p2d(points, K=None, is_pinhole=True):\n \"\"\"\n Args:\n points: (..., (N), 3)\n K: (..., 3, 3)\n Returns: shape is similar to points but without z\n \"\"\"\n points = points.clone()\n if is_pinhole:\n z = points[..., [-1]]\n z.masked_fill_(z.abs() < 1e-6, 1e-6)\n points_proj = points / z\n else: # orthogonal\n points_proj = F.pad(points[..., :2], (0, 1), value=1)\n\n if K is not None:\n # Handle N\n if len(points_proj.shape) == len(K.shape):\n p2d_h = torch.einsum(\"...ki,...ji->...jk\", K, points_proj)\n else:\n p2d_h = torch.einsum(\"...ki,...i->...k\", K, points_proj)\n else:\n p2d_h = points_proj[..., :2]\n\n return p2d_h[..., :2]\n\n\ndef gen_uv_from_HW(H, W, device=\"cpu\"):\n \"\"\"Returns: (H, W, 2), as float. Note: uv not ij\"\"\"\n grid_v, grid_u = torch.meshgrid(torch.arange(H), torch.arange(W))\n return (\n torch.stack(\n [grid_u, grid_v],\n dim=-1,\n )\n .float()\n .to(device)\n ) # (H, W, 2)\n\n\ndef unproject_p2d(uv, z, K):\n \"\"\"we assume a pinhole camera for unprojection\n uv: (B, N, 2)\n z: (B, N, 1)\n K: (B, 3, 3)\n Returns: (B, N, 3)\n \"\"\"\n xy_atz1 = (uv - K[:, None, :2, 2]) / K[:, None, [0, 1], [0, 1]] # (B, N, 2)\n xyz = torch.cat([xy_atz1 * z, z], dim=-1)\n return xyz\n\n\ndef cvt_p2d_from_i_to_c(uv, K):\n \"\"\"\n Args:\n uv: (..., 2) or (..., N, 2)\n K: (..., 3, 3)\n Returns: the same shape as input uv\n \"\"\"\n if len(uv.shape) == len(K.shape):\n xy = (uv - K[..., None, :2, 2]) / K[..., None, [0, 1], [0, 1]]\n else: # without N\n xy = (uv - K[..., :2, 2]) / K[..., [0, 1], [0, 1]]\n return xy\n\n\ndef cvt_to_bi01_p2d(p2d, bbx_lurb):\n \"\"\"\n p2d: (..., (N), 2)\n bbx_lurb: (..., 4)\n \"\"\"\n if len(p2d.shape) == len(bbx_lurb.shape) + 1:\n bbx_lurb = bbx_lurb[..., None, :]\n\n bbx_wh = bbx_lurb[..., 2:] - bbx_lurb[..., :2]\n bi01_p2d = (p2d - bbx_lurb[..., :2]) / bbx_wh\n return bi01_p2d\n\n\ndef cvt_from_bi01_p2d(bi01_p2d, bbx_lurb):\n \"\"\"Use bbx_lurb to resize bi01_p2d to p2d (image-coordinates)\n Args:\n p2d: (..., 2) or (..., N, 2)\n bbx_lurb: (..., 4)\n Returns:\n p2d: shape is the same as input\n \"\"\"\n bbx_wh = bbx_lurb[..., 2:] - bbx_lurb[..., :2] # (..., 2)\n if len(bi01_p2d.shape) == len(bbx_wh.shape) + 1:\n p2d = (bi01_p2d * bbx_wh.unsqueeze(-2)) + bbx_lurb[..., None, :2]\n else:\n p2d = (bi01_p2d * bbx_wh) + bbx_lurb[..., :2]\n return p2d\n\n\ndef cvt_p2d_from_bi01_to_c(bi01, bbxs_lurb, Ks):\n \"\"\"\n Args:\n bi01: (..., (N), 2), value in range (0,1), the point in the bbx image\n bbxs_lurb: (..., 4)\n Ks: (..., 3, 3)\n Returns:\n c: (..., (N), 2)\n \"\"\"\n i = cvt_from_bi01_p2d(bi01, bbxs_lurb)\n c = cvt_p2d_from_i_to_c(i, Ks)\n return c\n\n\ndef cvt_p2d_from_pm1_to_i(p2d_pm1, bbx_xys):\n \"\"\"\n Args:\n p2d_pm1: (..., (N), 2), value in range (-1,1), the point in the bbx image\n bbx_xys: (..., 3)\n Returns:\n p2d: (..., (N), 2)\n \"\"\"\n return bbx_xys[..., :2] + p2d_pm1 * bbx_xys[..., [2]] / 2\n\n\ndef uv2l_index(uv, W):\n return uv[..., 0] + uv[..., 1] * W\n\n\ndef l2uv_index(l, W):\n v = torch.div(l, W, rounding_mode=\"floor\")\n u = l % W\n return torch.stack([u, v], dim=-1)\n\n\ndef transform_mat(R, t):\n \"\"\"\n Args:\n R: Bx3x3 array of a batch of rotation matrices\n t: Bx3x(1) array of a batch of translation vectors\n Returns:\n T: Bx4x4 Transformation matrix\n \"\"\"\n # No padding left or right, only add an extra row\n if len(R.shape) > len(t.shape):\n t = t[..., None]\n return torch.cat([F.pad(R, [0, 0, 0, 1]), F.pad(t, [0, 0, 0, 1], value=1)], dim=-1)\n\n\ndef axis_angle_to_matrix_exp_map(aa):\n \"\"\"use pytorch3d so3_exp_map\n Args:\n aa: (*, 3)\n Returns:\n R: (*, 3, 3)\n \"\"\"\n print(\"Use pytorch3d.transforms.axis_angle_to_matrix instead!!!\")\n ori_shape = aa.shape[:-1]\n return so3_exp_map(aa.reshape(-1, 3)).reshape(*ori_shape, 3, 3)\n\n\ndef matrix_to_axis_angle_log_map(R):\n \"\"\"use pytorch3d so3_log_map\n Args:\n aa: (*, 3, 3)\n Returns:\n R: (*, 3)\n \"\"\"\n print(\"WARINING! I met singularity problem with this function, use matrix_to_axis_angle instead!\")\n ori_shape = R.shape[:-2]\n return so3_log_map(R.reshape(-1, 3, 3)).reshape(*ori_shape, 3)\n\n\ndef matrix_to_axis_angle(R):\n \"\"\"use pytorch3d so3_log_map\n Args:\n aa: (*, 3, 3)\n Returns:\n R: (*, 3)\n \"\"\"\n return quaternion_to_axis_angle(matrix_to_quaternion(R))\n\n\ndef ransac_PnP(K, pts_2d, pts_3d, err_thr=10):\n \"\"\"solve pnp\"\"\"\n dist_coeffs = np.zeros(shape=[8, 1], dtype=\"float64\")\n\n pts_2d = np.ascontiguousarray(pts_2d.astype(np.float64))\n pts_3d = np.ascontiguousarray(pts_3d.astype(np.float64))\n K = K.astype(np.float64)\n\n try:\n _, rvec, tvec, inliers = cv2.solvePnPRansac(\n pts_3d, pts_2d, K, dist_coeffs, reprojectionError=err_thr, iterationsCount=10000, flags=cv2.SOLVEPNP_EPNP\n )\n\n rotation = cv2.Rodrigues(rvec)[0]\n\n pose = np.concatenate([rotation, tvec], axis=-1)\n pose_homo = np.concatenate([pose, np.array([[0, 0, 0, 1]])], axis=0)\n\n inliers = [] if inliers is None else inliers\n\n return pose, pose_homo, inliers\n except cv2.error:\n print(\"CV ERROR\")\n return np.eye(4)[:3], np.eye(4), []\n\n\ndef ransac_PnP_batch(K_raw, pts_2d, pts_3d, err_thr=10):\n fit_R, fit_t = [], []\n for b in range(K_raw.shape[0]):\n pose, _, inliers = ransac_PnP(K_raw[b], pts_2d[b], pts_3d[b], err_thr=err_thr)\n fit_R.append(pose[:3, :3])\n fit_t.append(pose[:3, 3])\n fit_R = np.stack(fit_R, axis=0)\n fit_t = np.stack(fit_t, axis=0)\n return fit_R, fit_t\n\n\ndef triangulate_point(Ts_w2c, c_p2d, **kwargs):\n from hmr4d.utils.geo.triangulation import triangulate_persp\n\n print(\"Deprecated, please import from hmr4d.utils.geo.triangulation\")\n return triangulate_persp(Ts_w2c, c_p2d, **kwargs)\n\n\ndef triangulate_point_ortho(Ts_w2c, c_p2d, **kwargs):\n from hmr4d.utils.geo.triangulation import triangulate_ortho\n\n print(\"Deprecated, please import from hmr4d.utils.geo.triangulation\")\n return triangulate_ortho(Ts_w2c, c_p2d, **kwargs)\n\n\ndef get_nearby_points(points, query_verts, padding=0.0, p=1):\n \"\"\"\n points: (S, 3)\n query_verts: (V, 3)\n \"\"\"\n if p == 1:\n max_xyz = query_verts.max(0)[0] + padding\n min_xyz = query_verts.min(0)[0] - padding\n idx = (((points - min_xyz) > 0).all(dim=-1) * ((points - max_xyz) < 0).all(dim=-1)).nonzero().squeeze(-1)\n nearby_points = points[idx]\n elif p == 2:\n squared_dist, _, _ = knn.knn_points(points[None], query_verts[None], K=1, return_nn=False)\n mask = squared_dist[0, :, 0] < padding**2 # (S,)\n nearby_points = points[mask]\n\n return nearby_points\n\n\ndef unproj_bbx_to_fst(bbx_lurb, K, near_z=0.5, far_z=12.5):\n B = bbx_lurb.size(0)\n uv = bbx_lurb[:, [[0, 1], [2, 1], [2, 3], [0, 3], [0, 1], [2, 1], [2, 3], [0, 3]]]\n if isinstance(near_z, float):\n z = uv.new([near_z] * 4 + [far_z] * 4).reshape(1, 8, 1).repeat(B, 1, 1)\n else:\n z = torch.cat([near_z[:, None, None].repeat(1, 4, 1), far_z[:, None, None].repeat(1, 4, 1)], dim=1)\n c_frustum_points = unproject_p2d(uv, z, K) # (B, 8, 3)\n return c_frustum_points\n\n\ndef convert_bbx_xys_to_lurb(bbx_xys):\n \"\"\"\n Args: bbx_xys (..., 3) -> bbx_lurb (..., 4)\n \"\"\"\n size = bbx_xys[..., 2:]\n center = bbx_xys[..., :2]\n lurb = torch.cat([center - size / 2, center + size / 2], dim=-1)\n return lurb\n\n\ndef convert_lurb_to_bbx_xys(bbx_lurb):\n \"\"\"\n Args: bbx_lurb (..., 4) -> bbx_xys (..., 3) be aware that it is squared\n \"\"\"\n size = (bbx_lurb[..., 2:] - bbx_lurb[..., :2]).max(-1, keepdim=True)[0]\n center = (bbx_lurb[..., :2] + bbx_lurb[..., 2:]) / 2\n return torch.cat([center, size], dim=-1)\n\n\n# ================== AZ/AY Transformations ================== #\n\n\ndef compute_T_ayf2az(joints, inverse=False):\n \"\"\"\n Args:\n joints: (B, J, 3), in the start-frame, az-coordinate\n Returns:\n if inverse == False:\n T_af2az: (B, 4, 4)\n else :\n T_az2af: (B, 4, 4)\n \"\"\"\n\n t_ayf2az = joints[:, 0, :].detach().clone()\n t_ayf2az[:, 2] = 0 # do not modify z\n\n RL_xy_h = joints[:, 1, [0, 1]] - joints[:, 2, [0, 1]] # (B, 2), hip point to left side\n RL_xy_s = joints[:, 16, [0, 1]] - joints[:, 17, [0, 1]] # (B, 2), shoulder point to left side\n RL_xy = RL_xy_h + RL_xy_s\n I_mask = RL_xy.pow(2).sum(-1) < 1e-4 # do not rotate, when can't decided the face direction\n if I_mask.sum() > 0:\n Log.warn(\"{} samples can't decide the face direction\".format(I_mask.sum()))\n x_dir = F.pad(F.normalize(RL_xy, 2, -1), (0, 1), value=0) # (B, 3)\n y_dir = torch.zeros_like(x_dir)\n y_dir[..., 2] = 1\n z_dir = torch.cross(x_dir, y_dir, dim=-1)\n R_ayf2az = torch.stack([x_dir, y_dir, z_dir], dim=-1) # (B, 3, 3)\n R_ayf2az[I_mask] = torch.eye(3).to(R_ayf2az)\n\n if inverse:\n R_az2ayf = R_ayf2az.transpose(1, 2) # (B, 3, 3)\n t_az2ayf = -einsum(R_ayf2az, t_ayf2az, \"b i j , b i -> b j\") # (B, 3)\n return transform_mat(R_az2ayf, t_az2ayf)\n else:\n return transform_mat(R_ayf2az, t_ayf2az)\n\n\ndef compute_T_ayfz2ay(joints, inverse=False):\n \"\"\"\n Args:\n joints: (B, J, 3), in the start-frame, ay-coordinate\n Returns:\n if inverse == False:\n T_ayfz2ay: (B, 4, 4)\n else :\n T_ay2ayfz: (B, 4, 4)\n \"\"\"\n t_ayfz2ay = joints[:, 0, :].detach().clone()\n t_ayfz2ay[:, 1] = 0 # do not modify y\n\n RL_xz_h = joints[:, 1, [0, 2]] - joints[:, 2, [0, 2]] # (B, 2), hip point to left side\n RL_xz_s = joints[:, 16, [0, 2]] - joints[:, 17, [0, 2]] # (B, 2), shoulder point to left side\n RL_xz = RL_xz_h + RL_xz_s\n I_mask = RL_xz.pow(2).sum(-1) < 1e-4 # do not rotate, when can't decided the face direction\n if I_mask.sum() > 0:\n Log.warn(\"{} samples can't decide the face direction\".format(I_mask.sum()))\n\n x_dir = torch.zeros_like(t_ayfz2ay) # (B, 3)\n x_dir[:, [0, 2]] = F.normalize(RL_xz, 2, -1)\n y_dir = torch.zeros_like(x_dir)\n y_dir[..., 1] = 1 # (B, 3)\n z_dir = torch.cross(x_dir, y_dir, dim=-1)\n R_ayfz2ay = torch.stack([x_dir, y_dir, z_dir], dim=-1) # (B, 3, 3)\n R_ayfz2ay[I_mask] = torch.eye(3).to(R_ayfz2ay)\n\n if inverse:\n R_ay2ayfz = R_ayfz2ay.transpose(1, 2)\n t_ay2ayfz = -einsum(R_ayfz2ay, t_ayfz2ay, \"b i j , b i -> b j\")\n return transform_mat(R_ay2ayfz, t_ay2ayfz)\n else:\n return transform_mat(R_ayfz2ay, t_ayfz2ay)\n\n\ndef compute_T_ay2ayrot(joints):\n \"\"\"\n Args:\n joints: (B, J, 3), in the start-frame, ay-coordinate\n Returns:\n T_ay2ayrot: (B, 4, 4)\n \"\"\"\n t_ayrot2ay = joints[:, 0, :].detach().clone()\n t_ayrot2ay[:, 1] = 0 # do not modify y\n\n B = joints.shape[0]\n euler_angle = torch.zeros((B, 3), device=joints.device)\n yrot_angle = torch.rand((B,), device=joints.device) * 2 * torch.pi\n euler_angle[:, 0] = yrot_angle\n R_ay2ayrot = euler_angles_to_matrix(euler_angle, \"YXZ\") # (B, 3, 3)\n\n R_ayrot2ay = R_ay2ayrot.transpose(1, 2)\n t_ay2ayrot = -einsum(R_ayrot2ay, t_ayrot2ay, \"b i j , b i -> b j\")\n return transform_mat(R_ay2ayrot, t_ay2ayrot)\n\n\ndef compute_root_quaternion_ay(joints):\n \"\"\"\n Args:\n joints: (B, J, 3), in the start-frame, ay-coordinate\n Returns:\n root_quat: (B, 4) from z-axis to fz\n \"\"\"\n joints_shape = joints.shape\n joints = joints.reshape((-1,) + joints_shape[-2:])\n t_ayfz2ay = joints[:, 0, :].detach().clone()\n t_ayfz2ay[:, 1] = 0 # do not modify y\n\n RL_xz_h = joints[:, 1, [0, 2]] - joints[:, 2, [0, 2]] # (B, 2), hip point to left side\n RL_xz_s = joints[:, 16, [0, 2]] - joints[:, 17, [0, 2]] # (B, 2), shoulder point to left side\n RL_xz = RL_xz_h + RL_xz_s\n I_mask = RL_xz.pow(2).sum(-1) < 1e-4 # do not rotate, when can't decided the face direction\n if I_mask.sum() > 0:\n Log.warn(\"{} samples can't decide the face direction\".format(I_mask.sum()))\n\n x_dir = torch.zeros_like(t_ayfz2ay) # (B, 3)\n x_dir[:, [0, 2]] = F.normalize(RL_xz, 2, -1)\n y_dir = torch.zeros_like(x_dir)\n y_dir[..., 1] = 1 # (B, 3)\n z_dir = torch.cross(x_dir, y_dir, dim=-1)\n\n z_dir[..., 2] += 1e-9\n pos_z_vec = torch.tensor([0, 0, 1]).to(joints.device).float() # (3,)\n root_quat = qbetween(pos_z_vec[None], z_dir) # (B, 4)\n root_quat = root_quat.reshape(joints_shape[:-2] + (4,))\n return root_quat\n\n\n# ================== Transformations between two sets of features ================== #\n\n\ndef similarity_transform_batch(S1, S2):\n \"\"\"\n Computes a similarity transform (sR, t) that solves the orthogonal Procrutes problem.\n Args:\n S1, S2: (*, L, 3)\n \"\"\"\n assert S1.shape == S2.shape\n S_shape = S1.shape\n S1 = S1.reshape(-1, *S_shape[-2:])\n S2 = S2.reshape(-1, *S_shape[-2:])\n\n S1 = S1.transpose(-2, -1)\n S2 = S2.transpose(-2, -1)\n\n # --- The code is borrowed from WHAM ---\n # 1. Remove mean.\n mu1 = S1.mean(axis=-1, keepdims=True) # axis is along N, S1(B, 3, N)\n mu2 = S2.mean(axis=-1, keepdims=True)\n\n X1 = S1 - mu1\n X2 = S2 - mu2\n\n # 2. Compute variance of X1 used for scale.\n var1 = torch.sum(X1**2, dim=1).sum(dim=1)\n\n # 3. The outer product of X1 and X2.\n K = X1.bmm(X2.permute(0, 2, 1))\n\n # 4. Solution that Maximizes trace(R'K) is R=U*V', where U, V are\n # singular vectors of K.\n U, s, V = torch.svd(K)\n\n # Construct Z that fixes the orientation of R to get det(R)=1.\n Z = torch.eye(U.shape[1], device=S1.device).unsqueeze(0)\n Z = Z.repeat(U.shape[0], 1, 1)\n Z[:, -1, -1] *= torch.sign(torch.det(U.bmm(V.permute(0, 2, 1))))\n\n # Construct R.\n R = V.bmm(Z.bmm(U.permute(0, 2, 1)))\n\n # 5. Recover scale.\n scale = torch.cat([torch.trace(x).unsqueeze(0) for x in R.bmm(K)]) / var1\n\n # 6. Recover translation.\n t = mu2 - (scale.unsqueeze(-1).unsqueeze(-1) * (R.bmm(mu1)))\n\n # -------\n # reshape back\n # sR = scale[:, None, None] * R\n # sR = sR.reshape(*S_shape[:-2], 3, 3)\n scale = scale.reshape(*S_shape[:-2], 1, 1)\n R = R.reshape(*S_shape[:-2], 3, 3)\n t = t.reshape(*S_shape[:-2], 3, 1)\n\n return (scale, R), t\n\n\ndef kabsch_algorithm_batch(X1, X2):\n \"\"\"\n Computes a rigid transform (R, t)\n Args:\n X1, X2: (*, L, 3)\n \"\"\"\n assert X1.shape == X2.shape\n X_shape = X1.shape\n X1 = X1.reshape(-1, *X_shape[-2:])\n X2 = X2.reshape(-1, *X_shape[-2:])\n\n # 1. 计算质心\n centroid_X1 = torch.mean(X1, dim=-2, keepdim=True)\n centroid_X2 = torch.mean(X2, dim=-2, keepdim=True)\n\n # 2. 去中心化\n X1_centered = X1 - centroid_X1\n X2_centered = X2 - centroid_X2\n\n # 3. 计算协方差矩阵\n H = torch.matmul(X1_centered.transpose(-2, -1), X2_centered)\n\n # 4. 奇异值分解\n U, S, Vt = torch.linalg.svd(H)\n\n # 5. 计算旋转矩阵\n R = torch.matmul(Vt.transpose(-2, -1), U.transpose(-2, -1))\n\n # 修正反射矩阵\n d = (torch.det(R) < 0).unsqueeze(-1).unsqueeze(-1)\n Vt = torch.where(d, -Vt, Vt)\n R = torch.matmul(Vt.transpose(-2, -1), U.transpose(-2, -1))\n\n # 6. 计算平移向量\n t = centroid_X2.transpose(-2, -1) - torch.matmul(R, centroid_X1.transpose(-2, -1))\n\n # -------\n # reshape back\n R = R.reshape(*X_shape[:-2], 3, 3)\n t = t.reshape(*X_shape[:-2], 3, 1)\n\n return R, t\n\n\n# ===== WHAM cam_angvel ===== #\n\n\ndef compute_cam_angvel(R_w2c, padding_last=True):\n \"\"\"\n R_w2c : (F, 3, 3)\n \"\"\"\n # R @ R0 = R1, so R = R1 @ R0^T\n cam_angvel = matrix_to_rotation_6d(R_w2c[1:] @ R_w2c[:-1].transpose(-1, -2)) # (F-1, 6)\n # cam_angvel = (cam_angvel - torch.tensor([[1, 0, 0, 0, 1, 0]])) * FPS\n assert padding_last\n cam_angvel = torch.cat([cam_angvel, cam_angvel[-1:]], dim=0) # (F, 6)\n return cam_angvel.float()\n\n\ndef ransac_gravity_vec(xyz, num_iterations=100, threshold=0.05, verbose=False):\n # xyz: (L, 3)\n N = xyz.shape[0]\n max_inliers = []\n best_model = None\n norms = xyz.norm(dim=-1) # (L,)\n\n for _ in range(num_iterations):\n # 随机选择一个样本\n sample_index = np.random.randint(N)\n sample = xyz[sample_index] # (3,)\n\n # 计算所有点与样本点的角度差\n dot_product = (xyz * sample).sum(dim=-1) # (L,)\n angles = dot_product / norms * norms[sample_index] # (L,)\n angles = torch.clamp(angles, -1, 1) # 防止数值误差导致的异常\n angles = torch.acos(angles)\n\n # 确定内点\n inliers = xyz[angles < threshold]\n\n if len(inliers) > len(max_inliers):\n max_inliers = inliers\n best_model = sample\n if len(max_inliers) == N:\n break\n if verbose:\n print(f\"Inliers: {len(max_inliers)} / {N}\")\n result = max_inliers.mean(dim=0)\n\n return result, max_inliers\n\n\ndef sequence_best_cammat(w_j3d, c_j3d, cam_rot):\n # get best camera estimation along the sequence, requires static camera\n # w_j3d: (L, J, 3)\n # c_j3d: (L, J, 3)\n # cam_rot: (L, 3, 3)\n\n L, J, _ = w_j3d.shape\n\n root_in_w = w_j3d[:, 0] # (L, 3)\n root_in_c = c_j3d[:, 0] # (L, 3)\n cam_mat = matrix.get_TRS(cam_rot, root_in_w) # (L, 4, 4)\n cam_pos = matrix.get_position_from(-root_in_c[:, None], cam_mat)[:, 0] # (L, 3)\n cam_mat = matrix.set_position(cam_mat, cam_pos) # (L, 4, 4)\n\n w_j3d_expand = w_j3d[None].expand(L, -1, -1, -1) # (L, L, J, 3)\n w_j3d_expand = w_j3d_expand.reshape(L, -1, 3) # (L, L*J, 3)\n\n # get reproject error\n w_j3d_expand_in_c = matrix.get_relative_position_to(w_j3d_expand, cam_mat) # (L, L*J, 3)\n w_j2d_expand_in_c = project_p2d(w_j3d_expand_in_c) # (L, L*J, 2)\n w_j2d_expand_in_c = w_j2d_expand_in_c.reshape(L, L, J, 2) # (L, L, J, 2)\n c_j2d = project_p2d(c_j3d) # (L, J, 2)\n error = w_j2d_expand_in_c - c_j2d[None] # (L, L, J, 2)\n error = error.norm(dim=-1).mean(dim=-1) # (L, L)\n error = error.mean(dim=-1) # (L,)\n ind = error.argmin()\n return cam_mat[ind], ind\n\n\ndef get_sequence_cammat(w_j3d, c_j3d, cam_rot):\n # w_j3d: (L, J, 3)\n # c_j3d: (L, J, 3)\n # cam_rot: (L, 3, 3)\n\n L, J, _ = w_j3d.shape\n\n root_in_w = w_j3d[:, 0] # (L, 3)\n root_in_c = c_j3d[:, 0] # (L, 3)\n cam_mat = matrix.get_TRS(cam_rot, root_in_w) # (L, 4, 4)\n cam_pos = matrix.get_position_from(-root_in_c[:, None], cam_mat)[:, 0] # (L, 3)\n cam_mat = matrix.set_position(cam_mat, cam_pos) # (L, 4, 4)\n return cam_mat\n\n\ndef ransac_vec(vel, min_multiply=20, verbose=False):\n # xyz: (L, 3)\n # remove outlier velocity\n N = vel.shape[0]\n vel_1 = vel[None].expand(N, -1, -1) # (L, L, 3)\n vel_2 = vel[:, None].expand(-1, N, -1) # (L, L, 3)\n dist_mat = (vel_1 - vel_2).norm(dim=-1) # (L, L)\n big_identity = torch.eye(N, device=vel.device) * 1e6\n dist_mat_ = dist_mat + big_identity\n threshold = dist_mat_.min() * min_multiply\n inner_mask = dist_mat < threshold # (L, L)\n inner_num = inner_mask.sum(dim=-1) # (L, )\n ind = inner_num.argmax()\n result = vel[inner_mask[ind]].mean(dim=0) # (3,)\n if verbose:\n print(inner_mask[ind].sum().item())\n\n return result, inner_mask[ind]\n" }, { "path": "eval/GVHMR/hmr4d/utils/ik/ccd_ik.py", "content": "# Sebastian IK\nimport torch\nimport torch.nn as nn\nimport torch.nn.functional as F\nfrom einops import einsum, rearrange, repeat\n\nfrom pytorch3d.transforms import (\n matrix_to_rotation_6d,\n rotation_6d_to_matrix,\n axis_angle_to_matrix,\n matrix_to_axis_angle,\n quaternion_to_matrix,\n matrix_to_quaternion,\n)\nimport hmr4d.utils.matrix as matrix\nfrom hmr4d.utils.geo.quaternion import qbetween, qslerp, qinv, qmul, qrot\n\n\nclass CCD_IK:\n def __init__(\n self,\n local_mat,\n parent,\n target_ind,\n target_pos=None,\n target_rot=None,\n kinematic_chain=None,\n max_iter=2, # sebas sets 25 but with converged flag, 2 is enough\n threshold=0.001,\n pos_weight=1.0,\n rot_weight=0.0, # this makes optimization unstable, although sebas uses 1.0\n ):\n if kinematic_chain is None:\n kinematic_chain = range(local_mat.shape[-3])\n global_mat = matrix.forward_kinematics(local_mat, parent)\n\n # get kinematic chain only local mat and assign root mat (do not modify root during IK)\n local_mat = local_mat.clone()\n local_mat = local_mat[..., kinematic_chain, :, :]\n local_mat[..., 0, :, :] = global_mat[..., kinematic_chain[0], :, :]\n\n parent = [i - 1 for i in range(len(kinematic_chain))]\n self.local_mat = local_mat\n self.global_mat = matrix.forward_kinematics(local_mat, parent) # (*, J, 4, 4)\n self.parent = parent\n\n self.target_ind = target_ind\n if target_pos is not None:\n self.target_pos = target_pos # (*, O, 3)\n else:\n self.target_pos = None\n if target_rot is not None:\n self.target_q = matrix_to_quaternion(target_rot) # (*, O, 4)\n else:\n self.target_q = None\n\n self.threshold = threshold\n self.J_N = self.local_mat.shape[-3]\n self.target_N = len(target_ind)\n self.max_iter = max_iter\n self.pos_weight = pos_weight\n self.rot_weight = rot_weight\n\n def is_converged(self):\n end_pos = matrix.get_position(self.global_mat)[..., self.target_ind, :] # (*, OJ, 3)\n converged_mask = (self.target_pos - end_pos).norm(dim=-1) < self.threshold\n self.converged_mask = converged_mask\n if self.converged_mask.sum() > 0:\n return False\n return True\n\n def solve(self):\n for _ in range(self.max_iter):\n # if self.is_converged():\n # return self.local_mat\n # do not optimize root, so start from 1\n self.optimize(1)\n return self.local_mat\n\n def optimize(self, i):\n # i: joint_i\n if i == self.J_N - 1:\n return\n pos = matrix.get_position(self.global_mat)[..., i, :] # (*, 3)\n rot = matrix.get_rotation(self.global_mat)[..., i, :, :] # (*, 3, 3)\n quat = matrix_to_quaternion(rot) # (*, 4)\n x_vec = torch.zeros((quat.shape[:-1] + (3,)), device=quat.device)\n x_vec[..., 0] = 1.0\n x_vec_sum = torch.zeros_like(x_vec)\n y_vec = torch.zeros((quat.shape[:-1] + (3,)), device=quat.device)\n y_vec[..., 1] = 1.0\n y_vec_sum = torch.zeros_like(y_vec)\n\n count = 0\n\n for target_i, j in enumerate(self.target_ind):\n if i >= j:\n # do not optimise same joint or child joint of targets\n continue\n end_pos = matrix.get_position(self.global_mat)[..., j, :] # (*, 3)\n end_rot = matrix.get_rotation(self.global_mat)[..., j, :, :] # (*, 3, 3)\n end_quat = matrix_to_quaternion(end_rot) # (*, 4)\n\n if self.target_pos is not None:\n target_pos = self.target_pos[..., target_i, :] # (*, 3)\n # Solve objective position\n solved_pos_target_quat = qslerp(\n quat,\n qmul(qbetween(end_pos - pos, target_pos - pos), quat),\n self.get_weight(i),\n )\n\n x_vec_sum += qrot(solved_pos_target_quat, x_vec)\n y_vec_sum += qrot(solved_pos_target_quat, y_vec)\n if self.pos_weight > 0:\n count += 1\n\n if self.target_q is not None:\n if target_i < self.target_N - 1:\n # multiple rot target makes more unstable, only keep the last one\n continue\n # optimize rotation target is not stable\n target_q = self.target_q[..., target_i, :] # (*, 4)\n # Solve objective rotation\n solved_q_target_quat = qslerp(\n quat,\n qmul(qmul(target_q, qinv(end_quat)), quat),\n self.get_weight(i),\n )\n x_vec_sum += qrot(solved_q_target_quat, x_vec) * self.rot_weight\n y_vec_sum += qrot(solved_q_target_quat, y_vec) * self.rot_weight\n if self.rot_weight > 0:\n count += 1\n\n if count > 0:\n x_vec_avg = matrix.normalize(x_vec_sum / count)\n y_vec_avg = matrix.normalize(y_vec_sum / count)\n z_vec_avg = torch.cross(x_vec_avg, y_vec_avg, dim=-1)\n solved_rot = torch.stack([x_vec_avg, y_vec_avg, z_vec_avg], dim=-1) # column\n\n parent_rot = matrix.get_rotation(self.global_mat)[..., self.parent[i], :, :]\n solved_local_rot = matrix.get_mat_BtoA(parent_rot, solved_rot)\n self.local_mat[..., i, :-1, :-1] = solved_local_rot\n self.global_mat = matrix.forward_kinematics(self.local_mat, self.parent)\n self.optimize(i + 1)\n\n def get_weight(self, i):\n weight = (i + 1) / self.J_N\n return weight\n" }, { "path": "eval/GVHMR/hmr4d/utils/kpts/kp2d_utils.py", "content": "# Copyright (c) OpenMMLab. All rights reserved.\nimport warnings\nimport cv2\nimport numpy as np\n\n# expose _taylor to outside\n__all__ = [\"keypoints_from_heatmaps\"]\n\n\ndef _taylor(heatmap, coord):\n \"\"\"Distribution aware coordinate decoding method.\n\n Note:\n - heatmap height: H\n - heatmap width: W\n\n Args:\n heatmap (np.ndarray[H, W]): Heatmap of a particular joint type.\n coord (np.ndarray[2,]): Coordinates of the predicted keypoints.\n\n Returns:\n np.ndarray[2,]: Updated coordinates.\n \"\"\"\n H, W = heatmap.shape[:2]\n px, py = int(coord[0]), int(coord[1])\n if 1 < px < W - 2 and 1 < py < H - 2:\n dx = 0.5 * (heatmap[py][px + 1] - heatmap[py][px - 1])\n dy = 0.5 * (heatmap[py + 1][px] - heatmap[py - 1][px])\n dxx = 0.25 * (heatmap[py][px + 2] - 2 * heatmap[py][px] + heatmap[py][px - 2])\n dxy = 0.25 * (\n heatmap[py + 1][px + 1] - heatmap[py - 1][px + 1] - heatmap[py + 1][px - 1] + heatmap[py - 1][px - 1]\n )\n dyy = 0.25 * (heatmap[py + 2 * 1][px] - 2 * heatmap[py][px] + heatmap[py - 2 * 1][px])\n derivative = np.array([[dx], [dy]])\n hessian = np.array([[dxx, dxy], [dxy, dyy]])\n if dxx * dyy - dxy**2 != 0:\n hessianinv = np.linalg.inv(hessian)\n offset = -hessianinv @ derivative\n offset = np.squeeze(np.array(offset.T), axis=0)\n coord += offset\n return coord\n\n\ndef _get_max_preds(heatmaps):\n \"\"\"Get keypoint predictions from score maps.\n\n Note:\n batch_size: N\n num_keypoints: K\n heatmap height: H\n heatmap width: W\n\n Args:\n heatmaps (np.ndarray[N, K, H, W]): model predicted heatmaps.\n\n Returns:\n tuple: A tuple containing aggregated results.\n\n - preds (np.ndarray[N, K, 2]): Predicted keypoint location.\n - maxvals (np.ndarray[N, K, 1]): Scores (confidence) of the keypoints.\n \"\"\"\n assert isinstance(heatmaps, np.ndarray), \"heatmaps should be numpy.ndarray\"\n assert heatmaps.ndim == 4, \"batch_images should be 4-ndim\"\n\n N, K, _, W = heatmaps.shape\n heatmaps_reshaped = heatmaps.reshape((N, K, -1))\n idx = np.argmax(heatmaps_reshaped, 2).reshape((N, K, 1))\n maxvals = np.amax(heatmaps_reshaped, 2).reshape((N, K, 1))\n\n preds = np.tile(idx, (1, 1, 2)).astype(np.float32)\n preds[:, :, 0] = preds[:, :, 0] % W\n preds[:, :, 1] = preds[:, :, 1] // W\n\n preds = np.where(np.tile(maxvals, (1, 1, 2)) > 0.0, preds, -1)\n return preds, maxvals\n\n\ndef post_dark_udp(coords, batch_heatmaps, kernel=3):\n \"\"\"DARK post-pocessing. Implemented by udp. Paper ref: Huang et al. The\n Devil is in the Details: Delving into Unbiased Data Processing for Human\n Pose Estimation (CVPR 2020). Zhang et al. Distribution-Aware Coordinate\n Representation for Human Pose Estimation (CVPR 2020).\n\n Note:\n - batch size: B\n - num keypoints: K\n - num persons: N\n - height of heatmaps: H\n - width of heatmaps: W\n\n B=1 for bottom_up paradigm where all persons share the same heatmap.\n B=N for top_down paradigm where each person has its own heatmaps.\n\n Args:\n coords (np.ndarray[N, K, 2]): Initial coordinates of human pose.\n batch_heatmaps (np.ndarray[B, K, H, W]): batch_heatmaps\n kernel (int): Gaussian kernel size (K) for modulation.\n\n Returns:\n np.ndarray([N, K, 2]): Refined coordinates.\n \"\"\"\n if not isinstance(batch_heatmaps, np.ndarray):\n batch_heatmaps = batch_heatmaps.cpu().numpy()\n B, K, H, W = batch_heatmaps.shape\n N = coords.shape[0]\n assert B == 1 or B == N\n for heatmaps in batch_heatmaps:\n for heatmap in heatmaps:\n cv2.GaussianBlur(heatmap, (kernel, kernel), 0, heatmap)\n np.clip(batch_heatmaps, 0.001, 50, batch_heatmaps)\n np.log(batch_heatmaps, batch_heatmaps)\n\n batch_heatmaps_pad = np.pad(batch_heatmaps, ((0, 0), (0, 0), (1, 1), (1, 1)), mode=\"edge\").flatten()\n\n index = coords[..., 0] + 1 + (coords[..., 1] + 1) * (W + 2)\n index += (W + 2) * (H + 2) * np.arange(0, B * K).reshape(-1, K)\n index = index.astype(int).reshape(-1, 1)\n i_ = batch_heatmaps_pad[index]\n ix1 = batch_heatmaps_pad[index + 1]\n iy1 = batch_heatmaps_pad[index + W + 2]\n ix1y1 = batch_heatmaps_pad[index + W + 3]\n ix1_y1_ = batch_heatmaps_pad[index - W - 3]\n ix1_ = batch_heatmaps_pad[index - 1]\n iy1_ = batch_heatmaps_pad[index - 2 - W]\n\n dx = 0.5 * (ix1 - ix1_)\n dy = 0.5 * (iy1 - iy1_)\n derivative = np.concatenate([dx, dy], axis=1)\n derivative = derivative.reshape(N, K, 2, 1)\n dxx = ix1 - 2 * i_ + ix1_\n dyy = iy1 - 2 * i_ + iy1_\n dxy = 0.5 * (ix1y1 - ix1 - iy1 + i_ + i_ - ix1_ - iy1_ + ix1_y1_)\n hessian = np.concatenate([dxx, dxy, dxy, dyy], axis=1)\n hessian = hessian.reshape(N, K, 2, 2)\n hessian = np.linalg.inv(hessian + np.finfo(np.float32).eps * np.eye(2))\n coords -= np.einsum(\"ijmn,ijnk->ijmk\", hessian, derivative).squeeze()\n return coords\n\n\ndef _gaussian_blur(heatmaps, kernel=11):\n \"\"\"Modulate heatmap distribution with Gaussian.\n sigma = 0.3*((kernel_size-1)*0.5-1)+0.8\n sigma~=3 if k=17\n sigma=2 if k=11;\n sigma~=1.5 if k=7;\n sigma~=1 if k=3;\n\n Note:\n - batch_size: N\n - num_keypoints: K\n - heatmap height: H\n - heatmap width: W\n\n Args:\n heatmaps (np.ndarray[N, K, H, W]): model predicted heatmaps.\n kernel (int): Gaussian kernel size (K) for modulation, which should\n match the heatmap gaussian sigma when training.\n K=17 for sigma=3 and k=11 for sigma=2.\n\n Returns:\n np.ndarray ([N, K, H, W]): Modulated heatmap distribution.\n \"\"\"\n assert kernel % 2 == 1\n\n border = (kernel - 1) // 2\n batch_size = heatmaps.shape[0]\n num_joints = heatmaps.shape[1]\n height = heatmaps.shape[2]\n width = heatmaps.shape[3]\n for i in range(batch_size):\n for j in range(num_joints):\n origin_max = np.max(heatmaps[i, j])\n dr = np.zeros((height + 2 * border, width + 2 * border), dtype=np.float32)\n dr[border:-border, border:-border] = heatmaps[i, j].copy()\n dr = cv2.GaussianBlur(dr, (kernel, kernel), 0)\n heatmaps[i, j] = dr[border:-border, border:-border].copy()\n heatmaps[i, j] *= origin_max / np.max(heatmaps[i, j])\n return heatmaps\n\n\ndef keypoints_from_heatmaps(\n heatmaps,\n center,\n scale,\n unbiased=False,\n post_process=\"default\",\n kernel=11,\n valid_radius_factor=0.0546875,\n use_udp=False,\n target_type=\"GaussianHeatmap\",\n):\n \"\"\"Get final keypoint predictions from heatmaps and transform them back to\n the image.\n\n Note:\n - batch size: N\n - num keypoints: K\n - heatmap height: H\n - heatmap width: W\n\n Args:\n heatmaps (np.ndarray[N, K, H, W]): model predicted heatmaps.\n center (np.ndarray[N, 2]): Center of the bounding box (x, y).\n scale (np.ndarray[N, 2]): Scale of the bounding box\n wrt height/width.\n post_process (str/None): Choice of methods to post-process\n heatmaps. Currently supported: None, 'default', 'unbiased',\n 'megvii'.\n unbiased (bool): Option to use unbiased decoding. Mutually\n exclusive with megvii.\n Note: this arg is deprecated and unbiased=True can be replaced\n by post_process='unbiased'\n Paper ref: Zhang et al. Distribution-Aware Coordinate\n Representation for Human Pose Estimation (CVPR 2020).\n kernel (int): Gaussian kernel size (K) for modulation, which should\n match the heatmap gaussian sigma when training.\n K=17 for sigma=3 and k=11 for sigma=2.\n valid_radius_factor (float): The radius factor of the positive area\n in classification heatmap for UDP.\n use_udp (bool): Use unbiased data processing.\n target_type (str): 'GaussianHeatmap' or 'CombinedTarget'.\n GaussianHeatmap: Classification target with gaussian distribution.\n CombinedTarget: The combination of classification target\n (response map) and regression target (offset map).\n Paper ref: Huang et al. The Devil is in the Details: Delving into\n Unbiased Data Processing for Human Pose Estimation (CVPR 2020).\n\n Returns:\n tuple: A tuple containing keypoint predictions and scores.\n\n - preds (np.ndarray[N, K, 2]): Predicted keypoint location in images.\n - maxvals (np.ndarray[N, K, 1]): Scores (confidence) of the keypoints.\n \"\"\"\n # Avoid being affected\n heatmaps = heatmaps.copy()\n\n # detect conflicts\n if unbiased:\n assert post_process not in [False, None, \"megvii\"]\n if post_process in [\"megvii\", \"unbiased\"]:\n assert kernel > 0\n if use_udp:\n assert not post_process == \"megvii\"\n\n # normalize configs\n if post_process is False:\n warnings.warn(\"post_process=False is deprecated, \" \"please use post_process=None instead\", DeprecationWarning)\n post_process = None\n elif post_process is True:\n if unbiased is True:\n warnings.warn(\n \"post_process=True, unbiased=True is deprecated,\" \" please use post_process='unbiased' instead\",\n DeprecationWarning,\n )\n post_process = \"unbiased\"\n else:\n warnings.warn(\n \"post_process=True, unbiased=False is deprecated, \" \"please use post_process='default' instead\",\n DeprecationWarning,\n )\n post_process = \"default\"\n elif post_process == \"default\":\n if unbiased is True:\n warnings.warn(\n \"unbiased=True is deprecated, please use \" \"post_process='unbiased' instead\", DeprecationWarning\n )\n post_process = \"unbiased\"\n\n # start processing\n if post_process == \"megvii\":\n heatmaps = _gaussian_blur(heatmaps, kernel=kernel)\n\n N, K, H, W = heatmaps.shape\n if use_udp:\n if target_type.lower() == \"GaussianHeatMap\".lower():\n preds, maxvals = _get_max_preds(heatmaps)\n preds = post_dark_udp(preds, heatmaps, kernel=kernel)\n elif target_type.lower() == \"CombinedTarget\".lower():\n for person_heatmaps in heatmaps:\n for i, heatmap in enumerate(person_heatmaps):\n kt = 2 * kernel + 1 if i % 3 == 0 else kernel\n cv2.GaussianBlur(heatmap, (kt, kt), 0, heatmap)\n # valid radius is in direct proportion to the height of heatmap.\n valid_radius = valid_radius_factor * H\n offset_x = heatmaps[:, 1::3, :].flatten() * valid_radius\n offset_y = heatmaps[:, 2::3, :].flatten() * valid_radius\n heatmaps = heatmaps[:, ::3, :]\n preds, maxvals = _get_max_preds(heatmaps)\n index = preds[..., 0] + preds[..., 1] * W\n index += W * H * np.arange(0, N * K / 3)\n index = index.astype(int).reshape(N, K // 3, 1)\n preds += np.concatenate((offset_x[index], offset_y[index]), axis=2)\n else:\n raise ValueError(\"target_type should be either \" \"'GaussianHeatmap' or 'CombinedTarget'\")\n else:\n preds, maxvals = _get_max_preds(heatmaps)\n if post_process == \"unbiased\": # alleviate biased coordinate\n # apply Gaussian distribution modulation.\n heatmaps = np.log(np.maximum(_gaussian_blur(heatmaps, kernel), 1e-10))\n for n in range(N):\n for k in range(K):\n preds[n][k] = _taylor(heatmaps[n][k], preds[n][k])\n elif post_process is not None:\n # add +/-0.25 shift to the predicted locations for higher acc.\n for n in range(N):\n for k in range(K):\n heatmap = heatmaps[n][k]\n px = int(preds[n][k][0])\n py = int(preds[n][k][1])\n if 1 < px < W - 1 and 1 < py < H - 1:\n diff = np.array(\n [heatmap[py][px + 1] - heatmap[py][px - 1], heatmap[py + 1][px] - heatmap[py - 1][px]]\n )\n preds[n][k] += np.sign(diff) * 0.25\n if post_process == \"megvii\":\n preds[n][k] += 0.5\n\n # Transform back to the image\n for i in range(N):\n preds[i] = transform_preds(preds[i], center[i], scale[i], [W, H], use_udp=use_udp)\n\n if post_process == \"megvii\":\n maxvals = maxvals / 255.0 + 0.5\n\n return preds, maxvals\n\n\ndef transform_preds(coords, center, scale, output_size, use_udp=False):\n \"\"\"Get final keypoint predictions from heatmaps and apply scaling and\n translation to map them back to the image.\n\n Note:\n num_keypoints: K\n\n Args:\n coords (np.ndarray[K, ndims]):\n\n * If ndims=2, corrds are predicted keypoint location.\n * If ndims=4, corrds are composed of (x, y, scores, tags)\n * If ndims=5, corrds are composed of (x, y, scores, tags,\n flipped_tags)\n\n center (np.ndarray[2, ]): Center of the bounding box (x, y).\n scale (np.ndarray[2, ]): Scale of the bounding box\n wrt [width, height].\n output_size (np.ndarray[2, ] | list(2,)): Size of the\n destination heatmaps.\n use_udp (bool): Use unbiased data processing\n\n Returns:\n np.ndarray: Predicted coordinates in the images.\n \"\"\"\n assert coords.shape[1] in (2, 4, 5)\n assert len(center) == 2\n assert len(scale) == 2\n assert len(output_size) == 2\n\n # Recover the scale which is normalized by a factor of 200.\n scale = scale * 200.0\n\n if use_udp:\n scale_x = scale[0] / (output_size[0] - 1.0)\n scale_y = scale[1] / (output_size[1] - 1.0)\n else:\n scale_x = scale[0] / output_size[0]\n scale_y = scale[1] / output_size[1]\n\n target_coords = np.ones_like(coords)\n target_coords[:, 0] = coords[:, 0] * scale_x + center[0] - scale[0] * 0.5\n target_coords[:, 1] = coords[:, 1] * scale_y + center[1] - scale[1] * 0.5\n\n return target_coords\n" }, { "path": "eval/GVHMR/hmr4d/utils/matrix.py", "content": "import torch\nimport torch.nn as nn\nimport torch.nn.functional as F\nimport copy\nfrom typing import List, Optional\n\nimport numpy as np\n\nimport math\n\n\ndef identity_mat(x=None, device=\"cpu\", is_numpy=False):\n if x is not None:\n if isinstance(x, torch.Tensor):\n mat = torch.eye(4, device=device)\n mat = mat.repeat(x.shape[:-2] + (1, 1))\n elif isinstance(x, np.ndarray):\n mat = np.eye(4, dtype=np.float32)\n if x is not None:\n for _ in range(len(x.shape) - 2):\n mat = mat[None]\n mat = np.tile(mat, x.shape[:-2] + (1, 1))\n else:\n raise ValueError\n else:\n # (4, 4)\n if is_numpy:\n mat = np.eye(4, dtype=np.float32)\n else:\n mat = torch.eye(4, device=device)\n\n return mat\n\n\ndef vec2mat(vec):\n \"\"\"_summary_\n\n Args:\n vec (tensor): [12], pos, forward, up and right\n\n Returns:\n mat_world(tensor): [4, 4]\n \"\"\"\n # Assume bs = 1\n v = np.tile(np.array([[0, 0, 0, 1]]), (1, 1))\n if isinstance(vec, torch.Tensor):\n v = torch.tensor(\n v,\n device=vec.device,\n dtype=vec.dtype,\n )\n pos = vec[:3]\n forward = vec[3:6]\n up = vec[6:9]\n right = vec[9:12]\n\n if isinstance(vec, torch.Tensor):\n mat_world = torch.stack([right, up, forward, pos], dim=-1)\n mat_world = torch.cat([mat_world, v], dim=-2)\n elif isinstance(vec, np.ndarray):\n mat_world = np.stack([right, up, forward, pos], axis=-1)\n mat_world = np.concatenate([mat_world, v], axis=-2)\n else:\n raise ValueError\n mat_world = normalized_matrix(mat_world)\n return mat_world\n\n\ndef mat2vec(mat):\n \"\"\"_summary_\n\n Args:\n mat(tensor): [4, 4]\n\n Returns:\n vec (tensor): [12], pos, forward, up and right\n \"\"\"\n # Assume bs = 1\n pos = mat[:-1, 3]\n forward = normalized(mat[:-1, 2])\n up = normalized(mat[:-1, 1])\n right = normalized(mat[:-1, 0])\n if isinstance(mat, torch.Tensor):\n vec = torch.cat((pos, forward, up, right))\n elif isinstance(mat, np.ndarray):\n vec = np.concatenate((pos, forward, up, right))\n else:\n raise ValueError\n\n return vec\n\n\ndef vec2mat_batch(vec):\n \"\"\"_summary_\n\n Args:\n vec (tensor): [B, 12], pos, forward, up and right\n\n Returns:\n mat_world(tensor): [B, 4, 4]\n \"\"\"\n # Assume bs = 1\n\n v = np.tile(np.array([[0, 0, 0, 1]], dtype=np.float32), (vec.shape[0], 1, 1))\n if isinstance(vec, torch.Tensor):\n v = torch.tensor(\n v,\n device=vec.device,\n dtype=vec.dtype,\n )\n pos = vec[..., :3]\n forward = vec[..., 3:6]\n up = vec[..., 6:9]\n right = vec[..., 9:12]\n if isinstance(vec, torch.Tensor):\n mat_world = torch.stack([right, up, forward, pos], dim=-1)\n mat_world = torch.cat([mat_world, v], dim=-2)\n elif isinstance(vec, np.ndarray):\n mat_world = np.stack([right, up, forward, pos], axis=-1)\n mat_world = np.concatenate([mat_world, v], axis=-2)\n else:\n raise ValueError\n\n mat_world = normalized_matrix(mat_world)\n return mat_world\n\n\ndef rotmat2tan_norm(mat):\n \"\"\"_summary_\n\n Args:\n mat(tensor): [B, 3, 3]\n\n Returns:\n vec (tensor): [B, 6], tan norm\n \"\"\"\n if isinstance(mat, np.ndarray):\n tan = np.zeros_like(mat[..., 2])\n norm = np.zeros_like(mat[..., 0])\n elif isinstance(mat, torch.Tensor):\n tan = torch.zeros_like(mat[..., 2])\n norm = torch.zeros_like(mat[..., 0])\n else:\n raise ValueError\n tan[...] = mat[..., 2, ::-1]\n tan[..., -1] *= -1\n norm[...] = mat[..., 0, ::-1]\n norm[..., -1] *= -1\n if isinstance(mat, np.ndarray):\n tan_norm = np.concatenate((tan, norm), axis=-1)\n elif isinstance(mat, torch.Tensor):\n tan_norm = torch.cat((tan, norm), dim=-1)\n else:\n raise ValueError\n return tan_norm\n\n\ndef mat2tan_norm(mat):\n \"\"\"_summary_\n\n Args:\n mat(tensor): [B, 4, 4]\n\n Returns:\n vec (tensor): [B, 6], tan norm\n \"\"\"\n rot_mat = mat[..., :-1, :-1]\n return rotmat2tan_norm(rot_mat)\n\n\ndef rotmat2tan_norm(mat):\n \"\"\"_summary_\n\n Args:\n mat(tensor): [B, 3, 3]\n\n Returns:\n vec (tensor): [B, 6], tan norm\n \"\"\"\n if isinstance(mat, np.ndarray):\n tan = np.zeros_like(mat[..., 2])\n norm = np.zeros_like(mat[..., 0])\n tan[...] = mat[..., 2, ::-1]\n norm[...] = mat[..., 0, ::-1]\n elif isinstance(mat, torch.Tensor):\n tan = torch.zeros_like(mat[..., 2])\n norm = torch.zeros_like(mat[..., 0])\n tan[...] = torch.flip(mat[..., 2], dims=[-1])\n norm[...] = torch.flip(mat[..., 0], dims=[-1])\n else:\n raise ValueError\n tan[..., -1] *= -1\n norm[..., -1] *= -1\n if isinstance(mat, np.ndarray):\n tan_norm = np.concatenate((tan, norm), axis=-1)\n elif isinstance(mat, torch.Tensor):\n tan_norm = torch.cat((tan, norm), dim=-1)\n else:\n raise ValueError\n return tan_norm\n\n\ndef tan_norm2rotmat(tan_norm):\n \"\"\"_summary_\n\n Args:\n mat(tensor): [B, 6]\n\n Returns:\n vec (tensor): [B, 3]\n \"\"\"\n tan = copy.deepcopy(tan_norm[..., :3])\n norm = copy.deepcopy(tan_norm[..., 3:])\n tan[..., -1] *= -1\n norm[..., -1] *= -1\n if isinstance(tan_norm, np.ndarray):\n rotmat = np.zeros(tan_norm.shape[:-1] + (3, 3))\n tan = tan[..., ::-1]\n norm = norm[..., ::-1]\n other = np.cross(tan, norm)\n elif isinstance(tan_norm, torch.Tensor):\n rotmat = torch.zeros(tan_norm.shape[:-1] + (3, 3), device=tan_norm.device)\n tan = torch.flip(tan, dims=[-1])\n norm = torch.flip(norm, dims=[-1])\n other = torch.cross(tan, norm)\n else:\n raise ValueError\n rotmat[..., 2, :] = tan\n rotmat[..., 0, :] = norm\n rotmat[..., 1, :] = other\n return rotmat\n\n\ndef rotmat332vec_batch(mat):\n \"\"\"_summary_\n\n Args:\n mat(tensor): [B, 3, 3]\n\n Returns:\n vec (tensor): [B, 6], forward, up, right\n \"\"\"\n # Assume bs = 1\n mat = normalized_matrix(mat)\n forward = mat[..., :, 2]\n up = mat[..., :, 1]\n right = mat[..., :, 0]\n if isinstance(mat, torch.Tensor):\n vec = torch.cat((forward, up, right), dim=-1)\n elif isinstance(mat, np.ndarray):\n vec = np.concatenate((forward, up, right), axis=-1)\n else:\n raise ValueError\n return vec\n\n\ndef rotmat2vec_batch(mat):\n \"\"\"_summary_\n\n Args:\n mat(tensor): [B, 4, 4]\n\n Returns:\n vec (tensor): [B, 9], forward, up, right\n \"\"\"\n # Assume bs = 1\n mat = normalized_matrix(mat)\n forward = mat[..., :-1, 2]\n up = mat[..., :-1, 1]\n right = mat[..., :-1, 0]\n if isinstance(mat, torch.Tensor):\n vec = torch.cat((forward, up, right), dim=-1)\n elif isinstance(mat, np.ndarray):\n vec = np.concatenate((forward, up, right), axis=-1)\n else:\n raise ValueError\n return vec\n\n\ndef mat2vec_batch(mat):\n \"\"\"_summary_\n\n Args:\n mat(tensor): [B, 4, 4]\n\n Returns:\n vec (tensor): [B, 12], pos, forward, up and right\n \"\"\"\n # Assume bs = 1\n mat = normalized_matrix(mat)\n pos = mat[..., :-1, 3]\n forward = mat[..., :-1, 2]\n up = mat[..., :-1, 1]\n right = mat[..., :-1, 0]\n if isinstance(mat, torch.Tensor):\n vec = torch.cat((pos, forward, up, right), dim=-1)\n elif isinstance(mat, np.ndarray):\n vec = np.concatenate((pos, forward, up, right), axis=-1)\n else:\n raise ValueError\n return vec\n\n\ndef mat2pose_batch(mat, returnvel=True):\n \"\"\"_summary_\n\n Args:\n mat(tensor): [B, 4, 4]\n\n Returns:\n vec (tensor): [B, 12], pos, forward, up, zeros\n \"\"\"\n # Assume bs = 1\n mat = normalized_matrix(mat)\n pos = mat[..., :-1, 3]\n forward = mat[..., :-1, 2]\n up = mat[..., :-1, 1]\n if isinstance(mat, torch.Tensor):\n if returnvel:\n vel = torch.zeros_like(up)\n vec = torch.cat((pos, forward, up, vel), dim=-1)\n else:\n vec = torch.cat((pos, forward, up), dim=-1)\n elif isinstance(mat, np.ndarray):\n if returnvel:\n vel = np.zeros_like(up)\n vec = np.concatenate((pos, forward, up, vel), axis=-1)\n else:\n vec = np.concatenate((pos, forward, up), axis=-1)\n else:\n raise ValueError\n return vec\n\n\ndef get_mat_BinA(matCtoA, matCtoB):\n \"\"\"\n given matrix of the same object in two coordinate A and B,\n return matrix B in the coordinate of A\n\n Args:\n matCtoA (tensor): [4, 4] world matrix\n matCtoB (tensor): [4, 4] world matrix\n \"\"\"\n if isinstance(matCtoA, torch.Tensor):\n matCtoB_inv = torch.inverse(matCtoB)\n elif isinstance(matCtoA, np.ndarray):\n matCtoB_inv = np.linalg.inv(matCtoB)\n else:\n raise ValueError\n matCtoB_inv = normalized_matrix(matCtoB_inv)\n if isinstance(matCtoA, torch.Tensor):\n mat_BtoA = torch.matmul(matCtoA, matCtoB_inv)\n elif isinstance(matCtoA, np.ndarray):\n mat_BtoA = np.matmul(matCtoA, matCtoB_inv)\n mat_BtoA = normalized_matrix(mat_BtoA)\n return mat_BtoA\n\n\ndef get_mat_BtoA(matA, matB):\n \"\"\"\n return matrix B in the coordinate of A\n\n Args:\n matA (tensor): [4, 4] world matrix\n matB (tensor): [4, 4] world matrix\n \"\"\"\n if isinstance(matA, torch.Tensor):\n matA_inv = torch.inverse(matA)\n elif isinstance(matA, np.ndarray):\n matA_inv = np.linalg.inv(matA)\n else:\n raise ValueError\n matA_inv = normalized_matrix(matA_inv)\n if isinstance(matA, torch.Tensor):\n mat_BtoA = torch.matmul(matA_inv, matB)\n elif isinstance(matA, np.ndarray):\n mat_BtoA = np.matmul(matA_inv, matB)\n mat_BtoA = normalized_matrix(mat_BtoA)\n return mat_BtoA\n\n\ndef get_mat_BfromA(matA, matBtoA):\n \"\"\"\n return world matrix B given matrix A and mat B realtive to A\n\n Args:\n matA (_type_): [4, 4] world matrix\n matBtoA (_type_): [4, 4] matrix B relative to A\n \"\"\"\n if isinstance(matA, torch.Tensor):\n matB = torch.matmul(matA, matBtoA)\n if isinstance(matA, np.ndarray):\n matB = np.matmul(matA, matBtoA)\n matB = normalized_matrix(matB)\n return matB\n\n\ndef get_relative_position_to(pos, mat):\n \"\"\"_summary_\n\n Args:\n pos (_type_): [N, M, 3] or [N, 3]\n mat (_type_): [N, 4, 4] or [4, 4]\n\n Returns:\n _type_: _description_\n \"\"\"\n if isinstance(mat, torch.Tensor):\n mat_inv = torch.inverse(mat)\n elif isinstance(mat, np.ndarray):\n mat_inv = np.linalg.inv(mat)\n else:\n raise ValueError\n mat_inv = normalized_matrix(mat_inv)\n if isinstance(mat, torch.Tensor):\n rot_pos = torch.matmul(mat_inv[..., :-1, :-1], pos.transpose(-1, -2)).transpose(-1, -2)\n elif isinstance(mat, np.ndarray):\n rot_pos = np.matmul(mat_inv[..., :-1, :-1], pos.swapaxes(-1, -2)).swapaxes(-1, -2)\n world_pos = rot_pos + mat_inv[..., None, :-1, 3]\n return world_pos\n\n\ndef get_rotation(mat):\n \"\"\"_summary_\n\n Args:\n mat (_type_): [..., 4, 4]\n\n Returns:\n _type_: _description_\n \"\"\"\n return mat[..., :-1, :-1]\n\n\ndef set_rotation(mat, rotmat):\n \"\"\"_summary_\n\n Args:\n mat (_type_): [..., 4, 4]\n\n Returns:\n _type_: _description_\n \"\"\"\n mat[..., :-1, :-1] = rotmat\n return mat\n\n\ndef set_position(mat, pos):\n \"\"\"_summary_\n\n Args:\n mat (_type_): [..., 4, 4]\n\n Returns:\n _type_: _description_\n \"\"\"\n mat[..., :-1, 3] = pos\n return mat\n\n\ndef get_position(mat):\n \"\"\"_summary_\n\n Args:\n mat (_type_): [..., 4, 4]\n\n Returns:\n _type_: _description_\n \"\"\"\n return mat[..., :-1, 3]\n\n\ndef get_position_from(pos, mat):\n \"\"\"_summary_\n\n Args:\n pos (_type_): [N, M, 3] or [N, 3]\n mat (_type_): [N, 4, 4] or [4, 4]\n\n Returns:\n _type_: _description_\n \"\"\"\n if isinstance(mat, torch.Tensor):\n rot_pos = torch.matmul(mat[..., :-1, :-1], pos.transpose(-1, -2)).transpose(-1, -2)\n elif isinstance(mat, np.ndarray):\n rot_pos = np.matmul(mat[..., :-1, :-1], pos.swapaxes(-1, -2)).swapaxes(-1, -2)\n else:\n raise ValueError\n\n world_pos = rot_pos + mat[..., None, :-1, 3]\n return world_pos\n\n\ndef get_position_from_rotmat(pos, mat):\n \"\"\"_summary_\n\n Args:\n pos (_type_): [N, M, 3] or [N, 3]\n mat (_type_): [N, 4, 4] or [4, 4]\n\n Returns:\n _type_: _description_\n \"\"\"\n if isinstance(mat, torch.Tensor):\n rot_pos = torch.matmul(mat, pos.transpose(-1, -2)).transpose(-1, -2)\n elif isinstance(mat, np.ndarray):\n rot_pos = np.matmul(mat, pos.swapaxes(-1, -2)).swapaxes(-1, -2)\n else:\n raise ValueError\n return rot_pos\n\n\ndef get_relative_direction_to(dir, mat):\n \"\"\"_summary_\n\n Args:\n dir (_type_): [N, M, 3] or [N, 3]\n mat (_type_): [N, 4, 4] or [4, 4]\n\n Returns:\n _type_: _description_\n \"\"\"\n if isinstance(mat, torch.Tensor):\n mat_inv = torch.inverse(mat)\n elif isinstance(mat, np.ndarray):\n mat_inv = np.linalg.inv(mat)\n else:\n raise ValueError\n mat_inv = normalized_matrix(mat_inv)\n rot_mat_inv = mat_inv[..., :3, :3]\n if isinstance(mat, torch.Tensor):\n rel_dir = torch.matmul(rot_mat_inv, dir.transpose(-1, -2))\n return rel_dir.transpose(-1, -2)\n elif isinstance(mat, np.ndarray):\n rel_dir = np.matmul(rot_mat_inv, dir.swapaxes(-1, -2))\n return rel_dir.swapaxes(-1, -2)\n else:\n raise ValueError\n return\n\n\ndef get_direction_from(dir, mat):\n \"\"\"_summary_\n\n Args:\n dir (_type_): [N, M, 3] or [N, 3]\n mat (_type_): [N, 4, 4] or [4, 4]\n\n Returns:\n tensor: [N, M, 3] or [N, 3]\n \"\"\"\n rot_mat = mat[..., :3, :3]\n if isinstance(mat, torch.Tensor):\n world_dir = torch.matmul(rot_mat, dir.transpose(-1, -2))\n return world_dir.transpose(-1, -2)\n elif isinstance(mat, np.ndarray):\n world_dir = np.matmul(rot_mat, dir.swapaxes(-1, -2))\n return world_dir.swapaxes(-1, -2)\n else:\n raise ValueError\n return\n\n\ndef get_coord_vis(pos, rot_mat, scale=1.0):\n forward = rot_mat[..., :, 2]\n up = rot_mat[..., :, 1]\n right = rot_mat[..., :, 0]\n return pos + right * scale, pos + up * scale, pos + forward * scale\n\n\ndef project_vec(vec):\n \"\"\"_summary_\n\n Args:\n vec (tensor): [*, 12], pos, forward, up and right\n\n Returns:\n proj_vec (tensor): [*, 4], posx, posz, forwardx, forwardz\n \"\"\"\n posx = vec[..., 0:1]\n posz = vec[..., 2:3]\n forwardx = vec[..., 3:4]\n forwardz = vec[..., 5:6]\n if isinstance(vec, torch.Tensor):\n proj_vec = torch.cat((posx, posz, forwardx, forwardz), dim=-1)\n elif isinstance(vec, np.ndarray):\n proj_vec = np.concatenate((posx, posz, forwardx, forwardz), axis=-1)\n else:\n raise ValueError\n\n return proj_vec\n\n\ndef xz2xyz(vec):\n x = vec[..., 0:1]\n z = vec[..., 1:2]\n if isinstance(vec, torch.Tensor):\n y = torch.zeros(vec.shape[:-1] + (1,), device=vec.device)\n xyz_vec = torch.cat((x, y, z), dim=-1)\n elif isinstance(vec, np.ndarray):\n y = np.zeros(vec.shape[:-1] + (1,))\n xyz_vec = np.concatenate((x, y, z), axis=-1)\n else:\n raise ValueError\n\n return xyz_vec\n\n\ndef normalized(vec):\n if isinstance(vec, torch.Tensor):\n norm_vec = vec / (vec.norm(2, dim=-1, keepdim=True) + 1e-9)\n elif isinstance(vec, np.ndarray):\n norm_vec = vec / (np.linalg.norm(vec, ord=2, axis=-1, keepdims=True) + 1e-9)\n else:\n raise ValueError\n\n return norm_vec\n\n\ndef normalized_matrix(mat):\n if mat.shape[-1] == 4:\n rot_mat = mat[..., :-1, :-1]\n else:\n rot_mat = mat\n if isinstance(mat, torch.Tensor):\n rot_mat_norm = rot_mat / (rot_mat.norm(2, dim=-2, keepdim=True) + 1e-9)\n norm_mat = torch.zeros_like(mat)\n elif isinstance(mat, np.ndarray):\n rot_mat_norm = rot_mat / (np.linalg.norm(rot_mat, ord=2, axis=-2, keepdims=True) + 1e-9)\n norm_mat = np.zeros_like(mat)\n else:\n raise ValueError\n if mat.shape[-1] == 4:\n norm_mat[..., :-1, :-1] = rot_mat_norm\n norm_mat[..., :-1, -1] = mat[..., :-1, -1]\n norm_mat[..., -1, -1] = 1.0\n else:\n norm_mat = rot_mat_norm\n return norm_mat\n\n\ndef get_rot_mat_from_forward(forward):\n \"\"\"_summary_\n\n Args:\n forward (tensor): [N, M, 3]\n\n Returns:\n mat (tensor): [N, M, 3, 3]\n \"\"\"\n if isinstance(forward, torch.Tensor):\n mat = torch.eye(3, device=forward.device).repeat(forward.shape[:-1] + (1, 1))\n right = torch.zeros_like(forward)\n elif isinstance(forward, np.ndarray):\n mat = np.eye(3, dtype=np.float32)\n for _ in range(len(forward.shape) - 1):\n mat = mat[None]\n mat = np.tile(mat, forward.shape[:-1] + (1, 1))\n right = np.zeros_like(forward)\n else:\n raise ValueError\n\n right[..., 0] = forward[..., 2]\n right[..., 1] = 0.0\n right[..., 2] = -forward[..., 0]\n # right = torch.cross(mat[..., 1], forward) # cannot backward\n\n mat[..., 2] = normalized(forward)\n right = normalized(right)\n mat[..., 0] = right\n return mat\n\n\ndef get_rot_mat_from_forward_up(forward, up):\n \"\"\"_summary_\n\n Args:\n forward (tensor): [N, M, 3]\n up (tensor): [N, M, 3]\n\n Returns:\n mat (tensor): [N, M, 3, 3]\n \"\"\"\n if isinstance(forward, torch.Tensor):\n mat = torch.eye(3, device=forward.device).repeat(forward.shape[:-1] + (1, 1))\n right = torch.cross(up, forward)\n elif isinstance(forward, np.ndarray):\n mat = np.eye(3, dtype=np.float32)\n for _ in range(len(forward.shape) - 1):\n mat = mat[None]\n mat = np.tile(mat, forward.shape[:-1] + (1, 1))\n right = np.cross(up, forward)\n else:\n raise ValueError\n\n right = normalized(right)\n mat[..., 2] = normalized(forward)\n mat[..., 1] = normalized(up)\n mat[..., 0] = right\n return mat\n\n\ndef get_rot_mat_from_pose_vec(vec):\n \"\"\"_summary_\n\n Args:\n vec (tensor): [N, M, 6]\n\n Returns:\n mat (tensor): [N, M, 3, 3]\n \"\"\"\n forward = vec[..., :3]\n up = vec[..., 3:6]\n return get_rot_mat_from_forward_up(forward, up)\n\n\ndef get_TRS(rot_mat, pos):\n \"\"\"_summary_\n\n Args:\n rot_mat (tensor): [N, 3, 3]\n pos (tensor): [N, 3]\n\n Returns:\n mat (tensor): [N, 4, 4]\n \"\"\"\n if isinstance(rot_mat, torch.Tensor):\n mat = torch.eye(4, device=pos.device).repeat(pos.shape[:-1] + (1, 1))\n elif isinstance(rot_mat, np.ndarray):\n mat = np.eye(4, dtype=np.float32)\n for _ in range(len(pos.shape) - 1):\n mat = mat[None]\n mat = np.tile(mat, pos.shape[:-1] + (1, 1))\n else:\n raise ValueError\n mat[..., :3, :3] = rot_mat\n mat[..., :3, 3] = pos\n mat = normalized_matrix(mat)\n return mat\n\n\ndef xzvec2mat(vec):\n \"\"\"_summary_\n\n Args:\n vec (tensor): [N, 4]\n\n Returns:\n mat (tensor): [N, 4, 4]\n \"\"\"\n vec_shape = vec.shape[:-1]\n if isinstance(vec, torch.Tensor):\n pos = torch.zeros(vec_shape + (3,))\n forward = torch.zeros(vec_shape + (3,))\n elif isinstance(vec, np.ndarray):\n pos = np.zeros(vec_shape + (3,))\n forward = np.zeros(vec_shape + (3,))\n else:\n raise ValueError\n\n pos[..., 0] = vec[..., 0]\n pos[..., 2] = vec[..., 1]\n forward[..., 0] = vec[..., 2]\n forward[..., 2] = vec[..., 3]\n rot_mat = get_rot_mat_from_forward(forward)\n mat = get_TRS(rot_mat, pos)\n return mat\n\n\ndef distance(vec1, vec2):\n return ((vec1 - vec2) ** 2).sum() ** 0.5\n\n\ndef get_relative_pose_from_vec(pose, root, N):\n root_p_mat = xzvec2mat(root)\n pose = pose.reshape(-1, N, 12)\n pose[..., :3] = get_position_from(pose[..., :3], root_p_mat)\n pose[..., 3:6] = get_direction_from(pose[..., 3:6], root_p_mat)\n pose[..., 6:9] = get_direction_from(pose[..., 6:9], root_p_mat)\n pose[..., 9:] = get_direction_from(pose[..., 9:], root_p_mat)\n pos = pose[..., 0, :3]\n rot = pose[..., 3:9].reshape(-1, N * 6)\n pose = np.concatenate((pos, rot), axis=-1)\n return pose\n\n\ndef get_forward_from_pos(pos):\n \"\"\"_summary_\n\n Args:\n pos (N, J, 3): joints positions of each frame\n\n Returns:\n _type_: _description_\n \"\"\"\n\n pos_y_vec = torch.tensor([0, 1, 0], dtype=torch.float32).to(pos.device)\n face_joint_indx = [2, 1, 17, 16]\n r_hip, l_hip, r_sdr, l_sdr = face_joint_indx # use hip and shoulder to get the cross vector\n cross_hip = pos[..., 0, r_hip, :] - pos[..., 0, l_hip, :]\n cross_sdr = pos[..., 0, r_sdr, :] - pos[..., 0, l_sdr, :]\n cross_vec = cross_hip + cross_sdr # (3, )\n forward_vec = torch.cross(pos_y_vec, cross_vec, dim=-1)\n forward_vec = normalized(forward_vec)\n return forward_vec\n\n\ndef project_point_along_ray(p, ray, keepnorm=False):\n \"\"\"_summary_\n\n Args:\n p (*, 3): point positions\n ray (*, 3): ray direction\n keepnorm: False -> project point on the ray,\n True -> project point on the ray and keep the point length\n\n Returns:\n _type_: _description_\n \"\"\"\n ray = normalized(ray)\n if keepnorm:\n new_p = ray * p.norm(dim=-1, keepdim=True)\n else:\n dot_product = torch.sum(p * ray, dim=-1, keepdim=True)\n new_p = dot_product * ray\n return new_p\n\n\ndef solve_point_along_ray_with_constraint(c, ray, p, constraint=\"x\"):\n \"\"\"_summary_\n\n Args:\n c (*,): constraint value\n ray (*, 3): ray direction\n p (*, 3): start point of the ray\n\n Returns:\n _type_: _description_\n \"\"\"\n ray = normalized(ray)\n if constraint == \"x\":\n ind = 0\n elif constraint == \"y\":\n ind = 1\n elif constraint == \"z\":\n ind = 2\n else:\n raise ValueError\n t = (c - p[..., ind]) / ray[..., ind]\n out_p = ray * t[..., None] + p\n\n return out_p\n\n\ndef calc_cosine(vec1, vec2, return_angle=False):\n \"\"\"_summary_\n\n Args:\n vec1 (*, 3): vector\n vec2 (*, 3): vector\n return_angle: True -> return angle, False -> return cosine\n\n Returns:\n _type_: _description_\n \"\"\"\n vec1 = normalized(vec1)\n vec2 = normalized(vec2)\n cosine = torch.sum(vec1 * vec2, dim=-1)\n if return_angle:\n return torch.acos(cosine)\n return cosine\n\n\n############################################\n#\n# quaternion assumes xyzw\n#\n############################################\n\n\ndef quat_xyzw2wxyz(quat):\n new_quat = torch.cat([quat[..., 3:4], quat[..., :3]], dim=-1)\n return new_quat\n\n\ndef quat_wxyz2xyzw(quat):\n new_quat = torch.cat([quat[..., 1:4], quat[..., :1]], dim=-1)\n return new_quat\n\n\ndef quat_mul(a, b):\n \"\"\"\n quaternion multiplication\n \"\"\"\n x1, y1, z1, w1 = a[..., 0], a[..., 1], a[..., 2], a[..., 3]\n x2, y2, z2, w2 = b[..., 0], b[..., 1], b[..., 2], b[..., 3]\n\n w = w1 * w2 - x1 * x2 - y1 * y2 - z1 * z2\n x = w1 * x2 + x1 * w2 + y1 * z2 - z1 * y2\n y = w1 * y2 + y1 * w2 + z1 * x2 - x1 * z2\n z = w1 * z2 + z1 * w2 + x1 * y2 - y1 * x2\n\n return torch.stack([x, y, z, w], dim=-1)\n\n\ndef quat_pos(x):\n \"\"\"\n make all the real part of the quaternion positive\n \"\"\"\n q = x\n z = (q[..., 3:] < 0).float()\n q = (1 - 2 * z) * q\n return q\n\n\ndef quat_abs(x):\n \"\"\"\n quaternion norm (unit quaternion represents a 3D rotation, which has norm of 1)\n \"\"\"\n x = x.norm(p=2, dim=-1)\n return x\n\n\ndef quat_unit(x):\n \"\"\"\n normalized quaternion with norm of 1\n \"\"\"\n norm = quat_abs(x).unsqueeze(-1)\n return x / (norm.clamp(min=1e-4))\n\n\ndef quat_conjugate(x):\n \"\"\"\n quaternion with its imaginary part negated\n \"\"\"\n return torch.cat([-x[..., :3], x[..., 3:]], dim=-1)\n\n\ndef quat_real(x):\n \"\"\"\n real component of the quaternion\n \"\"\"\n return x[..., 3]\n\n\ndef quat_imaginary(x):\n \"\"\"\n imaginary components of the quaternion\n \"\"\"\n return x[..., :3]\n\n\ndef quat_norm_check(x):\n \"\"\"\n verify that a quaternion has norm 1\n \"\"\"\n assert bool((abs(x.norm(p=2, dim=-1) - 1) < 1e-3).all()), \"the quaternion is has non-1 norm: {}\".format(\n abs(x.norm(p=2, dim=-1) - 1)\n )\n assert bool((x[..., 3] >= 0).all()), \"the quaternion has negative real part\"\n\n\ndef quat_normalize(q):\n \"\"\"\n Construct 3D rotation from quaternion (the quaternion needs not to be normalized).\n \"\"\"\n q = quat_unit(quat_pos(q)) # normalized to positive and unit quaternion\n return q\n\n\ndef quat_from_xyz(xyz):\n \"\"\"\n Construct 3D rotation from the imaginary component\n \"\"\"\n w = (1.0 - xyz.norm()).unsqueeze(-1)\n assert bool((w >= 0).all()), \"xyz has its norm greater than 1\"\n return torch.cat([xyz, w], dim=-1)\n\n\ndef quat_identity(shape: List[int]):\n \"\"\"\n Construct 3D identity rotation given shape\n \"\"\"\n w = torch.ones(shape + (1,))\n xyz = torch.zeros(shape + (3,))\n q = torch.cat([xyz, w], dim=-1)\n return quat_normalize(q)\n\n\ndef tgm_quat_from_angle_axis(angle, axis, degree: bool = False):\n \"\"\"Create a 3D rotation from angle and axis of rotation. The rotation is counter-clockwise\n along the axis.\n\n The rotation can be interpreted as a_R_b where frame \"b\" is the new frame that\n gets rotated counter-clockwise along the axis from frame \"a\"\n\n :param angle: angle of rotation\n :type angle: Tensor\n :param axis: axis of rotation\n :type axis: Tensor\n :param degree: put True here if the angle is given by degree\n :type degree: bool, optional, default=False\n \"\"\"\n if degree:\n angle = angle / 180.0 * math.pi\n theta = (angle / 2).unsqueeze(-1)\n axis = axis / (axis.norm(p=2, dim=-1, keepdim=True).clamp(min=1e-4))\n xyz = axis * theta.sin()\n w = theta.cos()\n return quat_normalize(torch.cat([w, xyz], dim=-1))\n\n\ndef quat_from_rotation_matrix(m):\n \"\"\"\n Construct a 3D rotation from a valid 3x3 rotation matrices.\n Reference can be found here:\n http://www.cg.info.hiroshima-cu.ac.jp/~miyazaki/knowledge/teche52.html\n\n :param m: 3x3 orthogonal rotation matrices.\n :type m: Tensor\n\n :rtype: Tensor\n \"\"\"\n m = m.unsqueeze(0)\n diag0 = m[..., 0, 0]\n diag1 = m[..., 1, 1]\n diag2 = m[..., 2, 2]\n\n # Math stuff.\n w = (((diag0 + diag1 + diag2 + 1.0) / 4.0).clamp(0.0, None)) ** 0.5\n x = (((diag0 - diag1 - diag2 + 1.0) / 4.0).clamp(0.0, None)) ** 0.5\n y = (((-diag0 + diag1 - diag2 + 1.0) / 4.0).clamp(0.0, None)) ** 0.5\n z = (((-diag0 - diag1 + diag2 + 1.0) / 4.0).clamp(0.0, None)) ** 0.5\n\n # Only modify quaternions where w > x, y, z.\n c0 = (w >= x) & (w >= y) & (w >= z)\n x[c0] *= (m[..., 2, 1][c0] - m[..., 1, 2][c0]).sign()\n y[c0] *= (m[..., 0, 2][c0] - m[..., 2, 0][c0]).sign()\n z[c0] *= (m[..., 1, 0][c0] - m[..., 0, 1][c0]).sign()\n\n # Only modify quaternions where x > w, y, z\n c1 = (x >= w) & (x >= y) & (x >= z)\n w[c1] *= (m[..., 2, 1][c1] - m[..., 1, 2][c1]).sign()\n y[c1] *= (m[..., 1, 0][c1] + m[..., 0, 1][c1]).sign()\n z[c1] *= (m[..., 0, 2][c1] + m[..., 2, 0][c1]).sign()\n\n # Only modify quaternions where y > w, x, z.\n c2 = (y >= w) & (y >= x) & (y >= z)\n w[c2] *= (m[..., 0, 2][c2] - m[..., 2, 0][c2]).sign()\n x[c2] *= (m[..., 1, 0][c2] + m[..., 0, 1][c2]).sign()\n z[c2] *= (m[..., 2, 1][c2] + m[..., 1, 2][c2]).sign()\n\n # Only modify quaternions where z > w, x, y.\n c3 = (z >= w) & (z >= x) & (z >= y)\n w[c3] *= (m[..., 1, 0][c3] - m[..., 0, 1][c3]).sign()\n x[c3] *= (m[..., 2, 0][c3] + m[..., 0, 2][c3]).sign()\n y[c3] *= (m[..., 2, 1][c3] + m[..., 1, 2][c3]).sign()\n\n return quat_normalize(torch.stack([x, y, z, w], dim=-1)).squeeze(0)\n\n\ndef quat_mul_norm(x, y):\n \"\"\"\n Combine two set of 3D rotations together using \\**\\* operator. The shape needs to be\n broadcastable\n \"\"\"\n return quat_normalize(quat_mul(x, y))\n\n\ndef quat_rotate(rot, vec):\n \"\"\"\n Rotate a 3D vector with the 3D rotation\n \"\"\"\n other_q = torch.cat([vec, torch.zeros_like(vec[..., :1])], dim=-1)\n return quat_imaginary(quat_mul(quat_mul(rot, other_q), quat_conjugate(rot)))\n\n\ndef quat_inverse(x):\n \"\"\"\n The inverse of the rotation\n \"\"\"\n return quat_conjugate(x)\n\n\ndef quat_identity_like(x):\n \"\"\"\n Construct identity 3D rotation with the same shape\n \"\"\"\n return quat_identity(x.shape[:-1])\n\n\ndef quat_angle_axis(x):\n \"\"\"\n The (angle, axis) representation of the rotation. The axis is normalized to unit length.\n The angle is guaranteed to be between [0, pi].\n \"\"\"\n s = 2 * (x[..., 3] ** 2) - 1\n angle = s.clamp(-1, 1).arccos() # just to be safe\n axis = x[..., :3]\n axis /= axis.norm(p=2, dim=-1, keepdim=True).clamp(min=1e-4)\n return angle, axis\n\n\ndef quat_yaw_rotation(x, z_up: bool = True):\n \"\"\"\n Yaw rotation (rotation along z-axis)\n \"\"\"\n q = x\n if z_up:\n q = torch.cat([torch.zeros_like(q[..., 0:2]), q[..., 2:3], q[..., 3:]], dim=-1)\n else:\n q = torch.cat(\n [\n torch.zeros_like(q[..., 0:1]),\n q[..., 1:2],\n torch.zeros_like(q[..., 2:3]),\n q[..., 3:4],\n ],\n dim=-1,\n )\n return quat_normalize(q)\n\n\ndef transform_from_rotation_translation(r: Optional[torch.Tensor] = None, t: Optional[torch.Tensor] = None):\n \"\"\"\n Construct a transform from a quaternion and 3D translation. Only one of them can be None.\n \"\"\"\n assert r is not None or t is not None, \"rotation and translation can't be all None\"\n if r is None:\n assert t is not None\n r = quat_identity(list(t.shape))\n if t is None:\n t = torch.zeros(list(r.shape) + [3])\n return torch.cat([r, t], dim=-1)\n\n\ndef transform_identity(shape: List[int]):\n \"\"\"\n Identity transformation with given shape\n \"\"\"\n r = quat_identity(shape)\n t = torch.zeros(shape + [3])\n return transform_from_rotation_translation(r, t)\n\n\ndef transform_rotation(x):\n \"\"\"Get rotation from transform\"\"\"\n return x[..., :4]\n\n\ndef transform_translation(x):\n \"\"\"Get translation from transform\"\"\"\n return x[..., 4:]\n\n\ndef transform_inverse(x):\n \"\"\"\n Inverse transformation\n \"\"\"\n inv_so3 = quat_inverse(transform_rotation(x))\n return transform_from_rotation_translation(r=inv_so3, t=quat_rotate(inv_so3, -transform_translation(x)))\n\n\ndef transform_identity_like(x):\n \"\"\"\n identity transformation with the same shape\n \"\"\"\n return transform_identity(x.shape)\n\n\ndef transform_mul(x, y):\n \"\"\"\n Combine two transformation together\n \"\"\"\n z = transform_from_rotation_translation(\n r=quat_mul_norm(transform_rotation(x), transform_rotation(y)),\n t=quat_rotate(transform_rotation(x), transform_translation(y)) + transform_translation(x),\n )\n return z\n\n\ndef transform_apply(rot, vec):\n \"\"\"\n Transform a 3D vector\n \"\"\"\n assert isinstance(vec, torch.Tensor)\n return quat_rotate(transform_rotation(rot), vec) + transform_translation(rot)\n\n\ndef rot_matrix_det(x):\n \"\"\"\n Return the determinant of the 3x3 matrix. The shape of the tensor will be as same as the\n shape of the matrix\n \"\"\"\n a, b, c = x[..., 0, 0], x[..., 0, 1], x[..., 0, 2]\n d, e, f = x[..., 1, 0], x[..., 1, 1], x[..., 1, 2]\n g, h, i = x[..., 2, 0], x[..., 2, 1], x[..., 2, 2]\n t1 = a * (e * i - f * h)\n t2 = b * (d * i - f * g)\n t3 = c * (d * h - e * g)\n return t1 - t2 + t3\n\n\ndef rot_matrix_integrity_check(x):\n \"\"\"\n Verify that a rotation matrix has a determinant of one and is orthogonal\n \"\"\"\n det = rot_matrix_det(x)\n assert bool((abs(det - 1) < 1e-3).all()), \"the matrix has non-one determinant\"\n rtr = x @ x.permute(torch.arange(x.dim() - 2), -1, -2)\n rtr_gt = rtr.zeros_like()\n rtr_gt[..., 0, 0] = 1\n rtr_gt[..., 1, 1] = 1\n rtr_gt[..., 2, 2] = 1\n assert bool(((rtr - rtr_gt) < 1e-3).all()), \"the matrix is not orthogonal\"\n\n\ndef rot_matrix_from_quaternion(q):\n \"\"\"\n Construct rotation matrix from quaternion\n \"\"\"\n # Shortcuts for individual elements (using wikipedia's convention)\n qi, qj, qk, qr = q[..., 0], q[..., 1], q[..., 2], q[..., 3]\n\n # Set individual elements\n R00 = 1.0 - 2.0 * (qj**2 + qk**2)\n R01 = 2 * (qi * qj - qk * qr)\n R02 = 2 * (qi * qk + qj * qr)\n R10 = 2 * (qi * qj + qk * qr)\n R11 = 1.0 - 2.0 * (qi**2 + qk**2)\n R12 = 2 * (qj * qk - qi * qr)\n R20 = 2 * (qi * qk - qj * qr)\n R21 = 2 * (qj * qk + qi * qr)\n R22 = 1.0 - 2.0 * (qi**2 + qj**2)\n\n R0 = torch.stack([R00, R01, R02], dim=-1)\n R1 = torch.stack([R10, R11, R12], dim=-1)\n R2 = torch.stack([R20, R21, R22], dim=-1)\n\n R = torch.stack([R0, R1, R2], dim=-2)\n\n return R\n\n\ndef euclidean_to_rotation_matrix(x):\n \"\"\"\n Get the rotation matrix on the top-left corner of a Euclidean transformation matrix\n \"\"\"\n return x[..., :3, :3]\n\n\ndef euclidean_integrity_check(x):\n euclidean_to_rotation_matrix(x) # check 3d-rotation matrix\n assert bool((x[..., 3, :3] == 0).all()), \"the last row is illegal\"\n assert bool((x[..., 3, 3] == 1).all()), \"the last row is illegal\"\n\n\ndef euclidean_translation(x):\n \"\"\"\n Get the translation vector located at the last column of the matrix\n \"\"\"\n return x[..., :3, 3]\n\n\ndef euclidean_inverse(x):\n \"\"\"\n Compute the matrix that represents the inverse rotation\n \"\"\"\n s = x.zeros_like()\n irot = quat_inverse(quat_from_rotation_matrix(x))\n s[..., :3, :3] = irot\n s[..., :3, 4] = quat_rotate(irot, -euclidean_translation(x))\n return s\n\n\ndef euclidean_to_transform(transformation_matrix):\n \"\"\"\n Construct a transform from a Euclidean transformation matrix\n \"\"\"\n return transform_from_rotation_translation(\n r=quat_from_rotation_matrix(m=euclidean_to_rotation_matrix(transformation_matrix)),\n t=euclidean_translation(transformation_matrix),\n )\n\n\ndef to_torch(x, dtype=torch.float, device=\"cuda:0\", requires_grad=False):\n return torch.tensor(x, dtype=dtype, device=device, requires_grad=requires_grad)\n\n\ndef quat_mul(a, b):\n assert a.shape == b.shape\n shape = a.shape\n a = a.reshape(-1, 4)\n b = b.reshape(-1, 4)\n\n x1, y1, z1, w1 = a[:, 0], a[:, 1], a[:, 2], a[:, 3]\n x2, y2, z2, w2 = b[:, 0], b[:, 1], b[:, 2], b[:, 3]\n ww = (z1 + x1) * (x2 + y2)\n yy = (w1 - y1) * (w2 + z2)\n zz = (w1 + y1) * (w2 - z2)\n xx = ww + yy + zz\n qq = 0.5 * (xx + (z1 - x1) * (x2 - y2))\n w = qq - ww + (z1 - y1) * (y2 - z2)\n x = qq - xx + (x1 + w1) * (x2 + w2)\n y = qq - yy + (w1 - x1) * (y2 + z2)\n z = qq - zz + (z1 + y1) * (w2 - x2)\n\n quat = torch.stack([x, y, z, w], dim=-1).view(shape)\n\n return quat\n\n\ndef normalize(x, eps: float = 1e-9):\n return x / x.norm(p=2, dim=-1).clamp(min=eps, max=None).unsqueeze(-1)\n\n\ndef quat_apply(a, b):\n shape = b.shape\n a = a.reshape(-1, 4)\n b = b.reshape(-1, 3)\n xyz = a[:, :3]\n t = xyz.cross(b, dim=-1) * 2\n return (b + a[:, 3:] * t + xyz.cross(t, dim=-1)).view(shape)\n\n\ndef quat_rotate(q, v):\n shape = q.shape\n q_w = q[:, -1]\n q_vec = q[:, :3]\n a = v * (2.0 * q_w**2 - 1.0).unsqueeze(-1)\n b = torch.cross(q_vec, v, dim=-1) * q_w.unsqueeze(-1) * 2.0\n c = q_vec * torch.bmm(q_vec.view(shape[0], 1, 3), v.view(shape[0], 3, 1)).squeeze(-1) * 2.0\n return a + b + c\n\n\ndef quat_rotate_inverse(q, v):\n shape = q.shape\n q_w = q[:, -1]\n q_vec = q[:, :3]\n a = v * (2.0 * q_w**2 - 1.0).unsqueeze(-1)\n b = torch.cross(q_vec, v, dim=-1) * q_w.unsqueeze(-1) * 2.0\n c = q_vec * torch.bmm(q_vec.view(shape[0], 1, 3), v.view(shape[0], 3, 1)).squeeze(-1) * 2.0\n return a - b + c\n\n\ndef quat_conjugate(a):\n shape = a.shape\n a = a.reshape(-1, 4)\n return torch.cat((-a[:, :3], a[:, -1:]), dim=-1).view(shape)\n\n\ndef quat_unit(a):\n return normalize(a)\n\n\ndef quat_from_angle_axis(angle, axis):\n theta = (angle / 2).unsqueeze(-1)\n xyz = normalize(axis) * torch.sin(theta.clone())\n w = torch.cos(theta.clone())\n return quat_unit(torch.cat([xyz, w], dim=-1))\n\n\ndef normalize_angle(x):\n return torch.atan2(torch.sin(x.clone()), torch.cos(x.clone()))\n\n\ndef tf_inverse(q, t):\n q_inv = quat_conjugate(q)\n return q_inv, -quat_apply(q_inv, t)\n\n\ndef tf_apply(q, t, v):\n return quat_apply(q, v) + t\n\n\ndef tf_vector(q, v):\n return quat_apply(q, v)\n\n\ndef tf_combine(q1, t1, q2, t2):\n return quat_mul(q1, q2), quat_apply(q1, t2) + t1\n\n\ndef get_basis_vector(q, v):\n return quat_rotate(q, v)\n\n\ndef get_axis_params(value, axis_idx, x_value=0.0, dtype=float, n_dims=3):\n \"\"\"construct arguments to `Vec` according to axis index.\"\"\"\n zs = np.zeros((n_dims,))\n assert axis_idx < n_dims, \"the axis dim should be within the vector dimensions\"\n zs[axis_idx] = 1.0\n params = np.where(zs == 1.0, value, zs)\n params[0] = x_value\n return list(params.astype(dtype))\n\n\ndef copysign(a, b):\n # type: (float, Tensor) -> Tensor\n a = torch.tensor(a, device=b.device, dtype=torch.float).repeat(b.shape[0])\n return torch.abs(a) * torch.sign(b)\n\n\ndef get_euler_xyz(q):\n qx, qy, qz, qw = 0, 1, 2, 3\n # roll (x-axis rotation)\n sinr_cosp = 2.0 * (q[:, qw] * q[:, qx] + q[:, qy] * q[:, qz])\n cosr_cosp = q[:, qw] * q[:, qw] - q[:, qx] * q[:, qx] - q[:, qy] * q[:, qy] + q[:, qz] * q[:, qz]\n roll = torch.atan2(sinr_cosp, cosr_cosp)\n\n # pitch (y-axis rotation)\n sinp = 2.0 * (q[:, qw] * q[:, qy] - q[:, qz] * q[:, qx])\n pitch = torch.where(torch.abs(sinp) >= 1, copysign(np.pi / 2.0, sinp), torch.asin(sinp))\n\n # yaw (z-axis rotation)\n siny_cosp = 2.0 * (q[:, qw] * q[:, qz] + q[:, qx] * q[:, qy])\n cosy_cosp = q[:, qw] * q[:, qw] + q[:, qx] * q[:, qx] - q[:, qy] * q[:, qy] - q[:, qz] * q[:, qz]\n yaw = torch.atan2(siny_cosp, cosy_cosp)\n\n return roll % (2 * np.pi), pitch % (2 * np.pi), yaw % (2 * np.pi)\n\n\ndef quat_from_euler_xyz(roll, pitch, yaw):\n cy = torch.cos(yaw * 0.5)\n sy = torch.sin(yaw * 0.5)\n cr = torch.cos(roll * 0.5)\n sr = torch.sin(roll * 0.5)\n cp = torch.cos(pitch * 0.5)\n sp = torch.sin(pitch * 0.5)\n\n qw = cy * cr * cp + sy * sr * sp\n qx = cy * sr * cp - sy * cr * sp\n qy = cy * cr * sp + sy * sr * cp\n qz = sy * cr * cp - cy * sr * sp\n\n return torch.stack([qx, qy, qz, qw], dim=-1)\n\n\ndef torch_rand_float(lower, upper, shape, device):\n # type: (float, float, Tuple[int, int], str) -> Tensor\n return (upper - lower) * torch.rand(*shape, device=device) + lower\n\n\ndef torch_random_dir_2(shape, device):\n # type: (Tuple[int, int], str) -> Tensor\n angle = torch_rand_float(-np.pi, np.pi, shape, device).squeeze(-1)\n return torch.stack([torch.cos(angle), torch.sin(angle)], dim=-1)\n\n\ndef tensor_clamp(t, min_t, max_t):\n return torch.max(torch.min(t, max_t), min_t)\n\n\ndef scale(x, lower, upper):\n return 0.5 * (x + 1.0) * (upper - lower) + lower\n\n\ndef unscale(x, lower, upper):\n return (2.0 * x - upper - lower) / (upper - lower)\n\n\ndef unscale_np(x, lower, upper):\n return (2.0 * x - upper - lower) / (upper - lower)\n\n\ndef quat_to_angle_axis(q):\n # type: (Tensor) -> Tuple[Tensor, Tensor]\n # computes axis-angle representation from quaternion q\n # q must be normalized\n min_theta = 1e-5\n qx, qy, qz, qw = 0, 1, 2, 3\n\n sin_theta = torch.sqrt(1 - q[..., qw] * q[..., qw])\n angle = 2 * torch.acos(q[..., qw])\n angle = normalize_angle(angle)\n sin_theta_expand = sin_theta.unsqueeze(-1)\n axis = q[..., qx:qw] / sin_theta_expand\n\n mask = torch.abs(sin_theta) > min_theta\n default_axis = torch.zeros_like(axis)\n default_axis[..., -1] = 1\n\n angle = torch.where(mask, angle, torch.zeros_like(angle))\n mask_expand = mask.unsqueeze(-1)\n axis = torch.where(mask_expand, axis, default_axis)\n return angle, axis\n\n\ndef angle_axis_to_exp_map(angle, axis):\n # type: (Tensor, Tensor) -> Tensor\n # compute exponential map from axis-angle\n angle_expand = angle.unsqueeze(-1)\n exp_map = angle_expand * axis\n return exp_map\n\n\ndef quat_to_exp_map(q):\n # type: (Tensor) -> Tensor\n # compute exponential map from quaternion\n # q must be normalized\n angle, axis = quat_to_angle_axis(q)\n exp_map = angle_axis_to_exp_map(angle, axis)\n return exp_map\n\n\ndef quat_to_tan_norm(q):\n # type: (Tensor) -> Tensor\n # represents a rotation using the tangent and normal vectors\n ref_tan = torch.zeros_like(q[..., 0:3])\n ref_tan[..., 0] = 1\n tan = quat_rotate(q, ref_tan)\n\n ref_norm = torch.zeros_like(q[..., 0:3])\n ref_norm[..., -1] = 1\n norm = quat_rotate(q, ref_norm)\n\n norm_tan = torch.cat([tan, norm], dim=len(tan.shape) - 1)\n return norm_tan\n\n\ndef euler_xyz_to_exp_map(roll, pitch, yaw):\n # type: (Tensor, Tensor, Tensor) -> Tensor\n q = quat_from_euler_xyz(roll, pitch, yaw)\n exp_map = quat_to_exp_map(q)\n return exp_map\n\n\ndef exp_map_to_angle_axis(exp_map):\n min_theta = 1e-5\n\n angle = torch.norm(exp_map.clone(), dim=-1) + 1e-6\n angle_exp = torch.unsqueeze(angle, dim=-1)\n axis = exp_map.clone() / angle_exp.clone()\n angle = normalize_angle(angle)\n\n default_axis = torch.zeros_like(exp_map)\n default_axis[..., -1] = 1\n\n mask = torch.abs(angle) > min_theta\n angle = torch.where(mask, angle, torch.zeros_like(angle))\n mask_expand = mask.unsqueeze(-1)\n axis = torch.where(mask_expand, axis, default_axis)\n\n return angle, axis\n\n\ndef exp_map_to_quat(exp_map):\n angle, axis = exp_map_to_angle_axis(exp_map)\n q = quat_from_angle_axis(angle, axis)\n return q\n\n\ndef slerp(q0, q1, t):\n # type: (Tensor, Tensor, Tensor) -> Tensor\n cos_half_theta = torch.sum(q0 * q1, dim=-1)\n\n neg_mask = cos_half_theta < 0\n q1 = q1.clone()\n q1[neg_mask] = -q1[neg_mask]\n cos_half_theta = torch.abs(cos_half_theta)\n cos_half_theta = torch.unsqueeze(cos_half_theta, dim=-1)\n\n half_theta = torch.acos(cos_half_theta)\n sin_half_theta = torch.sqrt(1.0 - cos_half_theta * cos_half_theta)\n\n ratioA = torch.sin((1 - t) * half_theta) / sin_half_theta\n ratioB = torch.sin(t * half_theta) / sin_half_theta\n\n new_q = ratioA * q0 + ratioB * q1\n\n new_q = torch.where(torch.abs(sin_half_theta) < 0.001, 0.5 * q0 + 0.5 * q1, new_q)\n new_q = torch.where(torch.abs(cos_half_theta) >= 1, q0, new_q)\n\n return new_q\n\n\ndef calc_heading_vec(q, head_ind=0):\n # type: (Tensor, int) -> Tensor\n # calculate heading direction from quaternion\n # the heading is the direction vector\n # q must be normalized\n ref_dir = torch.zeros_like(q[..., 0:3])\n ref_dir[..., head_ind] = 1\n rot_dir = quat_rotate(q, ref_dir)\n\n return rot_dir\n\n\ndef calc_heading(q, head_ind=0, gravity_axis=\"z\"):\n # type: (Tensor, int, str) -> Tensor\n # calculate heading direction from quaternion\n # the heading is the direction on the xy plane\n # q must be normalized\n ref_dir = torch.zeros_like(q[..., 0:3])\n ref_dir[..., head_ind] = 1\n # ref_dir[..., 0] = 1\n shape = ref_dir.shape[:-1]\n q = q.reshape((-1, 4))\n ref_dir = ref_dir.reshape(-1, 3)\n rot_dir = quat_rotate(q, ref_dir)\n rot_dir = rot_dir.reshape(shape + (3,))\n if gravity_axis == \"z\":\n heading = torch.atan2(rot_dir[..., 1], rot_dir[..., 0])\n elif gravity_axis == \"y\":\n heading = torch.atan2(rot_dir[..., 0], rot_dir[..., 2])\n elif gravity_axis == \"x\":\n heading = torch.atan2(rot_dir[..., 2], rot_dir[..., 1])\n return heading\n\n\ndef calc_heading_quat(q, head_ind=0, gravity_axis=\"z\"):\n # type: (Tensor, int, str) -> Tensor\n # calculate heading rotation from quaternion\n # the heading is the direction on the xy plane\n # q must be normalized\n heading = calc_heading(q, head_ind, gravity_axis=gravity_axis)\n axis = torch.zeros_like(q[..., 0:3])\n if gravity_axis == \"z\":\n g_axis = 2\n elif gravity_axis == \"y\":\n g_axis = 1\n elif gravity_axis == \"x\":\n g_axis = 0\n axis[..., g_axis] = 1\n\n heading_q = quat_from_angle_axis(heading, axis)\n return heading_q\n\n\ndef calc_heading_quat_inv(q, head_ind=0):\n # type: (Tensor, int) -> Tensor\n # calculate heading rotation from quaternion\n # the heading is the direction on the xy plane\n # q must be normalized\n heading = calc_heading(q, head_ind)\n axis = torch.zeros_like(q[..., 0:3])\n axis[..., 2] = 1\n\n heading_q = quat_from_angle_axis(-heading, axis)\n return heading_q\n\n\ndef forward_kinematics(mat, parent):\n \"\"\"_summary_\n\n Args:\n mat ([..., N, 3, 3]): _description_\n parent (): _description_\n \"\"\"\n if isinstance(mat, torch.Tensor):\n rotations = torch.eye(mat.shape[-1], device=mat.device)\n rotations = rotations.repeat(mat.shape[:-2] + (1, 1))\n else:\n rotations = np.eye(mat.shape[-1], dtype=np.float32)\n rotations = np.tile(rotations, mat.shape[:-2] + (1, 1))\n for i in range(mat.shape[-3]):\n if parent[i] != -1:\n if isinstance(mat, torch.Tensor):\n # this way make gradient flow\n new_mat = get_mat_BfromA(rotations[..., parent[i], :, :], mat[..., i, :, :])\n rotations = torch.cat(\n (\n rotations[..., :i, :, :],\n new_mat[..., None, :, :],\n rotations[..., i + 1 :, :, :],\n ),\n dim=-3,\n )\n else:\n rotations[..., i, :, :] = get_mat_BfromA(rotations[..., parent[i], :, :], mat[..., i, :, :])\n else:\n if isinstance(mat, torch.Tensor):\n # this way make gradient flow\n rotations = torch.cat((mat[..., : i + 1, :, :], rotations[..., i + 1 :, :, :]), dim=-3)\n else:\n rotations[..., i, :, :] = mat[..., i, :, :]\n return rotations\n" }, { "path": "eval/GVHMR/hmr4d/utils/net_utils.py", "content": "import torch\nimport torch.nn as nn\nimport torch.nn.functional as F\nfrom pathlib import Path\nfrom hmr4d.utils.pylogger import Log\nfrom pytorch_lightning.utilities.memory import recursive_detach\nfrom einops import repeat, rearrange\nfrom scipy.ndimage._filters import _gaussian_kernel1d\n\n\ndef load_pretrained_model(model, ckpt_path):\n \"\"\"\n Load ckpt to model with strategy\n \"\"\"\n assert Path(ckpt_path).exists()\n # use model's own load_pretrained_model method\n if hasattr(model, \"load_pretrained_model\"):\n model.load_pretrained_model(ckpt_path)\n else:\n Log.info(f\"Loading ckpt: {ckpt_path}\")\n ckpt = torch.load(ckpt_path, \"cpu\")\n model.load_state_dict(ckpt, strict=True)\n\n\ndef find_last_ckpt_path(dirpath):\n \"\"\"\n Assume ckpt is named as e{}* or last*, following the convention of pytorch-lightning.\n \"\"\"\n assert dirpath is not None\n dirpath = Path(dirpath)\n assert dirpath.exists()\n # Priority 1: last.ckpt\n auto_last_ckpt_path = dirpath / \"last.ckpt\"\n if auto_last_ckpt_path.exists():\n return auto_last_ckpt_path\n\n # Priority 2\n model_paths = []\n for p in sorted(list(dirpath.glob(\"*.ckpt\"))):\n if \"last\" in p.name:\n continue\n model_paths.append(p)\n if len(model_paths) > 0:\n return model_paths[-1]\n else:\n Log.info(\"No checkpoint found, set model_path to None\")\n return None\n\n\ndef get_resume_ckpt_path(resume_mode, ckpt_dir=None):\n if Path(resume_mode).exists(): # This is a path\n return resume_mode\n assert resume_mode == \"last\"\n return find_last_ckpt_path(ckpt_dir)\n\n\ndef select_state_dict_by_prefix(state_dict, prefix, new_prefix=\"\"):\n \"\"\"\n For each weight that start with {old_prefix}, remove the {old_prefic} and form a new state_dict.\n Args:\n state_dict: dict\n prefix: str\n new_prefix: str, if exists, the new key will be {new_prefix} + {old_key[len(prefix):]}\n Returns:\n state_dict_new: dict\n \"\"\"\n state_dict_new = {}\n for k in list(state_dict.keys()):\n if k.startswith(prefix):\n new_key = new_prefix + k[len(prefix) :]\n state_dict_new[new_key] = state_dict[k]\n return state_dict_new\n\n\ndef detach_to_cpu(in_dict):\n return recursive_detach(in_dict, to_cpu=True)\n\n\ndef to_cuda(data):\n \"\"\"Move data in the batch to cuda(), carefully handle data that is not tensor\"\"\"\n if isinstance(data, torch.Tensor):\n return data.cuda()\n elif isinstance(data, dict):\n return {k: to_cuda(v) for k, v in data.items()}\n elif isinstance(data, list):\n return [to_cuda(v) for v in data]\n else:\n return data\n\n\ndef get_valid_mask(max_len, valid_len, device=\"cpu\"):\n mask = torch.zeros(max_len, dtype=torch.bool).to(device)\n mask[:valid_len] = True\n return mask\n\n\ndef length_to_mask(lengths, max_len):\n \"\"\"\n Returns: (B, max_len)\n \"\"\"\n mask = torch.arange(max_len, device=lengths.device).expand(len(lengths), max_len) < lengths.unsqueeze(1)\n return mask\n\n\ndef repeat_to_max_len(x, max_len, dim=0):\n \"\"\"Repeat last frame to max_len along dim\"\"\"\n assert isinstance(x, torch.Tensor)\n if x.shape[dim] == max_len:\n return x\n elif x.shape[dim] < max_len:\n x = x.clone()\n x = x.transpose(0, dim)\n x = torch.cat([x, repeat(x[-1:], \"b ... -> (b r) ...\", r=max_len - x.shape[0])])\n x = x.transpose(0, dim)\n return x\n else:\n raise ValueError(f\"Unexpected length v.s. max_len: {x.shape[0]} v.s. {max_len}\")\n\n\ndef repeat_to_max_len_dict(x_dict, max_len, dim=0):\n for k, v in x_dict.items():\n x_dict[k] = repeat_to_max_len(v, max_len, dim=dim)\n return x_dict\n\n\nclass Transpose(nn.Module):\n def __init__(self, dim1, dim2):\n super(Transpose, self).__init__()\n self.dim1 = dim1\n self.dim2 = dim2\n\n def forward(self, x):\n return x.transpose(self.dim1, self.dim2)\n\n\nclass GaussianSmooth(nn.Module):\n def __init__(self, sigma=3, dim=-1):\n super(GaussianSmooth, self).__init__()\n kernel_smooth = _gaussian_kernel1d(sigma=sigma, order=0, radius=int(4 * sigma + 0.5))\n kernel_smooth = torch.from_numpy(kernel_smooth).float()[None, None] # (1, 1, K)\n self.register_buffer(\"kernel_smooth\", kernel_smooth, persistent=False)\n self.dim = dim\n\n def forward(self, x):\n \"\"\"x (..., f, ...) f at dim\"\"\"\n rad = self.kernel_smooth.size(-1) // 2\n\n x = x.transpose(self.dim, -1)\n x_shape = x.shape[:-1]\n x = rearrange(x, \"... f -> (...) 1 f\") # (NB, 1, f)\n x = F.pad(x[None], (rad, rad, 0, 0), mode=\"replicate\")[0]\n x = F.conv1d(x, self.kernel_smooth)\n x = x.squeeze(1).reshape(*x_shape, -1) # (..., f)\n x = x.transpose(-1, self.dim)\n return x\n\n\ndef gaussian_smooth(x, sigma=3, dim=-1):\n kernel_smooth = _gaussian_kernel1d(sigma=sigma, order=0, radius=int(4 * sigma + 0.5))\n kernel_smooth = torch.from_numpy(kernel_smooth).float()[None, None].to(x) # (1, 1, K)\n rad = kernel_smooth.size(-1) // 2\n\n x = x.transpose(dim, -1)\n x_shape = x.shape[:-1]\n x = rearrange(x, \"... f -> (...) 1 f\") # (NB, 1, f)\n x = F.pad(x[None], (rad, rad, 0, 0), mode=\"replicate\")[0]\n x = F.conv1d(x, kernel_smooth)\n x = x.squeeze(1).reshape(*x_shape, -1) # (..., f)\n x = x.transpose(-1, dim)\n return x\n\n\ndef moving_average_smooth(x, window_size=5, dim=-1):\n kernel_smooth = torch.ones(window_size).float() / window_size\n kernel_smooth = kernel_smooth[None, None].to(x) # (1, 1, window_size)\n rad = kernel_smooth.size(-1) // 2\n\n x = x.transpose(dim, -1)\n x_shape = x.shape[:-1]\n x = rearrange(x, \"... f -> (...) 1 f\") # (NB, 1, f)\n x = F.pad(x[None], (rad, rad, 0, 0), mode=\"replicate\")[0]\n x = F.conv1d(x, kernel_smooth)\n x = x.squeeze(1).reshape(*x_shape, -1) # (..., f)\n x = x.transpose(-1, dim)\n return x\n" }, { "path": "eval/GVHMR/hmr4d/utils/preproc/__init__.py", "content": "try:\n from hmr4d.utils.preproc.tracker import Tracker\n from hmr4d.utils.preproc.vitfeat_extractor import Extractor\n from hmr4d.utils.preproc.vitpose import VitPoseExtractor\n from hmr4d.utils.preproc.slam import SLAMModel\nexcept:\n pass\n" }, { "path": "eval/GVHMR/hmr4d/utils/preproc/slam.py", "content": "import cv2\nimport time\nimport torch\nfrom multiprocessing import Process, Queue\n\ntry:\n from dpvo.utils import Timer\n from dpvo.dpvo import DPVO\n from dpvo.config import cfg\nexcept:\n pass\n\n\nfrom hmr4d import PROJ_ROOT\nfrom hmr4d.utils.geo.hmr_cam import estimate_focal_length\n\n\nclass SLAMModel(object):\n def __init__(self, video_path, width, height, intrinsics=None, stride=1, skip=0, buffer=2048, resize=0.5):\n \"\"\"\n Args:\n intrinsics: [fx, fy, cx, cy]\n \"\"\"\n if intrinsics is None:\n print(\"Estimating focal length\")\n focal_length = estimate_focal_length(width, height)\n intrinsics = torch.tensor([focal_length, focal_length, width / 2.0, height / 2.0])\n else:\n intrinsics = intrinsics.clone()\n\n self.dpvo_cfg = str(PROJ_ROOT / \"third-party/DPVO/config/default.yaml\")\n self.dpvo_ckpt = \"inputs/checkpoints/dpvo/dpvo.pth\"\n\n self.buffer = buffer\n self.times = []\n self.slam = None\n self.queue = Queue(maxsize=8)\n self.reader = Process(target=video_stream, args=(self.queue, video_path, intrinsics, stride, skip, resize))\n self.reader.start()\n\n def track(self):\n (t, image, intrinsics) = self.queue.get()\n\n if t < 0:\n return False\n\n image = torch.from_numpy(image).permute(2, 0, 1).cuda()\n intrinsics = intrinsics.cuda() # [fx, fy, cx, cy]\n\n if self.slam is None:\n cfg.merge_from_file(self.dpvo_cfg)\n cfg.BUFFER_SIZE = self.buffer\n self.slam = DPVO(cfg, self.dpvo_ckpt, ht=image.shape[1], wd=image.shape[2], viz=False)\n\n with Timer(\"SLAM\", enabled=False):\n t = time.time()\n self.slam(t, image, intrinsics)\n self.times.append(time.time() - t)\n\n return True\n\n def process(self):\n for _ in range(12):\n self.slam.update()\n\n self.reader.join()\n return self.slam.terminate()[0]\n\n\ndef video_stream(queue, imagedir, intrinsics, stride, skip=0, resize=0.5):\n \"\"\"video generator\"\"\"\n assert len(intrinsics) == 4, \"intrinsics should be [fx, fy, cx, cy]\"\n\n cap = cv2.VideoCapture(imagedir)\n t = 0\n for _ in range(skip):\n ret, image = cap.read()\n\n while True:\n # Capture frame-by-frame\n for _ in range(stride):\n ret, image = cap.read()\n # if frame is read correctly ret is True\n if not ret:\n break\n\n if not ret:\n break\n\n image = cv2.resize(image, None, fx=resize, fy=resize, interpolation=cv2.INTER_AREA)\n h, w, _ = image.shape\n image = image[: h - h % 16, : w - w % 16]\n\n intrinsics_ = intrinsics.clone() * resize\n queue.put((t, image, intrinsics_))\n\n t += 1\n\n queue.put((-1, image, intrinsics)) # -1 will terminate the process\n cap.release()\n\n # wait for the queue to be empty, otherwise the process will end immediately\n while not queue.empty():\n time.sleep(1)\n" }, { "path": "eval/GVHMR/hmr4d/utils/preproc/tracker.py", "content": "from ultralytics import YOLO\nfrom hmr4d import PROJ_ROOT\n\nimport torch\nimport numpy as np\nfrom tqdm import tqdm\nfrom collections import defaultdict\n\nfrom hmr4d.utils.seq_utils import (\n get_frame_id_list_from_mask,\n linear_interpolate_frame_ids,\n frame_id_to_mask,\n rearrange_by_mask,\n)\nfrom hmr4d.utils.video_io_utils import get_video_lwh\nfrom hmr4d.utils.net_utils import moving_average_smooth\n\n\nclass Tracker:\n def __init__(self) -> None:\n # https://docs.ultralytics.com/modes/predict/\n self.yolo = YOLO(PROJ_ROOT / \"inputs/checkpoints/yolo/yolov8x.pt\")\n\n def track(self, video_path):\n track_history = []\n cfg = {\n \"device\": \"cuda\",\n \"conf\": 0.5, # default 0.25, wham 0.5\n \"classes\": 0, # human\n \"verbose\": False,\n \"stream\": True,\n }\n results = self.yolo.track(video_path, **cfg)\n # frame-by-frame tracking\n track_history = []\n for result in tqdm(results, total=get_video_lwh(video_path)[0], desc=\"YoloV8 Tracking\"):\n if result.boxes.id is not None:\n track_ids = result.boxes.id.int().cpu().tolist() # (N)\n bbx_xyxy = result.boxes.xyxy.cpu().numpy() # (N, 4)\n result_frame = [{\"id\": track_ids[i], \"bbx_xyxy\": bbx_xyxy[i]} for i in range(len(track_ids))]\n else:\n result_frame = []\n track_history.append(result_frame)\n\n return track_history\n\n @staticmethod\n def sort_track_length(track_history, video_path):\n \"\"\"This handles the track history from YOLO tracker.\"\"\"\n id_to_frame_ids = defaultdict(list)\n id_to_bbx_xyxys = defaultdict(list)\n # parse to {det_id : [frame_id]}\n for frame_id, frame in enumerate(track_history):\n for det in frame:\n id_to_frame_ids[det[\"id\"]].append(frame_id)\n id_to_bbx_xyxys[det[\"id\"]].append(det[\"bbx_xyxy\"])\n for k, v in id_to_bbx_xyxys.items():\n id_to_bbx_xyxys[k] = np.array(v)\n\n # Sort by length of each track (max to min)\n id_length = {k: len(v) for k, v in id_to_frame_ids.items()}\n id2length = dict(sorted(id_length.items(), key=lambda item: item[1], reverse=True))\n\n # Sort by area sum (max to min)\n id_area_sum = {}\n l, w, h = get_video_lwh(video_path)\n for k, v in id_to_bbx_xyxys.items():\n bbx_wh = v[:, 2:] - v[:, :2]\n id_area_sum[k] = (bbx_wh[:, 0] * bbx_wh[:, 1] / w / h).sum()\n id2area_sum = dict(sorted(id_area_sum.items(), key=lambda item: item[1], reverse=True))\n id_sorted = list(id2area_sum.keys())\n\n return id_to_frame_ids, id_to_bbx_xyxys, id_sorted\n\n def get_one_track(self, video_path):\n # track\n track_history = self.track(video_path)\n\n # parse track_history & use top1 track\n id_to_frame_ids, id_to_bbx_xyxys, id_sorted = self.sort_track_length(track_history, video_path)\n track_id = id_sorted[0]\n frame_ids = torch.tensor(id_to_frame_ids[track_id]) # (N,)\n bbx_xyxys = torch.tensor(id_to_bbx_xyxys[track_id]) # (N, 4)\n\n # interpolate missing frames\n mask = frame_id_to_mask(frame_ids, get_video_lwh(video_path)[0])\n bbx_xyxy_one_track = rearrange_by_mask(bbx_xyxys, mask) # (F, 4), missing filled with 0\n missing_frame_id_list = get_frame_id_list_from_mask(~mask) # list of list\n bbx_xyxy_one_track = linear_interpolate_frame_ids(bbx_xyxy_one_track, missing_frame_id_list)\n assert (bbx_xyxy_one_track.sum(1) != 0).all()\n\n bbx_xyxy_one_track = moving_average_smooth(bbx_xyxy_one_track, window_size=5, dim=0)\n bbx_xyxy_one_track = moving_average_smooth(bbx_xyxy_one_track, window_size=5, dim=0)\n\n return bbx_xyxy_one_track\n" }, { "path": "eval/GVHMR/hmr4d/utils/preproc/vitfeat_extractor.py", "content": "import torch\nfrom hmr4d.network.hmr2 import load_hmr2, HMR2\n\n\nfrom hmr4d.utils.video_io_utils import read_video_np\nimport cv2\nimport numpy as np\n\nfrom hmr4d.network.hmr2.utils.preproc import crop_and_resize, IMAGE_MEAN, IMAGE_STD\nfrom tqdm import tqdm\n\n\ndef get_batch(input_path, bbx_xys, img_ds=0.5, img_dst_size=256, path_type=\"video\"):\n if path_type == \"video\":\n imgs = read_video_np(input_path, scale=img_ds)\n elif path_type == \"image\":\n imgs = cv2.imread(str(input_path))[..., ::-1]\n imgs = cv2.resize(imgs, (0, 0), fx=img_ds, fy=img_ds)\n imgs = imgs[None]\n elif path_type == \"np\":\n assert isinstance(input_path, np.ndarray)\n assert img_ds == 1.0 # this is safe\n imgs = input_path\n\n gt_center = bbx_xys[:, :2]\n gt_bbx_size = bbx_xys[:, 2]\n\n # Blur image to avoid aliasing artifacts\n if True:\n gt_bbx_size_ds = gt_bbx_size * img_ds\n ds_factors = ((gt_bbx_size_ds * 1.0) / img_dst_size / 2.0).numpy()\n imgs = np.stack(\n [\n # gaussian(v, sigma=(d - 1) / 2, channel_axis=2, preserve_range=True) if d > 1.1 else v\n cv2.GaussianBlur(v, (5, 5), (d - 1) / 2) if d > 1.1 else v\n for v, d in zip(imgs, ds_factors)\n ]\n )\n\n # Output\n imgs_list = []\n bbx_xys_ds_list = []\n for i in range(len(imgs)):\n img, bbx_xys_ds = crop_and_resize(\n imgs[i],\n gt_center[i] * img_ds,\n gt_bbx_size[i] * img_ds,\n img_dst_size,\n enlarge_ratio=1.0,\n )\n imgs_list.append(img)\n bbx_xys_ds_list.append(bbx_xys_ds)\n imgs = torch.from_numpy(np.stack(imgs_list)) # (F, 256, 256, 3), RGB\n bbx_xys = torch.from_numpy(np.stack(bbx_xys_ds_list)) / img_ds # (F, 3)\n\n imgs = ((imgs / 255.0 - IMAGE_MEAN) / IMAGE_STD).permute(0, 3, 1, 2) # (F, 3, 256, 256\n return imgs, bbx_xys\n\n\nclass Extractor:\n def __init__(self, tqdm_leave=True):\n self.extractor: HMR2 = load_hmr2().cuda().eval()\n self.tqdm_leave = tqdm_leave\n\n def extract_video_features(self, video_path, bbx_xys, img_ds=0.5):\n \"\"\"\n img_ds makes the image smaller, which is useful for faster processing\n \"\"\"\n # Get the batch\n if isinstance(video_path, str):\n imgs, bbx_xys = get_batch(video_path, bbx_xys, img_ds=img_ds)\n else:\n assert isinstance(video_path, torch.Tensor)\n imgs = video_path\n\n # Inference\n F, _, H, W = imgs.shape # (F, 3, H, W)\n imgs = imgs.cuda()\n batch_size = 16 # 5GB GPU memory, occupies all CUDA cores of 3090\n features = []\n for j in tqdm(range(0, F, batch_size), desc=\"HMR2 Feature\", leave=self.tqdm_leave):\n imgs_batch = imgs[j : j + batch_size]\n\n with torch.no_grad():\n feature = self.extractor({\"img\": imgs_batch})\n features.append(feature.detach().cpu())\n\n features = torch.cat(features, dim=0).clone() # (F, 1024)\n return features\n" }, { "path": "eval/GVHMR/hmr4d/utils/preproc/vitpose.py", "content": "import torch\nimport torch.nn.functional as F\nimport numpy as np\nfrom .vitpose_pytorch import build_model\nfrom .vitfeat_extractor import get_batch\nfrom tqdm import tqdm\n\nfrom hmr4d.utils.kpts.kp2d_utils import keypoints_from_heatmaps\nfrom hmr4d.utils.geo_transform import cvt_p2d_from_pm1_to_i\nfrom hmr4d.utils.geo.flip_utils import flip_heatmap_coco17\n\n\nclass VitPoseExtractor:\n def __init__(self, tqdm_leave=True):\n ckpt_path = \"inputs/checkpoints/vitpose/vitpose-h-multi-coco.pth\"\n self.pose = build_model(\"ViTPose_huge_coco_256x192\", ckpt_path)\n self.pose.cuda().eval()\n\n self.flip_test = True\n self.tqdm_leave = tqdm_leave\n\n @torch.no_grad()\n def extract(self, video_path, bbx_xys, img_ds=0.5):\n # Get the batch\n if isinstance(video_path, str):\n imgs, bbx_xys = get_batch(video_path, bbx_xys, img_ds=img_ds)\n else:\n assert isinstance(video_path, torch.Tensor)\n imgs = video_path\n\n # Inference\n L, _, H, W = imgs.shape # (L, 3, H, W)\n batch_size = 16\n vitpose = []\n for j in tqdm(range(0, L, batch_size), desc=\"ViTPose\", leave=self.tqdm_leave):\n # Heat map\n imgs_batch = imgs[j : j + batch_size, :, :, 32:224].cuda()\n if self.flip_test:\n heatmap, heatmap_flipped = self.pose(torch.cat([imgs_batch, imgs_batch.flip(3)], dim=0)).chunk(2)\n heatmap_flipped = flip_heatmap_coco17(heatmap_flipped)\n heatmap = (heatmap + heatmap_flipped) * 0.5\n del heatmap_flipped\n else:\n heatmap = self.pose(imgs_batch.clone()) # (B, J, 64, 48)\n\n if False:\n # Get joint\n bbx_xys_batch = bbx_xys[j : j + batch_size].cuda()\n method = \"hard\"\n if method == \"hard\":\n kp2d_pm1, conf = get_heatmap_preds(heatmap)\n elif method == \"soft\":\n kp2d_pm1, conf = get_heatmap_preds(heatmap, soft=True)\n\n # Convert 64, 48 to 64, 64\n kp2d_pm1[:, :, 0] *= 24 / 32\n kp2d = cvt_p2d_from_pm1_to_i(kp2d_pm1, bbx_xys_batch[:, None])\n kp2d = torch.cat([kp2d, conf], dim=-1)\n\n else: # postprocess from mmpose\n bbx_xys_batch = bbx_xys[j : j + batch_size]\n heatmap = heatmap.clone().cpu().numpy()\n center = bbx_xys_batch[:, :2].numpy()\n scale = (torch.cat((bbx_xys_batch[:, [2]] * 24 / 32, bbx_xys_batch[:, [2]]), dim=1) / 200).numpy()\n preds, maxvals = keypoints_from_heatmaps(heatmaps=heatmap, center=center, scale=scale, use_udp=True)\n kp2d = np.concatenate((preds, maxvals), axis=-1)\n kp2d = torch.from_numpy(kp2d)\n\n vitpose.append(kp2d.detach().cpu().clone())\n\n vitpose = torch.cat(vitpose, dim=0).clone() # (F, 17, 3)\n return vitpose\n\n\ndef get_heatmap_preds(heatmap, normalize_keypoints=True, thr=0.0, soft=False):\n \"\"\"\n heatmap: (B, J, H, W)\n \"\"\"\n assert heatmap.ndim == 4, \"batch_images should be 4-ndim\"\n\n B, J, H, W = heatmap.shape\n heatmaps_reshaped = heatmap.reshape((B, J, -1))\n\n maxvals, idx = torch.max(heatmaps_reshaped, 2)\n maxvals = maxvals.reshape((B, J, 1))\n idx = idx.reshape((B, J, 1))\n preds = idx.repeat(1, 1, 2).float()\n preds[:, :, 0] = (preds[:, :, 0]) % W\n preds[:, :, 1] = torch.floor((preds[:, :, 1]) / W)\n\n pred_mask = torch.gt(maxvals, thr).repeat(1, 1, 2)\n pred_mask = pred_mask.float()\n preds *= pred_mask\n\n # soft peak\n if soft:\n patch_size = 5\n patch_half = patch_size // 2\n patches = torch.zeros((B, J, patch_size, patch_size)).to(heatmap)\n default_patch = torch.zeros(patch_size, patch_size).to(heatmap)\n default_patch[patch_half, patch_half] = 1\n for b in range(B):\n for j in range(17):\n x, y = preds[b, j].int()\n if x >= patch_half and x <= W - patch_half and y >= patch_half and y <= H - patch_half:\n patches[b, j] = heatmap[\n b, j, y - patch_half : y + patch_half + 1, x - patch_half : x + patch_half + 1\n ]\n else:\n patches[b, j] = default_patch\n\n dx, dy = soft_patch_dx_dy(patches)\n preds[:, :, 0] += dx\n preds[:, :, 1] += dy\n\n if normalize_keypoints: # to [-1, 1]\n preds[:, :, 0] = preds[:, :, 0] / (W - 1) * 2 - 1\n preds[:, :, 1] = preds[:, :, 1] / (H - 1) * 2 - 1\n\n return preds, maxvals\n\n\ndef soft_patch_dx_dy(p):\n \"\"\"p (B,J,P,P)\"\"\"\n p_batch_shape = p.shape[:-2]\n patch_size = p.size(-1)\n temperature = 1.0\n score = F.softmax(p.view(-1, patch_size**2) * temperature, dim=-1)\n\n # get a offset_grid (BN, P, P, 2) for dx, dy\n offset_grid = torch.meshgrid(torch.arange(patch_size), torch.arange(patch_size))[::-1]\n offset_grid = torch.stack(offset_grid, dim=-1).float() - (patch_size - 1) / 2\n offset_grid = offset_grid.view(1, 1, patch_size, patch_size, 2).to(p.device)\n\n score = score.view(*p_batch_shape, patch_size, patch_size)\n dx = torch.sum(score * offset_grid[..., 0], dim=(-2, -1))\n dy = torch.sum(score * offset_grid[..., 1], dim=(-2, -1))\n\n if False:\n b, j = 0, 0\n print(torch.stack([dx[b, j], dy[b, j]]))\n print(p[b, j])\n\n return dx, dy\n" }, { "path": "eval/GVHMR/hmr4d/utils/preproc/vitpose_pytorch/__init__.py", "content": "from .src.vitpose_infer.model_builder import build_model\n" }, { "path": "eval/GVHMR/hmr4d/utils/preproc/vitpose_pytorch/src/vitpose_infer/__init__.py", "content": "" }, { "path": "eval/GVHMR/hmr4d/utils/preproc/vitpose_pytorch/src/vitpose_infer/builder/__init__.py", "content": "" }, { "path": "eval/GVHMR/hmr4d/utils/preproc/vitpose_pytorch/src/vitpose_infer/builder/backbones/__init__.py", "content": "# Copyright (c) OpenMMLab. All rights reserved.\n# from .alexnet import AlexNet\n# from .cpm import CPM\n# from .hourglass import HourglassNet\n# from .hourglass_ae import HourglassAENet\n# from .hrformer import HRFormer\n# from .hrnet import HRNet\n# from .litehrnet import LiteHRNet\n# from .mobilenet_v2 import MobileNetV2\n# from .mobilenet_v3 import MobileNetV3\n# from .mspn import MSPN\n# from .regnet import RegNet\n# from .resnest import ResNeSt\n# from .resnet import ResNet, ResNetV1d\n# from .resnext import ResNeXt\n# from .rsn import RSN\n# from .scnet import SCNet\n# from .seresnet import SEResNet\n# from .seresnext import SEResNeXt\n# from .shufflenet_v1 import ShuffleNetV1\n# from .shufflenet_v2 import ShuffleNetV2\n# from .tcn import TCN\n# from .v2v_net import V2VNet\n# from .vgg import VGG\n# from .vipnas_mbv3 import ViPNAS_MobileNetV3\n# from .vipnas_resnet import ViPNAS_ResNet\nfrom .vit import ViT\n\n# __all__ = [\n# 'AlexNet', 'HourglassNet', 'HourglassAENet', 'HRNet', 'MobileNetV2',\n# 'MobileNetV3', 'RegNet', 'ResNet', 'ResNetV1d', 'ResNeXt', 'SCNet',\n# 'SEResNet', 'SEResNeXt', 'ShuffleNetV1', 'ShuffleNetV2', 'CPM', 'RSN',\n# 'MSPN', 'ResNeSt', 'VGG', 'TCN', 'ViPNAS_ResNet', 'ViPNAS_MobileNetV3',\n# 'LiteHRNet', 'V2VNet', 'HRFormer', 'ViT'\n# ]\n" }, { "path": "eval/GVHMR/hmr4d/utils/preproc/vitpose_pytorch/src/vitpose_infer/builder/backbones/alexnet.py", "content": "# Copyright (c) OpenMMLab. All rights reserved.\nimport torch.nn as nn\n\nfrom ..builder import BACKBONES\nfrom .base_backbone import BaseBackbone\n\n\n@BACKBONES.register_module()\nclass AlexNet(BaseBackbone):\n \"\"\"`AlexNet `__ backbone.\n\n The input for AlexNet is a 224x224 RGB image.\n\n Args:\n num_classes (int): number of classes for classification.\n The default value is -1, which uses the backbone as\n a feature extractor without the top classifier.\n \"\"\"\n\n def __init__(self, num_classes=-1):\n super().__init__()\n self.num_classes = num_classes\n self.features = nn.Sequential(\n nn.Conv2d(3, 64, kernel_size=11, stride=4, padding=2),\n nn.ReLU(inplace=True),\n nn.MaxPool2d(kernel_size=3, stride=2),\n nn.Conv2d(64, 192, kernel_size=5, padding=2),\n nn.ReLU(inplace=True),\n nn.MaxPool2d(kernel_size=3, stride=2),\n nn.Conv2d(192, 384, kernel_size=3, padding=1),\n nn.ReLU(inplace=True),\n nn.Conv2d(384, 256, kernel_size=3, padding=1),\n nn.ReLU(inplace=True),\n nn.Conv2d(256, 256, kernel_size=3, padding=1),\n nn.ReLU(inplace=True),\n nn.MaxPool2d(kernel_size=3, stride=2),\n )\n if self.num_classes > 0:\n self.classifier = nn.Sequential(\n nn.Dropout(),\n nn.Linear(256 * 6 * 6, 4096),\n nn.ReLU(inplace=True),\n nn.Dropout(),\n nn.Linear(4096, 4096),\n nn.ReLU(inplace=True),\n nn.Linear(4096, num_classes),\n )\n\n def forward(self, x):\n\n x = self.features(x)\n if self.num_classes > 0:\n x = x.view(x.size(0), 256 * 6 * 6)\n x = self.classifier(x)\n\n return x\n" }, { "path": "eval/GVHMR/hmr4d/utils/preproc/vitpose_pytorch/src/vitpose_infer/builder/backbones/cpm.py", "content": "# Copyright (c) OpenMMLab. All rights reserved.\nimport copy\n\nimport torch\nimport torch.nn as nn\nfrom mmcv.cnn import ConvModule, constant_init, normal_init\nfrom torch.nn.modules.batchnorm import _BatchNorm\n\nfrom mmpose.utils import get_root_logger\nfrom ..builder import BACKBONES\nfrom .base_backbone import BaseBackbone\nfrom .utils import load_checkpoint\n\n\nclass CpmBlock(nn.Module):\n \"\"\"CpmBlock for Convolutional Pose Machine.\n\n Args:\n in_channels (int): Input channels of this block.\n channels (list): Output channels of each conv module.\n kernels (list): Kernel sizes of each conv module.\n \"\"\"\n\n def __init__(self,\n in_channels,\n channels=(128, 128, 128),\n kernels=(11, 11, 11),\n norm_cfg=None):\n super().__init__()\n\n assert len(channels) == len(kernels)\n layers = []\n for i in range(len(channels)):\n if i == 0:\n input_channels = in_channels\n else:\n input_channels = channels[i - 1]\n layers.append(\n ConvModule(\n input_channels,\n channels[i],\n kernels[i],\n padding=(kernels[i] - 1) // 2,\n norm_cfg=norm_cfg))\n self.model = nn.Sequential(*layers)\n\n def forward(self, x):\n \"\"\"Model forward function.\"\"\"\n out = self.model(x)\n return out\n\n\n@BACKBONES.register_module()\nclass CPM(BaseBackbone):\n \"\"\"CPM backbone.\n\n Convolutional Pose Machines.\n More details can be found in the `paper\n `__ .\n\n Args:\n in_channels (int): The input channels of the CPM.\n out_channels (int): The output channels of the CPM.\n feat_channels (int): Feature channel of each CPM stage.\n middle_channels (int): Feature channel of conv after the middle stage.\n num_stages (int): Number of stages.\n norm_cfg (dict): Dictionary to construct and config norm layer.\n\n Example:\n >>> from mmpose.models import CPM\n >>> import torch\n >>> self = CPM(3, 17)\n >>> self.eval()\n >>> inputs = torch.rand(1, 3, 368, 368)\n >>> level_outputs = self.forward(inputs)\n >>> for level_output in level_outputs:\n ... print(tuple(level_output.shape))\n (1, 17, 46, 46)\n (1, 17, 46, 46)\n (1, 17, 46, 46)\n (1, 17, 46, 46)\n (1, 17, 46, 46)\n (1, 17, 46, 46)\n \"\"\"\n\n def __init__(self,\n in_channels,\n out_channels,\n feat_channels=128,\n middle_channels=32,\n num_stages=6,\n norm_cfg=dict(type='BN', requires_grad=True)):\n # Protect mutable default arguments\n norm_cfg = copy.deepcopy(norm_cfg)\n super().__init__()\n\n assert in_channels == 3\n\n self.num_stages = num_stages\n assert self.num_stages >= 1\n\n self.stem = nn.Sequential(\n ConvModule(in_channels, 128, 9, padding=4, norm_cfg=norm_cfg),\n nn.MaxPool2d(kernel_size=3, stride=2, padding=1),\n ConvModule(128, 128, 9, padding=4, norm_cfg=norm_cfg),\n nn.MaxPool2d(kernel_size=3, stride=2, padding=1),\n ConvModule(128, 128, 9, padding=4, norm_cfg=norm_cfg),\n nn.MaxPool2d(kernel_size=3, stride=2, padding=1),\n ConvModule(128, 32, 5, padding=2, norm_cfg=norm_cfg),\n ConvModule(32, 512, 9, padding=4, norm_cfg=norm_cfg),\n ConvModule(512, 512, 1, padding=0, norm_cfg=norm_cfg),\n ConvModule(512, out_channels, 1, padding=0, act_cfg=None))\n\n self.middle = nn.Sequential(\n ConvModule(in_channels, 128, 9, padding=4, norm_cfg=norm_cfg),\n nn.MaxPool2d(kernel_size=3, stride=2, padding=1),\n ConvModule(128, 128, 9, padding=4, norm_cfg=norm_cfg),\n nn.MaxPool2d(kernel_size=3, stride=2, padding=1),\n ConvModule(128, 128, 9, padding=4, norm_cfg=norm_cfg),\n nn.MaxPool2d(kernel_size=3, stride=2, padding=1))\n\n self.cpm_stages = nn.ModuleList([\n CpmBlock(\n middle_channels + out_channels,\n channels=[feat_channels, feat_channels, feat_channels],\n kernels=[11, 11, 11],\n norm_cfg=norm_cfg) for _ in range(num_stages - 1)\n ])\n\n self.middle_conv = nn.ModuleList([\n nn.Sequential(\n ConvModule(\n 128, middle_channels, 5, padding=2, norm_cfg=norm_cfg))\n for _ in range(num_stages - 1)\n ])\n\n self.out_convs = nn.ModuleList([\n nn.Sequential(\n ConvModule(\n feat_channels,\n feat_channels,\n 1,\n padding=0,\n norm_cfg=norm_cfg),\n ConvModule(feat_channels, out_channels, 1, act_cfg=None))\n for _ in range(num_stages - 1)\n ])\n\n def init_weights(self, pretrained=None):\n \"\"\"Initialize the weights in backbone.\n\n Args:\n pretrained (str, optional): Path to pre-trained weights.\n Defaults to None.\n \"\"\"\n if isinstance(pretrained, str):\n logger = get_root_logger()\n load_checkpoint(self, pretrained, strict=False, logger=logger)\n elif pretrained is None:\n for m in self.modules():\n if isinstance(m, nn.Conv2d):\n normal_init(m, std=0.001)\n elif isinstance(m, (_BatchNorm, nn.GroupNorm)):\n constant_init(m, 1)\n else:\n raise TypeError('pretrained must be a str or None')\n\n def forward(self, x):\n \"\"\"Model forward function.\"\"\"\n stage1_out = self.stem(x)\n middle_out = self.middle(x)\n out_feats = []\n\n out_feats.append(stage1_out)\n\n for ind in range(self.num_stages - 1):\n single_stage = self.cpm_stages[ind]\n out_conv = self.out_convs[ind]\n\n inp_feat = torch.cat(\n [out_feats[-1], self.middle_conv[ind](middle_out)], 1)\n cpm_feat = single_stage(inp_feat)\n out_feat = out_conv(cpm_feat)\n out_feats.append(out_feat)\n\n return out_feats\n" }, { "path": "eval/GVHMR/hmr4d/utils/preproc/vitpose_pytorch/src/vitpose_infer/builder/backbones/hourglass.py", "content": "# Copyright (c) OpenMMLab. All rights reserved.\nimport copy\n\nimport torch.nn as nn\nfrom mmcv.cnn import ConvModule, constant_init, normal_init\nfrom torch.nn.modules.batchnorm import _BatchNorm\n\nfrom mmpose.utils import get_root_logger\nfrom ..builder import BACKBONES\nfrom .base_backbone import BaseBackbone\nfrom .resnet import BasicBlock, ResLayer\nfrom .utils import load_checkpoint\n\n\nclass HourglassModule(nn.Module):\n \"\"\"Hourglass Module for HourglassNet backbone.\n\n Generate module recursively and use BasicBlock as the base unit.\n\n Args:\n depth (int): Depth of current HourglassModule.\n stage_channels (list[int]): Feature channels of sub-modules in current\n and follow-up HourglassModule.\n stage_blocks (list[int]): Number of sub-modules stacked in current and\n follow-up HourglassModule.\n norm_cfg (dict): Dictionary to construct and config norm layer.\n \"\"\"\n\n def __init__(self,\n depth,\n stage_channels,\n stage_blocks,\n norm_cfg=dict(type='BN', requires_grad=True)):\n # Protect mutable default arguments\n norm_cfg = copy.deepcopy(norm_cfg)\n super().__init__()\n\n self.depth = depth\n\n cur_block = stage_blocks[0]\n next_block = stage_blocks[1]\n\n cur_channel = stage_channels[0]\n next_channel = stage_channels[1]\n\n self.up1 = ResLayer(\n BasicBlock, cur_block, cur_channel, cur_channel, norm_cfg=norm_cfg)\n\n self.low1 = ResLayer(\n BasicBlock,\n cur_block,\n cur_channel,\n next_channel,\n stride=2,\n norm_cfg=norm_cfg)\n\n if self.depth > 1:\n self.low2 = HourglassModule(depth - 1, stage_channels[1:],\n stage_blocks[1:])\n else:\n self.low2 = ResLayer(\n BasicBlock,\n next_block,\n next_channel,\n next_channel,\n norm_cfg=norm_cfg)\n\n self.low3 = ResLayer(\n BasicBlock,\n cur_block,\n next_channel,\n cur_channel,\n norm_cfg=norm_cfg,\n downsample_first=False)\n\n self.up2 = nn.Upsample(scale_factor=2)\n\n def forward(self, x):\n \"\"\"Model forward function.\"\"\"\n up1 = self.up1(x)\n low1 = self.low1(x)\n low2 = self.low2(low1)\n low3 = self.low3(low2)\n up2 = self.up2(low3)\n return up1 + up2\n\n\n@BACKBONES.register_module()\nclass HourglassNet(BaseBackbone):\n \"\"\"HourglassNet backbone.\n\n Stacked Hourglass Networks for Human Pose Estimation.\n More details can be found in the `paper\n `__ .\n\n Args:\n downsample_times (int): Downsample times in a HourglassModule.\n num_stacks (int): Number of HourglassModule modules stacked,\n 1 for Hourglass-52, 2 for Hourglass-104.\n stage_channels (list[int]): Feature channel of each sub-module in a\n HourglassModule.\n stage_blocks (list[int]): Number of sub-modules stacked in a\n HourglassModule.\n feat_channel (int): Feature channel of conv after a HourglassModule.\n norm_cfg (dict): Dictionary to construct and config norm layer.\n\n Example:\n >>> from mmpose.models import HourglassNet\n >>> import torch\n >>> self = HourglassNet()\n >>> self.eval()\n >>> inputs = torch.rand(1, 3, 511, 511)\n >>> level_outputs = self.forward(inputs)\n >>> for level_output in level_outputs:\n ... print(tuple(level_output.shape))\n (1, 256, 128, 128)\n (1, 256, 128, 128)\n \"\"\"\n\n def __init__(self,\n downsample_times=5,\n num_stacks=2,\n stage_channels=(256, 256, 384, 384, 384, 512),\n stage_blocks=(2, 2, 2, 2, 2, 4),\n feat_channel=256,\n norm_cfg=dict(type='BN', requires_grad=True)):\n # Protect mutable default arguments\n norm_cfg = copy.deepcopy(norm_cfg)\n super().__init__()\n\n self.num_stacks = num_stacks\n assert self.num_stacks >= 1\n assert len(stage_channels) == len(stage_blocks)\n assert len(stage_channels) > downsample_times\n\n cur_channel = stage_channels[0]\n\n self.stem = nn.Sequential(\n ConvModule(3, 128, 7, padding=3, stride=2, norm_cfg=norm_cfg),\n ResLayer(BasicBlock, 1, 128, 256, stride=2, norm_cfg=norm_cfg))\n\n self.hourglass_modules = nn.ModuleList([\n HourglassModule(downsample_times, stage_channels, stage_blocks)\n for _ in range(num_stacks)\n ])\n\n self.inters = ResLayer(\n BasicBlock,\n num_stacks - 1,\n cur_channel,\n cur_channel,\n norm_cfg=norm_cfg)\n\n self.conv1x1s = nn.ModuleList([\n ConvModule(\n cur_channel, cur_channel, 1, norm_cfg=norm_cfg, act_cfg=None)\n for _ in range(num_stacks - 1)\n ])\n\n self.out_convs = nn.ModuleList([\n ConvModule(\n cur_channel, feat_channel, 3, padding=1, norm_cfg=norm_cfg)\n for _ in range(num_stacks)\n ])\n\n self.remap_convs = nn.ModuleList([\n ConvModule(\n feat_channel, cur_channel, 1, norm_cfg=norm_cfg, act_cfg=None)\n for _ in range(num_stacks - 1)\n ])\n\n self.relu = nn.ReLU(inplace=True)\n\n def init_weights(self, pretrained=None):\n \"\"\"Initialize the weights in backbone.\n\n Args:\n pretrained (str, optional): Path to pre-trained weights.\n Defaults to None.\n \"\"\"\n if isinstance(pretrained, str):\n logger = get_root_logger()\n load_checkpoint(self, pretrained, strict=False, logger=logger)\n elif pretrained is None:\n for m in self.modules():\n if isinstance(m, nn.Conv2d):\n normal_init(m, std=0.001)\n elif isinstance(m, (_BatchNorm, nn.GroupNorm)):\n constant_init(m, 1)\n else:\n raise TypeError('pretrained must be a str or None')\n\n def forward(self, x):\n \"\"\"Model forward function.\"\"\"\n inter_feat = self.stem(x)\n out_feats = []\n\n for ind in range(self.num_stacks):\n single_hourglass = self.hourglass_modules[ind]\n out_conv = self.out_convs[ind]\n\n hourglass_feat = single_hourglass(inter_feat)\n out_feat = out_conv(hourglass_feat)\n out_feats.append(out_feat)\n\n if ind < self.num_stacks - 1:\n inter_feat = self.conv1x1s[ind](\n inter_feat) + self.remap_convs[ind](\n out_feat)\n inter_feat = self.inters[ind](self.relu(inter_feat))\n\n return out_feats\n" }, { "path": "eval/GVHMR/hmr4d/utils/preproc/vitpose_pytorch/src/vitpose_infer/builder/backbones/hourglass_ae.py", "content": "# Copyright (c) OpenMMLab. All rights reserved.\nimport copy\n\nimport torch.nn as nn\nfrom mmcv.cnn import ConvModule, MaxPool2d, constant_init, normal_init\nfrom torch.nn.modules.batchnorm import _BatchNorm\n\nfrom mmpose.utils import get_root_logger\nfrom ..builder import BACKBONES\nfrom .base_backbone import BaseBackbone\nfrom .utils import load_checkpoint\n\n\nclass HourglassAEModule(nn.Module):\n \"\"\"Modified Hourglass Module for HourglassNet_AE backbone.\n\n Generate module recursively and use BasicBlock as the base unit.\n\n Args:\n depth (int): Depth of current HourglassModule.\n stage_channels (list[int]): Feature channels of sub-modules in current\n and follow-up HourglassModule.\n norm_cfg (dict): Dictionary to construct and config norm layer.\n \"\"\"\n\n def __init__(self,\n depth,\n stage_channels,\n norm_cfg=dict(type='BN', requires_grad=True)):\n # Protect mutable default arguments\n norm_cfg = copy.deepcopy(norm_cfg)\n super().__init__()\n\n self.depth = depth\n\n cur_channel = stage_channels[0]\n next_channel = stage_channels[1]\n\n self.up1 = ConvModule(\n cur_channel, cur_channel, 3, padding=1, norm_cfg=norm_cfg)\n\n self.pool1 = MaxPool2d(2, 2)\n\n self.low1 = ConvModule(\n cur_channel, next_channel, 3, padding=1, norm_cfg=norm_cfg)\n\n if self.depth > 1:\n self.low2 = HourglassAEModule(depth - 1, stage_channels[1:])\n else:\n self.low2 = ConvModule(\n next_channel, next_channel, 3, padding=1, norm_cfg=norm_cfg)\n\n self.low3 = ConvModule(\n next_channel, cur_channel, 3, padding=1, norm_cfg=norm_cfg)\n\n self.up2 = nn.UpsamplingNearest2d(scale_factor=2)\n\n def forward(self, x):\n \"\"\"Model forward function.\"\"\"\n up1 = self.up1(x)\n pool1 = self.pool1(x)\n low1 = self.low1(pool1)\n low2 = self.low2(low1)\n low3 = self.low3(low2)\n up2 = self.up2(low3)\n return up1 + up2\n\n\n@BACKBONES.register_module()\nclass HourglassAENet(BaseBackbone):\n \"\"\"Hourglass-AE Network proposed by Newell et al.\n\n Associative Embedding: End-to-End Learning for Joint\n Detection and Grouping.\n\n More details can be found in the `paper\n `__ .\n\n Args:\n downsample_times (int): Downsample times in a HourglassModule.\n num_stacks (int): Number of HourglassModule modules stacked,\n 1 for Hourglass-52, 2 for Hourglass-104.\n stage_channels (list[int]): Feature channel of each sub-module in a\n HourglassModule.\n stage_blocks (list[int]): Number of sub-modules stacked in a\n HourglassModule.\n feat_channels (int): Feature channel of conv after a HourglassModule.\n norm_cfg (dict): Dictionary to construct and config norm layer.\n\n Example:\n >>> from mmpose.models import HourglassAENet\n >>> import torch\n >>> self = HourglassAENet()\n >>> self.eval()\n >>> inputs = torch.rand(1, 3, 512, 512)\n >>> level_outputs = self.forward(inputs)\n >>> for level_output in level_outputs:\n ... print(tuple(level_output.shape))\n (1, 34, 128, 128)\n \"\"\"\n\n def __init__(self,\n downsample_times=4,\n num_stacks=1,\n out_channels=34,\n stage_channels=(256, 384, 512, 640, 768),\n feat_channels=256,\n norm_cfg=dict(type='BN', requires_grad=True)):\n # Protect mutable default arguments\n norm_cfg = copy.deepcopy(norm_cfg)\n super().__init__()\n\n self.num_stacks = num_stacks\n assert self.num_stacks >= 1\n assert len(stage_channels) > downsample_times\n\n cur_channels = stage_channels[0]\n\n self.stem = nn.Sequential(\n ConvModule(3, 64, 7, padding=3, stride=2, norm_cfg=norm_cfg),\n ConvModule(64, 128, 3, padding=1, norm_cfg=norm_cfg),\n MaxPool2d(2, 2),\n ConvModule(128, 128, 3, padding=1, norm_cfg=norm_cfg),\n ConvModule(128, feat_channels, 3, padding=1, norm_cfg=norm_cfg),\n )\n\n self.hourglass_modules = nn.ModuleList([\n nn.Sequential(\n HourglassAEModule(\n downsample_times, stage_channels, norm_cfg=norm_cfg),\n ConvModule(\n feat_channels,\n feat_channels,\n 3,\n padding=1,\n norm_cfg=norm_cfg),\n ConvModule(\n feat_channels,\n feat_channels,\n 3,\n padding=1,\n norm_cfg=norm_cfg)) for _ in range(num_stacks)\n ])\n\n self.out_convs = nn.ModuleList([\n ConvModule(\n cur_channels,\n out_channels,\n 1,\n padding=0,\n norm_cfg=None,\n act_cfg=None) for _ in range(num_stacks)\n ])\n\n self.remap_out_convs = nn.ModuleList([\n ConvModule(\n out_channels,\n feat_channels,\n 1,\n norm_cfg=norm_cfg,\n act_cfg=None) for _ in range(num_stacks - 1)\n ])\n\n self.remap_feature_convs = nn.ModuleList([\n ConvModule(\n feat_channels,\n feat_channels,\n 1,\n norm_cfg=norm_cfg,\n act_cfg=None) for _ in range(num_stacks - 1)\n ])\n\n self.relu = nn.ReLU(inplace=True)\n\n def init_weights(self, pretrained=None):\n \"\"\"Initialize the weights in backbone.\n\n Args:\n pretrained (str, optional): Path to pre-trained weights.\n Defaults to None.\n \"\"\"\n if isinstance(pretrained, str):\n logger = get_root_logger()\n load_checkpoint(self, pretrained, strict=False, logger=logger)\n elif pretrained is None:\n for m in self.modules():\n if isinstance(m, nn.Conv2d):\n normal_init(m, std=0.001)\n elif isinstance(m, (_BatchNorm, nn.GroupNorm)):\n constant_init(m, 1)\n else:\n raise TypeError('pretrained must be a str or None')\n\n def forward(self, x):\n \"\"\"Model forward function.\"\"\"\n inter_feat = self.stem(x)\n out_feats = []\n\n for ind in range(self.num_stacks):\n single_hourglass = self.hourglass_modules[ind]\n out_conv = self.out_convs[ind]\n\n hourglass_feat = single_hourglass(inter_feat)\n out_feat = out_conv(hourglass_feat)\n out_feats.append(out_feat)\n\n if ind < self.num_stacks - 1:\n inter_feat = inter_feat + self.remap_out_convs[ind](\n out_feat) + self.remap_feature_convs[ind](\n hourglass_feat)\n\n return out_feats\n" }, { "path": "eval/GVHMR/hmr4d/utils/preproc/vitpose_pytorch/src/vitpose_infer/builder/backbones/hrformer.py", "content": "# Copyright (c) OpenMMLab. All rights reserved.\n\nimport math\n\nimport torch\nimport torch.nn as nn\n# from timm.models.layers import to_2tuple, trunc_normal_\nfrom mmcv.cnn import (build_activation_layer, build_conv_layer,\n build_norm_layer, trunc_normal_init)\nfrom mmcv.cnn.bricks.transformer import build_dropout\nfrom mmcv.runner import BaseModule\nfrom torch.nn.functional import pad\n\nfrom ..builder import BACKBONES\nfrom .hrnet import Bottleneck, HRModule, HRNet\n\n\ndef nlc_to_nchw(x, hw_shape):\n \"\"\"Convert [N, L, C] shape tensor to [N, C, H, W] shape tensor.\n\n Args:\n x (Tensor): The input tensor of shape [N, L, C] before conversion.\n hw_shape (Sequence[int]): The height and width of output feature map.\n\n Returns:\n Tensor: The output tensor of shape [N, C, H, W] after conversion.\n \"\"\"\n H, W = hw_shape\n assert len(x.shape) == 3\n B, L, C = x.shape\n assert L == H * W, 'The seq_len doesn\\'t match H, W'\n return x.transpose(1, 2).reshape(B, C, H, W)\n\n\ndef nchw_to_nlc(x):\n \"\"\"Flatten [N, C, H, W] shape tensor to [N, L, C] shape tensor.\n\n Args:\n x (Tensor): The input tensor of shape [N, C, H, W] before conversion.\n\n Returns:\n Tensor: The output tensor of shape [N, L, C] after conversion.\n \"\"\"\n assert len(x.shape) == 4\n return x.flatten(2).transpose(1, 2).contiguous()\n\n\ndef build_drop_path(drop_path_rate):\n \"\"\"Build drop path layer.\"\"\"\n return build_dropout(dict(type='DropPath', drop_prob=drop_path_rate))\n\n\nclass WindowMSA(BaseModule):\n \"\"\"Window based multi-head self-attention (W-MSA) module with relative\n position bias.\n\n Args:\n embed_dims (int): Number of input channels.\n num_heads (int): Number of attention heads.\n window_size (tuple[int]): The height and width of the window.\n qkv_bias (bool, optional): If True, add a learnable bias to q, k, v.\n Default: True.\n qk_scale (float | None, optional): Override default qk scale of\n head_dim ** -0.5 if set. Default: None.\n attn_drop_rate (float, optional): Dropout ratio of attention weight.\n Default: 0.0\n proj_drop_rate (float, optional): Dropout ratio of output. Default: 0.\n with_rpe (bool, optional): If True, use relative position bias.\n Default: True.\n init_cfg (dict | None, optional): The Config for initialization.\n Default: None.\n \"\"\"\n\n def __init__(self,\n embed_dims,\n num_heads,\n window_size,\n qkv_bias=True,\n qk_scale=None,\n attn_drop_rate=0.,\n proj_drop_rate=0.,\n with_rpe=True,\n init_cfg=None):\n\n super().__init__(init_cfg=init_cfg)\n self.embed_dims = embed_dims\n self.window_size = window_size # Wh, Ww\n self.num_heads = num_heads\n head_embed_dims = embed_dims // num_heads\n self.scale = qk_scale or head_embed_dims**-0.5\n\n self.with_rpe = with_rpe\n if self.with_rpe:\n # define a parameter table of relative position bias\n self.relative_position_bias_table = nn.Parameter(\n torch.zeros(\n (2 * window_size[0] - 1) * (2 * window_size[1] - 1),\n num_heads)) # 2*Wh-1 * 2*Ww-1, nH\n\n Wh, Ww = self.window_size\n rel_index_coords = self.double_step_seq(2 * Ww - 1, Wh, 1, Ww)\n rel_position_index = rel_index_coords + rel_index_coords.T\n rel_position_index = rel_position_index.flip(1).contiguous()\n self.register_buffer('relative_position_index', rel_position_index)\n\n self.qkv = nn.Linear(embed_dims, embed_dims * 3, bias=qkv_bias)\n self.attn_drop = nn.Dropout(attn_drop_rate)\n self.proj = nn.Linear(embed_dims, embed_dims)\n self.proj_drop = nn.Dropout(proj_drop_rate)\n\n self.softmax = nn.Softmax(dim=-1)\n\n def init_weights(self):\n trunc_normal_init(self.relative_position_bias_table, std=0.02)\n\n def forward(self, x, mask=None):\n \"\"\"\n Args:\n\n x (tensor): input features with shape of (B*num_windows, N, C)\n mask (tensor | None, Optional): mask with shape of (num_windows,\n Wh*Ww, Wh*Ww), value should be between (-inf, 0].\n \"\"\"\n B, N, C = x.shape\n qkv = self.qkv(x).reshape(B, N, 3, self.num_heads,\n C // self.num_heads).permute(2, 0, 3, 1, 4)\n q, k, v = qkv[0], qkv[1], qkv[2]\n\n q = q * self.scale\n attn = (q @ k.transpose(-2, -1))\n\n if self.with_rpe:\n relative_position_bias = self.relative_position_bias_table[\n self.relative_position_index.view(-1)].view(\n self.window_size[0] * self.window_size[1],\n self.window_size[0] * self.window_size[1],\n -1) # Wh*Ww,Wh*Ww,nH\n relative_position_bias = relative_position_bias.permute(\n 2, 0, 1).contiguous() # nH, Wh*Ww, Wh*Ww\n attn = attn + relative_position_bias.unsqueeze(0)\n\n if mask is not None:\n nW = mask.shape[0]\n attn = attn.view(B // nW, nW, self.num_heads, N,\n N) + mask.unsqueeze(1).unsqueeze(0)\n attn = attn.view(-1, self.num_heads, N, N)\n attn = self.softmax(attn)\n\n attn = self.attn_drop(attn)\n\n x = (attn @ v).transpose(1, 2).reshape(B, N, C)\n x = self.proj(x)\n x = self.proj_drop(x)\n return x\n\n @staticmethod\n def double_step_seq(step1, len1, step2, len2):\n seq1 = torch.arange(0, step1 * len1, step1)\n seq2 = torch.arange(0, step2 * len2, step2)\n return (seq1[:, None] + seq2[None, :]).reshape(1, -1)\n\n\nclass LocalWindowSelfAttention(BaseModule):\n r\"\"\" Local-window Self Attention (LSA) module with relative position bias.\n\n This module is the short-range self-attention module in the\n Interlaced Sparse Self-Attention `_.\n\n Args:\n embed_dims (int): Number of input channels.\n num_heads (int): Number of attention heads.\n window_size (tuple[int] | int): The height and width of the window.\n qkv_bias (bool, optional): If True, add a learnable bias to q, k, v.\n Default: True.\n qk_scale (float | None, optional): Override default qk scale of\n head_dim ** -0.5 if set. Default: None.\n attn_drop_rate (float, optional): Dropout ratio of attention weight.\n Default: 0.0\n proj_drop_rate (float, optional): Dropout ratio of output. Default: 0.\n with_rpe (bool, optional): If True, use relative position bias.\n Default: True.\n with_pad_mask (bool, optional): If True, mask out the padded tokens in\n the attention process. Default: False.\n init_cfg (dict | None, optional): The Config for initialization.\n Default: None.\n \"\"\"\n\n def __init__(self,\n embed_dims,\n num_heads,\n window_size,\n qkv_bias=True,\n qk_scale=None,\n attn_drop_rate=0.,\n proj_drop_rate=0.,\n with_rpe=True,\n with_pad_mask=False,\n init_cfg=None):\n super().__init__(init_cfg=init_cfg)\n if isinstance(window_size, int):\n window_size = (window_size, window_size)\n self.window_size = window_size\n self.with_pad_mask = with_pad_mask\n self.attn = WindowMSA(\n embed_dims=embed_dims,\n num_heads=num_heads,\n window_size=window_size,\n qkv_bias=qkv_bias,\n qk_scale=qk_scale,\n attn_drop_rate=attn_drop_rate,\n proj_drop_rate=proj_drop_rate,\n with_rpe=with_rpe,\n init_cfg=init_cfg)\n\n def forward(self, x, H, W, **kwargs):\n \"\"\"Forward function.\"\"\"\n B, N, C = x.shape\n x = x.view(B, H, W, C)\n Wh, Ww = self.window_size\n\n # center-pad the feature on H and W axes\n pad_h = math.ceil(H / Wh) * Wh - H\n pad_w = math.ceil(W / Ww) * Ww - W\n x = pad(x, (0, 0, pad_w // 2, pad_w - pad_w // 2, pad_h // 2,\n pad_h - pad_h // 2))\n\n # permute\n x = x.view(B, math.ceil(H / Wh), Wh, math.ceil(W / Ww), Ww, C)\n x = x.permute(0, 1, 3, 2, 4, 5)\n x = x.reshape(-1, Wh * Ww, C) # (B*num_window, Wh*Ww, C)\n\n # attention\n if self.with_pad_mask and pad_h > 0 and pad_w > 0:\n pad_mask = x.new_zeros(1, H, W, 1)\n pad_mask = pad(\n pad_mask, [\n 0, 0, pad_w // 2, pad_w - pad_w // 2, pad_h // 2,\n pad_h - pad_h // 2\n ],\n value=-float('inf'))\n pad_mask = pad_mask.view(1, math.ceil(H / Wh), Wh,\n math.ceil(W / Ww), Ww, 1)\n pad_mask = pad_mask.permute(1, 3, 0, 2, 4, 5)\n pad_mask = pad_mask.reshape(-1, Wh * Ww)\n pad_mask = pad_mask[:, None, :].expand([-1, Wh * Ww, -1])\n out = self.attn(x, pad_mask, **kwargs)\n else:\n out = self.attn(x, **kwargs)\n\n # reverse permutation\n out = out.reshape(B, math.ceil(H / Wh), math.ceil(W / Ww), Wh, Ww, C)\n out = out.permute(0, 1, 3, 2, 4, 5)\n out = out.reshape(B, H + pad_h, W + pad_w, C)\n\n # de-pad\n out = out[:, pad_h // 2:H + pad_h // 2, pad_w // 2:W + pad_w // 2]\n return out.reshape(B, N, C)\n\n\nclass CrossFFN(BaseModule):\n r\"\"\"FFN with Depthwise Conv of HRFormer.\n\n Args:\n in_features (int): The feature dimension.\n hidden_features (int, optional): The hidden dimension of FFNs.\n Defaults: The same as in_features.\n act_cfg (dict, optional): Config of activation layer.\n Default: dict(type='GELU').\n dw_act_cfg (dict, optional): Config of activation layer appended\n right after DW Conv. Default: dict(type='GELU').\n norm_cfg (dict, optional): Config of norm layer.\n Default: dict(type='SyncBN').\n init_cfg (dict | list | None, optional): The init config.\n Default: None.\n \"\"\"\n\n def __init__(self,\n in_features,\n hidden_features=None,\n out_features=None,\n act_cfg=dict(type='GELU'),\n dw_act_cfg=dict(type='GELU'),\n norm_cfg=dict(type='SyncBN'),\n init_cfg=None):\n super().__init__(init_cfg=init_cfg)\n out_features = out_features or in_features\n hidden_features = hidden_features or in_features\n self.fc1 = nn.Conv2d(in_features, hidden_features, kernel_size=1)\n self.act1 = build_activation_layer(act_cfg)\n self.norm1 = build_norm_layer(norm_cfg, hidden_features)[1]\n self.dw3x3 = nn.Conv2d(\n hidden_features,\n hidden_features,\n kernel_size=3,\n stride=1,\n groups=hidden_features,\n padding=1)\n self.act2 = build_activation_layer(dw_act_cfg)\n self.norm2 = build_norm_layer(norm_cfg, hidden_features)[1]\n self.fc2 = nn.Conv2d(hidden_features, out_features, kernel_size=1)\n self.act3 = build_activation_layer(act_cfg)\n self.norm3 = build_norm_layer(norm_cfg, out_features)[1]\n\n # put the modules togather\n self.layers = [\n self.fc1, self.norm1, self.act1, self.dw3x3, self.norm2, self.act2,\n self.fc2, self.norm3, self.act3\n ]\n\n def forward(self, x, H, W):\n \"\"\"Forward function.\"\"\"\n x = nlc_to_nchw(x, (H, W))\n for layer in self.layers:\n x = layer(x)\n x = nchw_to_nlc(x)\n return x\n\n\nclass HRFormerBlock(BaseModule):\n \"\"\"High-Resolution Block for HRFormer.\n\n Args:\n in_features (int): The input dimension.\n out_features (int): The output dimension.\n num_heads (int): The number of head within each LSA.\n window_size (int, optional): The window size for the LSA.\n Default: 7\n mlp_ratio (int, optional): The expansion ration of FFN.\n Default: 4\n act_cfg (dict, optional): Config of activation layer.\n Default: dict(type='GELU').\n norm_cfg (dict, optional): Config of norm layer.\n Default: dict(type='SyncBN').\n transformer_norm_cfg (dict, optional): Config of transformer norm\n layer. Default: dict(type='LN', eps=1e-6).\n init_cfg (dict | list | None, optional): The init config.\n Default: None.\n \"\"\"\n\n expansion = 1\n\n def __init__(self,\n in_features,\n out_features,\n num_heads,\n window_size=7,\n mlp_ratio=4.0,\n drop_path=0.0,\n act_cfg=dict(type='GELU'),\n norm_cfg=dict(type='SyncBN'),\n transformer_norm_cfg=dict(type='LN', eps=1e-6),\n init_cfg=None,\n **kwargs):\n super(HRFormerBlock, self).__init__(init_cfg=init_cfg)\n self.num_heads = num_heads\n self.window_size = window_size\n self.mlp_ratio = mlp_ratio\n\n self.norm1 = build_norm_layer(transformer_norm_cfg, in_features)[1]\n self.attn = LocalWindowSelfAttention(\n in_features,\n num_heads=num_heads,\n window_size=window_size,\n init_cfg=None,\n **kwargs)\n\n self.norm2 = build_norm_layer(transformer_norm_cfg, out_features)[1]\n self.ffn = CrossFFN(\n in_features=in_features,\n hidden_features=int(in_features * mlp_ratio),\n out_features=out_features,\n norm_cfg=norm_cfg,\n act_cfg=act_cfg,\n dw_act_cfg=act_cfg,\n init_cfg=None)\n\n self.drop_path = build_drop_path(\n drop_path) if drop_path > 0.0 else nn.Identity()\n\n def forward(self, x):\n \"\"\"Forward function.\"\"\"\n B, C, H, W = x.size()\n # Attention\n x = x.view(B, C, -1).permute(0, 2, 1)\n x = x + self.drop_path(self.attn(self.norm1(x), H, W))\n # FFN\n x = x + self.drop_path(self.ffn(self.norm2(x), H, W))\n x = x.permute(0, 2, 1).view(B, C, H, W)\n return x\n\n def extra_repr(self):\n \"\"\"(Optional) Set the extra information about this module.\"\"\"\n return 'num_heads={}, window_size={}, mlp_ratio={}'.format(\n self.num_heads, self.window_size, self.mlp_ratio)\n\n\nclass HRFomerModule(HRModule):\n \"\"\"High-Resolution Module for HRFormer.\n\n Args:\n num_branches (int): The number of branches in the HRFormerModule.\n block (nn.Module): The building block of HRFormer.\n The block should be the HRFormerBlock.\n num_blocks (tuple): The number of blocks in each branch.\n The length must be equal to num_branches.\n num_inchannels (tuple): The number of input channels in each branch.\n The length must be equal to num_branches.\n num_channels (tuple): The number of channels in each branch.\n The length must be equal to num_branches.\n num_heads (tuple): The number of heads within the LSAs.\n num_window_sizes (tuple): The window size for the LSAs.\n num_mlp_ratios (tuple): The expansion ratio for the FFNs.\n drop_path (int, optional): The drop path rate of HRFomer.\n Default: 0.0\n multiscale_output (bool, optional): Whether to output multi-level\n features produced by multiple branches. If False, only the first\n level feature will be output. Default: True.\n conv_cfg (dict, optional): Config of the conv layers.\n Default: None.\n norm_cfg (dict, optional): Config of the norm layers appended\n right after conv. Default: dict(type='SyncBN', requires_grad=True)\n transformer_norm_cfg (dict, optional): Config of the norm layers.\n Default: dict(type='LN', eps=1e-6)\n with_cp (bool): Use checkpoint or not. Using checkpoint will save some\n memory while slowing down the training speed. Default: False\n upsample_cfg(dict, optional): The config of upsample layers in fuse\n layers. Default: dict(mode='bilinear', align_corners=False)\n \"\"\"\n\n def __init__(self,\n num_branches,\n block,\n num_blocks,\n num_inchannels,\n num_channels,\n num_heads,\n num_window_sizes,\n num_mlp_ratios,\n multiscale_output=True,\n drop_paths=0.0,\n with_rpe=True,\n with_pad_mask=False,\n conv_cfg=None,\n norm_cfg=dict(type='SyncBN', requires_grad=True),\n transformer_norm_cfg=dict(type='LN', eps=1e-6),\n with_cp=False,\n upsample_cfg=dict(mode='bilinear', align_corners=False)):\n\n self.transformer_norm_cfg = transformer_norm_cfg\n self.drop_paths = drop_paths\n self.num_heads = num_heads\n self.num_window_sizes = num_window_sizes\n self.num_mlp_ratios = num_mlp_ratios\n self.with_rpe = with_rpe\n self.with_pad_mask = with_pad_mask\n\n super().__init__(num_branches, block, num_blocks, num_inchannels,\n num_channels, multiscale_output, with_cp, conv_cfg,\n norm_cfg, upsample_cfg)\n\n def _make_one_branch(self,\n branch_index,\n block,\n num_blocks,\n num_channels,\n stride=1):\n \"\"\"Build one branch.\"\"\"\n # HRFormerBlock does not support down sample layer yet.\n assert stride == 1 and self.in_channels[branch_index] == num_channels[\n branch_index]\n layers = []\n layers.append(\n block(\n self.in_channels[branch_index],\n num_channels[branch_index],\n num_heads=self.num_heads[branch_index],\n window_size=self.num_window_sizes[branch_index],\n mlp_ratio=self.num_mlp_ratios[branch_index],\n drop_path=self.drop_paths[0],\n norm_cfg=self.norm_cfg,\n transformer_norm_cfg=self.transformer_norm_cfg,\n init_cfg=None,\n with_rpe=self.with_rpe,\n with_pad_mask=self.with_pad_mask))\n\n self.in_channels[\n branch_index] = self.in_channels[branch_index] * block.expansion\n for i in range(1, num_blocks[branch_index]):\n layers.append(\n block(\n self.in_channels[branch_index],\n num_channels[branch_index],\n num_heads=self.num_heads[branch_index],\n window_size=self.num_window_sizes[branch_index],\n mlp_ratio=self.num_mlp_ratios[branch_index],\n drop_path=self.drop_paths[i],\n norm_cfg=self.norm_cfg,\n transformer_norm_cfg=self.transformer_norm_cfg,\n init_cfg=None,\n with_rpe=self.with_rpe,\n with_pad_mask=self.with_pad_mask))\n return nn.Sequential(*layers)\n\n def _make_fuse_layers(self):\n \"\"\"Build fuse layers.\"\"\"\n if self.num_branches == 1:\n return None\n num_branches = self.num_branches\n num_inchannels = self.in_channels\n fuse_layers = []\n for i in range(num_branches if self.multiscale_output else 1):\n fuse_layer = []\n for j in range(num_branches):\n if j > i:\n fuse_layer.append(\n nn.Sequential(\n build_conv_layer(\n self.conv_cfg,\n num_inchannels[j],\n num_inchannels[i],\n kernel_size=1,\n stride=1,\n bias=False),\n build_norm_layer(self.norm_cfg,\n num_inchannels[i])[1],\n nn.Upsample(\n scale_factor=2**(j - i),\n mode=self.upsample_cfg['mode'],\n align_corners=self.\n upsample_cfg['align_corners'])))\n elif j == i:\n fuse_layer.append(None)\n else:\n conv3x3s = []\n for k in range(i - j):\n if k == i - j - 1:\n num_outchannels_conv3x3 = num_inchannels[i]\n with_out_act = False\n else:\n num_outchannels_conv3x3 = num_inchannels[j]\n with_out_act = True\n sub_modules = [\n build_conv_layer(\n self.conv_cfg,\n num_inchannels[j],\n num_inchannels[j],\n kernel_size=3,\n stride=2,\n padding=1,\n groups=num_inchannels[j],\n bias=False,\n ),\n build_norm_layer(self.norm_cfg,\n num_inchannels[j])[1],\n build_conv_layer(\n self.conv_cfg,\n num_inchannels[j],\n num_outchannels_conv3x3,\n kernel_size=1,\n stride=1,\n bias=False,\n ),\n build_norm_layer(self.norm_cfg,\n num_outchannels_conv3x3)[1]\n ]\n if with_out_act:\n sub_modules.append(nn.ReLU(False))\n conv3x3s.append(nn.Sequential(*sub_modules))\n fuse_layer.append(nn.Sequential(*conv3x3s))\n fuse_layers.append(nn.ModuleList(fuse_layer))\n\n return nn.ModuleList(fuse_layers)\n\n def get_num_inchannels(self):\n \"\"\"Return the number of input channels.\"\"\"\n return self.in_channels\n\n\n@BACKBONES.register_module()\nclass HRFormer(HRNet):\n \"\"\"HRFormer backbone.\n\n This backbone is the implementation of `HRFormer: High-Resolution\n Transformer for Dense Prediction `_.\n\n Args:\n extra (dict): Detailed configuration for each stage of HRNet.\n There must be 4 stages, the configuration for each stage must have\n 5 keys:\n\n - num_modules (int): The number of HRModule in this stage.\n - num_branches (int): The number of branches in the HRModule.\n - block (str): The type of block.\n - num_blocks (tuple): The number of blocks in each branch.\n The length must be equal to num_branches.\n - num_channels (tuple): The number of channels in each branch.\n The length must be equal to num_branches.\n in_channels (int): Number of input image channels. Normally 3.\n conv_cfg (dict): Dictionary to construct and config conv layer.\n Default: None.\n norm_cfg (dict): Config of norm layer.\n Use `SyncBN` by default.\n transformer_norm_cfg (dict): Config of transformer norm layer.\n Use `LN` by default.\n norm_eval (bool): Whether to set norm layers to eval mode, namely,\n freeze running stats (mean and var). Note: Effect on Batch Norm\n and its variants only. Default: False.\n zero_init_residual (bool): Whether to use zero init for last norm layer\n in resblocks to let them behave as identity. Default: False.\n frozen_stages (int): Stages to be frozen (stop grad and set eval mode).\n -1 means not freezing any parameters. Default: -1.\n Example:\n >>> from mmpose.models import HRFormer\n >>> import torch\n >>> extra = dict(\n >>> stage1=dict(\n >>> num_modules=1,\n >>> num_branches=1,\n >>> block='BOTTLENECK',\n >>> num_blocks=(2, ),\n >>> num_channels=(64, )),\n >>> stage2=dict(\n >>> num_modules=1,\n >>> num_branches=2,\n >>> block='HRFORMER',\n >>> window_sizes=(7, 7),\n >>> num_heads=(1, 2),\n >>> mlp_ratios=(4, 4),\n >>> num_blocks=(2, 2),\n >>> num_channels=(32, 64)),\n >>> stage3=dict(\n >>> num_modules=4,\n >>> num_branches=3,\n >>> block='HRFORMER',\n >>> window_sizes=(7, 7, 7),\n >>> num_heads=(1, 2, 4),\n >>> mlp_ratios=(4, 4, 4),\n >>> num_blocks=(2, 2, 2),\n >>> num_channels=(32, 64, 128)),\n >>> stage4=dict(\n >>> num_modules=2,\n >>> num_branches=4,\n >>> block='HRFORMER',\n >>> window_sizes=(7, 7, 7, 7),\n >>> num_heads=(1, 2, 4, 8),\n >>> mlp_ratios=(4, 4, 4, 4),\n >>> num_blocks=(2, 2, 2, 2),\n >>> num_channels=(32, 64, 128, 256)))\n >>> self = HRFormer(extra, in_channels=1)\n >>> self.eval()\n >>> inputs = torch.rand(1, 1, 32, 32)\n >>> level_outputs = self.forward(inputs)\n >>> for level_out in level_outputs:\n ... print(tuple(level_out.shape))\n (1, 32, 8, 8)\n (1, 64, 4, 4)\n (1, 128, 2, 2)\n (1, 256, 1, 1)\n \"\"\"\n\n blocks_dict = {'BOTTLENECK': Bottleneck, 'HRFORMERBLOCK': HRFormerBlock}\n\n def __init__(self,\n extra,\n in_channels=3,\n conv_cfg=None,\n norm_cfg=dict(type='BN', requires_grad=True),\n transformer_norm_cfg=dict(type='LN', eps=1e-6),\n norm_eval=False,\n with_cp=False,\n zero_init_residual=False,\n frozen_stages=-1):\n\n # stochastic depth\n depths = [\n extra[stage]['num_blocks'][0] * extra[stage]['num_modules']\n for stage in ['stage2', 'stage3', 'stage4']\n ]\n depth_s2, depth_s3, _ = depths\n drop_path_rate = extra['drop_path_rate']\n dpr = [\n x.item() for x in torch.linspace(0, drop_path_rate, sum(depths))\n ]\n extra['stage2']['drop_path_rates'] = dpr[0:depth_s2]\n extra['stage3']['drop_path_rates'] = dpr[depth_s2:depth_s2 + depth_s3]\n extra['stage4']['drop_path_rates'] = dpr[depth_s2 + depth_s3:]\n\n # HRFormer use bilinear upsample as default\n upsample_cfg = extra.get('upsample', {\n 'mode': 'bilinear',\n 'align_corners': False\n })\n extra['upsample'] = upsample_cfg\n self.transformer_norm_cfg = transformer_norm_cfg\n self.with_rpe = extra.get('with_rpe', True)\n self.with_pad_mask = extra.get('with_pad_mask', False)\n\n super().__init__(extra, in_channels, conv_cfg, norm_cfg, norm_eval,\n with_cp, zero_init_residual, frozen_stages)\n\n def _make_stage(self,\n layer_config,\n num_inchannels,\n multiscale_output=True):\n \"\"\"Make each stage.\"\"\"\n num_modules = layer_config['num_modules']\n num_branches = layer_config['num_branches']\n num_blocks = layer_config['num_blocks']\n num_channels = layer_config['num_channels']\n block = self.blocks_dict[layer_config['block']]\n num_heads = layer_config['num_heads']\n num_window_sizes = layer_config['window_sizes']\n num_mlp_ratios = layer_config['mlp_ratios']\n drop_path_rates = layer_config['drop_path_rates']\n\n modules = []\n for i in range(num_modules):\n # multiscale_output is only used at the last module\n if not multiscale_output and i == num_modules - 1:\n reset_multiscale_output = False\n else:\n reset_multiscale_output = True\n\n modules.append(\n HRFomerModule(\n num_branches,\n block,\n num_blocks,\n num_inchannels,\n num_channels,\n num_heads,\n num_window_sizes,\n num_mlp_ratios,\n reset_multiscale_output,\n drop_paths=drop_path_rates[num_blocks[0] *\n i:num_blocks[0] * (i + 1)],\n with_rpe=self.with_rpe,\n with_pad_mask=self.with_pad_mask,\n conv_cfg=self.conv_cfg,\n norm_cfg=self.norm_cfg,\n transformer_norm_cfg=self.transformer_norm_cfg,\n with_cp=self.with_cp,\n upsample_cfg=self.upsample_cfg))\n num_inchannels = modules[-1].get_num_inchannels()\n\n return nn.Sequential(*modules), num_inchannels\n" }, { "path": "eval/GVHMR/hmr4d/utils/preproc/vitpose_pytorch/src/vitpose_infer/builder/backbones/litehrnet.py", "content": "# ------------------------------------------------------------------------------\n# Adapted from https://github.com/HRNet/Lite-HRNet\n# Original licence: Apache License 2.0.\n# ------------------------------------------------------------------------------\n\nimport mmcv\nimport torch\nimport torch.nn as nn\nimport torch.nn.functional as F\nimport torch.utils.checkpoint as cp\nfrom mmcv.cnn import (ConvModule, DepthwiseSeparableConvModule,\n build_conv_layer, build_norm_layer, constant_init,\n normal_init)\nfrom torch.nn.modules.batchnorm import _BatchNorm\n\nfrom mmpose.utils import get_root_logger\nfrom ..builder import BACKBONES\nfrom .utils import channel_shuffle, load_checkpoint\n\n\nclass SpatialWeighting(nn.Module):\n \"\"\"Spatial weighting module.\n\n Args:\n channels (int): The channels of the module.\n ratio (int): channel reduction ratio.\n conv_cfg (dict): Config dict for convolution layer.\n Default: None, which means using conv2d.\n norm_cfg (dict): Config dict for normalization layer.\n Default: None.\n act_cfg (dict): Config dict for activation layer.\n Default: (dict(type='ReLU'), dict(type='Sigmoid')).\n The last ConvModule uses Sigmoid by default.\n \"\"\"\n\n def __init__(self,\n channels,\n ratio=16,\n conv_cfg=None,\n norm_cfg=None,\n act_cfg=(dict(type='ReLU'), dict(type='Sigmoid'))):\n super().__init__()\n if isinstance(act_cfg, dict):\n act_cfg = (act_cfg, act_cfg)\n assert len(act_cfg) == 2\n assert mmcv.is_tuple_of(act_cfg, dict)\n self.global_avgpool = nn.AdaptiveAvgPool2d(1)\n self.conv1 = ConvModule(\n in_channels=channels,\n out_channels=int(channels / ratio),\n kernel_size=1,\n stride=1,\n conv_cfg=conv_cfg,\n norm_cfg=norm_cfg,\n act_cfg=act_cfg[0])\n self.conv2 = ConvModule(\n in_channels=int(channels / ratio),\n out_channels=channels,\n kernel_size=1,\n stride=1,\n conv_cfg=conv_cfg,\n norm_cfg=norm_cfg,\n act_cfg=act_cfg[1])\n\n def forward(self, x):\n out = self.global_avgpool(x)\n out = self.conv1(out)\n out = self.conv2(out)\n return x * out\n\n\nclass CrossResolutionWeighting(nn.Module):\n \"\"\"Cross-resolution channel weighting module.\n\n Args:\n channels (int): The channels of the module.\n ratio (int): channel reduction ratio.\n conv_cfg (dict): Config dict for convolution layer.\n Default: None, which means using conv2d.\n norm_cfg (dict): Config dict for normalization layer.\n Default: None.\n act_cfg (dict): Config dict for activation layer.\n Default: (dict(type='ReLU'), dict(type='Sigmoid')).\n The last ConvModule uses Sigmoid by default.\n \"\"\"\n\n def __init__(self,\n channels,\n ratio=16,\n conv_cfg=None,\n norm_cfg=None,\n act_cfg=(dict(type='ReLU'), dict(type='Sigmoid'))):\n super().__init__()\n if isinstance(act_cfg, dict):\n act_cfg = (act_cfg, act_cfg)\n assert len(act_cfg) == 2\n assert mmcv.is_tuple_of(act_cfg, dict)\n self.channels = channels\n total_channel = sum(channels)\n self.conv1 = ConvModule(\n in_channels=total_channel,\n out_channels=int(total_channel / ratio),\n kernel_size=1,\n stride=1,\n conv_cfg=conv_cfg,\n norm_cfg=norm_cfg,\n act_cfg=act_cfg[0])\n self.conv2 = ConvModule(\n in_channels=int(total_channel / ratio),\n out_channels=total_channel,\n kernel_size=1,\n stride=1,\n conv_cfg=conv_cfg,\n norm_cfg=norm_cfg,\n act_cfg=act_cfg[1])\n\n def forward(self, x):\n mini_size = x[-1].size()[-2:]\n out = [F.adaptive_avg_pool2d(s, mini_size) for s in x[:-1]] + [x[-1]]\n out = torch.cat(out, dim=1)\n out = self.conv1(out)\n out = self.conv2(out)\n out = torch.split(out, self.channels, dim=1)\n out = [\n s * F.interpolate(a, size=s.size()[-2:], mode='nearest')\n for s, a in zip(x, out)\n ]\n return out\n\n\nclass ConditionalChannelWeighting(nn.Module):\n \"\"\"Conditional channel weighting block.\n\n Args:\n in_channels (int): The input channels of the block.\n stride (int): Stride of the 3x3 convolution layer.\n reduce_ratio (int): channel reduction ratio.\n conv_cfg (dict): Config dict for convolution layer.\n Default: None, which means using conv2d.\n norm_cfg (dict): Config dict for normalization layer.\n Default: dict(type='BN').\n with_cp (bool): Use checkpoint or not. Using checkpoint will save some\n memory while slowing down the training speed. Default: False.\n \"\"\"\n\n def __init__(self,\n in_channels,\n stride,\n reduce_ratio,\n conv_cfg=None,\n norm_cfg=dict(type='BN'),\n with_cp=False):\n super().__init__()\n self.with_cp = with_cp\n self.stride = stride\n assert stride in [1, 2]\n\n branch_channels = [channel // 2 for channel in in_channels]\n\n self.cross_resolution_weighting = CrossResolutionWeighting(\n branch_channels,\n ratio=reduce_ratio,\n conv_cfg=conv_cfg,\n norm_cfg=norm_cfg)\n\n self.depthwise_convs = nn.ModuleList([\n ConvModule(\n channel,\n channel,\n kernel_size=3,\n stride=self.stride,\n padding=1,\n groups=channel,\n conv_cfg=conv_cfg,\n norm_cfg=norm_cfg,\n act_cfg=None) for channel in branch_channels\n ])\n\n self.spatial_weighting = nn.ModuleList([\n SpatialWeighting(channels=channel, ratio=4)\n for channel in branch_channels\n ])\n\n def forward(self, x):\n\n def _inner_forward(x):\n x = [s.chunk(2, dim=1) for s in x]\n x1 = [s[0] for s in x]\n x2 = [s[1] for s in x]\n\n x2 = self.cross_resolution_weighting(x2)\n x2 = [dw(s) for s, dw in zip(x2, self.depthwise_convs)]\n x2 = [sw(s) for s, sw in zip(x2, self.spatial_weighting)]\n\n out = [torch.cat([s1, s2], dim=1) for s1, s2 in zip(x1, x2)]\n out = [channel_shuffle(s, 2) for s in out]\n\n return out\n\n if self.with_cp and x.requires_grad:\n out = cp.checkpoint(_inner_forward, x)\n else:\n out = _inner_forward(x)\n\n return out\n\n\nclass Stem(nn.Module):\n \"\"\"Stem network block.\n\n Args:\n in_channels (int): The input channels of the block.\n stem_channels (int): Output channels of the stem layer.\n out_channels (int): The output channels of the block.\n expand_ratio (int): adjusts number of channels of the hidden layer\n in InvertedResidual by this amount.\n conv_cfg (dict): Config dict for convolution layer.\n Default: None, which means using conv2d.\n norm_cfg (dict): Config dict for normalization layer.\n Default: dict(type='BN').\n with_cp (bool): Use checkpoint or not. Using checkpoint will save some\n memory while slowing down the training speed. Default: False.\n \"\"\"\n\n def __init__(self,\n in_channels,\n stem_channels,\n out_channels,\n expand_ratio,\n conv_cfg=None,\n norm_cfg=dict(type='BN'),\n with_cp=False):\n super().__init__()\n self.in_channels = in_channels\n self.out_channels = out_channels\n self.conv_cfg = conv_cfg\n self.norm_cfg = norm_cfg\n self.with_cp = with_cp\n\n self.conv1 = ConvModule(\n in_channels=in_channels,\n out_channels=stem_channels,\n kernel_size=3,\n stride=2,\n padding=1,\n conv_cfg=self.conv_cfg,\n norm_cfg=self.norm_cfg,\n act_cfg=dict(type='ReLU'))\n\n mid_channels = int(round(stem_channels * expand_ratio))\n branch_channels = stem_channels // 2\n if stem_channels == self.out_channels:\n inc_channels = self.out_channels - branch_channels\n else:\n inc_channels = self.out_channels - stem_channels\n\n self.branch1 = nn.Sequential(\n ConvModule(\n branch_channels,\n branch_channels,\n kernel_size=3,\n stride=2,\n padding=1,\n groups=branch_channels,\n conv_cfg=conv_cfg,\n norm_cfg=norm_cfg,\n act_cfg=None),\n ConvModule(\n branch_channels,\n inc_channels,\n kernel_size=1,\n stride=1,\n padding=0,\n conv_cfg=conv_cfg,\n norm_cfg=norm_cfg,\n act_cfg=dict(type='ReLU')),\n )\n\n self.expand_conv = ConvModule(\n branch_channels,\n mid_channels,\n kernel_size=1,\n stride=1,\n padding=0,\n conv_cfg=conv_cfg,\n norm_cfg=norm_cfg,\n act_cfg=dict(type='ReLU'))\n self.depthwise_conv = ConvModule(\n mid_channels,\n mid_channels,\n kernel_size=3,\n stride=2,\n padding=1,\n groups=mid_channels,\n conv_cfg=conv_cfg,\n norm_cfg=norm_cfg,\n act_cfg=None)\n self.linear_conv = ConvModule(\n mid_channels,\n branch_channels\n if stem_channels == self.out_channels else stem_channels,\n kernel_size=1,\n stride=1,\n padding=0,\n conv_cfg=conv_cfg,\n norm_cfg=norm_cfg,\n act_cfg=dict(type='ReLU'))\n\n def forward(self, x):\n\n def _inner_forward(x):\n x = self.conv1(x)\n x1, x2 = x.chunk(2, dim=1)\n\n x2 = self.expand_conv(x2)\n x2 = self.depthwise_conv(x2)\n x2 = self.linear_conv(x2)\n\n out = torch.cat((self.branch1(x1), x2), dim=1)\n\n out = channel_shuffle(out, 2)\n\n return out\n\n if self.with_cp and x.requires_grad:\n out = cp.checkpoint(_inner_forward, x)\n else:\n out = _inner_forward(x)\n\n return out\n\n\nclass IterativeHead(nn.Module):\n \"\"\"Extra iterative head for feature learning.\n\n Args:\n in_channels (int): The input channels of the block.\n norm_cfg (dict): Config dict for normalization layer.\n Default: dict(type='BN').\n \"\"\"\n\n def __init__(self, in_channels, norm_cfg=dict(type='BN')):\n super().__init__()\n projects = []\n num_branchs = len(in_channels)\n self.in_channels = in_channels[::-1]\n\n for i in range(num_branchs):\n if i != num_branchs - 1:\n projects.append(\n DepthwiseSeparableConvModule(\n in_channels=self.in_channels[i],\n out_channels=self.in_channels[i + 1],\n kernel_size=3,\n stride=1,\n padding=1,\n norm_cfg=norm_cfg,\n act_cfg=dict(type='ReLU'),\n dw_act_cfg=None,\n pw_act_cfg=dict(type='ReLU')))\n else:\n projects.append(\n DepthwiseSeparableConvModule(\n in_channels=self.in_channels[i],\n out_channels=self.in_channels[i],\n kernel_size=3,\n stride=1,\n padding=1,\n norm_cfg=norm_cfg,\n act_cfg=dict(type='ReLU'),\n dw_act_cfg=None,\n pw_act_cfg=dict(type='ReLU')))\n self.projects = nn.ModuleList(projects)\n\n def forward(self, x):\n x = x[::-1]\n\n y = []\n last_x = None\n for i, s in enumerate(x):\n if last_x is not None:\n last_x = F.interpolate(\n last_x,\n size=s.size()[-2:],\n mode='bilinear',\n align_corners=True)\n s = s + last_x\n s = self.projects[i](s)\n y.append(s)\n last_x = s\n\n return y[::-1]\n\n\nclass ShuffleUnit(nn.Module):\n \"\"\"InvertedResidual block for ShuffleNetV2 backbone.\n\n Args:\n in_channels (int): The input channels of the block.\n out_channels (int): The output channels of the block.\n stride (int): Stride of the 3x3 convolution layer. Default: 1\n conv_cfg (dict): Config dict for convolution layer.\n Default: None, which means using conv2d.\n norm_cfg (dict): Config dict for normalization layer.\n Default: dict(type='BN').\n act_cfg (dict): Config dict for activation layer.\n Default: dict(type='ReLU').\n with_cp (bool): Use checkpoint or not. Using checkpoint will save some\n memory while slowing down the training speed. Default: False.\n \"\"\"\n\n def __init__(self,\n in_channels,\n out_channels,\n stride=1,\n conv_cfg=None,\n norm_cfg=dict(type='BN'),\n act_cfg=dict(type='ReLU'),\n with_cp=False):\n super().__init__()\n self.stride = stride\n self.with_cp = with_cp\n\n branch_features = out_channels // 2\n if self.stride == 1:\n assert in_channels == branch_features * 2, (\n f'in_channels ({in_channels}) should equal to '\n f'branch_features * 2 ({branch_features * 2}) '\n 'when stride is 1')\n\n if in_channels != branch_features * 2:\n assert self.stride != 1, (\n f'stride ({self.stride}) should not equal 1 when '\n f'in_channels != branch_features * 2')\n\n if self.stride > 1:\n self.branch1 = nn.Sequential(\n ConvModule(\n in_channels,\n in_channels,\n kernel_size=3,\n stride=self.stride,\n padding=1,\n groups=in_channels,\n conv_cfg=conv_cfg,\n norm_cfg=norm_cfg,\n act_cfg=None),\n ConvModule(\n in_channels,\n branch_features,\n kernel_size=1,\n stride=1,\n padding=0,\n conv_cfg=conv_cfg,\n norm_cfg=norm_cfg,\n act_cfg=act_cfg),\n )\n\n self.branch2 = nn.Sequential(\n ConvModule(\n in_channels if (self.stride > 1) else branch_features,\n branch_features,\n kernel_size=1,\n stride=1,\n padding=0,\n conv_cfg=conv_cfg,\n norm_cfg=norm_cfg,\n act_cfg=act_cfg),\n ConvModule(\n branch_features,\n branch_features,\n kernel_size=3,\n stride=self.stride,\n padding=1,\n groups=branch_features,\n conv_cfg=conv_cfg,\n norm_cfg=norm_cfg,\n act_cfg=None),\n ConvModule(\n branch_features,\n branch_features,\n kernel_size=1,\n stride=1,\n padding=0,\n conv_cfg=conv_cfg,\n norm_cfg=norm_cfg,\n act_cfg=act_cfg))\n\n def forward(self, x):\n\n def _inner_forward(x):\n if self.stride > 1:\n out = torch.cat((self.branch1(x), self.branch2(x)), dim=1)\n else:\n x1, x2 = x.chunk(2, dim=1)\n out = torch.cat((x1, self.branch2(x2)), dim=1)\n\n out = channel_shuffle(out, 2)\n\n return out\n\n if self.with_cp and x.requires_grad:\n out = cp.checkpoint(_inner_forward, x)\n else:\n out = _inner_forward(x)\n\n return out\n\n\nclass LiteHRModule(nn.Module):\n \"\"\"High-Resolution Module for LiteHRNet.\n\n It contains conditional channel weighting blocks and\n shuffle blocks.\n\n\n Args:\n num_branches (int): Number of branches in the module.\n num_blocks (int): Number of blocks in the module.\n in_channels (list(int)): Number of input image channels.\n reduce_ratio (int): Channel reduction ratio.\n module_type (str): 'LITE' or 'NAIVE'\n multiscale_output (bool): Whether to output multi-scale features.\n with_fuse (bool): Whether to use fuse layers.\n conv_cfg (dict): dictionary to construct and config conv layer.\n norm_cfg (dict): dictionary to construct and config norm layer.\n with_cp (bool): Use checkpoint or not. Using checkpoint will save some\n memory while slowing down the training speed.\n \"\"\"\n\n def __init__(\n self,\n num_branches,\n num_blocks,\n in_channels,\n reduce_ratio,\n module_type,\n multiscale_output=False,\n with_fuse=True,\n conv_cfg=None,\n norm_cfg=dict(type='BN'),\n with_cp=False,\n ):\n super().__init__()\n self._check_branches(num_branches, in_channels)\n\n self.in_channels = in_channels\n self.num_branches = num_branches\n\n self.module_type = module_type\n self.multiscale_output = multiscale_output\n self.with_fuse = with_fuse\n self.norm_cfg = norm_cfg\n self.conv_cfg = conv_cfg\n self.with_cp = with_cp\n\n if self.module_type.upper() == 'LITE':\n self.layers = self._make_weighting_blocks(num_blocks, reduce_ratio)\n elif self.module_type.upper() == 'NAIVE':\n self.layers = self._make_naive_branches(num_branches, num_blocks)\n else:\n raise ValueError(\"module_type should be either 'LITE' or 'NAIVE'.\")\n if self.with_fuse:\n self.fuse_layers = self._make_fuse_layers()\n self.relu = nn.ReLU()\n\n def _check_branches(self, num_branches, in_channels):\n \"\"\"Check input to avoid ValueError.\"\"\"\n if num_branches != len(in_channels):\n error_msg = f'NUM_BRANCHES({num_branches}) ' \\\n f'!= NUM_INCHANNELS({len(in_channels)})'\n raise ValueError(error_msg)\n\n def _make_weighting_blocks(self, num_blocks, reduce_ratio, stride=1):\n \"\"\"Make channel weighting blocks.\"\"\"\n layers = []\n for i in range(num_blocks):\n layers.append(\n ConditionalChannelWeighting(\n self.in_channels,\n stride=stride,\n reduce_ratio=reduce_ratio,\n conv_cfg=self.conv_cfg,\n norm_cfg=self.norm_cfg,\n with_cp=self.with_cp))\n\n return nn.Sequential(*layers)\n\n def _make_one_branch(self, branch_index, num_blocks, stride=1):\n \"\"\"Make one branch.\"\"\"\n layers = []\n layers.append(\n ShuffleUnit(\n self.in_channels[branch_index],\n self.in_channels[branch_index],\n stride=stride,\n conv_cfg=self.conv_cfg,\n norm_cfg=self.norm_cfg,\n act_cfg=dict(type='ReLU'),\n with_cp=self.with_cp))\n for i in range(1, num_blocks):\n layers.append(\n ShuffleUnit(\n self.in_channels[branch_index],\n self.in_channels[branch_index],\n stride=1,\n conv_cfg=self.conv_cfg,\n norm_cfg=self.norm_cfg,\n act_cfg=dict(type='ReLU'),\n with_cp=self.with_cp))\n\n return nn.Sequential(*layers)\n\n def _make_naive_branches(self, num_branches, num_blocks):\n \"\"\"Make branches.\"\"\"\n branches = []\n\n for i in range(num_branches):\n branches.append(self._make_one_branch(i, num_blocks))\n\n return nn.ModuleList(branches)\n\n def _make_fuse_layers(self):\n \"\"\"Make fuse layer.\"\"\"\n if self.num_branches == 1:\n return None\n\n num_branches = self.num_branches\n in_channels = self.in_channels\n fuse_layers = []\n num_out_branches = num_branches if self.multiscale_output else 1\n for i in range(num_out_branches):\n fuse_layer = []\n for j in range(num_branches):\n if j > i:\n fuse_layer.append(\n nn.Sequential(\n build_conv_layer(\n self.conv_cfg,\n in_channels[j],\n in_channels[i],\n kernel_size=1,\n stride=1,\n padding=0,\n bias=False),\n build_norm_layer(self.norm_cfg, in_channels[i])[1],\n nn.Upsample(\n scale_factor=2**(j - i), mode='nearest')))\n elif j == i:\n fuse_layer.append(None)\n else:\n conv_downsamples = []\n for k in range(i - j):\n if k == i - j - 1:\n conv_downsamples.append(\n nn.Sequential(\n build_conv_layer(\n self.conv_cfg,\n in_channels[j],\n in_channels[j],\n kernel_size=3,\n stride=2,\n padding=1,\n groups=in_channels[j],\n bias=False),\n build_norm_layer(self.norm_cfg,\n in_channels[j])[1],\n build_conv_layer(\n self.conv_cfg,\n in_channels[j],\n in_channels[i],\n kernel_size=1,\n stride=1,\n padding=0,\n bias=False),\n build_norm_layer(self.norm_cfg,\n in_channels[i])[1]))\n else:\n conv_downsamples.append(\n nn.Sequential(\n build_conv_layer(\n self.conv_cfg,\n in_channels[j],\n in_channels[j],\n kernel_size=3,\n stride=2,\n padding=1,\n groups=in_channels[j],\n bias=False),\n build_norm_layer(self.norm_cfg,\n in_channels[j])[1],\n build_conv_layer(\n self.conv_cfg,\n in_channels[j],\n in_channels[j],\n kernel_size=1,\n stride=1,\n padding=0,\n bias=False),\n build_norm_layer(self.norm_cfg,\n in_channels[j])[1],\n nn.ReLU(inplace=True)))\n fuse_layer.append(nn.Sequential(*conv_downsamples))\n fuse_layers.append(nn.ModuleList(fuse_layer))\n\n return nn.ModuleList(fuse_layers)\n\n def forward(self, x):\n \"\"\"Forward function.\"\"\"\n if self.num_branches == 1:\n return [self.layers[0](x[0])]\n\n if self.module_type.upper() == 'LITE':\n out = self.layers(x)\n elif self.module_type.upper() == 'NAIVE':\n for i in range(self.num_branches):\n x[i] = self.layers[i](x[i])\n out = x\n\n if self.with_fuse:\n out_fuse = []\n for i in range(len(self.fuse_layers)):\n # `y = 0` will lead to decreased accuracy (0.5~1 mAP)\n y = out[0] if i == 0 else self.fuse_layers[i][0](out[0])\n for j in range(self.num_branches):\n if i == j:\n y += out[j]\n else:\n y += self.fuse_layers[i][j](out[j])\n out_fuse.append(self.relu(y))\n out = out_fuse\n if not self.multiscale_output:\n out = [out[0]]\n return out\n\n\n@BACKBONES.register_module()\nclass LiteHRNet(nn.Module):\n \"\"\"Lite-HRNet backbone.\n\n `Lite-HRNet: A Lightweight High-Resolution Network\n `_.\n\n Code adapted from 'https://github.com/HRNet/Lite-HRNet'.\n\n Args:\n extra (dict): detailed configuration for each stage of HRNet.\n in_channels (int): Number of input image channels. Default: 3.\n conv_cfg (dict): dictionary to construct and config conv layer.\n norm_cfg (dict): dictionary to construct and config norm layer.\n norm_eval (bool): Whether to set norm layers to eval mode, namely,\n freeze running stats (mean and var). Note: Effect on Batch Norm\n and its variants only. Default: False\n with_cp (bool): Use checkpoint or not. Using checkpoint will save some\n memory while slowing down the training speed.\n\n Example:\n >>> from mmpose.models import LiteHRNet\n >>> import torch\n >>> extra=dict(\n >>> stem=dict(stem_channels=32, out_channels=32, expand_ratio=1),\n >>> num_stages=3,\n >>> stages_spec=dict(\n >>> num_modules=(2, 4, 2),\n >>> num_branches=(2, 3, 4),\n >>> num_blocks=(2, 2, 2),\n >>> module_type=('LITE', 'LITE', 'LITE'),\n >>> with_fuse=(True, True, True),\n >>> reduce_ratios=(8, 8, 8),\n >>> num_channels=(\n >>> (40, 80),\n >>> (40, 80, 160),\n >>> (40, 80, 160, 320),\n >>> )),\n >>> with_head=False)\n >>> self = LiteHRNet(extra, in_channels=1)\n >>> self.eval()\n >>> inputs = torch.rand(1, 1, 32, 32)\n >>> level_outputs = self.forward(inputs)\n >>> for level_out in level_outputs:\n ... print(tuple(level_out.shape))\n (1, 40, 8, 8)\n \"\"\"\n\n def __init__(self,\n extra,\n in_channels=3,\n conv_cfg=None,\n norm_cfg=dict(type='BN'),\n norm_eval=False,\n with_cp=False):\n super().__init__()\n self.extra = extra\n self.conv_cfg = conv_cfg\n self.norm_cfg = norm_cfg\n self.norm_eval = norm_eval\n self.with_cp = with_cp\n\n self.stem = Stem(\n in_channels,\n stem_channels=self.extra['stem']['stem_channels'],\n out_channels=self.extra['stem']['out_channels'],\n expand_ratio=self.extra['stem']['expand_ratio'],\n conv_cfg=self.conv_cfg,\n norm_cfg=self.norm_cfg)\n\n self.num_stages = self.extra['num_stages']\n self.stages_spec = self.extra['stages_spec']\n\n num_channels_last = [\n self.stem.out_channels,\n ]\n for i in range(self.num_stages):\n num_channels = self.stages_spec['num_channels'][i]\n num_channels = [num_channels[i] for i in range(len(num_channels))]\n setattr(\n self, f'transition{i}',\n self._make_transition_layer(num_channels_last, num_channels))\n\n stage, num_channels_last = self._make_stage(\n self.stages_spec, i, num_channels, multiscale_output=True)\n setattr(self, f'stage{i}', stage)\n\n self.with_head = self.extra['with_head']\n if self.with_head:\n self.head_layer = IterativeHead(\n in_channels=num_channels_last,\n norm_cfg=self.norm_cfg,\n )\n\n def _make_transition_layer(self, num_channels_pre_layer,\n num_channels_cur_layer):\n \"\"\"Make transition layer.\"\"\"\n num_branches_cur = len(num_channels_cur_layer)\n num_branches_pre = len(num_channels_pre_layer)\n\n transition_layers = []\n for i in range(num_branches_cur):\n if i < num_branches_pre:\n if num_channels_cur_layer[i] != num_channels_pre_layer[i]:\n transition_layers.append(\n nn.Sequential(\n build_conv_layer(\n self.conv_cfg,\n num_channels_pre_layer[i],\n num_channels_pre_layer[i],\n kernel_size=3,\n stride=1,\n padding=1,\n groups=num_channels_pre_layer[i],\n bias=False),\n build_norm_layer(self.norm_cfg,\n num_channels_pre_layer[i])[1],\n build_conv_layer(\n self.conv_cfg,\n num_channels_pre_layer[i],\n num_channels_cur_layer[i],\n kernel_size=1,\n stride=1,\n padding=0,\n bias=False),\n build_norm_layer(self.norm_cfg,\n num_channels_cur_layer[i])[1],\n nn.ReLU()))\n else:\n transition_layers.append(None)\n else:\n conv_downsamples = []\n for j in range(i + 1 - num_branches_pre):\n in_channels = num_channels_pre_layer[-1]\n out_channels = num_channels_cur_layer[i] \\\n if j == i - num_branches_pre else in_channels\n conv_downsamples.append(\n nn.Sequential(\n build_conv_layer(\n self.conv_cfg,\n in_channels,\n in_channels,\n kernel_size=3,\n stride=2,\n padding=1,\n groups=in_channels,\n bias=False),\n build_norm_layer(self.norm_cfg, in_channels)[1],\n build_conv_layer(\n self.conv_cfg,\n in_channels,\n out_channels,\n kernel_size=1,\n stride=1,\n padding=0,\n bias=False),\n build_norm_layer(self.norm_cfg, out_channels)[1],\n nn.ReLU()))\n transition_layers.append(nn.Sequential(*conv_downsamples))\n\n return nn.ModuleList(transition_layers)\n\n def _make_stage(self,\n stages_spec,\n stage_index,\n in_channels,\n multiscale_output=True):\n num_modules = stages_spec['num_modules'][stage_index]\n num_branches = stages_spec['num_branches'][stage_index]\n num_blocks = stages_spec['num_blocks'][stage_index]\n reduce_ratio = stages_spec['reduce_ratios'][stage_index]\n with_fuse = stages_spec['with_fuse'][stage_index]\n module_type = stages_spec['module_type'][stage_index]\n\n modules = []\n for i in range(num_modules):\n # multi_scale_output is only used last module\n if not multiscale_output and i == num_modules - 1:\n reset_multiscale_output = False\n else:\n reset_multiscale_output = True\n\n modules.append(\n LiteHRModule(\n num_branches,\n num_blocks,\n in_channels,\n reduce_ratio,\n module_type,\n multiscale_output=reset_multiscale_output,\n with_fuse=with_fuse,\n conv_cfg=self.conv_cfg,\n norm_cfg=self.norm_cfg,\n with_cp=self.with_cp))\n in_channels = modules[-1].in_channels\n\n return nn.Sequential(*modules), in_channels\n\n def init_weights(self, pretrained=None):\n \"\"\"Initialize the weights in backbone.\n\n Args:\n pretrained (str, optional): Path to pre-trained weights.\n Defaults to None.\n \"\"\"\n if isinstance(pretrained, str):\n logger = get_root_logger()\n load_checkpoint(self, pretrained, strict=False, logger=logger)\n elif pretrained is None:\n for m in self.modules():\n if isinstance(m, nn.Conv2d):\n normal_init(m, std=0.001)\n elif isinstance(m, (_BatchNorm, nn.GroupNorm)):\n constant_init(m, 1)\n else:\n raise TypeError('pretrained must be a str or None')\n\n def forward(self, x):\n \"\"\"Forward function.\"\"\"\n x = self.stem(x)\n\n y_list = [x]\n for i in range(self.num_stages):\n x_list = []\n transition = getattr(self, f'transition{i}')\n for j in range(self.stages_spec['num_branches'][i]):\n if transition[j]:\n if j >= len(y_list):\n x_list.append(transition[j](y_list[-1]))\n else:\n x_list.append(transition[j](y_list[j]))\n else:\n x_list.append(y_list[j])\n y_list = getattr(self, f'stage{i}')(x_list)\n\n x = y_list\n if self.with_head:\n x = self.head_layer(x)\n\n return [x[0]]\n\n def train(self, mode=True):\n \"\"\"Convert the model into training mode.\"\"\"\n super().train(mode)\n if mode and self.norm_eval:\n for m in self.modules():\n if isinstance(m, _BatchNorm):\n m.eval()\n" }, { "path": "eval/GVHMR/hmr4d/utils/preproc/vitpose_pytorch/src/vitpose_infer/builder/backbones/mobilenet_v2.py", "content": "# Copyright (c) OpenMMLab. All rights reserved.\nimport copy\nimport logging\n\nimport torch.nn as nn\nimport torch.utils.checkpoint as cp\nfrom mmcv.cnn import ConvModule, constant_init, kaiming_init\nfrom torch.nn.modules.batchnorm import _BatchNorm\n\nfrom ..builder import BACKBONES\nfrom .base_backbone import BaseBackbone\nfrom .utils import load_checkpoint, make_divisible\n\n\nclass InvertedResidual(nn.Module):\n \"\"\"InvertedResidual block for MobileNetV2.\n\n Args:\n in_channels (int): The input channels of the InvertedResidual block.\n out_channels (int): The output channels of the InvertedResidual block.\n stride (int): Stride of the middle (first) 3x3 convolution.\n expand_ratio (int): adjusts number of channels of the hidden layer\n in InvertedResidual by this amount.\n conv_cfg (dict): Config dict for convolution layer.\n Default: None, which means using conv2d.\n norm_cfg (dict): Config dict for normalization layer.\n Default: dict(type='BN').\n act_cfg (dict): Config dict for activation layer.\n Default: dict(type='ReLU6').\n with_cp (bool): Use checkpoint or not. Using checkpoint will save some\n memory while slowing down the training speed. Default: False.\n \"\"\"\n\n def __init__(self,\n in_channels,\n out_channels,\n stride,\n expand_ratio,\n conv_cfg=None,\n norm_cfg=dict(type='BN'),\n act_cfg=dict(type='ReLU6'),\n with_cp=False):\n # Protect mutable default arguments\n norm_cfg = copy.deepcopy(norm_cfg)\n act_cfg = copy.deepcopy(act_cfg)\n super().__init__()\n self.stride = stride\n assert stride in [1, 2], f'stride must in [1, 2]. ' \\\n f'But received {stride}.'\n self.with_cp = with_cp\n self.use_res_connect = self.stride == 1 and in_channels == out_channels\n hidden_dim = int(round(in_channels * expand_ratio))\n\n layers = []\n if expand_ratio != 1:\n layers.append(\n ConvModule(\n in_channels=in_channels,\n out_channels=hidden_dim,\n kernel_size=1,\n conv_cfg=conv_cfg,\n norm_cfg=norm_cfg,\n act_cfg=act_cfg))\n layers.extend([\n ConvModule(\n in_channels=hidden_dim,\n out_channels=hidden_dim,\n kernel_size=3,\n stride=stride,\n padding=1,\n groups=hidden_dim,\n conv_cfg=conv_cfg,\n norm_cfg=norm_cfg,\n act_cfg=act_cfg),\n ConvModule(\n in_channels=hidden_dim,\n out_channels=out_channels,\n kernel_size=1,\n conv_cfg=conv_cfg,\n norm_cfg=norm_cfg,\n act_cfg=None)\n ])\n self.conv = nn.Sequential(*layers)\n\n def forward(self, x):\n\n def _inner_forward(x):\n if self.use_res_connect:\n return x + self.conv(x)\n return self.conv(x)\n\n if self.with_cp and x.requires_grad:\n out = cp.checkpoint(_inner_forward, x)\n else:\n out = _inner_forward(x)\n\n return out\n\n\n@BACKBONES.register_module()\nclass MobileNetV2(BaseBackbone):\n \"\"\"MobileNetV2 backbone.\n\n Args:\n widen_factor (float): Width multiplier, multiply number of\n channels in each layer by this amount. Default: 1.0.\n out_indices (None or Sequence[int]): Output from which stages.\n Default: (7, ).\n frozen_stages (int): Stages to be frozen (all param fixed).\n Default: -1, which means not freezing any parameters.\n conv_cfg (dict): Config dict for convolution layer.\n Default: None, which means using conv2d.\n norm_cfg (dict): Config dict for normalization layer.\n Default: dict(type='BN').\n act_cfg (dict): Config dict for activation layer.\n Default: dict(type='ReLU6').\n norm_eval (bool): Whether to set norm layers to eval mode, namely,\n freeze running stats (mean and var). Note: Effect on Batch Norm\n and its variants only. Default: False.\n with_cp (bool): Use checkpoint or not. Using checkpoint will save some\n memory while slowing down the training speed. Default: False.\n \"\"\"\n\n # Parameters to build layers. 4 parameters are needed to construct a\n # layer, from left to right: expand_ratio, channel, num_blocks, stride.\n arch_settings = [[1, 16, 1, 1], [6, 24, 2, 2], [6, 32, 3, 2],\n [6, 64, 4, 2], [6, 96, 3, 1], [6, 160, 3, 2],\n [6, 320, 1, 1]]\n\n def __init__(self,\n widen_factor=1.,\n out_indices=(7, ),\n frozen_stages=-1,\n conv_cfg=None,\n norm_cfg=dict(type='BN'),\n act_cfg=dict(type='ReLU6'),\n norm_eval=False,\n with_cp=False):\n # Protect mutable default arguments\n norm_cfg = copy.deepcopy(norm_cfg)\n act_cfg = copy.deepcopy(act_cfg)\n super().__init__()\n self.widen_factor = widen_factor\n self.out_indices = out_indices\n for index in out_indices:\n if index not in range(0, 8):\n raise ValueError('the item in out_indices must in '\n f'range(0, 8). But received {index}')\n\n if frozen_stages not in range(-1, 8):\n raise ValueError('frozen_stages must be in range(-1, 8). '\n f'But received {frozen_stages}')\n self.out_indices = out_indices\n self.frozen_stages = frozen_stages\n self.conv_cfg = conv_cfg\n self.norm_cfg = norm_cfg\n self.act_cfg = act_cfg\n self.norm_eval = norm_eval\n self.with_cp = with_cp\n\n self.in_channels = make_divisible(32 * widen_factor, 8)\n\n self.conv1 = ConvModule(\n in_channels=3,\n out_channels=self.in_channels,\n kernel_size=3,\n stride=2,\n padding=1,\n conv_cfg=self.conv_cfg,\n norm_cfg=self.norm_cfg,\n act_cfg=self.act_cfg)\n\n self.layers = []\n\n for i, layer_cfg in enumerate(self.arch_settings):\n expand_ratio, channel, num_blocks, stride = layer_cfg\n out_channels = make_divisible(channel * widen_factor, 8)\n inverted_res_layer = self.make_layer(\n out_channels=out_channels,\n num_blocks=num_blocks,\n stride=stride,\n expand_ratio=expand_ratio)\n layer_name = f'layer{i + 1}'\n self.add_module(layer_name, inverted_res_layer)\n self.layers.append(layer_name)\n\n if widen_factor > 1.0:\n self.out_channel = int(1280 * widen_factor)\n else:\n self.out_channel = 1280\n\n layer = ConvModule(\n in_channels=self.in_channels,\n out_channels=self.out_channel,\n kernel_size=1,\n stride=1,\n padding=0,\n conv_cfg=self.conv_cfg,\n norm_cfg=self.norm_cfg,\n act_cfg=self.act_cfg)\n self.add_module('conv2', layer)\n self.layers.append('conv2')\n\n def make_layer(self, out_channels, num_blocks, stride, expand_ratio):\n \"\"\"Stack InvertedResidual blocks to build a layer for MobileNetV2.\n\n Args:\n out_channels (int): out_channels of block.\n num_blocks (int): number of blocks.\n stride (int): stride of the first block. Default: 1\n expand_ratio (int): Expand the number of channels of the\n hidden layer in InvertedResidual by this ratio. Default: 6.\n \"\"\"\n layers = []\n for i in range(num_blocks):\n if i >= 1:\n stride = 1\n layers.append(\n InvertedResidual(\n self.in_channels,\n out_channels,\n stride,\n expand_ratio=expand_ratio,\n conv_cfg=self.conv_cfg,\n norm_cfg=self.norm_cfg,\n act_cfg=self.act_cfg,\n with_cp=self.with_cp))\n self.in_channels = out_channels\n\n return nn.Sequential(*layers)\n\n def init_weights(self, pretrained=None):\n if isinstance(pretrained, str):\n logger = logging.getLogger()\n load_checkpoint(self, pretrained, strict=False, logger=logger)\n elif pretrained is None:\n for m in self.modules():\n if isinstance(m, nn.Conv2d):\n kaiming_init(m)\n elif isinstance(m, (_BatchNorm, nn.GroupNorm)):\n constant_init(m, 1)\n else:\n raise TypeError('pretrained must be a str or None')\n\n def forward(self, x):\n x = self.conv1(x)\n\n outs = []\n for i, layer_name in enumerate(self.layers):\n layer = getattr(self, layer_name)\n x = layer(x)\n if i in self.out_indices:\n outs.append(x)\n\n if len(outs) == 1:\n return outs[0]\n return tuple(outs)\n\n def _freeze_stages(self):\n if self.frozen_stages >= 0:\n for param in self.conv1.parameters():\n param.requires_grad = False\n for i in range(1, self.frozen_stages + 1):\n layer = getattr(self, f'layer{i}')\n layer.eval()\n for param in layer.parameters():\n param.requires_grad = False\n\n def train(self, mode=True):\n super().train(mode)\n self._freeze_stages()\n if mode and self.norm_eval:\n for m in self.modules():\n if isinstance(m, _BatchNorm):\n m.eval()\n" }, { "path": "eval/GVHMR/hmr4d/utils/preproc/vitpose_pytorch/src/vitpose_infer/builder/backbones/mobilenet_v3.py", "content": "# Copyright (c) OpenMMLab. All rights reserved.\nimport copy\nimport logging\n\nimport torch.nn as nn\nfrom mmcv.cnn import ConvModule, constant_init, kaiming_init\nfrom torch.nn.modules.batchnorm import _BatchNorm\n\nfrom ..builder import BACKBONES\nfrom .base_backbone import BaseBackbone\nfrom .utils import InvertedResidual, load_checkpoint\n\n\n@BACKBONES.register_module()\nclass MobileNetV3(BaseBackbone):\n \"\"\"MobileNetV3 backbone.\n\n Args:\n arch (str): Architecture of mobilnetv3, from {small, big}.\n Default: small.\n conv_cfg (dict): Config dict for convolution layer.\n Default: None, which means using conv2d.\n norm_cfg (dict): Config dict for normalization layer.\n Default: dict(type='BN').\n out_indices (None or Sequence[int]): Output from which stages.\n Default: (-1, ), which means output tensors from final stage.\n frozen_stages (int): Stages to be frozen (all param fixed).\n Default: -1, which means not freezing any parameters.\n norm_eval (bool): Whether to set norm layers to eval mode, namely,\n freeze running stats (mean and var). Note: Effect on Batch Norm\n and its variants only. Default: False.\n with_cp (bool): Use checkpoint or not. Using checkpoint will save\n some memory while slowing down the training speed.\n Default: False.\n \"\"\"\n # Parameters to build each block:\n # [kernel size, mid channels, out channels, with_se, act type, stride]\n arch_settings = {\n 'small': [[3, 16, 16, True, 'ReLU', 2],\n [3, 72, 24, False, 'ReLU', 2],\n [3, 88, 24, False, 'ReLU', 1],\n [5, 96, 40, True, 'HSwish', 2],\n [5, 240, 40, True, 'HSwish', 1],\n [5, 240, 40, True, 'HSwish', 1],\n [5, 120, 48, True, 'HSwish', 1],\n [5, 144, 48, True, 'HSwish', 1],\n [5, 288, 96, True, 'HSwish', 2],\n [5, 576, 96, True, 'HSwish', 1],\n [5, 576, 96, True, 'HSwish', 1]],\n 'big': [[3, 16, 16, False, 'ReLU', 1],\n [3, 64, 24, False, 'ReLU', 2],\n [3, 72, 24, False, 'ReLU', 1],\n [5, 72, 40, True, 'ReLU', 2],\n [5, 120, 40, True, 'ReLU', 1],\n [5, 120, 40, True, 'ReLU', 1],\n [3, 240, 80, False, 'HSwish', 2],\n [3, 200, 80, False, 'HSwish', 1],\n [3, 184, 80, False, 'HSwish', 1],\n [3, 184, 80, False, 'HSwish', 1],\n [3, 480, 112, True, 'HSwish', 1],\n [3, 672, 112, True, 'HSwish', 1],\n [5, 672, 160, True, 'HSwish', 1],\n [5, 672, 160, True, 'HSwish', 2],\n [5, 960, 160, True, 'HSwish', 1]]\n } # yapf: disable\n\n def __init__(self,\n arch='small',\n conv_cfg=None,\n norm_cfg=dict(type='BN'),\n out_indices=(-1, ),\n frozen_stages=-1,\n norm_eval=False,\n with_cp=False):\n # Protect mutable default arguments\n norm_cfg = copy.deepcopy(norm_cfg)\n super().__init__()\n assert arch in self.arch_settings\n for index in out_indices:\n if index not in range(-len(self.arch_settings[arch]),\n len(self.arch_settings[arch])):\n raise ValueError('the item in out_indices must in '\n f'range(0, {len(self.arch_settings[arch])}). '\n f'But received {index}')\n\n if frozen_stages not in range(-1, len(self.arch_settings[arch])):\n raise ValueError('frozen_stages must be in range(-1, '\n f'{len(self.arch_settings[arch])}). '\n f'But received {frozen_stages}')\n self.arch = arch\n self.conv_cfg = conv_cfg\n self.norm_cfg = norm_cfg\n self.out_indices = out_indices\n self.frozen_stages = frozen_stages\n self.norm_eval = norm_eval\n self.with_cp = with_cp\n\n self.in_channels = 16\n self.conv1 = ConvModule(\n in_channels=3,\n out_channels=self.in_channels,\n kernel_size=3,\n stride=2,\n padding=1,\n conv_cfg=conv_cfg,\n norm_cfg=norm_cfg,\n act_cfg=dict(type='HSwish'))\n\n self.layers = self._make_layer()\n self.feat_dim = self.arch_settings[arch][-1][2]\n\n def _make_layer(self):\n layers = []\n layer_setting = self.arch_settings[self.arch]\n for i, params in enumerate(layer_setting):\n (kernel_size, mid_channels, out_channels, with_se, act,\n stride) = params\n if with_se:\n se_cfg = dict(\n channels=mid_channels,\n ratio=4,\n act_cfg=(dict(type='ReLU'), dict(type='HSigmoid')))\n else:\n se_cfg = None\n\n layer = InvertedResidual(\n in_channels=self.in_channels,\n out_channels=out_channels,\n mid_channels=mid_channels,\n kernel_size=kernel_size,\n stride=stride,\n se_cfg=se_cfg,\n with_expand_conv=True,\n conv_cfg=self.conv_cfg,\n norm_cfg=self.norm_cfg,\n act_cfg=dict(type=act),\n with_cp=self.with_cp)\n self.in_channels = out_channels\n layer_name = f'layer{i + 1}'\n self.add_module(layer_name, layer)\n layers.append(layer_name)\n return layers\n\n def init_weights(self, pretrained=None):\n if isinstance(pretrained, str):\n logger = logging.getLogger()\n load_checkpoint(self, pretrained, strict=False, logger=logger)\n elif pretrained is None:\n for m in self.modules():\n if isinstance(m, nn.Conv2d):\n kaiming_init(m)\n elif isinstance(m, nn.BatchNorm2d):\n constant_init(m, 1)\n else:\n raise TypeError('pretrained must be a str or None')\n\n def forward(self, x):\n x = self.conv1(x)\n\n outs = []\n for i, layer_name in enumerate(self.layers):\n layer = getattr(self, layer_name)\n x = layer(x)\n if i in self.out_indices or \\\n i - len(self.layers) in self.out_indices:\n outs.append(x)\n\n if len(outs) == 1:\n return outs[0]\n return tuple(outs)\n\n def _freeze_stages(self):\n if self.frozen_stages >= 0:\n for param in self.conv1.parameters():\n param.requires_grad = False\n for i in range(1, self.frozen_stages + 1):\n layer = getattr(self, f'layer{i}')\n layer.eval()\n for param in layer.parameters():\n param.requires_grad = False\n\n def train(self, mode=True):\n super().train(mode)\n self._freeze_stages()\n if mode and self.norm_eval:\n for m in self.modules():\n if isinstance(m, _BatchNorm):\n m.eval()\n" }, { "path": "eval/GVHMR/hmr4d/utils/preproc/vitpose_pytorch/src/vitpose_infer/builder/backbones/mspn.py", "content": "# Copyright (c) OpenMMLab. All rights reserved.\nimport copy as cp\nfrom collections import OrderedDict\n\nimport torch.nn as nn\nimport torch.nn.functional as F\nfrom mmcv.cnn import (ConvModule, MaxPool2d, constant_init, kaiming_init,\n normal_init)\nfrom mmcv.runner.checkpoint import load_state_dict\n\nfrom mmpose.utils import get_root_logger\nfrom ..builder import BACKBONES\nfrom .base_backbone import BaseBackbone\nfrom .resnet import Bottleneck as _Bottleneck\nfrom .utils.utils import get_state_dict\n\n\nclass Bottleneck(_Bottleneck):\n expansion = 4\n \"\"\"Bottleneck block for MSPN.\n\n Args:\n in_channels (int): Input channels of this block.\n out_channels (int): Output channels of this block.\n stride (int): stride of the block. Default: 1\n downsample (nn.Module): downsample operation on identity branch.\n Default: None\n norm_cfg (dict): dictionary to construct and config norm layer.\n Default: dict(type='BN')\n \"\"\"\n\n def __init__(self, in_channels, out_channels, **kwargs):\n super().__init__(in_channels, out_channels * 4, **kwargs)\n\n\nclass DownsampleModule(nn.Module):\n \"\"\"Downsample module for MSPN.\n\n Args:\n block (nn.Module): Downsample block.\n num_blocks (list): Number of blocks in each downsample unit.\n num_units (int): Numbers of downsample units. Default: 4\n has_skip (bool): Have skip connections from prior upsample\n module or not. Default:False\n norm_cfg (dict): dictionary to construct and config norm layer.\n Default: dict(type='BN')\n in_channels (int): Number of channels of the input feature to\n downsample module. Default: 64\n \"\"\"\n\n def __init__(self,\n block,\n num_blocks,\n num_units=4,\n has_skip=False,\n norm_cfg=dict(type='BN'),\n in_channels=64):\n # Protect mutable default arguments\n norm_cfg = cp.deepcopy(norm_cfg)\n super().__init__()\n self.has_skip = has_skip\n self.in_channels = in_channels\n assert len(num_blocks) == num_units\n self.num_blocks = num_blocks\n self.num_units = num_units\n self.norm_cfg = norm_cfg\n self.layer1 = self._make_layer(block, in_channels, num_blocks[0])\n for i in range(1, num_units):\n module_name = f'layer{i + 1}'\n self.add_module(\n module_name,\n self._make_layer(\n block, in_channels * pow(2, i), num_blocks[i], stride=2))\n\n def _make_layer(self, block, out_channels, blocks, stride=1):\n downsample = None\n if stride != 1 or self.in_channels != out_channels * block.expansion:\n downsample = ConvModule(\n self.in_channels,\n out_channels * block.expansion,\n kernel_size=1,\n stride=stride,\n padding=0,\n norm_cfg=self.norm_cfg,\n act_cfg=None,\n inplace=True)\n\n units = list()\n units.append(\n block(\n self.in_channels,\n out_channels,\n stride=stride,\n downsample=downsample,\n norm_cfg=self.norm_cfg))\n self.in_channels = out_channels * block.expansion\n for _ in range(1, blocks):\n units.append(block(self.in_channels, out_channels))\n\n return nn.Sequential(*units)\n\n def forward(self, x, skip1, skip2):\n out = list()\n for i in range(self.num_units):\n module_name = f'layer{i + 1}'\n module_i = getattr(self, module_name)\n x = module_i(x)\n if self.has_skip:\n x = x + skip1[i] + skip2[i]\n out.append(x)\n out.reverse()\n\n return tuple(out)\n\n\nclass UpsampleUnit(nn.Module):\n \"\"\"Upsample unit for upsample module.\n\n Args:\n ind (int): Indicates whether to interpolate (>0) and whether to\n generate feature map for the next hourglass-like module.\n num_units (int): Number of units that form a upsample module. Along\n with ind and gen_cross_conv, nm_units is used to decide whether\n to generate feature map for the next hourglass-like module.\n in_channels (int): Channel number of the skip-in feature maps from\n the corresponding downsample unit.\n unit_channels (int): Channel number in this unit. Default:256.\n gen_skip: (bool): Whether or not to generate skips for the posterior\n downsample module. Default:False\n gen_cross_conv (bool): Whether to generate feature map for the next\n hourglass-like module. Default:False\n norm_cfg (dict): dictionary to construct and config norm layer.\n Default: dict(type='BN')\n out_channels (int): Number of channels of feature output by upsample\n module. Must equal to in_channels of downsample module. Default:64\n \"\"\"\n\n def __init__(self,\n ind,\n num_units,\n in_channels,\n unit_channels=256,\n gen_skip=False,\n gen_cross_conv=False,\n norm_cfg=dict(type='BN'),\n out_channels=64):\n # Protect mutable default arguments\n norm_cfg = cp.deepcopy(norm_cfg)\n super().__init__()\n self.num_units = num_units\n self.norm_cfg = norm_cfg\n self.in_skip = ConvModule(\n in_channels,\n unit_channels,\n kernel_size=1,\n stride=1,\n padding=0,\n norm_cfg=self.norm_cfg,\n act_cfg=None,\n inplace=True)\n self.relu = nn.ReLU(inplace=True)\n\n self.ind = ind\n if self.ind > 0:\n self.up_conv = ConvModule(\n unit_channels,\n unit_channels,\n kernel_size=1,\n stride=1,\n padding=0,\n norm_cfg=self.norm_cfg,\n act_cfg=None,\n inplace=True)\n\n self.gen_skip = gen_skip\n if self.gen_skip:\n self.out_skip1 = ConvModule(\n in_channels,\n in_channels,\n kernel_size=1,\n stride=1,\n padding=0,\n norm_cfg=self.norm_cfg,\n inplace=True)\n\n self.out_skip2 = ConvModule(\n unit_channels,\n in_channels,\n kernel_size=1,\n stride=1,\n padding=0,\n norm_cfg=self.norm_cfg,\n inplace=True)\n\n self.gen_cross_conv = gen_cross_conv\n if self.ind == num_units - 1 and self.gen_cross_conv:\n self.cross_conv = ConvModule(\n unit_channels,\n out_channels,\n kernel_size=1,\n stride=1,\n padding=0,\n norm_cfg=self.norm_cfg,\n inplace=True)\n\n def forward(self, x, up_x):\n out = self.in_skip(x)\n\n if self.ind > 0:\n up_x = F.interpolate(\n up_x,\n size=(x.size(2), x.size(3)),\n mode='bilinear',\n align_corners=True)\n up_x = self.up_conv(up_x)\n out = out + up_x\n out = self.relu(out)\n\n skip1 = None\n skip2 = None\n if self.gen_skip:\n skip1 = self.out_skip1(x)\n skip2 = self.out_skip2(out)\n\n cross_conv = None\n if self.ind == self.num_units - 1 and self.gen_cross_conv:\n cross_conv = self.cross_conv(out)\n\n return out, skip1, skip2, cross_conv\n\n\nclass UpsampleModule(nn.Module):\n \"\"\"Upsample module for MSPN.\n\n Args:\n unit_channels (int): Channel number in the upsample units.\n Default:256.\n num_units (int): Numbers of upsample units. Default: 4\n gen_skip (bool): Whether to generate skip for posterior downsample\n module or not. Default:False\n gen_cross_conv (bool): Whether to generate feature map for the next\n hourglass-like module. Default:False\n norm_cfg (dict): dictionary to construct and config norm layer.\n Default: dict(type='BN')\n out_channels (int): Number of channels of feature output by upsample\n module. Must equal to in_channels of downsample module. Default:64\n \"\"\"\n\n def __init__(self,\n unit_channels=256,\n num_units=4,\n gen_skip=False,\n gen_cross_conv=False,\n norm_cfg=dict(type='BN'),\n out_channels=64):\n # Protect mutable default arguments\n norm_cfg = cp.deepcopy(norm_cfg)\n super().__init__()\n self.in_channels = list()\n for i in range(num_units):\n self.in_channels.append(Bottleneck.expansion * out_channels *\n pow(2, i))\n self.in_channels.reverse()\n self.num_units = num_units\n self.gen_skip = gen_skip\n self.gen_cross_conv = gen_cross_conv\n self.norm_cfg = norm_cfg\n for i in range(num_units):\n module_name = f'up{i + 1}'\n self.add_module(\n module_name,\n UpsampleUnit(\n i,\n self.num_units,\n self.in_channels[i],\n unit_channels,\n self.gen_skip,\n self.gen_cross_conv,\n norm_cfg=self.norm_cfg,\n out_channels=64))\n\n def forward(self, x):\n out = list()\n skip1 = list()\n skip2 = list()\n cross_conv = None\n for i in range(self.num_units):\n module_i = getattr(self, f'up{i + 1}')\n if i == 0:\n outi, skip1_i, skip2_i, _ = module_i(x[i], None)\n elif i == self.num_units - 1:\n outi, skip1_i, skip2_i, cross_conv = module_i(x[i], out[i - 1])\n else:\n outi, skip1_i, skip2_i, _ = module_i(x[i], out[i - 1])\n out.append(outi)\n skip1.append(skip1_i)\n skip2.append(skip2_i)\n skip1.reverse()\n skip2.reverse()\n\n return out, skip1, skip2, cross_conv\n\n\nclass SingleStageNetwork(nn.Module):\n \"\"\"Single_stage Network.\n\n Args:\n unit_channels (int): Channel number in the upsample units. Default:256.\n num_units (int): Numbers of downsample/upsample units. Default: 4\n gen_skip (bool): Whether to generate skip for posterior downsample\n module or not. Default:False\n gen_cross_conv (bool): Whether to generate feature map for the next\n hourglass-like module. Default:False\n has_skip (bool): Have skip connections from prior upsample\n module or not. Default:False\n num_blocks (list): Number of blocks in each downsample unit.\n Default: [2, 2, 2, 2] Note: Make sure num_units==len(num_blocks)\n norm_cfg (dict): dictionary to construct and config norm layer.\n Default: dict(type='BN')\n in_channels (int): Number of channels of the feature from ResNetTop.\n Default: 64.\n \"\"\"\n\n def __init__(self,\n has_skip=False,\n gen_skip=False,\n gen_cross_conv=False,\n unit_channels=256,\n num_units=4,\n num_blocks=[2, 2, 2, 2],\n norm_cfg=dict(type='BN'),\n in_channels=64):\n # Protect mutable default arguments\n norm_cfg = cp.deepcopy(norm_cfg)\n num_blocks = cp.deepcopy(num_blocks)\n super().__init__()\n assert len(num_blocks) == num_units\n self.has_skip = has_skip\n self.gen_skip = gen_skip\n self.gen_cross_conv = gen_cross_conv\n self.num_units = num_units\n self.unit_channels = unit_channels\n self.num_blocks = num_blocks\n self.norm_cfg = norm_cfg\n\n self.downsample = DownsampleModule(Bottleneck, num_blocks, num_units,\n has_skip, norm_cfg, in_channels)\n self.upsample = UpsampleModule(unit_channels, num_units, gen_skip,\n gen_cross_conv, norm_cfg, in_channels)\n\n def forward(self, x, skip1, skip2):\n mid = self.downsample(x, skip1, skip2)\n out, skip1, skip2, cross_conv = self.upsample(mid)\n\n return out, skip1, skip2, cross_conv\n\n\nclass ResNetTop(nn.Module):\n \"\"\"ResNet top for MSPN.\n\n Args:\n norm_cfg (dict): dictionary to construct and config norm layer.\n Default: dict(type='BN')\n channels (int): Number of channels of the feature output by ResNetTop.\n \"\"\"\n\n def __init__(self, norm_cfg=dict(type='BN'), channels=64):\n # Protect mutable default arguments\n norm_cfg = cp.deepcopy(norm_cfg)\n super().__init__()\n self.top = nn.Sequential(\n ConvModule(\n 3,\n channels,\n kernel_size=7,\n stride=2,\n padding=3,\n norm_cfg=norm_cfg,\n inplace=True), MaxPool2d(kernel_size=3, stride=2, padding=1))\n\n def forward(self, img):\n return self.top(img)\n\n\n@BACKBONES.register_module()\nclass MSPN(BaseBackbone):\n \"\"\"MSPN backbone. Paper ref: Li et al. \"Rethinking on Multi-Stage Networks\n for Human Pose Estimation\" (CVPR 2020).\n\n Args:\n unit_channels (int): Number of Channels in an upsample unit.\n Default: 256\n num_stages (int): Number of stages in a multi-stage MSPN. Default: 4\n num_units (int): Number of downsample/upsample units in a single-stage\n network. Default: 4\n Note: Make sure num_units == len(self.num_blocks)\n num_blocks (list): Number of bottlenecks in each\n downsample unit. Default: [2, 2, 2, 2]\n norm_cfg (dict): dictionary to construct and config norm layer.\n Default: dict(type='BN')\n res_top_channels (int): Number of channels of feature from ResNetTop.\n Default: 64.\n\n Example:\n >>> from mmpose.models import MSPN\n >>> import torch\n >>> self = MSPN(num_stages=2,num_units=2,num_blocks=[2,2])\n >>> self.eval()\n >>> inputs = torch.rand(1, 3, 511, 511)\n >>> level_outputs = self.forward(inputs)\n >>> for level_output in level_outputs:\n ... for feature in level_output:\n ... print(tuple(feature.shape))\n ...\n (1, 256, 64, 64)\n (1, 256, 128, 128)\n (1, 256, 64, 64)\n (1, 256, 128, 128)\n \"\"\"\n\n def __init__(self,\n unit_channels=256,\n num_stages=4,\n num_units=4,\n num_blocks=[2, 2, 2, 2],\n norm_cfg=dict(type='BN'),\n res_top_channels=64):\n # Protect mutable default arguments\n norm_cfg = cp.deepcopy(norm_cfg)\n num_blocks = cp.deepcopy(num_blocks)\n super().__init__()\n self.unit_channels = unit_channels\n self.num_stages = num_stages\n self.num_units = num_units\n self.num_blocks = num_blocks\n self.norm_cfg = norm_cfg\n\n assert self.num_stages > 0\n assert self.num_units > 1\n assert self.num_units == len(self.num_blocks)\n self.top = ResNetTop(norm_cfg=norm_cfg)\n self.multi_stage_mspn = nn.ModuleList([])\n for i in range(self.num_stages):\n if i == 0:\n has_skip = False\n else:\n has_skip = True\n if i != self.num_stages - 1:\n gen_skip = True\n gen_cross_conv = True\n else:\n gen_skip = False\n gen_cross_conv = False\n self.multi_stage_mspn.append(\n SingleStageNetwork(has_skip, gen_skip, gen_cross_conv,\n unit_channels, num_units, num_blocks,\n norm_cfg, res_top_channels))\n\n def forward(self, x):\n \"\"\"Model forward function.\"\"\"\n out_feats = []\n skip1 = None\n skip2 = None\n x = self.top(x)\n for i in range(self.num_stages):\n out, skip1, skip2, x = self.multi_stage_mspn[i](x, skip1, skip2)\n out_feats.append(out)\n\n return out_feats\n\n def init_weights(self, pretrained=None):\n \"\"\"Initialize model weights.\"\"\"\n if isinstance(pretrained, str):\n logger = get_root_logger()\n state_dict_tmp = get_state_dict(pretrained)\n state_dict = OrderedDict()\n state_dict['top'] = OrderedDict()\n state_dict['bottlenecks'] = OrderedDict()\n for k, v in state_dict_tmp.items():\n if k.startswith('layer'):\n if 'downsample.0' in k:\n state_dict['bottlenecks'][k.replace(\n 'downsample.0', 'downsample.conv')] = v\n elif 'downsample.1' in k:\n state_dict['bottlenecks'][k.replace(\n 'downsample.1', 'downsample.bn')] = v\n else:\n state_dict['bottlenecks'][k] = v\n elif k.startswith('conv1'):\n state_dict['top'][k.replace('conv1', 'top.0.conv')] = v\n elif k.startswith('bn1'):\n state_dict['top'][k.replace('bn1', 'top.0.bn')] = v\n\n load_state_dict(\n self.top, state_dict['top'], strict=False, logger=logger)\n for i in range(self.num_stages):\n load_state_dict(\n self.multi_stage_mspn[i].downsample,\n state_dict['bottlenecks'],\n strict=False,\n logger=logger)\n else:\n for m in self.multi_stage_mspn.modules():\n if isinstance(m, nn.Conv2d):\n kaiming_init(m)\n elif isinstance(m, nn.BatchNorm2d):\n constant_init(m, 1)\n elif isinstance(m, nn.Linear):\n normal_init(m, std=0.01)\n\n for m in self.top.modules():\n if isinstance(m, nn.Conv2d):\n kaiming_init(m)\n" }, { "path": "eval/GVHMR/hmr4d/utils/preproc/vitpose_pytorch/src/vitpose_infer/builder/backbones/regnet.py", "content": "# Copyright (c) OpenMMLab. All rights reserved.\nimport copy\n\nimport numpy as np\nimport torch.nn as nn\nfrom mmcv.cnn import build_conv_layer, build_norm_layer\n\nfrom ..builder import BACKBONES\nfrom .resnet import ResNet\nfrom .resnext import Bottleneck\n\n\n@BACKBONES.register_module()\nclass RegNet(ResNet):\n \"\"\"RegNet backbone.\n\n More details can be found in `paper `__ .\n\n Args:\n arch (dict): The parameter of RegNets.\n - w0 (int): initial width\n - wa (float): slope of width\n - wm (float): quantization parameter to quantize the width\n - depth (int): depth of the backbone\n - group_w (int): width of group\n - bot_mul (float): bottleneck ratio, i.e. expansion of bottleneck.\n strides (Sequence[int]): Strides of the first block of each stage.\n base_channels (int): Base channels after stem layer.\n in_channels (int): Number of input image channels. Default: 3.\n dilations (Sequence[int]): Dilation of each stage.\n out_indices (Sequence[int]): Output from which stages.\n style (str): `pytorch` or `caffe`. If set to \"pytorch\", the stride-two\n layer is the 3x3 conv layer, otherwise the stride-two layer is\n the first 1x1 conv layer. Default: \"pytorch\".\n frozen_stages (int): Stages to be frozen (all param fixed). -1 means\n not freezing any parameters. Default: -1.\n norm_cfg (dict): dictionary to construct and config norm layer.\n Default: dict(type='BN', requires_grad=True).\n norm_eval (bool): Whether to set norm layers to eval mode, namely,\n freeze running stats (mean and var). Note: Effect on Batch Norm\n and its variants only. Default: False.\n with_cp (bool): Use checkpoint or not. Using checkpoint will save some\n memory while slowing down the training speed. Default: False.\n zero_init_residual (bool): whether to use zero init for last norm layer\n in resblocks to let them behave as identity. Default: True.\n\n Example:\n >>> from mmpose.models import RegNet\n >>> import torch\n >>> self = RegNet(\n arch=dict(\n w0=88,\n wa=26.31,\n wm=2.25,\n group_w=48,\n depth=25,\n bot_mul=1.0),\n out_indices=(0, 1, 2, 3))\n >>> self.eval()\n >>> inputs = torch.rand(1, 3, 32, 32)\n >>> level_outputs = self.forward(inputs)\n >>> for level_out in level_outputs:\n ... print(tuple(level_out.shape))\n (1, 96, 8, 8)\n (1, 192, 4, 4)\n (1, 432, 2, 2)\n (1, 1008, 1, 1)\n \"\"\"\n arch_settings = {\n 'regnetx_400mf':\n dict(w0=24, wa=24.48, wm=2.54, group_w=16, depth=22, bot_mul=1.0),\n 'regnetx_800mf':\n dict(w0=56, wa=35.73, wm=2.28, group_w=16, depth=16, bot_mul=1.0),\n 'regnetx_1.6gf':\n dict(w0=80, wa=34.01, wm=2.25, group_w=24, depth=18, bot_mul=1.0),\n 'regnetx_3.2gf':\n dict(w0=88, wa=26.31, wm=2.25, group_w=48, depth=25, bot_mul=1.0),\n 'regnetx_4.0gf':\n dict(w0=96, wa=38.65, wm=2.43, group_w=40, depth=23, bot_mul=1.0),\n 'regnetx_6.4gf':\n dict(w0=184, wa=60.83, wm=2.07, group_w=56, depth=17, bot_mul=1.0),\n 'regnetx_8.0gf':\n dict(w0=80, wa=49.56, wm=2.88, group_w=120, depth=23, bot_mul=1.0),\n 'regnetx_12gf':\n dict(w0=168, wa=73.36, wm=2.37, group_w=112, depth=19, bot_mul=1.0),\n }\n\n def __init__(self,\n arch,\n in_channels=3,\n stem_channels=32,\n base_channels=32,\n strides=(2, 2, 2, 2),\n dilations=(1, 1, 1, 1),\n out_indices=(3, ),\n style='pytorch',\n deep_stem=False,\n avg_down=False,\n frozen_stages=-1,\n conv_cfg=None,\n norm_cfg=dict(type='BN', requires_grad=True),\n norm_eval=False,\n with_cp=False,\n zero_init_residual=True):\n # Protect mutable default arguments\n norm_cfg = copy.deepcopy(norm_cfg)\n super(ResNet, self).__init__()\n\n # Generate RegNet parameters first\n if isinstance(arch, str):\n assert arch in self.arch_settings, \\\n f'\"arch\": \"{arch}\" is not one of the' \\\n ' arch_settings'\n arch = self.arch_settings[arch]\n elif not isinstance(arch, dict):\n raise TypeError('Expect \"arch\" to be either a string '\n f'or a dict, got {type(arch)}')\n\n widths, num_stages = self.generate_regnet(\n arch['w0'],\n arch['wa'],\n arch['wm'],\n arch['depth'],\n )\n # Convert to per stage format\n stage_widths, stage_blocks = self.get_stages_from_blocks(widths)\n # Generate group widths and bot muls\n group_widths = [arch['group_w'] for _ in range(num_stages)]\n self.bottleneck_ratio = [arch['bot_mul'] for _ in range(num_stages)]\n # Adjust the compatibility of stage_widths and group_widths\n stage_widths, group_widths = self.adjust_width_group(\n stage_widths, self.bottleneck_ratio, group_widths)\n\n # Group params by stage\n self.stage_widths = stage_widths\n self.group_widths = group_widths\n self.depth = sum(stage_blocks)\n self.stem_channels = stem_channels\n self.base_channels = base_channels\n self.num_stages = num_stages\n assert 1 <= num_stages <= 4\n self.strides = strides\n self.dilations = dilations\n assert len(strides) == len(dilations) == num_stages\n self.out_indices = out_indices\n assert max(out_indices) < num_stages\n self.style = style\n self.deep_stem = deep_stem\n if self.deep_stem:\n raise NotImplementedError(\n 'deep_stem has not been implemented for RegNet')\n self.avg_down = avg_down\n self.frozen_stages = frozen_stages\n self.conv_cfg = conv_cfg\n self.norm_cfg = norm_cfg\n self.with_cp = with_cp\n self.norm_eval = norm_eval\n self.zero_init_residual = zero_init_residual\n self.stage_blocks = stage_blocks[:num_stages]\n\n self._make_stem_layer(in_channels, stem_channels)\n\n _in_channels = stem_channels\n self.res_layers = []\n for i, num_blocks in enumerate(self.stage_blocks):\n stride = self.strides[i]\n dilation = self.dilations[i]\n group_width = self.group_widths[i]\n width = int(round(self.stage_widths[i] * self.bottleneck_ratio[i]))\n stage_groups = width // group_width\n\n res_layer = self.make_res_layer(\n block=Bottleneck,\n num_blocks=num_blocks,\n in_channels=_in_channels,\n out_channels=self.stage_widths[i],\n expansion=1,\n stride=stride,\n dilation=dilation,\n style=self.style,\n avg_down=self.avg_down,\n with_cp=self.with_cp,\n conv_cfg=self.conv_cfg,\n norm_cfg=self.norm_cfg,\n base_channels=self.stage_widths[i],\n groups=stage_groups,\n width_per_group=group_width)\n _in_channels = self.stage_widths[i]\n layer_name = f'layer{i + 1}'\n self.add_module(layer_name, res_layer)\n self.res_layers.append(layer_name)\n\n self._freeze_stages()\n\n self.feat_dim = stage_widths[-1]\n\n def _make_stem_layer(self, in_channels, base_channels):\n self.conv1 = build_conv_layer(\n self.conv_cfg,\n in_channels,\n base_channels,\n kernel_size=3,\n stride=2,\n padding=1,\n bias=False)\n self.norm1_name, norm1 = build_norm_layer(\n self.norm_cfg, base_channels, postfix=1)\n self.add_module(self.norm1_name, norm1)\n self.relu = nn.ReLU(inplace=True)\n\n @staticmethod\n def generate_regnet(initial_width,\n width_slope,\n width_parameter,\n depth,\n divisor=8):\n \"\"\"Generates per block width from RegNet parameters.\n\n Args:\n initial_width ([int]): Initial width of the backbone\n width_slope ([float]): Slope of the quantized linear function\n width_parameter ([int]): Parameter used to quantize the width.\n depth ([int]): Depth of the backbone.\n divisor (int, optional): The divisor of channels. Defaults to 8.\n\n Returns:\n list, int: return a list of widths of each stage and the number of\n stages\n \"\"\"\n assert width_slope >= 0\n assert initial_width > 0\n assert width_parameter > 1\n assert initial_width % divisor == 0\n widths_cont = np.arange(depth) * width_slope + initial_width\n ks = np.round(\n np.log(widths_cont / initial_width) / np.log(width_parameter))\n widths = initial_width * np.power(width_parameter, ks)\n widths = np.round(np.divide(widths, divisor)) * divisor\n num_stages = len(np.unique(widths))\n widths, widths_cont = widths.astype(int).tolist(), widths_cont.tolist()\n return widths, num_stages\n\n @staticmethod\n def quantize_float(number, divisor):\n \"\"\"Converts a float to closest non-zero int divisible by divior.\n\n Args:\n number (int): Original number to be quantized.\n divisor (int): Divisor used to quantize the number.\n\n Returns:\n int: quantized number that is divisible by devisor.\n \"\"\"\n return int(round(number / divisor) * divisor)\n\n def adjust_width_group(self, widths, bottleneck_ratio, groups):\n \"\"\"Adjusts the compatibility of widths and groups.\n\n Args:\n widths (list[int]): Width of each stage.\n bottleneck_ratio (float): Bottleneck ratio.\n groups (int): number of groups in each stage\n\n Returns:\n tuple(list): The adjusted widths and groups of each stage.\n \"\"\"\n bottleneck_width = [\n int(w * b) for w, b in zip(widths, bottleneck_ratio)\n ]\n groups = [min(g, w_bot) for g, w_bot in zip(groups, bottleneck_width)]\n bottleneck_width = [\n self.quantize_float(w_bot, g)\n for w_bot, g in zip(bottleneck_width, groups)\n ]\n widths = [\n int(w_bot / b)\n for w_bot, b in zip(bottleneck_width, bottleneck_ratio)\n ]\n return widths, groups\n\n def get_stages_from_blocks(self, widths):\n \"\"\"Gets widths/stage_blocks of network at each stage.\n\n Args:\n widths (list[int]): Width in each stage.\n\n Returns:\n tuple(list): width and depth of each stage\n \"\"\"\n width_diff = [\n width != width_prev\n for width, width_prev in zip(widths + [0], [0] + widths)\n ]\n stage_widths = [\n width for width, diff in zip(widths, width_diff[:-1]) if diff\n ]\n stage_blocks = np.diff([\n depth for depth, diff in zip(range(len(width_diff)), width_diff)\n if diff\n ]).tolist()\n return stage_widths, stage_blocks\n\n def forward(self, x):\n x = self.conv1(x)\n x = self.norm1(x)\n x = self.relu(x)\n\n outs = []\n for i, layer_name in enumerate(self.res_layers):\n res_layer = getattr(self, layer_name)\n x = res_layer(x)\n if i in self.out_indices:\n outs.append(x)\n\n if len(outs) == 1:\n return outs[0]\n return tuple(outs)\n" }, { "path": "eval/GVHMR/hmr4d/utils/preproc/vitpose_pytorch/src/vitpose_infer/builder/backbones/resnest.py", "content": "# Copyright (c) OpenMMLab. All rights reserved.\nimport torch\nimport torch.nn as nn\nimport torch.nn.functional as F\nimport torch.utils.checkpoint as cp\nfrom mmcv.cnn import build_conv_layer, build_norm_layer\n\nfrom ..builder import BACKBONES\nfrom .resnet import Bottleneck as _Bottleneck\nfrom .resnet import ResLayer, ResNetV1d\n\n\nclass RSoftmax(nn.Module):\n \"\"\"Radix Softmax module in ``SplitAttentionConv2d``.\n\n Args:\n radix (int): Radix of input.\n groups (int): Groups of input.\n \"\"\"\n\n def __init__(self, radix, groups):\n super().__init__()\n self.radix = radix\n self.groups = groups\n\n def forward(self, x):\n batch = x.size(0)\n if self.radix > 1:\n x = x.view(batch, self.groups, self.radix, -1).transpose(1, 2)\n x = F.softmax(x, dim=1)\n x = x.reshape(batch, -1)\n else:\n x = torch.sigmoid(x)\n return x\n\n\nclass SplitAttentionConv2d(nn.Module):\n \"\"\"Split-Attention Conv2d.\n\n Args:\n in_channels (int): Same as nn.Conv2d.\n out_channels (int): Same as nn.Conv2d.\n kernel_size (int | tuple[int]): Same as nn.Conv2d.\n stride (int | tuple[int]): Same as nn.Conv2d.\n padding (int | tuple[int]): Same as nn.Conv2d.\n dilation (int | tuple[int]): Same as nn.Conv2d.\n groups (int): Same as nn.Conv2d.\n radix (int): Radix of SpltAtConv2d. Default: 2\n reduction_factor (int): Reduction factor of SplitAttentionConv2d.\n Default: 4.\n conv_cfg (dict): Config dict for convolution layer. Default: None,\n which means using conv2d.\n norm_cfg (dict): Config dict for normalization layer. Default: None.\n \"\"\"\n\n def __init__(self,\n in_channels,\n channels,\n kernel_size,\n stride=1,\n padding=0,\n dilation=1,\n groups=1,\n radix=2,\n reduction_factor=4,\n conv_cfg=None,\n norm_cfg=dict(type='BN')):\n super().__init__()\n inter_channels = max(in_channels * radix // reduction_factor, 32)\n self.radix = radix\n self.groups = groups\n self.channels = channels\n self.conv = build_conv_layer(\n conv_cfg,\n in_channels,\n channels * radix,\n kernel_size,\n stride=stride,\n padding=padding,\n dilation=dilation,\n groups=groups * radix,\n bias=False)\n self.norm0_name, norm0 = build_norm_layer(\n norm_cfg, channels * radix, postfix=0)\n self.add_module(self.norm0_name, norm0)\n self.relu = nn.ReLU(inplace=True)\n self.fc1 = build_conv_layer(\n None, channels, inter_channels, 1, groups=self.groups)\n self.norm1_name, norm1 = build_norm_layer(\n norm_cfg, inter_channels, postfix=1)\n self.add_module(self.norm1_name, norm1)\n self.fc2 = build_conv_layer(\n None, inter_channels, channels * radix, 1, groups=self.groups)\n self.rsoftmax = RSoftmax(radix, groups)\n\n @property\n def norm0(self):\n return getattr(self, self.norm0_name)\n\n @property\n def norm1(self):\n return getattr(self, self.norm1_name)\n\n def forward(self, x):\n x = self.conv(x)\n x = self.norm0(x)\n x = self.relu(x)\n\n batch, rchannel = x.shape[:2]\n if self.radix > 1:\n splits = x.view(batch, self.radix, -1, *x.shape[2:])\n gap = splits.sum(dim=1)\n else:\n gap = x\n gap = F.adaptive_avg_pool2d(gap, 1)\n gap = self.fc1(gap)\n\n gap = self.norm1(gap)\n gap = self.relu(gap)\n\n atten = self.fc2(gap)\n atten = self.rsoftmax(atten).view(batch, -1, 1, 1)\n\n if self.radix > 1:\n attens = atten.view(batch, self.radix, -1, *atten.shape[2:])\n out = torch.sum(attens * splits, dim=1)\n else:\n out = atten * x\n return out.contiguous()\n\n\nclass Bottleneck(_Bottleneck):\n \"\"\"Bottleneck block for ResNeSt.\n\n Args:\n in_channels (int): Input channels of this block.\n out_channels (int): Output channels of this block.\n groups (int): Groups of conv2.\n width_per_group (int): Width per group of conv2. 64x4d indicates\n ``groups=64, width_per_group=4`` and 32x8d indicates\n ``groups=32, width_per_group=8``.\n radix (int): Radix of SpltAtConv2d. Default: 2\n reduction_factor (int): Reduction factor of SplitAttentionConv2d.\n Default: 4.\n avg_down_stride (bool): Whether to use average pool for stride in\n Bottleneck. Default: True.\n stride (int): stride of the block. Default: 1\n dilation (int): dilation of convolution. Default: 1\n downsample (nn.Module): downsample operation on identity branch.\n Default: None\n style (str): `pytorch` or `caffe`. If set to \"pytorch\", the stride-two\n layer is the 3x3 conv layer, otherwise the stride-two layer is\n the first 1x1 conv layer.\n conv_cfg (dict): dictionary to construct and config conv layer.\n Default: None\n norm_cfg (dict): dictionary to construct and config norm layer.\n Default: dict(type='BN')\n with_cp (bool): Use checkpoint or not. Using checkpoint will save some\n memory while slowing down the training speed.\n \"\"\"\n\n def __init__(self,\n in_channels,\n out_channels,\n groups=1,\n width_per_group=4,\n base_channels=64,\n radix=2,\n reduction_factor=4,\n avg_down_stride=True,\n **kwargs):\n super().__init__(in_channels, out_channels, **kwargs)\n\n self.groups = groups\n self.width_per_group = width_per_group\n\n # For ResNet bottleneck, middle channels are determined by expansion\n # and out_channels, but for ResNeXt bottleneck, it is determined by\n # groups and width_per_group and the stage it is located in.\n if groups != 1:\n assert self.mid_channels % base_channels == 0\n self.mid_channels = (\n groups * width_per_group * self.mid_channels // base_channels)\n\n self.avg_down_stride = avg_down_stride and self.conv2_stride > 1\n\n self.norm1_name, norm1 = build_norm_layer(\n self.norm_cfg, self.mid_channels, postfix=1)\n self.norm3_name, norm3 = build_norm_layer(\n self.norm_cfg, self.out_channels, postfix=3)\n\n self.conv1 = build_conv_layer(\n self.conv_cfg,\n self.in_channels,\n self.mid_channels,\n kernel_size=1,\n stride=self.conv1_stride,\n bias=False)\n self.add_module(self.norm1_name, norm1)\n self.conv2 = SplitAttentionConv2d(\n self.mid_channels,\n self.mid_channels,\n kernel_size=3,\n stride=1 if self.avg_down_stride else self.conv2_stride,\n padding=self.dilation,\n dilation=self.dilation,\n groups=groups,\n radix=radix,\n reduction_factor=reduction_factor,\n conv_cfg=self.conv_cfg,\n norm_cfg=self.norm_cfg)\n delattr(self, self.norm2_name)\n\n if self.avg_down_stride:\n self.avd_layer = nn.AvgPool2d(3, self.conv2_stride, padding=1)\n\n self.conv3 = build_conv_layer(\n self.conv_cfg,\n self.mid_channels,\n self.out_channels,\n kernel_size=1,\n bias=False)\n self.add_module(self.norm3_name, norm3)\n\n def forward(self, x):\n\n def _inner_forward(x):\n identity = x\n\n out = self.conv1(x)\n out = self.norm1(out)\n out = self.relu(out)\n\n out = self.conv2(out)\n\n if self.avg_down_stride:\n out = self.avd_layer(out)\n\n out = self.conv3(out)\n out = self.norm3(out)\n\n if self.downsample is not None:\n identity = self.downsample(x)\n\n out += identity\n\n return out\n\n if self.with_cp and x.requires_grad:\n out = cp.checkpoint(_inner_forward, x)\n else:\n out = _inner_forward(x)\n\n out = self.relu(out)\n\n return out\n\n\n@BACKBONES.register_module()\nclass ResNeSt(ResNetV1d):\n \"\"\"ResNeSt backbone.\n\n Please refer to the `paper `__\n for details.\n\n Args:\n depth (int): Network depth, from {50, 101, 152, 200}.\n groups (int): Groups of conv2 in Bottleneck. Default: 32.\n width_per_group (int): Width per group of conv2 in Bottleneck.\n Default: 4.\n radix (int): Radix of SpltAtConv2d. Default: 2\n reduction_factor (int): Reduction factor of SplitAttentionConv2d.\n Default: 4.\n avg_down_stride (bool): Whether to use average pool for stride in\n Bottleneck. Default: True.\n in_channels (int): Number of input image channels. Default: 3.\n stem_channels (int): Output channels of the stem layer. Default: 64.\n num_stages (int): Stages of the network. Default: 4.\n strides (Sequence[int]): Strides of the first block of each stage.\n Default: ``(1, 2, 2, 2)``.\n dilations (Sequence[int]): Dilation of each stage.\n Default: ``(1, 1, 1, 1)``.\n out_indices (Sequence[int]): Output from which stages. If only one\n stage is specified, a single tensor (feature map) is returned,\n otherwise multiple stages are specified, a tuple of tensors will\n be returned. Default: ``(3, )``.\n style (str): `pytorch` or `caffe`. If set to \"pytorch\", the stride-two\n layer is the 3x3 conv layer, otherwise the stride-two layer is\n the first 1x1 conv layer.\n deep_stem (bool): Replace 7x7 conv in input stem with 3 3x3 conv.\n Default: False.\n avg_down (bool): Use AvgPool instead of stride conv when\n downsampling in the bottleneck. Default: False.\n frozen_stages (int): Stages to be frozen (stop grad and set eval mode).\n -1 means not freezing any parameters. Default: -1.\n conv_cfg (dict | None): The config dict for conv layers. Default: None.\n norm_cfg (dict): The config dict for norm layers.\n norm_eval (bool): Whether to set norm layers to eval mode, namely,\n freeze running stats (mean and var). Note: Effect on Batch Norm\n and its variants only. Default: False.\n with_cp (bool): Use checkpoint or not. Using checkpoint will save some\n memory while slowing down the training speed. Default: False.\n zero_init_residual (bool): Whether to use zero init for last norm layer\n in resblocks to let them behave as identity. Default: True.\n \"\"\"\n\n arch_settings = {\n 50: (Bottleneck, (3, 4, 6, 3)),\n 101: (Bottleneck, (3, 4, 23, 3)),\n 152: (Bottleneck, (3, 8, 36, 3)),\n 200: (Bottleneck, (3, 24, 36, 3)),\n 269: (Bottleneck, (3, 30, 48, 8))\n }\n\n def __init__(self,\n depth,\n groups=1,\n width_per_group=4,\n radix=2,\n reduction_factor=4,\n avg_down_stride=True,\n **kwargs):\n self.groups = groups\n self.width_per_group = width_per_group\n self.radix = radix\n self.reduction_factor = reduction_factor\n self.avg_down_stride = avg_down_stride\n super().__init__(depth=depth, **kwargs)\n\n def make_res_layer(self, **kwargs):\n return ResLayer(\n groups=self.groups,\n width_per_group=self.width_per_group,\n base_channels=self.base_channels,\n radix=self.radix,\n reduction_factor=self.reduction_factor,\n avg_down_stride=self.avg_down_stride,\n **kwargs)\n" }, { "path": "eval/GVHMR/hmr4d/utils/preproc/vitpose_pytorch/src/vitpose_infer/builder/backbones/resnext.py", "content": "# Copyright (c) OpenMMLab. All rights reserved.\nfrom mmcv.cnn import build_conv_layer, build_norm_layer\n\nfrom ..builder import BACKBONES\nfrom .resnet import Bottleneck as _Bottleneck\nfrom .resnet import ResLayer, ResNet\n\n\nclass Bottleneck(_Bottleneck):\n \"\"\"Bottleneck block for ResNeXt.\n\n Args:\n in_channels (int): Input channels of this block.\n out_channels (int): Output channels of this block.\n groups (int): Groups of conv2.\n width_per_group (int): Width per group of conv2. 64x4d indicates\n ``groups=64, width_per_group=4`` and 32x8d indicates\n ``groups=32, width_per_group=8``.\n stride (int): stride of the block. Default: 1\n dilation (int): dilation of convolution. Default: 1\n downsample (nn.Module): downsample operation on identity branch.\n Default: None\n style (str): `pytorch` or `caffe`. If set to \"pytorch\", the stride-two\n layer is the 3x3 conv layer, otherwise the stride-two layer is\n the first 1x1 conv layer.\n conv_cfg (dict): dictionary to construct and config conv layer.\n Default: None\n norm_cfg (dict): dictionary to construct and config norm layer.\n Default: dict(type='BN')\n with_cp (bool): Use checkpoint or not. Using checkpoint will save some\n memory while slowing down the training speed.\n \"\"\"\n\n def __init__(self,\n in_channels,\n out_channels,\n base_channels=64,\n groups=32,\n width_per_group=4,\n **kwargs):\n super().__init__(in_channels, out_channels, **kwargs)\n self.groups = groups\n self.width_per_group = width_per_group\n\n # For ResNet bottleneck, middle channels are determined by expansion\n # and out_channels, but for ResNeXt bottleneck, it is determined by\n # groups and width_per_group and the stage it is located in.\n if groups != 1:\n assert self.mid_channels % base_channels == 0\n self.mid_channels = (\n groups * width_per_group * self.mid_channels // base_channels)\n\n self.norm1_name, norm1 = build_norm_layer(\n self.norm_cfg, self.mid_channels, postfix=1)\n self.norm2_name, norm2 = build_norm_layer(\n self.norm_cfg, self.mid_channels, postfix=2)\n self.norm3_name, norm3 = build_norm_layer(\n self.norm_cfg, self.out_channels, postfix=3)\n\n self.conv1 = build_conv_layer(\n self.conv_cfg,\n self.in_channels,\n self.mid_channels,\n kernel_size=1,\n stride=self.conv1_stride,\n bias=False)\n self.add_module(self.norm1_name, norm1)\n self.conv2 = build_conv_layer(\n self.conv_cfg,\n self.mid_channels,\n self.mid_channels,\n kernel_size=3,\n stride=self.conv2_stride,\n padding=self.dilation,\n dilation=self.dilation,\n groups=groups,\n bias=False)\n\n self.add_module(self.norm2_name, norm2)\n self.conv3 = build_conv_layer(\n self.conv_cfg,\n self.mid_channels,\n self.out_channels,\n kernel_size=1,\n bias=False)\n self.add_module(self.norm3_name, norm3)\n\n\n@BACKBONES.register_module()\nclass ResNeXt(ResNet):\n \"\"\"ResNeXt backbone.\n\n Please refer to the `paper `__ for\n details.\n\n Args:\n depth (int): Network depth, from {50, 101, 152}.\n groups (int): Groups of conv2 in Bottleneck. Default: 32.\n width_per_group (int): Width per group of conv2 in Bottleneck.\n Default: 4.\n in_channels (int): Number of input image channels. Default: 3.\n stem_channels (int): Output channels of the stem layer. Default: 64.\n num_stages (int): Stages of the network. Default: 4.\n strides (Sequence[int]): Strides of the first block of each stage.\n Default: ``(1, 2, 2, 2)``.\n dilations (Sequence[int]): Dilation of each stage.\n Default: ``(1, 1, 1, 1)``.\n out_indices (Sequence[int]): Output from which stages. If only one\n stage is specified, a single tensor (feature map) is returned,\n otherwise multiple stages are specified, a tuple of tensors will\n be returned. Default: ``(3, )``.\n style (str): `pytorch` or `caffe`. If set to \"pytorch\", the stride-two\n layer is the 3x3 conv layer, otherwise the stride-two layer is\n the first 1x1 conv layer.\n deep_stem (bool): Replace 7x7 conv in input stem with 3 3x3 conv.\n Default: False.\n avg_down (bool): Use AvgPool instead of stride conv when\n downsampling in the bottleneck. Default: False.\n frozen_stages (int): Stages to be frozen (stop grad and set eval mode).\n -1 means not freezing any parameters. Default: -1.\n conv_cfg (dict | None): The config dict for conv layers. Default: None.\n norm_cfg (dict): The config dict for norm layers.\n norm_eval (bool): Whether to set norm layers to eval mode, namely,\n freeze running stats (mean and var). Note: Effect on Batch Norm\n and its variants only. Default: False.\n with_cp (bool): Use checkpoint or not. Using checkpoint will save some\n memory while slowing down the training speed. Default: False.\n zero_init_residual (bool): Whether to use zero init for last norm layer\n in resblocks to let them behave as identity. Default: True.\n\n Example:\n >>> from mmpose.models import ResNeXt\n >>> import torch\n >>> self = ResNeXt(depth=50, out_indices=(0, 1, 2, 3))\n >>> self.eval()\n >>> inputs = torch.rand(1, 3, 32, 32)\n >>> level_outputs = self.forward(inputs)\n >>> for level_out in level_outputs:\n ... print(tuple(level_out.shape))\n (1, 256, 8, 8)\n (1, 512, 4, 4)\n (1, 1024, 2, 2)\n (1, 2048, 1, 1)\n \"\"\"\n\n arch_settings = {\n 50: (Bottleneck, (3, 4, 6, 3)),\n 101: (Bottleneck, (3, 4, 23, 3)),\n 152: (Bottleneck, (3, 8, 36, 3))\n }\n\n def __init__(self, depth, groups=32, width_per_group=4, **kwargs):\n self.groups = groups\n self.width_per_group = width_per_group\n super().__init__(depth, **kwargs)\n\n def make_res_layer(self, **kwargs):\n return ResLayer(\n groups=self.groups,\n width_per_group=self.width_per_group,\n base_channels=self.base_channels,\n **kwargs)\n" }, { "path": "eval/GVHMR/hmr4d/utils/preproc/vitpose_pytorch/src/vitpose_infer/builder/backbones/rsn.py", "content": "# Copyright (c) OpenMMLab. All rights reserved.\nimport copy as cp\n\nimport torch\nimport torch.nn as nn\nimport torch.nn.functional as F\nfrom mmcv.cnn import (ConvModule, MaxPool2d, constant_init, kaiming_init,\n normal_init)\n\nfrom ..builder import BACKBONES\nfrom .base_backbone import BaseBackbone\n\n\nclass RSB(nn.Module):\n \"\"\"Residual Steps block for RSN. Paper ref: Cai et al. \"Learning Delicate\n Local Representations for Multi-Person Pose Estimation\" (ECCV 2020).\n\n Args:\n in_channels (int): Input channels of this block.\n out_channels (int): Output channels of this block.\n num_steps (int): Numbers of steps in RSB\n stride (int): stride of the block. Default: 1\n downsample (nn.Module): downsample operation on identity branch.\n Default: None.\n norm_cfg (dict): dictionary to construct and config norm layer.\n Default: dict(type='BN')\n expand_times (int): Times by which the in_channels are expanded.\n Default:26.\n res_top_channels (int): Number of channels of feature output by\n ResNet_top. Default:64.\n \"\"\"\n\n expansion = 1\n\n def __init__(self,\n in_channels,\n out_channels,\n num_steps=4,\n stride=1,\n downsample=None,\n with_cp=False,\n norm_cfg=dict(type='BN'),\n expand_times=26,\n res_top_channels=64):\n # Protect mutable default arguments\n norm_cfg = cp.deepcopy(norm_cfg)\n super().__init__()\n assert num_steps > 1\n self.in_channels = in_channels\n self.branch_channels = self.in_channels * expand_times\n self.branch_channels //= res_top_channels\n self.out_channels = out_channels\n self.stride = stride\n self.downsample = downsample\n self.with_cp = with_cp\n self.norm_cfg = norm_cfg\n self.num_steps = num_steps\n self.conv_bn_relu1 = ConvModule(\n self.in_channels,\n self.num_steps * self.branch_channels,\n kernel_size=1,\n stride=self.stride,\n padding=0,\n norm_cfg=self.norm_cfg,\n inplace=False)\n for i in range(self.num_steps):\n for j in range(i + 1):\n module_name = f'conv_bn_relu2_{i + 1}_{j + 1}'\n self.add_module(\n module_name,\n ConvModule(\n self.branch_channels,\n self.branch_channels,\n kernel_size=3,\n stride=1,\n padding=1,\n norm_cfg=self.norm_cfg,\n inplace=False))\n self.conv_bn3 = ConvModule(\n self.num_steps * self.branch_channels,\n self.out_channels * self.expansion,\n kernel_size=1,\n stride=1,\n padding=0,\n act_cfg=None,\n norm_cfg=self.norm_cfg,\n inplace=False)\n self.relu = nn.ReLU(inplace=False)\n\n def forward(self, x):\n \"\"\"Forward function.\"\"\"\n\n identity = x\n x = self.conv_bn_relu1(x)\n spx = torch.split(x, self.branch_channels, 1)\n outputs = list()\n outs = list()\n for i in range(self.num_steps):\n outputs_i = list()\n outputs.append(outputs_i)\n for j in range(i + 1):\n if j == 0:\n inputs = spx[i]\n else:\n inputs = outputs[i][j - 1]\n if i > j:\n inputs = inputs + outputs[i - 1][j]\n module_name = f'conv_bn_relu2_{i + 1}_{j + 1}'\n module_i_j = getattr(self, module_name)\n outputs[i].append(module_i_j(inputs))\n\n outs.append(outputs[i][i])\n out = torch.cat(tuple(outs), 1)\n out = self.conv_bn3(out)\n\n if self.downsample is not None:\n identity = self.downsample(identity)\n out = out + identity\n\n out = self.relu(out)\n\n return out\n\n\nclass Downsample_module(nn.Module):\n \"\"\"Downsample module for RSN.\n\n Args:\n block (nn.Module): Downsample block.\n num_blocks (list): Number of blocks in each downsample unit.\n num_units (int): Numbers of downsample units. Default: 4\n has_skip (bool): Have skip connections from prior upsample\n module or not. Default:False\n num_steps (int): Number of steps in a block. Default:4\n norm_cfg (dict): dictionary to construct and config norm layer.\n Default: dict(type='BN')\n in_channels (int): Number of channels of the input feature to\n downsample module. Default: 64\n expand_times (int): Times by which the in_channels are expanded.\n Default:26.\n \"\"\"\n\n def __init__(self,\n block,\n num_blocks,\n num_steps=4,\n num_units=4,\n has_skip=False,\n norm_cfg=dict(type='BN'),\n in_channels=64,\n expand_times=26):\n # Protect mutable default arguments\n norm_cfg = cp.deepcopy(norm_cfg)\n super().__init__()\n self.has_skip = has_skip\n self.in_channels = in_channels\n assert len(num_blocks) == num_units\n self.num_blocks = num_blocks\n self.num_units = num_units\n self.num_steps = num_steps\n self.norm_cfg = norm_cfg\n self.layer1 = self._make_layer(\n block,\n in_channels,\n num_blocks[0],\n expand_times=expand_times,\n res_top_channels=in_channels)\n for i in range(1, num_units):\n module_name = f'layer{i + 1}'\n self.add_module(\n module_name,\n self._make_layer(\n block,\n in_channels * pow(2, i),\n num_blocks[i],\n stride=2,\n expand_times=expand_times,\n res_top_channels=in_channels))\n\n def _make_layer(self,\n block,\n out_channels,\n blocks,\n stride=1,\n expand_times=26,\n res_top_channels=64):\n downsample = None\n if stride != 1 or self.in_channels != out_channels * block.expansion:\n downsample = ConvModule(\n self.in_channels,\n out_channels * block.expansion,\n kernel_size=1,\n stride=stride,\n padding=0,\n norm_cfg=self.norm_cfg,\n act_cfg=None,\n inplace=True)\n\n units = list()\n units.append(\n block(\n self.in_channels,\n out_channels,\n num_steps=self.num_steps,\n stride=stride,\n downsample=downsample,\n norm_cfg=self.norm_cfg,\n expand_times=expand_times,\n res_top_channels=res_top_channels))\n self.in_channels = out_channels * block.expansion\n for _ in range(1, blocks):\n units.append(\n block(\n self.in_channels,\n out_channels,\n num_steps=self.num_steps,\n expand_times=expand_times,\n res_top_channels=res_top_channels))\n\n return nn.Sequential(*units)\n\n def forward(self, x, skip1, skip2):\n out = list()\n for i in range(self.num_units):\n module_name = f'layer{i + 1}'\n module_i = getattr(self, module_name)\n x = module_i(x)\n if self.has_skip:\n x = x + skip1[i] + skip2[i]\n out.append(x)\n out.reverse()\n\n return tuple(out)\n\n\nclass Upsample_unit(nn.Module):\n \"\"\"Upsample unit for upsample module.\n\n Args:\n ind (int): Indicates whether to interpolate (>0) and whether to\n generate feature map for the next hourglass-like module.\n num_units (int): Number of units that form a upsample module. Along\n with ind and gen_cross_conv, nm_units is used to decide whether\n to generate feature map for the next hourglass-like module.\n in_channels (int): Channel number of the skip-in feature maps from\n the corresponding downsample unit.\n unit_channels (int): Channel number in this unit. Default:256.\n gen_skip: (bool): Whether or not to generate skips for the posterior\n downsample module. Default:False\n gen_cross_conv (bool): Whether to generate feature map for the next\n hourglass-like module. Default:False\n norm_cfg (dict): dictionary to construct and config norm layer.\n Default: dict(type='BN')\n out_channels (in): Number of channels of feature output by upsample\n module. Must equal to in_channels of downsample module. Default:64\n \"\"\"\n\n def __init__(self,\n ind,\n num_units,\n in_channels,\n unit_channels=256,\n gen_skip=False,\n gen_cross_conv=False,\n norm_cfg=dict(type='BN'),\n out_channels=64):\n # Protect mutable default arguments\n norm_cfg = cp.deepcopy(norm_cfg)\n super().__init__()\n self.num_units = num_units\n self.norm_cfg = norm_cfg\n self.in_skip = ConvModule(\n in_channels,\n unit_channels,\n kernel_size=1,\n stride=1,\n padding=0,\n norm_cfg=self.norm_cfg,\n act_cfg=None,\n inplace=True)\n self.relu = nn.ReLU(inplace=True)\n\n self.ind = ind\n if self.ind > 0:\n self.up_conv = ConvModule(\n unit_channels,\n unit_channels,\n kernel_size=1,\n stride=1,\n padding=0,\n norm_cfg=self.norm_cfg,\n act_cfg=None,\n inplace=True)\n\n self.gen_skip = gen_skip\n if self.gen_skip:\n self.out_skip1 = ConvModule(\n in_channels,\n in_channels,\n kernel_size=1,\n stride=1,\n padding=0,\n norm_cfg=self.norm_cfg,\n inplace=True)\n\n self.out_skip2 = ConvModule(\n unit_channels,\n in_channels,\n kernel_size=1,\n stride=1,\n padding=0,\n norm_cfg=self.norm_cfg,\n inplace=True)\n\n self.gen_cross_conv = gen_cross_conv\n if self.ind == num_units - 1 and self.gen_cross_conv:\n self.cross_conv = ConvModule(\n unit_channels,\n out_channels,\n kernel_size=1,\n stride=1,\n padding=0,\n norm_cfg=self.norm_cfg,\n inplace=True)\n\n def forward(self, x, up_x):\n out = self.in_skip(x)\n\n if self.ind > 0:\n up_x = F.interpolate(\n up_x,\n size=(x.size(2), x.size(3)),\n mode='bilinear',\n align_corners=True)\n up_x = self.up_conv(up_x)\n out = out + up_x\n out = self.relu(out)\n\n skip1 = None\n skip2 = None\n if self.gen_skip:\n skip1 = self.out_skip1(x)\n skip2 = self.out_skip2(out)\n\n cross_conv = None\n if self.ind == self.num_units - 1 and self.gen_cross_conv:\n cross_conv = self.cross_conv(out)\n\n return out, skip1, skip2, cross_conv\n\n\nclass Upsample_module(nn.Module):\n \"\"\"Upsample module for RSN.\n\n Args:\n unit_channels (int): Channel number in the upsample units.\n Default:256.\n num_units (int): Numbers of upsample units. Default: 4\n gen_skip (bool): Whether to generate skip for posterior downsample\n module or not. Default:False\n gen_cross_conv (bool): Whether to generate feature map for the next\n hourglass-like module. Default:False\n norm_cfg (dict): dictionary to construct and config norm layer.\n Default: dict(type='BN')\n out_channels (int): Number of channels of feature output by upsample\n module. Must equal to in_channels of downsample module. Default:64\n \"\"\"\n\n def __init__(self,\n unit_channels=256,\n num_units=4,\n gen_skip=False,\n gen_cross_conv=False,\n norm_cfg=dict(type='BN'),\n out_channels=64):\n # Protect mutable default arguments\n norm_cfg = cp.deepcopy(norm_cfg)\n super().__init__()\n self.in_channels = list()\n for i in range(num_units):\n self.in_channels.append(RSB.expansion * out_channels * pow(2, i))\n self.in_channels.reverse()\n self.num_units = num_units\n self.gen_skip = gen_skip\n self.gen_cross_conv = gen_cross_conv\n self.norm_cfg = norm_cfg\n for i in range(num_units):\n module_name = f'up{i + 1}'\n self.add_module(\n module_name,\n Upsample_unit(\n i,\n self.num_units,\n self.in_channels[i],\n unit_channels,\n self.gen_skip,\n self.gen_cross_conv,\n norm_cfg=self.norm_cfg,\n out_channels=64))\n\n def forward(self, x):\n out = list()\n skip1 = list()\n skip2 = list()\n cross_conv = None\n for i in range(self.num_units):\n module_i = getattr(self, f'up{i + 1}')\n if i == 0:\n outi, skip1_i, skip2_i, _ = module_i(x[i], None)\n elif i == self.num_units - 1:\n outi, skip1_i, skip2_i, cross_conv = module_i(x[i], out[i - 1])\n else:\n outi, skip1_i, skip2_i, _ = module_i(x[i], out[i - 1])\n out.append(outi)\n skip1.append(skip1_i)\n skip2.append(skip2_i)\n skip1.reverse()\n skip2.reverse()\n\n return out, skip1, skip2, cross_conv\n\n\nclass Single_stage_RSN(nn.Module):\n \"\"\"Single_stage Residual Steps Network.\n\n Args:\n unit_channels (int): Channel number in the upsample units. Default:256.\n num_units (int): Numbers of downsample/upsample units. Default: 4\n gen_skip (bool): Whether to generate skip for posterior downsample\n module or not. Default:False\n gen_cross_conv (bool): Whether to generate feature map for the next\n hourglass-like module. Default:False\n has_skip (bool): Have skip connections from prior upsample\n module or not. Default:False\n num_steps (int): Number of steps in RSB. Default: 4\n num_blocks (list): Number of blocks in each downsample unit.\n Default: [2, 2, 2, 2] Note: Make sure num_units==len(num_blocks)\n norm_cfg (dict): dictionary to construct and config norm layer.\n Default: dict(type='BN')\n in_channels (int): Number of channels of the feature from ResNet_Top.\n Default: 64.\n expand_times (int): Times by which the in_channels are expanded in RSB.\n Default:26.\n \"\"\"\n\n def __init__(self,\n has_skip=False,\n gen_skip=False,\n gen_cross_conv=False,\n unit_channels=256,\n num_units=4,\n num_steps=4,\n num_blocks=[2, 2, 2, 2],\n norm_cfg=dict(type='BN'),\n in_channels=64,\n expand_times=26):\n # Protect mutable default arguments\n norm_cfg = cp.deepcopy(norm_cfg)\n num_blocks = cp.deepcopy(num_blocks)\n super().__init__()\n assert len(num_blocks) == num_units\n self.has_skip = has_skip\n self.gen_skip = gen_skip\n self.gen_cross_conv = gen_cross_conv\n self.num_units = num_units\n self.num_steps = num_steps\n self.unit_channels = unit_channels\n self.num_blocks = num_blocks\n self.norm_cfg = norm_cfg\n\n self.downsample = Downsample_module(RSB, num_blocks, num_steps,\n num_units, has_skip, norm_cfg,\n in_channels, expand_times)\n self.upsample = Upsample_module(unit_channels, num_units, gen_skip,\n gen_cross_conv, norm_cfg, in_channels)\n\n def forward(self, x, skip1, skip2):\n mid = self.downsample(x, skip1, skip2)\n out, skip1, skip2, cross_conv = self.upsample(mid)\n\n return out, skip1, skip2, cross_conv\n\n\nclass ResNet_top(nn.Module):\n \"\"\"ResNet top for RSN.\n\n Args:\n norm_cfg (dict): dictionary to construct and config norm layer.\n Default: dict(type='BN')\n channels (int): Number of channels of the feature output by ResNet_top.\n \"\"\"\n\n def __init__(self, norm_cfg=dict(type='BN'), channels=64):\n # Protect mutable default arguments\n norm_cfg = cp.deepcopy(norm_cfg)\n super().__init__()\n self.top = nn.Sequential(\n ConvModule(\n 3,\n channels,\n kernel_size=7,\n stride=2,\n padding=3,\n norm_cfg=norm_cfg,\n inplace=True), MaxPool2d(kernel_size=3, stride=2, padding=1))\n\n def forward(self, img):\n return self.top(img)\n\n\n@BACKBONES.register_module()\nclass RSN(BaseBackbone):\n \"\"\"Residual Steps Network backbone. Paper ref: Cai et al. \"Learning\n Delicate Local Representations for Multi-Person Pose Estimation\" (ECCV\n 2020).\n\n Args:\n unit_channels (int): Number of Channels in an upsample unit.\n Default: 256\n num_stages (int): Number of stages in a multi-stage RSN. Default: 4\n num_units (int): NUmber of downsample/upsample units in a single-stage\n RSN. Default: 4 Note: Make sure num_units == len(self.num_blocks)\n num_blocks (list): Number of RSBs (Residual Steps Block) in each\n downsample unit. Default: [2, 2, 2, 2]\n num_steps (int): Number of steps in a RSB. Default:4\n norm_cfg (dict): dictionary to construct and config norm layer.\n Default: dict(type='BN')\n res_top_channels (int): Number of channels of feature from ResNet_top.\n Default: 64.\n expand_times (int): Times by which the in_channels are expanded in RSB.\n Default:26.\n Example:\n >>> from mmpose.models import RSN\n >>> import torch\n >>> self = RSN(num_stages=2,num_units=2,num_blocks=[2,2])\n >>> self.eval()\n >>> inputs = torch.rand(1, 3, 511, 511)\n >>> level_outputs = self.forward(inputs)\n >>> for level_output in level_outputs:\n ... for feature in level_output:\n ... print(tuple(feature.shape))\n ...\n (1, 256, 64, 64)\n (1, 256, 128, 128)\n (1, 256, 64, 64)\n (1, 256, 128, 128)\n \"\"\"\n\n def __init__(self,\n unit_channels=256,\n num_stages=4,\n num_units=4,\n num_blocks=[2, 2, 2, 2],\n num_steps=4,\n norm_cfg=dict(type='BN'),\n res_top_channels=64,\n expand_times=26):\n # Protect mutable default arguments\n norm_cfg = cp.deepcopy(norm_cfg)\n num_blocks = cp.deepcopy(num_blocks)\n super().__init__()\n self.unit_channels = unit_channels\n self.num_stages = num_stages\n self.num_units = num_units\n self.num_blocks = num_blocks\n self.num_steps = num_steps\n self.norm_cfg = norm_cfg\n\n assert self.num_stages > 0\n assert self.num_steps > 1\n assert self.num_units > 1\n assert self.num_units == len(self.num_blocks)\n self.top = ResNet_top(norm_cfg=norm_cfg)\n self.multi_stage_rsn = nn.ModuleList([])\n for i in range(self.num_stages):\n if i == 0:\n has_skip = False\n else:\n has_skip = True\n if i != self.num_stages - 1:\n gen_skip = True\n gen_cross_conv = True\n else:\n gen_skip = False\n gen_cross_conv = False\n self.multi_stage_rsn.append(\n Single_stage_RSN(has_skip, gen_skip, gen_cross_conv,\n unit_channels, num_units, num_steps,\n num_blocks, norm_cfg, res_top_channels,\n expand_times))\n\n def forward(self, x):\n \"\"\"Model forward function.\"\"\"\n out_feats = []\n skip1 = None\n skip2 = None\n x = self.top(x)\n for i in range(self.num_stages):\n out, skip1, skip2, x = self.multi_stage_rsn[i](x, skip1, skip2)\n out_feats.append(out)\n\n return out_feats\n\n def init_weights(self, pretrained=None):\n \"\"\"Initialize model weights.\"\"\"\n for m in self.multi_stage_rsn.modules():\n if isinstance(m, nn.Conv2d):\n kaiming_init(m)\n elif isinstance(m, nn.BatchNorm2d):\n constant_init(m, 1)\n elif isinstance(m, nn.Linear):\n normal_init(m, std=0.01)\n\n for m in self.top.modules():\n if isinstance(m, nn.Conv2d):\n kaiming_init(m)\n" }, { "path": "eval/GVHMR/hmr4d/utils/preproc/vitpose_pytorch/src/vitpose_infer/builder/backbones/scnet.py", "content": "# Copyright (c) OpenMMLab. All rights reserved.\nimport copy\n\nimport torch\nimport torch.nn as nn\nimport torch.nn.functional as F\nimport torch.utils.checkpoint as cp\nfrom mmcv.cnn import build_conv_layer, build_norm_layer\n\nfrom ..builder import BACKBONES\nfrom .resnet import Bottleneck, ResNet\n\n\nclass SCConv(nn.Module):\n \"\"\"SCConv (Self-calibrated Convolution)\n\n Args:\n in_channels (int): The input channels of the SCConv.\n out_channels (int): The output channel of the SCConv.\n stride (int): stride of SCConv.\n pooling_r (int): size of pooling for scconv.\n conv_cfg (dict): dictionary to construct and config conv layer.\n Default: None\n norm_cfg (dict): dictionary to construct and config norm layer.\n Default: dict(type='BN')\n \"\"\"\n\n def __init__(self,\n in_channels,\n out_channels,\n stride,\n pooling_r,\n conv_cfg=None,\n norm_cfg=dict(type='BN', momentum=0.1)):\n # Protect mutable default arguments\n norm_cfg = copy.deepcopy(norm_cfg)\n super().__init__()\n\n assert in_channels == out_channels\n\n self.k2 = nn.Sequential(\n nn.AvgPool2d(kernel_size=pooling_r, stride=pooling_r),\n build_conv_layer(\n conv_cfg,\n in_channels,\n in_channels,\n kernel_size=3,\n stride=1,\n padding=1,\n bias=False),\n build_norm_layer(norm_cfg, in_channels)[1],\n )\n self.k3 = nn.Sequential(\n build_conv_layer(\n conv_cfg,\n in_channels,\n in_channels,\n kernel_size=3,\n stride=1,\n padding=1,\n bias=False),\n build_norm_layer(norm_cfg, in_channels)[1],\n )\n self.k4 = nn.Sequential(\n build_conv_layer(\n conv_cfg,\n in_channels,\n in_channels,\n kernel_size=3,\n stride=stride,\n padding=1,\n bias=False),\n build_norm_layer(norm_cfg, out_channels)[1],\n nn.ReLU(inplace=True),\n )\n\n def forward(self, x):\n \"\"\"Forward function.\"\"\"\n identity = x\n\n out = torch.sigmoid(\n torch.add(identity, F.interpolate(self.k2(x),\n identity.size()[2:])))\n out = torch.mul(self.k3(x), out)\n out = self.k4(out)\n\n return out\n\n\nclass SCBottleneck(Bottleneck):\n \"\"\"SC(Self-calibrated) Bottleneck.\n\n Args:\n in_channels (int): The input channels of the SCBottleneck block.\n out_channels (int): The output channel of the SCBottleneck block.\n \"\"\"\n\n pooling_r = 4\n\n def __init__(self, in_channels, out_channels, **kwargs):\n super().__init__(in_channels, out_channels, **kwargs)\n self.mid_channels = out_channels // self.expansion // 2\n\n self.norm1_name, norm1 = build_norm_layer(\n self.norm_cfg, self.mid_channels, postfix=1)\n self.norm2_name, norm2 = build_norm_layer(\n self.norm_cfg, self.mid_channels, postfix=2)\n self.norm3_name, norm3 = build_norm_layer(\n self.norm_cfg, out_channels, postfix=3)\n\n self.conv1 = build_conv_layer(\n self.conv_cfg,\n in_channels,\n self.mid_channels,\n kernel_size=1,\n stride=1,\n bias=False)\n self.add_module(self.norm1_name, norm1)\n\n self.k1 = nn.Sequential(\n build_conv_layer(\n self.conv_cfg,\n self.mid_channels,\n self.mid_channels,\n kernel_size=3,\n stride=self.stride,\n padding=1,\n bias=False),\n build_norm_layer(self.norm_cfg, self.mid_channels)[1],\n nn.ReLU(inplace=True))\n\n self.conv2 = build_conv_layer(\n self.conv_cfg,\n in_channels,\n self.mid_channels,\n kernel_size=1,\n stride=1,\n bias=False)\n self.add_module(self.norm2_name, norm2)\n\n self.scconv = SCConv(self.mid_channels, self.mid_channels, self.stride,\n self.pooling_r, self.conv_cfg, self.norm_cfg)\n\n self.conv3 = build_conv_layer(\n self.conv_cfg,\n self.mid_channels * 2,\n out_channels,\n kernel_size=1,\n stride=1,\n bias=False)\n self.add_module(self.norm3_name, norm3)\n\n def forward(self, x):\n \"\"\"Forward function.\"\"\"\n\n def _inner_forward(x):\n identity = x\n\n out_a = self.conv1(x)\n out_a = self.norm1(out_a)\n out_a = self.relu(out_a)\n\n out_a = self.k1(out_a)\n\n out_b = self.conv2(x)\n out_b = self.norm2(out_b)\n out_b = self.relu(out_b)\n\n out_b = self.scconv(out_b)\n\n out = self.conv3(torch.cat([out_a, out_b], dim=1))\n out = self.norm3(out)\n\n if self.downsample is not None:\n identity = self.downsample(x)\n\n out += identity\n\n return out\n\n if self.with_cp and x.requires_grad:\n out = cp.checkpoint(_inner_forward, x)\n else:\n out = _inner_forward(x)\n\n out = self.relu(out)\n\n return out\n\n\n@BACKBONES.register_module()\nclass SCNet(ResNet):\n \"\"\"SCNet backbone.\n\n Improving Convolutional Networks with Self-Calibrated Convolutions,\n Jiang-Jiang Liu, Qibin Hou, Ming-Ming Cheng, Changhu Wang, Jiashi Feng,\n IEEE CVPR, 2020.\n http://mftp.mmcheng.net/Papers/20cvprSCNet.pdf\n\n Args:\n depth (int): Depth of scnet, from {50, 101}.\n in_channels (int): Number of input image channels. Normally 3.\n base_channels (int): Number of base channels of hidden layer.\n num_stages (int): SCNet stages, normally 4.\n strides (Sequence[int]): Strides of the first block of each stage.\n dilations (Sequence[int]): Dilation of each stage.\n out_indices (Sequence[int]): Output from which stages.\n style (str): `pytorch` or `caffe`. If set to \"pytorch\", the stride-two\n layer is the 3x3 conv layer, otherwise the stride-two layer is\n the first 1x1 conv layer.\n deep_stem (bool): Replace 7x7 conv in input stem with 3 3x3 conv\n avg_down (bool): Use AvgPool instead of stride conv when\n downsampling in the bottleneck.\n frozen_stages (int): Stages to be frozen (stop grad and set eval mode).\n -1 means not freezing any parameters.\n norm_cfg (dict): Dictionary to construct and config norm layer.\n norm_eval (bool): Whether to set norm layers to eval mode, namely,\n freeze running stats (mean and var). Note: Effect on Batch Norm\n and its variants only.\n with_cp (bool): Use checkpoint or not. Using checkpoint will save some\n memory while slowing down the training speed.\n zero_init_residual (bool): Whether to use zero init for last norm layer\n in resblocks to let them behave as identity.\n\n Example:\n >>> from mmpose.models import SCNet\n >>> import torch\n >>> self = SCNet(depth=50, out_indices=(0, 1, 2, 3))\n >>> self.eval()\n >>> inputs = torch.rand(1, 3, 224, 224)\n >>> level_outputs = self.forward(inputs)\n >>> for level_out in level_outputs:\n ... print(tuple(level_out.shape))\n (1, 256, 56, 56)\n (1, 512, 28, 28)\n (1, 1024, 14, 14)\n (1, 2048, 7, 7)\n \"\"\"\n\n arch_settings = {\n 50: (SCBottleneck, [3, 4, 6, 3]),\n 101: (SCBottleneck, [3, 4, 23, 3])\n }\n\n def __init__(self, depth, **kwargs):\n if depth not in self.arch_settings:\n raise KeyError(f'invalid depth {depth} for SCNet')\n super().__init__(depth, **kwargs)\n" }, { "path": "eval/GVHMR/hmr4d/utils/preproc/vitpose_pytorch/src/vitpose_infer/builder/backbones/seresnet.py", "content": "# Copyright (c) OpenMMLab. All rights reserved.\nimport torch.utils.checkpoint as cp\n\nfrom ..builder import BACKBONES\nfrom .resnet import Bottleneck, ResLayer, ResNet\nfrom .utils.se_layer import SELayer\n\n\nclass SEBottleneck(Bottleneck):\n \"\"\"SEBottleneck block for SEResNet.\n\n Args:\n in_channels (int): The input channels of the SEBottleneck block.\n out_channels (int): The output channel of the SEBottleneck block.\n se_ratio (int): Squeeze ratio in SELayer. Default: 16\n \"\"\"\n\n def __init__(self, in_channels, out_channels, se_ratio=16, **kwargs):\n super().__init__(in_channels, out_channels, **kwargs)\n self.se_layer = SELayer(out_channels, ratio=se_ratio)\n\n def forward(self, x):\n\n def _inner_forward(x):\n identity = x\n\n out = self.conv1(x)\n out = self.norm1(out)\n out = self.relu(out)\n\n out = self.conv2(out)\n out = self.norm2(out)\n out = self.relu(out)\n\n out = self.conv3(out)\n out = self.norm3(out)\n\n out = self.se_layer(out)\n\n if self.downsample is not None:\n identity = self.downsample(x)\n\n out += identity\n\n return out\n\n if self.with_cp and x.requires_grad:\n out = cp.checkpoint(_inner_forward, x)\n else:\n out = _inner_forward(x)\n\n out = self.relu(out)\n\n return out\n\n\n@BACKBONES.register_module()\nclass SEResNet(ResNet):\n \"\"\"SEResNet backbone.\n\n Please refer to the `paper `__ for\n details.\n\n Args:\n depth (int): Network depth, from {50, 101, 152}.\n se_ratio (int): Squeeze ratio in SELayer. Default: 16.\n in_channels (int): Number of input image channels. Default: 3.\n stem_channels (int): Output channels of the stem layer. Default: 64.\n num_stages (int): Stages of the network. Default: 4.\n strides (Sequence[int]): Strides of the first block of each stage.\n Default: ``(1, 2, 2, 2)``.\n dilations (Sequence[int]): Dilation of each stage.\n Default: ``(1, 1, 1, 1)``.\n out_indices (Sequence[int]): Output from which stages. If only one\n stage is specified, a single tensor (feature map) is returned,\n otherwise multiple stages are specified, a tuple of tensors will\n be returned. Default: ``(3, )``.\n style (str): `pytorch` or `caffe`. If set to \"pytorch\", the stride-two\n layer is the 3x3 conv layer, otherwise the stride-two layer is\n the first 1x1 conv layer.\n deep_stem (bool): Replace 7x7 conv in input stem with 3 3x3 conv.\n Default: False.\n avg_down (bool): Use AvgPool instead of stride conv when\n downsampling in the bottleneck. Default: False.\n frozen_stages (int): Stages to be frozen (stop grad and set eval mode).\n -1 means not freezing any parameters. Default: -1.\n conv_cfg (dict | None): The config dict for conv layers. Default: None.\n norm_cfg (dict): The config dict for norm layers.\n norm_eval (bool): Whether to set norm layers to eval mode, namely,\n freeze running stats (mean and var). Note: Effect on Batch Norm\n and its variants only. Default: False.\n with_cp (bool): Use checkpoint or not. Using checkpoint will save some\n memory while slowing down the training speed. Default: False.\n zero_init_residual (bool): Whether to use zero init for last norm layer\n in resblocks to let them behave as identity. Default: True.\n\n Example:\n >>> from mmpose.models import SEResNet\n >>> import torch\n >>> self = SEResNet(depth=50, out_indices=(0, 1, 2, 3))\n >>> self.eval()\n >>> inputs = torch.rand(1, 3, 224, 224)\n >>> level_outputs = self.forward(inputs)\n >>> for level_out in level_outputs:\n ... print(tuple(level_out.shape))\n (1, 256, 56, 56)\n (1, 512, 28, 28)\n (1, 1024, 14, 14)\n (1, 2048, 7, 7)\n \"\"\"\n\n arch_settings = {\n 50: (SEBottleneck, (3, 4, 6, 3)),\n 101: (SEBottleneck, (3, 4, 23, 3)),\n 152: (SEBottleneck, (3, 8, 36, 3))\n }\n\n def __init__(self, depth, se_ratio=16, **kwargs):\n if depth not in self.arch_settings:\n raise KeyError(f'invalid depth {depth} for SEResNet')\n self.se_ratio = se_ratio\n super().__init__(depth, **kwargs)\n\n def make_res_layer(self, **kwargs):\n return ResLayer(se_ratio=self.se_ratio, **kwargs)\n" }, { "path": "eval/GVHMR/hmr4d/utils/preproc/vitpose_pytorch/src/vitpose_infer/builder/backbones/seresnext.py", "content": "# Copyright (c) OpenMMLab. All rights reserved.\nfrom mmcv.cnn import build_conv_layer, build_norm_layer\n\nfrom ..builder import BACKBONES\nfrom .resnet import ResLayer\nfrom .seresnet import SEBottleneck as _SEBottleneck\nfrom .seresnet import SEResNet\n\n\nclass SEBottleneck(_SEBottleneck):\n \"\"\"SEBottleneck block for SEResNeXt.\n\n Args:\n in_channels (int): Input channels of this block.\n out_channels (int): Output channels of this block.\n base_channels (int): Middle channels of the first stage. Default: 64.\n groups (int): Groups of conv2.\n width_per_group (int): Width per group of conv2. 64x4d indicates\n ``groups=64, width_per_group=4`` and 32x8d indicates\n ``groups=32, width_per_group=8``.\n stride (int): stride of the block. Default: 1\n dilation (int): dilation of convolution. Default: 1\n downsample (nn.Module): downsample operation on identity branch.\n Default: None\n se_ratio (int): Squeeze ratio in SELayer. Default: 16\n style (str): `pytorch` or `caffe`. If set to \"pytorch\", the stride-two\n layer is the 3x3 conv layer, otherwise the stride-two layer is\n the first 1x1 conv layer.\n conv_cfg (dict): dictionary to construct and config conv layer.\n Default: None\n norm_cfg (dict): dictionary to construct and config norm layer.\n Default: dict(type='BN')\n with_cp (bool): Use checkpoint or not. Using checkpoint will save some\n memory while slowing down the training speed.\n \"\"\"\n\n def __init__(self,\n in_channels,\n out_channels,\n base_channels=64,\n groups=32,\n width_per_group=4,\n se_ratio=16,\n **kwargs):\n super().__init__(in_channels, out_channels, se_ratio, **kwargs)\n self.groups = groups\n self.width_per_group = width_per_group\n\n # We follow the same rational of ResNext to compute mid_channels.\n # For SEResNet bottleneck, middle channels are determined by expansion\n # and out_channels, but for SEResNeXt bottleneck, it is determined by\n # groups and width_per_group and the stage it is located in.\n if groups != 1:\n assert self.mid_channels % base_channels == 0\n self.mid_channels = (\n groups * width_per_group * self.mid_channels // base_channels)\n\n self.norm1_name, norm1 = build_norm_layer(\n self.norm_cfg, self.mid_channels, postfix=1)\n self.norm2_name, norm2 = build_norm_layer(\n self.norm_cfg, self.mid_channels, postfix=2)\n self.norm3_name, norm3 = build_norm_layer(\n self.norm_cfg, self.out_channels, postfix=3)\n\n self.conv1 = build_conv_layer(\n self.conv_cfg,\n self.in_channels,\n self.mid_channels,\n kernel_size=1,\n stride=self.conv1_stride,\n bias=False)\n self.add_module(self.norm1_name, norm1)\n self.conv2 = build_conv_layer(\n self.conv_cfg,\n self.mid_channels,\n self.mid_channels,\n kernel_size=3,\n stride=self.conv2_stride,\n padding=self.dilation,\n dilation=self.dilation,\n groups=groups,\n bias=False)\n\n self.add_module(self.norm2_name, norm2)\n self.conv3 = build_conv_layer(\n self.conv_cfg,\n self.mid_channels,\n self.out_channels,\n kernel_size=1,\n bias=False)\n self.add_module(self.norm3_name, norm3)\n\n\n@BACKBONES.register_module()\nclass SEResNeXt(SEResNet):\n \"\"\"SEResNeXt backbone.\n\n Please refer to the `paper `__ for\n details.\n\n Args:\n depth (int): Network depth, from {50, 101, 152}.\n groups (int): Groups of conv2 in Bottleneck. Default: 32.\n width_per_group (int): Width per group of conv2 in Bottleneck.\n Default: 4.\n se_ratio (int): Squeeze ratio in SELayer. Default: 16.\n in_channels (int): Number of input image channels. Default: 3.\n stem_channels (int): Output channels of the stem layer. Default: 64.\n num_stages (int): Stages of the network. Default: 4.\n strides (Sequence[int]): Strides of the first block of each stage.\n Default: ``(1, 2, 2, 2)``.\n dilations (Sequence[int]): Dilation of each stage.\n Default: ``(1, 1, 1, 1)``.\n out_indices (Sequence[int]): Output from which stages. If only one\n stage is specified, a single tensor (feature map) is returned,\n otherwise multiple stages are specified, a tuple of tensors will\n be returned. Default: ``(3, )``.\n style (str): `pytorch` or `caffe`. If set to \"pytorch\", the stride-two\n layer is the 3x3 conv layer, otherwise the stride-two layer is\n the first 1x1 conv layer.\n deep_stem (bool): Replace 7x7 conv in input stem with 3 3x3 conv.\n Default: False.\n avg_down (bool): Use AvgPool instead of stride conv when\n downsampling in the bottleneck. Default: False.\n frozen_stages (int): Stages to be frozen (stop grad and set eval mode).\n -1 means not freezing any parameters. Default: -1.\n conv_cfg (dict | None): The config dict for conv layers. Default: None.\n norm_cfg (dict): The config dict for norm layers.\n norm_eval (bool): Whether to set norm layers to eval mode, namely,\n freeze running stats (mean and var). Note: Effect on Batch Norm\n and its variants only. Default: False.\n with_cp (bool): Use checkpoint or not. Using checkpoint will save some\n memory while slowing down the training speed. Default: False.\n zero_init_residual (bool): Whether to use zero init for last norm layer\n in resblocks to let them behave as identity. Default: True.\n\n Example:\n >>> from mmpose.models import SEResNeXt\n >>> import torch\n >>> self = SEResNet(depth=50, out_indices=(0, 1, 2, 3))\n >>> self.eval()\n >>> inputs = torch.rand(1, 3, 224, 224)\n >>> level_outputs = self.forward(inputs)\n >>> for level_out in level_outputs:\n ... print(tuple(level_out.shape))\n (1, 256, 56, 56)\n (1, 512, 28, 28)\n (1, 1024, 14, 14)\n (1, 2048, 7, 7)\n \"\"\"\n\n arch_settings = {\n 50: (SEBottleneck, (3, 4, 6, 3)),\n 101: (SEBottleneck, (3, 4, 23, 3)),\n 152: (SEBottleneck, (3, 8, 36, 3))\n }\n\n def __init__(self, depth, groups=32, width_per_group=4, **kwargs):\n self.groups = groups\n self.width_per_group = width_per_group\n super().__init__(depth, **kwargs)\n\n def make_res_layer(self, **kwargs):\n return ResLayer(\n groups=self.groups,\n width_per_group=self.width_per_group,\n base_channels=self.base_channels,\n **kwargs)\n" }, { "path": "eval/GVHMR/hmr4d/utils/preproc/vitpose_pytorch/src/vitpose_infer/builder/backbones/shufflenet_v1.py", "content": "# Copyright (c) OpenMMLab. All rights reserved.\nimport copy\nimport logging\n\nimport torch\nimport torch.nn as nn\nimport torch.utils.checkpoint as cp\nfrom mmcv.cnn import (ConvModule, build_activation_layer, constant_init,\n normal_init)\nfrom torch.nn.modules.batchnorm import _BatchNorm\n\nfrom ..builder import BACKBONES\nfrom .base_backbone import BaseBackbone\nfrom .utils import channel_shuffle, load_checkpoint, make_divisible\n\n\nclass ShuffleUnit(nn.Module):\n \"\"\"ShuffleUnit block.\n\n ShuffleNet unit with pointwise group convolution (GConv) and channel\n shuffle.\n\n Args:\n in_channels (int): The input channels of the ShuffleUnit.\n out_channels (int): The output channels of the ShuffleUnit.\n groups (int, optional): The number of groups to be used in grouped 1x1\n convolutions in each ShuffleUnit. Default: 3\n first_block (bool, optional): Whether it is the first ShuffleUnit of a\n sequential ShuffleUnits. Default: True, which means not using the\n grouped 1x1 convolution.\n combine (str, optional): The ways to combine the input and output\n branches. Default: 'add'.\n conv_cfg (dict): Config dict for convolution layer. Default: None,\n which means using conv2d.\n norm_cfg (dict): Config dict for normalization layer.\n Default: dict(type='BN').\n act_cfg (dict): Config dict for activation layer.\n Default: dict(type='ReLU').\n with_cp (bool, optional): Use checkpoint or not. Using checkpoint\n will save some memory while slowing down the training speed.\n Default: False.\n\n Returns:\n Tensor: The output tensor.\n \"\"\"\n\n def __init__(self,\n in_channels,\n out_channels,\n groups=3,\n first_block=True,\n combine='add',\n conv_cfg=None,\n norm_cfg=dict(type='BN'),\n act_cfg=dict(type='ReLU'),\n with_cp=False):\n # Protect mutable default arguments\n norm_cfg = copy.deepcopy(norm_cfg)\n act_cfg = copy.deepcopy(act_cfg)\n super().__init__()\n self.in_channels = in_channels\n self.out_channels = out_channels\n self.first_block = first_block\n self.combine = combine\n self.groups = groups\n self.bottleneck_channels = self.out_channels // 4\n self.with_cp = with_cp\n\n if self.combine == 'add':\n self.depthwise_stride = 1\n self._combine_func = self._add\n assert in_channels == out_channels, (\n 'in_channels must be equal to out_channels when combine '\n 'is add')\n elif self.combine == 'concat':\n self.depthwise_stride = 2\n self._combine_func = self._concat\n self.out_channels -= self.in_channels\n self.avgpool = nn.AvgPool2d(kernel_size=3, stride=2, padding=1)\n else:\n raise ValueError(f'Cannot combine tensors with {self.combine}. '\n 'Only \"add\" and \"concat\" are supported')\n\n self.first_1x1_groups = 1 if first_block else self.groups\n self.g_conv_1x1_compress = ConvModule(\n in_channels=self.in_channels,\n out_channels=self.bottleneck_channels,\n kernel_size=1,\n groups=self.first_1x1_groups,\n conv_cfg=conv_cfg,\n norm_cfg=norm_cfg,\n act_cfg=act_cfg)\n\n self.depthwise_conv3x3_bn = ConvModule(\n in_channels=self.bottleneck_channels,\n out_channels=self.bottleneck_channels,\n kernel_size=3,\n stride=self.depthwise_stride,\n padding=1,\n groups=self.bottleneck_channels,\n conv_cfg=conv_cfg,\n norm_cfg=norm_cfg,\n act_cfg=None)\n\n self.g_conv_1x1_expand = ConvModule(\n in_channels=self.bottleneck_channels,\n out_channels=self.out_channels,\n kernel_size=1,\n groups=self.groups,\n conv_cfg=conv_cfg,\n norm_cfg=norm_cfg,\n act_cfg=None)\n\n self.act = build_activation_layer(act_cfg)\n\n @staticmethod\n def _add(x, out):\n # residual connection\n return x + out\n\n @staticmethod\n def _concat(x, out):\n # concatenate along channel axis\n return torch.cat((x, out), 1)\n\n def forward(self, x):\n\n def _inner_forward(x):\n residual = x\n\n out = self.g_conv_1x1_compress(x)\n out = self.depthwise_conv3x3_bn(out)\n\n if self.groups > 1:\n out = channel_shuffle(out, self.groups)\n\n out = self.g_conv_1x1_expand(out)\n\n if self.combine == 'concat':\n residual = self.avgpool(residual)\n out = self.act(out)\n out = self._combine_func(residual, out)\n else:\n out = self._combine_func(residual, out)\n out = self.act(out)\n return out\n\n if self.with_cp and x.requires_grad:\n out = cp.checkpoint(_inner_forward, x)\n else:\n out = _inner_forward(x)\n\n return out\n\n\n@BACKBONES.register_module()\nclass ShuffleNetV1(BaseBackbone):\n \"\"\"ShuffleNetV1 backbone.\n\n Args:\n groups (int, optional): The number of groups to be used in grouped 1x1\n convolutions in each ShuffleUnit. Default: 3.\n widen_factor (float, optional): Width multiplier - adjusts the number\n of channels in each layer by this amount. Default: 1.0.\n out_indices (Sequence[int]): Output from which stages.\n Default: (2, )\n frozen_stages (int): Stages to be frozen (all param fixed).\n Default: -1, which means not freezing any parameters.\n conv_cfg (dict): Config dict for convolution layer. Default: None,\n which means using conv2d.\n norm_cfg (dict): Config dict for normalization layer.\n Default: dict(type='BN').\n act_cfg (dict): Config dict for activation layer.\n Default: dict(type='ReLU').\n norm_eval (bool): Whether to set norm layers to eval mode, namely,\n freeze running stats (mean and var). Note: Effect on Batch Norm\n and its variants only. Default: False.\n with_cp (bool): Use checkpoint or not. Using checkpoint will save some\n memory while slowing down the training speed. Default: False.\n \"\"\"\n\n def __init__(self,\n groups=3,\n widen_factor=1.0,\n out_indices=(2, ),\n frozen_stages=-1,\n conv_cfg=None,\n norm_cfg=dict(type='BN'),\n act_cfg=dict(type='ReLU'),\n norm_eval=False,\n with_cp=False):\n # Protect mutable default arguments\n norm_cfg = copy.deepcopy(norm_cfg)\n act_cfg = copy.deepcopy(act_cfg)\n super().__init__()\n self.stage_blocks = [4, 8, 4]\n self.groups = groups\n\n for index in out_indices:\n if index not in range(0, 3):\n raise ValueError('the item in out_indices must in '\n f'range(0, 3). But received {index}')\n\n if frozen_stages not in range(-1, 3):\n raise ValueError('frozen_stages must be in range(-1, 3). '\n f'But received {frozen_stages}')\n self.out_indices = out_indices\n self.frozen_stages = frozen_stages\n self.conv_cfg = conv_cfg\n self.norm_cfg = norm_cfg\n self.act_cfg = act_cfg\n self.norm_eval = norm_eval\n self.with_cp = with_cp\n\n if groups == 1:\n channels = (144, 288, 576)\n elif groups == 2:\n channels = (200, 400, 800)\n elif groups == 3:\n channels = (240, 480, 960)\n elif groups == 4:\n channels = (272, 544, 1088)\n elif groups == 8:\n channels = (384, 768, 1536)\n else:\n raise ValueError(f'{groups} groups is not supported for 1x1 '\n 'Grouped Convolutions')\n\n channels = [make_divisible(ch * widen_factor, 8) for ch in channels]\n\n self.in_channels = int(24 * widen_factor)\n\n self.conv1 = ConvModule(\n in_channels=3,\n out_channels=self.in_channels,\n kernel_size=3,\n stride=2,\n padding=1,\n conv_cfg=conv_cfg,\n norm_cfg=norm_cfg,\n act_cfg=act_cfg)\n self.maxpool = nn.MaxPool2d(kernel_size=3, stride=2, padding=1)\n\n self.layers = nn.ModuleList()\n for i, num_blocks in enumerate(self.stage_blocks):\n first_block = (i == 0)\n layer = self.make_layer(channels[i], num_blocks, first_block)\n self.layers.append(layer)\n\n def _freeze_stages(self):\n if self.frozen_stages >= 0:\n for param in self.conv1.parameters():\n param.requires_grad = False\n for i in range(self.frozen_stages):\n layer = self.layers[i]\n layer.eval()\n for param in layer.parameters():\n param.requires_grad = False\n\n def init_weights(self, pretrained=None):\n if isinstance(pretrained, str):\n logger = logging.getLogger()\n load_checkpoint(self, pretrained, strict=False, logger=logger)\n elif pretrained is None:\n for name, m in self.named_modules():\n if isinstance(m, nn.Conv2d):\n if 'conv1' in name:\n normal_init(m, mean=0, std=0.01)\n else:\n normal_init(m, mean=0, std=1.0 / m.weight.shape[1])\n elif isinstance(m, (_BatchNorm, nn.GroupNorm)):\n constant_init(m, val=1, bias=0.0001)\n if isinstance(m, _BatchNorm):\n if m.running_mean is not None:\n nn.init.constant_(m.running_mean, 0)\n else:\n raise TypeError('pretrained must be a str or None. But received '\n f'{type(pretrained)}')\n\n def make_layer(self, out_channels, num_blocks, first_block=False):\n \"\"\"Stack ShuffleUnit blocks to make a layer.\n\n Args:\n out_channels (int): out_channels of the block.\n num_blocks (int): Number of blocks.\n first_block (bool, optional): Whether is the first ShuffleUnit of a\n sequential ShuffleUnits. Default: False, which means using\n the grouped 1x1 convolution.\n \"\"\"\n layers = []\n for i in range(num_blocks):\n first_block = first_block if i == 0 else False\n combine_mode = 'concat' if i == 0 else 'add'\n layers.append(\n ShuffleUnit(\n self.in_channels,\n out_channels,\n groups=self.groups,\n first_block=first_block,\n combine=combine_mode,\n conv_cfg=self.conv_cfg,\n norm_cfg=self.norm_cfg,\n act_cfg=self.act_cfg,\n with_cp=self.with_cp))\n self.in_channels = out_channels\n\n return nn.Sequential(*layers)\n\n def forward(self, x):\n x = self.conv1(x)\n x = self.maxpool(x)\n\n outs = []\n for i, layer in enumerate(self.layers):\n x = layer(x)\n if i in self.out_indices:\n outs.append(x)\n\n if len(outs) == 1:\n return outs[0]\n return tuple(outs)\n\n def train(self, mode=True):\n super().train(mode)\n self._freeze_stages()\n if mode and self.norm_eval:\n for m in self.modules():\n if isinstance(m, _BatchNorm):\n m.eval()\n" }, { "path": "eval/GVHMR/hmr4d/utils/preproc/vitpose_pytorch/src/vitpose_infer/builder/backbones/shufflenet_v2.py", "content": "# Copyright (c) OpenMMLab. All rights reserved.\nimport copy\nimport logging\n\nimport torch\nimport torch.nn as nn\nimport torch.utils.checkpoint as cp\nfrom mmcv.cnn import ConvModule, constant_init, normal_init\nfrom torch.nn.modules.batchnorm import _BatchNorm\n\nfrom ..builder import BACKBONES\nfrom .base_backbone import BaseBackbone\nfrom .utils import channel_shuffle, load_checkpoint\n\n\nclass InvertedResidual(nn.Module):\n \"\"\"InvertedResidual block for ShuffleNetV2 backbone.\n\n Args:\n in_channels (int): The input channels of the block.\n out_channels (int): The output channels of the block.\n stride (int): Stride of the 3x3 convolution layer. Default: 1\n conv_cfg (dict): Config dict for convolution layer.\n Default: None, which means using conv2d.\n norm_cfg (dict): Config dict for normalization layer.\n Default: dict(type='BN').\n act_cfg (dict): Config dict for activation layer.\n Default: dict(type='ReLU').\n with_cp (bool): Use checkpoint or not. Using checkpoint will save some\n memory while slowing down the training speed. Default: False.\n \"\"\"\n\n def __init__(self,\n in_channels,\n out_channels,\n stride=1,\n conv_cfg=None,\n norm_cfg=dict(type='BN'),\n act_cfg=dict(type='ReLU'),\n with_cp=False):\n # Protect mutable default arguments\n norm_cfg = copy.deepcopy(norm_cfg)\n act_cfg = copy.deepcopy(act_cfg)\n super().__init__()\n self.stride = stride\n self.with_cp = with_cp\n\n branch_features = out_channels // 2\n if self.stride == 1:\n assert in_channels == branch_features * 2, (\n f'in_channels ({in_channels}) should equal to '\n f'branch_features * 2 ({branch_features * 2}) '\n 'when stride is 1')\n\n if in_channels != branch_features * 2:\n assert self.stride != 1, (\n f'stride ({self.stride}) should not equal 1 when '\n f'in_channels != branch_features * 2')\n\n if self.stride > 1:\n self.branch1 = nn.Sequential(\n ConvModule(\n in_channels,\n in_channels,\n kernel_size=3,\n stride=self.stride,\n padding=1,\n groups=in_channels,\n conv_cfg=conv_cfg,\n norm_cfg=norm_cfg,\n act_cfg=None),\n ConvModule(\n in_channels,\n branch_features,\n kernel_size=1,\n stride=1,\n padding=0,\n conv_cfg=conv_cfg,\n norm_cfg=norm_cfg,\n act_cfg=act_cfg),\n )\n\n self.branch2 = nn.Sequential(\n ConvModule(\n in_channels if (self.stride > 1) else branch_features,\n branch_features,\n kernel_size=1,\n stride=1,\n padding=0,\n conv_cfg=conv_cfg,\n norm_cfg=norm_cfg,\n act_cfg=act_cfg),\n ConvModule(\n branch_features,\n branch_features,\n kernel_size=3,\n stride=self.stride,\n padding=1,\n groups=branch_features,\n conv_cfg=conv_cfg,\n norm_cfg=norm_cfg,\n act_cfg=None),\n ConvModule(\n branch_features,\n branch_features,\n kernel_size=1,\n stride=1,\n padding=0,\n conv_cfg=conv_cfg,\n norm_cfg=norm_cfg,\n act_cfg=act_cfg))\n\n def forward(self, x):\n\n def _inner_forward(x):\n if self.stride > 1:\n out = torch.cat((self.branch1(x), self.branch2(x)), dim=1)\n else:\n x1, x2 = x.chunk(2, dim=1)\n out = torch.cat((x1, self.branch2(x2)), dim=1)\n\n out = channel_shuffle(out, 2)\n\n return out\n\n if self.with_cp and x.requires_grad:\n out = cp.checkpoint(_inner_forward, x)\n else:\n out = _inner_forward(x)\n\n return out\n\n\n@BACKBONES.register_module()\nclass ShuffleNetV2(BaseBackbone):\n \"\"\"ShuffleNetV2 backbone.\n\n Args:\n widen_factor (float): Width multiplier - adjusts the number of\n channels in each layer by this amount. Default: 1.0.\n out_indices (Sequence[int]): Output from which stages.\n Default: (0, 1, 2, 3).\n frozen_stages (int): Stages to be frozen (all param fixed).\n Default: -1, which means not freezing any parameters.\n conv_cfg (dict): Config dict for convolution layer.\n Default: None, which means using conv2d.\n norm_cfg (dict): Config dict for normalization layer.\n Default: dict(type='BN').\n act_cfg (dict): Config dict for activation layer.\n Default: dict(type='ReLU').\n norm_eval (bool): Whether to set norm layers to eval mode, namely,\n freeze running stats (mean and var). Note: Effect on Batch Norm\n and its variants only. Default: False.\n with_cp (bool): Use checkpoint or not. Using checkpoint will save some\n memory while slowing down the training speed. Default: False.\n \"\"\"\n\n def __init__(self,\n widen_factor=1.0,\n out_indices=(3, ),\n frozen_stages=-1,\n conv_cfg=None,\n norm_cfg=dict(type='BN'),\n act_cfg=dict(type='ReLU'),\n norm_eval=False,\n with_cp=False):\n # Protect mutable default arguments\n norm_cfg = copy.deepcopy(norm_cfg)\n act_cfg = copy.deepcopy(act_cfg)\n super().__init__()\n self.stage_blocks = [4, 8, 4]\n for index in out_indices:\n if index not in range(0, 4):\n raise ValueError('the item in out_indices must in '\n f'range(0, 4). But received {index}')\n\n if frozen_stages not in range(-1, 4):\n raise ValueError('frozen_stages must be in range(-1, 4). '\n f'But received {frozen_stages}')\n self.out_indices = out_indices\n self.frozen_stages = frozen_stages\n self.conv_cfg = conv_cfg\n self.norm_cfg = norm_cfg\n self.act_cfg = act_cfg\n self.norm_eval = norm_eval\n self.with_cp = with_cp\n\n if widen_factor == 0.5:\n channels = [48, 96, 192, 1024]\n elif widen_factor == 1.0:\n channels = [116, 232, 464, 1024]\n elif widen_factor == 1.5:\n channels = [176, 352, 704, 1024]\n elif widen_factor == 2.0:\n channels = [244, 488, 976, 2048]\n else:\n raise ValueError('widen_factor must be in [0.5, 1.0, 1.5, 2.0]. '\n f'But received {widen_factor}')\n\n self.in_channels = 24\n self.conv1 = ConvModule(\n in_channels=3,\n out_channels=self.in_channels,\n kernel_size=3,\n stride=2,\n padding=1,\n conv_cfg=conv_cfg,\n norm_cfg=norm_cfg,\n act_cfg=act_cfg)\n\n self.maxpool = nn.MaxPool2d(kernel_size=3, stride=2, padding=1)\n\n self.layers = nn.ModuleList()\n for i, num_blocks in enumerate(self.stage_blocks):\n layer = self._make_layer(channels[i], num_blocks)\n self.layers.append(layer)\n\n output_channels = channels[-1]\n self.layers.append(\n ConvModule(\n in_channels=self.in_channels,\n out_channels=output_channels,\n kernel_size=1,\n conv_cfg=conv_cfg,\n norm_cfg=norm_cfg,\n act_cfg=act_cfg))\n\n def _make_layer(self, out_channels, num_blocks):\n \"\"\"Stack blocks to make a layer.\n\n Args:\n out_channels (int): out_channels of the block.\n num_blocks (int): number of blocks.\n \"\"\"\n layers = []\n for i in range(num_blocks):\n stride = 2 if i == 0 else 1\n layers.append(\n InvertedResidual(\n in_channels=self.in_channels,\n out_channels=out_channels,\n stride=stride,\n conv_cfg=self.conv_cfg,\n norm_cfg=self.norm_cfg,\n act_cfg=self.act_cfg,\n with_cp=self.with_cp))\n self.in_channels = out_channels\n\n return nn.Sequential(*layers)\n\n def _freeze_stages(self):\n if self.frozen_stages >= 0:\n for param in self.conv1.parameters():\n param.requires_grad = False\n\n for i in range(self.frozen_stages):\n m = self.layers[i]\n m.eval()\n for param in m.parameters():\n param.requires_grad = False\n\n def init_weights(self, pretrained=None):\n if isinstance(pretrained, str):\n logger = logging.getLogger()\n load_checkpoint(self, pretrained, strict=False, logger=logger)\n elif pretrained is None:\n for name, m in self.named_modules():\n if isinstance(m, nn.Conv2d):\n if 'conv1' in name:\n normal_init(m, mean=0, std=0.01)\n else:\n normal_init(m, mean=0, std=1.0 / m.weight.shape[1])\n elif isinstance(m, (_BatchNorm, nn.GroupNorm)):\n constant_init(m.weight, val=1, bias=0.0001)\n if isinstance(m, _BatchNorm):\n if m.running_mean is not None:\n nn.init.constant_(m.running_mean, 0)\n else:\n raise TypeError('pretrained must be a str or None. But received '\n f'{type(pretrained)}')\n\n def forward(self, x):\n x = self.conv1(x)\n x = self.maxpool(x)\n\n outs = []\n for i, layer in enumerate(self.layers):\n x = layer(x)\n if i in self.out_indices:\n outs.append(x)\n\n if len(outs) == 1:\n return outs[0]\n return tuple(outs)\n\n def train(self, mode=True):\n super().train(mode)\n self._freeze_stages()\n if mode and self.norm_eval:\n for m in self.modules():\n if isinstance(m, nn.BatchNorm2d):\n m.eval()\n" }, { "path": "eval/GVHMR/hmr4d/utils/preproc/vitpose_pytorch/src/vitpose_infer/builder/backbones/tcn.py", "content": "# Copyright (c) OpenMMLab. All rights reserved.\nimport copy\n\nimport torch.nn as nn\nfrom mmcv.cnn import ConvModule, build_conv_layer, constant_init, kaiming_init\nfrom mmcv.utils.parrots_wrapper import _BatchNorm\n\nfrom mmpose.core import WeightNormClipHook\nfrom ..builder import BACKBONES\nfrom .base_backbone import BaseBackbone\n\n\nclass BasicTemporalBlock(nn.Module):\n \"\"\"Basic block for VideoPose3D.\n\n Args:\n in_channels (int): Input channels of this block.\n out_channels (int): Output channels of this block.\n mid_channels (int): The output channels of conv1. Default: 1024.\n kernel_size (int): Size of the convolving kernel. Default: 3.\n dilation (int): Spacing between kernel elements. Default: 3.\n dropout (float): Dropout rate. Default: 0.25.\n causal (bool): Use causal convolutions instead of symmetric\n convolutions (for real-time applications). Default: False.\n residual (bool): Use residual connection. Default: True.\n use_stride_conv (bool): Use optimized TCN that designed\n specifically for single-frame batching, i.e. where batches have\n input length = receptive field, and output length = 1. This\n implementation replaces dilated convolutions with strided\n convolutions to avoid generating unused intermediate results.\n Default: False.\n conv_cfg (dict): dictionary to construct and config conv layer.\n Default: dict(type='Conv1d').\n norm_cfg (dict): dictionary to construct and config norm layer.\n Default: dict(type='BN1d').\n \"\"\"\n\n def __init__(self,\n in_channels,\n out_channels,\n mid_channels=1024,\n kernel_size=3,\n dilation=3,\n dropout=0.25,\n causal=False,\n residual=True,\n use_stride_conv=False,\n conv_cfg=dict(type='Conv1d'),\n norm_cfg=dict(type='BN1d')):\n # Protect mutable default arguments\n conv_cfg = copy.deepcopy(conv_cfg)\n norm_cfg = copy.deepcopy(norm_cfg)\n super().__init__()\n self.in_channels = in_channels\n self.out_channels = out_channels\n self.mid_channels = mid_channels\n self.kernel_size = kernel_size\n self.dilation = dilation\n self.dropout = dropout\n self.causal = causal\n self.residual = residual\n self.use_stride_conv = use_stride_conv\n\n self.pad = (kernel_size - 1) * dilation // 2\n if use_stride_conv:\n self.stride = kernel_size\n self.causal_shift = kernel_size // 2 if causal else 0\n self.dilation = 1\n else:\n self.stride = 1\n self.causal_shift = kernel_size // 2 * dilation if causal else 0\n\n self.conv1 = nn.Sequential(\n ConvModule(\n in_channels,\n mid_channels,\n kernel_size=kernel_size,\n stride=self.stride,\n dilation=self.dilation,\n bias='auto',\n conv_cfg=conv_cfg,\n norm_cfg=norm_cfg))\n self.conv2 = nn.Sequential(\n ConvModule(\n mid_channels,\n out_channels,\n kernel_size=1,\n bias='auto',\n conv_cfg=conv_cfg,\n norm_cfg=norm_cfg))\n\n if residual and in_channels != out_channels:\n self.short_cut = build_conv_layer(conv_cfg, in_channels,\n out_channels, 1)\n else:\n self.short_cut = None\n\n self.dropout = nn.Dropout(dropout) if dropout > 0 else None\n\n def forward(self, x):\n \"\"\"Forward function.\"\"\"\n if self.use_stride_conv:\n assert self.causal_shift + self.kernel_size // 2 < x.shape[2]\n else:\n assert 0 <= self.pad + self.causal_shift < x.shape[2] - \\\n self.pad + self.causal_shift <= x.shape[2]\n\n out = self.conv1(x)\n if self.dropout is not None:\n out = self.dropout(out)\n\n out = self.conv2(out)\n if self.dropout is not None:\n out = self.dropout(out)\n\n if self.residual:\n if self.use_stride_conv:\n res = x[:, :, self.causal_shift +\n self.kernel_size // 2::self.kernel_size]\n else:\n res = x[:, :,\n (self.pad + self.causal_shift):(x.shape[2] - self.pad +\n self.causal_shift)]\n\n if self.short_cut is not None:\n res = self.short_cut(res)\n out = out + res\n\n return out\n\n\n@BACKBONES.register_module()\nclass TCN(BaseBackbone):\n \"\"\"TCN backbone.\n\n Temporal Convolutional Networks.\n More details can be found in the\n `paper `__ .\n\n Args:\n in_channels (int): Number of input channels, which equals to\n num_keypoints * num_features.\n stem_channels (int): Number of feature channels. Default: 1024.\n num_blocks (int): NUmber of basic temporal convolutional blocks.\n Default: 2.\n kernel_sizes (Sequence[int]): Sizes of the convolving kernel of\n each basic block. Default: ``(3, 3, 3)``.\n dropout (float): Dropout rate. Default: 0.25.\n causal (bool): Use causal convolutions instead of symmetric\n convolutions (for real-time applications).\n Default: False.\n residual (bool): Use residual connection. Default: True.\n use_stride_conv (bool): Use TCN backbone optimized for\n single-frame batching, i.e. where batches have input length =\n receptive field, and output length = 1. This implementation\n replaces dilated convolutions with strided convolutions to avoid\n generating unused intermediate results. The weights are\n interchangeable with the reference implementation. Default: False\n conv_cfg (dict): dictionary to construct and config conv layer.\n Default: dict(type='Conv1d').\n norm_cfg (dict): dictionary to construct and config norm layer.\n Default: dict(type='BN1d').\n max_norm (float|None): if not None, the weight of convolution layers\n will be clipped to have a maximum norm of max_norm.\n\n Example:\n >>> from mmpose.models import TCN\n >>> import torch\n >>> self = TCN(in_channels=34)\n >>> self.eval()\n >>> inputs = torch.rand(1, 34, 243)\n >>> level_outputs = self.forward(inputs)\n >>> for level_out in level_outputs:\n ... print(tuple(level_out.shape))\n (1, 1024, 235)\n (1, 1024, 217)\n \"\"\"\n\n def __init__(self,\n in_channels,\n stem_channels=1024,\n num_blocks=2,\n kernel_sizes=(3, 3, 3),\n dropout=0.25,\n causal=False,\n residual=True,\n use_stride_conv=False,\n conv_cfg=dict(type='Conv1d'),\n norm_cfg=dict(type='BN1d'),\n max_norm=None):\n # Protect mutable default arguments\n conv_cfg = copy.deepcopy(conv_cfg)\n norm_cfg = copy.deepcopy(norm_cfg)\n super().__init__()\n self.in_channels = in_channels\n self.stem_channels = stem_channels\n self.num_blocks = num_blocks\n self.kernel_sizes = kernel_sizes\n self.dropout = dropout\n self.causal = causal\n self.residual = residual\n self.use_stride_conv = use_stride_conv\n self.max_norm = max_norm\n\n assert num_blocks == len(kernel_sizes) - 1\n for ks in kernel_sizes:\n assert ks % 2 == 1, 'Only odd filter widths are supported.'\n\n self.expand_conv = ConvModule(\n in_channels,\n stem_channels,\n kernel_size=kernel_sizes[0],\n stride=kernel_sizes[0] if use_stride_conv else 1,\n bias='auto',\n conv_cfg=conv_cfg,\n norm_cfg=norm_cfg)\n\n dilation = kernel_sizes[0]\n self.tcn_blocks = nn.ModuleList()\n for i in range(1, num_blocks + 1):\n self.tcn_blocks.append(\n BasicTemporalBlock(\n in_channels=stem_channels,\n out_channels=stem_channels,\n mid_channels=stem_channels,\n kernel_size=kernel_sizes[i],\n dilation=dilation,\n dropout=dropout,\n causal=causal,\n residual=residual,\n use_stride_conv=use_stride_conv,\n conv_cfg=conv_cfg,\n norm_cfg=norm_cfg))\n dilation *= kernel_sizes[i]\n\n if self.max_norm is not None:\n # Apply weight norm clip to conv layers\n weight_clip = WeightNormClipHook(self.max_norm)\n for module in self.modules():\n if isinstance(module, nn.modules.conv._ConvNd):\n weight_clip.register(module)\n\n self.dropout = nn.Dropout(dropout) if dropout > 0 else None\n\n def forward(self, x):\n \"\"\"Forward function.\"\"\"\n x = self.expand_conv(x)\n\n if self.dropout is not None:\n x = self.dropout(x)\n\n outs = []\n for i in range(self.num_blocks):\n x = self.tcn_blocks[i](x)\n outs.append(x)\n\n return tuple(outs)\n\n def init_weights(self, pretrained=None):\n \"\"\"Initialize the weights.\"\"\"\n super().init_weights(pretrained)\n if pretrained is None:\n for m in self.modules():\n if isinstance(m, nn.modules.conv._ConvNd):\n kaiming_init(m, mode='fan_in', nonlinearity='relu')\n elif isinstance(m, _BatchNorm):\n constant_init(m, 1)\n" }, { "path": "eval/GVHMR/hmr4d/utils/preproc/vitpose_pytorch/src/vitpose_infer/builder/backbones/test_torch.py", "content": "import torch\nimport torch.nn as nn\nimport torch.nn.functional as F\n\n\nclass Net(nn.Module):\n\n def __init__(self):\n super(Net, self).__init__()\n # 1 input image channel, 6 output channels, 5x5 square convolution\n # kernel\n self.conv1 = nn.Conv2d(1, 6, 5)\n self.conv2 = nn.Conv2d(6, 16, 5)\n # an affine operation: y = Wx + b\n self.fc1 = nn.Linear(16 * 5 * 5, 120) # 5*5 from image dimension\n self.fc2 = nn.Linear(120, 84)\n self.fc3 = nn.Linear(84, 10)\n\n def forward(self, x):\n # Max pooling over a (2, 2) window\n x = F.max_pool2d(F.relu(self.conv1(x)), (2, 2))\n # If the size is a square, you can specify with a single number\n x = F.max_pool2d(F.relu(self.conv2(x)), 2)\n x = torch.flatten(x, 1) # flatten all dimensions except the batch dimension\n x = F.relu(self.fc1(x))\n x = F.relu(self.fc2(x))\n x = self.fc3(x)\n return x\n\n\nnet = Net()\n# print(net)\n\nnet.train()\n\ninput = torch.randn(1, 1, 32, 32)\n# out = net(input)\n# print(out)\noutput = net(input)\ntarget = torch.randn(10) # a dummy target, for example\ntarget = target.view(1, -1) # make it the same shape as output\ncriterion = nn.MSELoss()\n\n# loss = criterion(output.cuda(), target.cuda())\n\nimport torch.optim as optim\n\n# create your optimizer\noptimizer = optim.SGD(net.parameters(), lr=0.01)\n\n# in your training loop:\noptimizer.zero_grad() # zero the gradient buffers\noutput = net(input)\nloss = criterion(output, target)\n\nloss.backward()\n\noptimizer.step() \n\n# print(loss)" }, { "path": "eval/GVHMR/hmr4d/utils/preproc/vitpose_pytorch/src/vitpose_infer/builder/backbones/utils/__init__.py", "content": "# Copyright (c) OpenMMLab. All rights reserved.\nfrom .channel_shuffle import channel_shuffle\nfrom .inverted_residual import InvertedResidual\nfrom .make_divisible import make_divisible\nfrom .se_layer import SELayer\nfrom .utils import load_checkpoint\n\n__all__ = [\n 'channel_shuffle', 'make_divisible', 'InvertedResidual', 'SELayer',\n 'load_checkpoint'\n]\n" }, { "path": "eval/GVHMR/hmr4d/utils/preproc/vitpose_pytorch/src/vitpose_infer/builder/backbones/utils/channel_shuffle.py", "content": "# Copyright (c) OpenMMLab. All rights reserved.\nimport torch\n\n\ndef channel_shuffle(x, groups):\n \"\"\"Channel Shuffle operation.\n\n This function enables cross-group information flow for multiple groups\n convolution layers.\n\n Args:\n x (Tensor): The input tensor.\n groups (int): The number of groups to divide the input tensor\n in the channel dimension.\n\n Returns:\n Tensor: The output tensor after channel shuffle operation.\n \"\"\"\n\n batch_size, num_channels, height, width = x.size()\n assert (num_channels % groups == 0), ('num_channels should be '\n 'divisible by groups')\n channels_per_group = num_channels // groups\n\n x = x.view(batch_size, groups, channels_per_group, height, width)\n x = torch.transpose(x, 1, 2).contiguous()\n x = x.view(batch_size, -1, height, width)\n\n return x\n" }, { "path": "eval/GVHMR/hmr4d/utils/preproc/vitpose_pytorch/src/vitpose_infer/builder/backbones/utils/inverted_residual.py", "content": "# Copyright (c) OpenMMLab. All rights reserved.\nimport copy\n\nimport torch.nn as nn\nimport torch.utils.checkpoint as cp\nfrom mmcv.cnn import ConvModule\n\nfrom .se_layer import SELayer\n\n\nclass InvertedResidual(nn.Module):\n \"\"\"Inverted Residual Block.\n\n Args:\n in_channels (int): The input channels of this Module.\n out_channels (int): The output channels of this Module.\n mid_channels (int): The input channels of the depthwise convolution.\n kernel_size (int): The kernel size of the depthwise convolution.\n Default: 3.\n groups (None or int): The group number of the depthwise convolution.\n Default: None, which means group number = mid_channels.\n stride (int): The stride of the depthwise convolution. Default: 1.\n se_cfg (dict): Config dict for se layer. Default: None, which means no\n se layer.\n with_expand_conv (bool): Use expand conv or not. If set False,\n mid_channels must be the same with in_channels.\n Default: True.\n conv_cfg (dict): Config dict for convolution layer. Default: None,\n which means using conv2d.\n norm_cfg (dict): Config dict for normalization layer.\n Default: dict(type='BN').\n act_cfg (dict): Config dict for activation layer.\n Default: dict(type='ReLU').\n with_cp (bool): Use checkpoint or not. Using checkpoint will save some\n memory while slowing down the training speed. Default: False.\n\n Returns:\n Tensor: The output tensor.\n \"\"\"\n\n def __init__(self,\n in_channels,\n out_channels,\n mid_channels,\n kernel_size=3,\n groups=None,\n stride=1,\n se_cfg=None,\n with_expand_conv=True,\n conv_cfg=None,\n norm_cfg=dict(type='BN'),\n act_cfg=dict(type='ReLU'),\n with_cp=False):\n # Protect mutable default arguments\n norm_cfg = copy.deepcopy(norm_cfg)\n act_cfg = copy.deepcopy(act_cfg)\n super().__init__()\n self.with_res_shortcut = (stride == 1 and in_channels == out_channels)\n assert stride in [1, 2]\n self.with_cp = with_cp\n self.with_se = se_cfg is not None\n self.with_expand_conv = with_expand_conv\n\n if groups is None:\n groups = mid_channels\n\n if self.with_se:\n assert isinstance(se_cfg, dict)\n if not self.with_expand_conv:\n assert mid_channels == in_channels\n\n if self.with_expand_conv:\n self.expand_conv = ConvModule(\n in_channels=in_channels,\n out_channels=mid_channels,\n kernel_size=1,\n stride=1,\n padding=0,\n conv_cfg=conv_cfg,\n norm_cfg=norm_cfg,\n act_cfg=act_cfg)\n self.depthwise_conv = ConvModule(\n in_channels=mid_channels,\n out_channels=mid_channels,\n kernel_size=kernel_size,\n stride=stride,\n padding=kernel_size // 2,\n groups=groups,\n conv_cfg=conv_cfg,\n norm_cfg=norm_cfg,\n act_cfg=act_cfg)\n if self.with_se:\n self.se = SELayer(**se_cfg)\n self.linear_conv = ConvModule(\n in_channels=mid_channels,\n out_channels=out_channels,\n kernel_size=1,\n stride=1,\n padding=0,\n conv_cfg=conv_cfg,\n norm_cfg=norm_cfg,\n act_cfg=None)\n\n def forward(self, x):\n\n def _inner_forward(x):\n out = x\n\n if self.with_expand_conv:\n out = self.expand_conv(out)\n\n out = self.depthwise_conv(out)\n\n if self.with_se:\n out = self.se(out)\n\n out = self.linear_conv(out)\n\n if self.with_res_shortcut:\n return x + out\n return out\n\n if self.with_cp and x.requires_grad:\n out = cp.checkpoint(_inner_forward, x)\n else:\n out = _inner_forward(x)\n\n return out\n" }, { "path": "eval/GVHMR/hmr4d/utils/preproc/vitpose_pytorch/src/vitpose_infer/builder/backbones/utils/make_divisible.py", "content": "# Copyright (c) OpenMMLab. All rights reserved.\ndef make_divisible(value, divisor, min_value=None, min_ratio=0.9):\n \"\"\"Make divisible function.\n\n This function rounds the channel number down to the nearest value that can\n be divisible by the divisor.\n\n Args:\n value (int): The original channel number.\n divisor (int): The divisor to fully divide the channel number.\n min_value (int, optional): The minimum value of the output channel.\n Default: None, means that the minimum value equal to the divisor.\n min_ratio (float, optional): The minimum ratio of the rounded channel\n number to the original channel number. Default: 0.9.\n Returns:\n int: The modified output channel number\n \"\"\"\n\n if min_value is None:\n min_value = divisor\n new_value = max(min_value, int(value + divisor / 2) // divisor * divisor)\n # Make sure that round down does not go down by more than (1-min_ratio).\n if new_value < min_ratio * value:\n new_value += divisor\n return new_value\n" }, { "path": "eval/GVHMR/hmr4d/utils/preproc/vitpose_pytorch/src/vitpose_infer/builder/backbones/utils/se_layer.py", "content": "# Copyright (c) OpenMMLab. All rights reserved.\nimport mmcv\nimport torch.nn as nn\nfrom mmcv.cnn import ConvModule\n\n\nclass SELayer(nn.Module):\n \"\"\"Squeeze-and-Excitation Module.\n\n Args:\n channels (int): The input (and output) channels of the SE layer.\n ratio (int): Squeeze ratio in SELayer, the intermediate channel will be\n ``int(channels/ratio)``. Default: 16.\n conv_cfg (None or dict): Config dict for convolution layer.\n Default: None, which means using conv2d.\n act_cfg (dict or Sequence[dict]): Config dict for activation layer.\n If act_cfg is a dict, two activation layers will be configurated\n by this dict. If act_cfg is a sequence of dicts, the first\n activation layer will be configurated by the first dict and the\n second activation layer will be configurated by the second dict.\n Default: (dict(type='ReLU'), dict(type='Sigmoid'))\n \"\"\"\n\n def __init__(self,\n channels,\n ratio=16,\n conv_cfg=None,\n act_cfg=(dict(type='ReLU'), dict(type='Sigmoid'))):\n super().__init__()\n if isinstance(act_cfg, dict):\n act_cfg = (act_cfg, act_cfg)\n assert len(act_cfg) == 2\n assert mmcv.is_tuple_of(act_cfg, dict)\n self.global_avgpool = nn.AdaptiveAvgPool2d(1)\n self.conv1 = ConvModule(\n in_channels=channels,\n out_channels=int(channels / ratio),\n kernel_size=1,\n stride=1,\n conv_cfg=conv_cfg,\n act_cfg=act_cfg[0])\n self.conv2 = ConvModule(\n in_channels=int(channels / ratio),\n out_channels=channels,\n kernel_size=1,\n stride=1,\n conv_cfg=conv_cfg,\n act_cfg=act_cfg[1])\n\n def forward(self, x):\n out = self.global_avgpool(x)\n out = self.conv1(out)\n out = self.conv2(out)\n return x * out\n" }, { "path": "eval/GVHMR/hmr4d/utils/preproc/vitpose_pytorch/src/vitpose_infer/builder/backbones/utils/utils.py", "content": "# Copyright (c) OpenMMLab. All rights reserved.\nfrom collections import OrderedDict\n\nfrom mmcv.runner.checkpoint import _load_checkpoint, load_state_dict\n\n\ndef load_checkpoint(model,\n filename,\n map_location='cpu',\n strict=False,\n logger=None):\n \"\"\"Load checkpoint from a file or URI.\n\n Args:\n model (Module): Module to load checkpoint.\n filename (str): Accept local filepath, URL, ``torchvision://xxx``,\n ``open-mmlab://xxx``.\n map_location (str): Same as :func:`torch.load`.\n strict (bool): Whether to allow different params for the model and\n checkpoint.\n logger (:mod:`logging.Logger` or None): The logger for error message.\n\n Returns:\n dict or OrderedDict: The loaded checkpoint.\n \"\"\"\n checkpoint = _load_checkpoint(filename, map_location)\n # OrderedDict is a subclass of dict\n if not isinstance(checkpoint, dict):\n raise RuntimeError(\n f'No state_dict found in checkpoint file {filename}')\n # get state_dict from checkpoint\n if 'state_dict' in checkpoint:\n state_dict_tmp = checkpoint['state_dict']\n else:\n state_dict_tmp = checkpoint\n\n state_dict = OrderedDict()\n # strip prefix of state_dict\n for k, v in state_dict_tmp.items():\n if k.startswith('module.backbone.'):\n state_dict[k[16:]] = v\n elif k.startswith('module.'):\n state_dict[k[7:]] = v\n elif k.startswith('backbone.'):\n state_dict[k[9:]] = v\n else:\n state_dict[k] = v\n # load state_dict\n load_state_dict(model, state_dict, strict, logger)\n return checkpoint\n\n\ndef get_state_dict(filename, map_location='cpu'):\n \"\"\"Get state_dict from a file or URI.\n\n Args:\n filename (str): Accept local filepath, URL, ``torchvision://xxx``,\n ``open-mmlab://xxx``.\n map_location (str): Same as :func:`torch.load`.\n\n Returns:\n OrderedDict: The state_dict.\n \"\"\"\n checkpoint = _load_checkpoint(filename, map_location)\n # OrderedDict is a subclass of dict\n if not isinstance(checkpoint, dict):\n raise RuntimeError(\n f'No state_dict found in checkpoint file {filename}')\n # get state_dict from checkpoint\n if 'state_dict' in checkpoint:\n state_dict_tmp = checkpoint['state_dict']\n else:\n state_dict_tmp = checkpoint\n\n state_dict = OrderedDict()\n # strip prefix of state_dict\n for k, v in state_dict_tmp.items():\n if k.startswith('module.backbone.'):\n state_dict[k[16:]] = v\n elif k.startswith('module.'):\n state_dict[k[7:]] = v\n elif k.startswith('backbone.'):\n state_dict[k[9:]] = v\n else:\n state_dict[k] = v\n\n return state_dict\n" }, { "path": "eval/GVHMR/hmr4d/utils/preproc/vitpose_pytorch/src/vitpose_infer/builder/backbones/vgg.py", "content": "# Copyright (c) OpenMMLab. All rights reserved.\nimport torch.nn as nn\nfrom mmcv.cnn import ConvModule, constant_init, kaiming_init, normal_init\nfrom mmcv.utils.parrots_wrapper import _BatchNorm\n\nfrom ..builder import BACKBONES\nfrom .base_backbone import BaseBackbone\n\n\ndef make_vgg_layer(in_channels,\n out_channels,\n num_blocks,\n conv_cfg=None,\n norm_cfg=None,\n act_cfg=dict(type='ReLU'),\n dilation=1,\n with_norm=False,\n ceil_mode=False):\n layers = []\n for _ in range(num_blocks):\n layer = ConvModule(\n in_channels=in_channels,\n out_channels=out_channels,\n kernel_size=3,\n dilation=dilation,\n padding=dilation,\n bias=True,\n conv_cfg=conv_cfg,\n norm_cfg=norm_cfg,\n act_cfg=act_cfg)\n layers.append(layer)\n in_channels = out_channels\n layers.append(nn.MaxPool2d(kernel_size=2, stride=2, ceil_mode=ceil_mode))\n\n return layers\n\n\n@BACKBONES.register_module()\nclass VGG(BaseBackbone):\n \"\"\"VGG backbone.\n\n Args:\n depth (int): Depth of vgg, from {11, 13, 16, 19}.\n with_norm (bool): Use BatchNorm or not.\n num_classes (int): number of classes for classification.\n num_stages (int): VGG stages, normally 5.\n dilations (Sequence[int]): Dilation of each stage.\n out_indices (Sequence[int]): Output from which stages. If only one\n stage is specified, a single tensor (feature map) is returned,\n otherwise multiple stages are specified, a tuple of tensors will\n be returned. When it is None, the default behavior depends on\n whether num_classes is specified. If num_classes <= 0, the default\n value is (4, ), outputting the last feature map before classifier.\n If num_classes > 0, the default value is (5, ), outputting the\n classification score. Default: None.\n frozen_stages (int): Stages to be frozen (all param fixed). -1 means\n not freezing any parameters.\n norm_eval (bool): Whether to set norm layers to eval mode, namely,\n freeze running stats (mean and var). Note: Effect on Batch Norm\n and its variants only. Default: False.\n ceil_mode (bool): Whether to use ceil_mode of MaxPool. Default: False.\n with_last_pool (bool): Whether to keep the last pooling before\n classifier. Default: True.\n \"\"\"\n\n # Parameters to build layers. Each element specifies the number of conv in\n # each stage. For example, VGG11 contains 11 layers with learnable\n # parameters. 11 is computed as 11 = (1 + 1 + 2 + 2 + 2) + 3,\n # where 3 indicates the last three fully-connected layers.\n arch_settings = {\n 11: (1, 1, 2, 2, 2),\n 13: (2, 2, 2, 2, 2),\n 16: (2, 2, 3, 3, 3),\n 19: (2, 2, 4, 4, 4)\n }\n\n def __init__(self,\n depth,\n num_classes=-1,\n num_stages=5,\n dilations=(1, 1, 1, 1, 1),\n out_indices=None,\n frozen_stages=-1,\n conv_cfg=None,\n norm_cfg=None,\n act_cfg=dict(type='ReLU'),\n norm_eval=False,\n ceil_mode=False,\n with_last_pool=True):\n super().__init__()\n if depth not in self.arch_settings:\n raise KeyError(f'invalid depth {depth} for vgg')\n assert num_stages >= 1 and num_stages <= 5\n stage_blocks = self.arch_settings[depth]\n self.stage_blocks = stage_blocks[:num_stages]\n assert len(dilations) == num_stages\n\n self.num_classes = num_classes\n self.frozen_stages = frozen_stages\n self.norm_eval = norm_eval\n with_norm = norm_cfg is not None\n\n if out_indices is None:\n out_indices = (5, ) if num_classes > 0 else (4, )\n assert max(out_indices) <= num_stages\n self.out_indices = out_indices\n\n self.in_channels = 3\n start_idx = 0\n vgg_layers = []\n self.range_sub_modules = []\n for i, num_blocks in enumerate(self.stage_blocks):\n num_modules = num_blocks + 1\n end_idx = start_idx + num_modules\n dilation = dilations[i]\n out_channels = 64 * 2**i if i < 4 else 512\n vgg_layer = make_vgg_layer(\n self.in_channels,\n out_channels,\n num_blocks,\n conv_cfg=conv_cfg,\n norm_cfg=norm_cfg,\n act_cfg=act_cfg,\n dilation=dilation,\n with_norm=with_norm,\n ceil_mode=ceil_mode)\n vgg_layers.extend(vgg_layer)\n self.in_channels = out_channels\n self.range_sub_modules.append([start_idx, end_idx])\n start_idx = end_idx\n if not with_last_pool:\n vgg_layers.pop(-1)\n self.range_sub_modules[-1][1] -= 1\n self.module_name = 'features'\n self.add_module(self.module_name, nn.Sequential(*vgg_layers))\n\n if self.num_classes > 0:\n self.classifier = nn.Sequential(\n nn.Linear(512 * 7 * 7, 4096),\n nn.ReLU(True),\n nn.Dropout(),\n nn.Linear(4096, 4096),\n nn.ReLU(True),\n nn.Dropout(),\n nn.Linear(4096, num_classes),\n )\n\n def init_weights(self, pretrained=None):\n super().init_weights(pretrained)\n if pretrained is None:\n for m in self.modules():\n if isinstance(m, nn.Conv2d):\n kaiming_init(m)\n elif isinstance(m, _BatchNorm):\n constant_init(m, 1)\n elif isinstance(m, nn.Linear):\n normal_init(m, std=0.01)\n\n def forward(self, x):\n outs = []\n vgg_layers = getattr(self, self.module_name)\n for i in range(len(self.stage_blocks)):\n for j in range(*self.range_sub_modules[i]):\n vgg_layer = vgg_layers[j]\n x = vgg_layer(x)\n if i in self.out_indices:\n outs.append(x)\n if self.num_classes > 0:\n x = x.view(x.size(0), -1)\n x = self.classifier(x)\n outs.append(x)\n if len(outs) == 1:\n return outs[0]\n else:\n return tuple(outs)\n\n def _freeze_stages(self):\n vgg_layers = getattr(self, self.module_name)\n for i in range(self.frozen_stages):\n for j in range(*self.range_sub_modules[i]):\n m = vgg_layers[j]\n m.eval()\n for param in m.parameters():\n param.requires_grad = False\n\n def train(self, mode=True):\n super().train(mode)\n self._freeze_stages()\n if mode and self.norm_eval:\n for m in self.modules():\n # trick: eval have effect on BatchNorm only\n if isinstance(m, _BatchNorm):\n m.eval()\n" }, { "path": "eval/GVHMR/hmr4d/utils/preproc/vitpose_pytorch/src/vitpose_infer/builder/backbones/vipnas_mbv3.py", "content": "# Copyright (c) OpenMMLab. All rights reserved.\nimport copy\nimport logging\n\nimport torch.nn as nn\nfrom mmcv.cnn import ConvModule\nfrom torch.nn.modules.batchnorm import _BatchNorm\n\nfrom ..builder import BACKBONES\nfrom .base_backbone import BaseBackbone\nfrom .utils import InvertedResidual, load_checkpoint\n\n\n@BACKBONES.register_module()\nclass ViPNAS_MobileNetV3(BaseBackbone):\n \"\"\"ViPNAS_MobileNetV3 backbone.\n\n \"ViPNAS: Efficient Video Pose Estimation via Neural Architecture Search\"\n More details can be found in the `paper\n `__ .\n\n Args:\n wid (list(int)): Searched width config for each stage.\n expan (list(int)): Searched expansion ratio config for each stage.\n dep (list(int)): Searched depth config for each stage.\n ks (list(int)): Searched kernel size config for each stage.\n group (list(int)): Searched group number config for each stage.\n att (list(bool)): Searched attention config for each stage.\n stride (list(int)): Stride config for each stage.\n act (list(dict)): Activation config for each stage.\n conv_cfg (dict): Config dict for convolution layer.\n Default: None, which means using conv2d.\n norm_cfg (dict): Config dict for normalization layer.\n Default: dict(type='BN').\n frozen_stages (int): Stages to be frozen (all param fixed).\n Default: -1, which means not freezing any parameters.\n norm_eval (bool): Whether to set norm layers to eval mode, namely,\n freeze running stats (mean and var). Note: Effect on Batch Norm\n and its variants only. Default: False.\n with_cp (bool): Use checkpoint or not. Using checkpoint will save\n some memory while slowing down the training speed.\n Default: False.\n \"\"\"\n\n def __init__(self,\n wid=[16, 16, 24, 40, 80, 112, 160],\n expan=[None, 1, 5, 4, 5, 5, 6],\n dep=[None, 1, 4, 4, 4, 4, 4],\n ks=[3, 3, 7, 7, 5, 7, 5],\n group=[None, 8, 120, 20, 100, 280, 240],\n att=[None, True, True, False, True, True, True],\n stride=[2, 1, 2, 2, 2, 1, 2],\n act=[\n 'HSwish', 'ReLU', 'ReLU', 'ReLU', 'HSwish', 'HSwish',\n 'HSwish'\n ],\n conv_cfg=None,\n norm_cfg=dict(type='BN'),\n frozen_stages=-1,\n norm_eval=False,\n with_cp=False):\n # Protect mutable default arguments\n norm_cfg = copy.deepcopy(norm_cfg)\n super().__init__()\n self.wid = wid\n self.expan = expan\n self.dep = dep\n self.ks = ks\n self.group = group\n self.att = att\n self.stride = stride\n self.act = act\n self.conv_cfg = conv_cfg\n self.norm_cfg = norm_cfg\n self.frozen_stages = frozen_stages\n self.norm_eval = norm_eval\n self.with_cp = with_cp\n\n self.conv1 = ConvModule(\n in_channels=3,\n out_channels=self.wid[0],\n kernel_size=self.ks[0],\n stride=self.stride[0],\n padding=self.ks[0] // 2,\n conv_cfg=conv_cfg,\n norm_cfg=norm_cfg,\n act_cfg=dict(type=self.act[0]))\n\n self.layers = self._make_layer()\n\n def _make_layer(self):\n layers = []\n layer_index = 0\n for i, dep in enumerate(self.dep[1:]):\n mid_channels = self.wid[i + 1] * self.expan[i + 1]\n\n if self.att[i + 1]:\n se_cfg = dict(\n channels=mid_channels,\n ratio=4,\n act_cfg=(dict(type='ReLU'), dict(type='HSigmoid')))\n else:\n se_cfg = None\n\n if self.expan[i + 1] == 1:\n with_expand_conv = False\n else:\n with_expand_conv = True\n\n for j in range(dep):\n if j == 0:\n stride = self.stride[i + 1]\n in_channels = self.wid[i]\n else:\n stride = 1\n in_channels = self.wid[i + 1]\n\n layer = InvertedResidual(\n in_channels=in_channels,\n out_channels=self.wid[i + 1],\n mid_channels=mid_channels,\n kernel_size=self.ks[i + 1],\n groups=self.group[i + 1],\n stride=stride,\n se_cfg=se_cfg,\n with_expand_conv=with_expand_conv,\n conv_cfg=self.conv_cfg,\n norm_cfg=self.norm_cfg,\n act_cfg=dict(type=self.act[i + 1]),\n with_cp=self.with_cp)\n layer_index += 1\n layer_name = f'layer{layer_index}'\n self.add_module(layer_name, layer)\n layers.append(layer_name)\n return layers\n\n def init_weights(self, pretrained=None):\n if isinstance(pretrained, str):\n logger = logging.getLogger()\n load_checkpoint(self, pretrained, strict=False, logger=logger)\n elif pretrained is None:\n for m in self.modules():\n if isinstance(m, nn.Conv2d):\n nn.init.normal_(m.weight, std=0.001)\n for name, _ in m.named_parameters():\n if name in ['bias']:\n nn.init.constant_(m.bias, 0)\n elif isinstance(m, nn.BatchNorm2d):\n nn.init.constant_(m.weight, 1)\n nn.init.constant_(m.bias, 0)\n else:\n raise TypeError('pretrained must be a str or None')\n\n def forward(self, x):\n x = self.conv1(x)\n\n for i, layer_name in enumerate(self.layers):\n layer = getattr(self, layer_name)\n x = layer(x)\n\n return x\n\n def _freeze_stages(self):\n if self.frozen_stages >= 0:\n for param in self.conv1.parameters():\n param.requires_grad = False\n for i in range(1, self.frozen_stages + 1):\n layer = getattr(self, f'layer{i}')\n layer.eval()\n for param in layer.parameters():\n param.requires_grad = False\n\n def train(self, mode=True):\n super().train(mode)\n self._freeze_stages()\n if mode and self.norm_eval:\n for m in self.modules():\n if isinstance(m, _BatchNorm):\n m.eval()\n" }, { "path": "eval/GVHMR/hmr4d/utils/preproc/vitpose_pytorch/src/vitpose_infer/builder/backbones/vipnas_resnet.py", "content": "# Copyright (c) OpenMMLab. All rights reserved.\nimport copy\n\nimport torch.nn as nn\nimport torch.utils.checkpoint as cp\nfrom mmcv.cnn import ConvModule, build_conv_layer, build_norm_layer\nfrom mmcv.cnn.bricks import ContextBlock\nfrom mmcv.utils.parrots_wrapper import _BatchNorm\n\nfrom ..builder import BACKBONES\nfrom .base_backbone import BaseBackbone\n\n\nclass ViPNAS_Bottleneck(nn.Module):\n \"\"\"Bottleneck block for ViPNAS_ResNet.\n\n Args:\n in_channels (int): Input channels of this block.\n out_channels (int): Output channels of this block.\n expansion (int): The ratio of ``out_channels/mid_channels`` where\n ``mid_channels`` is the input/output channels of conv2. Default: 4.\n stride (int): stride of the block. Default: 1\n dilation (int): dilation of convolution. Default: 1\n downsample (nn.Module): downsample operation on identity branch.\n Default: None.\n style (str): ``\"pytorch\"`` or ``\"caffe\"``. If set to \"pytorch\", the\n stride-two layer is the 3x3 conv layer, otherwise the stride-two\n layer is the first 1x1 conv layer. Default: \"pytorch\".\n with_cp (bool): Use checkpoint or not. Using checkpoint will save some\n memory while slowing down the training speed.\n conv_cfg (dict): dictionary to construct and config conv layer.\n Default: None\n norm_cfg (dict): dictionary to construct and config norm layer.\n Default: dict(type='BN')\n kernel_size (int): kernel size of conv2 searched in ViPANS.\n groups (int): group number of conv2 searched in ViPNAS.\n attention (bool): whether to use attention module in the end of\n the block.\n \"\"\"\n\n def __init__(self,\n in_channels,\n out_channels,\n expansion=4,\n stride=1,\n dilation=1,\n downsample=None,\n style='pytorch',\n with_cp=False,\n conv_cfg=None,\n norm_cfg=dict(type='BN'),\n kernel_size=3,\n groups=1,\n attention=False):\n # Protect mutable default arguments\n norm_cfg = copy.deepcopy(norm_cfg)\n super().__init__()\n assert style in ['pytorch', 'caffe']\n\n self.in_channels = in_channels\n self.out_channels = out_channels\n self.expansion = expansion\n assert out_channels % expansion == 0\n self.mid_channels = out_channels // expansion\n self.stride = stride\n self.dilation = dilation\n self.style = style\n self.with_cp = with_cp\n self.conv_cfg = conv_cfg\n self.norm_cfg = norm_cfg\n\n if self.style == 'pytorch':\n self.conv1_stride = 1\n self.conv2_stride = stride\n else:\n self.conv1_stride = stride\n self.conv2_stride = 1\n\n self.norm1_name, norm1 = build_norm_layer(\n norm_cfg, self.mid_channels, postfix=1)\n self.norm2_name, norm2 = build_norm_layer(\n norm_cfg, self.mid_channels, postfix=2)\n self.norm3_name, norm3 = build_norm_layer(\n norm_cfg, out_channels, postfix=3)\n\n self.conv1 = build_conv_layer(\n conv_cfg,\n in_channels,\n self.mid_channels,\n kernel_size=1,\n stride=self.conv1_stride,\n bias=False)\n self.add_module(self.norm1_name, norm1)\n self.conv2 = build_conv_layer(\n conv_cfg,\n self.mid_channels,\n self.mid_channels,\n kernel_size=kernel_size,\n stride=self.conv2_stride,\n padding=kernel_size // 2,\n groups=groups,\n dilation=dilation,\n bias=False)\n\n self.add_module(self.norm2_name, norm2)\n self.conv3 = build_conv_layer(\n conv_cfg,\n self.mid_channels,\n out_channels,\n kernel_size=1,\n bias=False)\n self.add_module(self.norm3_name, norm3)\n\n if attention:\n self.attention = ContextBlock(out_channels,\n max(1.0 / 16, 16.0 / out_channels))\n else:\n self.attention = None\n\n self.relu = nn.ReLU(inplace=True)\n self.downsample = downsample\n\n @property\n def norm1(self):\n \"\"\"nn.Module: the normalization layer named \"norm1\" \"\"\"\n return getattr(self, self.norm1_name)\n\n @property\n def norm2(self):\n \"\"\"nn.Module: the normalization layer named \"norm2\" \"\"\"\n return getattr(self, self.norm2_name)\n\n @property\n def norm3(self):\n \"\"\"nn.Module: the normalization layer named \"norm3\" \"\"\"\n return getattr(self, self.norm3_name)\n\n def forward(self, x):\n \"\"\"Forward function.\"\"\"\n\n def _inner_forward(x):\n identity = x\n\n out = self.conv1(x)\n out = self.norm1(out)\n out = self.relu(out)\n\n out = self.conv2(out)\n out = self.norm2(out)\n out = self.relu(out)\n\n out = self.conv3(out)\n out = self.norm3(out)\n\n if self.attention is not None:\n out = self.attention(out)\n\n if self.downsample is not None:\n identity = self.downsample(x)\n\n out += identity\n\n return out\n\n if self.with_cp and x.requires_grad:\n out = cp.checkpoint(_inner_forward, x)\n else:\n out = _inner_forward(x)\n\n out = self.relu(out)\n\n return out\n\n\ndef get_expansion(block, expansion=None):\n \"\"\"Get the expansion of a residual block.\n\n The block expansion will be obtained by the following order:\n\n 1. If ``expansion`` is given, just return it.\n 2. If ``block`` has the attribute ``expansion``, then return\n ``block.expansion``.\n 3. Return the default value according the the block type:\n 4 for ``ViPNAS_Bottleneck``.\n\n Args:\n block (class): The block class.\n expansion (int | None): The given expansion ratio.\n\n Returns:\n int: The expansion of the block.\n \"\"\"\n if isinstance(expansion, int):\n assert expansion > 0\n elif expansion is None:\n if hasattr(block, 'expansion'):\n expansion = block.expansion\n elif issubclass(block, ViPNAS_Bottleneck):\n expansion = 1\n else:\n raise TypeError(f'expansion is not specified for {block.__name__}')\n else:\n raise TypeError('expansion must be an integer or None')\n\n return expansion\n\n\nclass ViPNAS_ResLayer(nn.Sequential):\n \"\"\"ViPNAS_ResLayer to build ResNet style backbone.\n\n Args:\n block (nn.Module): Residual block used to build ViPNAS ResLayer.\n num_blocks (int): Number of blocks.\n in_channels (int): Input channels of this block.\n out_channels (int): Output channels of this block.\n expansion (int, optional): The expansion for BasicBlock/Bottleneck.\n If not specified, it will firstly be obtained via\n ``block.expansion``. If the block has no attribute \"expansion\",\n the following default values will be used: 1 for BasicBlock and\n 4 for Bottleneck. Default: None.\n stride (int): stride of the first block. Default: 1.\n avg_down (bool): Use AvgPool instead of stride conv when\n downsampling in the bottleneck. Default: False\n conv_cfg (dict): dictionary to construct and config conv layer.\n Default: None\n norm_cfg (dict): dictionary to construct and config norm layer.\n Default: dict(type='BN')\n downsample_first (bool): Downsample at the first block or last block.\n False for Hourglass, True for ResNet. Default: True\n kernel_size (int): Kernel Size of the corresponding convolution layer\n searched in the block.\n groups (int): Group number of the corresponding convolution layer\n searched in the block.\n attention (bool): Whether to use attention module in the end of the\n block.\n \"\"\"\n\n def __init__(self,\n block,\n num_blocks,\n in_channels,\n out_channels,\n expansion=None,\n stride=1,\n avg_down=False,\n conv_cfg=None,\n norm_cfg=dict(type='BN'),\n downsample_first=True,\n kernel_size=3,\n groups=1,\n attention=False,\n **kwargs):\n # Protect mutable default arguments\n norm_cfg = copy.deepcopy(norm_cfg)\n self.block = block\n self.expansion = get_expansion(block, expansion)\n\n downsample = None\n if stride != 1 or in_channels != out_channels:\n downsample = []\n conv_stride = stride\n if avg_down and stride != 1:\n conv_stride = 1\n downsample.append(\n nn.AvgPool2d(\n kernel_size=stride,\n stride=stride,\n ceil_mode=True,\n count_include_pad=False))\n downsample.extend([\n build_conv_layer(\n conv_cfg,\n in_channels,\n out_channels,\n kernel_size=1,\n stride=conv_stride,\n bias=False),\n build_norm_layer(norm_cfg, out_channels)[1]\n ])\n downsample = nn.Sequential(*downsample)\n\n layers = []\n if downsample_first:\n layers.append(\n block(\n in_channels=in_channels,\n out_channels=out_channels,\n expansion=self.expansion,\n stride=stride,\n downsample=downsample,\n conv_cfg=conv_cfg,\n norm_cfg=norm_cfg,\n kernel_size=kernel_size,\n groups=groups,\n attention=attention,\n **kwargs))\n in_channels = out_channels\n for _ in range(1, num_blocks):\n layers.append(\n block(\n in_channels=in_channels,\n out_channels=out_channels,\n expansion=self.expansion,\n stride=1,\n conv_cfg=conv_cfg,\n norm_cfg=norm_cfg,\n kernel_size=kernel_size,\n groups=groups,\n attention=attention,\n **kwargs))\n else: # downsample_first=False is for HourglassModule\n for i in range(0, num_blocks - 1):\n layers.append(\n block(\n in_channels=in_channels,\n out_channels=in_channels,\n expansion=self.expansion,\n stride=1,\n conv_cfg=conv_cfg,\n norm_cfg=norm_cfg,\n kernel_size=kernel_size,\n groups=groups,\n attention=attention,\n **kwargs))\n layers.append(\n block(\n in_channels=in_channels,\n out_channels=out_channels,\n expansion=self.expansion,\n stride=stride,\n downsample=downsample,\n conv_cfg=conv_cfg,\n norm_cfg=norm_cfg,\n kernel_size=kernel_size,\n groups=groups,\n attention=attention,\n **kwargs))\n\n super().__init__(*layers)\n\n\n@BACKBONES.register_module()\nclass ViPNAS_ResNet(BaseBackbone):\n \"\"\"ViPNAS_ResNet backbone.\n\n \"ViPNAS: Efficient Video Pose Estimation via Neural Architecture Search\"\n More details can be found in the `paper\n `__ .\n\n Args:\n depth (int): Network depth, from {18, 34, 50, 101, 152}.\n in_channels (int): Number of input image channels. Default: 3.\n num_stages (int): Stages of the network. Default: 4.\n strides (Sequence[int]): Strides of the first block of each stage.\n Default: ``(1, 2, 2, 2)``.\n dilations (Sequence[int]): Dilation of each stage.\n Default: ``(1, 1, 1, 1)``.\n out_indices (Sequence[int]): Output from which stages. If only one\n stage is specified, a single tensor (feature map) is returned,\n otherwise multiple stages are specified, a tuple of tensors will\n be returned. Default: ``(3, )``.\n style (str): `pytorch` or `caffe`. If set to \"pytorch\", the stride-two\n layer is the 3x3 conv layer, otherwise the stride-two layer is\n the first 1x1 conv layer.\n deep_stem (bool): Replace 7x7 conv in input stem with 3 3x3 conv.\n Default: False.\n avg_down (bool): Use AvgPool instead of stride conv when\n downsampling in the bottleneck. Default: False.\n frozen_stages (int): Stages to be frozen (stop grad and set eval mode).\n -1 means not freezing any parameters. Default: -1.\n conv_cfg (dict | None): The config dict for conv layers. Default: None.\n norm_cfg (dict): The config dict for norm layers.\n norm_eval (bool): Whether to set norm layers to eval mode, namely,\n freeze running stats (mean and var). Note: Effect on Batch Norm\n and its variants only. Default: False.\n with_cp (bool): Use checkpoint or not. Using checkpoint will save some\n memory while slowing down the training speed. Default: False.\n zero_init_residual (bool): Whether to use zero init for last norm layer\n in resblocks to let them behave as identity. Default: True.\n wid (list(int)): Searched width config for each stage.\n expan (list(int)): Searched expansion ratio config for each stage.\n dep (list(int)): Searched depth config for each stage.\n ks (list(int)): Searched kernel size config for each stage.\n group (list(int)): Searched group number config for each stage.\n att (list(bool)): Searched attention config for each stage.\n \"\"\"\n\n arch_settings = {\n 50: ViPNAS_Bottleneck,\n }\n\n def __init__(self,\n depth,\n in_channels=3,\n num_stages=4,\n strides=(1, 2, 2, 2),\n dilations=(1, 1, 1, 1),\n out_indices=(3, ),\n style='pytorch',\n deep_stem=False,\n avg_down=False,\n frozen_stages=-1,\n conv_cfg=None,\n norm_cfg=dict(type='BN', requires_grad=True),\n norm_eval=False,\n with_cp=False,\n zero_init_residual=True,\n wid=[48, 80, 160, 304, 608],\n expan=[None, 1, 1, 1, 1],\n dep=[None, 4, 6, 7, 3],\n ks=[7, 3, 5, 5, 5],\n group=[None, 16, 16, 16, 16],\n att=[None, True, False, True, True]):\n # Protect mutable default arguments\n norm_cfg = copy.deepcopy(norm_cfg)\n super().__init__()\n if depth not in self.arch_settings:\n raise KeyError(f'invalid depth {depth} for resnet')\n self.depth = depth\n self.stem_channels = dep[0]\n self.num_stages = num_stages\n assert 1 <= num_stages <= 4\n self.strides = strides\n self.dilations = dilations\n assert len(strides) == len(dilations) == num_stages\n self.out_indices = out_indices\n assert max(out_indices) < num_stages\n self.style = style\n self.deep_stem = deep_stem\n self.avg_down = avg_down\n self.frozen_stages = frozen_stages\n self.conv_cfg = conv_cfg\n self.norm_cfg = norm_cfg\n self.with_cp = with_cp\n self.norm_eval = norm_eval\n self.zero_init_residual = zero_init_residual\n self.block = self.arch_settings[depth]\n self.stage_blocks = dep[1:1 + num_stages]\n\n self._make_stem_layer(in_channels, wid[0], ks[0])\n\n self.res_layers = []\n _in_channels = wid[0]\n for i, num_blocks in enumerate(self.stage_blocks):\n expansion = get_expansion(self.block, expan[i + 1])\n _out_channels = wid[i + 1] * expansion\n stride = strides[i]\n dilation = dilations[i]\n res_layer = self.make_res_layer(\n block=self.block,\n num_blocks=num_blocks,\n in_channels=_in_channels,\n out_channels=_out_channels,\n expansion=expansion,\n stride=stride,\n dilation=dilation,\n style=self.style,\n avg_down=self.avg_down,\n with_cp=with_cp,\n conv_cfg=conv_cfg,\n norm_cfg=norm_cfg,\n kernel_size=ks[i + 1],\n groups=group[i + 1],\n attention=att[i + 1])\n _in_channels = _out_channels\n layer_name = f'layer{i + 1}'\n self.add_module(layer_name, res_layer)\n self.res_layers.append(layer_name)\n\n self._freeze_stages()\n\n self.feat_dim = res_layer[-1].out_channels\n\n def make_res_layer(self, **kwargs):\n \"\"\"Make a ViPNAS ResLayer.\"\"\"\n return ViPNAS_ResLayer(**kwargs)\n\n @property\n def norm1(self):\n \"\"\"nn.Module: the normalization layer named \"norm1\" \"\"\"\n return getattr(self, self.norm1_name)\n\n def _make_stem_layer(self, in_channels, stem_channels, kernel_size):\n \"\"\"Make stem layer.\"\"\"\n if self.deep_stem:\n self.stem = nn.Sequential(\n ConvModule(\n in_channels,\n stem_channels // 2,\n kernel_size=3,\n stride=2,\n padding=1,\n conv_cfg=self.conv_cfg,\n norm_cfg=self.norm_cfg,\n inplace=True),\n ConvModule(\n stem_channels // 2,\n stem_channels // 2,\n kernel_size=3,\n stride=1,\n padding=1,\n conv_cfg=self.conv_cfg,\n norm_cfg=self.norm_cfg,\n inplace=True),\n ConvModule(\n stem_channels // 2,\n stem_channels,\n kernel_size=3,\n stride=1,\n padding=1,\n conv_cfg=self.conv_cfg,\n norm_cfg=self.norm_cfg,\n inplace=True))\n else:\n self.conv1 = build_conv_layer(\n self.conv_cfg,\n in_channels,\n stem_channels,\n kernel_size=kernel_size,\n stride=2,\n padding=kernel_size // 2,\n bias=False)\n self.norm1_name, norm1 = build_norm_layer(\n self.norm_cfg, stem_channels, postfix=1)\n self.add_module(self.norm1_name, norm1)\n self.relu = nn.ReLU(inplace=True)\n self.maxpool = nn.MaxPool2d(kernel_size=3, stride=2, padding=1)\n\n def _freeze_stages(self):\n \"\"\"Freeze parameters.\"\"\"\n if self.frozen_stages >= 0:\n if self.deep_stem:\n self.stem.eval()\n for param in self.stem.parameters():\n param.requires_grad = False\n else:\n self.norm1.eval()\n for m in [self.conv1, self.norm1]:\n for param in m.parameters():\n param.requires_grad = False\n\n for i in range(1, self.frozen_stages + 1):\n m = getattr(self, f'layer{i}')\n m.eval()\n for param in m.parameters():\n param.requires_grad = False\n\n def init_weights(self, pretrained=None):\n \"\"\"Initialize model weights.\"\"\"\n super().init_weights(pretrained)\n if pretrained is None:\n for m in self.modules():\n if isinstance(m, nn.Conv2d):\n nn.init.normal_(m.weight, std=0.001)\n for name, _ in m.named_parameters():\n if name in ['bias']:\n nn.init.constant_(m.bias, 0)\n elif isinstance(m, nn.BatchNorm2d):\n nn.init.constant_(m.weight, 1)\n nn.init.constant_(m.bias, 0)\n\n def forward(self, x):\n \"\"\"Forward function.\"\"\"\n if self.deep_stem:\n x = self.stem(x)\n else:\n x = self.conv1(x)\n x = self.norm1(x)\n x = self.relu(x)\n x = self.maxpool(x)\n outs = []\n for i, layer_name in enumerate(self.res_layers):\n res_layer = getattr(self, layer_name)\n x = res_layer(x)\n if i in self.out_indices:\n outs.append(x)\n if len(outs) == 1:\n return outs[0]\n return tuple(outs)\n\n def train(self, mode=True):\n \"\"\"Convert the model into training mode.\"\"\"\n super().train(mode)\n self._freeze_stages()\n if mode and self.norm_eval:\n for m in self.modules():\n # trick: eval have effect on BatchNorm only\n if isinstance(m, _BatchNorm):\n m.eval()\n" }, { "path": "eval/GVHMR/hmr4d/utils/preproc/vitpose_pytorch/src/vitpose_infer/builder/backbones/vit.py", "content": "# Copyright (c) OpenMMLab. All rights reserved.\nimport math\n\nimport torch\nfrom functools import partial\nimport torch.nn as nn\nimport torch.nn.functional as F\nimport torch.utils.checkpoint as checkpoint\n\nfrom timm.models.layers import drop_path, to_2tuple, trunc_normal_\n\n# from ..builder import BACKBONES\n# from .base_backbone import BaseBackbone\n\nclass DropPath(nn.Module):\n \"\"\"Drop paths (Stochastic Depth) per sample (when applied in main path of residual blocks).\n \"\"\"\n def __init__(self, drop_prob=None):\n super(DropPath, self).__init__()\n self.drop_prob = drop_prob\n\n def forward(self, x):\n return drop_path(x, self.drop_prob, self.training)\n \n def extra_repr(self):\n return 'p={}'.format(self.drop_prob)\n\nclass Mlp(nn.Module):\n def __init__(self, in_features, hidden_features=None, out_features=None, act_layer=nn.GELU, drop=0.):\n super().__init__()\n out_features = out_features or in_features\n hidden_features = hidden_features or in_features\n self.fc1 = nn.Linear(in_features, hidden_features)\n self.act = act_layer()\n self.fc2 = nn.Linear(hidden_features, out_features)\n self.drop = nn.Dropout(drop)\n\n def forward(self, x):\n x = self.fc1(x)\n x = self.act(x)\n x = self.fc2(x)\n x = self.drop(x)\n return x\n\nclass Attention(nn.Module):\n def __init__(\n self, dim, num_heads=8, qkv_bias=False, qk_scale=None, attn_drop=0.,\n proj_drop=0., attn_head_dim=None,):\n super().__init__()\n self.num_heads = num_heads\n head_dim = dim // num_heads\n self.dim = dim\n\n if attn_head_dim is not None:\n head_dim = attn_head_dim\n all_head_dim = head_dim * self.num_heads\n\n self.scale = qk_scale or head_dim ** -0.5\n\n self.qkv = nn.Linear(dim, all_head_dim * 3, bias=qkv_bias)\n\n self.attn_drop = nn.Dropout(attn_drop)\n self.proj = nn.Linear(all_head_dim, dim)\n self.proj_drop = nn.Dropout(proj_drop)\n\n def forward(self, x):\n B, N, C = x.shape\n qkv = self.qkv(x)\n qkv = qkv.reshape(B, N, 3, self.num_heads, -1).permute(2, 0, 3, 1, 4)\n q, k, v = qkv[0], qkv[1], qkv[2] # make torchscript happy (cannot use tensor as tuple)\n\n q = q * self.scale\n attn = (q @ k.transpose(-2, -1))\n\n attn = attn.softmax(dim=-1)\n attn = self.attn_drop(attn)\n\n x = (attn @ v).transpose(1, 2).reshape(B, N, -1)\n x = self.proj(x)\n x = self.proj_drop(x)\n\n return x\n\nclass Block(nn.Module):\n\n def __init__(self, dim, num_heads, mlp_ratio=4., qkv_bias=False, qk_scale=None, \n drop=0., attn_drop=0., drop_path=0., act_layer=nn.GELU, \n norm_layer=nn.LayerNorm, attn_head_dim=None\n ):\n super().__init__()\n \n self.norm1 = norm_layer(dim)\n self.attn = Attention(\n dim, num_heads=num_heads, qkv_bias=qkv_bias, qk_scale=qk_scale,\n attn_drop=attn_drop, proj_drop=drop, attn_head_dim=attn_head_dim\n )\n\n # NOTE: drop path for stochastic depth, we shall see if this is better than dropout here\n self.drop_path = DropPath(drop_path) if drop_path > 0. else nn.Identity()\n self.norm2 = norm_layer(dim)\n mlp_hidden_dim = int(dim * mlp_ratio)\n self.mlp = Mlp(in_features=dim, hidden_features=mlp_hidden_dim, act_layer=act_layer, drop=drop)\n\n def forward(self, x):\n x = x + self.drop_path(self.attn(self.norm1(x)))\n x = x + self.drop_path(self.mlp(self.norm2(x)))\n return x\n\n\nclass PatchEmbed(nn.Module):\n \"\"\" Image to Patch Embedding\n \"\"\"\n def __init__(self, img_size=224, patch_size=16, in_chans=3, embed_dim=768, ratio=1):\n super().__init__()\n img_size = to_2tuple(img_size)\n patch_size = to_2tuple(patch_size)\n num_patches = (img_size[1] // patch_size[1]) * (img_size[0] // patch_size[0]) * (ratio ** 2)\n self.patch_shape = (int(img_size[0] // patch_size[0] * ratio), int(img_size[1] // patch_size[1] * ratio))\n self.origin_patch_shape = (int(img_size[0] // patch_size[0]), int(img_size[1] // patch_size[1]))\n self.img_size = img_size\n self.patch_size = patch_size\n self.num_patches = num_patches\n\n self.proj = nn.Conv2d(in_chans, embed_dim, kernel_size=patch_size, stride=(patch_size[0] // ratio), padding=4 + 2 * (ratio//2-1))\n\n def forward(self, x, **kwargs):\n B, C, H, W = x.shape\n x = self.proj(x)\n Hp, Wp = x.shape[2], x.shape[3]\n\n x = x.flatten(2).transpose(1, 2)\n return x, (Hp, Wp)\n\n\nclass HybridEmbed(nn.Module):\n \"\"\" CNN Feature Map Embedding\n Extract feature map from CNN, flatten, project to embedding dim.\n \"\"\"\n def __init__(self, backbone, img_size=224, feature_size=None, in_chans=3, embed_dim=768):\n super().__init__()\n assert isinstance(backbone, nn.Module)\n img_size = to_2tuple(img_size)\n self.img_size = img_size\n self.backbone = backbone\n if feature_size is None:\n with torch.no_grad():\n training = backbone.training\n if training:\n backbone.eval()\n o = self.backbone(torch.zeros(1, in_chans, img_size[0], img_size[1]))[-1]\n feature_size = o.shape[-2:]\n feature_dim = o.shape[1]\n backbone.train(training)\n else:\n feature_size = to_2tuple(feature_size)\n feature_dim = self.backbone.feature_info.channels()[-1]\n self.num_patches = feature_size[0] * feature_size[1]\n self.proj = nn.Linear(feature_dim, embed_dim)\n\n def forward(self, x):\n x = self.backbone(x)[-1]\n x = x.flatten(2).transpose(1, 2)\n x = self.proj(x)\n return x\n\n\n# @BACKBONES.register_module()\nclass ViT(nn.Module):\n\n def __init__(self,\n img_size=224, patch_size=16, in_chans=3, num_classes=80, embed_dim=768, depth=12,\n num_heads=12, mlp_ratio=4., qkv_bias=False, qk_scale=None, drop_rate=0., attn_drop_rate=0.,\n drop_path_rate=0., hybrid_backbone=None, norm_layer=None, use_checkpoint=False, \n frozen_stages=-1, ratio=1, last_norm=True,\n patch_padding='pad', freeze_attn=False, freeze_ffn=False,\n ):\n # Protect mutable default arguments\n super(ViT, self).__init__()\n norm_layer = norm_layer or partial(nn.LayerNorm, eps=1e-6)\n self.num_classes = num_classes\n self.num_features = self.embed_dim = embed_dim # num_features for consistency with other models\n self.frozen_stages = frozen_stages\n self.use_checkpoint = use_checkpoint\n self.patch_padding = patch_padding\n self.freeze_attn = freeze_attn\n self.freeze_ffn = freeze_ffn\n self.depth = depth\n\n if hybrid_backbone is not None:\n self.patch_embed = HybridEmbed(\n hybrid_backbone, img_size=img_size, in_chans=in_chans, embed_dim=embed_dim)\n else:\n self.patch_embed = PatchEmbed(\n img_size=img_size, patch_size=patch_size, in_chans=in_chans, embed_dim=embed_dim, ratio=ratio)\n num_patches = self.patch_embed.num_patches\n\n # since the pretraining model has class token\n self.pos_embed = nn.Parameter(torch.zeros(1, num_patches + 1, embed_dim))\n\n dpr = [x.item() for x in torch.linspace(0, drop_path_rate, depth)] # stochastic depth decay rule\n\n self.blocks = nn.ModuleList([\n Block(\n dim=embed_dim, num_heads=num_heads, mlp_ratio=mlp_ratio, qkv_bias=qkv_bias, qk_scale=qk_scale,\n drop=drop_rate, attn_drop=attn_drop_rate, drop_path=dpr[i], norm_layer=norm_layer,\n )\n for i in range(depth)])\n\n self.last_norm = norm_layer(embed_dim) if last_norm else nn.Identity()\n\n if self.pos_embed is not None:\n trunc_normal_(self.pos_embed, std=.02)\n\n self._freeze_stages()\n\n def _freeze_stages(self):\n \"\"\"Freeze parameters.\"\"\"\n if self.frozen_stages >= 0:\n self.patch_embed.eval()\n for param in self.patch_embed.parameters():\n param.requires_grad = False\n\n for i in range(1, self.frozen_stages + 1):\n m = self.blocks[i]\n m.eval()\n for param in m.parameters():\n param.requires_grad = False\n\n if self.freeze_attn:\n for i in range(0, self.depth):\n m = self.blocks[i]\n m.attn.eval()\n m.norm1.eval()\n for param in m.attn.parameters():\n param.requires_grad = False\n for param in m.norm1.parameters():\n param.requires_grad = False\n\n if self.freeze_ffn:\n self.pos_embed.requires_grad = False\n self.patch_embed.eval()\n for param in self.patch_embed.parameters():\n param.requires_grad = False\n for i in range(0, self.depth):\n m = self.blocks[i]\n m.mlp.eval()\n m.norm2.eval()\n for param in m.mlp.parameters():\n param.requires_grad = False\n for param in m.norm2.parameters():\n param.requires_grad = False\n\n def init_weights(self, pretrained=None):\n \"\"\"Initialize the weights in backbone.\n Args:\n pretrained (str, optional): Path to pre-trained weights.\n Defaults to None.\n \"\"\"\n super().init_weights(pretrained, patch_padding=self.patch_padding)\n\n if pretrained is None:\n def _init_weights(m):\n if isinstance(m, nn.Linear):\n trunc_normal_(m.weight, std=.02)\n if isinstance(m, nn.Linear) and m.bias is not None:\n nn.init.constant_(m.bias, 0)\n elif isinstance(m, nn.LayerNorm):\n nn.init.constant_(m.bias, 0)\n nn.init.constant_(m.weight, 1.0)\n\n self.apply(_init_weights)\n\n def get_num_layers(self):\n return len(self.blocks)\n\n @torch.jit.ignore\n def no_weight_decay(self):\n return {'pos_embed', 'cls_token'}\n\n def forward_features(self, x):\n B, C, H, W = x.shape\n x, (Hp, Wp) = self.patch_embed(x)\n\n if self.pos_embed is not None:\n # fit for multiple GPU training\n # since the first element for pos embed (sin-cos manner) is zero, it will cause no difference\n x = x + self.pos_embed[:, 1:] + self.pos_embed[:, :1]\n\n for blk in self.blocks:\n if self.use_checkpoint:\n x = checkpoint.checkpoint(blk, x)\n else:\n x = blk(x)\n\n x = self.last_norm(x)\n\n xp = x.permute(0, 2, 1).reshape(B, -1, Hp, Wp).contiguous()\n\n return xp\n\n def forward(self, x):\n x = self.forward_features(x)\n return x\n\n def train(self, mode=True):\n \"\"\"Convert the model into training mode.\"\"\"\n super().train(mode)\n self._freeze_stages()\n" }, { "path": "eval/GVHMR/hmr4d/utils/preproc/vitpose_pytorch/src/vitpose_infer/builder/configs/coco/ViTPose_base_coco_256x192.py", "content": "_base_ = [\n '../../../../_base_/default_runtime.py',\n '../../../../_base_/datasets/coco.py'\n]\nevaluation = dict(interval=10, metric='mAP', save_best='AP')\n\noptimizer = dict(type='AdamW', lr=5e-4, betas=(0.9, 0.999), weight_decay=0.1,\n constructor='LayerDecayOptimizerConstructor', \n paramwise_cfg=dict(\n num_layers=12, \n layer_decay_rate=0.75,\n custom_keys={\n 'bias': dict(decay_multi=0.),\n 'pos_embed': dict(decay_mult=0.),\n 'relative_position_bias_table': dict(decay_mult=0.),\n 'norm': dict(decay_mult=0.)\n }\n )\n )\n\noptimizer_config = dict(grad_clip=dict(max_norm=1., norm_type=2))\n\n# learning policy\nlr_config = dict(\n policy='step',\n warmup='linear',\n warmup_iters=500,\n warmup_ratio=0.001,\n step=[170, 200])\ntotal_epochs = 210\ntarget_type = 'GaussianHeatmap'\nchannel_cfg = dict(\n num_output_channels=17,\n dataset_joints=17,\n dataset_channel=[\n [0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16],\n ],\n inference_channel=[\n 0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16\n ])\n\n# model settings\nmodel = dict(\n type='TopDown',\n pretrained=None,\n backbone=dict(\n type='ViT',\n img_size=(256, 192),\n patch_size=16,\n embed_dim=768,\n depth=12,\n num_heads=12,\n ratio=1,\n use_checkpoint=False,\n mlp_ratio=4,\n qkv_bias=True,\n drop_path_rate=0.3,\n ),\n keypoint_head=dict(\n type='TopdownHeatmapSimpleHead',\n in_channels=768,\n num_deconv_layers=2,\n num_deconv_filters=(256, 256),\n num_deconv_kernels=(4, 4),\n extra=dict(final_conv_kernel=1, ),\n out_channels=channel_cfg['num_output_channels'],\n loss_keypoint=dict(type='JointsMSELoss', use_target_weight=True)),\n train_cfg=dict(),\n test_cfg=dict())\n\ndata_cfg = dict(\n image_size=[192, 256],\n heatmap_size=[48, 64],\n num_output_channels=channel_cfg['num_output_channels'],\n num_joints=channel_cfg['dataset_joints'],\n dataset_channel=channel_cfg['dataset_channel'],\n inference_channel=channel_cfg['inference_channel'],\n soft_nms=False,\n nms_thr=1.0,\n oks_thr=0.9,\n vis_thr=0.2,\n use_gt_bbox=False,\n det_bbox_thr=0.0,\n bbox_file='data/coco/person_detection_results/'\n 'COCO_val2017_detections_AP_H_56_person.json',\n)\n\ntrain_pipeline = [\n dict(type='LoadImageFromFile'),\n dict(type='TopDownRandomFlip', flip_prob=0.5),\n dict(\n type='TopDownHalfBodyTransform',\n num_joints_half_body=8,\n prob_half_body=0.3),\n dict(\n type='TopDownGetRandomScaleRotation', rot_factor=40, scale_factor=0.5),\n dict(type='TopDownAffine', use_udp=True),\n dict(type='ToTensor'),\n dict(\n type='NormalizeTensor',\n mean=[0.485, 0.456, 0.406],\n std=[0.229, 0.224, 0.225]),\n dict(\n type='TopDownGenerateTarget',\n sigma=2,\n encoding='UDP',\n target_type=target_type),\n dict(\n type='Collect',\n keys=['img', 'target', 'target_weight'],\n meta_keys=[\n 'image_file', 'joints_3d', 'joints_3d_visible', 'center', 'scale',\n 'rotation', 'bbox_score', 'flip_pairs'\n ]),\n]\n\nval_pipeline = [\n dict(type='LoadImageFromFile'),\n dict(type='TopDownAffine', use_udp=True),\n dict(type='ToTensor'),\n dict(\n type='NormalizeTensor',\n mean=[0.485, 0.456, 0.406],\n std=[0.229, 0.224, 0.225]),\n dict(\n type='Collect',\n keys=['img'],\n meta_keys=[\n 'image_file', 'center', 'scale', 'rotation', 'bbox_score',\n 'flip_pairs'\n ]),\n]\n\ntest_pipeline = val_pipeline\n\ndata_root = 'data/coco'\n# data = dict(\n# samples_per_gpu=64,\n# workers_per_gpu=4,\n# val_dataloader=dict(samples_per_gpu=32),\n# test_dataloader=dict(samples_per_gpu=32),\n# train=dict(\n# type='TopDownCocoDataset',\n# ann_file=f'{data_root}/annotations/person_keypoints_train2017.json',\n# img_prefix=f'{data_root}/train2017/',\n# data_cfg=data_cfg,\n# pipeline=train_pipeline,\n# dataset_info={{_base_.dataset_info}}),\n# val=dict(\n# type='TopDownCocoDataset',\n# ann_file=f'{data_root}/annotations/person_keypoints_val2017.json',\n# img_prefix=f'{data_root}/val2017/',\n# data_cfg=data_cfg,\n# pipeline=val_pipeline,\n# dataset_info={{_base_.dataset_info}}),\n# test=dict(\n# type='TopDownCocoDataset',\n# ann_file=f'{data_root}/annotations/person_keypoints_val2017.json',\n# img_prefix=f'{data_root}/val2017/',\n# data_cfg=data_cfg,\n# pipeline=test_pipeline,\n# dataset_info={{_base_.dataset_info}}),\n# )\n\ndef make_cfg(model=model,data_cfg=data_cfg):\n cfg={}\n cfg['model'] = model\n cfg['data_cfg'] = data_cfg" }, { "path": "eval/GVHMR/hmr4d/utils/preproc/vitpose_pytorch/src/vitpose_infer/builder/configs/coco/ViTPose_base_simple_coco_256x192.py", "content": "_base_ = [\n '../../../../_base_/default_runtime.py',\n '../../../../_base_/datasets/coco.py'\n]\nevaluation = dict(interval=10, metric='mAP', save_best='AP')\n\noptimizer = dict(type='AdamW', lr=5e-4, betas=(0.9, 0.999), weight_decay=0.1,\n constructor='LayerDecayOptimizerConstructor', \n paramwise_cfg=dict(\n num_layers=12, \n layer_decay_rate=0.75,\n custom_keys={\n 'bias': dict(decay_multi=0.),\n 'pos_embed': dict(decay_mult=0.),\n 'relative_position_bias_table': dict(decay_mult=0.),\n 'norm': dict(decay_mult=0.)\n }\n )\n )\n\noptimizer_config = dict(grad_clip=dict(max_norm=1., norm_type=2))\n\n# learning policy\nlr_config = dict(\n policy='step',\n warmup='linear',\n warmup_iters=500,\n warmup_ratio=0.001,\n step=[170, 200])\ntotal_epochs = 210\ntarget_type = 'GaussianHeatmap'\nchannel_cfg = dict(\n num_output_channels=17,\n dataset_joints=17,\n dataset_channel=[\n [0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16],\n ],\n inference_channel=[\n 0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16\n ])\n\n# model settings\nmodel = dict(\n type='TopDown',\n pretrained=None,\n backbone=dict(\n type='ViT',\n img_size=(256, 192),\n patch_size=16,\n embed_dim=768,\n depth=12,\n num_heads=12,\n ratio=1,\n use_checkpoint=False,\n mlp_ratio=4,\n qkv_bias=True,\n drop_path_rate=0.3,\n ),\n keypoint_head=dict(\n type='TopdownHeatmapSimpleHead',\n in_channels=768,\n num_deconv_layers=0,\n num_deconv_filters=[],\n num_deconv_kernels=[],\n upsample=4,\n extra=dict(final_conv_kernel=3, ),\n out_channels=channel_cfg['num_output_channels'],\n loss_keypoint=dict(type='JointsMSELoss', use_target_weight=True)),\n train_cfg=dict(),\n test_cfg=dict(\n flip_test=True,\n post_process='default',\n shift_heatmap=False,\n target_type=target_type,\n modulate_kernel=11,\n use_udp=True))\n\ndata_cfg = dict(\n image_size=[192, 256],\n heatmap_size=[48, 64],\n num_output_channels=channel_cfg['num_output_channels'],\n num_joints=channel_cfg['dataset_joints'],\n dataset_channel=channel_cfg['dataset_channel'],\n inference_channel=channel_cfg['inference_channel'],\n soft_nms=False,\n nms_thr=1.0,\n oks_thr=0.9,\n vis_thr=0.2,\n use_gt_bbox=False,\n det_bbox_thr=0.0,\n bbox_file='data/coco/person_detection_results/'\n 'COCO_val2017_detections_AP_H_56_person.json',\n)\n\ntrain_pipeline = [\n dict(type='LoadImageFromFile'),\n dict(type='TopDownRandomFlip', flip_prob=0.5),\n dict(\n type='TopDownHalfBodyTransform',\n num_joints_half_body=8,\n prob_half_body=0.3),\n dict(\n type='TopDownGetRandomScaleRotation', rot_factor=40, scale_factor=0.5),\n dict(type='TopDownAffine', use_udp=True),\n dict(type='ToTensor'),\n dict(\n type='NormalizeTensor',\n mean=[0.485, 0.456, 0.406],\n std=[0.229, 0.224, 0.225]),\n dict(\n type='TopDownGenerateTarget',\n sigma=2,\n encoding='UDP',\n target_type=target_type),\n dict(\n type='Collect',\n keys=['img', 'target', 'target_weight'],\n meta_keys=[\n 'image_file', 'joints_3d', 'joints_3d_visible', 'center', 'scale',\n 'rotation', 'bbox_score', 'flip_pairs'\n ]),\n]\n\nval_pipeline = [\n dict(type='LoadImageFromFile'),\n dict(type='TopDownAffine', use_udp=True),\n dict(type='ToTensor'),\n dict(\n type='NormalizeTensor',\n mean=[0.485, 0.456, 0.406],\n std=[0.229, 0.224, 0.225]),\n dict(\n type='Collect',\n keys=['img'],\n meta_keys=[\n 'image_file', 'center', 'scale', 'rotation', 'bbox_score',\n 'flip_pairs'\n ]),\n]\n\ntest_pipeline = val_pipeline\n\ndata_root = 'data/coco'\ndata = dict(\n samples_per_gpu=64,\n workers_per_gpu=4,\n val_dataloader=dict(samples_per_gpu=32),\n test_dataloader=dict(samples_per_gpu=32),\n train=dict(\n type='TopDownCocoDataset',\n ann_file=f'{data_root}/annotations/person_keypoints_train2017.json',\n img_prefix=f'{data_root}/train2017/',\n data_cfg=data_cfg,\n pipeline=train_pipeline,\n dataset_info={{_base_.dataset_info}}),\n val=dict(\n type='TopDownCocoDataset',\n ann_file=f'{data_root}/annotations/person_keypoints_val2017.json',\n img_prefix=f'{data_root}/val2017/',\n data_cfg=data_cfg,\n pipeline=val_pipeline,\n dataset_info={{_base_.dataset_info}}),\n test=dict(\n type='TopDownCocoDataset',\n ann_file=f'{data_root}/annotations/person_keypoints_val2017.json',\n img_prefix=f'{data_root}/val2017/',\n data_cfg=data_cfg,\n pipeline=test_pipeline,\n dataset_info={{_base_.dataset_info}}),\n)\n\n" }, { "path": "eval/GVHMR/hmr4d/utils/preproc/vitpose_pytorch/src/vitpose_infer/builder/configs/coco/ViTPose_huge_coco_256x192.py", "content": "_base_ = [\n '../../../../_base_/default_runtime.py',\n '../../../../_base_/datasets/coco.py'\n]\nevaluation = dict(interval=10, metric='mAP', save_best='AP')\n\noptimizer = dict(type='AdamW', lr=5e-4, betas=(0.9, 0.999), weight_decay=0.1,\n constructor='LayerDecayOptimizerConstructor', \n paramwise_cfg=dict(\n num_layers=32, \n layer_decay_rate=0.85,\n custom_keys={\n 'bias': dict(decay_multi=0.),\n 'pos_embed': dict(decay_mult=0.),\n 'relative_position_bias_table': dict(decay_mult=0.),\n 'norm': dict(decay_mult=0.)\n }\n )\n )\n\noptimizer_config = dict(grad_clip=dict(max_norm=1., norm_type=2))\n\n# learning policy\nlr_config = dict(\n policy='step',\n warmup='linear',\n warmup_iters=500,\n warmup_ratio=0.001,\n step=[170, 200])\ntotal_epochs = 210\ntarget_type = 'GaussianHeatmap'\nchannel_cfg = dict(\n num_output_channels=17,\n dataset_joints=17,\n dataset_channel=[\n [0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16],\n ],\n inference_channel=[\n 0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16\n ])\n\n# model settings\nmodel = dict(\n type='TopDown',\n pretrained=None,\n backbone=dict(\n type='ViT',\n img_size=(256, 192),\n patch_size=16,\n embed_dim=1280,\n depth=32,\n num_heads=16,\n ratio=1,\n use_checkpoint=False,\n mlp_ratio=4,\n qkv_bias=True,\n drop_path_rate=0.55,\n ),\n keypoint_head=dict(\n type='TopdownHeatmapSimpleHead',\n in_channels=1280,\n num_deconv_layers=2,\n num_deconv_filters=(256, 256),\n num_deconv_kernels=(4, 4),\n extra=dict(final_conv_kernel=1, ),\n out_channels=channel_cfg['num_output_channels'],\n loss_keypoint=dict(type='JointsMSELoss', use_target_weight=True)),\n train_cfg=dict(),\n test_cfg=dict(\n flip_test=True,\n post_process='default',\n shift_heatmap=False,\n target_type=target_type,\n modulate_kernel=11,\n use_udp=True))\n\ndata_cfg = dict(\n image_size=[192, 256],\n heatmap_size=[48, 64],\n num_output_channels=channel_cfg['num_output_channels'],\n num_joints=channel_cfg['dataset_joints'],\n dataset_channel=channel_cfg['dataset_channel'],\n inference_channel=channel_cfg['inference_channel'],\n soft_nms=False,\n nms_thr=1.0,\n oks_thr=0.9,\n vis_thr=0.2,\n use_gt_bbox=False,\n det_bbox_thr=0.0,\n bbox_file='data/coco/person_detection_results/'\n 'COCO_val2017_detections_AP_H_56_person.json',\n)\n\ntrain_pipeline = [\n dict(type='LoadImageFromFile'),\n dict(type='TopDownRandomFlip', flip_prob=0.5),\n dict(\n type='TopDownHalfBodyTransform',\n num_joints_half_body=8,\n prob_half_body=0.3),\n dict(\n type='TopDownGetRandomScaleRotation', rot_factor=40, scale_factor=0.5),\n dict(type='TopDownAffine', use_udp=True),\n dict(type='ToTensor'),\n dict(\n type='NormalizeTensor',\n mean=[0.485, 0.456, 0.406],\n std=[0.229, 0.224, 0.225]),\n dict(\n type='TopDownGenerateTarget',\n sigma=2,\n encoding='UDP',\n target_type=target_type),\n dict(\n type='Collect',\n keys=['img', 'target', 'target_weight'],\n meta_keys=[\n 'image_file', 'joints_3d', 'joints_3d_visible', 'center', 'scale',\n 'rotation', 'bbox_score', 'flip_pairs'\n ]),\n]\n\nval_pipeline = [\n dict(type='LoadImageFromFile'),\n dict(type='TopDownAffine', use_udp=True),\n dict(type='ToTensor'),\n dict(\n type='NormalizeTensor',\n mean=[0.485, 0.456, 0.406],\n std=[0.229, 0.224, 0.225]),\n dict(\n type='Collect',\n keys=['img'],\n meta_keys=[\n 'image_file', 'center', 'scale', 'rotation', 'bbox_score',\n 'flip_pairs'\n ]),\n]\n\ntest_pipeline = val_pipeline\n\ndata_root = 'data/coco'\ndata = dict(\n samples_per_gpu=64,\n workers_per_gpu=4,\n val_dataloader=dict(samples_per_gpu=32),\n test_dataloader=dict(samples_per_gpu=32),\n train=dict(\n type='TopDownCocoDataset',\n ann_file=f'{data_root}/annotations/person_keypoints_train2017.json',\n img_prefix=f'{data_root}/train2017/',\n data_cfg=data_cfg,\n pipeline=train_pipeline,\n dataset_info={{_base_.dataset_info}}),\n val=dict(\n type='TopDownCocoDataset',\n ann_file=f'{data_root}/annotations/person_keypoints_val2017.json',\n img_prefix=f'{data_root}/val2017/',\n data_cfg=data_cfg,\n pipeline=val_pipeline,\n dataset_info={{_base_.dataset_info}}),\n test=dict(\n type='TopDownCocoDataset',\n ann_file=f'{data_root}/annotations/person_keypoints_val2017.json',\n img_prefix=f'{data_root}/val2017/',\n data_cfg=data_cfg,\n pipeline=test_pipeline,\n dataset_info={{_base_.dataset_info}}),\n)\n\n" }, { "path": "eval/GVHMR/hmr4d/utils/preproc/vitpose_pytorch/src/vitpose_infer/builder/configs/coco/ViTPose_huge_simple_coco_256x192.py", "content": "_base_ = [\n '../../../../_base_/default_runtime.py',\n '../../../../_base_/datasets/coco.py'\n]\nevaluation = dict(interval=10, metric='mAP', save_best='AP')\n\noptimizer = dict(type='AdamW', lr=5e-4, betas=(0.9, 0.999), weight_decay=0.1,\n constructor='LayerDecayOptimizerConstructor', \n paramwise_cfg=dict(\n num_layers=32, \n layer_decay_rate=0.85,\n custom_keys={\n 'bias': dict(decay_multi=0.),\n 'pos_embed': dict(decay_mult=0.),\n 'relative_position_bias_table': dict(decay_mult=0.),\n 'norm': dict(decay_mult=0.)\n }\n )\n )\n\noptimizer_config = dict(grad_clip=dict(max_norm=1., norm_type=2))\n\n# learning policy\nlr_config = dict(\n policy='step',\n warmup='linear',\n warmup_iters=500,\n warmup_ratio=0.001,\n step=[170, 200])\ntotal_epochs = 210\ntarget_type = 'GaussianHeatmap'\nchannel_cfg = dict(\n num_output_channels=17,\n dataset_joints=17,\n dataset_channel=[\n [0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16],\n ],\n inference_channel=[\n 0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16\n ])\n\n# model settings\nmodel = dict(\n type='TopDown',\n pretrained=None,\n backbone=dict(\n type='ViT',\n img_size=(256, 192),\n patch_size=16,\n embed_dim=1280,\n depth=32,\n num_heads=16,\n ratio=1,\n use_checkpoint=False,\n mlp_ratio=4,\n qkv_bias=True,\n drop_path_rate=0.55,\n ),\n keypoint_head=dict(\n type='TopdownHeatmapSimpleHead',\n in_channels=1280,\n num_deconv_layers=0,\n num_deconv_filters=[],\n num_deconv_kernels=[],\n upsample=4,\n extra=dict(final_conv_kernel=3, ),\n out_channels=channel_cfg['num_output_channels'],\n loss_keypoint=dict(type='JointsMSELoss', use_target_weight=True)),\n train_cfg=dict(),\n test_cfg=dict(\n flip_test=True,\n post_process='default',\n shift_heatmap=False,\n target_type=target_type,\n modulate_kernel=11,\n use_udp=True))\n\ndata_cfg = dict(\n image_size=[192, 256],\n heatmap_size=[48, 64],\n num_output_channels=channel_cfg['num_output_channels'],\n num_joints=channel_cfg['dataset_joints'],\n dataset_channel=channel_cfg['dataset_channel'],\n inference_channel=channel_cfg['inference_channel'],\n soft_nms=False,\n nms_thr=1.0,\n oks_thr=0.9,\n vis_thr=0.2,\n use_gt_bbox=False,\n det_bbox_thr=0.0,\n bbox_file='data/coco/person_detection_results/'\n 'COCO_val2017_detections_AP_H_56_person.json',\n)\n\ntrain_pipeline = [\n dict(type='LoadImageFromFile'),\n dict(type='TopDownRandomFlip', flip_prob=0.5),\n dict(\n type='TopDownHalfBodyTransform',\n num_joints_half_body=8,\n prob_half_body=0.3),\n dict(\n type='TopDownGetRandomScaleRotation', rot_factor=40, scale_factor=0.5),\n dict(type='TopDownAffine', use_udp=True),\n dict(type='ToTensor'),\n dict(\n type='NormalizeTensor',\n mean=[0.485, 0.456, 0.406],\n std=[0.229, 0.224, 0.225]),\n dict(\n type='TopDownGenerateTarget',\n sigma=2,\n encoding='UDP',\n target_type=target_type),\n dict(\n type='Collect',\n keys=['img', 'target', 'target_weight'],\n meta_keys=[\n 'image_file', 'joints_3d', 'joints_3d_visible', 'center', 'scale',\n 'rotation', 'bbox_score', 'flip_pairs'\n ]),\n]\n\nval_pipeline = [\n dict(type='LoadImageFromFile'),\n dict(type='TopDownAffine', use_udp=True),\n dict(type='ToTensor'),\n dict(\n type='NormalizeTensor',\n mean=[0.485, 0.456, 0.406],\n std=[0.229, 0.224, 0.225]),\n dict(\n type='Collect',\n keys=['img'],\n meta_keys=[\n 'image_file', 'center', 'scale', 'rotation', 'bbox_score',\n 'flip_pairs'\n ]),\n]\n\ntest_pipeline = val_pipeline\n\ndata_root = 'data/coco'\ndata = dict(\n samples_per_gpu=64,\n workers_per_gpu=4,\n val_dataloader=dict(samples_per_gpu=32),\n test_dataloader=dict(samples_per_gpu=32),\n train=dict(\n type='TopDownCocoDataset',\n ann_file=f'{data_root}/annotations/person_keypoints_train2017.json',\n img_prefix=f'{data_root}/train2017/',\n data_cfg=data_cfg,\n pipeline=train_pipeline,\n dataset_info={{_base_.dataset_info}}),\n val=dict(\n type='TopDownCocoDataset',\n ann_file=f'{data_root}/annotations/person_keypoints_val2017.json',\n img_prefix=f'{data_root}/val2017/',\n data_cfg=data_cfg,\n pipeline=val_pipeline,\n dataset_info={{_base_.dataset_info}}),\n test=dict(\n type='TopDownCocoDataset',\n ann_file=f'{data_root}/annotations/person_keypoints_val2017.json',\n img_prefix=f'{data_root}/val2017/',\n data_cfg=data_cfg,\n pipeline=test_pipeline,\n dataset_info={{_base_.dataset_info}}),\n)\n\n" }, { "path": "eval/GVHMR/hmr4d/utils/preproc/vitpose_pytorch/src/vitpose_infer/builder/configs/coco/ViTPose_large_coco_256x192.py", "content": "_base_ = [\n '../../../../_base_/default_runtime.py',\n '../../../../_base_/datasets/coco.py'\n]\nevaluation = dict(interval=10, metric='mAP', save_best='AP')\n\noptimizer = dict(type='AdamW', lr=5e-4, betas=(0.9, 0.999), weight_decay=0.1,\n constructor='LayerDecayOptimizerConstructor', \n paramwise_cfg=dict(\n num_layers=16, \n layer_decay_rate=0.8,\n custom_keys={\n 'bias': dict(decay_multi=0.),\n 'pos_embed': dict(decay_mult=0.),\n 'relative_position_bias_table': dict(decay_mult=0.),\n 'norm': dict(decay_mult=0.)\n }\n )\n )\n\noptimizer_config = dict(grad_clip=dict(max_norm=1., norm_type=2))\n\n# learning policy\nlr_config = dict(\n policy='step',\n warmup='linear',\n warmup_iters=500,\n warmup_ratio=0.001,\n step=[170, 200])\ntotal_epochs = 210\ntarget_type = 'GaussianHeatmap'\nchannel_cfg = dict(\n num_output_channels=17,\n dataset_joints=17,\n dataset_channel=[\n [0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16],\n ],\n inference_channel=[\n 0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16\n ])\n\n# model settings\nmodel = dict(\n type='TopDown',\n pretrained=None,\n backbone=dict(\n type='ViT',\n img_size=(256, 192),\n patch_size=16,\n embed_dim=1024,\n depth=24,\n num_heads=16,\n ratio=1,\n use_checkpoint=False,\n mlp_ratio=4,\n qkv_bias=True,\n drop_path_rate=0.5,\n ),\n keypoint_head=dict(\n type='TopdownHeatmapSimpleHead',\n in_channels=1024,\n num_deconv_layers=2,\n num_deconv_filters=(256, 256),\n num_deconv_kernels=(4, 4),\n extra=dict(final_conv_kernel=1, ),\n out_channels=channel_cfg['num_output_channels'],\n loss_keypoint=dict(type='JointsMSELoss', use_target_weight=True)),\n train_cfg=dict(),\n test_cfg=dict(\n flip_test=True,\n post_process='default',\n shift_heatmap=False,\n target_type=target_type,\n modulate_kernel=11,\n use_udp=True))\n\ndata_cfg = dict(\n image_size=[192, 256],\n heatmap_size=[48, 64],\n num_output_channels=channel_cfg['num_output_channels'],\n num_joints=channel_cfg['dataset_joints'],\n dataset_channel=channel_cfg['dataset_channel'],\n inference_channel=channel_cfg['inference_channel'],\n soft_nms=False,\n nms_thr=1.0,\n oks_thr=0.9,\n vis_thr=0.2,\n use_gt_bbox=False,\n det_bbox_thr=0.0,\n bbox_file='data/coco/person_detection_results/'\n 'COCO_val2017_detections_AP_H_56_person.json',\n)\n\ntrain_pipeline = [\n dict(type='LoadImageFromFile'),\n dict(type='TopDownRandomFlip', flip_prob=0.5),\n dict(\n type='TopDownHalfBodyTransform',\n num_joints_half_body=8,\n prob_half_body=0.3),\n dict(\n type='TopDownGetRandomScaleRotation', rot_factor=40, scale_factor=0.5),\n dict(type='TopDownAffine', use_udp=True),\n dict(type='ToTensor'),\n dict(\n type='NormalizeTensor',\n mean=[0.485, 0.456, 0.406],\n std=[0.229, 0.224, 0.225]),\n dict(\n type='TopDownGenerateTarget',\n sigma=2,\n encoding='UDP',\n target_type=target_type),\n dict(\n type='Collect',\n keys=['img', 'target', 'target_weight'],\n meta_keys=[\n 'image_file', 'joints_3d', 'joints_3d_visible', 'center', 'scale',\n 'rotation', 'bbox_score', 'flip_pairs'\n ]),\n]\n\nval_pipeline = [\n dict(type='LoadImageFromFile'),\n dict(type='TopDownAffine', use_udp=True),\n dict(type='ToTensor'),\n dict(\n type='NormalizeTensor',\n mean=[0.485, 0.456, 0.406],\n std=[0.229, 0.224, 0.225]),\n dict(\n type='Collect',\n keys=['img'],\n meta_keys=[\n 'image_file', 'center', 'scale', 'rotation', 'bbox_score',\n 'flip_pairs'\n ]),\n]\n\ntest_pipeline = val_pipeline\n\ndata_root = 'data/coco'\ndata = dict(\n samples_per_gpu=64,\n workers_per_gpu=4,\n val_dataloader=dict(samples_per_gpu=32),\n test_dataloader=dict(samples_per_gpu=32),\n train=dict(\n type='TopDownCocoDataset',\n ann_file=f'{data_root}/annotations/person_keypoints_train2017.json',\n img_prefix=f'{data_root}/train2017/',\n data_cfg=data_cfg,\n pipeline=train_pipeline,\n dataset_info={{_base_.dataset_info}}),\n val=dict(\n type='TopDownCocoDataset',\n ann_file=f'{data_root}/annotations/person_keypoints_val2017.json',\n img_prefix=f'{data_root}/val2017/',\n data_cfg=data_cfg,\n pipeline=val_pipeline,\n dataset_info={{_base_.dataset_info}}),\n test=dict(\n type='TopDownCocoDataset',\n ann_file=f'{data_root}/annotations/person_keypoints_val2017.json',\n img_prefix=f'{data_root}/val2017/',\n data_cfg=data_cfg,\n pipeline=test_pipeline,\n dataset_info={{_base_.dataset_info}}),\n)\n\n" }, { "path": "eval/GVHMR/hmr4d/utils/preproc/vitpose_pytorch/src/vitpose_infer/builder/configs/coco/ViTPose_large_simple_coco_256x192.py", "content": "_base_ = [\n '../../../../_base_/default_runtime.py',\n '../../../../_base_/datasets/coco.py'\n]\nevaluation = dict(interval=10, metric='mAP', save_best='AP')\n\noptimizer = dict(type='AdamW', lr=5e-4, betas=(0.9, 0.999), weight_decay=0.1,\n constructor='LayerDecayOptimizerConstructor', \n paramwise_cfg=dict(\n num_layers=24, \n layer_decay_rate=0.8,\n custom_keys={\n 'bias': dict(decay_multi=0.),\n 'pos_embed': dict(decay_mult=0.),\n 'relative_position_bias_table': dict(decay_mult=0.),\n 'norm': dict(decay_mult=0.)\n }\n )\n )\n\noptimizer_config = dict(grad_clip=dict(max_norm=1., norm_type=2))\n\n# learning policy\nlr_config = dict(\n policy='step',\n warmup='linear',\n warmup_iters=500,\n warmup_ratio=0.001,\n step=[170, 200])\ntotal_epochs = 210\ntarget_type = 'GaussianHeatmap'\nchannel_cfg = dict(\n num_output_channels=17,\n dataset_joints=17,\n dataset_channel=[\n [0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16],\n ],\n inference_channel=[\n 0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16\n ])\n\n# model settings\nmodel = dict(\n type='TopDown',\n pretrained=None,\n backbone=dict(\n type='ViT',\n img_size=(256, 192),\n patch_size=16,\n embed_dim=1024,\n depth=24,\n num_heads=16,\n ratio=1,\n use_checkpoint=False,\n mlp_ratio=4,\n qkv_bias=True,\n drop_path_rate=0.5,\n ),\n keypoint_head=dict(\n type='TopdownHeatmapSimpleHead',\n in_channels=1024,\n num_deconv_layers=0,\n num_deconv_filters=[],\n num_deconv_kernels=[],\n upsample=4,\n extra=dict(final_conv_kernel=3, ),\n out_channels=channel_cfg['num_output_channels'],\n loss_keypoint=dict(type='JointsMSELoss', use_target_weight=True)),\n train_cfg=dict(),\n test_cfg=dict(\n flip_test=True,\n post_process='default',\n shift_heatmap=False,\n target_type=target_type,\n modulate_kernel=11,\n use_udp=True))\n\ndata_cfg = dict(\n image_size=[192, 256],\n heatmap_size=[48, 64],\n num_output_channels=channel_cfg['num_output_channels'],\n num_joints=channel_cfg['dataset_joints'],\n dataset_channel=channel_cfg['dataset_channel'],\n inference_channel=channel_cfg['inference_channel'],\n soft_nms=False,\n nms_thr=1.0,\n oks_thr=0.9,\n vis_thr=0.2,\n use_gt_bbox=False,\n det_bbox_thr=0.0,\n bbox_file='data/coco/person_detection_results/'\n 'COCO_val2017_detections_AP_H_56_person.json',\n)\n\ntrain_pipeline = [\n dict(type='LoadImageFromFile'),\n dict(type='TopDownRandomFlip', flip_prob=0.5),\n dict(\n type='TopDownHalfBodyTransform',\n num_joints_half_body=8,\n prob_half_body=0.3),\n dict(\n type='TopDownGetRandomScaleRotation', rot_factor=40, scale_factor=0.5),\n dict(type='TopDownAffine', use_udp=True),\n dict(type='ToTensor'),\n dict(\n type='NormalizeTensor',\n mean=[0.485, 0.456, 0.406],\n std=[0.229, 0.224, 0.225]),\n dict(\n type='TopDownGenerateTarget',\n sigma=2,\n encoding='UDP',\n target_type=target_type),\n dict(\n type='Collect',\n keys=['img', 'target', 'target_weight'],\n meta_keys=[\n 'image_file', 'joints_3d', 'joints_3d_visible', 'center', 'scale',\n 'rotation', 'bbox_score', 'flip_pairs'\n ]),\n]\n\nval_pipeline = [\n dict(type='LoadImageFromFile'),\n dict(type='TopDownAffine', use_udp=True),\n dict(type='ToTensor'),\n dict(\n type='NormalizeTensor',\n mean=[0.485, 0.456, 0.406],\n std=[0.229, 0.224, 0.225]),\n dict(\n type='Collect',\n keys=['img'],\n meta_keys=[\n 'image_file', 'center', 'scale', 'rotation', 'bbox_score',\n 'flip_pairs'\n ]),\n]\n\ntest_pipeline = val_pipeline\n\ndata_root = 'data/coco'\ndata = dict(\n samples_per_gpu=64,\n workers_per_gpu=4,\n val_dataloader=dict(samples_per_gpu=32),\n test_dataloader=dict(samples_per_gpu=32),\n train=dict(\n type='TopDownCocoDataset',\n ann_file=f'{data_root}/annotations/person_keypoints_train2017.json',\n img_prefix=f'{data_root}/train2017/',\n data_cfg=data_cfg,\n pipeline=train_pipeline,\n dataset_info={{_base_.dataset_info}}),\n val=dict(\n type='TopDownCocoDataset',\n ann_file=f'{data_root}/annotations/person_keypoints_val2017.json',\n img_prefix=f'{data_root}/val2017/',\n data_cfg=data_cfg,\n pipeline=val_pipeline,\n dataset_info={{_base_.dataset_info}}),\n test=dict(\n type='TopDownCocoDataset',\n ann_file=f'{data_root}/annotations/person_keypoints_val2017.json',\n img_prefix=f'{data_root}/val2017/',\n data_cfg=data_cfg,\n pipeline=test_pipeline,\n dataset_info={{_base_.dataset_info}}),\n)\n\n" }, { "path": "eval/GVHMR/hmr4d/utils/preproc/vitpose_pytorch/src/vitpose_infer/builder/configs/coco/__init__.py", "content": "" }, { "path": "eval/GVHMR/hmr4d/utils/preproc/vitpose_pytorch/src/vitpose_infer/builder/heads/__init__.py", "content": "# Copyright (c) OpenMMLab. All rights reserved.\n# from .ae_higher_resolution_head import AEHigherResolutionHead\n# from .ae_multi_stage_head import AEMultiStageHead\n# from .ae_simple_head import AESimpleHead\n# from .deconv_head import DeconvHead\n# from .deeppose_regression_head import DeepposeRegressionHead\n# from .hmr_head import HMRMeshHead\n# from .interhand_3d_head import Interhand3DHead\n# from .temporal_regression_head import TemporalRegressionHead\nfrom .topdown_heatmap_base_head import TopdownHeatmapBaseHead\n# from .topdown_heatmap_multi_stage_head import (TopdownHeatmapMSMUHead,\n# TopdownHeatmapMultiStageHead)\nfrom .topdown_heatmap_simple_head import TopdownHeatmapSimpleHead\n# from .vipnas_heatmap_simple_head import ViPNASHeatmapSimpleHead\n# from .voxelpose_head import CuboidCenterHead, CuboidPoseHead\n\n# __all__ = [\n# 'TopdownHeatmapSimpleHead', 'TopdownHeatmapMultiStageHead',\n# 'TopdownHeatmapMSMUHead', 'TopdownHeatmapBaseHead',\n# 'AEHigherResolutionHead', 'AESimpleHead', 'AEMultiStageHead',\n# 'DeepposeRegressionHead', 'TemporalRegressionHead', 'Interhand3DHead',\n# 'HMRMeshHead', 'DeconvHead', 'ViPNASHeatmapSimpleHead', 'CuboidCenterHead',\n# 'CuboidPoseHead'\n# ]\n" }, { "path": "eval/GVHMR/hmr4d/utils/preproc/vitpose_pytorch/src/vitpose_infer/builder/heads/deconv_head.py", "content": "# Copyright (c) OpenMMLab. All rights reserved.\nimport torch\nimport torch.nn as nn\nfrom mmcv.cnn import (build_conv_layer, build_norm_layer, build_upsample_layer,\n constant_init, normal_init)\n\nfrom mmpose.models.builder import HEADS, build_loss\nfrom mmpose.models.utils.ops import resize\n\n\n@HEADS.register_module()\nclass DeconvHead(nn.Module):\n \"\"\"Simple deconv head.\n\n Args:\n in_channels (int): Number of input channels.\n out_channels (int): Number of output channels.\n num_deconv_layers (int): Number of deconv layers.\n num_deconv_layers should >= 0. Note that 0 means\n no deconv layers.\n num_deconv_filters (list|tuple): Number of filters.\n If num_deconv_layers > 0, the length of\n num_deconv_kernels (list|tuple): Kernel sizes.\n in_index (int|Sequence[int]): Input feature index. Default: 0\n input_transform (str|None): Transformation type of input features.\n Options: 'resize_concat', 'multiple_select', None.\n Default: None.\n\n - 'resize_concat': Multiple feature maps will be resized to the\n same size as the first one and then concat together.\n Usually used in FCN head of HRNet.\n - 'multiple_select': Multiple feature maps will be bundle into\n a list and passed into decode head.\n - None: Only one select feature map is allowed.\n align_corners (bool): align_corners argument of F.interpolate.\n Default: False.\n loss_keypoint (dict): Config for loss. Default: None.\n \"\"\"\n\n def __init__(self,\n in_channels=3,\n out_channels=17,\n num_deconv_layers=3,\n num_deconv_filters=(256, 256, 256),\n num_deconv_kernels=(4, 4, 4),\n extra=None,\n in_index=0,\n input_transform=None,\n align_corners=False,\n loss_keypoint=None):\n super().__init__()\n\n self.in_channels = in_channels\n self.loss = build_loss(loss_keypoint)\n\n self._init_inputs(in_channels, in_index, input_transform)\n self.in_index = in_index\n self.align_corners = align_corners\n\n if extra is not None and not isinstance(extra, dict):\n raise TypeError('extra should be dict or None.')\n\n if num_deconv_layers > 0:\n self.deconv_layers = self._make_deconv_layer(\n num_deconv_layers,\n num_deconv_filters,\n num_deconv_kernels,\n )\n elif num_deconv_layers == 0:\n self.deconv_layers = nn.Identity()\n else:\n raise ValueError(\n f'num_deconv_layers ({num_deconv_layers}) should >= 0.')\n\n identity_final_layer = False\n if extra is not None and 'final_conv_kernel' in extra:\n assert extra['final_conv_kernel'] in [0, 1, 3]\n if extra['final_conv_kernel'] == 3:\n padding = 1\n elif extra['final_conv_kernel'] == 1:\n padding = 0\n else:\n # 0 for Identity mapping.\n identity_final_layer = True\n kernel_size = extra['final_conv_kernel']\n else:\n kernel_size = 1\n padding = 0\n\n if identity_final_layer:\n self.final_layer = nn.Identity()\n else:\n conv_channels = num_deconv_filters[\n -1] if num_deconv_layers > 0 else self.in_channels\n\n layers = []\n if extra is not None:\n num_conv_layers = extra.get('num_conv_layers', 0)\n num_conv_kernels = extra.get('num_conv_kernels',\n [1] * num_conv_layers)\n\n for i in range(num_conv_layers):\n layers.append(\n build_conv_layer(\n dict(type='Conv2d'),\n in_channels=conv_channels,\n out_channels=conv_channels,\n kernel_size=num_conv_kernels[i],\n stride=1,\n padding=(num_conv_kernels[i] - 1) // 2))\n layers.append(\n build_norm_layer(dict(type='BN'), conv_channels)[1])\n layers.append(nn.ReLU(inplace=True))\n\n layers.append(\n build_conv_layer(\n cfg=dict(type='Conv2d'),\n in_channels=conv_channels,\n out_channels=out_channels,\n kernel_size=kernel_size,\n stride=1,\n padding=padding))\n\n if len(layers) > 1:\n self.final_layer = nn.Sequential(*layers)\n else:\n self.final_layer = layers[0]\n\n def _init_inputs(self, in_channels, in_index, input_transform):\n \"\"\"Check and initialize input transforms.\n\n The in_channels, in_index and input_transform must match.\n Specifically, when input_transform is None, only single feature map\n will be selected. So in_channels and in_index must be of type int.\n When input_transform is not None, in_channels and in_index must be\n list or tuple, with the same length.\n\n Args:\n in_channels (int|Sequence[int]): Input channels.\n in_index (int|Sequence[int]): Input feature index.\n input_transform (str|None): Transformation type of input features.\n Options: 'resize_concat', 'multiple_select', None.\n\n - 'resize_concat': Multiple feature maps will be resize to the\n same size as first one and than concat together.\n Usually used in FCN head of HRNet.\n - 'multiple_select': Multiple feature maps will be bundle into\n a list and passed into decode head.\n - None: Only one select feature map is allowed.\n \"\"\"\n\n if input_transform is not None:\n assert input_transform in ['resize_concat', 'multiple_select']\n self.input_transform = input_transform\n self.in_index = in_index\n if input_transform is not None:\n assert isinstance(in_channels, (list, tuple))\n assert isinstance(in_index, (list, tuple))\n assert len(in_channels) == len(in_index)\n if input_transform == 'resize_concat':\n self.in_channels = sum(in_channels)\n else:\n self.in_channels = in_channels\n else:\n assert isinstance(in_channels, int)\n assert isinstance(in_index, int)\n self.in_channels = in_channels\n\n def _transform_inputs(self, inputs):\n \"\"\"Transform inputs for decoder.\n\n Args:\n inputs (list[Tensor] | Tensor): multi-level img features.\n\n Returns:\n Tensor: The transformed inputs\n \"\"\"\n if not isinstance(inputs, list):\n return inputs\n\n if self.input_transform == 'resize_concat':\n inputs = [inputs[i] for i in self.in_index]\n upsampled_inputs = [\n resize(\n input=x,\n size=inputs[0].shape[2:],\n mode='bilinear',\n align_corners=self.align_corners) for x in inputs\n ]\n inputs = torch.cat(upsampled_inputs, dim=1)\n elif self.input_transform == 'multiple_select':\n inputs = [inputs[i] for i in self.in_index]\n else:\n inputs = inputs[self.in_index]\n\n return inputs\n\n def _make_deconv_layer(self, num_layers, num_filters, num_kernels):\n \"\"\"Make deconv layers.\"\"\"\n if num_layers != len(num_filters):\n error_msg = f'num_layers({num_layers}) ' \\\n f'!= length of num_filters({len(num_filters)})'\n raise ValueError(error_msg)\n if num_layers != len(num_kernels):\n error_msg = f'num_layers({num_layers}) ' \\\n f'!= length of num_kernels({len(num_kernels)})'\n raise ValueError(error_msg)\n\n layers = []\n for i in range(num_layers):\n kernel, padding, output_padding = \\\n self._get_deconv_cfg(num_kernels[i])\n\n planes = num_filters[i]\n layers.append(\n build_upsample_layer(\n dict(type='deconv'),\n in_channels=self.in_channels,\n out_channels=planes,\n kernel_size=kernel,\n stride=2,\n padding=padding,\n output_padding=output_padding,\n bias=False))\n layers.append(nn.BatchNorm2d(planes))\n layers.append(nn.ReLU(inplace=True))\n self.in_channels = planes\n\n return nn.Sequential(*layers)\n\n @staticmethod\n def _get_deconv_cfg(deconv_kernel):\n \"\"\"Get configurations for deconv layers.\"\"\"\n if deconv_kernel == 4:\n padding = 1\n output_padding = 0\n elif deconv_kernel == 3:\n padding = 1\n output_padding = 1\n elif deconv_kernel == 2:\n padding = 0\n output_padding = 0\n else:\n raise ValueError(f'Not supported num_kernels ({deconv_kernel}).')\n\n return deconv_kernel, padding, output_padding\n\n def get_loss(self, outputs, targets, masks):\n \"\"\"Calculate bottom-up masked mse loss.\n\n Note:\n - batch_size: N\n - num_channels: C\n - heatmaps height: H\n - heatmaps weight: W\n\n Args:\n outputs (List(torch.Tensor[N,C,H,W])): Multi-scale outputs.\n targets (List(torch.Tensor[N,C,H,W])): Multi-scale targets.\n masks (List(torch.Tensor[N,H,W])): Masks of multi-scale targets.\n \"\"\"\n\n losses = dict()\n\n for idx in range(len(targets)):\n if 'loss' not in losses:\n losses['loss'] = self.loss(outputs[idx], targets[idx],\n masks[idx])\n else:\n losses['loss'] += self.loss(outputs[idx], targets[idx],\n masks[idx])\n\n return losses\n\n def forward(self, x):\n \"\"\"Forward function.\"\"\"\n x = self._transform_inputs(x)\n final_outputs = []\n x = self.deconv_layers(x)\n y = self.final_layer(x)\n final_outputs.append(y)\n return final_outputs\n\n def init_weights(self):\n \"\"\"Initialize model weights.\"\"\"\n for _, m in self.deconv_layers.named_modules():\n if isinstance(m, nn.ConvTranspose2d):\n normal_init(m, std=0.001)\n elif isinstance(m, nn.BatchNorm2d):\n constant_init(m, 1)\n for m in self.final_layer.modules():\n if isinstance(m, nn.Conv2d):\n normal_init(m, std=0.001, bias=0)\n elif isinstance(m, nn.BatchNorm2d):\n constant_init(m, 1)\n" }, { "path": "eval/GVHMR/hmr4d/utils/preproc/vitpose_pytorch/src/vitpose_infer/builder/heads/deeppose_regression_head.py", "content": "# Copyright (c) OpenMMLab. All rights reserved.\nimport numpy as np\nimport torch.nn as nn\nfrom mmcv.cnn import normal_init\n\nfrom mmpose.core.evaluation import (keypoint_pck_accuracy,\n keypoints_from_regression)\nfrom mmpose.core.post_processing import fliplr_regression\nfrom mmpose.models.builder import HEADS, build_loss\n\n\n@HEADS.register_module()\nclass DeepposeRegressionHead(nn.Module):\n \"\"\"Deeppose regression head with fully connected layers.\n\n \"DeepPose: Human Pose Estimation via Deep Neural Networks\".\n\n Args:\n in_channels (int): Number of input channels\n num_joints (int): Number of joints\n loss_keypoint (dict): Config for keypoint loss. Default: None.\n \"\"\"\n\n def __init__(self,\n in_channels,\n num_joints,\n loss_keypoint=None,\n train_cfg=None,\n test_cfg=None):\n super().__init__()\n\n self.in_channels = in_channels\n self.num_joints = num_joints\n\n self.loss = build_loss(loss_keypoint)\n\n self.train_cfg = {} if train_cfg is None else train_cfg\n self.test_cfg = {} if test_cfg is None else test_cfg\n\n self.fc = nn.Linear(self.in_channels, self.num_joints * 2)\n\n def forward(self, x):\n \"\"\"Forward function.\"\"\"\n output = self.fc(x)\n N, C = output.shape\n return output.reshape([N, C // 2, 2])\n\n def get_loss(self, output, target, target_weight):\n \"\"\"Calculate top-down keypoint loss.\n\n Note:\n - batch_size: N\n - num_keypoints: K\n\n Args:\n output (torch.Tensor[N, K, 2]): Output keypoints.\n target (torch.Tensor[N, K, 2]): Target keypoints.\n target_weight (torch.Tensor[N, K, 2]):\n Weights across different joint types.\n \"\"\"\n\n losses = dict()\n assert not isinstance(self.loss, nn.Sequential)\n assert target.dim() == 3 and target_weight.dim() == 3\n losses['reg_loss'] = self.loss(output, target, target_weight)\n\n return losses\n\n def get_accuracy(self, output, target, target_weight):\n \"\"\"Calculate accuracy for top-down keypoint loss.\n\n Note:\n - batch_size: N\n - num_keypoints: K\n\n Args:\n output (torch.Tensor[N, K, 2]): Output keypoints.\n target (torch.Tensor[N, K, 2]): Target keypoints.\n target_weight (torch.Tensor[N, K, 2]):\n Weights across different joint types.\n \"\"\"\n\n accuracy = dict()\n\n N = output.shape[0]\n\n _, avg_acc, cnt = keypoint_pck_accuracy(\n output.detach().cpu().numpy(),\n target.detach().cpu().numpy(),\n target_weight[:, :, 0].detach().cpu().numpy() > 0,\n thr=0.05,\n normalize=np.ones((N, 2), dtype=np.float32))\n accuracy['acc_pose'] = avg_acc\n\n return accuracy\n\n def inference_model(self, x, flip_pairs=None):\n \"\"\"Inference function.\n\n Returns:\n output_regression (np.ndarray): Output regression.\n\n Args:\n x (torch.Tensor[N, K, 2]): Input features.\n flip_pairs (None | list[tuple()):\n Pairs of keypoints which are mirrored.\n \"\"\"\n output = self.forward(x)\n\n if flip_pairs is not None:\n output_regression = fliplr_regression(\n output.detach().cpu().numpy(), flip_pairs)\n else:\n output_regression = output.detach().cpu().numpy()\n return output_regression\n\n def decode(self, img_metas, output, **kwargs):\n \"\"\"Decode the keypoints from output regression.\n\n Args:\n img_metas (list(dict)): Information about data augmentation\n By default this includes:\n\n - \"image_file: path to the image file\n - \"center\": center of the bbox\n - \"scale\": scale of the bbox\n - \"rotation\": rotation of the bbox\n - \"bbox_score\": score of bbox\n output (np.ndarray[N, K, 2]): predicted regression vector.\n kwargs: dict contains 'img_size'.\n img_size (tuple(img_width, img_height)): input image size.\n \"\"\"\n batch_size = len(img_metas)\n\n if 'bbox_id' in img_metas[0]:\n bbox_ids = []\n else:\n bbox_ids = None\n\n c = np.zeros((batch_size, 2), dtype=np.float32)\n s = np.zeros((batch_size, 2), dtype=np.float32)\n image_paths = []\n score = np.ones(batch_size)\n for i in range(batch_size):\n c[i, :] = img_metas[i]['center']\n s[i, :] = img_metas[i]['scale']\n image_paths.append(img_metas[i]['image_file'])\n\n if 'bbox_score' in img_metas[i]:\n score[i] = np.array(img_metas[i]['bbox_score']).reshape(-1)\n if bbox_ids is not None:\n bbox_ids.append(img_metas[i]['bbox_id'])\n\n preds, maxvals = keypoints_from_regression(output, c, s,\n kwargs['img_size'])\n\n all_preds = np.zeros((batch_size, preds.shape[1], 3), dtype=np.float32)\n all_boxes = np.zeros((batch_size, 6), dtype=np.float32)\n all_preds[:, :, 0:2] = preds[:, :, 0:2]\n all_preds[:, :, 2:3] = maxvals\n all_boxes[:, 0:2] = c[:, 0:2]\n all_boxes[:, 2:4] = s[:, 0:2]\n all_boxes[:, 4] = np.prod(s * 200.0, axis=1)\n all_boxes[:, 5] = score\n\n result = {}\n\n result['preds'] = all_preds\n result['boxes'] = all_boxes\n result['image_paths'] = image_paths\n result['bbox_ids'] = bbox_ids\n\n return result\n\n def init_weights(self):\n normal_init(self.fc, mean=0, std=0.01, bias=0)\n" }, { "path": "eval/GVHMR/hmr4d/utils/preproc/vitpose_pytorch/src/vitpose_infer/builder/heads/hmr_head.py", "content": "# Copyright (c) OpenMMLab. All rights reserved.\nimport numpy as np\nimport torch\nimport torch.nn as nn\nfrom mmcv.cnn import xavier_init\n\nfrom ..builder import HEADS\nfrom ..utils.geometry import rot6d_to_rotmat\n\n\n@HEADS.register_module()\nclass HMRMeshHead(nn.Module):\n \"\"\"SMPL parameters regressor head of simple baseline. \"End-to-end Recovery\n of Human Shape and Pose\", CVPR'2018.\n\n Args:\n in_channels (int): Number of input channels\n smpl_mean_params (str): The file name of the mean SMPL parameters\n n_iter (int): The iterations of estimating delta parameters\n \"\"\"\n\n def __init__(self, in_channels, smpl_mean_params=None, n_iter=3):\n super().__init__()\n\n self.in_channels = in_channels\n self.n_iter = n_iter\n\n npose = 24 * 6\n nbeta = 10\n ncam = 3\n hidden_dim = 1024\n\n self.fc1 = nn.Linear(in_channels + npose + nbeta + ncam, hidden_dim)\n self.drop1 = nn.Dropout()\n self.fc2 = nn.Linear(hidden_dim, hidden_dim)\n self.drop2 = nn.Dropout()\n self.decpose = nn.Linear(hidden_dim, npose)\n self.decshape = nn.Linear(hidden_dim, nbeta)\n self.deccam = nn.Linear(hidden_dim, ncam)\n\n # Load mean SMPL parameters\n if smpl_mean_params is None:\n init_pose = torch.zeros([1, npose])\n init_shape = torch.zeros([1, nbeta])\n init_cam = torch.FloatTensor([[1, 0, 0]])\n else:\n mean_params = np.load(smpl_mean_params)\n init_pose = torch.from_numpy(\n mean_params['pose'][:]).unsqueeze(0).float()\n init_shape = torch.from_numpy(\n mean_params['shape'][:]).unsqueeze(0).float()\n init_cam = torch.from_numpy(\n mean_params['cam']).unsqueeze(0).float()\n self.register_buffer('init_pose', init_pose)\n self.register_buffer('init_shape', init_shape)\n self.register_buffer('init_cam', init_cam)\n\n def forward(self, x):\n \"\"\"Forward function.\n\n x is the image feature map and is expected to be in shape (batch size x\n channel number x height x width)\n \"\"\"\n batch_size = x.shape[0]\n # extract the global feature vector by average along\n # spatial dimension.\n x = x.mean(dim=-1).mean(dim=-1)\n\n init_pose = self.init_pose.expand(batch_size, -1)\n init_shape = self.init_shape.expand(batch_size, -1)\n init_cam = self.init_cam.expand(batch_size, -1)\n\n pred_pose = init_pose\n pred_shape = init_shape\n pred_cam = init_cam\n for _ in range(self.n_iter):\n xc = torch.cat([x, pred_pose, pred_shape, pred_cam], 1)\n xc = self.fc1(xc)\n xc = self.drop1(xc)\n xc = self.fc2(xc)\n xc = self.drop2(xc)\n pred_pose = self.decpose(xc) + pred_pose\n pred_shape = self.decshape(xc) + pred_shape\n pred_cam = self.deccam(xc) + pred_cam\n\n pred_rotmat = rot6d_to_rotmat(pred_pose).view(batch_size, 24, 3, 3)\n out = (pred_rotmat, pred_shape, pred_cam)\n return out\n\n def init_weights(self):\n \"\"\"Initialize model weights.\"\"\"\n xavier_init(self.decpose, gain=0.01)\n xavier_init(self.decshape, gain=0.01)\n xavier_init(self.deccam, gain=0.01)\n" }, { "path": "eval/GVHMR/hmr4d/utils/preproc/vitpose_pytorch/src/vitpose_infer/builder/heads/interhand_3d_head.py", "content": "# Copyright (c) OpenMMLab. All rights reserved.\nimport numpy as np\nimport torch\nimport torch.nn as nn\nimport torch.nn.functional as F\nfrom mmcv.cnn import (build_conv_layer, build_norm_layer, build_upsample_layer,\n constant_init, normal_init)\n\nfrom mmpose.core.evaluation.top_down_eval import (\n keypoints_from_heatmaps3d, multilabel_classification_accuracy)\nfrom mmpose.core.post_processing import flip_back\nfrom mmpose.models.builder import build_loss\nfrom mmpose.models.necks import GlobalAveragePooling\nfrom ..builder import HEADS\n\n\nclass Heatmap3DHead(nn.Module):\n \"\"\"Heatmap3DHead is a sub-module of Interhand3DHead, and outputs 3D\n heatmaps. Heatmap3DHead is composed of (>=0) number of deconv layers and a\n simple conv2d layer.\n\n Args:\n in_channels (int): Number of input channels\n out_channels (int): Number of output channels\n depth_size (int): Number of depth discretization size\n num_deconv_layers (int): Number of deconv layers.\n num_deconv_layers should >= 0. Note that 0 means no deconv layers.\n num_deconv_filters (list|tuple): Number of filters.\n num_deconv_kernels (list|tuple): Kernel sizes.\n extra (dict): Configs for extra conv layers. Default: None\n \"\"\"\n\n def __init__(self,\n in_channels,\n out_channels,\n depth_size=64,\n num_deconv_layers=3,\n num_deconv_filters=(256, 256, 256),\n num_deconv_kernels=(4, 4, 4),\n extra=None):\n\n super().__init__()\n\n assert out_channels % depth_size == 0\n self.depth_size = depth_size\n self.in_channels = in_channels\n\n if extra is not None and not isinstance(extra, dict):\n raise TypeError('extra should be dict or None.')\n\n if num_deconv_layers > 0:\n self.deconv_layers = self._make_deconv_layer(\n num_deconv_layers,\n num_deconv_filters,\n num_deconv_kernels,\n )\n elif num_deconv_layers == 0:\n self.deconv_layers = nn.Identity()\n else:\n raise ValueError(\n f'num_deconv_layers ({num_deconv_layers}) should >= 0.')\n\n identity_final_layer = False\n if extra is not None and 'final_conv_kernel' in extra:\n assert extra['final_conv_kernel'] in [0, 1, 3]\n if extra['final_conv_kernel'] == 3:\n padding = 1\n elif extra['final_conv_kernel'] == 1:\n padding = 0\n else:\n # 0 for Identity mapping.\n identity_final_layer = True\n kernel_size = extra['final_conv_kernel']\n else:\n kernel_size = 1\n padding = 0\n\n if identity_final_layer:\n self.final_layer = nn.Identity()\n else:\n conv_channels = num_deconv_filters[\n -1] if num_deconv_layers > 0 else self.in_channels\n\n layers = []\n if extra is not None:\n num_conv_layers = extra.get('num_conv_layers', 0)\n num_conv_kernels = extra.get('num_conv_kernels',\n [1] * num_conv_layers)\n\n for i in range(num_conv_layers):\n layers.append(\n build_conv_layer(\n dict(type='Conv2d'),\n in_channels=conv_channels,\n out_channels=conv_channels,\n kernel_size=num_conv_kernels[i],\n stride=1,\n padding=(num_conv_kernels[i] - 1) // 2))\n layers.append(\n build_norm_layer(dict(type='BN'), conv_channels)[1])\n layers.append(nn.ReLU(inplace=True))\n\n layers.append(\n build_conv_layer(\n cfg=dict(type='Conv2d'),\n in_channels=conv_channels,\n out_channels=out_channels,\n kernel_size=kernel_size,\n stride=1,\n padding=padding))\n\n if len(layers) > 1:\n self.final_layer = nn.Sequential(*layers)\n else:\n self.final_layer = layers[0]\n\n def _make_deconv_layer(self, num_layers, num_filters, num_kernels):\n \"\"\"Make deconv layers.\"\"\"\n if num_layers != len(num_filters):\n error_msg = f'num_layers({num_layers}) ' \\\n f'!= length of num_filters({len(num_filters)})'\n raise ValueError(error_msg)\n if num_layers != len(num_kernels):\n error_msg = f'num_layers({num_layers}) ' \\\n f'!= length of num_kernels({len(num_kernels)})'\n raise ValueError(error_msg)\n\n layers = []\n for i in range(num_layers):\n kernel, padding, output_padding = \\\n self._get_deconv_cfg(num_kernels[i])\n\n planes = num_filters[i]\n layers.append(\n build_upsample_layer(\n dict(type='deconv'),\n in_channels=self.in_channels,\n out_channels=planes,\n kernel_size=kernel,\n stride=2,\n padding=padding,\n output_padding=output_padding,\n bias=False))\n layers.append(nn.BatchNorm2d(planes))\n layers.append(nn.ReLU(inplace=True))\n self.in_channels = planes\n\n return nn.Sequential(*layers)\n\n @staticmethod\n def _get_deconv_cfg(deconv_kernel):\n \"\"\"Get configurations for deconv layers.\"\"\"\n if deconv_kernel == 4:\n padding = 1\n output_padding = 0\n elif deconv_kernel == 3:\n padding = 1\n output_padding = 1\n elif deconv_kernel == 2:\n padding = 0\n output_padding = 0\n else:\n raise ValueError(f'Not supported num_kernels ({deconv_kernel}).')\n\n return deconv_kernel, padding, output_padding\n\n def forward(self, x):\n \"\"\"Forward function.\"\"\"\n x = self.deconv_layers(x)\n x = self.final_layer(x)\n N, C, H, W = x.shape\n # reshape the 2D heatmap to 3D heatmap\n x = x.reshape(N, C // self.depth_size, self.depth_size, H, W)\n return x\n\n def init_weights(self):\n \"\"\"Initialize model weights.\"\"\"\n for _, m in self.deconv_layers.named_modules():\n if isinstance(m, nn.ConvTranspose2d):\n normal_init(m, std=0.001)\n elif isinstance(m, nn.BatchNorm2d):\n constant_init(m, 1)\n for m in self.final_layer.modules():\n if isinstance(m, nn.Conv2d):\n normal_init(m, std=0.001, bias=0)\n elif isinstance(m, nn.BatchNorm2d):\n constant_init(m, 1)\n\n\nclass Heatmap1DHead(nn.Module):\n \"\"\"Heatmap1DHead is a sub-module of Interhand3DHead, and outputs 1D\n heatmaps.\n\n Args:\n in_channels (int): Number of input channels\n heatmap_size (int): Heatmap size\n hidden_dims (list|tuple): Number of feature dimension of FC layers.\n \"\"\"\n\n def __init__(self, in_channels=2048, heatmap_size=64, hidden_dims=(512, )):\n super().__init__()\n\n self.in_channels = in_channels\n self.heatmap_size = heatmap_size\n\n feature_dims = [in_channels, *hidden_dims, heatmap_size]\n self.fc = self._make_linear_layers(feature_dims, relu_final=False)\n\n def soft_argmax_1d(self, heatmap1d):\n heatmap1d = F.softmax(heatmap1d, 1)\n accu = heatmap1d * torch.arange(\n self.heatmap_size, dtype=heatmap1d.dtype,\n device=heatmap1d.device)[None, :]\n coord = accu.sum(dim=1)\n return coord\n\n def _make_linear_layers(self, feat_dims, relu_final=False):\n \"\"\"Make linear layers.\"\"\"\n layers = []\n for i in range(len(feat_dims) - 1):\n layers.append(nn.Linear(feat_dims[i], feat_dims[i + 1]))\n if i < len(feat_dims) - 2 or \\\n (i == len(feat_dims) - 2 and relu_final):\n layers.append(nn.ReLU(inplace=True))\n return nn.Sequential(*layers)\n\n def forward(self, x):\n \"\"\"Forward function.\"\"\"\n heatmap1d = self.fc(x)\n value = self.soft_argmax_1d(heatmap1d).view(-1, 1)\n return value\n\n def init_weights(self):\n \"\"\"Initialize model weights.\"\"\"\n for m in self.fc.modules():\n if isinstance(m, nn.Linear):\n normal_init(m, mean=0, std=0.01, bias=0)\n\n\nclass MultilabelClassificationHead(nn.Module):\n \"\"\"MultilabelClassificationHead is a sub-module of Interhand3DHead, and\n outputs hand type classification.\n\n Args:\n in_channels (int): Number of input channels\n num_labels (int): Number of labels\n hidden_dims (list|tuple): Number of hidden dimension of FC layers.\n \"\"\"\n\n def __init__(self, in_channels=2048, num_labels=2, hidden_dims=(512, )):\n super().__init__()\n\n self.in_channels = in_channels\n self.num_labesl = num_labels\n\n feature_dims = [in_channels, *hidden_dims, num_labels]\n self.fc = self._make_linear_layers(feature_dims, relu_final=False)\n\n def _make_linear_layers(self, feat_dims, relu_final=False):\n \"\"\"Make linear layers.\"\"\"\n layers = []\n for i in range(len(feat_dims) - 1):\n layers.append(nn.Linear(feat_dims[i], feat_dims[i + 1]))\n if i < len(feat_dims) - 2 or \\\n (i == len(feat_dims) - 2 and relu_final):\n layers.append(nn.ReLU(inplace=True))\n return nn.Sequential(*layers)\n\n def forward(self, x):\n \"\"\"Forward function.\"\"\"\n labels = torch.sigmoid(self.fc(x))\n return labels\n\n def init_weights(self):\n for m in self.fc.modules():\n if isinstance(m, nn.Linear):\n normal_init(m, mean=0, std=0.01, bias=0)\n\n\n@HEADS.register_module()\nclass Interhand3DHead(nn.Module):\n \"\"\"Interhand 3D head of paper ref: Gyeongsik Moon. \"InterHand2.6M: A\n Dataset and Baseline for 3D Interacting Hand Pose Estimation from a Single\n RGB Image\".\n\n Args:\n keypoint_head_cfg (dict): Configs of Heatmap3DHead for hand\n keypoint estimation.\n root_head_cfg (dict): Configs of Heatmap1DHead for relative\n hand root depth estimation.\n hand_type_head_cfg (dict): Configs of MultilabelClassificationHead\n for hand type classification.\n loss_keypoint (dict): Config for keypoint loss. Default: None.\n loss_root_depth (dict): Config for relative root depth loss.\n Default: None.\n loss_hand_type (dict): Config for hand type classification\n loss. Default: None.\n \"\"\"\n\n def __init__(self,\n keypoint_head_cfg,\n root_head_cfg,\n hand_type_head_cfg,\n loss_keypoint=None,\n loss_root_depth=None,\n loss_hand_type=None,\n train_cfg=None,\n test_cfg=None):\n super().__init__()\n\n # build sub-module heads\n self.right_hand_head = Heatmap3DHead(**keypoint_head_cfg)\n self.left_hand_head = Heatmap3DHead(**keypoint_head_cfg)\n self.root_head = Heatmap1DHead(**root_head_cfg)\n self.hand_type_head = MultilabelClassificationHead(\n **hand_type_head_cfg)\n self.neck = GlobalAveragePooling()\n\n # build losses\n self.keypoint_loss = build_loss(loss_keypoint)\n self.root_depth_loss = build_loss(loss_root_depth)\n self.hand_type_loss = build_loss(loss_hand_type)\n self.train_cfg = {} if train_cfg is None else train_cfg\n self.test_cfg = {} if test_cfg is None else test_cfg\n self.target_type = self.test_cfg.get('target_type', 'GaussianHeatmap')\n\n def init_weights(self):\n self.left_hand_head.init_weights()\n self.right_hand_head.init_weights()\n self.root_head.init_weights()\n self.hand_type_head.init_weights()\n\n def get_loss(self, output, target, target_weight):\n \"\"\"Calculate loss for hand keypoint heatmaps, relative root depth and\n hand type.\n\n Args:\n output (list[Tensor]): a list of outputs from multiple heads.\n target (list[Tensor]): a list of targets for multiple heads.\n target_weight (list[Tensor]): a list of targets weight for\n multiple heads.\n \"\"\"\n losses = dict()\n\n # hand keypoint loss\n assert not isinstance(self.keypoint_loss, nn.Sequential)\n out, tar, tar_weight = output[0], target[0], target_weight[0]\n assert tar.dim() == 5 and tar_weight.dim() == 3\n losses['hand_loss'] = self.keypoint_loss(out, tar, tar_weight)\n\n # relative root depth loss\n assert not isinstance(self.root_depth_loss, nn.Sequential)\n out, tar, tar_weight = output[1], target[1], target_weight[1]\n assert tar.dim() == 2 and tar_weight.dim() == 2\n losses['rel_root_loss'] = self.root_depth_loss(out, tar, tar_weight)\n\n # hand type loss\n assert not isinstance(self.hand_type_loss, nn.Sequential)\n out, tar, tar_weight = output[2], target[2], target_weight[2]\n assert tar.dim() == 2 and tar_weight.dim() in [1, 2]\n losses['hand_type_loss'] = self.hand_type_loss(out, tar, tar_weight)\n\n return losses\n\n def get_accuracy(self, output, target, target_weight):\n \"\"\"Calculate accuracy for hand type.\n\n Args:\n output (list[Tensor]): a list of outputs from multiple heads.\n target (list[Tensor]): a list of targets for multiple heads.\n target_weight (list[Tensor]): a list of targets weight for\n multiple heads.\n \"\"\"\n accuracy = dict()\n avg_acc = multilabel_classification_accuracy(\n output[2].detach().cpu().numpy(),\n target[2].detach().cpu().numpy(),\n target_weight[2].detach().cpu().numpy(),\n )\n accuracy['acc_classification'] = float(avg_acc)\n return accuracy\n\n def forward(self, x):\n \"\"\"Forward function.\"\"\"\n outputs = []\n outputs.append(\n torch.cat([self.right_hand_head(x),\n self.left_hand_head(x)], dim=1))\n x = self.neck(x)\n outputs.append(self.root_head(x))\n outputs.append(self.hand_type_head(x))\n return outputs\n\n def inference_model(self, x, flip_pairs=None):\n \"\"\"Inference function.\n\n Returns:\n output (list[np.ndarray]): list of output hand keypoint\n heatmaps, relative root depth and hand type.\n\n Args:\n x (torch.Tensor[N,K,H,W]): Input features.\n flip_pairs (None | list[tuple()):\n Pairs of keypoints which are mirrored.\n \"\"\"\n\n output = self.forward(x)\n\n if flip_pairs is not None:\n # flip 3D heatmap\n heatmap_3d = output[0]\n N, K, D, H, W = heatmap_3d.shape\n # reshape 3D heatmap to 2D heatmap\n heatmap_3d = heatmap_3d.reshape(N, K * D, H, W)\n # 2D heatmap flip\n heatmap_3d_flipped_back = flip_back(\n heatmap_3d.detach().cpu().numpy(),\n flip_pairs,\n target_type=self.target_type)\n # reshape back to 3D heatmap\n heatmap_3d_flipped_back = heatmap_3d_flipped_back.reshape(\n N, K, D, H, W)\n # feature is not aligned, shift flipped heatmap for higher accuracy\n if self.test_cfg.get('shift_heatmap', False):\n heatmap_3d_flipped_back[...,\n 1:] = heatmap_3d_flipped_back[..., :-1]\n output[0] = heatmap_3d_flipped_back\n\n # flip relative hand root depth\n output[1] = -output[1].detach().cpu().numpy()\n\n # flip hand type\n hand_type = output[2].detach().cpu().numpy()\n hand_type_flipped_back = hand_type.copy()\n hand_type_flipped_back[:, 0] = hand_type[:, 1]\n hand_type_flipped_back[:, 1] = hand_type[:, 0]\n output[2] = hand_type_flipped_back\n else:\n output = [out.detach().cpu().numpy() for out in output]\n\n return output\n\n def decode(self, img_metas, output, **kwargs):\n \"\"\"Decode hand keypoint, relative root depth and hand type.\n\n Args:\n img_metas (list(dict)): Information about data augmentation\n By default this includes:\n\n - \"image_file: path to the image file\n - \"center\": center of the bbox\n - \"scale\": scale of the bbox\n - \"rotation\": rotation of the bbox\n - \"bbox_score\": score of bbox\n - \"heatmap3d_depth_bound\": depth bound of hand keypoint\n 3D heatmap\n - \"root_depth_bound\": depth bound of relative root depth\n 1D heatmap\n output (list[np.ndarray]): model predicted 3D heatmaps, relative\n root depth and hand type.\n \"\"\"\n\n batch_size = len(img_metas)\n result = {}\n\n heatmap3d_depth_bound = np.ones(batch_size, dtype=np.float32)\n root_depth_bound = np.ones(batch_size, dtype=np.float32)\n center = np.zeros((batch_size, 2), dtype=np.float32)\n scale = np.zeros((batch_size, 2), dtype=np.float32)\n image_paths = []\n score = np.ones(batch_size, dtype=np.float32)\n if 'bbox_id' in img_metas[0]:\n bbox_ids = []\n else:\n bbox_ids = None\n\n for i in range(batch_size):\n heatmap3d_depth_bound[i] = img_metas[i]['heatmap3d_depth_bound']\n root_depth_bound[i] = img_metas[i]['root_depth_bound']\n center[i, :] = img_metas[i]['center']\n scale[i, :] = img_metas[i]['scale']\n image_paths.append(img_metas[i]['image_file'])\n\n if 'bbox_score' in img_metas[i]:\n score[i] = np.array(img_metas[i]['bbox_score']).reshape(-1)\n if bbox_ids is not None:\n bbox_ids.append(img_metas[i]['bbox_id'])\n\n all_boxes = np.zeros((batch_size, 6), dtype=np.float32)\n all_boxes[:, 0:2] = center[:, 0:2]\n all_boxes[:, 2:4] = scale[:, 0:2]\n # scale is defined as: bbox_size / 200.0, so we\n # need multiply 200.0 to get bbox size\n all_boxes[:, 4] = np.prod(scale * 200.0, axis=1)\n all_boxes[:, 5] = score\n result['boxes'] = all_boxes\n result['image_paths'] = image_paths\n result['bbox_ids'] = bbox_ids\n\n # decode 3D heatmaps of hand keypoints\n heatmap3d = output[0]\n preds, maxvals = keypoints_from_heatmaps3d(heatmap3d, center, scale)\n keypoints_3d = np.zeros((batch_size, preds.shape[1], 4),\n dtype=np.float32)\n keypoints_3d[:, :, 0:3] = preds[:, :, 0:3]\n keypoints_3d[:, :, 3:4] = maxvals\n # transform keypoint depth to camera space\n keypoints_3d[:, :, 2] = \\\n (keypoints_3d[:, :, 2] / self.right_hand_head.depth_size - 0.5) \\\n * heatmap3d_depth_bound[:, np.newaxis]\n\n result['preds'] = keypoints_3d\n\n # decode relative hand root depth\n # transform relative root depth to camera space\n result['rel_root_depth'] = (output[1] / self.root_head.heatmap_size -\n 0.5) * root_depth_bound\n\n # decode hand type\n result['hand_type'] = output[2] > 0.5\n return result\n" }, { "path": "eval/GVHMR/hmr4d/utils/preproc/vitpose_pytorch/src/vitpose_infer/builder/heads/temporal_regression_head.py", "content": "# Copyright (c) OpenMMLab. All rights reserved.\nimport numpy as np\nimport torch.nn as nn\nfrom mmcv.cnn import build_conv_layer, constant_init, kaiming_init\nfrom mmcv.utils.parrots_wrapper import _BatchNorm\n\nfrom mmpose.core import (WeightNormClipHook, compute_similarity_transform,\n fliplr_regression)\nfrom mmpose.models.builder import HEADS, build_loss\n\n\n@HEADS.register_module()\nclass TemporalRegressionHead(nn.Module):\n \"\"\"Regression head of VideoPose3D.\n\n \"3D human pose estimation in video with temporal convolutions and\n semi-supervised training\", CVPR'2019.\n\n Args:\n in_channels (int): Number of input channels\n num_joints (int): Number of joints\n loss_keypoint (dict): Config for keypoint loss. Default: None.\n max_norm (float|None): if not None, the weight of convolution layers\n will be clipped to have a maximum norm of max_norm.\n is_trajectory (bool): If the model only predicts root joint\n position, then this arg should be set to True. In this case,\n traj_loss will be calculated. Otherwise, it should be set to\n False. Default: False.\n \"\"\"\n\n def __init__(self,\n in_channels,\n num_joints,\n max_norm=None,\n loss_keypoint=None,\n is_trajectory=False,\n train_cfg=None,\n test_cfg=None):\n super().__init__()\n\n self.in_channels = in_channels\n self.num_joints = num_joints\n self.max_norm = max_norm\n self.loss = build_loss(loss_keypoint)\n self.is_trajectory = is_trajectory\n if self.is_trajectory:\n assert self.num_joints == 1\n\n self.train_cfg = {} if train_cfg is None else train_cfg\n self.test_cfg = {} if test_cfg is None else test_cfg\n\n self.conv = build_conv_layer(\n dict(type='Conv1d'), in_channels, num_joints * 3, 1)\n\n if self.max_norm is not None:\n # Apply weight norm clip to conv layers\n weight_clip = WeightNormClipHook(self.max_norm)\n for module in self.modules():\n if isinstance(module, nn.modules.conv._ConvNd):\n weight_clip.register(module)\n\n @staticmethod\n def _transform_inputs(x):\n \"\"\"Transform inputs for decoder.\n\n Args:\n inputs (tuple or list of Tensor | Tensor): multi-level features.\n\n Returns:\n Tensor: The transformed inputs\n \"\"\"\n if not isinstance(x, (list, tuple)):\n return x\n\n assert len(x) > 0\n\n # return the top-level feature of the 1D feature pyramid\n return x[-1]\n\n def forward(self, x):\n \"\"\"Forward function.\"\"\"\n x = self._transform_inputs(x)\n\n assert x.ndim == 3 and x.shape[2] == 1, f'Invalid shape {x.shape}'\n output = self.conv(x)\n N = output.shape[0]\n return output.reshape(N, self.num_joints, 3)\n\n def get_loss(self, output, target, target_weight):\n \"\"\"Calculate keypoint loss.\n\n Note:\n - batch_size: N\n - num_keypoints: K\n\n Args:\n output (torch.Tensor[N, K, 3]): Output keypoints.\n target (torch.Tensor[N, K, 3]): Target keypoints.\n target_weight (torch.Tensor[N, K, 3]):\n Weights across different joint types.\n If self.is_trajectory is True and target_weight is None,\n target_weight will be set inversely proportional to joint\n depth.\n \"\"\"\n losses = dict()\n assert not isinstance(self.loss, nn.Sequential)\n\n # trajectory model\n if self.is_trajectory:\n if target.dim() == 2:\n target.unsqueeze_(1)\n\n if target_weight is None:\n target_weight = (1 / target[:, :, 2:]).expand(target.shape)\n assert target.dim() == 3 and target_weight.dim() == 3\n\n losses['traj_loss'] = self.loss(output, target, target_weight)\n\n # pose model\n else:\n if target_weight is None:\n target_weight = target.new_ones(target.shape)\n assert target.dim() == 3 and target_weight.dim() == 3\n losses['reg_loss'] = self.loss(output, target, target_weight)\n\n return losses\n\n def get_accuracy(self, output, target, target_weight, metas):\n \"\"\"Calculate accuracy for keypoint loss.\n\n Note:\n - batch_size: N\n - num_keypoints: K\n\n Args:\n output (torch.Tensor[N, K, 3]): Output keypoints.\n target (torch.Tensor[N, K, 3]): Target keypoints.\n target_weight (torch.Tensor[N, K, 3]):\n Weights across different joint types.\n metas (list(dict)): Information about data augmentation including:\n\n - target_image_path (str): Optional, path to the image file\n - target_mean (float): Optional, normalization parameter of\n the target pose.\n - target_std (float): Optional, normalization parameter of the\n target pose.\n - root_position (np.ndarray[3,1]): Optional, global\n position of the root joint.\n - root_index (torch.ndarray[1,]): Optional, original index of\n the root joint before root-centering.\n \"\"\"\n\n accuracy = dict()\n\n N = output.shape[0]\n output_ = output.detach().cpu().numpy()\n target_ = target.detach().cpu().numpy()\n # Denormalize the predicted pose\n if 'target_mean' in metas[0] and 'target_std' in metas[0]:\n target_mean = np.stack([m['target_mean'] for m in metas])\n target_std = np.stack([m['target_std'] for m in metas])\n output_ = self._denormalize_joints(output_, target_mean,\n target_std)\n target_ = self._denormalize_joints(target_, target_mean,\n target_std)\n\n # Restore global position\n if self.test_cfg.get('restore_global_position', False):\n root_pos = np.stack([m['root_position'] for m in metas])\n root_idx = metas[0].get('root_position_index', None)\n output_ = self._restore_global_position(output_, root_pos,\n root_idx)\n target_ = self._restore_global_position(target_, root_pos,\n root_idx)\n # Get target weight\n if target_weight is None:\n target_weight_ = np.ones_like(target_)\n else:\n target_weight_ = target_weight.detach().cpu().numpy()\n if self.test_cfg.get('restore_global_position', False):\n root_idx = metas[0].get('root_position_index', None)\n root_weight = metas[0].get('root_joint_weight', 1.0)\n target_weight_ = self._restore_root_target_weight(\n target_weight_, root_weight, root_idx)\n\n mpjpe = np.mean(\n np.linalg.norm((output_ - target_) * target_weight_, axis=-1))\n\n transformed_output = np.zeros_like(output_)\n for i in range(N):\n transformed_output[i, :, :] = compute_similarity_transform(\n output_[i, :, :], target_[i, :, :])\n p_mpjpe = np.mean(\n np.linalg.norm(\n (transformed_output - target_) * target_weight_, axis=-1))\n\n accuracy['mpjpe'] = output.new_tensor(mpjpe)\n accuracy['p_mpjpe'] = output.new_tensor(p_mpjpe)\n\n return accuracy\n\n def inference_model(self, x, flip_pairs=None):\n \"\"\"Inference function.\n\n Returns:\n output_regression (np.ndarray): Output regression.\n\n Args:\n x (torch.Tensor[N, K, 2]): Input features.\n flip_pairs (None | list[tuple()):\n Pairs of keypoints which are mirrored.\n \"\"\"\n output = self.forward(x)\n\n if flip_pairs is not None:\n output_regression = fliplr_regression(\n output.detach().cpu().numpy(),\n flip_pairs,\n center_mode='static',\n center_x=0)\n else:\n output_regression = output.detach().cpu().numpy()\n return output_regression\n\n def decode(self, metas, output):\n \"\"\"Decode the keypoints from output regression.\n\n Args:\n metas (list(dict)): Information about data augmentation.\n By default this includes:\n\n - \"target_image_path\": path to the image file\n output (np.ndarray[N, K, 3]): predicted regression vector.\n metas (list(dict)): Information about data augmentation including:\n\n - target_image_path (str): Optional, path to the image file\n - target_mean (float): Optional, normalization parameter of\n the target pose.\n - target_std (float): Optional, normalization parameter of the\n target pose.\n - root_position (np.ndarray[3,1]): Optional, global\n position of the root joint.\n - root_index (torch.ndarray[1,]): Optional, original index of\n the root joint before root-centering.\n \"\"\"\n\n # Denormalize the predicted pose\n if 'target_mean' in metas[0] and 'target_std' in metas[0]:\n target_mean = np.stack([m['target_mean'] for m in metas])\n target_std = np.stack([m['target_std'] for m in metas])\n output = self._denormalize_joints(output, target_mean, target_std)\n\n # Restore global position\n if self.test_cfg.get('restore_global_position', False):\n root_pos = np.stack([m['root_position'] for m in metas])\n root_idx = metas[0].get('root_position_index', None)\n output = self._restore_global_position(output, root_pos, root_idx)\n\n target_image_paths = [m.get('target_image_path', None) for m in metas]\n result = {'preds': output, 'target_image_paths': target_image_paths}\n\n return result\n\n @staticmethod\n def _denormalize_joints(x, mean, std):\n \"\"\"Denormalize joint coordinates with given statistics mean and std.\n\n Args:\n x (np.ndarray[N, K, 3]): Normalized joint coordinates.\n mean (np.ndarray[K, 3]): Mean value.\n std (np.ndarray[K, 3]): Std value.\n \"\"\"\n assert x.ndim == 3\n assert x.shape == mean.shape == std.shape\n\n return x * std + mean\n\n @staticmethod\n def _restore_global_position(x, root_pos, root_idx=None):\n \"\"\"Restore global position of the root-centered joints.\n\n Args:\n x (np.ndarray[N, K, 3]): root-centered joint coordinates\n root_pos (np.ndarray[N,1,3]): The global position of the\n root joint.\n root_idx (int|None): If not none, the root joint will be inserted\n back to the pose at the given index.\n \"\"\"\n x = x + root_pos\n if root_idx is not None:\n x = np.insert(x, root_idx, root_pos.squeeze(1), axis=1)\n return x\n\n @staticmethod\n def _restore_root_target_weight(target_weight, root_weight, root_idx=None):\n \"\"\"Restore the target weight of the root joint after the restoration of\n the global position.\n\n Args:\n target_weight (np.ndarray[N, K, 1]): Target weight of relativized\n joints.\n root_weight (float): The target weight value of the root joint.\n root_idx (int|None): If not none, the root joint weight will be\n inserted back to the target weight at the given index.\n \"\"\"\n if root_idx is not None:\n root_weight = np.full(\n target_weight.shape[0], root_weight, dtype=target_weight.dtype)\n target_weight = np.insert(\n target_weight, root_idx, root_weight[:, None], axis=1)\n return target_weight\n\n def init_weights(self):\n \"\"\"Initialize the weights.\"\"\"\n for m in self.modules():\n if isinstance(m, nn.modules.conv._ConvNd):\n kaiming_init(m, mode='fan_in', nonlinearity='relu')\n elif isinstance(m, _BatchNorm):\n constant_init(m, 1)\n" }, { "path": "eval/GVHMR/hmr4d/utils/preproc/vitpose_pytorch/src/vitpose_infer/builder/heads/topdown_heatmap_base_head.py", "content": "# Copyright (c) OpenMMLab. All rights reserved.\nfrom abc import ABCMeta, abstractmethod\n\nimport numpy as np\nimport torch.nn as nn\n\n# from mmpose.core.evaluation.top_down_eval import keypoints_from_heatmaps\n\n\nclass TopdownHeatmapBaseHead(nn.Module):\n \"\"\"Base class for top-down heatmap heads.\n\n All top-down heatmap heads should subclass it.\n All subclass should overwrite:\n\n Methods:`get_loss`, supporting to calculate loss.\n Methods:`get_accuracy`, supporting to calculate accuracy.\n Methods:`forward`, supporting to forward model.\n Methods:`inference_model`, supporting to inference model.\n \"\"\"\n\n __metaclass__ = ABCMeta\n\n @abstractmethod\n def get_loss(self, **kwargs):\n \"\"\"Gets the loss.\"\"\"\n\n @abstractmethod\n def get_accuracy(self, **kwargs):\n \"\"\"Gets the accuracy.\"\"\"\n\n @abstractmethod\n def forward(self, **kwargs):\n \"\"\"Forward function.\"\"\"\n\n @abstractmethod\n def inference_model(self, **kwargs):\n \"\"\"Inference function.\"\"\"\n\n def decode(self, img_metas, output, **kwargs):\n \"\"\"Decode keypoints from heatmaps.\n\n Args:\n img_metas (list(dict)): Information about data augmentation\n By default this includes:\n\n - \"image_file: path to the image file\n - \"center\": center of the bbox\n - \"scale\": scale of the bbox\n - \"rotation\": rotation of the bbox\n - \"bbox_score\": score of bbox\n output (np.ndarray[N, K, H, W]): model predicted heatmaps.\n \"\"\"\n # batch_size = len(img_metas)\n\n # if 'bbox_id' in img_metas[0]:\n # bbox_ids = []\n # else:\n # bbox_ids = None\n\n # c = np.zeros((batch_size, 2), dtype=np.float32)\n # s = np.zeros((batch_size, 2), dtype=np.float32)\n # image_paths = []\n # score = np.ones(batch_size)\n # for i in range(batch_size):\n # c[i, :] = img_metas[i]['center']\n # s[i, :] = img_metas[i]['scale']\n # image_paths.append(img_metas[i]['image_file'])\n\n # if 'bbox_score' in img_metas[i]:\n # score[i] = np.array(img_metas[i]['bbox_score']).reshape(-1)\n # if bbox_ids is not None:\n # bbox_ids.append(img_metas[i]['bbox_id'])\n\n # preds, maxvals = keypoints_from_heatmaps(\n # output,\n # c,\n # s,\n # unbiased=self.test_cfg.get('unbiased_decoding', False),\n # post_process=self.test_cfg.get('post_process', 'default'),\n # kernel=self.test_cfg.get('modulate_kernel', 11),\n # valid_radius_factor=self.test_cfg.get('valid_radius_factor',\n # 0.0546875),\n # use_udp=self.test_cfg.get('use_udp', False),\n # target_type=self.test_cfg.get('target_type', 'GaussianHeatmap'))\n\n # all_preds = np.zeros((batch_size, preds.shape[1], 3), dtype=np.float32)\n # all_boxes = np.zeros((batch_size, 6), dtype=np.float32)\n # all_preds[:, :, 0:2] = preds[:, :, 0:2]\n # all_preds[:, :, 2:3] = maxvals\n # all_boxes[:, 0:2] = c[:, 0:2]\n # all_boxes[:, 2:4] = s[:, 0:2]\n # all_boxes[:, 4] = np.prod(s * 200.0, axis=1)\n # all_boxes[:, 5] = score\n\n # result = {}\n\n # result['preds'] = all_preds\n # result['boxes'] = all_boxes\n # result['image_paths'] = image_paths\n # result['bbox_ids'] = bbox_ids\n\n return None\n\n @staticmethod\n def _get_deconv_cfg(deconv_kernel):\n \"\"\"Get configurations for deconv layers.\"\"\"\n if deconv_kernel == 4:\n padding = 1\n output_padding = 0\n elif deconv_kernel == 3:\n padding = 1\n output_padding = 1\n elif deconv_kernel == 2:\n padding = 0\n output_padding = 0\n else:\n raise ValueError(f'Not supported num_kernels ({deconv_kernel}).')\n\n return deconv_kernel, padding, output_padding\n" }, { "path": "eval/GVHMR/hmr4d/utils/preproc/vitpose_pytorch/src/vitpose_infer/builder/heads/topdown_heatmap_multi_stage_head.py", "content": "# Copyright (c) OpenMMLab. All rights reserved.\nimport copy as cp\n\nimport torch.nn as nn\nfrom mmcv.cnn import (ConvModule, DepthwiseSeparableConvModule, Linear,\n build_activation_layer, build_conv_layer,\n build_norm_layer, build_upsample_layer, constant_init,\n kaiming_init, normal_init)\n\nfrom mmpose.core.evaluation import pose_pck_accuracy\nfrom mmpose.core.post_processing import flip_back\nfrom mmpose.models.builder import build_loss\nfrom ..builder import HEADS\nfrom .topdown_heatmap_base_head import TopdownHeatmapBaseHead\n\n\n@HEADS.register_module()\nclass TopdownHeatmapMultiStageHead(TopdownHeatmapBaseHead):\n \"\"\"Top-down heatmap multi-stage head.\n\n TopdownHeatmapMultiStageHead is consisted of multiple branches,\n each of which has num_deconv_layers(>=0) number of deconv layers\n and a simple conv2d layer.\n\n Args:\n in_channels (int): Number of input channels.\n out_channels (int): Number of output channels.\n num_stages (int): Number of stages.\n num_deconv_layers (int): Number of deconv layers.\n num_deconv_layers should >= 0. Note that 0 means\n no deconv layers.\n num_deconv_filters (list|tuple): Number of filters.\n If num_deconv_layers > 0, the length of\n num_deconv_kernels (list|tuple): Kernel sizes.\n loss_keypoint (dict): Config for keypoint loss. Default: None.\n \"\"\"\n\n def __init__(self,\n in_channels=512,\n out_channels=17,\n num_stages=1,\n num_deconv_layers=3,\n num_deconv_filters=(256, 256, 256),\n num_deconv_kernels=(4, 4, 4),\n extra=None,\n loss_keypoint=None,\n train_cfg=None,\n test_cfg=None):\n super().__init__()\n\n self.in_channels = in_channels\n self.num_stages = num_stages\n self.loss = build_loss(loss_keypoint)\n\n self.train_cfg = {} if train_cfg is None else train_cfg\n self.test_cfg = {} if test_cfg is None else test_cfg\n self.target_type = self.test_cfg.get('target_type', 'GaussianHeatmap')\n\n if extra is not None and not isinstance(extra, dict):\n raise TypeError('extra should be dict or None.')\n\n # build multi-stage deconv layers\n self.multi_deconv_layers = nn.ModuleList([])\n for _ in range(self.num_stages):\n if num_deconv_layers > 0:\n deconv_layers = self._make_deconv_layer(\n num_deconv_layers,\n num_deconv_filters,\n num_deconv_kernels,\n )\n elif num_deconv_layers == 0:\n deconv_layers = nn.Identity()\n else:\n raise ValueError(\n f'num_deconv_layers ({num_deconv_layers}) should >= 0.')\n self.multi_deconv_layers.append(deconv_layers)\n\n identity_final_layer = False\n if extra is not None and 'final_conv_kernel' in extra:\n assert extra['final_conv_kernel'] in [0, 1, 3]\n if extra['final_conv_kernel'] == 3:\n padding = 1\n elif extra['final_conv_kernel'] == 1:\n padding = 0\n else:\n # 0 for Identity mapping.\n identity_final_layer = True\n kernel_size = extra['final_conv_kernel']\n else:\n kernel_size = 1\n padding = 0\n\n # build multi-stage final layers\n self.multi_final_layers = nn.ModuleList([])\n for i in range(self.num_stages):\n if identity_final_layer:\n final_layer = nn.Identity()\n else:\n final_layer = build_conv_layer(\n cfg=dict(type='Conv2d'),\n in_channels=num_deconv_filters[-1]\n if num_deconv_layers > 0 else in_channels,\n out_channels=out_channels,\n kernel_size=kernel_size,\n stride=1,\n padding=padding)\n self.multi_final_layers.append(final_layer)\n\n def get_loss(self, output, target, target_weight):\n \"\"\"Calculate top-down keypoint loss.\n\n Note:\n - batch_size: N\n - num_keypoints: K\n - num_outputs: O\n - heatmaps height: H\n - heatmaps weight: W\n\n Args:\n output (torch.Tensor[N,K,H,W]):\n Output heatmaps.\n target (torch.Tensor[N,K,H,W]):\n Target heatmaps.\n target_weight (torch.Tensor[N,K,1]):\n Weights across different joint types.\n \"\"\"\n\n losses = dict()\n\n assert isinstance(output, list)\n assert target.dim() == 4 and target_weight.dim() == 3\n\n if isinstance(self.loss, nn.Sequential):\n assert len(self.loss) == len(output)\n for i in range(len(output)):\n target_i = target\n target_weight_i = target_weight\n if isinstance(self.loss, nn.Sequential):\n loss_func = self.loss[i]\n else:\n loss_func = self.loss\n loss_i = loss_func(output[i], target_i, target_weight_i)\n if 'heatmap_loss' not in losses:\n losses['heatmap_loss'] = loss_i\n else:\n losses['heatmap_loss'] += loss_i\n\n return losses\n\n def get_accuracy(self, output, target, target_weight):\n \"\"\"Calculate accuracy for top-down keypoint loss.\n\n Note:\n - batch_size: N\n - num_keypoints: K\n - heatmaps height: H\n - heatmaps weight: W\n\n Args:\n output (torch.Tensor[N,K,H,W]): Output heatmaps.\n target (torch.Tensor[N,K,H,W]): Target heatmaps.\n target_weight (torch.Tensor[N,K,1]):\n Weights across different joint types.\n \"\"\"\n\n accuracy = dict()\n\n if self.target_type == 'GaussianHeatmap':\n _, avg_acc, _ = pose_pck_accuracy(\n output[-1].detach().cpu().numpy(),\n target.detach().cpu().numpy(),\n target_weight.detach().cpu().numpy().squeeze(-1) > 0)\n accuracy['acc_pose'] = float(avg_acc)\n\n return accuracy\n\n def forward(self, x):\n \"\"\"Forward function.\n\n Returns:\n out (list[Tensor]): a list of heatmaps from multiple stages.\n \"\"\"\n out = []\n assert isinstance(x, list)\n for i in range(self.num_stages):\n y = self.multi_deconv_layers[i](x[i])\n y = self.multi_final_layers[i](y)\n out.append(y)\n return out\n\n def inference_model(self, x, flip_pairs=None):\n \"\"\"Inference function.\n\n Returns:\n output_heatmap (np.ndarray): Output heatmaps.\n\n Args:\n x (List[torch.Tensor[NxKxHxW]]): Input features.\n flip_pairs (None | list[tuple()):\n Pairs of keypoints which are mirrored.\n \"\"\"\n output = self.forward(x)\n assert isinstance(output, list)\n output = output[-1]\n\n if flip_pairs is not None:\n # perform flip\n output_heatmap = flip_back(\n output.detach().cpu().numpy(),\n flip_pairs,\n target_type=self.target_type)\n # feature is not aligned, shift flipped heatmap for higher accuracy\n if self.test_cfg.get('shift_heatmap', False):\n output_heatmap[:, :, :, 1:] = output_heatmap[:, :, :, :-1]\n else:\n output_heatmap = output.detach().cpu().numpy()\n\n return output_heatmap\n\n def _make_deconv_layer(self, num_layers, num_filters, num_kernels):\n \"\"\"Make deconv layers.\"\"\"\n if num_layers != len(num_filters):\n error_msg = f'num_layers({num_layers}) ' \\\n f'!= length of num_filters({len(num_filters)})'\n raise ValueError(error_msg)\n if num_layers != len(num_kernels):\n error_msg = f'num_layers({num_layers}) ' \\\n f'!= length of num_kernels({len(num_kernels)})'\n raise ValueError(error_msg)\n\n layers = []\n for i in range(num_layers):\n kernel, padding, output_padding = \\\n self._get_deconv_cfg(num_kernels[i])\n\n planes = num_filters[i]\n layers.append(\n build_upsample_layer(\n dict(type='deconv'),\n in_channels=self.in_channels,\n out_channels=planes,\n kernel_size=kernel,\n stride=2,\n padding=padding,\n output_padding=output_padding,\n bias=False))\n layers.append(nn.BatchNorm2d(planes))\n layers.append(nn.ReLU(inplace=True))\n self.in_channels = planes\n\n return nn.Sequential(*layers)\n\n def init_weights(self):\n \"\"\"Initialize model weights.\"\"\"\n for _, m in self.multi_deconv_layers.named_modules():\n if isinstance(m, nn.ConvTranspose2d):\n normal_init(m, std=0.001)\n elif isinstance(m, nn.BatchNorm2d):\n constant_init(m, 1)\n for m in self.multi_final_layers.modules():\n if isinstance(m, nn.Conv2d):\n normal_init(m, std=0.001, bias=0)\n\n\nclass PredictHeatmap(nn.Module):\n \"\"\"Predict the heat map for an input feature.\n\n Args:\n unit_channels (int): Number of input channels.\n out_channels (int): Number of output channels.\n out_shape (tuple): Shape of the output heatmap.\n use_prm (bool): Whether to use pose refine machine. Default: False.\n norm_cfg (dict): dictionary to construct and config norm layer.\n Default: dict(type='BN')\n \"\"\"\n\n def __init__(self,\n unit_channels,\n out_channels,\n out_shape,\n use_prm=False,\n norm_cfg=dict(type='BN')):\n # Protect mutable default arguments\n norm_cfg = cp.deepcopy(norm_cfg)\n super().__init__()\n self.unit_channels = unit_channels\n self.out_channels = out_channels\n self.out_shape = out_shape\n self.use_prm = use_prm\n if use_prm:\n self.prm = PRM(out_channels, norm_cfg=norm_cfg)\n self.conv_layers = nn.Sequential(\n ConvModule(\n unit_channels,\n unit_channels,\n kernel_size=1,\n stride=1,\n padding=0,\n norm_cfg=norm_cfg,\n inplace=False),\n ConvModule(\n unit_channels,\n out_channels,\n kernel_size=3,\n stride=1,\n padding=1,\n norm_cfg=norm_cfg,\n act_cfg=None,\n inplace=False))\n\n def forward(self, feature):\n feature = self.conv_layers(feature)\n output = nn.functional.interpolate(\n feature, size=self.out_shape, mode='bilinear', align_corners=True)\n if self.use_prm:\n output = self.prm(output)\n return output\n\n\nclass PRM(nn.Module):\n \"\"\"Pose Refine Machine.\n\n Please refer to \"Learning Delicate Local Representations\n for Multi-Person Pose Estimation\" (ECCV 2020).\n\n Args:\n out_channels (int): Channel number of the output. Equals to\n the number of key points.\n norm_cfg (dict): dictionary to construct and config norm layer.\n Default: dict(type='BN')\n \"\"\"\n\n def __init__(self, out_channels, norm_cfg=dict(type='BN')):\n # Protect mutable default arguments\n norm_cfg = cp.deepcopy(norm_cfg)\n super().__init__()\n self.out_channels = out_channels\n self.global_pooling = nn.AdaptiveAvgPool2d((1, 1))\n self.middle_path = nn.Sequential(\n Linear(self.out_channels, self.out_channels),\n build_norm_layer(dict(type='BN1d'), out_channels)[1],\n build_activation_layer(dict(type='ReLU')),\n Linear(self.out_channels, self.out_channels),\n build_norm_layer(dict(type='BN1d'), out_channels)[1],\n build_activation_layer(dict(type='ReLU')),\n build_activation_layer(dict(type='Sigmoid')))\n\n self.bottom_path = nn.Sequential(\n ConvModule(\n self.out_channels,\n self.out_channels,\n kernel_size=1,\n stride=1,\n padding=0,\n norm_cfg=norm_cfg,\n inplace=False),\n DepthwiseSeparableConvModule(\n self.out_channels,\n 1,\n kernel_size=9,\n stride=1,\n padding=4,\n norm_cfg=norm_cfg,\n inplace=False), build_activation_layer(dict(type='Sigmoid')))\n self.conv_bn_relu_prm_1 = ConvModule(\n self.out_channels,\n self.out_channels,\n kernel_size=3,\n stride=1,\n padding=1,\n norm_cfg=norm_cfg,\n inplace=False)\n\n def forward(self, x):\n out = self.conv_bn_relu_prm_1(x)\n out_1 = out\n\n out_2 = self.global_pooling(out_1)\n out_2 = out_2.view(out_2.size(0), -1)\n out_2 = self.middle_path(out_2)\n out_2 = out_2.unsqueeze(2)\n out_2 = out_2.unsqueeze(3)\n\n out_3 = self.bottom_path(out_1)\n out = out_1 * (1 + out_2 * out_3)\n\n return out\n\n\n@HEADS.register_module()\nclass TopdownHeatmapMSMUHead(TopdownHeatmapBaseHead):\n \"\"\"Heads for multi-stage multi-unit heads used in Multi-Stage Pose\n estimation Network (MSPN), and Residual Steps Networks (RSN).\n\n Args:\n unit_channels (int): Number of input channels.\n out_channels (int): Number of output channels.\n out_shape (tuple): Shape of the output heatmap.\n num_stages (int): Number of stages.\n num_units (int): Number of units in each stage.\n use_prm (bool): Whether to use pose refine machine (PRM).\n Default: False.\n norm_cfg (dict): dictionary to construct and config norm layer.\n Default: dict(type='BN')\n loss_keypoint (dict): Config for keypoint loss. Default: None.\n \"\"\"\n\n def __init__(self,\n out_shape,\n unit_channels=256,\n out_channels=17,\n num_stages=4,\n num_units=4,\n use_prm=False,\n norm_cfg=dict(type='BN'),\n loss_keypoint=None,\n train_cfg=None,\n test_cfg=None):\n # Protect mutable default arguments\n norm_cfg = cp.deepcopy(norm_cfg)\n super().__init__()\n\n self.train_cfg = {} if train_cfg is None else train_cfg\n self.test_cfg = {} if test_cfg is None else test_cfg\n self.target_type = self.test_cfg.get('target_type', 'GaussianHeatmap')\n\n self.out_shape = out_shape\n self.unit_channels = unit_channels\n self.out_channels = out_channels\n self.num_stages = num_stages\n self.num_units = num_units\n\n self.loss = build_loss(loss_keypoint)\n\n self.predict_layers = nn.ModuleList([])\n for i in range(self.num_stages):\n for j in range(self.num_units):\n self.predict_layers.append(\n PredictHeatmap(\n unit_channels,\n out_channels,\n out_shape,\n use_prm,\n norm_cfg=norm_cfg))\n\n def get_loss(self, output, target, target_weight):\n \"\"\"Calculate top-down keypoint loss.\n\n Note:\n - batch_size: N\n - num_keypoints: K\n - num_outputs: O\n - heatmaps height: H\n - heatmaps weight: W\n\n Args:\n output (torch.Tensor[N,O,K,H,W]): Output heatmaps.\n target (torch.Tensor[N,O,K,H,W]): Target heatmaps.\n target_weight (torch.Tensor[N,O,K,1]):\n Weights across different joint types.\n \"\"\"\n\n losses = dict()\n\n assert isinstance(output, list)\n assert target.dim() == 5 and target_weight.dim() == 4\n assert target.size(1) == len(output)\n\n if isinstance(self.loss, nn.Sequential):\n assert len(self.loss) == len(output)\n for i in range(len(output)):\n target_i = target[:, i, :, :, :]\n target_weight_i = target_weight[:, i, :, :]\n\n if isinstance(self.loss, nn.Sequential):\n loss_func = self.loss[i]\n else:\n loss_func = self.loss\n\n loss_i = loss_func(output[i], target_i, target_weight_i)\n if 'heatmap_loss' not in losses:\n losses['heatmap_loss'] = loss_i\n else:\n losses['heatmap_loss'] += loss_i\n\n return losses\n\n def get_accuracy(self, output, target, target_weight):\n \"\"\"Calculate accuracy for top-down keypoint loss.\n\n Note:\n - batch_size: N\n - num_keypoints: K\n - heatmaps height: H\n - heatmaps weight: W\n\n Args:\n output (torch.Tensor[N,K,H,W]): Output heatmaps.\n target (torch.Tensor[N,K,H,W]): Target heatmaps.\n target_weight (torch.Tensor[N,K,1]):\n Weights across different joint types.\n \"\"\"\n\n accuracy = dict()\n\n if self.target_type == 'GaussianHeatmap':\n assert isinstance(output, list)\n assert target.dim() == 5 and target_weight.dim() == 4\n _, avg_acc, _ = pose_pck_accuracy(\n output[-1].detach().cpu().numpy(),\n target[:, -1, ...].detach().cpu().numpy(),\n target_weight[:, -1,\n ...].detach().cpu().numpy().squeeze(-1) > 0)\n accuracy['acc_pose'] = float(avg_acc)\n\n return accuracy\n\n def forward(self, x):\n \"\"\"Forward function.\n\n Returns:\n out (list[Tensor]): a list of heatmaps from multiple stages\n and units.\n \"\"\"\n out = []\n assert isinstance(x, list)\n assert len(x) == self.num_stages\n assert isinstance(x[0], list)\n assert len(x[0]) == self.num_units\n assert x[0][0].shape[1] == self.unit_channels\n for i in range(self.num_stages):\n for j in range(self.num_units):\n y = self.predict_layers[i * self.num_units + j](x[i][j])\n out.append(y)\n\n return out\n\n def inference_model(self, x, flip_pairs=None):\n \"\"\"Inference function.\n\n Returns:\n output_heatmap (np.ndarray): Output heatmaps.\n\n Args:\n x (list[torch.Tensor[N,K,H,W]]): Input features.\n flip_pairs (None | list[tuple]):\n Pairs of keypoints which are mirrored.\n \"\"\"\n output = self.forward(x)\n assert isinstance(output, list)\n output = output[-1]\n if flip_pairs is not None:\n output_heatmap = flip_back(\n output.detach().cpu().numpy(),\n flip_pairs,\n target_type=self.target_type)\n # feature is not aligned, shift flipped heatmap for higher accuracy\n if self.test_cfg.get('shift_heatmap', False):\n output_heatmap[:, :, :, 1:] = output_heatmap[:, :, :, :-1]\n else:\n output_heatmap = output.detach().cpu().numpy()\n return output_heatmap\n\n def init_weights(self):\n \"\"\"Initialize model weights.\"\"\"\n for m in self.predict_layers.modules():\n if isinstance(m, nn.Conv2d):\n kaiming_init(m)\n elif isinstance(m, nn.BatchNorm2d):\n constant_init(m, 1)\n elif isinstance(m, nn.Linear):\n normal_init(m, std=0.01)\n" }, { "path": "eval/GVHMR/hmr4d/utils/preproc/vitpose_pytorch/src/vitpose_infer/builder/heads/topdown_heatmap_simple_head.py", "content": "# Copyright (c) OpenMMLab. All rights reserved.\nimport torch\nimport torch.nn as nn\n# from mmcv.cnn import (build_conv_layer, build_norm_layer, build_upsample_layer,\n# constant_init, normal_init)\n\n# from mmpose.core.evaluation import pose_pck_accuracy\n# from mmpose.core.post_processing import flip_back\n# from mmpose.models.builder import build_loss\n# from mmpose.models.utils.ops import resize\n# from ..builder import HEADS\nimport torch.nn.functional as F\nfrom .topdown_heatmap_base_head import TopdownHeatmapBaseHead\n\ndef build_conv_layer(cfg, *args, **kwargs) -> nn.Module:\n \"\"\"LICENSE\"\"\"\n \n if cfg is None:\n cfg_ = dict(type='Conv2d')\n else:\n if not isinstance(cfg, dict):\n raise TypeError('cfg must be a dict')\n if 'type' not in cfg:\n raise KeyError('the cfg dict must contain the key \"type\"')\n cfg_ = cfg.copy()\n\n layer_type = cfg_.pop('type')\n if layer_type !='Conv2d':\n raise KeyError(f'Unrecognized layer type {layer_type}')\n else:\n conv_layer = nn.Conv2d\n\n layer = conv_layer(*args, **kwargs, **cfg_)\n\n return layer\n\ndef build_upsample_layer(cfg, *args, **kwargs) -> nn.Module:\n\n if not isinstance(cfg, dict):\n raise TypeError(f'cfg must be a dict, but got {type(cfg)}')\n if 'type' not in cfg:\n raise KeyError(\n f'the cfg dict must contain the key \"type\", but got {cfg}')\n cfg_ = cfg.copy()\n\n layer_type = cfg_.pop('type')\n if layer_type !='deconv':\n raise KeyError(f'Unrecognized upsample type {layer_type}')\n else:\n upsample = nn.ConvTranspose2d\n\n if upsample is nn.Upsample:\n cfg_['mode'] = layer_type\n layer = upsample(*args, **kwargs, **cfg_)\n return layer\n\n# @HEADS.register_module()\nclass TopdownHeatmapSimpleHead(TopdownHeatmapBaseHead):\n \"\"\"Top-down heatmap simple head. paper ref: Bin Xiao et al. ``Simple\n Baselines for Human Pose Estimation and Tracking``.\n\n TopdownHeatmapSimpleHead is consisted of (>=0) number of deconv layers\n and a simple conv2d layer.\n\n Args:\n in_channels (int): Number of input channels\n out_channels (int): Number of output channels\n num_deconv_layers (int): Number of deconv layers.\n num_deconv_layers should >= 0. Note that 0 means\n no deconv layers.\n num_deconv_filters (list|tuple): Number of filters.\n If num_deconv_layers > 0, the length of\n num_deconv_kernels (list|tuple): Kernel sizes.\n in_index (int|Sequence[int]): Input feature index. Default: 0\n input_transform (str|None): Transformation type of input features.\n Options: 'resize_concat', 'multiple_select', None.\n Default: None.\n\n - 'resize_concat': Multiple feature maps will be resized to the\n same size as the first one and then concat together.\n Usually used in FCN head of HRNet.\n - 'multiple_select': Multiple feature maps will be bundle into\n a list and passed into decode head.\n - None: Only one select feature map is allowed.\n align_corners (bool): align_corners argument of F.interpolate.\n Default: False.\n loss_keypoint (dict): Config for keypoint loss. Default: None.\n \"\"\"\n\n def __init__(self,\n in_channels,\n out_channels,\n num_deconv_layers=3,\n num_deconv_filters=(256, 256, 256),\n num_deconv_kernels=(4, 4, 4),\n extra=None,\n in_index=0,\n input_transform=None,\n align_corners=False,\n loss_keypoint=None,\n train_cfg=None,\n test_cfg=None,\n upsample=0,):\n super().__init__()\n\n self.in_channels = in_channels\n self.loss = None\n self.upsample = upsample\n\n self.train_cfg = {} if train_cfg is None else train_cfg\n self.test_cfg = {} if test_cfg is None else test_cfg\n self.target_type = self.test_cfg.get('target_type', 'GaussianHeatmap')\n\n self._init_inputs(in_channels, in_index, input_transform)\n self.in_index = in_index\n self.align_corners = align_corners\n\n if extra is not None and not isinstance(extra, dict):\n raise TypeError('extra should be dict or None.')\n\n if num_deconv_layers > 0:\n self.deconv_layers = self._make_deconv_layer(\n num_deconv_layers,\n num_deconv_filters,\n num_deconv_kernels,\n )\n elif num_deconv_layers == 0:\n self.deconv_layers = nn.Identity()\n else:\n raise ValueError(\n f'num_deconv_layers ({num_deconv_layers}) should >= 0.')\n\n identity_final_layer = False\n if extra is not None and 'final_conv_kernel' in extra:\n assert extra['final_conv_kernel'] in [0, 1, 3]\n if extra['final_conv_kernel'] == 3:\n padding = 1\n elif extra['final_conv_kernel'] == 1:\n padding = 0\n else:\n # 0 for Identity mapping.\n identity_final_layer = True\n kernel_size = extra['final_conv_kernel']\n else:\n kernel_size = 1\n padding = 0\n\n if identity_final_layer:\n self.final_layer = nn.Identity()\n else:\n conv_channels = num_deconv_filters[\n -1] if num_deconv_layers > 0 else self.in_channels\n\n layers = []\n if extra is not None:\n num_conv_layers = extra.get('num_conv_layers', 0)\n num_conv_kernels = extra.get('num_conv_kernels',\n [1] * num_conv_layers)\n\n for i in range(num_conv_layers):\n layers.append(\n build_conv_layer(\n dict(type='Conv2d'),\n in_channels=conv_channels,\n out_channels=conv_channels,\n kernel_size=num_conv_kernels[i],\n stride=1,\n padding=(num_conv_kernels[i] - 1) // 2))\n layers.append(\n nn.BatchNorm2d(conv_channels)\n)\n layers.append(nn.ReLU(inplace=True))\n\n layers.append(\n build_conv_layer(\n cfg=dict(type='Conv2d'),\n in_channels=conv_channels,\n out_channels=out_channels,\n kernel_size=kernel_size,\n stride=1,\n padding=padding))\n\n if len(layers) > 1:\n self.final_layer = nn.Sequential(*layers)\n else:\n self.final_layer = layers[0]\n\n def get_loss(self, output, target, target_weight):\n \"\"\"Calculate top-down keypoint loss.\n\n Note:\n - batch_size: N\n - num_keypoints: K\n - heatmaps height: H\n - heatmaps weight: W\n\n Args:\n output (torch.Tensor[N,K,H,W]): Output heatmaps.\n target (torch.Tensor[N,K,H,W]): Target heatmaps.\n target_weight (torch.Tensor[N,K,1]):\n Weights across different joint types.\n \"\"\"\n\n losses = dict()\n\n assert not isinstance(self.loss, nn.Sequential)\n assert target.dim() == 4 and target_weight.dim() == 3\n losses['heatmap_loss'] = self.loss(output, target, target_weight)\n\n return losses\n\n def get_accuracy(self, output, target, target_weight):\n \"\"\"Calculate accuracy for top-down keypoint loss.\n\n Note:\n - batch_size: N\n - num_keypoints: K\n - heatmaps height: H\n - heatmaps weight: W\n\n Args:\n output (torch.Tensor[N,K,H,W]): Output heatmaps.\n target (torch.Tensor[N,K,H,W]): Target heatmaps.\n target_weight (torch.Tensor[N,K,1]):\n Weights across different joint types.\n \"\"\"\n\n accuracy = dict()\n\n if self.target_type == 'GaussianHeatmap':\n _, avg_acc, _ = pose_pck_accuracy(\n output.detach().cpu().numpy(),\n target.detach().cpu().numpy(),\n target_weight.detach().cpu().numpy().squeeze(-1) > 0)\n accuracy['acc_pose'] = float(avg_acc)\n\n return accuracy\n\n def forward(self, x):\n \"\"\"Forward function.\"\"\"\n x = self._transform_inputs(x)\n x = self.deconv_layers(x)\n x = self.final_layer(x)\n return x\n\n def inference_model(self, x, flip_pairs=None):\n \"\"\"Inference function.\n\n Returns:\n output_heatmap (np.ndarray): Output heatmaps.\n\n Args:\n x (torch.Tensor[N,K,H,W]): Input features.\n flip_pairs (None | list[tuple]):\n Pairs of keypoints which are mirrored.\n \"\"\"\n output = self.forward(x)\n\n if flip_pairs is not None:\n output_heatmap = flip_back(\n output.detach().cpu().numpy(),\n flip_pairs,\n target_type=self.target_type)\n # feature is not aligned, shift flipped heatmap for higher accuracy\n if self.test_cfg.get('shift_heatmap', False):\n output_heatmap[:, :, :, 1:] = output_heatmap[:, :, :, :-1]\n else:\n output_heatmap = output.detach().cpu().numpy()\n return output_heatmap\n\n def _init_inputs(self, in_channels, in_index, input_transform):\n \"\"\"Check and initialize input transforms.\n\n The in_channels, in_index and input_transform must match.\n Specifically, when input_transform is None, only single feature map\n will be selected. So in_channels and in_index must be of type int.\n When input_transform is not None, in_channels and in_index must be\n list or tuple, with the same length.\n\n Args:\n in_channels (int|Sequence[int]): Input channels.\n in_index (int|Sequence[int]): Input feature index.\n input_transform (str|None): Transformation type of input features.\n Options: 'resize_concat', 'multiple_select', None.\n\n - 'resize_concat': Multiple feature maps will be resize to the\n same size as first one and than concat together.\n Usually used in FCN head of HRNet.\n - 'multiple_select': Multiple feature maps will be bundle into\n a list and passed into decode head.\n - None: Only one select feature map is allowed.\n \"\"\"\n\n if input_transform is not None:\n assert input_transform in ['resize_concat', 'multiple_select']\n self.input_transform = input_transform\n self.in_index = in_index\n if input_transform is not None:\n assert isinstance(in_channels, (list, tuple))\n assert isinstance(in_index, (list, tuple))\n assert len(in_channels) == len(in_index)\n if input_transform == 'resize_concat':\n self.in_channels = sum(in_channels)\n else:\n self.in_channels = in_channels\n else:\n assert isinstance(in_channels, int)\n assert isinstance(in_index, int)\n self.in_channels = in_channels\n\n def _transform_inputs(self, inputs):\n \"\"\"Transform inputs for decoder.\n\n Args:\n inputs (list[Tensor] | Tensor): multi-level img features.\n\n Returns:\n Tensor: The transformed inputs\n \"\"\"\n if not isinstance(inputs, list):\n if not isinstance(inputs, list):\n if self.upsample > 0:\n inputs = resize(\n input=F.relu(inputs),\n scale_factor=self.upsample,\n mode='bilinear',\n align_corners=self.align_corners\n )\n return inputs\n\n if self.input_transform == 'resize_concat':\n inputs = [inputs[i] for i in self.in_index]\n upsampled_inputs = [\n resize(\n input=x,\n size=inputs[0].shape[2:],\n mode='bilinear',\n align_corners=self.align_corners) for x in inputs\n ]\n inputs = torch.cat(upsampled_inputs, dim=1)\n elif self.input_transform == 'multiple_select':\n inputs = [inputs[i] for i in self.in_index]\n else:\n inputs = inputs[self.in_index]\n\n return inputs\n\n def _make_deconv_layer(self, num_layers, num_filters, num_kernels):\n \"\"\"Make deconv layers.\"\"\"\n if num_layers != len(num_filters):\n error_msg = f'num_layers({num_layers}) ' \\\n f'!= length of num_filters({len(num_filters)})'\n raise ValueError(error_msg)\n if num_layers != len(num_kernels):\n error_msg = f'num_layers({num_layers}) ' \\\n f'!= length of num_kernels({len(num_kernels)})'\n raise ValueError(error_msg)\n\n layers = []\n for i in range(num_layers):\n kernel, padding, output_padding = \\\n self._get_deconv_cfg(num_kernels[i])\n\n planes = num_filters[i]\n layers.append(\n build_upsample_layer(\n dict(type='deconv'),\n in_channels=self.in_channels,\n out_channels=planes,\n kernel_size=kernel,\n stride=2,\n padding=padding,\n output_padding=output_padding,\n bias=False))\n layers.append(nn.BatchNorm2d(planes))\n layers.append(nn.ReLU(inplace=True))\n self.in_channels = planes\n\n return nn.Sequential(*layers)\n\n def init_weights(self):\n \"\"\"Initialize model weights.\"\"\"\n for _, m in self.deconv_layers.named_modules():\n if isinstance(m, nn.ConvTranspose2d):\n normal_init(m, std=0.001)\n elif isinstance(m, nn.BatchNorm2d):\n constant_init(m, 1)\n for m in self.final_layer.modules():\n if isinstance(m, nn.Conv2d):\n normal_init(m, std=0.001, bias=0)\n elif isinstance(m, nn.BatchNorm2d):\n constant_init(m, 1)\n" }, { "path": "eval/GVHMR/hmr4d/utils/preproc/vitpose_pytorch/src/vitpose_infer/builder/heads/vipnas_heatmap_simple_head.py", "content": "# Copyright (c) OpenMMLab. All rights reserved.\nimport torch\nimport torch.nn as nn\n# from mmcv.cnn import (build_conv_layer, build_norm_layer, build_upsample_layer,\n# constant_init, normal_init)\n\n# from mmpose.core.evaluation import pose_pck_accuracy\n# from mmpose.core.post_processing import flip_back\n# from mmpose.models.builder import build_loss\n# from mmpose.models.utils.ops import resize\n# from ..builder import HEADS\n# from .topdown_heatmap_base_head import TopdownHeatmapBaseHead\n\n\n# @HEADS.register_module()\nclass ViPNASHeatmapSimpleHead(TopdownHeatmapBaseHead):\n \"\"\"ViPNAS heatmap simple head.\n\n ViPNAS: Efficient Video Pose Estimation via Neural Architecture Search.\n More details can be found in the `paper\n `__ .\n\n TopdownHeatmapSimpleHead is consisted of (>=0) number of deconv layers\n and a simple conv2d layer.\n\n Args:\n in_channels (int): Number of input channels\n out_channels (int): Number of output channels\n num_deconv_layers (int): Number of deconv layers.\n num_deconv_layers should >= 0. Note that 0 means\n no deconv layers.\n num_deconv_filters (list|tuple): Number of filters.\n If num_deconv_layers > 0, the length of\n num_deconv_kernels (list|tuple): Kernel sizes.\n num_deconv_groups (list|tuple): Group number.\n in_index (int|Sequence[int]): Input feature index. Default: -1\n input_transform (str|None): Transformation type of input features.\n Options: 'resize_concat', 'multiple_select', None.\n Default: None.\n\n - 'resize_concat': Multiple feature maps will be resize to the\n same size as first one and than concat together.\n Usually used in FCN head of HRNet.\n - 'multiple_select': Multiple feature maps will be bundle into\n a list and passed into decode head.\n - None: Only one select feature map is allowed.\n align_corners (bool): align_corners argument of F.interpolate.\n Default: False.\n loss_keypoint (dict): Config for keypoint loss. Default: None.\n \"\"\"\n\n def __init__(self,\n in_channels,\n out_channels,\n num_deconv_layers=3,\n num_deconv_filters=(144, 144, 144),\n num_deconv_kernels=(4, 4, 4),\n num_deconv_groups=(16, 16, 16),\n extra=None,\n in_index=0,\n input_transform=None,\n align_corners=False,\n loss_keypoint=None,\n train_cfg=None,\n test_cfg=None):\n super().__init__()\n\n self.in_channels = in_channels\n self.loss = build_loss(loss_keypoint)\n\n self.train_cfg = {} if train_cfg is None else train_cfg\n self.test_cfg = {} if test_cfg is None else test_cfg\n self.target_type = self.test_cfg.get('target_type', 'GaussianHeatmap')\n\n self._init_inputs(in_channels, in_index, input_transform)\n self.in_index = in_index\n self.align_corners = align_corners\n\n if extra is not None and not isinstance(extra, dict):\n raise TypeError('extra should be dict or None.')\n\n if num_deconv_layers > 0:\n self.deconv_layers = self._make_deconv_layer(\n num_deconv_layers, num_deconv_filters, num_deconv_kernels,\n num_deconv_groups)\n elif num_deconv_layers == 0:\n self.deconv_layers = nn.Identity()\n else:\n raise ValueError(\n f'num_deconv_layers ({num_deconv_layers}) should >= 0.')\n\n identity_final_layer = False\n if extra is not None and 'final_conv_kernel' in extra:\n assert extra['final_conv_kernel'] in [0, 1, 3]\n if extra['final_conv_kernel'] == 3:\n padding = 1\n elif extra['final_conv_kernel'] == 1:\n padding = 0\n else:\n # 0 for Identity mapping.\n identity_final_layer = True\n kernel_size = extra['final_conv_kernel']\n else:\n kernel_size = 1\n padding = 0\n\n if identity_final_layer:\n self.final_layer = nn.Identity()\n else:\n conv_channels = num_deconv_filters[\n -1] if num_deconv_layers > 0 else self.in_channels\n\n layers = []\n if extra is not None:\n num_conv_layers = extra.get('num_conv_layers', 0)\n num_conv_kernels = extra.get('num_conv_kernels',\n [1] * num_conv_layers)\n\n for i in range(num_conv_layers):\n layers.append(\n build_conv_layer(\n dict(type='Conv2d'),\n in_channels=conv_channels,\n out_channels=conv_channels,\n kernel_size=num_conv_kernels[i],\n stride=1,\n padding=(num_conv_kernels[i] - 1) // 2))\n layers.append(\n build_norm_layer(dict(type='BN'), conv_channels)[1])\n layers.append(nn.ReLU(inplace=True))\n\n layers.append(\n build_conv_layer(\n cfg=dict(type='Conv2d'),\n in_channels=conv_channels,\n out_channels=out_channels,\n kernel_size=kernel_size,\n stride=1,\n padding=padding))\n\n if len(layers) > 1:\n self.final_layer = nn.Sequential(*layers)\n else:\n self.final_layer = layers[0]\n\n def get_loss(self, output, target, target_weight):\n \"\"\"Calculate top-down keypoint loss.\n\n Note:\n - batch_size: N\n - num_keypoints: K\n - heatmaps height: H\n - heatmaps weight: W\n\n Args:\n output (torch.Tensor[N,K,H,W]): Output heatmaps.\n target (torch.Tensor[N,K,H,W]): Target heatmaps.\n target_weight (torch.Tensor[N,K,1]):\n Weights across different joint types.\n \"\"\"\n\n losses = dict()\n\n assert not isinstance(self.loss, nn.Sequential)\n assert target.dim() == 4 and target_weight.dim() == 3\n losses['heatmap_loss'] = self.loss(output, target, target_weight)\n\n return losses\n\n def get_accuracy(self, output, target, target_weight):\n \"\"\"Calculate accuracy for top-down keypoint loss.\n\n Note:\n - batch_size: N\n - num_keypoints: K\n - heatmaps height: H\n - heatmaps weight: W\n\n Args:\n output (torch.Tensor[N,K,H,W]): Output heatmaps.\n target (torch.Tensor[N,K,H,W]): Target heatmaps.\n target_weight (torch.Tensor[N,K,1]):\n Weights across different joint types.\n \"\"\"\n\n accuracy = dict()\n\n if self.target_type.lower() == 'GaussianHeatmap'.lower():\n _, avg_acc, _ = pose_pck_accuracy(\n output.detach().cpu().numpy(),\n target.detach().cpu().numpy(),\n target_weight.detach().cpu().numpy().squeeze(-1) > 0)\n accuracy['acc_pose'] = float(avg_acc)\n\n return accuracy\n\n def forward(self, x):\n \"\"\"Forward function.\"\"\"\n x = self._transform_inputs(x)\n x = self.deconv_layers(x)\n x = self.final_layer(x)\n return x\n\n def inference_model(self, x, flip_pairs=None):\n \"\"\"Inference function.\n\n Returns:\n output_heatmap (np.ndarray): Output heatmaps.\n\n Args:\n x (torch.Tensor[N,K,H,W]): Input features.\n flip_pairs (None | list[tuple]):\n Pairs of keypoints which are mirrored.\n \"\"\"\n output = self.forward(x)\n\n if flip_pairs is not None:\n output_heatmap = flip_back(\n output.detach().cpu().numpy(),\n flip_pairs,\n target_type=self.target_type)\n # feature is not aligned, shift flipped heatmap for higher accuracy\n if self.test_cfg.get('shift_heatmap', False):\n output_heatmap[:, :, :, 1:] = output_heatmap[:, :, :, :-1]\n else:\n output_heatmap = output.detach().cpu().numpy()\n return output_heatmap\n\n def _init_inputs(self, in_channels, in_index, input_transform):\n \"\"\"Check and initialize input transforms.\n\n The in_channels, in_index and input_transform must match.\n Specifically, when input_transform is None, only single feature map\n will be selected. So in_channels and in_index must be of type int.\n When input_transform is not None, in_channels and in_index must be\n list or tuple, with the same length.\n\n Args:\n in_channels (int|Sequence[int]): Input channels.\n in_index (int|Sequence[int]): Input feature index.\n input_transform (str|None): Transformation type of input features.\n Options: 'resize_concat', 'multiple_select', None.\n\n - 'resize_concat': Multiple feature maps will be resize to the\n same size as first one and than concat together.\n Usually used in FCN head of HRNet.\n - 'multiple_select': Multiple feature maps will be bundle into\n a list and passed into decode head.\n - None: Only one select feature map is allowed.\n \"\"\"\n\n if input_transform is not None:\n assert input_transform in ['resize_concat', 'multiple_select']\n self.input_transform = input_transform\n self.in_index = in_index\n if input_transform is not None:\n assert isinstance(in_channels, (list, tuple))\n assert isinstance(in_index, (list, tuple))\n assert len(in_channels) == len(in_index)\n if input_transform == 'resize_concat':\n self.in_channels = sum(in_channels)\n else:\n self.in_channels = in_channels\n else:\n assert isinstance(in_channels, int)\n assert isinstance(in_index, int)\n self.in_channels = in_channels\n\n def _transform_inputs(self, inputs):\n \"\"\"Transform inputs for decoder.\n\n Args:\n inputs (list[Tensor] | Tensor): multi-level img features.\n\n Returns:\n Tensor: The transformed inputs\n \"\"\"\n if not isinstance(inputs, list):\n return inputs\n\n if self.input_transform == 'resize_concat':\n inputs = [inputs[i] for i in self.in_index]\n upsampled_inputs = [\n resize(\n input=x,\n size=inputs[0].shape[2:],\n mode='bilinear',\n align_corners=self.align_corners) for x in inputs\n ]\n inputs = torch.cat(upsampled_inputs, dim=1)\n elif self.input_transform == 'multiple_select':\n inputs = [inputs[i] for i in self.in_index]\n else:\n inputs = inputs[self.in_index]\n\n return inputs\n\n def _make_deconv_layer(self, num_layers, num_filters, num_kernels,\n num_groups):\n \"\"\"Make deconv layers.\"\"\"\n if num_layers != len(num_filters):\n error_msg = f'num_layers({num_layers}) ' \\\n f'!= length of num_filters({len(num_filters)})'\n raise ValueError(error_msg)\n if num_layers != len(num_kernels):\n error_msg = f'num_layers({num_layers}) ' \\\n f'!= length of num_kernels({len(num_kernels)})'\n raise ValueError(error_msg)\n if num_layers != len(num_groups):\n error_msg = f'num_layers({num_layers}) ' \\\n f'!= length of num_groups({len(num_groups)})'\n raise ValueError(error_msg)\n\n layers = []\n for i in range(num_layers):\n kernel, padding, output_padding = \\\n self._get_deconv_cfg(num_kernels[i])\n\n planes = num_filters[i]\n groups = num_groups[i]\n layers.append(\n build_upsample_layer(\n dict(type='deconv'),\n in_channels=self.in_channels,\n out_channels=planes,\n kernel_size=kernel,\n groups=groups,\n stride=2,\n padding=padding,\n output_padding=output_padding,\n bias=False))\n layers.append(nn.BatchNorm2d(planes))\n layers.append(nn.ReLU(inplace=True))\n self.in_channels = planes\n\n return nn.Sequential(*layers)\n\n def init_weights(self):\n \"\"\"Initialize model weights.\"\"\"\n for _, m in self.deconv_layers.named_modules():\n if isinstance(m, nn.ConvTranspose2d):\n normal_init(m, std=0.001)\n elif isinstance(m, nn.BatchNorm2d):\n constant_init(m, 1)\n for m in self.final_layer.modules():\n if isinstance(m, nn.Conv2d):\n normal_init(m, std=0.001, bias=0)\n elif isinstance(m, nn.BatchNorm2d):\n constant_init(m, 1)\n" }, { "path": "eval/GVHMR/hmr4d/utils/preproc/vitpose_pytorch/src/vitpose_infer/builder/heads/voxelpose_head.py", "content": "# ------------------------------------------------------------------------------\n# Copyright and License Information\n# https://github.com/microsoft/voxelpose-pytorch/blob/main/lib/models\n# Original Licence: MIT License\n# ------------------------------------------------------------------------------\n\nimport torch\nimport torch.nn as nn\nimport torch.nn.functional as F\n\nfrom ..builder import HEADS\n\n\n@HEADS.register_module()\nclass CuboidCenterHead(nn.Module):\n \"\"\"Get results from the 3D human center heatmap. In this module, human 3D\n centers are local maximums obtained from the 3D heatmap via NMS (max-\n pooling).\n\n Args:\n space_size (list[3]): The size of the 3D space.\n cube_size (list[3]): The size of the heatmap volume.\n space_center (list[3]): The coordinate of space center.\n max_num (int): Maximum of human center detections.\n max_pool_kernel (int): Kernel size of the max-pool kernel in nms.\n \"\"\"\n\n def __init__(self,\n space_size,\n space_center,\n cube_size,\n max_num=10,\n max_pool_kernel=3):\n super(CuboidCenterHead, self).__init__()\n # use register_buffer\n self.register_buffer('grid_size', torch.tensor(space_size))\n self.register_buffer('cube_size', torch.tensor(cube_size))\n self.register_buffer('grid_center', torch.tensor(space_center))\n\n self.num_candidates = max_num\n self.max_pool_kernel = max_pool_kernel\n self.loss = nn.MSELoss()\n\n def _get_real_locations(self, indices):\n \"\"\"\n Args:\n indices (torch.Tensor(NXP)): Indices of points in the 3D tensor\n\n Returns:\n real_locations (torch.Tensor(NXPx3)): Locations of points\n in the world coordinate system\n \"\"\"\n real_locations = indices.float() / (\n self.cube_size - 1) * self.grid_size + \\\n self.grid_center - self.grid_size / 2.0\n return real_locations\n\n def _nms_by_max_pool(self, heatmap_volumes):\n max_num = self.num_candidates\n batch_size = heatmap_volumes.shape[0]\n root_cubes_nms = self._max_pool(heatmap_volumes)\n root_cubes_nms_reshape = root_cubes_nms.reshape(batch_size, -1)\n topk_values, topk_index = root_cubes_nms_reshape.topk(max_num)\n topk_unravel_index = self._get_3d_indices(topk_index,\n heatmap_volumes[0].shape)\n\n return topk_values, topk_unravel_index\n\n def _max_pool(self, inputs):\n kernel = self.max_pool_kernel\n padding = (kernel - 1) // 2\n max = F.max_pool3d(\n inputs, kernel_size=kernel, stride=1, padding=padding)\n keep = (inputs == max).float()\n return keep * inputs\n\n @staticmethod\n def _get_3d_indices(indices, shape):\n \"\"\"Get indices in the 3-D tensor.\n\n Args:\n indices (torch.Tensor(NXp)): Indices of points in the 1D tensor\n shape (torch.Size(3)): The shape of the original 3D tensor\n\n Returns:\n indices: Indices of points in the original 3D tensor\n \"\"\"\n batch_size = indices.shape[0]\n num_people = indices.shape[1]\n indices_x = (indices //\n (shape[1] * shape[2])).reshape(batch_size, num_people, -1)\n indices_y = ((indices % (shape[1] * shape[2])) //\n shape[2]).reshape(batch_size, num_people, -1)\n indices_z = (indices % shape[2]).reshape(batch_size, num_people, -1)\n indices = torch.cat([indices_x, indices_y, indices_z], dim=2)\n return indices\n\n def forward(self, heatmap_volumes):\n \"\"\"\n\n Args:\n heatmap_volumes (torch.Tensor(NXLXWXH)):\n 3D human center heatmaps predicted by the network.\n Returns:\n human_centers (torch.Tensor(NXPX5)):\n Coordinates of human centers.\n \"\"\"\n batch_size = heatmap_volumes.shape[0]\n\n topk_values, topk_unravel_index = self._nms_by_max_pool(\n heatmap_volumes.detach())\n\n topk_unravel_index = self._get_real_locations(topk_unravel_index)\n\n human_centers = torch.zeros(\n batch_size, self.num_candidates, 5, device=heatmap_volumes.device)\n human_centers[:, :, 0:3] = topk_unravel_index\n human_centers[:, :, 4] = topk_values\n\n return human_centers\n\n def get_loss(self, pred_cubes, gt):\n\n return dict(loss_center=self.loss(pred_cubes, gt))\n\n\n@HEADS.register_module()\nclass CuboidPoseHead(nn.Module):\n\n def __init__(self, beta):\n \"\"\"Get results from the 3D human pose heatmap. Instead of obtaining\n maximums on the heatmap, this module regresses the coordinates of\n keypoints via integral pose regression. Refer to `paper.\n\n ` for more details.\n\n Args:\n beta: Constant to adjust the magnification of soft-maxed heatmap.\n \"\"\"\n super(CuboidPoseHead, self).__init__()\n self.beta = beta\n self.loss = nn.L1Loss()\n\n def forward(self, heatmap_volumes, grid_coordinates):\n \"\"\"\n\n Args:\n heatmap_volumes (torch.Tensor(NxKxLxWxH)):\n 3D human pose heatmaps predicted by the network.\n grid_coordinates (torch.Tensor(Nx(LxWxH)x3)):\n Coordinates of the grids in the heatmap volumes.\n Returns:\n human_poses (torch.Tensor(NxKx3)): Coordinates of human poses.\n \"\"\"\n batch_size = heatmap_volumes.size(0)\n channel = heatmap_volumes.size(1)\n x = heatmap_volumes.reshape(batch_size, channel, -1, 1)\n x = F.softmax(self.beta * x, dim=2)\n grid_coordinates = grid_coordinates.unsqueeze(1)\n x = torch.mul(x, grid_coordinates)\n human_poses = torch.sum(x, dim=2)\n\n return human_poses\n\n def get_loss(self, preds, targets, weights):\n\n return dict(loss_pose=self.loss(preds * weights, targets * weights))\n" }, { "path": "eval/GVHMR/hmr4d/utils/preproc/vitpose_pytorch/src/vitpose_infer/builder/model_builder.py", "content": "import torch\n\n# from configs.coco.ViTPose_base_coco_256x192 import model\nfrom .heads.topdown_heatmap_simple_head import TopdownHeatmapSimpleHead\n\n# import TopdownHeatmapSimpleHead\nfrom .backbones import ViT\n\n# print(model)\nimport torch\nfrom functools import partial\nimport torch.nn as nn\nimport torch.nn.functional as F\nfrom importlib import import_module\n\n\ndef build_model(model_name, checkpoint=None):\n try:\n path = \".configs.coco.\" + model_name\n mod = import_module(path, package=\"src.vitpose_infer\")\n\n model = getattr(mod, \"model\")\n # from path import model\n except:\n raise ValueError(\"not a correct config\")\n\n head = TopdownHeatmapSimpleHead(\n in_channels=model[\"keypoint_head\"][\"in_channels\"],\n out_channels=model[\"keypoint_head\"][\"out_channels\"],\n num_deconv_filters=model[\"keypoint_head\"][\"num_deconv_filters\"],\n num_deconv_kernels=model[\"keypoint_head\"][\"num_deconv_kernels\"],\n num_deconv_layers=model[\"keypoint_head\"][\"num_deconv_layers\"],\n extra=model[\"keypoint_head\"][\"extra\"],\n )\n # print(head)\n backbone = ViT(\n img_size=model[\"backbone\"][\"img_size\"],\n patch_size=model[\"backbone\"][\"patch_size\"],\n embed_dim=model[\"backbone\"][\"embed_dim\"],\n depth=model[\"backbone\"][\"depth\"],\n num_heads=model[\"backbone\"][\"num_heads\"],\n ratio=model[\"backbone\"][\"ratio\"],\n mlp_ratio=model[\"backbone\"][\"mlp_ratio\"],\n qkv_bias=model[\"backbone\"][\"qkv_bias\"],\n drop_path_rate=model[\"backbone\"][\"drop_path_rate\"],\n )\n\n class VitPoseModel(nn.Module):\n def __init__(self, backbone, keypoint_head):\n super(VitPoseModel, self).__init__()\n self.backbone = backbone\n self.keypoint_head = keypoint_head\n\n def forward(self, x):\n x = self.backbone(x)\n x = self.keypoint_head(x)\n return x\n\n pose = VitPoseModel(backbone, head)\n if checkpoint is not None:\n check = torch.load(checkpoint)\n\n pose.load_state_dict(check[\"state_dict\"])\n return pose\n\n\n# pose = build_model('ViTPose_base_coco_256x192','./models/vitpose-b-multi-coco.pth')\n" }, { "path": "eval/GVHMR/hmr4d/utils/preproc/vitpose_pytorch/src/vitpose_infer/model_builder.py", "content": "import torch\n\n# from configs.coco.ViTPose_base_coco_256x192 import model\nfrom .builder.heads.topdown_heatmap_simple_head import TopdownHeatmapSimpleHead\n\n# import TopdownHeatmapSimpleHead\nfrom .builder.backbones import ViT\n\n# print(model)\nimport torch\nfrom functools import partial\nimport torch.nn as nn\nimport torch.nn.functional as F\nfrom importlib import import_module\n\nmodels = {\n \"ViTPose_huge_coco_256x192\": dict(\n type=\"TopDown\",\n pretrained=None,\n backbone=dict(\n type=\"ViT\",\n img_size=(256, 192),\n patch_size=16,\n embed_dim=1280,\n depth=32,\n num_heads=16,\n ratio=1,\n use_checkpoint=False,\n mlp_ratio=4,\n qkv_bias=True,\n drop_path_rate=0.55,\n ),\n keypoint_head=dict(\n type=\"TopdownHeatmapSimpleHead\",\n in_channels=1280,\n num_deconv_layers=2,\n num_deconv_filters=(256, 256),\n num_deconv_kernels=(4, 4),\n extra=dict(\n final_conv_kernel=1,\n ),\n out_channels=17,\n loss_keypoint=dict(type=\"JointsMSELoss\", use_target_weight=True),\n ),\n train_cfg=dict(),\n test_cfg=dict(),\n ),\n \"ViTPose_base_coco_256x192\": dict(\n type=\"TopDown\",\n pretrained=None,\n backbone=dict(\n type=\"ViT\",\n img_size=(256, 192),\n patch_size=16,\n embed_dim=768,\n depth=12,\n num_heads=12,\n ratio=1,\n use_checkpoint=False,\n mlp_ratio=4,\n qkv_bias=True,\n drop_path_rate=0.3,\n ),\n keypoint_head=dict(\n type=\"TopdownHeatmapSimpleHead\",\n in_channels=768,\n num_deconv_layers=2,\n num_deconv_filters=(256, 256),\n num_deconv_kernels=(4, 4),\n extra=dict(\n final_conv_kernel=1,\n ),\n out_channels=17,\n loss_keypoint=dict(type=\"JointsMSELoss\", use_target_weight=True),\n ),\n train_cfg=dict(),\n test_cfg=dict(),\n ),\n \"ViTPose_base_simple_coco_256x192\": dict(\n type=\"TopDown\",\n pretrained=None,\n backbone=dict(\n type=\"ViT\",\n img_size=(256, 192),\n patch_size=16,\n embed_dim=768,\n depth=12,\n num_heads=12,\n ratio=1,\n use_checkpoint=False,\n mlp_ratio=4,\n qkv_bias=True,\n drop_path_rate=0.3,\n ),\n keypoint_head=dict(\n type=\"TopdownHeatmapSimpleHead\",\n in_channels=768,\n num_deconv_layers=0,\n num_deconv_filters=[],\n num_deconv_kernels=[],\n upsample=4,\n extra=dict(\n final_conv_kernel=3,\n ),\n out_channels=17,\n loss_keypoint=dict(type=\"JointsMSELoss\", use_target_weight=True),\n ),\n train_cfg=dict(),\n test_cfg=dict(\n flip_test=True,\n post_process=\"default\",\n shift_heatmap=False,\n target_type=\"GaussianHeatmap\",\n modulate_kernel=11,\n use_udp=True,\n ),\n ),\n}\n\n\ndef build_model(model_name, checkpoint=None):\n try:\n model = models[model_name]\n except:\n raise ValueError(\"not a correct config\")\n\n head = TopdownHeatmapSimpleHead(\n in_channels=model[\"keypoint_head\"][\"in_channels\"],\n out_channels=model[\"keypoint_head\"][\"out_channels\"],\n num_deconv_filters=model[\"keypoint_head\"][\"num_deconv_filters\"],\n num_deconv_kernels=model[\"keypoint_head\"][\"num_deconv_kernels\"],\n num_deconv_layers=model[\"keypoint_head\"][\"num_deconv_layers\"],\n extra=model[\"keypoint_head\"][\"extra\"],\n )\n # print(head)\n backbone = ViT(\n img_size=model[\"backbone\"][\"img_size\"],\n patch_size=model[\"backbone\"][\"patch_size\"],\n embed_dim=model[\"backbone\"][\"embed_dim\"],\n depth=model[\"backbone\"][\"depth\"],\n num_heads=model[\"backbone\"][\"num_heads\"],\n ratio=model[\"backbone\"][\"ratio\"],\n mlp_ratio=model[\"backbone\"][\"mlp_ratio\"],\n qkv_bias=model[\"backbone\"][\"qkv_bias\"],\n drop_path_rate=model[\"backbone\"][\"drop_path_rate\"],\n )\n\n class VitPoseModel(nn.Module):\n def __init__(self, backbone, keypoint_head):\n super(VitPoseModel, self).__init__()\n self.backbone = backbone\n self.keypoint_head = keypoint_head\n\n def forward(self, x):\n x = self.backbone(x)\n x = self.keypoint_head(x)\n return x\n\n pose = VitPoseModel(backbone, head)\n if checkpoint is not None:\n check = torch.load(checkpoint)\n\n pose.load_state_dict(check[\"state_dict\"])\n return pose\n\n\n# pose = build_model('ViTPose_base_coco_256x192','./models/vitpose-b-multi-coco.pth')\n" }, { "path": "eval/GVHMR/hmr4d/utils/preproc/vitpose_pytorch/src/vitpose_infer/pose_utils/ViTPose_trt.py", "content": "import tensorrt as trt\nimport torch.nn\nfrom collections import OrderedDict, namedtuple\nimport numpy as np\n\ndef torch_device_from_trt(device):\n if device == trt.TensorLocation.DEVICE:\n return torch.device(\"cuda\")\n elif device == trt.TensorLocation.HOST:\n return torch.device(\"cpu\")\n else:\n return TypeError(\"%s is not supported by torch\" % device)\ndef torch_dtype_from_trt(dtype):\n if dtype == trt.int8:\n return torch.int8\n elif trt.__version__ >= '7.0' and dtype == trt.bool:\n return torch.bool\n elif dtype == trt.int32:\n return torch.int32\n elif dtype == trt.float16:\n return torch.float16\n elif dtype == trt.float32:\n return torch.float32\n else:\n raise TypeError(\"%s is not supported by torch\" % dtype)\nclass TRTModule_ViTPose(torch.nn.Module):\n def __init__(self, engine=None, input_names=None, output_names=None, input_flattener=None, output_flattener=None,path=None,device=None):\n super(TRTModule_ViTPose, self).__init__()\n # self._register_state_dict_hook(TRTModule._on_state_dict)\n # self.engine = engine\n logger = trt.Logger(trt.Logger.INFO)\n with open(path, 'rb') as f, trt.Runtime(logger) as runtime:\n self.engine = runtime.deserialize_cuda_engine(f.read())\n if self.engine is not None:\n self.context = self.engine.create_execution_context()\n self.input_names = ['images']\n self.output_names = []\n self.input_flattener = input_flattener\n self.output_flattener = output_flattener\n Binding = namedtuple('Binding', ('name', 'dtype', 'shape', 'data', 'ptr'))\n \n # with open(path, 'rb') as f, trt.Runtime(logger) as runtime:\n # self.model = runtime.deserialize_cuda_engine(f.read())\n # self.context = self.model.create_execution_context()\n self.bindings = OrderedDict()\n # self.output_names = []\n fp16 = False # default updated below\n dynamic = False\n for i in range(self.engine.num_bindings):\n name = self.engine.get_binding_name(i)\n dtype = trt.nptype(self.engine.get_binding_dtype(i))\n if self.engine.binding_is_input(i):\n if -1 in tuple(self.engine.get_binding_shape(i)): # dynamic\n dynamic = True\n self.context.set_binding_shape(i, tuple(self.engine.get_profile_shape(0, i)[2]))\n if dtype == np.float16:\n fp16 = True\n else: # output\n self.output_names.append(name)\n shape = tuple(self.context.get_binding_shape(i))\n im = torch.from_numpy(np.empty(shape, dtype=dtype)).to(device)\n self.bindings[name] = Binding(name, dtype, shape, im, int(im.data_ptr()))\n self.binding_addrs = OrderedDict((n, d.ptr) for n, d in self.bindings.items())\n self.batch_size = self.bindings['images'].shape[0] \n\n \n\n def forward(self, *inputs):\n bindings = [None] * (len(self.input_names) + len(self.output_names))\n \n if self.input_flattener is not None:\n inputs = self.input_flattener.flatten(inputs)\n\n for i, input_name in enumerate(self.input_names):\n idx = self.engine.get_binding_index(input_name)\n shape = tuple(inputs[i].shape)\n bindings[idx] = inputs[i].contiguous().data_ptr()\n self.context.set_binding_shape(idx, shape)\n\n # create output tensors\n outputs = [None] * len(self.output_names)\n for i, output_name in enumerate(self.output_names):\n idx = self.engine.get_binding_index(output_name)\n dtype = torch_dtype_from_trt(self.engine.get_binding_dtype(idx))\n shape = tuple(self.context.get_binding_shape(idx))\n device = torch_device_from_trt(self.engine.get_location(idx))\n output = torch.empty(size=shape, dtype=dtype, device=device)\n outputs[i] = output\n bindings[idx] = output.data_ptr()\n\n self.context.execute_async_v2(\n bindings, torch.cuda.current_stream().cuda_stream\n )\n\n if self.output_flattener is not None:\n outputs = self.output_flattener.unflatten(outputs)\n else:\n outputs = tuple(outputs)\n if len(outputs) == 1:\n outputs = outputs[0]\n\n return outputs" }, { "path": "eval/GVHMR/hmr4d/utils/preproc/vitpose_pytorch/src/vitpose_infer/pose_utils/__init__.py", "content": "" }, { "path": "eval/GVHMR/hmr4d/utils/preproc/vitpose_pytorch/src/vitpose_infer/pose_utils/convert_to_trt.py", "content": "from torch2trt import TRTModule,torch2trt\nfrom builder import build_model\nimport torch\npose = build_model('ViTPose_base_coco_256x192','./models/vitpose-b.pth')\npose.cuda().eval()\n\nx = torch.ones(1,3,256,192).cuda()\nnet_trt = torch2trt(pose, [x],max_batch_size=10, fp16_mode=True)\ntorch.save(net_trt.state_dict(), 'vitpose_trt.pth')" }, { "path": "eval/GVHMR/hmr4d/utils/preproc/vitpose_pytorch/src/vitpose_infer/pose_utils/general_utils.py", "content": "#!/usr/bin/env python3\n# -*- coding: utf-8 -*-\n\"\"\"\nCreated on Wed Jun 15 15:49:22 2022\n\n@author: gpastal\n\"\"\"\nimport numpy as np\nimport numpy.ma as ma\n# import pika\nimport json\nfrom collections import OrderedDict\nfrom collections.abc import Iterable\nfrom itertools import chain\nimport argparse\n\ndef make_parser():\n parser = argparse.ArgumentParser(\"ByteTrack Demo!\")\n # exp file\n # tracking args\n parser.add_argument(\"--track_thresh\", type=float, default=0.2, help=\"tracking confidence threshold\")\n parser.add_argument(\"--track_buffer\", type=int, default=240, help=\"the frames for keep lost tracks\")\n parser.add_argument(\"--match_thresh\", type=float, default=0.8, help=\"matching threshold for tracking\")\n parser.add_argument(\n \"--aspect_ratio_thresh\", type=float, default=1.6,\n help=\"threshold for filtering out boxes of which aspect ratio are above the given value.\"\n )\n parser.add_argument('--min_box_area', type=float, default=10, help='filter out tiny boxes')\n parser.add_argument(\"--mot20\", dest=\"mot20\", default=False, action=\"store_true\", help=\"test mot20.\")\n return parser\n\ndef jitter(tracking,temp,id1):\n pass\ndef jitter2(tracking,temp,id1) :\n pass\n \ndef create_json_rabbitmq( FRAME_ID,pose):\n pass\n\ndef producer_rabbitmq():\n pass\ndef fix_head(xyz):\n pass\n\ndef flatten_lst(x):\n if isinstance(x, Iterable):\n return [a for i in x for a in flatten_lst(i)]\n else:\n return [x]\n\ndef polys_from_pose(pts):\n seg=[]\n for ind, i in enumerate(pts):\n list_=[]\n list_sc=[]\n # list1 = [i[0][1],i[0][0]]\n\n # list2 = [i[0][1],i[0][0]]\n # print(i)\n for j in i:\n \n temp_ = [j[1],j[0]]\n \n if j[2]>0.4:\n\n temp2_ = [1]\n else:\n temp2_ =[0]\n list_.append(temp_)\n list_sc.append(temp2_)\n # print(list_sc)\n # list2 = [i[6][1],i[6][0]]\n # list3 = [i[11][1],i[11][0]]\n # list4 = [i[12][1],i[12][0]]\n \n # list_ = flatten_lst(list_)\n # print(list_)\n list_=fix_list_order(list_,list_sc)\n # print(list_)\n # list_=list(list_)\n # list_ = list_.to_list()\n # print(list_)\n # temp__=list(chain(*list_))\n seg.append(list_) # temp_ = list(chain(list1,list2,list3,list4,list1))\n return seg\ndef fix_list_order(list_,list2):\n # for index,values in enumerate(list_):\n myorder = [0, 2, 4, 6, 8,10,12,14,16,15,13,11,9,7,5,3,1]\n cor_list = [list_[i] for i in myorder]\n cor_list2 = [list2[i] for i in myorder]\n # print(cor_list)\n # result = list(set(map(tuple,cor_list)) & set(map(tuple,cor_list2)))\n # arr = np.array([x for x in cor_list])\n # print(cor_list)\n data = np.asarray(cor_list)\n # print(data)\n mask = np.column_stack((cor_list2, cor_list2))\n # masked = ma.masked_array(data, mask=np.column_stack((cor_list2, cor_list2)))#[cor_list2,cor_list2])\n # result = list(set(masked[~masked.mask]))\n # print(result)\n # print(data)\n # print(mask)\n result2 = []\n for inde,i in enumerate(data):\n # print(mask[inde])\n if mask[inde].all()==1:\n result2.append(i[0])\n result2.append(i[1])\n # result = [int(result[i] for i in result)]\n return result2\n\n" }, { "path": "eval/GVHMR/hmr4d/utils/preproc/vitpose_pytorch/src/vitpose_infer/pose_utils/inference_test.py", "content": "from builder import build_model\nimport torch\nfrom ViTPose_trt import TRTModule_ViTPose\n# pose = TRTModule_ViTPose(path='pose_higher_hrnet_w32_512.engine',device='cuda:0')\npose = build_model('ViTPose_base_coco_256x192','./models/vitpose-b.pth')\npose.cuda().eval()\nif pose.training:\n print('train')\nelse:\n print('eval')\ndevice = torch.device(\"cuda\")\n# pose.to(device)\ndummy_input = torch.randn(10, 3,256,192, dtype=torch.float).to(device)\nrepetitions=100\ntotal_time = 0\nstarter, ender = torch.cuda.Event(enable_timing=True), torch.cuda.Event(enable_timing=True)\nwith torch.no_grad():\n for rep in range(repetitions):\n # starter, ender = torch.cuda.Event(enable_timing=True), torch.cuda.Event(enable_timing=True)\n starter.record()\n # for k in range(10):\n _ = pose(dummy_input)\n ender.record()\n torch.cuda.synchronize()\n curr_time = starter.elapsed_time(ender)/1000\n total_time += curr_time\nThroughput = repetitions*10/total_time\nprint('Final Throughput:',Throughput)\nprint('Total time',total_time)" }, { "path": "eval/GVHMR/hmr4d/utils/preproc/vitpose_pytorch/src/vitpose_infer/pose_utils/logger_helper.py", "content": "import logging\n\nclass CustomFormatter(logging.Formatter):\n\n grey = \"\\x1b[38;20m\"\n yellow = \"\\x1b[33;20m\"\n red = \"\\x1b[31;20m\"\n bold_red = \"\\x1b[31;1m\"\n reset = \"\\x1b[0m\"\n format = \"%(asctime)s - %(name)s - %(levelname)s - %(message)s (%(filename)s:%(lineno)d)\"\n\n FORMATS = {\n logging.DEBUG: grey + format + reset,\n logging.INFO: grey + format + reset,\n logging.WARNING: yellow + format + reset,\n logging.ERROR: red + format + reset,\n logging.CRITICAL: bold_red + format + reset\n }\n\n def format(self, record):\n log_fmt = self.FORMATS.get(record.levelno)\n formatter = logging.Formatter(log_fmt)\n return formatter.format(record)" }, { "path": "eval/GVHMR/hmr4d/utils/preproc/vitpose_pytorch/src/vitpose_infer/pose_utils/pose_utils.py", "content": "#!/usr/bin/env python3\n# -*- coding: utf-8 -*-\n\"\"\"\nCreated on Wed Jun 15 15:45:33 2022\n\n@author: gpastal\n\"\"\"\nimport torch\nimport torchvision\nimport torch.nn.functional as F\nfrom torchvision import transforms as TR\nimport numpy as np\nimport cv2\nimport logging\n# from simpleHRNet.models_.hrnet import HRNet\n# from torch2trt import torch2trt,TRTModule\nlogger = logging.getLogger(\"Tracker !\")\nfrom .timerr import Timer\nfrom pathlib import Path\n# import gdown\ntimer_det = Timer()\ntimer_track = Timer()\ntimer_pose = Timer()\n\ndef pose_points_yolo5(detector,image,pose,tracker,tensorrt):\n timer_det.tic()\n # starter, ender = torch.cuda.Event(enable_timing=True), torch.cuda.Event(enable_timing=True)\n \n transform = TR.Compose([\n TR.ToPILImage(),\n # Padd(),\n TR.Resize((256, 192)), # (height, width)\n TR.ToTensor(),\n TR.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),\n ])\n # image = cv2.cvtColor(image, cv2.COLOR_BGRA2BGR)\n\n detections = detector(image)\n timer_det.toc()\n logger.info('DET FPS -- %s',1./timer_det.average_time)\n # print(detections.shape)\n dets = detections.xyxy[0]\n dets = dets[dets[:,5] == 0.]\n # dets = dets[dets[:,4] > 0.3]\n # logger.warning(len(dets))\n \n # if len(dets)>0:\n # image_gpu = torch.tensor(image).cuda()/255.\n # print(image_gpu.size())\n timer_track.tic()\n online_targets=tracker.update(dets,[image.shape[0],image.shape[1]],image.shape)\n\n online_tlwhs = []\n online_ids = []\n online_scores = []\n for t in online_targets:\n tlwh = t.tlwh\n tid = t.track_id\n # vertical = tlwh[2] / tlwh[3] > args.aspect_ratio_threshs\n if tlwh[2] * tlwh[3] > 10 :#and not vertical:\n online_tlwhs.append(tlwh)\n online_ids.append(tid)\n online_scores.append(t.score)\n # tracker.update()\n timer_track.toc()\n logger.info('TRACKING FPS --%s',1./timer_track.average_time)\n device='cuda'\n nof_people = len(online_ids) if online_ids is not None else 0\n # nof_people=1\n # print(dets)\n # print(nof_people)\n boxes = torch.empty((nof_people, 4), dtype=torch.int32,device= 'cuda')\n # boxes = []\n images = torch.empty((nof_people, 3, 256, 192)) # (height, width)\n heatmaps = np.zeros((nof_people, 17, 64, 48),\n dtype=np.float32)\n # starter.record()\n # print(online_tlwhs)\n if len(online_tlwhs):\n for i, (x1, y1, x2, y2) in enumerate(online_tlwhs):\n # for i, (x1, y1, x2, y2) in enumerate(np.array([[55,399,424-55,479-399]])):\n # if i<1:\n x1 = x1.astype(np.int32)\n x2 = x1+x2.astype(np.int32)\n y1 = y1.astype(np.int32)\n y2 = y1+ y2.astype(np.int32)\n if x2>image.shape[1]:x2=image.shape[1]-1\n if y2>image.shape[0]:y2=image.shape[0]-1\n if y1<0: y1=0\n if x1<0 : x1=0\n # print([x1,x2,y1,y2])\n # image = cv2.rectangle(image, (x1,y1), (x2,y2), (0,0,0), 1)\n # cv2.imwrite('saved.png',image)\n # # Adapt detections to match HRNet input aspect ratio (as suggested by xtyDoge in issue #14)\n correction_factor = 256 / 192 * (x2 - x1) / (y2 - y1)\n if correction_factor > 1:\n # increase y side\n center = y1 + (y2 - y1) // 2\n length = int(round((y2 - y1) * correction_factor))\n y1_new = int( center - length // 2)\n y2_new = int( center + length // 2)\n image_crop = image[y1:y2, x1:x2, ::-1]\n # print(y1,y2,x1,x2)\n pad = (int(abs(y1_new-y1))), int(abs(y2_new-y2))\n image_crop = np.pad(image_crop,((pad), (0, 0), (0, 0)))\n images[i] = transform(image_crop)\n boxes[i]= torch.tensor([x1, y1_new, x2, y2_new])\n \n elif correction_factor < 1:\n # increase x side\n center = x1 + (x2 - x1) // 2\n length = int(round((x2 - x1) * 1 / correction_factor))\n x1_new = int( center - length // 2)\n x2_new = int( center + length // 2)\n # images[i] = transform(image[y1:y2, x1:x2, ::-1])\n image_crop = image[y1:y2, x1:x2, ::-1]\n pad = (abs(x1_new-x1)), int(abs(x2_new-x2))\n image_crop = np.pad(image_crop,((0, 0), (pad), (0, 0)))\n images[i] = transform(image_crop)\n boxes[i]= torch.tensor([x1_new, y1, x2_new, y2])\n \n\n if images.shape[0] > 0:\n images = images.to(device)\n\n if tensorrt:\n out = torch.zeros((images.shape[0],17,64,48),device=device)\n with torch.no_grad():\n timer_pose.tic()\n\n for i in range(images.shape[0]):\n # timer_pose.tic()\n # print(images[i].size())\n \n out[i] = pose(images[i].unsqueeze(0))\n timer_pose.toc()\n logger.info('POSE FPS -- %s',1./timer_pose.average_time)\n else:\n with torch.no_grad():\n \n timer_pose.tic()\n\n \n \n out = pose(images)\n timer_pose.toc()\n logger.info('POSE FPS -- %s',1./timer_pose.average_time)\n \n\n pts = torch.empty((out.shape[0], out.shape[1], 3), dtype=torch.float32,device=device)\n pts2 = np.empty((out.shape[0], out.shape[1], 3), dtype=np.float32)\n\n (b,indices)=torch.max(out,dim=2)\n (b,indices)=torch.max(b,dim=2)\n \n (c,indicesc)=torch.max(out,dim=3)\n (c,indicesc)=torch.max(c,dim=2)\n dim1= torch.tensor(1. / 64,device=device)\n dim2= torch.tensor(1. / 48,device=device)\n\n for i in range(0,out.shape[0]):\n\n pts[i, :, 0] = indicesc[i,:] * dim1 * (boxes[i][3] - boxes[i][1]) + boxes[i][1]\n pts[i, :, 1] = indices[i,:] *dim2* (boxes[i][2] - boxes[i][0]) + boxes[i][0]\n pts[i, :, 2] = c[i,:]\n\n pts=pts.cpu().numpy()\n # print(pts)\n else:\n pts = np.empty((0, 0, 3), dtype=np.float32)\n online_tlwhs = []\n online_ids = []\n online_scores=[]\n res = list()\n\n res.append(pts)\n\n\n if len(res) > 1:\n return res,online_tlwhs,online_ids,online_scores#,pts2\n else:\n return res[0],online_tlwhs,online_ids,online_scores#,pts2\n\n" }, { "path": "eval/GVHMR/hmr4d/utils/preproc/vitpose_pytorch/src/vitpose_infer/pose_utils/pose_viz.py", "content": "import cv2\r\nimport matplotlib.pyplot as plt\r\nimport numpy as np\r\nimport torch\r\nimport torchvision\r\nimport ffmpeg\r\n\r\n\r\ndef joints_dict():\r\n joints = {\r\n \"coco\": {\r\n \"keypoints\": {\r\n 0: \"nose\",\r\n 1: \"left_eye\",\r\n 2: \"right_eye\",\r\n 3: \"left_ear\",\r\n 4: \"right_ear\",\r\n 5: \"left_shoulder\",\r\n 6: \"right_shoulder\",\r\n 7: \"left_elbow\",\r\n 8: \"right_elbow\",\r\n 9: \"left_wrist\",\r\n 10: \"right_wrist\",\r\n 11: \"left_hip\",\r\n 12: \"right_hip\",\r\n 13: \"left_knee\",\r\n 14: \"right_knee\",\r\n 15: \"left_ankle\",\r\n 16: \"right_ankle\"\r\n },\r\n \"skeleton\": [\r\n # # [16, 14], [14, 12], [17, 15], [15, 13], [12, 13], [6, 12], [7, 13], [6, 7], [6, 8],\r\n # # [7, 9], [8, 10], [9, 11], [2, 3], [1, 2], [1, 3], [2, 4], [3, 5], [4, 6], [5, 7]\r\n # [15, 13], [13, 11], [16, 14], [14, 12], [11, 12], [5, 11], [6, 12], [5, 6], [5, 7],\r\n # [6, 8], [7, 9], [8, 10], [1, 2], [0, 1], [0, 2], [1, 3], [2, 4], [3, 5], [4, 6]\r\n [15, 13], [13, 11], [16, 14], [14, 12], [11, 12], [5, 11], [6, 12], [5, 6], [5, 7],\r\n [6, 8], [7, 9], [8, 10], [1, 2], [0, 1], [0, 2], [1, 3], [2, 4], # [3, 5], [4, 6]\r\n [0, 5], [0, 6]\r\n ]\r\n },\r\n \"mpii\": {\r\n \"keypoints\": {\r\n 0: \"right_ankle\",\r\n 1: \"right_knee\",\r\n 2: \"right_hip\",\r\n 3: \"left_hip\",\r\n 4: \"left_knee\",\r\n 5: \"left_ankle\",\r\n 6: \"pelvis\",\r\n 7: \"thorax\",\r\n 8: \"upper_neck\",\r\n 9: \"head top\",\r\n 10: \"right_wrist\",\r\n 11: \"right_elbow\",\r\n 12: \"right_shoulder\",\r\n 13: \"left_shoulder\",\r\n 14: \"left_elbow\",\r\n 15: \"left_wrist\"\r\n },\r\n \"skeleton\": [\r\n # [5, 4], [4, 3], [0, 1], [1, 2], [3, 2], [13, 3], [12, 2], [13, 12], [13, 14],\r\n # [12, 11], [14, 15], [11, 10], # [2, 3], [1, 2], [1, 3], [2, 4], [3, 5], [4, 6], [5, 7]\r\n [5, 4], [4, 3], [0, 1], [1, 2], [3, 2], [3, 6], [2, 6], [6, 7], [7, 8], [8, 9],\r\n [13, 7], [12, 7], [13, 14], [12, 11], [14, 15], [11, 10],\r\n ]\r\n },\r\n }\r\n return joints\r\n\r\n\r\ndef draw_points(image, points, color_palette='tab20', palette_samples=16, confidence_threshold=0.5):\r\n \"\"\"\r\n Draws `points` on `image`.\r\n\r\n Args:\r\n image: image in opencv format\r\n points: list of points to be drawn.\r\n Shape: (nof_points, 3)\r\n Format: each point should contain (y, x, confidence)\r\n color_palette: name of a matplotlib color palette\r\n Default: 'tab20'\r\n palette_samples: number of different colors sampled from the `color_palette`\r\n Default: 16\r\n confidence_threshold: only points with a confidence higher than this threshold will be drawn. Range: [0, 1]\r\n Default: 0.5\r\n\r\n Returns:\r\n A new image with overlaid points\r\n\r\n \"\"\"\r\n try:\r\n colors = np.round(\r\n np.array(plt.get_cmap(color_palette).colors) * 255\r\n ).astype(np.uint8)[:, ::-1].tolist()\r\n except AttributeError: # if palette has not pre-defined colors\r\n colors = np.round(\r\n np.array(plt.get_cmap(color_palette)(np.linspace(0, 1, palette_samples))) * 255\r\n ).astype(np.uint8)[:, -2::-1].tolist()\r\n\r\n circle_size = max(1, min(image.shape[:2]) // 160) # ToDo Shape it taking into account the size of the detection\r\n # circle_size = max(2, int(np.sqrt(np.max(np.max(points, axis=0) - np.min(points, axis=0)) // 16)))\r\n\r\n for i, pt in enumerate(points):\r\n if pt[2] > confidence_threshold:\r\n image = cv2.circle(image, (int(pt[1]), int(pt[0])), circle_size, tuple(colors[i % len(colors)]), -1)\r\n\r\n return image\r\n\r\n\r\ndef draw_skeleton(image, points, skeleton, color_palette='Set2', palette_samples=8, person_index=0,\r\n confidence_threshold=0.5):\r\n \"\"\"\r\n Draws a `skeleton` on `image`.\r\n\r\n Args:\r\n image: image in opencv format\r\n points: list of points to be drawn.\r\n Shape: (nof_points, 3)\r\n Format: each point should contain (y, x, confidence)\r\n skeleton: list of joints to be drawn\r\n Shape: (nof_joints, 2)\r\n Format: each joint should contain (point_a, point_b) where `point_a` and `point_b` are an index in `points`\r\n color_palette: name of a matplotlib color palette\r\n Default: 'Set2'\r\n palette_samples: number of different colors sampled from the `color_palette`\r\n Default: 8\r\n person_index: index of the person in `image`\r\n Default: 0\r\n confidence_threshold: only points with a confidence higher than this threshold will be drawn. Range: [0, 1]\r\n Default: 0.5\r\n\r\n Returns:\r\n A new image with overlaid joints\r\n\r\n \"\"\"\r\n try:\r\n colors = np.round(\r\n np.array(plt.get_cmap(color_palette).colors) * 255\r\n ).astype(np.uint8)[:, ::-1].tolist()\r\n except AttributeError: # if palette has not pre-defined colors\r\n colors = np.round(\r\n np.array(plt.get_cmap(color_palette)(np.linspace(0, 1, palette_samples))) * 255\r\n ).astype(np.uint8)[:, -2::-1].tolist()\r\n\r\n for i, joint in enumerate(skeleton):\r\n pt1, pt2 = points[joint]\r\n if pt1[2] > confidence_threshold and pt2[2] > confidence_threshold:\r\n image = cv2.line(\r\n image, (int(pt1[1]), int(pt1[0])), (int(pt2[1]), int(pt2[0])),\r\n tuple(colors[person_index % len(colors)]), 2\r\n )\r\n\r\n return image\r\n\r\n\r\ndef draw_points_and_skeleton(image, points, skeleton, points_color_palette='tab20', points_palette_samples=16,\r\n skeleton_color_palette='Set2', skeleton_palette_samples=8, person_index=0,\r\n confidence_threshold=0.5):\r\n \"\"\"\r\n Draws `points` and `skeleton` on `image`.\r\n\r\n Args:\r\n image: image in opencv format\r\n points: list of points to be drawn.\r\n Shape: (nof_points, 3)\r\n Format: each point should contain (y, x, confidence)\r\n skeleton: list of joints to be drawn\r\n Shape: (nof_joints, 2)\r\n Format: each joint should contain (point_a, point_b) where `point_a` and `point_b` are an index in `points`\r\n points_color_palette: name of a matplotlib color palette\r\n Default: 'tab20'\r\n points_palette_samples: number of different colors sampled from the `color_palette`\r\n Default: 16\r\n skeleton_color_palette: name of a matplotlib color palette\r\n Default: 'Set2'\r\n skeleton_palette_samples: number of different colors sampled from the `color_palette`\r\n Default: 8\r\n person_index: index of the person in `image`\r\n Default: 0\r\n confidence_threshold: only points with a confidence higher than this threshold will be drawn. Range: [0, 1]\r\n Default: 0.5\r\n\r\n Returns:\r\n A new image with overlaid joints\r\n\r\n \"\"\"\r\n image = draw_skeleton(image, points, skeleton, color_palette=skeleton_color_palette,\r\n palette_samples=skeleton_palette_samples, person_index=person_index,\r\n confidence_threshold=confidence_threshold)\r\n image = draw_points(image, points, color_palette=points_color_palette, palette_samples=points_palette_samples,\r\n confidence_threshold=confidence_threshold)\r\n return image\r\n\r\n\r\ndef save_images(images, target, joint_target, output, joint_output, joint_visibility, summary_writer=None, step=0,\r\n prefix=''):\r\n \"\"\"\r\n Creates a grid of images with gt joints and a grid with predicted joints.\r\n This is a basic function for debugging purposes only.\r\n\r\n If summary_writer is not None, the grid will be written in that SummaryWriter with name \"{prefix}_images\" and\r\n \"{prefix}_predictions\".\r\n\r\n Args:\r\n images (torch.Tensor): a tensor of images with shape (batch x channels x height x width).\r\n target (torch.Tensor): a tensor of gt heatmaps with shape (batch x channels x height x width).\r\n joint_target (torch.Tensor): a tensor of gt joints with shape (batch x joints x 2).\r\n output (torch.Tensor): a tensor of predicted heatmaps with shape (batch x channels x height x width).\r\n joint_output (torch.Tensor): a tensor of predicted joints with shape (batch x joints x 2).\r\n joint_visibility (torch.Tensor): a tensor of joint visibility with shape (batch x joints).\r\n summary_writer (tb.SummaryWriter): a SummaryWriter where write the grids.\r\n Default: None\r\n step (int): summary_writer step.\r\n Default: 0\r\n prefix (str): summary_writer name prefix.\r\n Default: \"\"\r\n\r\n Returns:\r\n A pair of images which are built from torchvision.utils.make_grid\r\n \"\"\"\r\n # Input images with gt\r\n images_ok = images.detach().clone()\r\n images_ok[:, 0].mul_(0.229).add_(0.485)\r\n images_ok[:, 1].mul_(0.224).add_(0.456)\r\n images_ok[:, 2].mul_(0.225).add_(0.406)\r\n for i in range(images.shape[0]):\r\n joints = joint_target[i] * 4.\r\n joints_vis = joint_visibility[i]\r\n\r\n for joint, joint_vis in zip(joints, joints_vis):\r\n if joint_vis[0]:\r\n a = int(joint[1].item())\r\n b = int(joint[0].item())\r\n # images_ok[i][:, a-1:a+1, b-1:b+1] = torch.tensor([1, 0, 0])\r\n images_ok[i][0, a - 1:a + 1, b - 1:b + 1] = 1\r\n images_ok[i][1:, a - 1:a + 1, b - 1:b + 1] = 0\r\n grid_gt = torchvision.utils.make_grid(images_ok, nrow=int(images_ok.shape[0] ** 0.5), padding=2, normalize=False)\r\n if summary_writer is not None:\r\n summary_writer.add_image(prefix + 'images', grid_gt, global_step=step)\r\n\r\n # Input images with prediction\r\n images_ok = images.detach().clone()\r\n images_ok[:, 0].mul_(0.229).add_(0.485)\r\n images_ok[:, 1].mul_(0.224).add_(0.456)\r\n images_ok[:, 2].mul_(0.225).add_(0.406)\r\n for i in range(images.shape[0]):\r\n joints = joint_output[i] * 4.\r\n joints_vis = joint_visibility[i]\r\n\r\n for joint, joint_vis in zip(joints, joints_vis):\r\n if joint_vis[0]:\r\n a = int(joint[1].item())\r\n b = int(joint[0].item())\r\n # images_ok[i][:, a-1:a+1, b-1:b+1] = torch.tensor([1, 0, 0])\r\n images_ok[i][0, a - 1:a + 1, b - 1:b + 1] = 1\r\n images_ok[i][1:, a - 1:a + 1, b - 1:b + 1] = 0\r\n grid_pred = torchvision.utils.make_grid(images_ok, nrow=int(images_ok.shape[0] ** 0.5), padding=2, normalize=False)\r\n if summary_writer is not None:\r\n summary_writer.add_image(prefix + 'predictions', grid_pred, global_step=step)\r\n\r\n # Heatmaps\r\n # ToDo\r\n # for h in range(0,17):\r\n # heatmap = torchvision.utils.make_grid(output[h].detach(), nrow=int(np.sqrt(output.shape[0])),\r\n # padding=2, normalize=True, range=(0, 1))\r\n # summary_writer.add_image('train_heatmap_%d' % h, heatmap, global_step=step + epoch*len_dl_train)\r\n\r\n return grid_gt, grid_pred\r\n\r\n\r\ndef check_video_rotation(filename):\r\n # thanks to\r\n # https://stackoverflow.com/questions/53097092/frame-from-video-is-upside-down-after-extracting/55747773#55747773\r\n\r\n # this returns meta-data of the video file in form of a dictionary\r\n meta_dict = ffmpeg.probe(filename)\r\n\r\n # from the dictionary, meta_dict['streams'][0]['tags']['rotate'] is the key\r\n # we are looking for\r\n rotation_code = None\r\n try:\r\n if int(meta_dict['streams'][0]['tags']['rotate']) == 90:\r\n rotation_code = cv2.ROTATE_90_CLOCKWISE\r\n elif int(meta_dict['streams'][0]['tags']['rotate']) == 180:\r\n rotation_code = cv2.ROTATE_180\r\n elif int(meta_dict['streams'][0]['tags']['rotate']) == 270:\r\n rotation_code = cv2.ROTATE_90_COUNTERCLOCKWISE\r\n else:\r\n raise ValueError\r\n except KeyError:\r\n pass\r\n\r\n return rotation_code\r\n" }, { "path": "eval/GVHMR/hmr4d/utils/preproc/vitpose_pytorch/src/vitpose_infer/pose_utils/timerr.py", "content": "import time\n\n\nclass Timer(object):\n \"\"\"A simple timer.\"\"\"\n def __init__(self):\n self.total_time = 0.\n self.calls = 0\n self.start_time = 0.\n self.diff = 0.\n self.average_time = 0.\n\n self.duration = 0.\n\n def tic(self):\n # using time.time instead of time.clock because time time.clock\n # does not normalize for multithreading\n self.start_time = time.time()\n\n def toc(self, average=True):\n self.diff = time.time() - self.start_time\n self.total_time += self.diff\n self.calls += 1\n self.average_time = self.total_time / self.calls\n if average:\n self.duration = self.average_time\n else:\n self.duration = self.diff\n return self.duration\n\n def clear(self):\n self.total_time = 0.\n self.calls = 0\n self.start_time = 0.\n self.diff = 0.\n self.average_time = 0.\n self.duration = 0." }, { "path": "eval/GVHMR/hmr4d/utils/preproc/vitpose_pytorch/src/vitpose_infer/pose_utils/visualizer.py", "content": "import cv2\nimport numpy as np\n\n__all__ = [\"vis\"]\n\n\ndef vis(img, boxes, scores, cls_ids, conf=0.5, class_names=None):\n\n for i in range(len(boxes)):\n box = boxes[i]\n cls_id = int(cls_ids[i])\n score = scores[i]\n if score < conf:\n continue\n x0 = int(box[0])\n y0 = int(box[1])\n x1 = int(box[2])\n y1 = int(box[3])\n\n color = (_COLORS[cls_id] * 255).astype(np.uint8).tolist()\n text = '{}:{:.1f}%'.format(class_names[cls_id], score * 100)\n txt_color = (0, 0, 0) if np.mean(_COLORS[cls_id]) > 0.5 else (255, 255, 255)\n font = cv2.FONT_HERSHEY_SIMPLEX\n\n txt_size = cv2.getTextSize(text, font, 0.4, 1)[0]\n cv2.rectangle(img, (x0, y0), (x1, y1), color, 2)\n\n txt_bk_color = (_COLORS[cls_id] * 255 * 0.7).astype(np.uint8).tolist()\n cv2.rectangle(\n img,\n (x0, y0 + 1),\n (x0 + txt_size[0] + 1, y0 + int(1.5*txt_size[1])),\n txt_bk_color,\n -1\n )\n cv2.putText(img, text, (x0, y0 + txt_size[1]), font, 0.4, txt_color, thickness=1)\n\n return img\n\n\ndef get_color(idx):\n idx = idx * 3\n color = ((37 * idx) % 255, (17 * idx) % 255, (29 * idx) % 255)\n\n return color\n\n\ndef plot_tracking(image, tlwhs, obj_ids, scores=None, frame_id=0, fps=0., ids2=None):\n im = np.ascontiguousarray(np.copy(image))\n im_h, im_w = im.shape[:2]\n\n top_view = np.zeros([im_w, im_w, 3], dtype=np.uint8) + 255\n\n #text_scale = max(1, image.shape[1] / 1600.)\n #text_thickness = 2\n #line_thickness = max(1, int(image.shape[1] / 500.))\n text_scale = 2\n text_thickness = 2\n line_thickness = 3\n\n radius = max(5, int(im_w/140.))\n cv2.putText(im, 'frame: %d fps: %.2f num: %d' % (frame_id, fps, len(tlwhs)),\n (0, int(15 * text_scale)), cv2.FONT_HERSHEY_PLAIN, 2, (0, 0, 255), thickness=2)\n\n for i, tlwh in enumerate(tlwhs):\n x1, y1, w, h = tlwh\n intbox = tuple(map(int, (x1, y1, x1 + w, y1 + h)))\n obj_id = int(obj_ids[i])\n id_text = '{}'.format(int(obj_id))\n if ids2 is not None:\n id_text = id_text + ', {}'.format(int(ids2[i]))\n color = get_color(abs(obj_id))\n cv2.rectangle(im, intbox[0:2], intbox[2:4], color=color, thickness=line_thickness)\n cv2.putText(im, id_text, (intbox[0], intbox[1]), cv2.FONT_HERSHEY_PLAIN, text_scale, (0, 0, 255),\n thickness=text_thickness)\n return im\n\n\n_COLORS = np.array(\n [\n 0.000, 0.447, 0.741,\n 0.850, 0.325, 0.098,\n 0.929, 0.694, 0.125,\n 0.494, 0.184, 0.556,\n 0.466, 0.674, 0.188,\n 0.301, 0.745, 0.933,\n 0.635, 0.078, 0.184,\n 0.300, 0.300, 0.300,\n 0.600, 0.600, 0.600,\n 1.000, 0.000, 0.000,\n 1.000, 0.500, 0.000,\n 0.749, 0.749, 0.000,\n 0.000, 1.000, 0.000,\n 0.000, 0.000, 1.000,\n 0.667, 0.000, 1.000,\n 0.333, 0.333, 0.000,\n 0.333, 0.667, 0.000,\n 0.333, 1.000, 0.000,\n 0.667, 0.333, 0.000,\n 0.667, 0.667, 0.000,\n 0.667, 1.000, 0.000,\n 1.000, 0.333, 0.000,\n 1.000, 0.667, 0.000,\n 1.000, 1.000, 0.000,\n 0.000, 0.333, 0.500,\n 0.000, 0.667, 0.500,\n 0.000, 1.000, 0.500,\n 0.333, 0.000, 0.500,\n 0.333, 0.333, 0.500,\n 0.333, 0.667, 0.500,\n 0.333, 1.000, 0.500,\n 0.667, 0.000, 0.500,\n 0.667, 0.333, 0.500,\n 0.667, 0.667, 0.500,\n 0.667, 1.000, 0.500,\n 1.000, 0.000, 0.500,\n 1.000, 0.333, 0.500,\n 1.000, 0.667, 0.500,\n 1.000, 1.000, 0.500,\n 0.000, 0.333, 1.000,\n 0.000, 0.667, 1.000,\n 0.000, 1.000, 1.000,\n 0.333, 0.000, 1.000,\n 0.333, 0.333, 1.000,\n 0.333, 0.667, 1.000,\n 0.333, 1.000, 1.000,\n 0.667, 0.000, 1.000,\n 0.667, 0.333, 1.000,\n 0.667, 0.667, 1.000,\n 0.667, 1.000, 1.000,\n 1.000, 0.000, 1.000,\n 1.000, 0.333, 1.000,\n 1.000, 0.667, 1.000,\n 0.333, 0.000, 0.000,\n 0.500, 0.000, 0.000,\n 0.667, 0.000, 0.000,\n 0.833, 0.000, 0.000,\n 1.000, 0.000, 0.000,\n 0.000, 0.167, 0.000,\n 0.000, 0.333, 0.000,\n 0.000, 0.500, 0.000,\n 0.000, 0.667, 0.000,\n 0.000, 0.833, 0.000,\n 0.000, 1.000, 0.000,\n 0.000, 0.000, 0.167,\n 0.000, 0.000, 0.333,\n 0.000, 0.000, 0.500,\n 0.000, 0.000, 0.667,\n 0.000, 0.000, 0.833,\n 0.000, 0.000, 1.000,\n 0.000, 0.000, 0.000,\n 0.143, 0.143, 0.143,\n 0.286, 0.286, 0.286,\n 0.429, 0.429, 0.429,\n 0.571, 0.571, 0.571,\n 0.714, 0.714, 0.714,\n 0.857, 0.857, 0.857,\n 0.000, 0.447, 0.741,\n 0.314, 0.717, 0.741,\n 0.50, 0.5, 0\n ]\n).astype(np.float32).reshape(-1, 3)" }, { "path": "eval/GVHMR/hmr4d/utils/pylogger.py", "content": "from time import time\nimport logging\nimport torch\nfrom colorlog import ColoredFormatter\n\n\ndef sync_time():\n torch.cuda.synchronize()\n return time()\n\n\nLog = logging.getLogger()\nLog.time = time\nLog.sync_time = sync_time\n\n# Set default\nLog.setLevel(logging.INFO)\nch = logging.StreamHandler()\nch.setLevel(logging.INFO)\n# Use colorlog\nformatstring = \"[%(cyan)s%(asctime)s%(reset)s][%(log_color)s%(levelname)s%(reset)s] %(message)s\"\ndatefmt = \"%m/%d %H:%M:%S\"\nch.setFormatter(ColoredFormatter(formatstring, datefmt=datefmt))\n\nLog.addHandler(ch)\n# Log.info(\"Init-Logger\")\n\n\ndef timer(sync_cuda=False, mem=False, loop=1):\n \"\"\"\n Args:\n func: function\n sync_cuda: bool, whether to synchronize cuda\n mem: bool, whether to log memory\n \"\"\"\n\n def decorator(func):\n def wrapper(*args, **kwargs):\n if mem:\n start_mem = torch.cuda.memory_allocated() / 1024**2\n if sync_cuda:\n torch.cuda.synchronize()\n\n start = Log.time()\n for _ in range(loop):\n result = func(*args, **kwargs)\n\n if sync_cuda:\n torch.cuda.synchronize()\n if loop == 1:\n message = f\"{func.__name__} took {Log.time() - start:.3f} s.\"\n else:\n message = f\"{func.__name__} took {((Log.time() - start))/loop:.3f} s. (loop={loop})\"\n\n if mem:\n end_mem = torch.cuda.memory_allocated() / 1024**2\n end_max_mem = torch.cuda.max_memory_allocated() / 1024**2\n message += f\" Start_Mem {start_mem:.1f} Max {end_max_mem:.1f} MB\"\n Log.info(message)\n\n return result\n\n return wrapper\n\n return decorator\n\n\ndef timed(fn):\n \"\"\"example usage: timed(lambda: model(inp))\"\"\"\n start = torch.cuda.Event(enable_timing=True)\n end = torch.cuda.Event(enable_timing=True)\n start.record()\n result = fn()\n end.record()\n torch.cuda.synchronize()\n return result, start.elapsed_time(end) / 1000\n" }, { "path": "eval/GVHMR/hmr4d/utils/seq_utils.py", "content": "import torch\nimport numpy as np\n\n# def get_frame_id_list_from_mask(mask):\n# \"\"\"\n# Args:\n# mask (F,), bool.\n# Return:\n# frame_id_list: List of frame_ids.\n# \"\"\"\n# frame_id_list = []\n# i = 0\n# while i < len(mask):\n# if not mask[i]:\n# i += 1\n# else:\n# j = i\n# while j < len(mask) and mask[j]:\n# j += 1\n# frame_id_list.append(torch.arange(i, j))\n# i = j\n\n# return frame_id_list\n\n\n# From GPT\ndef get_frame_id_list_from_mask(mask):\n # batch=64, 0.13s\n \"\"\"\n Vectorized approach to get frame id list from a boolean mask.\n\n Args:\n mask (F,), bool tensor: Mask array where `True` indicates a frame to be processed.\n\n Returns:\n frame_id_list: List of torch.Tensors, each tensor containing continuous indices where mask is True.\n \"\"\"\n # Find the indices where the mask changes from False to True and vice versa\n padded_mask = torch.cat(\n [torch.tensor([False], device=mask.device), mask, torch.tensor([False], device=mask.device)]\n )\n diffs = torch.diff(padded_mask.int())\n starts = (diffs == 1).nonzero(as_tuple=False).squeeze()\n ends = (diffs == -1).nonzero(as_tuple=False).squeeze()\n if starts.numel() == 0:\n return []\n if starts.numel() == 1:\n starts = starts.reshape(-1)\n ends = ends.reshape(-1)\n\n # Create list of ranges\n frame_id_list = [torch.arange(start, end) for start, end in zip(starts, ends)]\n return frame_id_list\n\n\ndef get_batch_frame_id_lists_from_mask_BLC(masks):\n # batch=64, 0.10s\n \"\"\"\n 处理三维掩码数组，为每个批次和通道提取连续True区段的索引列表。\n\n 参数:\n masks (B, L, C), 布尔张量：每个元素代表一个掩码，True表示需要处理的帧。\n\n 返回:\n batch_frame_id_lists: 对应于每个批次和每个通道的帧id列表的嵌套列表。\n \"\"\"\n B, L, C = masks.size()\n # 在序列长度两端添加一个False\n padded_masks = torch.cat(\n [\n torch.zeros((B, 1, C), dtype=torch.bool, device=masks.device),\n masks,\n torch.zeros((B, 1, C), dtype=torch.bool, device=masks.device),\n ],\n dim=1,\n )\n # 计算差分来找到True区段的起始和结束点\n diffs = torch.diff(padded_masks.int(), dim=1)\n starts = (diffs == 1).nonzero(as_tuple=True)\n ends = (diffs == -1).nonzero(as_tuple=True)\n\n # 初始化返回列表\n batch_frame_id_lists = [[[] for _ in range(C)] for _ in range(B)]\n for b in range(B):\n for c in range(C):\n batch_start = starts[0][(starts[0] == b) & (starts[2] == c)]\n batch_end = ends[0][(ends[0] == b) & (ends[2] == c)]\n # 确保start和end都是1维张量\n batch_frame_id_lists[b][c] = [\n torch.arange(start.item(), end.item()) for start, end in zip(batch_start, batch_end)\n ]\n\n return batch_frame_id_lists\n\n\ndef get_frame_id_list_from_frame_id(frame_id):\n mask = torch.zeros(frame_id[-1] + 1, dtype=torch.bool)\n mask[frame_id] = True\n frame_id_list = get_frame_id_list_from_mask(mask)\n return frame_id_list\n\n\ndef rearrange_by_mask(x, mask):\n \"\"\"\n x (L, *)\n mask (M,), M >= L\n \"\"\"\n M = mask.size(0)\n L = x.size(0)\n if M == L:\n return x\n assert M > L\n assert mask.sum() == L\n x_rearranged = torch.zeros((M, *x.size()[1:]), dtype=x.dtype, device=x.device)\n x_rearranged[mask] = x\n return x_rearranged\n\n\ndef frame_id_to_mask(frame_id, max_len):\n mask = torch.zeros(max_len, dtype=torch.bool)\n mask[frame_id] = True\n return mask\n\n\ndef mask_to_frame_id(mask):\n frame_id = torch.where(mask)[0]\n return frame_id\n\n\ndef linear_interpolate_frame_ids(data, frame_id_list):\n data = data.clone()\n for i, invalid_frame_ids in enumerate(frame_id_list):\n # interplate between prev, next\n # if at beginning or end, use the same value\n if invalid_frame_ids[0] - 1 < 0 or invalid_frame_ids[-1] + 1 >= len(data):\n if invalid_frame_ids[0] - 1 < 0:\n data[invalid_frame_ids] = data[invalid_frame_ids[-1] + 1].clone()\n else:\n data[invalid_frame_ids] = data[invalid_frame_ids[0] - 1].clone()\n else:\n prev = data[invalid_frame_ids[0] - 1]\n next = data[invalid_frame_ids[-1] + 1]\n data[invalid_frame_ids] = (\n torch.linspace(0, 1, len(invalid_frame_ids) + 2)[1:-1][:, None] * (next - prev)[None] + prev[None]\n )\n return data\n\n\ndef linear_interpolate(data, N_middle_frames):\n \"\"\"\n Args:\n data: (2, C)\n Returns:\n data_interpolated: (1+N+1, C)\n \"\"\"\n prev = data[0]\n next = data[1]\n middle = torch.linspace(0, 1, N_middle_frames + 2)[1:-1][:, None] * (next - prev)[None] + prev[None] # (N, C)\n data_interpolated = torch.cat([data[0][None], middle, data[1][None]], dim=0) # (1+N+1, C)\n return data_interpolated\n\n\ndef find_top_k_span(mask, k=3):\n \"\"\"\n Args:\n mask: (L,)\n Return:\n topk_span: List of tuple, usage: [start, end)\n \"\"\"\n if isinstance(mask, np.ndarray):\n mask = torch.from_numpy(mask)\n if mask.sum() == 0:\n return []\n mask = mask.clone().float()\n mask = torch.cat([mask.new([0]), mask, mask.new([0])])\n diff = mask[1:] - mask[:-1]\n start = torch.where(diff == 1)[0]\n end = torch.where(diff == -1)[0]\n assert len(start) == len(end)\n span_lengths = end - start\n span_lengths, idx = span_lengths.sort(descending=True)\n start = start[idx]\n end = end[idx]\n return list(zip(start.tolist(), end.tolist()))[:k]\n" }, { "path": "eval/GVHMR/hmr4d/utils/smplx_utils.py", "content": "import torch\nimport torch.nn.functional as F\nimport numpy as np\nimport smplx\nimport pickle\nfrom smplx import SMPL, SMPLX, SMPLXLayer\nfrom hmr4d.utils.body_model import BodyModelSMPLH, BodyModelSMPLX\nfrom hmr4d.utils.body_model.smplx_lite import SmplxLiteCoco17, SmplxLiteV437Coco17, SmplxLiteSmplN24\nfrom hmr4d import PROJ_ROOT\n\n# fmt: off\nSMPLH_PARENTS = torch.tensor([-1, 0, 0, 0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 9, 9, 12, 13, 14,\n 16, 17, 18, 19, 20, 22, 23, 20, 25, 26, 20, 28, 29, 20, 31, 32, 20, 34,\n 35, 21, 37, 38, 21, 40, 41, 21, 43, 44, 21, 46, 47, 21, 49, 50])\n# fmt: on\n\n\ndef make_smplx(type=\"neu_fullpose\", **kwargs):\n if type == \"neu_fullpose\":\n model = smplx.create(\n model_path=\"inputs/models/smplx/SMPLX_NEUTRAL.npz\", use_pca=False, flat_hand_mean=True, **kwargs\n )\n elif type == \"supermotion\":\n # SuperMotion is trained on BEDLAM dataset, the smplx config is the same except only 10 betas are used\n bm_kwargs = {\n \"model_type\": \"smplx\",\n \"gender\": \"neutral\",\n \"num_pca_comps\": 12,\n \"flat_hand_mean\": False,\n }\n bm_kwargs.update(kwargs)\n model = BodyModelSMPLX(model_path=PROJ_ROOT / \"inputs/checkpoints/body_models\", **bm_kwargs)\n elif type == \"supermotion_EVAL3DPW\":\n # SuperMotion is trained on BEDLAM dataset, the smplx config is the same except only 10 betas are used\n bm_kwargs = {\n \"model_type\": \"smplx\",\n \"gender\": \"neutral\",\n \"num_pca_comps\": 12,\n \"flat_hand_mean\": True,\n }\n bm_kwargs.update(kwargs)\n model = BodyModelSMPLX(model_path=\"inputs/checkpoints/body_models\", **bm_kwargs)\n elif type == \"supermotion_coco17\":\n # Fast but only predicts 17 joints\n model = SmplxLiteCoco17()\n elif type == \"supermotion_v437coco17\":\n # Predicts 437 verts and 17 joints\n model = SmplxLiteV437Coco17()\n elif type == \"supermotion_smpl24\":\n model = SmplxLiteSmplN24()\n elif type == \"rich-smplx\":\n # https://github.com/paulchhuang/rich_toolkit/blob/main/smplx2images.py\n bm_kwargs = {\n \"model_type\": \"smplx\",\n \"gender\": kwargs.get(\"gender\", \"male\"),\n \"num_pca_comps\": 12,\n \"flat_hand_mean\": False,\n # create_expression=True, create_jaw_pose=Ture\n }\n # A /smplx folder should exist under the model_path\n model = BodyModelSMPLX(model_path=\"inputs/checkpoints/body_models\", **bm_kwargs)\n elif type == \"rich-smplh\":\n bm_kwargs = {\n \"model_type\": \"smplh\",\n \"gender\": kwargs.get(\"gender\", \"male\"),\n \"use_pca\": False,\n \"flat_hand_mean\": True,\n }\n model = BodyModelSMPLH(model_path=\"inputs/checkpoints/body_models\", **bm_kwargs)\n\n elif type in [\"smplx-circle\", \"smplx-groundlink\"]:\n # don't use hand\n bm_kwargs = {\n \"model_path\": \"inputs/checkpoints/body_models\",\n \"model_type\": \"smplx\",\n \"gender\": kwargs.get(\"gender\"),\n \"num_betas\": 16,\n \"num_expression\": 0,\n }\n model = BodyModelSMPLX(**bm_kwargs)\n\n elif type == \"smplx-motionx\":\n layer_args = {\n \"create_global_orient\": False,\n \"create_body_pose\": False,\n \"create_left_hand_pose\": False,\n \"create_right_hand_pose\": False,\n \"create_jaw_pose\": False,\n \"create_leye_pose\": False,\n \"create_reye_pose\": False,\n \"create_betas\": False,\n \"create_expression\": False,\n \"create_transl\": False,\n }\n\n bm_kwargs = {\n \"model_type\": \"smplx\",\n \"model_path\": \"inputs/checkpoints/body_models\",\n \"gender\": \"neutral\",\n \"use_pca\": False,\n \"use_face_contour\": True,\n **layer_args,\n }\n model = smplx.create(**bm_kwargs)\n\n elif type == \"smplx-samp\":\n # don't use hand\n bm_kwargs = {\n \"model_path\": \"inputs/checkpoints/body_models\",\n \"model_type\": \"smplx\",\n \"gender\": kwargs.get(\"gender\"),\n \"num_betas\": 10,\n \"num_expression\": 0,\n }\n model = BodyModelSMPLX(**bm_kwargs)\n\n elif type == \"smplx-bedlam\":\n # don't use hand\n bm_kwargs = {\n \"model_path\": \"inputs/checkpoints/body_models\",\n \"model_type\": \"smplx\",\n \"gender\": kwargs.get(\"gender\"),\n \"num_betas\": 11,\n \"num_expression\": 0,\n }\n model = BodyModelSMPLX(**bm_kwargs)\n\n elif type in [\"smplx-layer\", \"smplx-fit3d\"]:\n # Use layer\n if type == \"smplx-fit3d\":\n assert (\n kwargs.get(\"gender\") == \"neutral\"\n ), \"smplx-fit3d use neutral model: https://github.com/sminchisescu-research/imar_vision_datasets_tools/blob/e8c8f83ffac23cc36adf8ec8d0fd1c55679484ef/util/smplx_util.py#L15C34-L15C34\"\n\n bm_kwargs = {\n \"model_path\": \"inputs/checkpoints/body_models/smplx\",\n \"gender\": kwargs.get(\"gender\"),\n \"num_betas\": 10,\n \"num_expression\": 10,\n }\n model = SMPLXLayer(**bm_kwargs)\n\n elif type == \"smpl\":\n bm_kwargs = {\n \"model_path\": PROJ_ROOT / \"inputs/checkpoints/body_models\",\n \"model_type\": \"smpl\",\n \"gender\": \"neutral\",\n \"num_betas\": 10,\n \"create_body_pose\": False,\n \"create_betas\": False,\n \"create_global_orient\": False,\n \"create_transl\": False,\n }\n bm_kwargs.update(kwargs)\n # model = SMPL(**bm_kwargs)\n model = BodyModelSMPLH(**bm_kwargs)\n elif type == \"smplh\":\n bm_kwargs = {\n \"model_type\": \"smplh\",\n \"gender\": kwargs.get(\"gender\", \"male\"),\n \"use_pca\": False,\n \"flat_hand_mean\": False,\n }\n model = BodyModelSMPLH(model_path=\"inputs/checkpoints/body_models\", **bm_kwargs)\n\n else:\n raise NotImplementedError\n\n return model\n\n\ndef load_parents(npz_path=\"models/smplx/SMPLX_NEUTRAL.npz\"):\n smplx_struct = np.load(\"models/smplx/SMPLX_NEUTRAL.npz\", allow_pickle=True)\n parents = smplx_struct[\"kintree_table\"][0].astype(np.long)\n parents[0] = -1\n return parents\n\n\ndef load_smpl_faces(npz_path=\"models/smplh/SMPLH_FEMALE.pkl\"):\n with open(npz_path, \"rb\") as f:\n smpl_model = pickle.load(f, encoding=\"latin1\")\n faces = np.array(smpl_model[\"f\"].astype(np.int64))\n return faces\n\n\ndef decompose_fullpose(fullpose, model_type=\"smplx\"):\n assert model_type == \"smplx\"\n\n fullpose_dict = {\n \"global_orient\": fullpose[..., :3],\n \"body_pose\": fullpose[..., 3:66],\n \"jaw_pose\": fullpose[..., 66:69],\n \"leye_pose\": fullpose[..., 69:72],\n \"reye_pose\": fullpose[..., 72:75],\n \"left_hand_pose\": fullpose[..., 75:120],\n \"right_hand_pose\": fullpose[..., 120:165],\n }\n\n return fullpose_dict\n\n\ndef compose_fullpose(fullpose_dict, model_type=\"smplx\"):\n assert model_type == \"smplx\"\n fullpose = torch.cat(\n [\n fullpose_dict[k]\n for k in [\n \"global_orient\",\n \"body_pose\",\n \"jaw_pose\",\n \"leye_pose\",\n \"reye_pose\",\n \"left_hand_pose\",\n \"right_hand_pose\",\n ]\n ],\n dim=-1,\n )\n return fullpose\n\n\ndef compute_R_from_kinetree(rot_mats, parents):\n \"\"\"operation of lbs/batch_rigid_transform, focus on 3x3 R only\n Parameters\n ----------\n rot_mats: torch.tensor BxNx3x3\n Tensor of rotation matrices\n parents : torch.tensor BxN\n The kinematic tree of each object\n\n Returns\n -------\n R : torch.tensor BxNx3x3\n Tensor of rotation matrices\n \"\"\"\n rot_mat_chain = [rot_mats[:, 0]]\n for i in range(1, parents.shape[0]):\n curr_res = torch.matmul(rot_mat_chain[parents[i]], rot_mats[:, i])\n rot_mat_chain.append(curr_res)\n\n R = torch.stack(rot_mat_chain, dim=1)\n return R\n\n\ndef compute_relR_from_kinetree(R, parents):\n \"\"\"Inverse operation of lbs/batch_rigid_transform, focus on 3x3 R only\n Parameters\n ----------\n R : torch.tensor BxNx4x4 or BxNx3x3\n Tensor of rotation matrices\n parents : torch.tensor BxN\n The kinematic tree of each object\n\n Returns\n -------\n rot_mats: torch.tensor BxNx3x3\n Tensor of rotation matrices\n \"\"\"\n R = R[:, :, :3, :3]\n\n Rp = R[:, parents] # Rp[:, 0] is invalid\n rot_mats = Rp.transpose(2, 3) @ R\n rot_mats[:, 0] = R[:, 0]\n\n return rot_mats\n\n\ndef quat_mul(x, y):\n \"\"\"\n Performs quaternion multiplication on arrays of quaternions\n\n :param x: tensor of quaternions of shape (..., Nb of joints, 4)\n :param y: tensor of quaternions of shape (..., Nb of joints, 4)\n :return: The resulting quaternions\n \"\"\"\n x0, x1, x2, x3 = x[..., 0:1], x[..., 1:2], x[..., 2:3], x[..., 3:4]\n y0, y1, y2, y3 = y[..., 0:1], y[..., 1:2], y[..., 2:3], y[..., 3:4]\n\n # res = np.concatenate(\n # [\n # y0 * x0 - y1 * x1 - y2 * x2 - y3 * x3,\n # y0 * x1 + y1 * x0 - y2 * x3 + y3 * x2,\n # y0 * x2 + y1 * x3 + y2 * x0 - y3 * x1,\n # y0 * x3 - y1 * x2 + y2 * x1 + y3 * x0,\n # ],\n # axis=-1,\n # )\n res = torch.cat(\n [\n y0 * x0 - y1 * x1 - y2 * x2 - y3 * x3,\n y0 * x1 + y1 * x0 - y2 * x3 + y3 * x2,\n y0 * x2 + y1 * x3 + y2 * x0 - y3 * x1,\n y0 * x3 - y1 * x2 + y2 * x1 + y3 * x0,\n ],\n axis=-1,\n )\n\n return res\n\n\ndef quat_inv(q):\n \"\"\"\n Inverts a tensor of quaternions\n\n :param q: quaternion tensor\n :return: tensor of inverted quaternions\n \"\"\"\n # res = np.asarray([1, -1, -1, -1], dtype=np.float32) * q\n res = torch.tensor([1, -1, -1, -1], device=q.device).float() * q\n return res\n\n\ndef quat_mul_vec(q, x):\n \"\"\"\n Performs multiplication of an array of 3D vectors by an array of quaternions (rotation).\n\n :param q: tensor of quaternions of shape (..., Nb of joints, 4)\n :param x: tensor of vectors of shape (..., Nb of joints, 3)\n :return: the resulting array of rotated vectors\n \"\"\"\n # t = 2.0 * np.cross(q[..., 1:], x)\n t = 2.0 * torch.cross(q[..., 1:], x)\n # res = x + q[..., 0][..., np.newaxis] * t + np.cross(q[..., 1:], t)\n res = x + q[..., 0][..., None] * t + torch.cross(q[..., 1:], t)\n\n return res\n\n\ndef inverse_kinematics_motion(\n global_pos,\n global_rot,\n parents=SMPLH_PARENTS,\n):\n \"\"\"\n Args:\n global_pos : (B, T, J-1, 3)\n global_rot (q) : (B, T, J-1, 4)\n parents : SMPLH_PARENTS\n Returns:\n local_pos : (B, T, J-1, 3)\n local_rot (q) : (B, T, J-1, 4)\n \"\"\"\n J = 22\n local_pos = quat_mul_vec(\n quat_inv(global_rot[..., parents[1:J], :]),\n global_pos - global_pos[..., parents[1:J], :],\n )\n local_rot = (quat_mul(quat_inv(global_rot[..., parents[1:J], :]), global_rot),)\n return local_pos, local_rot\n\n\ndef transform_mat(R, t):\n \"\"\"Creates a batch of transformation matrices\n Args:\n - R: Bx3x3 array of a batch of rotation matrices\n - t: Bx3x1 array of a batch of translation vectors\n Returns:\n - T: Bx4x4 Transformation matrix\n \"\"\"\n # No padding left or right, only add an extra row\n return torch.cat([F.pad(R, [0, 0, 0, 1]), F.pad(t, [0, 0, 0, 1], value=1)], dim=2)\n\n\ndef normalize_joints(joints):\n \"\"\"\n Args:\n joints: (B, *, J, 3)\n \"\"\"\n LR_hips_xy = joints[..., 2, [0, 1]] - joints[..., 1, [0, 1]]\n LR_shoulders_xy = joints[..., 17, [0, 1]] - joints[..., 16, [0, 1]]\n LR_xy = (LR_hips_xy + LR_shoulders_xy) / 2 # (B, *, J, 2)\n\n x_dir = F.pad(F.normalize(LR_xy, 2, -1), (0, 1), \"constant\", 0) # (B, *, 3)\n z_dir = torch.zeros_like(x_dir) # (B, *, 3)\n z_dir[..., 2] = 1\n y_dir = torch.cross(z_dir, x_dir, dim=-1)\n\n joints_normalized = (joints - joints[..., [0], :]) @ torch.stack([x_dir, y_dir, z_dir], dim=-1)\n return joints_normalized\n\n\n@torch.no_grad()\ndef compute_Rt_af2az(joints, inverse=False):\n \"\"\"Assume z coord is upward\n Args:\n joints: (B, J, 3), in the start-frame\n Returns:\n R_af2az: (B, 3, 3)\n t_af2az: (B, 3)\n \"\"\"\n t_af2az = joints[:, 0, :].detach().clone()\n t_af2az[:, 2] = 0 # do not modify z\n\n LR_xy = joints[:, 2, [0, 1]] - joints[:, 1, [0, 1]] # (B, 2)\n I_mask = LR_xy.pow(2).sum(-1) < 1e-4 # do not rotate, when can't decided the face direction\n x_dir = F.pad(F.normalize(LR_xy, 2, -1), (0, 1), \"constant\", 0) # (B, 3)\n z_dir = torch.zeros_like(x_dir)\n z_dir[..., 2] = 1\n y_dir = torch.cross(z_dir, x_dir, dim=-1)\n R_af2az = torch.stack([x_dir, y_dir, z_dir], dim=-1) # (B, 3, 3)\n R_af2az[I_mask] = torch.eye(3).to(R_af2az)\n\n if inverse:\n R_az2af = R_af2az.transpose(1, 2)\n t_az2af = -(R_az2af @ t_af2az.unsqueeze(2)).squeeze(2)\n return R_az2af, t_az2af\n else:\n return R_af2az, t_af2az\n\n\ndef finite_difference_forward(x, dim_t=1, dup_last=True):\n if dim_t == 1:\n v = x[:, 1:] - x[:, :-1]\n if dup_last:\n v = torch.cat([v, v[:, [-1]]], dim=1)\n else:\n raise NotImplementedError\n\n return v\n\n\ndef compute_joints_zero(betas, gender):\n \"\"\"\n Args:\n betas: (16)\n gender: 'male' or 'female'\n Returns:\n joints_zero: (22, 3)\n \"\"\"\n body_model = {\n \"male\": make_smplx(type=\"humor\", gender=\"male\"),\n \"female\": make_smplx(type=\"humor\", gender=\"female\"),\n }\n\n smpl_params = {\n \"root_orient\": torch.zeros((1, 3)),\n \"pose_body\": torch.zeros((1, 63)),\n \"betas\": betas[None],\n \"trans\": torch.zeros(1, 3),\n }\n joints_zero = body_model[gender](**smpl_params).Jtr[0, :22]\n return joints_zero\n" }, { "path": "eval/GVHMR/hmr4d/utils/video_io_utils.py", "content": "import imageio.v3 as iio\nimport numpy as np\nimport torch\nfrom pathlib import Path\nimport shutil\nimport ffmpeg\nfrom tqdm import tqdm\nimport cv2\n\n\ndef get_video_lwh(video_path):\n L, H, W, _ = iio.improps(video_path, plugin=\"pyav\").shape\n return L, W, H\n\n\ndef read_video_np(video_path, start_frame=0, end_frame=-1, scale=1.0):\n \"\"\"\n Args:\n video_path: str\n Returns:\n frames: np.array, (N, H, W, 3) RGB, uint8\n \"\"\"\n # If video path not exists, an error will be raised by ffmpegs\n filter_args = []\n should_check_length = False\n\n # 1. Trim\n if not (start_frame == 0 and end_frame == -1):\n if end_frame == -1:\n filter_args.append((\"trim\", f\"start_frame={start_frame}\"))\n else:\n should_check_length = True\n filter_args.append((\"trim\", f\"start_frame={start_frame}:end_frame={end_frame}\"))\n\n # 2. Scale\n if scale != 1.0:\n filter_args.append((\"scale\", f\"iw*{scale}:ih*{scale}\"))\n\n # Excute then check\n frames = iio.imread(video_path, plugin=\"pyav\", filter_sequence=filter_args)\n if should_check_length:\n assert len(frames) == end_frame - start_frame\n\n return frames\n\n\ndef get_video_reader(video_path):\n return iio.imiter(video_path, plugin=\"pyav\")\n\n\ndef read_images_np(image_paths, verbose=False):\n \"\"\"\n Args:\n image_paths: list of str\n Returns:\n images: np.array, (N, H, W, 3) RGB, uint8\n \"\"\"\n if verbose:\n images = [cv2.imread(str(img_path))[..., ::-1] for img_path in tqdm(image_paths)]\n else:\n images = [cv2.imread(str(img_path))[..., ::-1] for img_path in image_paths]\n images = np.stack(images, axis=0)\n return images\n\n\ndef save_video(images, video_path, fps=30, crf=17):\n \"\"\"\n Args:\n images: (N, H, W, 3) RGB, uint8\n crf: 17 is visually lossless, 23 is default, +6 results in half the bitrate\n 0 is lossless, https://trac.ffmpeg.org/wiki/Encode/H.264#crf\n \"\"\"\n if isinstance(images, torch.Tensor):\n images = images.cpu().numpy().astype(np.uint8)\n elif isinstance(images, list):\n images = np.array(images).astype(np.uint8)\n\n with iio.imopen(video_path, \"w\", plugin=\"pyav\") as writer:\n writer.init_video_stream(\"libx264\", fps=fps)\n writer._video_stream.options = {\"crf\": str(crf)}\n writer.write(images)\n\n\ndef get_writer(video_path, fps=30, crf=17):\n \"\"\"remember to .close()\"\"\"\n writer = iio.imopen(video_path, \"w\", plugin=\"pyav\")\n writer.init_video_stream(\"libx264\", fps=fps)\n writer._video_stream.options = {\"crf\": str(crf)}\n return writer\n\n\ndef copy_file(video_path, out_video_path, overwrite=True):\n if not overwrite and Path(out_video_path).exists():\n return\n shutil.copy(video_path, out_video_path)\n\n\ndef merge_videos_horizontal(in_video_paths: list, out_video_path: str):\n if len(in_video_paths) < 2:\n raise ValueError(\"At least two video paths are required for merging.\")\n inputs = [ffmpeg.input(path) for path in in_video_paths]\n merged_video = ffmpeg.filter(inputs, \"hstack\", inputs=len(inputs))\n output = ffmpeg.output(merged_video, out_video_path)\n ffmpeg.run(output, overwrite_output=True, quiet=True)\n\n\ndef merge_videos_vertical(in_video_paths: list, out_video_path: str):\n if len(in_video_paths) < 2:\n raise ValueError(\"At least two video paths are required for merging.\")\n inputs = [ffmpeg.input(path) for path in in_video_paths]\n merged_video = ffmpeg.filter(inputs, \"vstack\", inputs=len(inputs))\n output = ffmpeg.output(merged_video, out_video_path)\n ffmpeg.run(output, overwrite_output=True, quiet=True)\n" }, { "path": "eval/GVHMR/hmr4d/utils/vis/README.md", "content": "## Pytorch3D Renderer\n\nExample:\n```python\nfrom hmr4d.utils.vis.renderer import Renderer\nimport imageio\n\nfps = 30\nfocal_length = data[\"cam_int\"][0][0, 0]\nwidth, height = img_hw\nfaces = smplh[data[\"gender\"]].bm.faces\nrenderer = Renderer(width, height, focal_length, \"cuda\", faces)\nwriter = imageio.get_writer(\"tmp_debug.mp4\", fps=fps, mode=\"I\", format=\"FFMPEG\", macro_block_size=1)\n\nfor i in tqdm(range(length)):\n img = np.zeros((height, width, 3), dtype=np.uint8)\n img = renderer.render_mesh(smplh_out.vertices[i].cuda(), img)\n writer.append_data(img)\nwriter.close()\n```" }, { "path": "eval/GVHMR/hmr4d/utils/vis/cv2_utils.py", "content": "import torch\nimport cv2\nimport numpy as np\nfrom hmr4d.utils.wis3d_utils import get_colors_by_conf\n\n\ndef to_numpy(x):\n if isinstance(x, np.ndarray):\n return x.copy()\n elif isinstance(x, list):\n return np.array(x)\n return x.clone().cpu().numpy()\n\n\ndef draw_bbx_xys_on_image(bbx_xys, image, conf=True):\n assert isinstance(bbx_xys, np.ndarray)\n assert isinstance(image, np.ndarray)\n image = image.copy()\n lu_point = (bbx_xys[:2] - bbx_xys[2:] / 2).astype(int)\n rd_point = (bbx_xys[:2] + bbx_xys[2:] / 2).astype(int)\n color = (255, 178, 102) if conf == True else (128, 128, 128) # orange or gray\n image = cv2.rectangle(image, lu_point, rd_point, color, 2)\n return image\n\n\ndef draw_bbx_xys_on_image_batch(bbx_xys_batch, image_batch, conf=None):\n \"\"\"conf: if provided, list of bool\"\"\"\n use_conf = conf is not None\n bbx_xys_batch = to_numpy(bbx_xys_batch)\n assert len(bbx_xys_batch) == len(image_batch)\n image_batch_out = []\n for i in range(len(bbx_xys_batch)):\n if use_conf:\n image_batch_out.append(draw_bbx_xys_on_image(bbx_xys_batch[i], image_batch[i], conf[i]))\n else:\n image_batch_out.append(draw_bbx_xys_on_image(bbx_xys_batch[i], image_batch[i]))\n return image_batch_out\n\n\ndef draw_bbx_xyxy_on_image(bbx_xys, image, conf=True):\n bbx_xys = to_numpy(bbx_xys)\n image = to_numpy(image)\n color = (255, 178, 102) if conf == True else (128, 128, 128) # orange or gray\n image = cv2.rectangle(image, (int(bbx_xys[0]), int(bbx_xys[1])), (int(bbx_xys[2]), int(bbx_xys[3])), color, 2)\n return image\n\n\ndef draw_bbx_xyxy_on_image_batch(bbx_xyxy_batch, image_batch, mask=None, conf=None):\n \"\"\"\n Args:\n conf: if provided, list of bool, mutually exclusive with mask\n mask: whether to draw, historically used\n \"\"\"\n if mask is not None:\n assert conf is None\n if conf is not None:\n assert mask is None\n use_conf = conf is not None\n bbx_xyxy_batch = to_numpy(bbx_xyxy_batch)\n image_batch = to_numpy(image_batch)\n assert len(bbx_xyxy_batch) == len(image_batch)\n image_batch_out = []\n for i in range(len(bbx_xyxy_batch)):\n if use_conf:\n image_batch_out.append(draw_bbx_xyxy_on_image(bbx_xyxy_batch[i], image_batch[i], conf[i]))\n else:\n if mask is None or mask[i]:\n image_batch_out.append(draw_bbx_xyxy_on_image(bbx_xyxy_batch[i], image_batch[i]))\n else:\n image_batch_out.append(image_batch[i])\n return image_batch_out\n\n\ndef draw_kpts(frame, keypoints, color=(0, 255, 0), thickness=2):\n frame_ = frame.copy()\n for x, y in keypoints:\n cv2.circle(frame_, (int(x), int(y)), thickness, color, -1)\n return frame_\n\n\ndef draw_kpts_with_conf(frame, kp2d, conf, thickness=2):\n \"\"\"\n Args:\n kp2d: (J, 2),\n conf: (J,)\n \"\"\"\n frame_ = frame.copy()\n conf = conf.reshape(-1)\n colors = get_colors_by_conf(conf) # (J, 3)\n colors = colors[:, [2, 1, 0]].int().numpy().tolist()\n for j in range(kp2d.shape[0]):\n x, y = kp2d[j, :2]\n c = colors[j]\n cv2.circle(frame_, (int(x), int(y)), thickness, c, -1)\n return frame_\n\n\ndef draw_kpts_with_conf_batch(frames, kp2d_batch, conf_batch, thickness=2):\n \"\"\"\n Args:\n kp2d_batch: (B, J, 2),\n conf_batch: (B, J)\n \"\"\"\n assert len(frames) == len(kp2d_batch)\n assert len(frames) == len(conf_batch)\n frames_ = []\n for i in range(len(frames)):\n frames_.append(draw_kpts_with_conf(frames[i], kp2d_batch[i], conf_batch[i], thickness))\n return frames_\n\n\ndef draw_coco17_skeleton(img, keypoints, conf_thr=0):\n use_conf_thr = True if keypoints.shape[1] == 3 else False\n img = img.copy()\n # fmt:off\n coco_skel = [[15, 13], [13, 11], [16, 14], [14, 12], [11, 12], [5, 11], [6, 12], [5, 6], [5, 7], [6, 8], [7, 9], [8, 10], [1, 2], [0, 1], [0, 2], [1, 3], [2, 4], [3, 5], [4, 6]] \n # fmt:on\n for bone in coco_skel:\n if use_conf_thr:\n kp1 = keypoints[bone[0]][:2].astype(int)\n kp2 = keypoints[bone[1]][:2].astype(int)\n kp1_c = keypoints[bone[0]][2]\n kp2_c = keypoints[bone[1]][2]\n if kp1_c > conf_thr and kp2_c > conf_thr:\n img = cv2.line(img, (kp1[0], kp1[1]), (kp2[0], kp2[1]), (0, 255, 0), 4)\n if kp1_c > conf_thr:\n img = cv2.circle(img, (kp1[0], kp1[1]), 6, (0, 255, 0), -1)\n if kp2_c > conf_thr:\n img = cv2.circle(img, (kp2[0], kp2[1]), 6, (0, 255, 0), -1)\n\n else:\n kp1 = keypoints[bone[0]][:2].astype(int)\n kp2 = keypoints[bone[1]][:2].astype(int)\n img = cv2.line(img, (kp1[0], kp1[1]), (kp2[0], kp2[1]), (0, 255, 0), 4)\n return img\n\n\ndef draw_coco17_skeleton_batch(imgs, keypoints_batch, conf_thr=0):\n assert len(imgs) == len(keypoints_batch)\n keypoints_batch = to_numpy(keypoints_batch)\n imgs_out = []\n for i in range(len(imgs)):\n imgs_out.append(draw_coco17_skeleton(imgs[i], keypoints_batch[i], conf_thr))\n return imgs_out\n" }, { "path": "eval/GVHMR/hmr4d/utils/vis/renderer.py", "content": "import cv2\nimport torch\nimport numpy as np\n\nfrom pytorch3d.renderer import (\n PerspectiveCameras,\n TexturesVertex,\n PointLights,\n Materials,\n RasterizationSettings,\n MeshRenderer,\n MeshRasterizer,\n SoftPhongShader,\n)\nfrom pytorch3d.structures import Meshes\nfrom pytorch3d.structures.meshes import join_meshes_as_scene\nfrom pytorch3d.renderer.cameras import look_at_rotation\nfrom pytorch3d.transforms import axis_angle_to_matrix\n\nfrom .renderer_tools import get_colors, checkerboard_geometry\n\n\ncolors_str_map = {\n \"gray\": [0.8, 0.8, 0.8],\n \"green\": [39, 194, 128],\n}\n\n\ndef overlay_image_onto_background(image, mask, bbox, background):\n if isinstance(image, torch.Tensor):\n image = image.detach().cpu().numpy()\n if isinstance(mask, torch.Tensor):\n mask = mask.detach().cpu().numpy()\n\n out_image = background.copy()\n bbox = bbox[0].int().cpu().numpy().copy()\n roi_image = out_image[bbox[1] : bbox[3], bbox[0] : bbox[2]]\n\n roi_image[mask] = image[mask]\n out_image[bbox[1] : bbox[3], bbox[0] : bbox[2]] = roi_image\n\n return out_image\n\n\ndef update_intrinsics_from_bbox(K_org, bbox):\n device, dtype = K_org.device, K_org.dtype\n\n K = torch.zeros((K_org.shape[0], 4, 4)).to(device=device, dtype=dtype)\n K[:, :3, :3] = K_org.clone()\n K[:, 2, 2] = 0\n K[:, 2, -1] = 1\n K[:, -1, 2] = 1\n\n image_sizes = []\n for idx, bbox in enumerate(bbox):\n left, upper, right, lower = bbox\n cx, cy = K[idx, 0, 2], K[idx, 1, 2]\n\n new_cx = cx - left\n new_cy = cy - upper\n new_height = max(lower - upper, 1)\n new_width = max(right - left, 1)\n new_cx = new_width - new_cx\n new_cy = new_height - new_cy\n\n K[idx, 0, 2] = new_cx\n K[idx, 1, 2] = new_cy\n image_sizes.append((int(new_height), int(new_width)))\n\n return K, image_sizes\n\n\ndef perspective_projection(x3d, K, R=None, T=None):\n if R != None:\n x3d = torch.matmul(R, x3d.transpose(1, 2)).transpose(1, 2)\n if T != None:\n x3d = x3d + T.transpose(1, 2)\n\n x2d = torch.div(x3d, x3d[..., 2:])\n x2d = torch.matmul(K, x2d.transpose(-1, -2)).transpose(-1, -2)[..., :2]\n return x2d\n\n\ndef compute_bbox_from_points(X, img_w, img_h, scaleFactor=1.2):\n left = torch.clamp(X.min(1)[0][:, 0], min=0, max=img_w)\n right = torch.clamp(X.max(1)[0][:, 0], min=0, max=img_w)\n top = torch.clamp(X.min(1)[0][:, 1], min=0, max=img_h)\n bottom = torch.clamp(X.max(1)[0][:, 1], min=0, max=img_h)\n\n cx = (left + right) / 2\n cy = (top + bottom) / 2\n width = right - left\n height = bottom - top\n\n new_left = torch.clamp(cx - width / 2 * scaleFactor, min=0, max=img_w - 1)\n new_right = torch.clamp(cx + width / 2 * scaleFactor, min=1, max=img_w)\n new_top = torch.clamp(cy - height / 2 * scaleFactor, min=0, max=img_h - 1)\n new_bottom = torch.clamp(cy + height / 2 * scaleFactor, min=1, max=img_h)\n\n bbox = torch.stack((new_left.detach(), new_top.detach(), new_right.detach(), new_bottom.detach())).int().float().T\n\n return bbox\n\n\nclass Renderer:\n def __init__(self, width, height, focal_length=None, device=\"cuda\", faces=None, K=None, bin_size=None):\n \"\"\"set bin_size to 0 for no binning\"\"\"\n self.width = width\n self.height = height\n self.bin_size = bin_size\n assert (focal_length is not None) ^ (K is not None), \"focal_length and K are mutually exclusive\"\n\n self.device = device\n if faces is not None:\n if isinstance(faces, np.ndarray):\n faces = torch.from_numpy((faces).astype(\"int\"))\n self.faces = faces.unsqueeze(0).to(self.device)\n\n self.initialize_camera_params(focal_length, K)\n self.lights = PointLights(device=device, location=[[0.0, 0.0, -10.0]])\n self.create_renderer()\n\n def create_renderer(self):\n self.renderer = MeshRenderer(\n rasterizer=MeshRasterizer(\n raster_settings=RasterizationSettings(\n image_size=self.image_sizes[0], blur_radius=1e-5, bin_size=self.bin_size\n ),\n ),\n shader=SoftPhongShader(\n device=self.device,\n lights=self.lights,\n ),\n )\n\n def create_camera(self, R=None, T=None):\n if R is not None:\n self.R = R.clone().view(1, 3, 3).to(self.device)\n if T is not None:\n self.T = T.clone().view(1, 3).to(self.device)\n\n return PerspectiveCameras(\n device=self.device, R=self.R.mT, T=self.T, K=self.K_full, image_size=self.image_sizes, in_ndc=False\n )\n\n def initialize_camera_params(self, focal_length, K):\n # Extrinsics\n self.R = torch.diag(torch.tensor([1, 1, 1])).float().to(self.device).unsqueeze(0)\n\n self.T = torch.tensor([0, 0, 0]).unsqueeze(0).float().to(self.device)\n\n # Intrinsics\n if K is not None:\n self.K = K.float().reshape(1, 3, 3).to(self.device)\n else:\n assert focal_length is not None, \"focal_length or K should be provided\"\n self.K = (\n torch.tensor([[focal_length, 0, self.width / 2], [0, focal_length, self.height / 2], [0, 0, 1]])\n .float()\n .reshape(1, 3, 3)\n .to(self.device)\n )\n self.bboxes = torch.tensor([[0, 0, self.width, self.height]]).float()\n self.K_full, self.image_sizes = update_intrinsics_from_bbox(self.K, self.bboxes)\n self.cameras = self.create_camera()\n\n def set_intrinsic(self, K):\n self.K = K.reshape(1, 3, 3)\n\n def set_ground(self, length, center_x, center_z):\n device = self.device\n length, center_x, center_z = map(float, (length, center_x, center_z))\n v, f, vc, fc = map(torch.from_numpy, checkerboard_geometry(length=length, c1=center_x, c2=center_z, up=\"y\"))\n v, f, vc = v.to(device), f.to(device), vc.to(device)\n self.ground_geometry = [v, f, vc]\n\n def update_bbox(self, x3d, scale=2.0, mask=None):\n \"\"\"Update bbox of cameras from the given 3d points\n\n x3d: input 3D keypoints (or vertices), (num_frames, num_points, 3)\n \"\"\"\n\n if x3d.size(-1) != 3:\n x2d = x3d.unsqueeze(0)\n else:\n x2d = perspective_projection(x3d.unsqueeze(0), self.K, self.R, self.T.reshape(1, 3, 1))\n\n if mask is not None:\n x2d = x2d[:, ~mask]\n\n bbox = compute_bbox_from_points(x2d, self.width, self.height, scale)\n self.bboxes = bbox\n\n self.K_full, self.image_sizes = update_intrinsics_from_bbox(self.K, bbox)\n self.cameras = self.create_camera()\n self.create_renderer()\n\n def reset_bbox(\n self,\n ):\n bbox = torch.zeros((1, 4)).float().to(self.device)\n bbox[0, 2] = self.width\n bbox[0, 3] = self.height\n self.bboxes = bbox\n\n self.K_full, self.image_sizes = update_intrinsics_from_bbox(self.K, bbox)\n self.cameras = self.create_camera()\n self.create_renderer()\n\n def render_mesh(self, vertices, background=None, colors=[0.8, 0.8, 0.8], VI=50):\n self.update_bbox(vertices[::VI], scale=1.2)\n vertices = vertices.unsqueeze(0)\n\n if isinstance(colors, torch.Tensor):\n # per-vertex color\n verts_features = colors.to(device=vertices.device, dtype=vertices.dtype)\n colors = [0.8, 0.8, 0.8]\n else:\n if colors[0] > 1:\n colors = [c / 255.0 for c in colors]\n verts_features = torch.tensor(colors).reshape(1, 1, 3).to(device=vertices.device, dtype=vertices.dtype)\n verts_features = verts_features.repeat(1, vertices.shape[1], 1)\n textures = TexturesVertex(verts_features=verts_features)\n\n mesh = Meshes(\n verts=vertices,\n faces=self.faces,\n textures=textures,\n )\n\n materials = Materials(device=self.device, specular_color=(colors,), shininess=0)\n\n results = torch.flip(self.renderer(mesh, materials=materials, cameras=self.cameras, lights=self.lights), [1, 2])\n image = results[0, ..., :3] * 255\n mask = results[0, ..., -1] > 1e-3\n\n if background is None:\n background = np.ones((self.height, self.width, 3)).astype(np.uint8) * 255\n\n image = overlay_image_onto_background(image, mask, self.bboxes, background.copy())\n self.reset_bbox()\n return image\n\n def render_with_ground(self, verts, colors, cameras, lights, faces=None):\n \"\"\"\n :param verts (N, V, 3), potential multiple people\n :param colors (N, 3) or (N, V, 3)\n :param faces (N, F, 3), optional, otherwise self.faces is used will be used\n \"\"\"\n # Sanity check of input verts, colors and faces: (B, V, 3), (B, F, 3), (B, V, 3)\n N, V, _ = verts.shape\n if faces is None:\n faces = self.faces.clone().expand(N, -1, -1)\n else:\n assert len(faces.shape) == 3, \"faces should have shape of (N, F, 3)\"\n\n assert len(colors.shape) in [2, 3]\n if len(colors.shape) == 2:\n assert len(colors) == N, \"colors of shape 2 should be (N, 3)\"\n colors = colors[:, None]\n colors = colors.expand(N, V, -1)[..., :3]\n\n # (V, 3), (F, 3), (V, 3)\n gv, gf, gc = self.ground_geometry\n verts = list(torch.unbind(verts, dim=0)) + [gv]\n faces = list(torch.unbind(faces, dim=0)) + [gf]\n colors = list(torch.unbind(colors, dim=0)) + [gc[..., :3]]\n mesh = create_meshes(verts, faces, colors)\n\n materials = Materials(device=self.device, shininess=0)\n\n results = self.renderer(mesh, cameras=cameras, lights=lights, materials=materials)\n image = (results[0, ..., :3].cpu().numpy() * 255).astype(np.uint8)\n\n return image\n\n\ndef create_meshes(verts, faces, colors):\n \"\"\"\n :param verts (B, V, 3)\n :param faces (B, F, 3)\n :param colors (B, V, 3)\n \"\"\"\n textures = TexturesVertex(verts_features=colors)\n meshes = Meshes(verts=verts, faces=faces, textures=textures)\n return join_meshes_as_scene(meshes)\n\n\ndef get_global_cameras(verts, device=\"cuda\", distance=5, position=(-5.0, 5.0, 0.0)):\n \"\"\"This always put object at the center of view\"\"\"\n positions = torch.tensor([position]).repeat(len(verts), 1)\n targets = verts.mean(1)\n\n directions = targets - positions\n directions = directions / torch.norm(directions, dim=-1).unsqueeze(-1) * distance\n positions = targets - directions\n\n rotation = look_at_rotation(positions, targets).mT\n translation = -(rotation @ positions.unsqueeze(-1)).squeeze(-1)\n\n lights = PointLights(device=device, location=[position])\n return rotation, translation, lights\n\n\ndef get_global_cameras_static(\n verts, beta=4.0, cam_height_degree=30, target_center_height=1.0, use_long_axis=False, vec_rot=45, device=\"cuda\"\n):\n L, V, _ = verts.shape\n\n # Compute target trajectory, denote as center + scale\n targets = verts.mean(1) # (L, 3)\n targets[:, 1] = 0 # project to xz-plane\n target_center = targets.mean(0) # (3,)\n target_scale, target_idx = torch.norm(targets - target_center, dim=-1).max(0)\n\n # a 45 degree vec from longest axis\n if use_long_axis:\n long_vec = targets[target_idx] - target_center # (x, 0, z)\n long_vec = long_vec / torch.norm(long_vec)\n R = axis_angle_to_matrix(torch.tensor([0, np.pi / 4, 0])).to(long_vec)\n vec = R @ long_vec\n else:\n vec_rad = vec_rot / 180 * np.pi\n vec = torch.tensor([np.sin(vec_rad), 0, np.cos(vec_rad)]).float()\n vec = vec / torch.norm(vec)\n\n # Compute camera position (center + scale * vec * beta) + y=4\n target_scale = max(target_scale, 1.0) * beta\n position = target_center + vec * target_scale\n position[1] = target_scale * np.tan(np.pi * cam_height_degree / 180) + target_center_height\n\n # Compute camera rotation and translation\n positions = position.unsqueeze(0).repeat(L, 1)\n target_centers = target_center.unsqueeze(0).repeat(L, 1)\n target_centers[:, 1] = target_center_height\n rotation = look_at_rotation(positions, target_centers).mT\n translation = -(rotation @ positions.unsqueeze(-1)).squeeze(-1)\n\n lights = PointLights(device=device, location=[position.tolist()])\n return rotation, translation, lights\n\n\ndef get_ground_params_from_points(root_points, vert_points):\n \"\"\"xz-plane is the ground plane\n Args:\n root_points: (L, 3), to decide center\n vert_points: (L, V, 3), to decide scale\n \"\"\"\n root_max = root_points.max(0)[0] # (3,)\n root_min = root_points.min(0)[0] # (3,)\n cx, _, cz = (root_max + root_min) / 2.0\n\n vert_max = vert_points.reshape(-1, 3).max(0)[0] # (L, 3)\n vert_min = vert_points.reshape(-1, 3).min(0)[0] # (L, 3)\n scale = (vert_max - vert_min)[[0, 2]].max()\n return float(scale), float(cx), float(cz)\n" }, { "path": "eval/GVHMR/hmr4d/utils/vis/renderer_tools.py", "content": "import os\nimport cv2\nimport numpy as np\nimport torch\nfrom PIL import Image\n\n\ndef read_image(path, scale=1):\n im = Image.open(path)\n if scale == 1:\n return np.array(im)\n W, H = im.size\n w, h = int(scale * W), int(scale * H)\n return np.array(im.resize((w, h), Image.ANTIALIAS))\n\n\ndef transform_torch3d(T_c2w):\n \"\"\"\n :param T_c2w (*, 4, 4)\n returns (*, 3, 3), (*, 3)\n \"\"\"\n R1 = torch.tensor(\n [\n [-1.0, 0.0, 0.0],\n [0.0, -1.0, 0.0],\n [0.0, 0.0, 1.0],\n ],\n device=T_c2w.device,\n )\n R2 = torch.tensor(\n [\n [1.0, 0.0, 0.0],\n [0.0, -1.0, 0.0],\n [0.0, 0.0, -1.0],\n ],\n device=T_c2w.device,\n )\n cam_R, cam_t = T_c2w[..., :3, :3], T_c2w[..., :3, 3]\n cam_R = torch.einsum(\"...ij,jk->...ik\", cam_R, R1)\n cam_t = torch.einsum(\"ij,...j->...i\", R2, cam_t)\n return cam_R, cam_t\n\n\ndef transform_pyrender(T_c2w):\n \"\"\"\n :param T_c2w (*, 4, 4)\n \"\"\"\n T_vis = torch.tensor(\n [\n [1.0, 0.0, 0.0, 0.0],\n [0.0, -1.0, 0.0, 0.0],\n [0.0, 0.0, -1.0, 0.0],\n [0.0, 0.0, 0.0, 1.0],\n ],\n device=T_c2w.device,\n )\n return torch.einsum(\"...ij,jk->...ik\", torch.einsum(\"ij,...jk->...ik\", T_vis, T_c2w), T_vis)\n\n\ndef smpl_to_geometry(verts, faces, vis_mask=None, track_ids=None):\n \"\"\"\n :param verts (B, T, V, 3)\n :param faces (F, 3)\n :param vis_mask (optional) (B, T) visibility of each person\n :param track_ids (optional) (B,)\n returns list of T verts (B, V, 3), faces (F, 3), colors (B, 3)\n where B is different depending on the visibility of the people\n \"\"\"\n B, T = verts.shape[:2]\n device = verts.device\n\n # (B, 3)\n colors = track_to_colors(track_ids) if track_ids is not None else torch.ones(B, 3, device) * 0.5\n\n # list T (B, V, 3), T (B, 3), T (F, 3)\n return filter_visible_meshes(verts, colors, faces, vis_mask)\n\n\ndef filter_visible_meshes(verts, colors, faces, vis_mask=None, vis_opacity=False):\n \"\"\"\n :param verts (B, T, V, 3)\n :param colors (B, 3)\n :param faces (F, 3)\n :param vis_mask (optional tensor, default None) (B, T) ternary mask\n -1 if not in frame\n 0 if temporarily occluded\n 1 if visible\n :param vis_opacity (optional bool, default False)\n if True, make occluded people alpha=0.5, otherwise alpha=1\n returns a list of T lists verts (Bi, V, 3), colors (Bi, 4), faces (F, 3)\n \"\"\"\n # import ipdb; ipdb.set_trace()\n B, T = verts.shape[:2]\n faces = [faces for t in range(T)]\n if vis_mask is None:\n verts = [verts[:, t] for t in range(T)]\n colors = [colors for t in range(T)]\n return verts, colors, faces\n\n # render occluded and visible, but not removed\n vis_mask = vis_mask >= 0\n if vis_opacity:\n alpha = 0.5 * (vis_mask[..., None] + 1)\n else:\n alpha = (vis_mask[..., None] >= 0).float()\n vert_list = [verts[vis_mask[:, t], t] for t in range(T)]\n colors = [torch.cat([colors[vis_mask[:, t]], alpha[vis_mask[:, t], t]], dim=-1) for t in range(T)]\n bounds = get_bboxes(verts, vis_mask)\n return vert_list, colors, faces, bounds\n\n\ndef get_bboxes(verts, vis_mask):\n \"\"\"\n return bb_min, bb_max, and mean for each track (B, 3) over entire trajectory\n :param verts (B, T, V, 3)\n :param vis_mask (B, T)\n \"\"\"\n B, T, *_ = verts.shape\n bb_min, bb_max, mean = [], [], []\n for b in range(B):\n v = verts[b, vis_mask[b, :T]] # (Tb, V, 3)\n bb_min.append(v.amin(dim=(0, 1)))\n bb_max.append(v.amax(dim=(0, 1)))\n mean.append(v.mean(dim=(0, 1)))\n bb_min = torch.stack(bb_min, dim=0)\n bb_max = torch.stack(bb_max, dim=0)\n mean = torch.stack(mean, dim=0)\n # point to a track that's long and close to the camera\n zs = mean[:, 2]\n counts = vis_mask[:, :T].sum(dim=-1) # (B,)\n mask = counts < 0.8 * T\n zs[mask] = torch.inf\n sel = torch.argmin(zs)\n return bb_min.amin(dim=0), bb_max.amax(dim=0), mean[sel]\n\n\ndef track_to_colors(track_ids):\n \"\"\"\n :param track_ids (B)\n \"\"\"\n color_map = torch.from_numpy(get_colors()).to(track_ids)\n return color_map[track_ids] / 255 # (B, 3)\n\n\ndef get_colors():\n # color_file = os.path.abspath(os.path.join(__file__, \"../colors_phalp.txt\"))\n color_file = os.path.abspath(os.path.join(__file__, \"../colors.txt\"))\n RGB_tuples = np.vstack(\n [\n np.loadtxt(color_file, skiprows=0),\n # np.loadtxt(color_file, skiprows=1),\n np.random.uniform(0, 255, size=(10000, 3)),\n [[0, 0, 0]],\n ]\n )\n b = np.where(RGB_tuples == 0)\n RGB_tuples[b] = 1\n return RGB_tuples.astype(np.float32)\n\n\ndef checkerboard_geometry(\n length=12.0,\n color0=[0.8, 0.9, 0.9],\n color1=[0.6, 0.7, 0.7],\n tile_width=0.5,\n alpha=1.0,\n up=\"y\",\n c1=0.0,\n c2=0.0,\n):\n assert up == \"y\" or up == \"z\"\n color0 = np.array(color0 + [alpha])\n color1 = np.array(color1 + [alpha])\n num_rows = num_cols = max(2, int(length / tile_width))\n radius = float(num_rows * tile_width) / 2.0\n vertices = []\n vert_colors = []\n faces = []\n face_colors = []\n for i in range(num_rows):\n for j in range(num_cols):\n u0, v0 = j * tile_width - radius, i * tile_width - radius\n us = np.array([u0, u0, u0 + tile_width, u0 + tile_width])\n vs = np.array([v0, v0 + tile_width, v0 + tile_width, v0])\n zs = np.zeros(4)\n if up == \"y\":\n cur_verts = np.stack([us, zs, vs], axis=-1) # (4, 3)\n cur_verts[:, 0] += c1\n cur_verts[:, 2] += c2\n else:\n cur_verts = np.stack([us, vs, zs], axis=-1) # (4, 3)\n cur_verts[:, 0] += c1\n cur_verts[:, 1] += c2\n\n cur_faces = np.array([[0, 1, 3], [1, 2, 3], [0, 3, 1], [1, 3, 2]], dtype=np.int64)\n cur_faces += 4 * (i * num_cols + j) # the number of previously added verts\n use_color0 = (i % 2 == 0 and j % 2 == 0) or (i % 2 == 1 and j % 2 == 1)\n cur_color = color0 if use_color0 else color1\n cur_colors = np.array([cur_color, cur_color, cur_color, cur_color])\n\n vertices.append(cur_verts)\n faces.append(cur_faces)\n vert_colors.append(cur_colors)\n face_colors.append(cur_colors)\n\n vertices = np.concatenate(vertices, axis=0).astype(np.float32)\n vert_colors = np.concatenate(vert_colors, axis=0).astype(np.float32)\n faces = np.concatenate(faces, axis=0).astype(np.float32)\n face_colors = np.concatenate(face_colors, axis=0).astype(np.float32)\n\n return vertices, faces, vert_colors, face_colors\n\n\ndef camera_marker_geometry(radius, height, up):\n assert up == \"y\" or up == \"z\"\n if up == \"y\":\n vertices = np.array(\n [\n [-radius, -radius, 0],\n [radius, -radius, 0],\n [radius, radius, 0],\n [-radius, radius, 0],\n [0, 0, height],\n ]\n )\n else:\n vertices = np.array(\n [\n [-radius, 0, -radius],\n [radius, 0, -radius],\n [radius, 0, radius],\n [-radius, 0, radius],\n [0, -height, 0],\n ]\n )\n\n faces = np.array(\n [\n [0, 3, 1],\n [1, 3, 2],\n [0, 1, 4],\n [1, 2, 4],\n [2, 3, 4],\n [3, 0, 4],\n ]\n )\n\n face_colors = np.array(\n [\n [1.0, 1.0, 1.0, 1.0],\n [1.0, 1.0, 1.0, 1.0],\n [0.0, 1.0, 0.0, 1.0],\n [1.0, 0.0, 0.0, 1.0],\n [0.0, 1.0, 0.0, 1.0],\n [1.0, 0.0, 0.0, 1.0],\n ]\n )\n return vertices, faces, face_colors\n\n\ndef vis_keypoints(\n keypts_list,\n img_size,\n radius=6,\n thickness=3,\n kpt_score_thr=0.3,\n dataset=\"TopDownCocoDataset\",\n):\n \"\"\"\n Visualize keypoints\n From ViTPose/mmpose/apis/inference.py\n \"\"\"\n palette = np.array(\n [\n [255, 128, 0],\n [255, 153, 51],\n [255, 178, 102],\n [230, 230, 0],\n [255, 153, 255],\n [153, 204, 255],\n [255, 102, 255],\n [255, 51, 255],\n [102, 178, 255],\n [51, 153, 255],\n [255, 153, 153],\n [255, 102, 102],\n [255, 51, 51],\n [153, 255, 153],\n [102, 255, 102],\n [51, 255, 51],\n [0, 255, 0],\n [0, 0, 255],\n [255, 0, 0],\n [255, 255, 255],\n ]\n )\n\n if dataset in (\n \"TopDownCocoDataset\",\n \"BottomUpCocoDataset\",\n \"TopDownOCHumanDataset\",\n \"AnimalMacaqueDataset\",\n ):\n # show the results\n skeleton = [\n [15, 13],\n [13, 11],\n [16, 14],\n [14, 12],\n [11, 12],\n [5, 11],\n [6, 12],\n [5, 6],\n [5, 7],\n [6, 8],\n [7, 9],\n [8, 10],\n [1, 2],\n [0, 1],\n [0, 2],\n [1, 3],\n [2, 4],\n [3, 5],\n [4, 6],\n ]\n\n pose_link_color = palette[[0, 0, 0, 0, 7, 7, 7, 9, 9, 9, 9, 9, 16, 16, 16, 16, 16, 16, 16]]\n pose_kpt_color = palette[[16, 16, 16, 16, 16, 9, 9, 9, 9, 9, 9, 0, 0, 0, 0, 0, 0]]\n\n elif dataset == \"TopDownCocoWholeBodyDataset\":\n # show the results\n skeleton = [\n [15, 13],\n [13, 11],\n [16, 14],\n [14, 12],\n [11, 12],\n [5, 11],\n [6, 12],\n [5, 6],\n [5, 7],\n [6, 8],\n [7, 9],\n [8, 10],\n [1, 2],\n [0, 1],\n [0, 2],\n [1, 3],\n [2, 4],\n [3, 5],\n [4, 6],\n [15, 17],\n [15, 18],\n [15, 19],\n [16, 20],\n [16, 21],\n [16, 22],\n [91, 92],\n [92, 93],\n [93, 94],\n [94, 95],\n [91, 96],\n [96, 97],\n [97, 98],\n [98, 99],\n [91, 100],\n [100, 101],\n [101, 102],\n [102, 103],\n [91, 104],\n [104, 105],\n [105, 106],\n [106, 107],\n [91, 108],\n [108, 109],\n [109, 110],\n [110, 111],\n [112, 113],\n [113, 114],\n [114, 115],\n [115, 116],\n [112, 117],\n [117, 118],\n [118, 119],\n [119, 120],\n [112, 121],\n [121, 122],\n [122, 123],\n [123, 124],\n [112, 125],\n [125, 126],\n [126, 127],\n [127, 128],\n [112, 129],\n [129, 130],\n [130, 131],\n [131, 132],\n ]\n\n pose_link_color = palette[\n [0, 0, 0, 0, 7, 7, 7, 9, 9, 9, 9, 9, 16, 16, 16, 16, 16, 16, 16]\n + [16, 16, 16, 16, 16, 16]\n + [0, 0, 0, 0, 4, 4, 4, 4, 8, 8, 8, 8, 12, 12, 12, 12, 16, 16, 16, 16]\n + [0, 0, 0, 0, 4, 4, 4, 4, 8, 8, 8, 8, 12, 12, 12, 12, 16, 16, 16, 16]\n ]\n pose_kpt_color = palette[\n [16, 16, 16, 16, 16, 9, 9, 9, 9, 9, 9, 0, 0, 0, 0, 0, 0] + [0, 0, 0, 0, 0, 0] + [19] * (68 + 42)\n ]\n\n elif dataset == \"TopDownAicDataset\":\n skeleton = [\n [2, 1],\n [1, 0],\n [0, 13],\n [13, 3],\n [3, 4],\n [4, 5],\n [8, 7],\n [7, 6],\n [6, 9],\n [9, 10],\n [10, 11],\n [12, 13],\n [0, 6],\n [3, 9],\n ]\n\n pose_link_color = palette[[9, 9, 9, 9, 9, 9, 16, 16, 16, 16, 16, 0, 7, 7]]\n pose_kpt_color = palette[[9, 9, 9, 9, 9, 9, 16, 16, 16, 16, 16, 16, 0, 0]]\n\n elif dataset == \"TopDownMpiiDataset\":\n skeleton = [\n [0, 1],\n [1, 2],\n [2, 6],\n [6, 3],\n [3, 4],\n [4, 5],\n [6, 7],\n [7, 8],\n [8, 9],\n [8, 12],\n [12, 11],\n [11, 10],\n [8, 13],\n [13, 14],\n [14, 15],\n ]\n\n pose_link_color = palette[[16, 16, 16, 16, 16, 16, 7, 7, 0, 9, 9, 9, 9, 9, 9]]\n pose_kpt_color = palette[[16, 16, 16, 16, 16, 16, 7, 7, 0, 0, 9, 9, 9, 9, 9, 9]]\n\n elif dataset == \"TopDownMpiiTrbDataset\":\n skeleton = [\n [12, 13],\n [13, 0],\n [13, 1],\n [0, 2],\n [1, 3],\n [2, 4],\n [3, 5],\n [0, 6],\n [1, 7],\n [6, 7],\n [6, 8],\n [7, 9],\n [8, 10],\n [9, 11],\n [14, 15],\n [16, 17],\n [18, 19],\n [20, 21],\n [22, 23],\n [24, 25],\n [26, 27],\n [28, 29],\n [30, 31],\n [32, 33],\n [34, 35],\n [36, 37],\n [38, 39],\n ]\n\n pose_link_color = palette[[16] * 14 + [19] * 13]\n pose_kpt_color = palette[[16] * 14 + [0] * 26]\n\n elif dataset in (\"OneHand10KDataset\", \"FreiHandDataset\", \"PanopticDataset\"):\n skeleton = [\n [0, 1],\n [1, 2],\n [2, 3],\n [3, 4],\n [0, 5],\n [5, 6],\n [6, 7],\n [7, 8],\n [0, 9],\n [9, 10],\n [10, 11],\n [11, 12],\n [0, 13],\n [13, 14],\n [14, 15],\n [15, 16],\n [0, 17],\n [17, 18],\n [18, 19],\n [19, 20],\n ]\n\n pose_link_color = palette[[0, 0, 0, 0, 4, 4, 4, 4, 8, 8, 8, 8, 12, 12, 12, 12, 16, 16, 16, 16]]\n pose_kpt_color = palette[[0, 0, 0, 0, 0, 4, 4, 4, 4, 8, 8, 8, 8, 12, 12, 12, 12, 16, 16, 16, 16]]\n\n elif dataset == \"InterHand2DDataset\":\n skeleton = [\n [0, 1],\n [1, 2],\n [2, 3],\n [4, 5],\n [5, 6],\n [6, 7],\n [8, 9],\n [9, 10],\n [10, 11],\n [12, 13],\n [13, 14],\n [14, 15],\n [16, 17],\n [17, 18],\n [18, 19],\n [3, 20],\n [7, 20],\n [11, 20],\n [15, 20],\n [19, 20],\n ]\n\n pose_link_color = palette[[0, 0, 0, 4, 4, 4, 8, 8, 8, 12, 12, 12, 16, 16, 16, 0, 4, 8, 12, 16]]\n pose_kpt_color = palette[[0, 0, 0, 0, 4, 4, 4, 4, 8, 8, 8, 8, 12, 12, 12, 12, 16, 16, 16, 16, 0]]\n\n elif dataset == \"Face300WDataset\":\n # show the results\n skeleton = []\n\n pose_link_color = palette[[]]\n pose_kpt_color = palette[[19] * 68]\n kpt_score_thr = 0\n\n elif dataset == \"FaceAFLWDataset\":\n # show the results\n skeleton = []\n\n pose_link_color = palette[[]]\n pose_kpt_color = palette[[19] * 19]\n kpt_score_thr = 0\n\n elif dataset == \"FaceCOFWDataset\":\n # show the results\n skeleton = []\n\n pose_link_color = palette[[]]\n pose_kpt_color = palette[[19] * 29]\n kpt_score_thr = 0\n\n elif dataset == \"FaceWFLWDataset\":\n # show the results\n skeleton = []\n\n pose_link_color = palette[[]]\n pose_kpt_color = palette[[19] * 98]\n kpt_score_thr = 0\n\n elif dataset == \"AnimalHorse10Dataset\":\n skeleton = [\n [0, 1],\n [1, 12],\n [12, 16],\n [16, 21],\n [21, 17],\n [17, 11],\n [11, 10],\n [10, 8],\n [8, 9],\n [9, 12],\n [2, 3],\n [3, 4],\n [5, 6],\n [6, 7],\n [13, 14],\n [14, 15],\n [18, 19],\n [19, 20],\n ]\n\n pose_link_color = palette[[4] * 10 + [6] * 2 + [6] * 2 + [7] * 2 + [7] * 2]\n pose_kpt_color = palette[[4, 4, 6, 6, 6, 6, 6, 6, 4, 4, 4, 4, 4, 7, 7, 7, 4, 4, 7, 7, 7, 4]]\n\n elif dataset == \"AnimalFlyDataset\":\n skeleton = [\n [1, 0],\n [2, 0],\n [3, 0],\n [4, 3],\n [5, 4],\n [7, 6],\n [8, 7],\n [9, 8],\n [11, 10],\n [12, 11],\n [13, 12],\n [15, 14],\n [16, 15],\n [17, 16],\n [19, 18],\n [20, 19],\n [21, 20],\n [23, 22],\n [24, 23],\n [25, 24],\n [27, 26],\n [28, 27],\n [29, 28],\n [30, 3],\n [31, 3],\n ]\n\n pose_link_color = palette[[0] * 25]\n pose_kpt_color = palette[[0] * 32]\n\n elif dataset == \"AnimalLocustDataset\":\n skeleton = [\n [1, 0],\n [2, 1],\n [3, 2],\n [4, 3],\n [6, 5],\n [7, 6],\n [9, 8],\n [10, 9],\n [11, 10],\n [13, 12],\n [14, 13],\n [15, 14],\n [17, 16],\n [18, 17],\n [19, 18],\n [21, 20],\n [22, 21],\n [24, 23],\n [25, 24],\n [26, 25],\n [28, 27],\n [29, 28],\n [30, 29],\n [32, 31],\n [33, 32],\n [34, 33],\n ]\n\n pose_link_color = palette[[0] * 26]\n pose_kpt_color = palette[[0] * 35]\n\n elif dataset == \"AnimalZebraDataset\":\n skeleton = [[1, 0], [2, 1], [3, 2], [4, 2], [5, 7], [6, 7], [7, 2], [8, 7]]\n\n pose_link_color = palette[[0] * 8]\n pose_kpt_color = palette[[0] * 9]\n\n elif dataset in \"AnimalPoseDataset\":\n skeleton = [\n [0, 1],\n [0, 2],\n [1, 3],\n [0, 4],\n [1, 4],\n [4, 5],\n [5, 7],\n [6, 7],\n [5, 8],\n [8, 12],\n [12, 16],\n [5, 9],\n [9, 13],\n [13, 17],\n [6, 10],\n [10, 14],\n [14, 18],\n [6, 11],\n [11, 15],\n [15, 19],\n ]\n\n pose_link_color = palette[[0] * 20]\n pose_kpt_color = palette[[0] * 20]\n else:\n NotImplementedError()\n\n img_w, img_h = img_size\n img = 255 * np.ones((img_h, img_w, 3), dtype=np.uint8)\n img = imshow_keypoints(\n img,\n keypts_list,\n skeleton,\n kpt_score_thr,\n pose_kpt_color,\n pose_link_color,\n radius,\n thickness,\n )\n alpha = 255 * (img != 255).any(axis=-1, keepdims=True).astype(np.uint8)\n return np.concatenate([img, alpha], axis=-1)\n\n\ndef imshow_keypoints(\n img,\n pose_result,\n skeleton=None,\n kpt_score_thr=0.3,\n pose_kpt_color=None,\n pose_link_color=None,\n radius=4,\n thickness=1,\n show_keypoint_weight=False,\n):\n \"\"\"Draw keypoints and links on an image.\n From ViTPose/mmpose/core/visualization/image.py\n\n Args:\n img (H, W, 3) array\n pose_result (list[kpts]): The poses to draw. Each element kpts is\n a set of K keypoints as an Kx3 numpy.ndarray, where each\n keypoint is represented as x, y, score.\n kpt_score_thr (float, optional): Minimum score of keypoints\n to be shown. Default: 0.3.\n pose_kpt_color (np.array[Nx3]`): Color of N keypoints. If None,\n the keypoint will not be drawn.\n pose_link_color (np.array[Mx3]): Color of M links. If None, the\n links will not be drawn.\n thickness (int): Thickness of lines.\n show_keypoint_weight (bool): If True, opacity indicates keypoint score\n \"\"\"\n img_h, img_w, _ = img.shape\n idcs = [0, 16, 15, 18, 17, 5, 2, 6, 3, 7, 4, 12, 9, 13, 10, 14, 11]\n for kpts in pose_result:\n kpts = np.array(kpts, copy=False)[idcs]\n\n # draw each point on image\n if pose_kpt_color is not None:\n assert len(pose_kpt_color) == len(kpts)\n for kid, kpt in enumerate(kpts):\n x_coord, y_coord, kpt_score = int(kpt[0]), int(kpt[1]), kpt[2]\n if kpt_score > kpt_score_thr:\n color = tuple(int(c) for c in pose_kpt_color[kid])\n if show_keypoint_weight:\n img_copy = img.copy()\n cv2.circle(img_copy, (int(x_coord), int(y_coord)), radius, color, -1)\n transparency = max(0, min(1, kpt_score))\n cv2.addWeighted(img_copy, transparency, img, 1 - transparency, 0, dst=img)\n else:\n cv2.circle(img, (int(x_coord), int(y_coord)), radius, color, -1)\n\n # draw links\n if skeleton is not None and pose_link_color is not None:\n assert len(pose_link_color) == len(skeleton)\n for sk_id, sk in enumerate(skeleton):\n pos1 = (int(kpts[sk[0], 0]), int(kpts[sk[0], 1]))\n pos2 = (int(kpts[sk[1], 0]), int(kpts[sk[1], 1]))\n if (\n pos1[0] > 0\n and pos1[0] < img_w\n and pos1[1] > 0\n and pos1[1] < img_h\n and pos2[0] > 0\n and pos2[0] < img_w\n and pos2[1] > 0\n and pos2[1] < img_h\n and kpts[sk[0], 2] > kpt_score_thr\n and kpts[sk[1], 2] > kpt_score_thr\n ):\n color = tuple(int(c) for c in pose_link_color[sk_id])\n if show_keypoint_weight:\n img_copy = img.copy()\n X = (pos1[0], pos2[0])\n Y = (pos1[1], pos2[1])\n mX = np.mean(X)\n mY = np.mean(Y)\n length = ((Y[0] - Y[1]) ** 2 + (X[0] - X[1]) ** 2) ** 0.5\n angle = math.degrees(math.atan2(Y[0] - Y[1], X[0] - X[1]))\n stickwidth = 2\n polygon = cv2.ellipse2Poly(\n (int(mX), int(mY)),\n (int(length / 2), int(stickwidth)),\n int(angle),\n 0,\n 360,\n 1,\n )\n cv2.fillConvexPoly(img_copy, polygon, color)\n transparency = max(0, min(1, 0.5 * (kpts[sk[0], 2] + kpts[sk[1], 2])))\n cv2.addWeighted(img_copy, transparency, img, 1 - transparency, 0, dst=img)\n else:\n cv2.line(img, pos1, pos2, color, thickness=thickness)\n\n return img\n" }, { "path": "eval/GVHMR/hmr4d/utils/vis/renderer_utils.py", "content": "from hmr4d.utils.vis.renderer import Renderer\nfrom tqdm import tqdm\nimport numpy as np\n\n\ndef simple_render_mesh(render_dict):\n \"\"\"Render an camera-space mesh, blank background\"\"\"\n width, height, focal_length = render_dict[\"whf\"]\n faces = render_dict[\"faces\"]\n verts = render_dict[\"verts\"]\n\n renderer = Renderer(width, height, focal_length, device=\"cuda\", faces=faces)\n outputs = []\n for i in tqdm(range(len(verts)), desc=f\"Rendering\"):\n img = renderer.render_mesh(verts[i].cuda(), colors=[0.8, 0.8, 0.8])\n outputs.append(img)\n outputs = np.stack(outputs, axis=0)\n return outputs\n\n\ndef simple_render_mesh_background(render_dict, VI=50, colors=[0.8, 0.8, 0.8]):\n \"\"\"Render an camera-space mesh, blank background\"\"\"\n K = render_dict[\"K\"]\n faces = render_dict[\"faces\"]\n verts = render_dict[\"verts\"]\n background = render_dict[\"background\"]\n N_frames = len(verts)\n if len(background.shape) == 3:\n background = [background] * N_frames\n height, width = background[0].shape[:2]\n\n renderer = Renderer(width, height, device=\"cuda\", faces=faces, K=K)\n outputs = []\n for i in tqdm(range(len(verts)), desc=f\"Rendering\"):\n img = renderer.render_mesh(verts[i].cuda(), colors=colors, background=background[i], VI=VI)\n outputs.append(img)\n outputs = np.stack(outputs, axis=0)\n return outputs\n" }, { "path": "eval/GVHMR/hmr4d/utils/vis/rich_logger.py", "content": "from pytorch_lightning.utilities import rank_zero_only\nfrom omegaconf import DictConfig, OmegaConf\nimport rich\nimport rich.tree\nimport rich.syntax\nfrom hmr4d.utils.pylogger import Log\n\n\n@rank_zero_only\ndef print_cfg(cfg: DictConfig, use_rich: bool = False):\n if use_rich:\n print_order = (\"data\", \"model\", \"callbacks\", \"logger\", \"pl_trainer\")\n style = \"dim\"\n tree = rich.tree.Tree(\"CONFIG\", style=style, guide_style=style)\n\n # add fields from `print_order` to queue\n # add all the other fields to queue (not specified in `print_order`)\n queue = []\n for field in print_order:\n queue.append(field) if field in cfg else Log.warn(f\"Field '{field}' not found in config. Skipping.\")\n for field in cfg:\n if field not in queue:\n queue.append(field)\n\n # generate config tree from queue\n for field in queue:\n branch = tree.add(field, style=style, guide_style=style)\n config_group = cfg[field]\n if isinstance(config_group, DictConfig):\n branch_content = OmegaConf.to_yaml(config_group, resolve=False)\n else:\n branch_content = str(config_group)\n branch.add(rich.syntax.Syntax(branch_content, \"yaml\"))\n rich.print(tree)\n else:\n Log.info(OmegaConf.to_yaml(cfg, resolve=False))\n" }, { "path": "eval/GVHMR/hmr4d/utils/wis3d_utils.py", "content": "from wis3d import Wis3D\nfrom pathlib import Path\nfrom datetime import datetime\nimport torch\nimport numpy as np\nfrom einops import einsum\nfrom pytorch3d.transforms import axis_angle_to_matrix\n\n\ndef make_wis3d(output_dir=\"outputs/wis3d\", name=\"debug\", time_postfix=False):\n \"\"\"\n Make a Wis3D instance. e.g.:\n from hmr4d.utils.wis3d_utils import make_wis3d\n wis3d = make_wis3d(time_postfix=True)\n \"\"\"\n output_dir = Path(output_dir)\n output_dir.mkdir(parents=True, exist_ok=True)\n if time_postfix:\n time_str = datetime.now().strftime(\"%m%d-%H%M-%S\")\n name = f\"{name}_{time_str}\"\n print(f\"Creating Wis3D {name}\")\n wis3d = Wis3D(output_dir.absolute(), name)\n return wis3d\n\n\ncolor_schemes = {\n \"red\": ([255, 168, 154], [153, 17, 1]),\n \"green\": ([183, 255, 191], [0, 171, 8]),\n \"blue\": ([183, 255, 255], [0, 0, 255]),\n \"cyan\": ([183, 255, 255], [0, 255, 255]),\n \"magenta\": ([255, 183, 255], [255, 0, 255]),\n \"black\": ([0, 0, 0], [0, 0, 0]),\n \"orange\": ([255, 183, 0], [255, 128, 0]),\n \"grey\": ([203, 203, 203], [203, 203, 203]),\n}\n\n\ndef get_gradient_colors(scheme=\"red\", num_points=120, alpha=1.0):\n \"\"\"\n Return a list of colors that are gradient from start to end.\n \"\"\"\n start_rgba = torch.tensor(color_schemes[scheme][0] + [255 * alpha]) / 255\n end_rgba = torch.tensor(color_schemes[scheme][1] + [255 * alpha]) / 255\n colors = torch.stack([torch.linspace(s, e, steps=num_points) for s, e in zip(start_rgba, end_rgba)], dim=-1)\n return colors\n\n\ndef get_const_colors(name=\"red\", partial_shape=(120, 5), alpha=1.0):\n \"\"\"\n Return colors (partial_shape, 4)\n \"\"\"\n rgba = torch.tensor(color_schemes[name][1] + [255 * alpha]) / 255\n partial_shape = tuple(partial_shape)\n colors = rgba[None].repeat(*partial_shape, 1)\n return colors\n\n\ndef get_colors_by_conf(conf, low=\"red\", high=\"green\"):\n colors = torch.stack([conf] * 3, dim=-1)\n colors = colors * torch.tensor(color_schemes[high][1]) + (1 - colors) * torch.tensor(color_schemes[low][1])\n return colors\n\n\n# ================== Colored Motion Sequence ================== #\n\n\nKINEMATIC_CHAINS = {\n \"smpl22\": [\n [0, 2, 5, 8, 11], # right-leg\n [0, 1, 4, 7, 10], # left-leg\n [0, 3, 6, 9, 12, 15], # spine\n [9, 14, 17, 19, 21], # right-arm\n [9, 13, 16, 18, 20], # left-arm\n ],\n \"h36m17\": [\n [0, 1, 2, 3], # right-leg\n [0, 4, 5, 6], # left-leg\n [0, 7, 8, 9, 10], # spine\n [8, 14, 15, 16], # right-arm\n [8, 11, 12, 13], # left-arm\n ],\n \"coco17\": [\n [12, 14, 16], # right-leg\n [11, 13, 15], # left-leg\n [4, 2, 0, 1, 3], # replace spine with head\n [6, 8, 10], # right-arm\n [5, 7, 9], # left-arm\n ],\n}\n\n\ndef convert_motion_as_line_mesh(motion, skeleton_type=\"smpl22\", const_color=None):\n if isinstance(motion, np.ndarray):\n motion = torch.from_numpy(motion)\n motion = motion.detach().cpu()\n kinematic_chain = KINEMATIC_CHAINS[skeleton_type]\n color_names = [\"red\", \"green\", \"blue\", \"cyan\", \"magenta\"]\n s_points = []\n e_points = []\n m_colors = []\n length = motion.shape[0]\n device = motion.device\n for chain, color_name in zip(kinematic_chain, color_names):\n num_line = len(chain) - 1\n s_points.append(motion[:, chain[:-1]])\n e_points.append(motion[:, chain[1:]])\n if const_color is not None:\n color_name = const_color\n color_ = get_const_colors(color_name, partial_shape=(length, num_line), alpha=1.0).to(device) # (L, 4, 4)\n m_colors.append(color_[..., :3] * 255) # (L, 4, 3)\n\n s_points = torch.cat(s_points, dim=1) # (L, ?, 3)\n e_points = torch.cat(e_points, dim=1)\n m_colors = torch.cat(m_colors, dim=1)\n\n vertices = []\n for f in range(length):\n vertices_, faces, vertex_colors = create_skeleton_mesh(s_points[f], e_points[f], radius=0.02, color=m_colors[f])\n vertices.append(vertices_)\n vertices = torch.stack(vertices, dim=0)\n return vertices, faces, vertex_colors\n\n\ndef add_motion_as_lines(motion, wis3d, name=\"joints22\", skeleton_type=\"smpl22\", const_color=None, offset=0):\n \"\"\"\n Args:\n motion (tensor): (L, J, 3)\n \"\"\"\n vertices, faces, vertex_colors = convert_motion_as_line_mesh(\n motion, skeleton_type=skeleton_type, const_color=const_color\n )\n for f in range(len(vertices)):\n wis3d.set_scene_id(f + offset)\n wis3d.add_mesh(vertices[f], faces, vertex_colors, name=name) # Add skeleton as cylinders\n # Old way to add lines, this may cause problems when the number of lines is large\n # wis3d.add_lines(s_points[f], e_points[f], m_colors[f], name=name)\n\n\ndef add_prog_motion_as_lines(motion, wis3d, name=\"joints22\", skeleton_type=\"smpl22\"):\n \"\"\"\n Args:\n motion (tensor): (P, L, J, 3)\n \"\"\"\n if isinstance(motion, np.ndarray):\n motion = torch.from_numpy(motion)\n P, L, J, _ = motion.shape\n device = motion.device\n\n kinematic_chain = KINEMATIC_CHAINS[skeleton_type]\n color_names = [\"red\", \"green\", \"blue\", \"cyan\", \"magenta\"]\n s_points = []\n e_points = []\n m_colors = []\n for chain, color_name in zip(kinematic_chain, color_names):\n num_line = len(chain) - 1\n s_points.append(motion[:, :, chain[:-1]])\n e_points.append(motion[:, :, chain[1:]])\n color_ = get_gradient_colors(color_name, L, alpha=1.0).to(device) # (L, 4)\n color_ = color_[None, :, None, :].repeat(P, 1, num_line, 1) # (P, L, num_line, 4)\n m_colors.append(color_[..., :3] * 255) # (P, L, num_line, 3)\n s_points = torch.cat(s_points, dim=-2) # (L, ?, 3)\n e_points = torch.cat(e_points, dim=-2)\n m_colors = torch.cat(m_colors, dim=-2)\n\n s_points = s_points.reshape(P, -1, 3)\n e_points = e_points.reshape(P, -1, 3)\n m_colors = m_colors.reshape(P, -1, 3)\n\n for p in range(P):\n wis3d.set_scene_id(p)\n wis3d.add_lines(s_points[p], e_points[p], m_colors[p], name=name)\n\n\ndef add_joints_motion_as_spheres(joints, wis3d, radius=0.05, name=\"joints\", label_each_joint=False):\n \"\"\"Visualize skeleton as spheres to explore the skeleton.\n Args:\n joints: (NF, NJ, 3)\n wis3d\n radius: radius of the spheres\n name\n label_each_joint: if True, each joints will have a label in wis3d (then you can interact with it, but it's slower)\n \"\"\"\n colors = torch.zeros_like(joints).float()\n n_frames = joints.shape[0]\n n_joints = joints.shape[1]\n for i in range(n_joints):\n colors[:, i, 1] = 255 / n_joints * i\n colors[:, i, 2] = 255 / n_joints * (n_joints - i)\n for f in range(n_frames):\n wis3d.set_scene_id(f)\n if label_each_joint:\n for i in range(n_joints):\n wis3d.add_spheres(\n joints[f, i].float(),\n radius=radius,\n colors=colors[f, i],\n name=f\"{name}-j{i}\",\n )\n else:\n wis3d.add_spheres(\n joints[f].float(),\n radius=radius,\n colors=colors[f],\n name=f\"{name}\",\n )\n\n\ndef create_skeleton_mesh(p1, p2, radius, color, resolution=4, return_merged=True):\n \"\"\"\n Create mesh between p1 and p2.\n Args:\n p1 (torch.Tensor): (N, 3),\n p2 (torch.Tensor): (N, 3),\n radius (float): radius,\n color (torch.Tensor): (N, 3)\n resolution (int): number of vertices in one circle, denoted as Q\n Returns:\n vertices (torch.Tensor): (N * 2Q, 3), if return_merged is False (N, 2Q, 3)\n faces (torch.Tensor): (M', 3), if return_merged is False (N, M, 3)\n vertex_colors (torch.Tensor): (N * 2Q, 3), if return_merged is False (N, 2Q, 3)\n \"\"\"\n N = p1.shape[0]\n\n # Calculate segment direction\n seg_dir = p2 - p1 # (N, 3)\n unit_seg_dir = seg_dir / seg_dir.norm(dim=-1, keepdim=True) # (N, 3)\n\n # Compute an orthogonal vector\n x_vec = torch.tensor([1, 0, 0], device=p1.device).float().unsqueeze(0).repeat(N, 1) # (N, 3)\n y_vec = torch.tensor([0, 1, 0], device=p1.device).float().unsqueeze(0).repeat(N, 1)\n ortho_vec = torch.cross(unit_seg_dir, x_vec, dim=-1) # (N, 3)\n ortho_vec_ = torch.cross(unit_seg_dir, y_vec, dim=-1) # (N, 3) backup\n ortho_vec = torch.where(ortho_vec.norm(dim=-1, keepdim=True) > 1e-3, ortho_vec, ortho_vec_)\n\n # Get circle points on two ends\n unit_ortho_vec = ortho_vec / ortho_vec.norm(dim=-1, keepdim=True) # (N, 3)\n theta = torch.linspace(0, 2 * np.pi, resolution, device=p1.device)\n rotation_matrix = axis_angle_to_matrix(unit_seg_dir[:, None] * theta[None, :, None]) # (N, Q, 3, 3)\n rotated_points = einsum(rotation_matrix, unit_ortho_vec, \"n q i j, n i -> n q j\") * radius # (N, Q, 3)\n bottom_points = rotated_points + p1.unsqueeze(1) # (N, Q, 3)\n top_points = rotated_points + p2.unsqueeze(1) # (N, Q, 3)\n\n # Combine bottom and top points\n vertices = torch.cat([bottom_points, top_points], dim=1) # (N, 2Q, 3)\n\n # Generate face\n indices = torch.arange(0, resolution, device=p1.device)\n bottom_indices = indices\n top_indices = indices + resolution\n\n # outside face\n face_bottom = torch.stack([bottom_indices[:-2], bottom_indices[1:-1], bottom_indices[-1].repeat(resolution - 2)], 1)\n face_top = torch.stack([top_indices[1:-1], top_indices[:-2], top_indices[-1].repeat(resolution - 2)], 1)\n faces = torch.cat(\n [\n torch.stack([bottom_indices[1:], bottom_indices[:-1], top_indices[:-1]], 1), # out face\n torch.stack([bottom_indices[1:], top_indices[:-1], top_indices[1:]], 1), # out face\n face_bottom,\n face_top,\n ]\n )\n faces = faces.unsqueeze(0).repeat(p1.shape[0], 1, 1) # (N, M, 3)\n\n # Assign colors\n vertex_colors = color.unsqueeze(1).repeat(1, resolution * 2, 1)\n\n if return_merged:\n # manully adjust face ids\n N, V = vertices.shape[:2]\n faces = faces + torch.arange(0, N, device=p1.device).unsqueeze(1).unsqueeze(1) * V\n faces = faces.reshape(-1, 3)\n vertices = vertices.reshape(-1, 3)\n vertex_colors = vertex_colors.reshape(-1, 3)\n\n return vertices, faces, vertex_colors\n\n\ndef get_lines_of_my_frustum(frustum_points):\n \"\"\"\n frustum_points: (B, 8, 3), in (near {lu ru rd ld}, far {lu ru rd ld})\n \"\"\"\n start_points = frustum_points[:, [0, 1, 2, 3, 0, 1, 2, 3, 4, 5, 6, 7]].cpu().numpy()\n end_points = frustum_points[:, [4, 5, 6, 7, 1, 2, 3, 0, 5, 6, 7, 4]].cpu().numpy()\n return start_points, end_points\n\n\ndef draw_colored_vec(wis3d, vec, name, radius=0.02, colors=\"r\", starts=None, l=1.0):\n \"\"\"\n Args:\n vec: (3) or (L, 3), should be the same length as colors, like 'rgb'\n \"\"\"\n if len(vec.shape) == 1:\n vec = vec[None]\n else:\n assert len(vec.shape) == 2\n\n assert len(vec) == len(colors)\n # split colors, 'rgb' to 'r', 'g', 'b'\n color_tensor = torch.zeros((len(colors), 3))\n c2rgb = {\n \"r\": torch.tensor([1, 0, 0]).float(),\n \"g\": torch.tensor([0, 1, 0]).float(),\n \"b\": torch.tensor([0, 0, 1]).float(),\n }\n for i, c in enumerate(colors):\n color_tensor[i] = c2rgb[c]\n\n if starts is None:\n starts = torch.zeros_like(vec)\n ends = starts + vec * l\n vertices, faces, vertex_colors = create_skeleton_mesh(starts, ends, radius, color_tensor, resolution=10)\n wis3d.add_mesh(vertices, faces, vertex_colors, name=name)\n\n\ndef draw_T_w2c(wis3d, T_w2c, name, radius=0.01, all_in_one=True, l=0.1):\n \"\"\"\n Draw a camera trajectory in world coordinate.\n Args:\n T_w2c: (L, 4, 4)\n \"\"\"\n color_tensor = torch.eye(3)\n if all_in_one:\n starts = -T_w2c[:, :3, :3].mT @ T_w2c[:, :3, [3]] # (L, 3, 1)\n starts = starts[:, None, :, 0].expand(-1, 3, -1).reshape(-1, 3) # (L*3, 3)\n vec = T_w2c[:, :3, :3].reshape(-1, 3) # (L * 3, 3)\n ends = starts + vec * l\n color_tensor = color_tensor[None].expand(T_w2c.size(0), -1, -1).reshape(-1, 3)\n\n vertices, faces, vertex_colors = create_skeleton_mesh(starts, ends, radius, color_tensor, resolution=10)\n else:\n raise NotImplementedError\n wis3d.add_mesh(vertices, faces, vertex_colors, name=name)\n\n\ndef create_checkerboard_mesh(y=0.0, grid_size=1.0, bounds=((-3, -3), (3, 3))):\n \"\"\"\n example usage:\n vertices, faces, vertex_colors = create_checkerboard_mesh()\n wis3d.add_mesh(vertices=vertices, faces=faces, vertex_colors=vertex_colors, name=\"one\")\n \"\"\"\n color1 = np.array([236, 240, 241], np.uint8) # light\n color2 = np.array([120, 120, 120], np.uint8) # dark\n\n # 扩大范围\n min_x, min_z = bounds[0]\n max_x, max_z = bounds[1]\n min_x = grid_size * np.floor(min_x / grid_size)\n min_z = grid_size * np.floor(min_z / grid_size)\n max_x = grid_size * np.ceil(max_x / grid_size)\n max_z = grid_size * np.ceil(max_z / grid_size)\n\n vertices = []\n faces = []\n vertex_colors = []\n eps = 1e-4 # HACK: disable smooth color & double-side color artifacts of wis3d\n\n for i, x in enumerate(np.arange(min_x, max_x, grid_size)):\n for j, z in enumerate(np.arange(min_z, max_z, grid_size)):\n\n # Right-hand rule for normal direction\n x += ((i % 2 * 2) - 1) * eps\n z += ((j % 2 * 2) - 1) * eps\n v1 = np.array([x, y, z])\n v2 = np.array([x, y, z + grid_size])\n v3 = np.array([x + grid_size, y, z + grid_size])\n v4 = np.array([x + grid_size, y, z])\n offset = np.array([0, -eps, 0]) # For visualizing the down-side of the mesh\n\n vertices.extend([v1, v2, v3, v4, v1 + offset, v2 + offset, v3 + offset, v4 + offset])\n idx = len(vertices) - 8\n faces.extend(\n [\n [idx, idx + 1, idx + 2],\n [idx + 2, idx + 3, idx],\n [idx + 4, idx + 7, idx + 6], # double-sided\n [idx + 6, idx + 5, idx + 4], # double-sided\n ]\n )\n vertex_color = color1 if (i + j) % 2 == 0 else color2\n vertex_colors.extend([vertex_color] * 8)\n\n # To numpy.array and the shape should be (n, 3)\n vertices = np.array(vertices)\n faces = np.array(faces)\n vertex_colors = np.array(vertex_colors)\n assert len(vertices.shape) == 2 and vertices.shape[1] == 3\n assert len(faces.shape) == 2 and faces.shape[1] == 3\n assert len(vertex_colors.shape) == 2 and vertex_colors.shape[1] == 3 and vertex_colors.dtype == np.uint8\n\n return vertices, faces, vertex_colors\n\n\ndef add_a_trimesh(mesh, wis3d, name):\n mesh.apply_transform(wis3d.three_to_world)\n\n # filename = wis3d.__get_export_file_name(\"mesh\", name)\n export_dir = Path(wis3d.out_folder) / wis3d.sequence_name / f\"{wis3d.scene_id:05d}\" / \"meshes\"\n export_dir.mkdir(parents=True, exist_ok=True)\n assert name is not None\n filename = export_dir / f\"{name}.ply\"\n wis3d.counters[\"mesh\"] += 1\n\n mesh.export(filename)\n" }, { "path": "eval/GVHMR/pyproject.toml", "content": "[tool.black]\nline-length = 120\ninclude = '\\.pyi?$'\nexclude = '''\n/(\n \\.git\n | \\.hg\n | \\.mypy_cache\n | \\.tox\n | \\.venv\n | _build\n | buck-out\n | build\n | dist\n)/\n'''\n" }, { "path": "eval/GVHMR/pyrightconfig.json", "content": "{\n \"exclude\": [\n \"./inputs\",\n \"./outputs\"\n ],\n \"typeCheckingMode\": \"off\",\n}\n" }, { "path": "eval/GVHMR/requirements.txt", "content": "# PyTorch\n--extra-index-url https://download.pytorch.org/whl/cu121\ntorch==2.3.0+cu121\ntorchvision==0.18.0+cu121\ntimm==0.9.12 # For HMR2.0a feature extraction\n\n# Lightning + Hydra\nlightning==2.3.0\nhydra-core==1.3\nhydra-zen\nhydra_colorlog\nrich\n\n# Common utilities\nnumpy==1.23.5\njupyter\nmatplotlib\nipdb\nsetuptools>=68.0\nblack\ntensorboardX\nopencv-python\nffmpeg-python\nscikit-image\ntermcolor\neinops\nimageio==2.34.1\nav # imageio[pyav], improved performance over imageio[ffmpeg]\njoblib\n\n# Diffusion\n# diffusers[torch]==0.19.3\n# transformers==4.31.0\n\n# 3D-Vision\npytorch3d @ https://dl.fbaipublicfiles.com/pytorch3d/packaging/wheels/py310_cu121_pyt230/pytorch3d-0.7.6-cp310-cp310-linux_x86_64.whl\ntrimesh\nchumpy\nsmplx\n# open3d==0.17.0\nwis3d\n\n# 2D-Pose\nultralytics==8.2.42 # YOLO\ncython_bbox\nlapx" }, { "path": "eval/GVHMR/setup.py", "content": "from setuptools import setup, find_packages\n\n\nsetup(\n name=\"gvhmr\",\n version=\"1.0.0\",\n packages=find_packages(),\n author=\"Zehong Shen\",\n description=[\"GVHMR training and inference\"],\n url=\"https://github.com/zju3dv/GVHMR\",\n)\n" }, { "path": "eval/GVHMR/tools/demo/colab_demo.ipynb", "content": "{\n \"cells\": [\n {\n \"cell_type\": \"markdown\",\n \"metadata\": {\n \"id\": \"Dv4XCJqiKtun\"\n },\n \"source\": [\n \"
GVHMR
\\n\",\n \"\\n\",\n \"
World-Grounded Human Motion Recovery via Gravity-View Coordinates
\\n\",\n \"\\n\",\n \"![](https://zju3dv.github.io/gvhmr/img/teaser_v1.png)\\n\",\n \"\\n\",\n \"\\n\",\n \"
\\n\",\n \"Project Page\\n\",\n \"|\\n\",\n \"Paper\\n\",\n \"
\\n\",\n \"\\n\",\n \"> World-Grounded Human Motion Recovery via Gravity-View Coordinates \\n\",\n \"> [Zehong Shen](https://zehongs.github.io/)\\\\*,\\n\",\n \"[Huaijin Pi](https://phj128.github.io/)\\\\*,\\n\",\n \"[Yan Xia](https://isshikihugh.github.io/scholar),\\n\",\n \"[Zhi Cen](https://scholar.google.com/citations?user=Xyy-uFMAAAAJ),\\n\",\n \"[Sida Peng](https://pengsida.net/)†,\\n\",\n \"[Zechen Hu](https://zju3dv.github.io/gvhmr),\\n\",\n \"[Hujun Bao](http://www.cad.zju.edu.cn/home/bao/),\\n\",\n \"[Ruizhen Hu](https://csse.szu.edu.cn/staff/ruizhenhu/),\\n\",\n \"[Xiaowei Zhou](https://xzhou.me/) \\n\",\n \"> SIGGRAPH Asia 2024\\n\"\n ]\n },\n {\n \"cell_type\": \"markdown\",\n \"metadata\": {\n \"id\": \"zK2VFE-CEk-l\"\n },\n \"source\": [\n \"## Installation\\n\",\n \"\\n\",\n \"Check [INSTALL.md](https://github.com/IsshikiHugh/GVHMR/blob/main/docs/INSTALL.md) if you want to install GVHMR in your own machine.\\n\",\n \"\\n\",\n \"> Tips: you can fold the section and run the whole installation section at once.\"\n ]\n },\n {\n \"cell_type\": \"code\",\n \"execution_count\": 1,\n \"metadata\": {\n \"colab\": {\n \"base_uri\": \"https://localhost:8080/\"\n },\n \"id\": \"PIPJRQZFGplJ\",\n \"outputId\": \"44c58662-6b8e-4bda-fdb3-d796b6744bd6\"\n },\n \"outputs\": [\n {\n \"name\": \"stdout\",\n \"output_type\": \"stream\",\n \"text\": [\n \"Sat Sep 14 17:16:45 2024 \\n\",\n \"+---------------------------------------------------------------------------------------+\\n\",\n \"| NVIDIA-SMI 535.104.05 Driver Version: 535.104.05 CUDA Version: 12.2 |\\n\",\n \"|-----------------------------------------+----------------------+----------------------+\\n\",\n \"| GPU Name Persistence-M | Bus-Id Disp.A | Volatile Uncorr. ECC |\\n\",\n \"| Fan Temp Perf Pwr:Usage/Cap | Memory-Usage | GPU-Util Compute M. |\\n\",\n \"| | | MIG M. |\\n\",\n \"|=========================================+======================+======================|\\n\",\n \"| 0 Tesla T4 Off | 00000000:00:04.0 Off | 0 |\\n\",\n \"| N/A 52C P8 9W / 70W | 0MiB / 15360MiB | 0% Default |\\n\",\n \"| | | N/A |\\n\",\n \"+-----------------------------------------+----------------------+----------------------+\\n\",\n \" \\n\",\n \"+---------------------------------------------------------------------------------------+\\n\",\n \"| Processes: |\\n\",\n \"| GPU GI CI PID Type Process name GPU Memory |\\n\",\n \"| ID ID Usage |\\n\",\n \"|=======================================================================================|\\n\",\n \"| No running processes found |\\n\",\n \"+---------------------------------------------------------------------------------------+\\n\"\n ]\n }\n ],\n \"source\": [\n \"# Controlling notebook is connected to NVIDIA drivers with CUDA. If this doesn't load check that GPU is selected as hardware accelerator under Edit -> Notebook settings.\\n\",\n \"# Google may shut down the GPU if usage has surpassed the allocation\\n\",\n \"\\n\",\n \"!nvidia-smi\"\n ]\n },\n {\n \"cell_type\": \"markdown\",\n \"metadata\": {\n \"id\": \"nUM1o8lsGU0n\"\n },\n \"source\": [\n \"### Environment Prepration ~ 15 min\\n\",\n \"\\n\",\n \"Check [INSTALL.md#Environment](https://github.com/IsshikiHugh/GVHMR/blob/main/docs/INSTALL.md#environment) for further information.\"\n ]\n },\n {\n \"cell_type\": \"code\",\n \"execution_count\": null,\n \"metadata\": {\n \"id\": \"cKvUncvK-943\"\n },\n \"outputs\": [],\n \"source\": [\n \"import os\\n\",\n \"from pathlib import Path\\n\",\n \"\\n\",\n \"# Clone the repo.\\n\",\n \"!git clone https://github.com/zju3dv/GVHMR --recursive\\n\",\n \"proj_root = str(Path('GVHMR').absolute())\\n\",\n \"\\n\",\n \"# Install GVHMR. (If Colab asks you to restart, just click cancel and rerun the block.)\\n\",\n \"%cd {proj_root}\\n\",\n \"%pip install -r requirements.txt\\n\",\n \"%pip install -e .\"\n ]\n },\n {\n \"cell_type\": \"code\",\n \"execution_count\": null,\n \"metadata\": {\n \"id\": \"vvyrlqQOIz2g\"\n },\n \"outputs\": [],\n \"source\": [\n \"# Install DPVO.\\n\",\n \"%cd third-party/DPVO\\n\",\n \"\\n\",\n \"!wget https://gitlab.com/libeigen/eigen/-/archive/3.4.0/eigen-3.4.0.zip\\n\",\n \"!unzip -o eigen-3.4.0.zip -d thirdparty && rm -rf eigen-3.4.0.zip\\n\",\n \"\\n\",\n \"%pip install torch-scatter -f \\\"https://data.pyg.org/whl/torch-2.3.0+cu121.html\\\"\\n\",\n \"%pip install numba pypose\\n\",\n \"\\n\",\n \"if 'cuda_home' not in locals():\\n\",\n \" cuda_home = '/usr/local/cuda-12'\\n\",\n \" if not Path(cuda_home).exists():\\n\",\n \" raise FileNotFoundError('CUDA_HOME for cuda 12.x not found!')\\n\",\n \"\\n\",\n \" os.environ['CUDA_HOME'] = cuda_home\\n\",\n \" os.environ['PATH'] = os.environ['PATH'] + f':{cuda_home}/bin'\\n\",\n \"\\n\",\n \"%pip install -e .\\n\",\n \"%cd {proj_root}\"\n ]\n },\n {\n \"cell_type\": \"markdown\",\n \"metadata\": {\n \"id\": \"vziM8bI6E5ur\"\n },\n \"source\": [\n \"### Data Prepration ~ 1 min\\n\",\n \"\\n\",\n \"Check [INSTALL.md#Inputs&Outputs](https://github.com/IsshikiHugh/GVHMR/blob/main/docs/INSTALL.md#inputs--outputs) for further information.\"\n ]\n },\n {\n \"cell_type\": \"code\",\n \"execution_count\": 11,\n \"metadata\": {\n \"colab\": {\n \"base_uri\": \"https://localhost:8080/\"\n },\n \"id\": \"3SYDPesR_1m9\",\n \"outputId\": \"8658d240-9471-4779-c67e-61f9d92142dd\"\n },\n \"outputs\": [\n {\n \"name\": \"stdout\",\n \"output_type\": \"stream\",\n \"text\": [\n \"/content/GVHMR\\n\",\n \"mkdir: cannot create directory ‘inputs’: File exists\\n\",\n \"mkdir: cannot create directory ‘outputs’: File exists\\n\",\n \"aria2 is already the newest version (1.36.0-1).\\n\",\n \"0 upgraded, 0 newly installed, 0 to remove and 49 not upgraded.\\n\",\n \"\\n\",\n \"Download Results:\\n\",\n \"gid |stat|avg speed |path/URI\\n\",\n \"======+====+===========+=======================================================\\n\",\n \"858436|\\u001b[1;32mOK\\u001b[0m | 0B/s|inputs/checkpoints/body_models/smpl/SMPL_NEUTRAL.pkl\\n\",\n \"\\n\",\n \"Status Legend:\\n\",\n \"(OK):download completed.\\n\",\n \"\\n\",\n \"Download Results:\\n\",\n \"gid |stat|avg speed |path/URI\\n\",\n \"======+====+===========+=======================================================\\n\",\n \"da93ed|\\u001b[1;32mOK\\u001b[0m | 0B/s|inputs/checkpoints/body_models/smplx/SMPLX_NEUTRAL.npz\\n\",\n \"\\n\",\n \"Status Legend:\\n\",\n \"(OK):download completed.\\n\",\n \"\\n\",\n \"Download Results:\\n\",\n \"gid |stat|avg speed |path/URI\\n\",\n \"======+====+===========+=======================================================\\n\",\n \"b2dfd7|\\u001b[1;32mOK\\u001b[0m | 0B/s|inputs/checkpoints/dpvo/dpvo.pth\\n\",\n \"\\n\",\n \"Status Legend:\\n\",\n \"(OK):download completed.\\n\",\n \"\\n\",\n \"Download Results:\\n\",\n \"gid |stat|avg speed |path/URI\\n\",\n \"======+====+===========+=======================================================\\n\",\n \"5ce060|\\u001b[1;32mOK\\u001b[0m | 0B/s|inputs/checkpoints/gvhmr/gvhmr_siga24_release.ckpt\\n\",\n \"\\n\",\n \"Status Legend:\\n\",\n \"(OK):download completed.\\n\",\n \"\\n\",\n \"Download Results:\\n\",\n \"gid |stat|avg speed |path/URI\\n\",\n \"======+====+===========+=======================================================\\n\",\n \"d41b3e|\\u001b[1;32mOK\\u001b[0m | 0B/s|inputs/checkpoints/hmr2/epoch=10-step=25000.ckpt\\n\",\n \"\\n\",\n \"Status Legend:\\n\",\n \"(OK):download completed.\\n\",\n \"\\n\",\n \"Download Results:\\n\",\n \"gid |stat|avg speed |path/URI\\n\",\n \"======+====+===========+=======================================================\\n\",\n \"3fe18e|\\u001b[1;32mOK\\u001b[0m | 0B/s|inputs/checkpoints/vitpose/vitpose-h-multi-coco.pth\\n\",\n \"\\n\",\n \"Status Legend:\\n\",\n \"(OK):download completed.\\n\",\n \"\\n\",\n \"Download Results:\\n\",\n \"gid |stat|avg speed |path/URI\\n\",\n \"======+====+===========+=======================================================\\n\",\n \"4fe1eb|\\u001b[1;32mOK\\u001b[0m | 0B/s|inputs/checkpoints/yolo/yolov8x.pt\\n\",\n \"\\n\",\n \"Status Legend:\\n\",\n \"(OK):download completed.\\n\"\n ]\n }\n ],\n \"source\": [\n \"%cd {proj_root}\\n\",\n \"!mkdir inputs\\n\",\n \"!mkdir outputs\\n\",\n \"\\n\",\n \"!apt install -y -qq aria2\\n\",\n \"!aria2c --console-log-level=error -c -x 16 -s 16 -k 1M https://huggingface.co/camenduru/SMPLer-X/resolve/main/SMPL_NEUTRAL.pkl -d inputs/checkpoints/body_models/smpl -o SMPL_NEUTRAL.pkl\\n\",\n \"!aria2c --console-log-level=error -c -x 16 -s 16 -k 1M https://huggingface.co/camenduru/SMPLer-X/resolve/main/SMPLX_NEUTRAL.npz -d inputs/checkpoints/body_models/smplx -o SMPLX_NEUTRAL.npz\\n\",\n \"!aria2c --console-log-level=error -c -x 16 -s 16 -k 1M https://huggingface.co/camenduru/GVHMR/resolve/main/dpvo/dpvo.pth -d inputs/checkpoints/dpvo -o dpvo.pth\\n\",\n \"!aria2c --console-log-level=error -c -x 16 -s 16 -k 1M https://huggingface.co/camenduru/GVHMR/resolve/main/gvhmr/gvhmr_siga24_release.ckpt -d inputs/checkpoints/gvhmr -o gvhmr_siga24_release.ckpt\\n\",\n \"!aria2c --console-log-level=error -c -x 16 -s 16 -k 1M https://huggingface.co/camenduru/GVHMR/resolve/main/hmr2/epoch%3D10-step%3D25000.ckpt -d inputs/checkpoints/hmr2 -o epoch=10-step=25000.ckpt\\n\",\n \"!aria2c --console-log-level=error -c -x 16 -s 16 -k 1M https://huggingface.co/camenduru/GVHMR/resolve/main/vitpose/vitpose-h-multi-coco.pth -d inputs/checkpoints/vitpose -o vitpose-h-multi-coco.pth\\n\",\n \"!aria2c --console-log-level=error -c -x 16 -s 16 -k 1M https://huggingface.co/camenduru/GVHMR/resolve/main/yolo/yolov8x.pt -d inputs/checkpoints/yolo -o yolov8x.pt\"\n ]\n },\n {\n \"cell_type\": \"markdown\",\n \"metadata\": {\n \"id\": \"IxHWylQ-KJdG\"\n },\n \"source\": [\n \"## Run Demo\\n\",\n \"\\n\",\n \"Use `-s` to skip visual odometry if you know the camera is static, otherwise the camera will be estimated by DPVO.\\n\",\n \"\\n\",\n \"We also provide a script demo_folder.py to inference a entire folder.\\n\",\n \"\\n\",\n \"\\n\",\n \"```shell\\n\",\n \"python tools/demo/demo.py --video=docs/example_video/tennis.mp4 -s\\n\",\n \"python tools/demo/demo_folder.py -f inputs/demo/folder_in -d outputs/demo/folder_out -s\\n\",\n \"\\n\",\n \"```\"\n ]\n },\n {\n \"cell_type\": \"code\",\n \"execution_count\": 35,\n \"metadata\": {\n \"id\": \"YHaW4uybOQgG\"\n },\n \"outputs\": [],\n \"source\": [\n \"import io\\n\",\n \"import base64\\n\",\n \"from IPython.display import HTML\\n\",\n \"from hmr4d.utils.video_io_utils import get_video_lwh\\n\",\n \"\\n\",\n \"def display_video(fn):\\n\",\n \" L, W, H = get_video_lwh(fn)\\n\",\n \" scale = min(W, 1080) / W\\n\",\n \" W, H = int(W * scale), int(H * scale)\\n\",\n \" video_encoded = base64.b64encode(io.open(fn, 'rb').read())\\n\",\n \" return HTML(data=''''''.format(W, H, video_encoded.decode('ascii')))\"\n ]\n },\n {\n \"cell_type\": \"markdown\",\n \"metadata\": {\n \"id\": \"3EOGJBa_OCWu\"\n },\n \"source\": [\n \"### Demo Tennis\"\n ]\n },\n {\n \"cell_type\": \"code\",\n \"execution_count\": 12,\n \"metadata\": {\n \"colab\": {\n \"base_uri\": \"https://localhost:8080/\"\n },\n \"id\": \"NQIVoNKMKlhb\",\n \"outputId\": \"e125532f-f095-46b2-dd70-11f227b9307d\"\n },\n \"outputs\": [\n {\n \"name\": \"stdout\",\n \"output_type\": \"stream\",\n \"text\": [\n \"[\\u001b[36m09/14 17:37:05\\u001b[0m][\\u001b[32mINFO\\u001b[0m] [Input]: /content/GVHMR/docs/example_video/tennis.mp4\\u001b[0m\\n\",\n \"[\\u001b[36m09/14 17:37:05\\u001b[0m][\\u001b[32mINFO\\u001b[0m] (L, W, H) = (312, 812, 720)\\u001b[0m\\n\",\n \"[\\u001b[36m09/14 17:37:06\\u001b[0m][\\u001b[32mINFO\\u001b[0m] [Output Dir]: outputs/demo/tennis\\u001b[0m\\n\",\n \"[\\u001b[36m09/14 17:37:06\\u001b[0m][\\u001b[32mINFO\\u001b[0m] [Copy Video] /content/GVHMR/docs/example_video/tennis.mp4 -> outputs/demo/tennis/0_input_video.mp4\\u001b[0m\\n\",\n \"Copy: 100% 312/312 [00:09<00:00, 33.08it/s]\\n\",\n \"[\\u001b[36m09/14 17:37:17\\u001b[0m][\\u001b[32mINFO\\u001b[0m] [GPU]: Tesla T4\\u001b[0m\\n\",\n \"[\\u001b[36m09/14 17:37:17\\u001b[0m][\\u001b[32mINFO\\u001b[0m] [GPU]: _CudaDeviceProperties(name='Tesla T4', major=7, minor=5, total_memory=15102MB, multi_processor_count=40)\\u001b[0m\\n\",\n \"[\\u001b[36m09/14 17:37:17\\u001b[0m][\\u001b[32mINFO\\u001b[0m] [Preprocess] Start!\\u001b[0m\\n\",\n \"YoloV8 Tracking: 100% 312/312 [00:27<00:00, 11.14it/s]\\n\",\n \"ViTPose: 100% 20/20 [00:42<00:00, 2.13s/it]\\n\",\n \"HMR2 Feature: 100% 20/20 [00:22<00:00, 1.11s/it]\\n\",\n \"[\\u001b[36m09/14 17:39:36\\u001b[0m][\\u001b[32mINFO\\u001b[0m] [Preprocess] End. Time elapsed: 139.28s\\u001b[0m\\n\",\n \"[\\u001b[36m09/14 17:39:36\\u001b[0m][\\u001b[32mINFO\\u001b[0m] [HMR4D] Predicting\\u001b[0m\\n\",\n \"[\\u001b[36m09/14 17:39:36\\u001b[0m][\\u001b[32mINFO\\u001b[0m] [EnDecoder] Use MM_V1_AMASS_LOCAL_BEDLAM_CAM for statistics!\\u001b[0m\\n\",\n \"[\\u001b[36m09/14 17:39:38\\u001b[0m][\\u001b[32mINFO\\u001b[0m] [PL-Trainer] Loading ckpt type: inputs/checkpoints/gvhmr/gvhmr_siga24_release.ckpt\\u001b[0m\\n\",\n \"[\\u001b[36m09/14 17:39:40\\u001b[0m][\\u001b[32mINFO\\u001b[0m] [HMR4D] Elapsed: 1.32s for data-length=10.4s\\u001b[0m\\n\",\n \"Rendering Incam: 100% 312/312 [00:24<00:00, 12.71it/s]\\n\",\n \"Rendering Global: 100% 312/312 [00:16<00:00, 18.68it/s]\\n\",\n \"[\\u001b[36m09/14 17:40:29\\u001b[0m][\\u001b[32mINFO\\u001b[0m] [Merge Videos]\\u001b[0m\\n\"\n ]\n }\n ],\n \"source\": [\n \"# Run demo.\\n\",\n \"video_fn = f'{proj_root}/docs/example_video/tennis.mp4'\\n\",\n \"!python {proj_root}/tools/demo/demo.py --video={video_fn} -s\"\n ]\n },\n {\n \"cell_type\": \"code\",\n \"execution_count\": 36,\n \"metadata\": {\n \"colab\": {\n \"base_uri\": \"https://localhost:8080/\",\n \"height\": 499\n },\n \"id\": \"wPam0XZfQLL5\",\n \"outputId\": \"02b63472-4539-42a1-b8fd-c9208f3ded72\"\n },\n \"outputs\": [\n {\n \"data\": {\n \"text/html\": [\n \"\"\n ],\n \"text/plain\": [\n \"\"\n ]\n },\n \"execution_count\": 36,\n \"metadata\": {},\n \"output_type\": \"execute_result\"\n }\n ],\n \"source\": [\n \"# Visualize the result.\\n\",\n \"display_video('outputs/demo/tennis/tennis_3_incam_global_horiz.mp4')\"\n ]\n },\n {\n \"cell_type\": \"markdown\",\n \"metadata\": {\n \"id\": \"CnOwTAfyRo57\"\n },\n \"source\": [\n \"### Custom Demo\\n\",\n \"\\n\",\n \"In order to run your own video, you need to follow instructions below:\\n\",\n \"\\n\",\n \"0. Run the code block below to initialize `demo_google_drive_video()`.\\n\",\n \"1. Upload your video on Google drive and set accessbility as \\\"Anyone with link (Viewer)\\\".\\n\",\n \"2. Copy the link and extract the ID. The link should follow this pattern: `https://drive.google.com/file/d//view?usp=sharing`.\\n\",\n \"3. Call `demo_google_drive_video()` with the URL, and pass `True` if the camera is static.\\n\"\n ]\n },\n {\n \"cell_type\": \"code\",\n \"execution_count\": 31,\n \"metadata\": {\n \"colab\": {\n \"base_uri\": \"https://localhost:8080/\"\n },\n \"id\": \"jqOda6ONRm_C\",\n \"outputId\": \"93672ab6-a490-407c-d6d9-a8174be9bb5a\"\n },\n \"outputs\": [\n {\n \"name\": \"stdout\",\n \"output_type\": \"stream\",\n \"text\": [\n \"/content/GVHMR\\n\"\n ]\n }\n ],\n \"source\": [\n \"%cd {proj_root}\\n\",\n \"!mkdir -p inputs/demo\\n\",\n \"\\n\",\n \"def demo_google_drive_video(url:str, static_camera:bool):\\n\",\n \" ''' URL should be like https://drive.google.com/file/d/xxxxxxxx/view?usp=drive_link '''\\n\",\n \"\\n\",\n \" print(f'[1/3] 📥 Downloading video...')\\n\",\n \" gdid = url.split('/')[5]\\n\",\n \" video_name = f'custom_{gdid}'\\n\",\n \" download_url = f'\\\\'https://drive.google.com/uc?id={gdid}&export=download&confirm=t\\\\''\\n\",\n \" !gdown {download_url} -O inputs/demo/{video_name}.mp4\\n\",\n \"\\n\",\n \" print(f'[2/3] 💃 Start running GVHMR...')\\n\",\n \" flag = '-s' if static_camera else ''\\n\",\n \" !python {proj_root}/tools/demo/demo.py --video=inputs/demo/{video_name}.mp4 {flag}\\n\",\n \"\\n\",\n \" print(f'[3/3] 📺 Displaying result...')\\n\",\n \" return display_video(f'outputs/demo/{video_name}/{video_name}_3_incam_global_horiz.mp4')\"\n ]\n },\n {\n \"cell_type\": \"code\",\n \"execution_count\": 39,\n \"metadata\": {\n \"colab\": {\n \"base_uri\": \"https://localhost:8080/\",\n \"height\": 1000\n },\n \"id\": \"ElCtnGFQVHU1\",\n \"outputId\": \"d957a268-8648-4f2e-c360-c82ee56a7b25\"\n },\n \"outputs\": [\n {\n \"name\": \"stdout\",\n \"output_type\": \"stream\",\n \"text\": [\n \"[1/3] 📥 Downloading video...\\n\",\n \"Downloading...\\n\",\n \"From: https://drive.google.com/uc?id=1KkTsoAuj9yq5JCJ-qg4_MnNBn56wR47F&export=download&confirm=t\\n\",\n \"To: /content/GVHMR/inputs/demo/custom_1KkTsoAuj9yq5JCJ-qg4_MnNBn56wR47F.mp4\\n\",\n \"100% 936k/936k [00:00<00:00, 9.14MB/s]\\n\",\n \"[2/3] 💃 Start running GVHMR...\\n\",\n \"[\\u001b[36m09/14 18:09:12\\u001b[0m][\\u001b[32mINFO\\u001b[0m] [Input]: inputs/demo/custom_1KkTsoAuj9yq5JCJ-qg4_MnNBn56wR47F.mp4\\u001b[0m\\n\",\n \"[\\u001b[36m09/14 18:09:12\\u001b[0m][\\u001b[32mINFO\\u001b[0m] (L, W, H) = (367, 682, 666)\\u001b[0m\\n\",\n \"[\\u001b[36m09/14 18:09:13\\u001b[0m][\\u001b[32mINFO\\u001b[0m] [Output Dir]: outputs/demo/custom_1KkTsoAuj9yq5JCJ-qg4_MnNBn56wR47F\\u001b[0m\\n\",\n \"[\\u001b[36m09/14 18:09:13\\u001b[0m][\\u001b[32mINFO\\u001b[0m] [Copy Video] inputs/demo/custom_1KkTsoAuj9yq5JCJ-qg4_MnNBn56wR47F.mp4 -> outputs/demo/custom_1KkTsoAuj9yq5JCJ-qg4_MnNBn56wR47F/0_input_video.mp4\\u001b[0m\\n\",\n \"Copy: 100% 367/367 [00:09<00:00, 39.06it/s]\\n\",\n \"[\\u001b[36m09/14 18:09:23\\u001b[0m][\\u001b[32mINFO\\u001b[0m] [GPU]: Tesla T4\\u001b[0m\\n\",\n \"[\\u001b[36m09/14 18:09:23\\u001b[0m][\\u001b[32mINFO\\u001b[0m] [GPU]: _CudaDeviceProperties(name='Tesla T4', major=7, minor=5, total_memory=15102MB, multi_processor_count=40)\\u001b[0m\\n\",\n \"[\\u001b[36m09/14 18:09:23\\u001b[0m][\\u001b[32mINFO\\u001b[0m] [Preprocess] Start!\\u001b[0m\\n\",\n \"YoloV8 Tracking: 100% 367/367 [00:29<00:00, 12.65it/s]\\n\",\n \"ViTPose: 100% 23/23 [00:52<00:00, 2.27s/it]\\n\",\n \"HMR2 Feature: 100% 23/23 [00:25<00:00, 1.12s/it]\\n\",\n \"[\\u001b[36m09/14 18:11:57\\u001b[0m][\\u001b[32mINFO\\u001b[0m] [Preprocess] End. Time elapsed: 154.35s\\u001b[0m\\n\",\n \"[\\u001b[36m09/14 18:11:57\\u001b[0m][\\u001b[32mINFO\\u001b[0m] [HMR4D] Predicting\\u001b[0m\\n\",\n \"[\\u001b[36m09/14 18:11:57\\u001b[0m][\\u001b[32mINFO\\u001b[0m] [EnDecoder] Use MM_V1_AMASS_LOCAL_BEDLAM_CAM for statistics!\\u001b[0m\\n\",\n \"[\\u001b[36m09/14 18:11:59\\u001b[0m][\\u001b[32mINFO\\u001b[0m] [PL-Trainer] Loading ckpt type: inputs/checkpoints/gvhmr/gvhmr_siga24_release.ckpt\\u001b[0m\\n\",\n \"[\\u001b[36m09/14 18:12:02\\u001b[0m][\\u001b[32mINFO\\u001b[0m] [HMR4D] Elapsed: 1.88s for data-length=12.2s\\u001b[0m\\n\",\n \"Rendering Incam: 100% 367/367 [00:21<00:00, 16.96it/s]\\n\",\n \"Rendering Global: 100% 367/367 [00:17<00:00, 21.44it/s]\\n\",\n \"[\\u001b[36m09/14 18:12:47\\u001b[0m][\\u001b[32mINFO\\u001b[0m] [Merge Videos]\\u001b[0m\\n\",\n \"[3/3] 📺 Displaying result...\\n\"\n ]\n },\n {\n \"data\": {\n \"text/html\": [\n \"\"\n ],\n \"text/plain\": [\n \"\"\n ]\n },\n \"execution_count\": 39,\n \"metadata\": {},\n \"output_type\": \"execute_result\"\n }\n ],\n \"source\": [\n \"demo_google_drive_video(url='https://drive.google.com/file/d/1KkTsoAuj9yq5JCJ-qg4_MnNBn56wR47F/view?usp=drive_link', static_camera=True)\"\n ]\n },\n {\n \"cell_type\": \"code\",\n \"execution_count\": null,\n \"metadata\": {\n \"id\": \"irhhFd76pH5c\"\n },\n \"outputs\": [],\n \"source\": [\n \"# Try it yourself!\\n\",\n \"demo_google_drive_video(...)\"\n ]\n }\n ],\n \"metadata\": {\n \"accelerator\": \"GPU\",\n \"colab\": {\n \"gpuType\": \"T4\",\n \"provenance\": []\n },\n \"kernelspec\": {\n \"display_name\": \"Python 3\",\n \"name\": \"python3\"\n },\n \"language_info\": {\n \"name\": \"python\"\n }\n },\n \"nbformat\": 4,\n \"nbformat_minor\": 0\n}\n" }, { "path": "eval/GVHMR/tools/demo/demo.py", "content": "import cv2\nimport torch\nimport pytorch_lightning as pl\nimport numpy as np\nimport argparse\nfrom hmr4d.utils.pylogger import Log\nimport hydra\nfrom hydra import initialize_config_module, compose\nfrom pathlib import Path\nfrom pytorch3d.transforms import quaternion_to_matrix\nimport os\n\nfrom hmr4d.configs import register_store_gvhmr\nfrom hmr4d.utils.video_io_utils import (\n get_video_lwh,\n read_video_np,\n save_video,\n merge_videos_horizontal,\n get_writer,\n get_video_reader,\n)\nfrom hmr4d.utils.vis.cv2_utils import draw_bbx_xyxy_on_image_batch, draw_coco17_skeleton_batch\n\nfrom hmr4d.utils.preproc import Tracker, Extractor, VitPoseExtractor, SLAMModel\n\nfrom hmr4d.utils.geo.hmr_cam import get_bbx_xys_from_xyxy, estimate_K, convert_K_to_K4, create_camera_sensor\nfrom hmr4d.utils.geo_transform import compute_cam_angvel\nfrom hmr4d.model.gvhmr.gvhmr_pl_demo import DemoPL\nfrom hmr4d.utils.net_utils import detach_to_cpu, to_cuda\nfrom hmr4d.utils.smplx_utils import make_smplx\nfrom hmr4d.utils.vis.renderer import Renderer, get_global_cameras_static, get_ground_params_from_points\nfrom tqdm import tqdm\nfrom hmr4d.utils.geo_transform import apply_T_on_points, compute_T_ayfz2ay\nfrom einops import einsum, rearrange\n\n\nCRF = 23 # 17 is lossless, every +6 halves the mp4 size\n\n\ndef parse_args_to_cfg():\n # Put all args to cfg\n parser = argparse.ArgumentParser()\n parser.add_argument(\"--video\", type=str, default=\"inputs/demo/dance_3.mp4\")\n parser.add_argument(\"--output_root\", type=str, default=None, help=\"by default to outputs/demo\")\n parser.add_argument(\"-s\", \"--static_cam\", action=\"store_true\", help=\"If true, skip DPVO\")\n parser.add_argument(\"--verbose\", action=\"store_true\", help=\"If true, draw intermediate results\")\n args = parser.parse_args()\n\n # Input\n video_path = Path(args.video)\n assert video_path.exists(), f\"Video not found at {video_path}\"\n length, width, height = get_video_lwh(video_path)\n Log.info(f\"[Input]: {video_path}\")\n Log.info(f\"(L, W, H) = ({length}, {width}, {height})\")\n # Cfg\n with initialize_config_module(version_base=\"1.3\", config_module=f\"hmr4d.configs\"):\n overrides = [\n f\"video_name={video_path.stem}\",\n f\"static_cam={args.static_cam}\",\n f\"verbose={args.verbose}\",\n ]\n\n # Allow to change output root\n if args.output_root is not None:\n overrides.append(f\"output_root={args.output_root}\")\n register_store_gvhmr()\n cfg = compose(config_name=\"demo\", overrides=overrides)\n\n # Output\n Log.info(f\"[Output Dir]: {cfg.output_dir}\")\n Path(cfg.output_dir).mkdir(parents=True, exist_ok=True)\n Path(cfg.preprocess_dir).mkdir(parents=True, exist_ok=True)\n\n # Copy raw-input-video to video_path\n Log.info(f\"[Copy Video] {video_path} -> {cfg.video_path}\")\n if not Path(cfg.video_path).exists() or get_video_lwh(video_path)[0] != get_video_lwh(cfg.video_path)[0]:\n reader = get_video_reader(video_path)\n writer = get_writer(cfg.video_path, fps=30, crf=CRF)\n for img in tqdm(reader, total=get_video_lwh(video_path)[0], desc=f\"Copy\"):\n writer.write_frame(img)\n writer.close()\n reader.close()\n\n return cfg\n\n\n@torch.no_grad()\ndef run_preprocess(cfg):\n Log.info(f\"[Preprocess] Start!\")\n tic = Log.time()\n video_path = cfg.video_path\n paths = cfg.paths\n static_cam = cfg.static_cam\n verbose = cfg.verbose\n\n # Get bbx tracking result\n if not Path(paths.bbx).exists():\n tracker = Tracker()\n bbx_xyxy = tracker.get_one_track(video_path).float() # (L, 4)\n bbx_xys = get_bbx_xys_from_xyxy(bbx_xyxy, base_enlarge=1.2).float() # (L, 3) apply aspect ratio and enlarge\n torch.save({\"bbx_xyxy\": bbx_xyxy, \"bbx_xys\": bbx_xys}, paths.bbx)\n del tracker\n else:\n bbx_xys = torch.load(paths.bbx)[\"bbx_xys\"]\n Log.info(f\"[Preprocess] bbx (xyxy, xys) from {paths.bbx}\")\n if verbose:\n video = read_video_np(video_path)\n bbx_xyxy = torch.load(paths.bbx)[\"bbx_xyxy\"]\n video_overlay = draw_bbx_xyxy_on_image_batch(bbx_xyxy, video)\n save_video(video_overlay, cfg.paths.bbx_xyxy_video_overlay)\n\n # Get VitPose\n if not Path(paths.vitpose).exists():\n vitpose_extractor = VitPoseExtractor()\n vitpose = vitpose_extractor.extract(video_path, bbx_xys)\n torch.save(vitpose, paths.vitpose)\n del vitpose_extractor\n else:\n vitpose = torch.load(paths.vitpose)\n Log.info(f\"[Preprocess] vitpose from {paths.vitpose}\")\n if verbose:\n video = read_video_np(video_path)\n video_overlay = draw_coco17_skeleton_batch(video, vitpose, 0.5)\n save_video(video_overlay, paths.vitpose_video_overlay)\n\n # Get vit features\n if not Path(paths.vit_features).exists():\n extractor = Extractor()\n vit_features = extractor.extract_video_features(video_path, bbx_xys)\n torch.save(vit_features, paths.vit_features)\n del extractor\n else:\n Log.info(f\"[Preprocess] vit_features from {paths.vit_features}\")\n\n # Get DPVO results\n if not static_cam: # use slam to get cam rotation\n if not Path(paths.slam).exists():\n length, width, height = get_video_lwh(cfg.video_path)\n K_fullimg = estimate_K(width, height)\n intrinsics = convert_K_to_K4(K_fullimg)\n slam = SLAMModel(video_path, width, height, intrinsics, buffer=4000, resize=0.5)\n bar = tqdm(total=length, desc=\"DPVO\")\n while True:\n ret = slam.track()\n if ret:\n bar.update()\n else:\n break\n slam_results = slam.process() # (L, 7), numpy\n torch.save(slam_results, paths.slam)\n else:\n Log.info(f\"[Preprocess] slam results from {paths.slam}\")\n\n Log.info(f\"[Preprocess] End. Time elapsed: {Log.time()-tic:.2f}s\")\n\n\ndef load_data_dict(cfg):\n paths = cfg.paths\n length, width, height = get_video_lwh(cfg.video_path)\n if cfg.static_cam:\n R_w2c = torch.eye(3).repeat(length, 1, 1)\n else:\n traj = torch.load(cfg.paths.slam)\n traj_quat = torch.from_numpy(traj[:, [6, 3, 4, 5]])\n R_w2c = quaternion_to_matrix(traj_quat).mT\n K_fullimg = estimate_K(width, height).repeat(length, 1, 1)\n # K_fullimg = create_camera_sensor(width, height, 26)[2].repeat(length, 1, 1)\n\n data = {\n \"length\": torch.tensor(length),\n \"bbx_xys\": torch.load(paths.bbx)[\"bbx_xys\"],\n \"kp2d\": torch.load(paths.vitpose),\n \"K_fullimg\": K_fullimg,\n \"cam_angvel\": compute_cam_angvel(R_w2c),\n \"f_imgseq\": torch.load(paths.vit_features),\n }\n return data\n\n\ndef render_incam(cfg):\n incam_video_path = Path(cfg.paths.incam_video)\n if incam_video_path.exists():\n Log.info(f\"[Render Incam] Video already exists at {incam_video_path}\")\n return\n\n pred = torch.load(cfg.paths.hmr4d_results)\n smplx = make_smplx(\"supermotion\").cuda()\n smplx2smpl = torch.load(\"hmr4d/utils/body_model/smplx2smpl_sparse.pt\").cuda()\n faces_smpl = make_smplx(\"smpl\").faces\n\n # smpl\n smplx_out = smplx(**to_cuda(pred[\"smpl_params_incam\"]))\n pred_c_verts = torch.stack([torch.matmul(smplx2smpl, v_) for v_ in smplx_out.vertices])\n\n # -- rendering code -- #\n video_path = cfg.video_path\n length, width, height = get_video_lwh(video_path)\n K = pred[\"K_fullimg\"][0]\n\n # renderer\n renderer = Renderer(width, height, device=\"cuda\", faces=faces_smpl, K=K)\n reader = get_video_reader(video_path) # (F, H, W, 3), uint8, numpy\n bbx_xys_render = torch.load(cfg.paths.bbx)[\"bbx_xys\"]\n\n # -- render mesh -- #\n verts_incam = pred_c_verts\n writer = get_writer(incam_video_path, fps=30, crf=CRF)\n for i, img_raw in tqdm(enumerate(reader), total=get_video_lwh(video_path)[0], desc=f\"Rendering Incam\"):\n img = renderer.render_mesh(verts_incam[i].cuda(), img_raw, [0.8, 0.8, 0.8])\n\n # # bbx\n # bbx_xys_ = bbx_xys_render[i].cpu().numpy()\n # lu_point = (bbx_xys_[:2] - bbx_xys_[2:] / 2).astype(int)\n # rd_point = (bbx_xys_[:2] + bbx_xys_[2:] / 2).astype(int)\n # img = cv2.rectangle(img, lu_point, rd_point, (255, 178, 102), 2)\n\n writer.write_frame(img)\n writer.close()\n reader.close()\n\n\ndef render_global(cfg):\n global_video_path = Path(cfg.paths.global_video)\n if global_video_path.exists():\n Log.info(f\"[Render Global] Video already exists at {global_video_path}\")\n return\n\n debug_cam = False\n pred = torch.load(cfg.paths.hmr4d_results)\n smplx = make_smplx(\"supermotion\").cuda()\n smplx2smpl = torch.load(\"hmr4d/utils/body_model/smplx2smpl_sparse.pt\").cuda()\n faces_smpl = make_smplx(\"smpl\").faces\n J_regressor = torch.load(\"hmr4d/utils/body_model/smpl_neutral_J_regressor.pt\").cuda()\n\n # smpl\n smplx_out = smplx(**to_cuda(pred[\"smpl_params_global\"]))\n pred_ay_verts = torch.stack([torch.matmul(smplx2smpl, v_) for v_ in smplx_out.vertices])\n\n def move_to_start_point_face_z(verts):\n \"XZ to origin, Start from the ground, Face-Z\"\n # position\n verts = verts.clone() # (L, V, 3)\n offset = einsum(J_regressor, verts[0], \"j v, v i -> j i\")[0] # (3)\n offset[1] = verts[:, :, [1]].min()\n verts = verts - offset\n # face direction\n T_ay2ayfz = compute_T_ayfz2ay(einsum(J_regressor, verts[[0]], \"j v, l v i -> l j i\"), inverse=True)\n verts = apply_T_on_points(verts, T_ay2ayfz)\n return verts\n\n verts_glob = move_to_start_point_face_z(pred_ay_verts)\n joints_glob = einsum(J_regressor, verts_glob, \"j v, l v i -> l j i\") # (L, J, 3)\n global_R, global_T, global_lights = get_global_cameras_static(\n verts_glob.cpu(),\n beta=2.0,\n cam_height_degree=20,\n target_center_height=1.0,\n )\n\n # -- rendering code -- #\n video_path = cfg.video_path\n length, width, height = get_video_lwh(video_path)\n _, _, K = create_camera_sensor(width, height, 24) # render as 24mm lens\n\n # renderer\n renderer = Renderer(width, height, device=\"cuda\", faces=faces_smpl, K=K)\n # renderer = Renderer(width, height, device=\"cuda\", faces=faces_smpl, K=K, bin_size=0)\n\n # -- render mesh -- #\n scale, cx, cz = get_ground_params_from_points(joints_glob[:, 0], verts_glob)\n renderer.set_ground(scale * 1.5, cx, cz)\n color = torch.ones(3).float().cuda() * 0.8\n\n render_length = length if not debug_cam else 8\n writer = get_writer(global_video_path, fps=30, crf=CRF)\n for i in tqdm(range(render_length), desc=f\"Rendering Global\"):\n cameras = renderer.create_camera(global_R[i], global_T[i])\n img = renderer.render_with_ground(verts_glob[[i]], color[None], cameras, global_lights)\n writer.write_frame(img)\n writer.close()\n\n\nif __name__ == \"__main__\":\n cfg = parse_args_to_cfg()\n paths = cfg.paths\n Log.info(f\"[GPU]: {torch.cuda.get_device_name()}\")\n Log.info(f'[GPU]: {torch.cuda.get_device_properties(\"cuda\")}')\n\n # ===== Preprocess and save to disk ===== #\n run_preprocess(cfg)\n data = load_data_dict(cfg)\n\n # ===== HMR4D ===== #\n if not Path(paths.hmr4d_results).exists():\n Log.info(\"[HMR4D] Predicting\")\n model: DemoPL = hydra.utils.instantiate(cfg.model, _recursive_=False)\n model.load_pretrained_model(cfg.ckpt_path)\n model = model.eval().cuda()\n tic = Log.sync_time()\n pred = model.predict(data, static_cam=cfg.static_cam)\n pred = detach_to_cpu(pred)\n np.save(os.path.join(cfg.output_dir, 'smpl_orient.npy'), pred['smpl_params_incam']['global_orient'].numpy())\n np.save(os.path.join(cfg.output_dir, 'smpl_transl.npy'), pred['smpl_params_incam']['transl'].numpy())\n data_time = data[\"length\"] / 30\n Log.info(f\"[HMR4D] Elapsed: {Log.sync_time() - tic:.2f}s for data-length={data_time:.1f}s\")\n torch.save(pred, paths.hmr4d_results)\n\n # ===== Render ===== #\n render_incam(cfg)\n render_global(cfg)\n if not Path(paths.incam_global_horiz_video).exists():\n Log.info(\"[Merge Videos]\")\n merge_videos_horizontal([paths.incam_video, paths.global_video], paths.incam_global_horiz_video)\n" }, { "path": "eval/GVHMR/tools/demo/demo_folder.py", "content": "import argparse\nfrom pathlib import Path\nfrom tqdm import tqdm\nfrom hmr4d.utils.pylogger import Log\nimport subprocess\nimport os\n\n\nif __name__ == \"__main__\":\n parser = argparse.ArgumentParser()\n parser.add_argument(\"-f\", \"--folder\", type=str)\n parser.add_argument(\"-d\", \"--output_root\", type=str, default=None)\n parser.add_argument(\"-s\", \"--static_cam\", action=\"store_true\", help=\"If true, skip DPVO\")\n args = parser.parse_args()\n\n output_root = args.output_root\n\n sub_folders = os.listdir(args.folder)\n mp4_paths = []\n for sub_folder in sub_folders:\n files = os.listdir(os.path.join(args.folder, sub_folder))\n for file in files:\n if file.endswith('.mp4'):\n mp4_path = os.path.join(args.folder, sub_folder, file)\n mp4_paths.append(mp4_path)\n\n # Run demo.py for each .mp4 file\n Log.info(f\"Found {len(mp4_paths)} .mp4 files in {args.folder}\")\n for mp4_path in tqdm(mp4_paths):\n try:\n command = [\"python\", \"tools/demo/demo.py\", \"--video\", str(mp4_path)]\n if output_root is not None:\n command += [\"--output_root\", output_root]\n if args.static_cam:\n command += [\"-s\"]\n Log.info(f\"Running: {' '.join(command)}\")\n subprocess.run(command, env=dict(os.environ), check=True)\n except:\n continue\n" }, { "path": "eval/GVHMR/tools/eval_pose.py", "content": "# ParticleSfM\n# Copyright (C) 2022 ByteDance Inc.\n\n# This program is free software: you can redistribute it and/or modify\n# it under the terms of the GNU General Public License as published by\n# the Free Software Foundation, either version 3 of the License, or\n# (at your option) any later version.\n\n# This program is distributed in the hope that it will be useful,\n# but WITHOUT ANY WARRANTY; without even the implied warranty of\n# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the\n# GNU General Public License for more details.\n\n# You should have received a copy of the GNU General Public License\n# along with this program. If not, see .\n\nimport os\nfrom os.path import dirname\nimport numpy as np\nimport argparse\nimport tqdm\n\nimport numpy as np\nimport imageio\nimport os\nfrom glob import glob\nfrom multiprocessing import Pool\nimport json\nimport math\nfrom scipy.spatial.transform import Rotation as R\n\ndef batch_rotation_matrix_angle_error(R1_batch, R2_batch):\n\n assert R1_batch.shape == R2_batch.shape\n assert R1_batch.shape[1:] == (3, 3)\n \n B = R1_batch.shape[0]\n angle_errors = np.zeros(B) \n \n for i in range(B):\n R_relative = np.dot(R1_batch[i].T, R2_batch[i])\n \n r = R.from_matrix(R_relative)\n angle_error = r.magnitude()\n \n # angle_errors[i] = np.degrees(angle_error)\n angle_errors[i] = angle_error\n \n return angle_errors\n\ndef normalize(x):\n return x / np.linalg.norm(x)\n\ndef viewmatrix(z, up, pos):\n vec2 = normalize(z)\n vec1_avg = up\n vec0 = normalize(np.cross(vec1_avg, vec2))\n vec1 = normalize(np.cross(vec2, vec0))\n m = np.stack([vec0, vec1, vec2, pos], 1)\n return m\n\ndef matrix_to_euler_angles(matrix):\n sy = math.sqrt(matrix[0][0] * matrix[0][0] + matrix[1][0] * matrix[1][0])\n singular = sy < 1e-6\n\n if not singular:\n x = math.atan2(matrix[2][1], matrix[2][2])\n y = math.atan2(-matrix[2][0], sy)\n z = math.atan2(matrix[1][0], matrix[0][0])\n else:\n x = math.atan2(-matrix[1][2], matrix[1][1])\n y = math.atan2(-matrix[2][0], sy)\n z = 0\n\n return math.degrees(x), math.degrees(y), math.degrees(z)\n\ndef eul2rot(theta) :\n\n R = np.array([[np.cos(theta[1])*np.cos(theta[2]), np.sin(theta[0])*np.sin(theta[1])*np.cos(theta[2]) - np.sin(theta[2])*np.cos(theta[0]), np.sin(theta[1])*np.cos(theta[0])*np.cos(theta[2]) + np.sin(theta[0])*np.sin(theta[2])],\n [np.sin(theta[2])*np.cos(theta[1]), np.sin(theta[0])*np.sin(theta[1])*np.sin(theta[2]) + np.cos(theta[0])*np.cos(theta[2]), np.sin(theta[1])*np.sin(theta[2])*np.cos(theta[0]) - np.sin(theta[0])*np.cos(theta[2])],\n [-np.sin(theta[1]), np.sin(theta[0])*np.cos(theta[1]), np.cos(theta[0])*np.cos(theta[1])]])\n\n return R.T\n\ndef extract_location_rotation(data):\n results = {}\n for key, value in data.items():\n matrix = parse_matrix(value)\n location = np.array([matrix[3][0], matrix[3][1], matrix[3][2]])\n rotation = eul2rot(matrix_to_euler_angles(matrix))\n transofmed_matrix = np.identity(4)\n transofmed_matrix[:3,3] = location\n transofmed_matrix[:3,:3] = rotation\n results[key] = transofmed_matrix\n return results\n\ndef parse_matrix(matrix_str):\n rows = matrix_str.strip().split('] [')\n matrix = []\n for row in rows:\n row = row.replace('[', '').replace(']', '')\n matrix.append(list(map(float, row.split())))\n return np.array(matrix)\n\ndef batch_axis_angle_to_rotation_matrix(r_batch):\n batch_size = r_batch.shape[0]\n rotation_matrices = []\n \n for i in range(batch_size):\n r = r_batch[i]\n theta = np.linalg.norm(r)\n if theta == 0:\n rotation_matrices.append(np.eye(3))\n else:\n k = r / theta \n kx, ky, kz = k\n \n K = np.array([\n [0, -kz, ky],\n [kz, 0, -kx],\n [-ky, kx, 0]\n ])\n \n # Rodrigues formulation\n R = np.eye(3) + np.sin(theta) * K + (1 - np.cos(theta)) * np.dot(K, K)\n rotation_matrices.append(R)\n \n return np.array(rotation_matrices)\n\n\nif __name__ == '__main__' :\n\n parser = argparse.ArgumentParser()\n parser.add_argument(\"-f\", \"--folder\", type=str)\n args = parser.parse_args()\n\n folder_path = args.folder\n video_files = os.listdir(folder_path)\n for video_file in video_files:\n if video_file.endswith('txt'):\n video_files.remove(video_file)\n num_video_files = len(video_files)\n\n with open(('eval_gt_poses.json'), 'r') as f: eval_gt_poses = json.load(f)\n transl_err_all, rotat_err_all = 0, 0\n\n for video_file in tqdm.tqdm(sorted(video_files)):\n\n obj_poses = np.array(eval_gt_poses[video_file])\n \n start_frame_ind = 10\n sample_n_frames = 77\n frame_indices = np.linspace(start_frame_ind, start_frame_ind + sample_n_frames - 1, sample_n_frames, dtype=int)\n obj_poses = obj_poses[frame_indices]\n \n # load smpl pose\n video_path = os.path.join(folder_path, video_file)\n smpl_poses = np.zeros_like(obj_poses)\n smpl_poses[:,3,3] = 1.\n obj_rotats = np.load(os.path.join(video_path, 'smpl_orient.npy'))\n smpl_poses[:,:3,:3] = batch_axis_angle_to_rotation_matrix(obj_rotats)\n smpl_poses[:,:3,3] = np.load(os.path.join(video_path, 'smpl_transl.npy'))\n\n # align y-axis orientation\n smpl_poses[:,:3,3][:,1] *= -1.\n smpl_poses[:,:,:2] *= -1.\n\n # align pose translation\n translation_bias = smpl_poses[0,:3,3] - obj_poses[0,:3,3]\n smpl_poses[:,:3,3] -= translation_bias\n\n # evaluation\n transl_err = np.linalg.norm(smpl_poses[:,:3,3] - obj_poses[:,:3,3],ord=2,axis=1).mean()\n rotat_err = batch_rotation_matrix_angle_error(obj_poses[:,:3,:3],smpl_poses[:,:3,:3]).mean()\n\n transl_err_all += transl_err\n rotat_err_all += rotat_err\n\n print(video_path)\n print('transl_err:{:.3f}'.format(transl_err))\n print('rotat_err:{:.3f}'.format(rotat_err))\n\n print('transl_err_all:{:.3f}'.format(transl_err_all/num_video_files))\n print('rotat_err_all:{:.3f}'.format(rotat_err_all/num_video_files))\n" }, { "path": "eval/GVHMR/tools/train.py", "content": "import hydra\nimport pytorch_lightning as pl\nfrom omegaconf import DictConfig, OmegaConf\nfrom pytorch_lightning.callbacks.checkpoint import Checkpoint\n\nfrom hmr4d.utils.pylogger import Log\nfrom hmr4d.configs import register_store_gvhmr\nfrom hmr4d.utils.vis.rich_logger import print_cfg\nfrom hmr4d.utils.net_utils import load_pretrained_model, get_resume_ckpt_path\n\n\ndef get_callbacks(cfg: DictConfig) -> list:\n \"\"\"Parse and instantiate all the callbacks in the config.\"\"\"\n if not hasattr(cfg, \"callbacks\") or cfg.callbacks is None:\n return None\n # Handle special callbacks\n enable_checkpointing = cfg.pl_trainer.get(\"enable_checkpointing\", True)\n # Instantiate all the callbacks\n callbacks = []\n for callback in cfg.callbacks.values():\n if callback is not None:\n cb = hydra.utils.instantiate(callback, _recursive_=False)\n # skip when disable checkpointing and the callback is Checkpoint\n if not enable_checkpointing and isinstance(cb, Checkpoint):\n continue\n else:\n callbacks.append(cb)\n return callbacks\n\n\ndef train(cfg: DictConfig) -> None:\n \"\"\"Train/Test\"\"\"\n Log.info(f\"[Exp Name]: {cfg.exp_name}\")\n if cfg.task == \"fit\":\n Log.info(f\"[GPU x Batch] = {cfg.pl_trainer.devices} x {cfg.data.loader_opts.train.batch_size}\")\n pl.seed_everything(cfg.seed)\n\n # preparation\n datamodule: pl.LightningDataModule = hydra.utils.instantiate(cfg.data, _recursive_=False)\n model: pl.LightningModule = hydra.utils.instantiate(cfg.model, _recursive_=False)\n if cfg.ckpt_path is not None:\n load_pretrained_model(model, cfg.ckpt_path)\n\n # PL callbacks and logger\n callbacks = get_callbacks(cfg)\n has_ckpt_cb = any([isinstance(cb, Checkpoint) for cb in callbacks])\n if not has_ckpt_cb and cfg.pl_trainer.get(\"enable_checkpointing\", True):\n Log.warning(\"No checkpoint-callback found. Disabling PL auto checkpointing.\")\n cfg.pl_trainer = {**cfg.pl_trainer, \"enable_checkpointing\": False}\n logger = hydra.utils.instantiate(cfg.logger, _recursive_=False)\n\n # PL-Trainer\n if cfg.task == \"test\":\n Log.info(\"Test mode forces full-precision.\")\n cfg.pl_trainer = {**cfg.pl_trainer, \"precision\": 32}\n trainer = pl.Trainer(\n accelerator=\"gpu\",\n logger=logger if logger is not None else False,\n callbacks=callbacks,\n **cfg.pl_trainer,\n )\n\n if cfg.task == \"fit\":\n resume_path = None\n if cfg.resume_mode is not None:\n resume_path = get_resume_ckpt_path(cfg.resume_mode, ckpt_dir=cfg.callbacks.model_checkpoint.dirpath)\n Log.info(f\"Resume training from {resume_path}\")\n Log.info(\"Start Fitiing...\")\n trainer.fit(model, datamodule.train_dataloader(), datamodule.val_dataloader(), ckpt_path=resume_path)\n elif cfg.task == \"test\":\n Log.info(\"Start Testing...\")\n trainer.test(model, datamodule.test_dataloader())\n else:\n raise ValueError(f\"Unknown task: {cfg.task}\")\n\n Log.info(\"End of script.\")\n\n\n@hydra.main(version_base=\"1.3\", config_path=\"../hmr4d/configs\", config_name=\"train\")\ndef main(cfg) -> None:\n print_cfg(cfg, use_rich=True)\n train(cfg)\n\n\nif __name__ == \"__main__\":\n register_store_gvhmr()\n main()\n" }, { "path": "eval/GVHMR/tools/unitest/make_hydra_cfg.py", "content": "from hmr4d.configs import parse_args_to_cfg, register_store_gvhmr\nfrom hmr4d.utils.vis.rich_logger import print_cfg\n\nif __name__ == \"__main__\":\n register_store_gvhmr()\n cfg = parse_args_to_cfg()\n print_cfg(cfg, use_rich=True)\n" }, { "path": "eval/GVHMR/tools/unitest/run_dataset.py", "content": "import torch\nfrom torch.utils.data import DataLoader\nfrom tqdm import tqdm\n\n\ndef get_dataset(DATA_TYPE):\n if DATA_TYPE == \"BEDLAM_V2\":\n from hmr4d.dataset.bedlam.bedlam import BedlamDatasetV2\n\n return BedlamDatasetV2()\n\n if DATA_TYPE == \"3DPW_TRAIN\":\n from hmr4d.dataset.threedpw.threedpw_motion_train import ThreedpwSmplDataset\n\n return ThreedpwSmplDataset()\n\nif __name__ == \"__main__\":\n DATA_TYPE = \"3DPW_TRAIN\"\n dataset = get_dataset(DATA_TYPE)\n print(len(dataset))\n\n data = dataset[0]\n\n from hmr4d.datamodule.mocap_trainX_testY import collate_fn\n\n loader = DataLoader(\n dataset,\n shuffle=False,\n num_workers=0,\n persistent_workers=False,\n pin_memory=False,\n batch_size=1,\n collate_fn=collate_fn,\n )\n i = 0\n for batch in tqdm(loader):\n i += 1\n # if i == 20:\n # raise AssertionError\n # time.sleep(0.2)\n pass\n" }, { "path": "eval/GVHMR/tools/video/merge_folder.py", "content": "\"\"\"This script will glob two folder, check the mp4 files are one-to-one match precisely, then call merge_horizontal.py to merge them one by one\"\"\"\n\nimport os\nimport argparse\nfrom pathlib import Path\n\n\ndef main():\n parser = argparse.ArgumentParser()\n parser.add_argument(\"input_dir1\", type=str)\n parser.add_argument(\"input_dir2\", type=str)\n parser.add_argument(\"output_dir\", type=str)\n parser.add_argument(\"--vertical\", action=\"store_true\") # By default use horizontal\n args = parser.parse_args()\n\n # Check input\n input_dir1 = Path(args.input_dir1)\n input_dir2 = Path(args.input_dir2)\n assert input_dir1.exists()\n assert input_dir2.exists()\n video_paths1 = sorted(input_dir1.glob(\"*.mp4\"))\n video_paths2 = sorted(input_dir2.glob(\"*.mp4\"))\n assert len(video_paths1) == len(video_paths2)\n for path1, path2 in zip(video_paths1, video_paths2):\n assert path1.stem == path2.stem\n\n # Merge to output\n output_dir = Path(args.output_dir)\n output_dir.mkdir(parents=True, exist_ok=True)\n\n for path1, path2 in zip(video_paths1, video_paths2):\n out_path = output_dir / f\"{path1.stem}.mp4\"\n in_paths = [str(path1), str(path2)]\n print(f\"Merging {in_paths} to {out_path}\")\n if args.vertical:\n os.system(f\"python tools/video/merge_vertical.py {' '.join(in_paths)} -o {out_path}\")\n else:\n os.system(f\"python tools/video/merge_horizontal.py {' '.join(in_paths)} -o {out_path}\")\n\n\nif __name__ == \"__main__\":\n main()\n" }, { "path": "eval/GVHMR/tools/video/merge_horizontal.py", "content": "import argparse\nfrom hmr4d.utils.video_io_utils import merge_videos_horizontal\n\n\ndef parse_args():\n \"\"\"python tools/video/merge_horizontal.py a.mp4 b.mp4 c.mp4 -o out.mp4\"\"\"\n parser = argparse.ArgumentParser()\n parser.add_argument(\"input_videos\", nargs=\"+\", help=\"Input video paths\")\n parser.add_argument(\"-o\", \"--output\", type=str, required=True, help=\"Output video path\")\n return parser.parse_args()\n\n\nif __name__ == \"__main__\":\n args = parse_args()\n merge_videos_horizontal(args.input_videos, args.output)\n" }, { "path": "eval/GVHMR/tools/video/merge_vertical.py", "content": "import argparse\nfrom hmr4d.utils.video_io_utils import merge_videos_vertical\n\n\ndef parse_args():\n \"\"\"python tools/video/merge_vertical.py a.mp4 b.mp4 c.mp4 -o out.mp4\"\"\"\n parser = argparse.ArgumentParser()\n parser.add_argument(\"input_videos\", nargs=\"+\", help=\"Input video paths\")\n parser.add_argument(\"-o\", \"--output\", type=str, required=True, help=\"Output video path\")\n return parser.parse_args()\n\n\nif __name__ == \"__main__\":\n args = parse_args()\n merge_videos_vertical(args.input_videos, args.output)\n" }, { "path": "eval/common_metrics_on_video_quality/.gitignore", "content": "__pycache__" }, { "path": "eval/common_metrics_on_video_quality/README.md", "content": "# common_metrics_on_video_quality\n\nYou can easily calculate the following video quality metrics:\n\n- **FVD**: Frechét Video Distance\n- **SSIM**: structural similarity index measure\n- **LPIPS**: learned perceptual image patch similarity\n- **PSNR**: peak-signal-to-noise ratio\n\nAs for FVD\n 1. The codebase refers to [MVCD](https://github.com/voletiv/mcvd-pytorch) and other websites and projects, I've just extracted the part of it that's relevant to the calculation. This code can be used to evaluate FVD scores for generative or predictive models. \n 2. Now **we have supported 2 pytorch-based FVD implementations** ([videogpt](https://github.com/wilson1yan/VideoGPT) and [styleganv](https://github.com/universome/stylegan-v), see issue [#4](https://github.com/JunyaoHu/common_metrics_on_video_quality/issues/4)). Their calculations are almost identical, and the difference is negligible.\n 3. FVD calculates the feature distance between two sets of videos. (the I3D features of each video are do not go through the softmax() function, and the size of the last dimension is 400, not 1024)\n\nAnd...\n\n- This project supports grayscale and RGB videos.\n- This project supports Ubuntu, but maybe something is wrong with Windows. If you can solve it, welcome any PR.\n- **If the project cannot run correctly, please give me an issue or PR~**\n- For more details see below Notice.\n\n# Example\n\n8 videos of a batch, 10 frames, 3 channels, 64x64 size.\n\n```\nimport torch\nfrom calculate_fvd import calculate_fvd\nfrom calculate_psnr import calculate_psnr\nfrom calculate_ssim import calculate_ssim\nfrom calculate_lpips import calculate_lpips\n\nNUMBER_OF_VIDEOS = 8\nVIDEO_LENGTH = 30\nCHANNEL = 3\nSIZE = 64\nvideos1 = torch.zeros(NUMBER_OF_VIDEOS, VIDEO_LENGTH, CHANNEL, SIZE, SIZE, requires_grad=False)\nvideos2 = torch.ones(NUMBER_OF_VIDEOS, VIDEO_LENGTH, CHANNEL, SIZE, SIZE, requires_grad=False)\ndevice = torch.device(\"cuda\")\ndevice = torch.device(\"cpu\")\n\nimport json\nresult = {}\nresult['fvd'] = calculate_fvd(videos1, videos2, device, method='styleganv')\n# result['fvd'] = calculate_fvd(videos1, videos2, device, method='videogpt')\nresult['ssim'] = calculate_ssim(videos1, videos2)\nresult['psnr'] = calculate_psnr(videos1, videos2)\nresult['lpips'] = calculate_lpips(videos1, videos2, device)\nprint(json.dumps(result, indent=4))\n```\n\nIt means we calculate:\n \n- `FVD-frames[:10]`, `FVD-frames[:11]`, ..., `FVD-frames[:30]` \n- `avg-PSNR/SSIM/LPIPS-frame[0]`, `avg-PSNR/SSIM/LPIPS-frame[1]`, ..., `avg-PSNR/SSIM/LPIPS-frame[:30]`, and their std.\n\nWe cannot calculate `FVD-frames[:8]`, and it will pass when calculating, see ps.6.\n\nThe result shows: a all-zero matrix and a all-one matrix, their FVD-30 (FVD[:30]) is 151.17 (styleganv method). We also calculate their standard deviation. Other metrics are the same. And we use the calculation method of styleganv.\n\n```\n{\n \"fvd\": {\n \"value\": {\n \"10\": 570.07320378183,\n \"11\": 486.1906542471159,\n \"12\": 552.3373915075898,\n \"13\": 146.6242330185728,\n \"14\": 172.57268402948895,\n \"15\": 133.88932632116126,\n \"16\": 153.11023578170108,\n \"17\": 357.56400892781204,\n \"18\": 382.1335612721498,\n \"19\": 306.7100176942531,\n \"20\": 338.18221898178774,\n \"21\": 77.95587603163293,\n \"22\": 82.49997632357349,\n \"23\": 64.41624523513073,\n \"24\": 66.08097153313875,\n \"25\": 314.4341061962642,\n \"26\": 316.8616746151064,\n \"27\": 288.884418528541,\n \"28\": 287.8192683223724,\n \"29\": 152.15076552354864,\n \"30\": 151.16806952692093\n },\n \"video_setting\": [\n 8,\n 3,\n 30,\n 64,\n 64\n ],\n \"video_setting_name\": \"batch_size, channel, time, heigth, width\"\n },\n \"video_setting\": [\n 8,\n 3,\n 30,\n 64,\n 64\n ],\n \"video_setting_name\": \"batch_size, channel, time, heigth, width\"\n },\n \"ssim\": {\n \"value\": {\n \"0\": 9.999000099990664e-05,\n ...,\n \"29\": 9.999000099990664e-05\n },\n \"value_std\": {\n \"0\": 0.0,\n ...,\n \"29\": 0.0\n },\n \"video_setting\": [\n 30,\n 3,\n 64,\n 64\n ],\n \"video_setting_name\": \"time, channel, heigth, width\"\n },\n \"psnr\": {\n \"value\": {\n \"0\": 0.0,\n ...,\n \"29\": 0.0\n },\n \"value_std\": {\n \"0\": 0.0,\n ...,\n \"29\": 0.0\n },\n \"video_setting\": [\n 30,\n 3,\n 64,\n 64\n ],\n \"video_setting_name\": \"time, channel, heigth, width\"\n },\n \"lpips\": {\n \"value\": {\n \"0\": 0.8140146732330322,\n ...,\n \"29\": 0.8140146732330322\n },\n \"value_std\": {\n \"0\": 0.0,\n ...,\n \"29\": 0.0\n },\n \"video_setting\": [\n 30,\n 3,\n 64,\n 64\n ],\n \"video_setting_name\": \"time, channel, heigth, width\"\n }\n}\n```\n\n# Notice\n\n1. You should `pip install lpips` first.\n3. Make sure the pixel value of videos should be in [0, 1].\n2. If you have something wrong with downloading FVD pre-trained model, you should manually download any of the following and put it into FVD folder. \n - `i3d_torchscript.pt` from [here](https://www.dropbox.com/s/ge9e5ujwgetktms/i3d_torchscript.pt) \n - `i3d_pretrained_400.pt` from [here](https://onedrive.live.com/download?cid=78EEF3EB6AE7DBCB&resid=78EEF3EB6AE7DBCB%21199&authkey=AApKdFHPXzWLNyI)\n4. For grayscale videos, we multiply to 3 channels [as it says](https://github.com/richzhang/PerceptualSimilarity/issues/23#issuecomment-492368812).\n5. We average SSIM when images have 3 channels, ssim is the only metric extremely sensitive to gray being compared to b/w.\n6. Because the i3d model downsamples in the time dimension, `frames_num` should > 10 when calculating FVD, so FVD calculation begins from 10-th frame, like upper example.\n7. You had better use `scipy==1.7.3/1.9.3`, if you use 1.11.3, **you will calculate a WRONG FVD VALUE!!!**\n8. If you are running demo.py on a multi-GPU machine, remember to export CUDA_VISIBLE_DEVICES=0, see [here](https://github.com/JunyaoHu/common_metrics_on_video_quality/issues/13).\n\n# Star Trend\n## Star History\n\n[![Star History Chart](https://api.star-history.com/svg?repos=JunyaoHu/common_metrics_on_video_quality&type=Date)](https://star-history.com/#JunyaoHu/common_metrics_on_video_quality&Date)\n" }, { "path": "eval/common_metrics_on_video_quality/calculate_clip.py", "content": "import cv2\nfrom PIL import Image\nimport torch\nfrom transformers import CLIPProcessor, CLIPModel\nimport json\nimport os\nfrom tqdm import tqdm\nimport torch\nimport clip\nfrom PIL import Image\nimport cv2\nimport numpy as np\nimport os\nimport argparse\n\ndevice = \"cuda\" if torch.cuda.is_available() else \"cpu\"\n\nmodel = CLIPModel.from_pretrained(\"openai/clip-vit-base-patch32\").to(device)\nprocessor = CLIPProcessor.from_pretrained(\"openai/clip-vit-base-patch32\")\n\ndef get_video_scores(video_path, prompt):\n video = cv2.VideoCapture(video_path)\n texts = [prompt]\n clip_score_list = []\n while True:\n ret, frame = video.read()\n\n if ret:\n image = Image.fromarray(cv2.cvtColor(frame, cv2.COLOR_BGR2RGB))\n inputs = processor(text=texts, images=[image], return_tensors=\"pt\", padding=True, truncation=True).to(device)\n logits_per_image = model(**inputs).logits_per_image\n clip_score = logits_per_image.item()\n clip_score_list.append(clip_score)\n else:\n break\n\n video.release()\n return sum(clip_score_list) / len(clip_score_list)\n\n\nparser = argparse.ArgumentParser()\nparser.add_argument(\"-v_f\", \"--videos_folder\", type=str)\nargs = parser.parse_args()\n\nvideos_folder_path = args.videos_folder\nprompts_path = '/ytech_m2v2_hdd/fuxiao/scenectrl/common_metrics_on_video_quality/eval_prompts.json'\nwith open(prompts_path, \"r\", encoding=\"utf-8\") as f: prompts_dict = json.load(f)\n\nsub_folders = os.listdir(videos_folder_path)\nvideos_name = []\nfor sub_folder in sub_folders:\n files = os.listdir(os.path.join(videos_folder_path, sub_folder))\n for file in files:\n if file.endswith('.mp4'):\n video_name = os.path.join(sub_folder, file)\n videos_name.append(video_name)\n\nnum_videos = len(videos_name)\n\nprompts = []\nvideo_paths = []\nfor video_name in videos_name:\n prompt = prompts_dict[video_name.split('/')[0]]\n video_path = os.path.join(videos_folder_path, video_name)\n prompts.append(prompt)\n video_paths.append(video_path)\n\nimport csv\nCLIP_T = True\nif CLIP_T:\n scores = []\n for i in tqdm(range(num_videos)):\n # 加载图片\n video_path = video_paths[i]\n \n # 准备文本\n texts = prompts[i]\n score = get_video_scores(video_path, texts)\n scores.append(score)\n\n print(f\"CLIP-SIM: {sum(scores)/len(scores)/100.}\")\n #### CLIP-T ####\n # basemodel: 33.44" }, { "path": "eval/common_metrics_on_video_quality/calculate_fvd.py", "content": "import numpy as np\nimport torch\nfrom tqdm import tqdm\n\ndef trans(x):\n # if greyscale images add channel\n if x.shape[-3] == 1:\n x = x.repeat(1, 1, 3, 1, 1)\n\n # permute BTCHW -> BCTHW\n x = x.permute(0, 2, 1, 3, 4) \n\n return x\n\ndef calculate_fvd(videos1, videos2, device, method='styleganv'):\n\n if method == 'styleganv':\n from fvd.styleganv.fvd import get_fvd_feats, frechet_distance, load_i3d_pretrained\n elif method == 'videogpt':\n from fvd.videogpt.fvd import load_i3d_pretrained\n from fvd.videogpt.fvd import get_fvd_logits as get_fvd_feats\n from fvd.videogpt.fvd import frechet_distance\n\n print(\"calculate_fvd...\")\n\n # videos [batch_size, timestamps, channel, h, w]\n \n assert videos1.shape == videos2.shape\n\n i3d = load_i3d_pretrained(device=device)\n fvd_results = []\n\n # support grayscale input, if grayscale -> channel*3\n # BTCHW -> BCTHW\n # videos -> [batch_size, channel, timestamps, h, w]\n\n videos1 = trans(videos1)\n videos2 = trans(videos2)\n\n fvd_results = {}\n\n # for calculate FVD, each clip_timestamp must >= 10\n\n # get a video clip\n # videos_clip [batch_size, channel, timestamps[:clip], h, w]\n videos_clip1 = videos1[:, :, :]\n videos_clip2 = videos2[:, :, :]\n\n # get FVD features\n feats1 = get_fvd_feats(videos_clip1, i3d=i3d, device=device)\n feats2 = get_fvd_feats(videos_clip2, i3d=i3d, device=device)\n \n # calculate FVD when timestamps[:clip]\n fvd_results = frechet_distance(feats1, feats2)\n \n result = {\n \"value\": fvd_results,\n \"video_setting\": videos1.shape,\n \"video_setting_name\": \"batch_size, channel, time, heigth, width\",\n }\n\n return result\n\n# test code / using example\n\ndef main():\n NUMBER_OF_VIDEOS = 8\n VIDEO_LENGTH = 50\n CHANNEL = 3\n SIZE = 64\n videos1 = torch.zeros(NUMBER_OF_VIDEOS, VIDEO_LENGTH, CHANNEL, SIZE, SIZE, requires_grad=False)\n videos2 = torch.ones(NUMBER_OF_VIDEOS, VIDEO_LENGTH, CHANNEL, SIZE, SIZE, requires_grad=False)\n device = torch.device(\"cuda\")\n # device = torch.device(\"cpu\")\n\n import json\n result = calculate_fvd(videos1, videos2, device, method='videogpt')\n print(json.dumps(result, indent=4))\n\n result = calculate_fvd(videos1, videos2, device, method='styleganv')\n print(json.dumps(result, indent=4))\n\nif __name__ == \"__main__\":\n main()" }, { "path": "eval/common_metrics_on_video_quality/calculate_fvd_styleganv.py", "content": "import torch\nfrom calculate_fvd import calculate_fvd\nfrom calculate_psnr import calculate_psnr\nfrom calculate_ssim import calculate_ssim\nfrom calculate_lpips import calculate_lpips\nimport argparse\nimport os\nimport cv2\nimport decord\nimport numpy as np\nimport tqdm\nimport glob\nimport copy\nos.environ[\"CUDA_VISIBLE_DEVICES\"]=\"0\"\n\nparser = argparse.ArgumentParser()\nparser.add_argument(\"-v1_f\", \"--videos1_folder\", type=str)\nparser.add_argument(\"-v2_f\", \"--videos2_folder\", type=str)\nargs = parser.parse_args()\n\nvideos1_folder_path = args.videos1_folder\nvideos2_folder_path = args.videos2_folder\n\nsub_folders = os.listdir(videos1_folder_path)\nvideos_name = []\nfor sub_folder in sub_folders:\n files = os.listdir(os.path.join(videos1_folder_path, sub_folder))\n for file in files:\n if file.endswith('.mp4'):\n video_name = os.path.join(sub_folder, file)\n videos_name.append(video_name)\n\nbase_dir = os.path.dirname(videos2_folder_path)\nbase_name = os.path.basename(videos2_folder_path)\n\nos.makedirs(f'{base_dir}/eval_1', exist_ok=True)\nos.makedirs(f'{base_dir}/eval_2', exist_ok=True)\n# ps: pixel value should be in [0, 1]!\nNUMBER_OF_VIDEOS = len(videos_name)\nVIDEO_LENGTH = 77\nCHANNEL = 3\nH_SIZE = 384\nW_SIZE = 672\nvideos1 = torch.zeros(NUMBER_OF_VIDEOS, VIDEO_LENGTH, CHANNEL, H_SIZE, W_SIZE, requires_grad=False)\nvideos2 = torch.ones(NUMBER_OF_VIDEOS, VIDEO_LENGTH, CHANNEL, H_SIZE, W_SIZE, requires_grad=False)\n\nfor video_idx, video_name in tqdm.tqdm(enumerate(videos_name)):\n\n print(video_name)\n video_frames_path = os.path.join(videos1_folder_path, video_name)\n cap = cv2.VideoCapture(video_frames_path)\n height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))\n width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))\n ctx = decord.cpu(0)\n reader = decord.VideoReader(video_frames_path, ctx=ctx, height=height, width=width)\n frame_indexes = [frame_idx for frame_idx in range(VIDEO_LENGTH)]\n try:\n video_chunk = reader.get_batch(frame_indexes).asnumpy() \n except:\n video_chunk = reader.get_batch(frame_indexes).numpy()\n for frame_idx in range(VIDEO_LENGTH):\n cv2.imwrite(f'{base_dir}/eval_1/{video_idx:03d}_{frame_idx:02d}.png', video_chunk[frame_idx][:,:,::-1])\n video_chunk = video_chunk.transpose(0,3,1,2)/255.\n\n videos1[video_idx] = torch.from_numpy(video_chunk)\n\n video_frames_path = os.path.join(videos2_folder_path, video_name) \n cap = cv2.VideoCapture(video_frames_path)\n height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))\n width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))\n ctx = decord.cpu(0)\n reader = decord.VideoReader(video_frames_path, ctx=ctx, height=height, width=width)\n frame_indexes = [frame_idx for frame_idx in range(VIDEO_LENGTH)]\n try:\n video_chunk = reader.get_batch(frame_indexes).asnumpy() \n except:\n video_chunk = reader.get_batch(frame_indexes).numpy()\n for frame_idx in range(VIDEO_LENGTH):\n cv2.imwrite(f'{base_dir}/eval_2/{video_idx:03d}_{frame_idx:02d}.png', video_chunk[frame_idx][:,:,::-1])\n \n video_chunk = video_chunk.transpose(0,3,1,2)/255.\n videos2[video_idx] = torch.from_numpy(video_chunk)\n\nif NUMBER_OF_VIDEOS == 1:\n videos1 = videos1.repeat(2,1,1,1,1)\n videos2 = videos2.repeat(2,1,1,1,1)\n\nprint('load videos done')\ndevice = torch.device(\"cuda\")\n\nimport json\nresult = {}\nresult['fvd_styleganv'] = calculate_fvd(videos1, videos2, device, method='styleganv')\n# result['fvd_videogpt'] = calculate_fvd(videos1, videos2, device, method='videogpt')\n\nfvd_value = result['fvd_styleganv']['value']\nprint(f'FVD: {fvd_value}')" }, { "path": "eval/common_metrics_on_video_quality/calculate_lpips.py", "content": "import numpy as np\nimport torch\nfrom tqdm import tqdm\nimport math\n\nimport torch\nimport lpips\n\nspatial = True # Return a spatial map of perceptual distance.\n\n# Linearly calibrated models (LPIPS)\nloss_fn = lpips.LPIPS(net='alex', spatial=spatial) # Can also set net = 'squeeze' or 'vgg'\n# loss_fn = lpips.LPIPS(net='alex', spatial=spatial, lpips=False) # Can also set net = 'squeeze' or 'vgg'\n\ndef trans(x):\n # if greyscale images add channel\n if x.shape[-3] == 1:\n x = x.repeat(1, 1, 3, 1, 1)\n\n # value range [0, 1] -> [-1, 1]\n x = x * 2 - 1\n\n return x\n\ndef calculate_lpips(videos1, videos2, device):\n # image should be RGB, IMPORTANT: normalized to [-1,1]\n print(\"calculate_lpips...\")\n\n assert videos1.shape == videos2.shape\n\n # videos [batch_size, timestamps, channel, h, w]\n\n # support grayscale input, if grayscale -> channel*3\n # value range [0, 1] -> [-1, 1]\n videos1 = trans(videos1)\n videos2 = trans(videos2)\n\n lpips_results = []\n\n for video_num in tqdm(range(videos1.shape[0])):\n # get a video\n # video [timestamps, channel, h, w]\n video1 = videos1[video_num]\n video2 = videos2[video_num]\n\n lpips_results_of_a_video = []\n for clip_timestamp in range(len(video1)):\n # get a img\n # img [timestamps[x], channel, h, w]\n # img [channel, h, w] tensor\n\n img1 = video1[clip_timestamp].unsqueeze(0).to(device)\n img2 = video2[clip_timestamp].unsqueeze(0).to(device)\n \n loss_fn.to(device)\n\n # calculate lpips of a video\n lpips_results_of_a_video.append(loss_fn.forward(img1, img2).mean().detach().cpu().tolist())\n lpips_results.append(lpips_results_of_a_video)\n \n lpips_results = np.array(lpips_results)\n \n lpips = {}\n lpips_std = {}\n\n for clip_timestamp in range(len(video1)):\n lpips[clip_timestamp] = np.mean(lpips_results[:,clip_timestamp])\n lpips_std[clip_timestamp] = np.std(lpips_results[:,clip_timestamp])\n\n\n result = {\n \"value\": lpips,\n \"value_std\": lpips_std,\n \"video_setting\": video1.shape,\n \"video_setting_name\": \"time, channel, heigth, width\",\n }\n\n return result\n\n# test code / using example\n\ndef main():\n NUMBER_OF_VIDEOS = 8\n VIDEO_LENGTH = 50\n CHANNEL = 3\n SIZE = 64\n videos1 = torch.zeros(NUMBER_OF_VIDEOS, VIDEO_LENGTH, CHANNEL, SIZE, SIZE, requires_grad=False)\n videos2 = torch.ones(NUMBER_OF_VIDEOS, VIDEO_LENGTH, CHANNEL, SIZE, SIZE, requires_grad=False)\n device = torch.device(\"cuda\")\n # device = torch.device(\"cpu\")\n\n import json\n result = calculate_lpips(videos1, videos2, device)\n print(json.dumps(result, indent=4))\n\nif __name__ == \"__main__\":\n main()" }, { "path": "eval/common_metrics_on_video_quality/calculate_psnr.py", "content": "import numpy as np\nimport torch\nfrom tqdm import tqdm\nimport math\n\ndef img_psnr(img1, img2):\n # [0,1]\n # compute mse\n # mse = np.mean((img1-img2)**2)\n mse = np.mean((img1 / 1.0 - img2 / 1.0) ** 2)\n # compute psnr\n if mse < 1e-10:\n return 100\n psnr = 20 * math.log10(1 / math.sqrt(mse))\n return psnr\n\ndef trans(x):\n return x\n\ndef calculate_psnr(videos1, videos2):\n print(\"calculate_psnr...\")\n\n # videos [batch_size, timestamps, channel, h, w]\n \n assert videos1.shape == videos2.shape\n\n videos1 = trans(videos1)\n videos2 = trans(videos2)\n\n psnr_results = []\n \n for video_num in tqdm(range(videos1.shape[0])):\n # get a video\n # video [timestamps, channel, h, w]\n video1 = videos1[video_num]\n video2 = videos2[video_num]\n\n psnr_results_of_a_video = []\n for clip_timestamp in range(len(video1)):\n # get a img\n # img [timestamps[x], channel, h, w]\n # img [channel, h, w] numpy\n\n img1 = video1[clip_timestamp].numpy()\n img2 = video2[clip_timestamp].numpy()\n \n # calculate psnr of a video\n psnr_results_of_a_video.append(img_psnr(img1, img2))\n\n psnr_results.append(psnr_results_of_a_video)\n \n psnr_results = np.array(psnr_results)\n \n psnr = {}\n psnr_std = {}\n\n for clip_timestamp in range(len(video1)):\n psnr[clip_timestamp] = np.mean(psnr_results[:,clip_timestamp])\n psnr_std[clip_timestamp] = np.std(psnr_results[:,clip_timestamp])\n\n result = {\n \"value\": psnr,\n \"value_std\": psnr_std,\n \"video_setting\": video1.shape,\n \"video_setting_name\": \"time, channel, heigth, width\",\n }\n\n return result\n\n# test code / using example\n\ndef main():\n NUMBER_OF_VIDEOS = 8\n VIDEO_LENGTH = 50\n CHANNEL = 3\n SIZE = 64\n videos1 = torch.zeros(NUMBER_OF_VIDEOS, VIDEO_LENGTH, CHANNEL, SIZE, SIZE, requires_grad=False)\n videos2 = torch.zeros(NUMBER_OF_VIDEOS, VIDEO_LENGTH, CHANNEL, SIZE, SIZE, requires_grad=False)\n\n import json\n result = calculate_psnr(videos1, videos2)\n print(json.dumps(result, indent=4))\n\nif __name__ == \"__main__\":\n main()" }, { "path": "eval/common_metrics_on_video_quality/calculate_ssim.py", "content": "import numpy as np\nimport torch\nfrom tqdm import tqdm\nimport cv2\n \ndef ssim(img1, img2):\n C1 = 0.01 ** 2\n C2 = 0.03 ** 2\n img1 = img1.astype(np.float64)\n img2 = img2.astype(np.float64)\n kernel = cv2.getGaussianKernel(11, 1.5)\n window = np.outer(kernel, kernel.transpose())\n mu1 = cv2.filter2D(img1, -1, window)[5:-5, 5:-5] # valid\n mu2 = cv2.filter2D(img2, -1, window)[5:-5, 5:-5]\n mu1_sq = mu1 ** 2\n mu2_sq = mu2 ** 2\n mu1_mu2 = mu1 * mu2\n sigma1_sq = cv2.filter2D(img1 ** 2, -1, window)[5:-5, 5:-5] - mu1_sq\n sigma2_sq = cv2.filter2D(img2 ** 2, -1, window)[5:-5, 5:-5] - mu2_sq\n sigma12 = cv2.filter2D(img1 * img2, -1, window)[5:-5, 5:-5] - mu1_mu2\n ssim_map = ((2 * mu1_mu2 + C1) * (2 * sigma12 + C2)) / ((mu1_sq + mu2_sq + C1) *\n (sigma1_sq + sigma2_sq + C2))\n return ssim_map.mean()\n \n \ndef calculate_ssim_function(img1, img2):\n # [0,1]\n # ssim is the only metric extremely sensitive to gray being compared to b/w \n if not img1.shape == img2.shape:\n raise ValueError('Input images must have the same dimensions.')\n if img1.ndim == 2:\n return ssim(img1, img2)\n elif img1.ndim == 3:\n if img1.shape[0] == 3:\n ssims = []\n for i in range(3):\n ssims.append(ssim(img1[i], img2[i]))\n return np.array(ssims).mean() \n elif img1.shape[0] == 1:\n return ssim(np.squeeze(img1), np.squeeze(img2))\n else:\n raise ValueError('Wrong input image dimensions.')\n\ndef trans(x):\n return x\n\ndef calculate_ssim(videos1, videos2):\n print(\"calculate_ssim...\")\n\n # videos [batch_size, timestamps, channel, h, w]\n \n assert videos1.shape == videos2.shape\n\n videos1 = trans(videos1)\n videos2 = trans(videos2)\n\n ssim_results = []\n \n for video_num in tqdm(range(videos1.shape[0])):\n # get a video\n # video [timestamps, channel, h, w]\n video1 = videos1[video_num]\n video2 = videos2[video_num]\n\n ssim_results_of_a_video = []\n for clip_timestamp in range(len(video1)):\n # get a img\n # img [timestamps[x], channel, h, w]\n # img [channel, h, w] numpy\n\n img1 = video1[clip_timestamp].numpy()\n img2 = video2[clip_timestamp].numpy()\n \n # calculate ssim of a video\n ssim_results_of_a_video.append(calculate_ssim_function(img1, img2))\n\n ssim_results.append(ssim_results_of_a_video)\n\n ssim_results = np.array(ssim_results)\n\n ssim = {}\n ssim_std = {}\n\n for clip_timestamp in range(len(video1)):\n ssim[clip_timestamp] = np.mean(ssim_results[:,clip_timestamp])\n ssim_std[clip_timestamp] = np.std(ssim_results[:,clip_timestamp])\n\n result = {\n \"value\": ssim,\n \"value_std\": ssim_std,\n \"video_setting\": video1.shape,\n \"video_setting_name\": \"time, channel, heigth, width\",\n }\n\n return result\n\n# test code / using example\n\ndef main():\n NUMBER_OF_VIDEOS = 8\n VIDEO_LENGTH = 50\n CHANNEL = 3\n SIZE = 64\n videos1 = torch.zeros(NUMBER_OF_VIDEOS, VIDEO_LENGTH, CHANNEL, SIZE, SIZE, requires_grad=False)\n videos2 = torch.zeros(NUMBER_OF_VIDEOS, VIDEO_LENGTH, CHANNEL, SIZE, SIZE, requires_grad=False)\n device = torch.device(\"cuda\")\n\n import json\n result = calculate_ssim(videos1, videos2)\n print(json.dumps(result, indent=4))\n\nif __name__ == \"__main__\":\n main()" }, { "path": "eval/common_metrics_on_video_quality/download_eval_visual.sh", "content": "gdown https://drive.google.com/uc\\?id\\=1U2hd6qvwKLfp7c8yGgcTqdqrP_lKJElB\ngdown https://drive.google.com/uc\\?id\\=1jMH2-ZC0ZBgtqej5Sp-E5ebBIX7mk3Xz\ngdown https://drive.google.com/uc\\?id\\=1kfdCDA5koYh9g3IkCCHb4XPch2CJAwek\n\nunzip fvd.zip\nunzip eval_sets.zip\nunzip base_t2v_eval_sets.zip\n\nmv eval_sets eval_folder/\nmv base_t2v_eval_sets eval_folder/\n\nrm -rf *.zip" }, { "path": "eval/common_metrics_on_video_quality/eval_prompts.json", "content": "{\n \"0_D_loc5_530_t1n2_0212_Hemi12_1\": \"a man with short spiky brown hair, athletic build, a navy blue jacket, beige cargo pants, and black sneakers is moving in the fjord\",\n \"1_D_loc1_31_t1n6_001f_Hemi12_1\": \"a woman with long wavy blonde hair, petite figure, a red floral dress, white sandals, and a yellow shoulder bag is moving in the sunset beach\",\n \"2_D_loc1_296_t1n8_0128_Hemi12_1\": \"a man with a shaved head, broad shoulders, a gray graphic t-shirt, dark jeans, and brown leather boots is moving in the cave\",\n \"3_D_loc5_1100_t1n9_044c_Hemi12_1\": \"a woman with shoulder-length straight auburn hair, a slender figure, a green button-up blouse, black leggings, and white sneakers is moving in the snowy tundra\",\n \"4_D_loc5_590_t1n11_024e_Hemi12_1\": \"a man with messy black hair, tall frame, a plaid red and black shirt, faded blue jeans, and tan hiking boots is moving in the prairie\",\n \"5_D_loc5_1125_t1n13_0465_Hemi12_1\": \"a man with medium-length straight brown hair, tall and slender, a gray crew-neck t-shirt, beige trousers, and dark green sneakers is moving in the asian town\",\n \"6_D_loc3_1158_t1n15_0486_Hemi12_1\": \"a woman with short curly black hair, slender build, a pink hoodie, light gray joggers, and blue sneakers is moving in the rainforest\",\n \"7_D_loc4_1469_t1n17_05bd_Hemi12_1\": \"a man with short black wavy hair, lean figure, a green and yellow plaid shirt, dark brown pants, and black suede shoes is moving in the canyon\",\n \"8_D_loc1_351_t1n18_015f_Hemi12_1\": \"a man with curly black hair, muscular build, a dark green hoodie, gray joggers, and white running shoes is moving in the savanna\",\n \"9_D_loc5_1475_t1n31_05c3_Hemi12_1\": \"a woman with short blonde hair, slim athletic build, a red leather jacket, dark blue jeans, and white sneakers is moving in the urban rooftop garden\",\n \"10_D_loc5_730_t1n34_02da_Hemi12_1\": \"a man with medium-length wavy brown hair, lean build, a black bomber jacket, olive green cargo pants, and brown hiking boots is moving in the swamp\",\n \"11_D_loc3_1028_t1n38_0404_Hemi12_1\": \"a man with buzz-cut blonde hair, stocky build, a gray zip-up sweater, black shorts, and red basketball shoes is moving in the riverbank\",\n \"12_D_loc1_01_t2n1_0001_Hemi12_1\": \"a man with short spiky brown hair, athletic build, a navy blue jacket, beige cargo pants, and black sneakers and a dog with a fluffy coat, wagging tail, and warm golden-brown fur, exuding a gentle and friendly charm are moving in the fjord\",\n \"13_D_loc1_136_t2n2_0088_Hemi12_1\": \"a woman with long wavy blonde hair, petite figure, a red floral dress, white sandals, and a yellow shoulder bag and a tiger with vibrant orange and black stripes, piercing yellow eyes, and a powerful stance, exuding strength and grace are moving in the sunset beach\",\n \"14_D_loc2_137_t2n3_0089_Hemi12_1\": \"a man with a shaved head, broad shoulders, a gray graphic t-shirt, dark jeans, and brown leather boots and a giraffe with golden-yellow fur, long legs, a tall slender neck, and patches of brown spots, exuding elegance and calm are moving in the cave\",\n \"15_D_loc1_141_t2n6_008d_Hemi12_1\": \"a woman with shoulder-length straight auburn hair, a slender figure, a green button-up blouse, black leggings, and white sneakers and an alpaca with soft white wool, short legs, a thick neck, and a fluffy head of fur, radiating gentle charm are moving in the snowy tundra\",\n \"16_D_loc2_142_t2n7_008e_Hemi12_1\": \"a man with messy black hair, tall frame, a plaid red and black shirt, faded blue jeans, and tan hiking boots and a zebra with black and white stripes, sturdy legs, a short neck, and a sleek mane running down its back are moving in the prairie\",\n \"17_D_loc1_506_t2n9_01fa_Hemi12_1\": \"a man with medium-length straight brown hair, tall and slender, a gray crew-neck t-shirt, beige trousers, and dark green sneakers and an alpaca with soft white wool, short legs, a thick neck, and a fluffy head of fur, radiating gentle charm are moving in the asian town\",\n \"18_D_loc1_11_t2n10_000b_Hemi12_1\": \"a woman with short curly black hair, slender build, a pink hoodie, light gray joggers, and blue sneakers and a gazelle with light golden fur, long slender legs, a thin neck, and short, sharp horns, embodying elegance and agility are moving in the rainforest\",\n \"19_D_loc1_81_t2n11_0051_Hemi12_1\": \"a man with short black wavy hair, lean figure, a green and yellow plaid shirt, dark brown pants, and black suede shoes and a horse with chestnut brown fur, muscular legs, a slim neck, and a flowing mane, exuding strength and grace are moving in the canyon\",\n \"20_D_loc1_16_t2n13_0010_Hemi12_1\": \"a man with curly black hair, muscular build, a dark green hoodie, gray joggers, and white running shoes and a sleek black panther with a smooth, glossy coat, emerald green eyes, and a powerful stance are moving in the savanna\",\n \"21_D_loc1_86_t2n14_0056_Hemi12_1\": \"a woman with short blonde hair, slim athletic build, a red leather jacket, dark blue jeans, and white sneakers and a cheetah with golden fur covered in black spots, intense amber eyes, and a slender, agile body are moving in the urban rooftop garden\",\n \"22_D_loc1_1446_t2n15_05a6_Hemi12_1\": \"a man with medium-length wavy brown hair, lean build, a black bomber jacket, olive green cargo pants, and brown hiking boots and a regal lion with a thick, flowing golden mane, sharp brown eyes, and a powerful muscular frame are moving in the swamp\",\n \"23_D_loc2_92_t2n17_005c_Hemi12_1\": \"a man with buzz-cut blonde hair, stocky build, a gray zip-up sweater, black shorts, and red basketball shoes and a snow leopard with pale gray fur adorned with dark rosettes, icy blue eyes, and a stealthy, poised posture are moving in the riverbank\",\n \"24_D_loc1_156_t2n19_009c_Hemi12_1\": \"a woman with long straight black hair, toned build, a blue denim jacket, light gray leggings, and black slip-on shoes and a jaguar with a golden-yellow coat dotted with intricate black rosettes, deep green eyes, and a muscular build are moving in the coral reef\",\n \"25_D_loc1_101_t2n21_0065_Hemi12_1\": \"a man with short curly red hair, average build, a black leather jacket, dark blue cargo pants, and white sneakers and a wolf with thick silver-gray fur, alert golden eyes, and a lean yet strong body, exuding confidence and boldness are moving in the volcanic landscape\",\n \"26_D_loc1_161_t2n22_00a1_Hemi12_1\": \"a woman with shoulder-length wavy brown hair, slim build, a green parka, black leggings, and gray hiking boots and a tiger with a pristine white coat marked by bold black stripes, bright blue eyes, and a graceful, poised form are moving in the wind farm\",\n \"27_D_loc1_36_t2n25_0024_Hemi12_1\": \"a man with short straight black hair, tall and lean build, a navy blue sweater, khaki shorts, and brown sandals and a lynx with tufted ears, soft reddish-brown fur with faint spots, and intense yellow-green eyes are moving in the town street\",\n \"28_D_loc1_166_t2n26_00a6_Hemi12_1\": \"a woman with pixie-cut blonde hair, athletic build, a red windbreaker, blue ripped jeans, and black combat boots and a bear with dark brown fur, small but fierce black eyes, and a broad and muscular build, radiating power are moving in the night city square\",\n \"29_D_loc1_221_t2n29_00dd_Hemi12_1\": \"a man with medium-length wavy gray hair, muscular build, a maroon t-shirt, beige chinos, and brown loafers and a swift fox with reddish-orange fur, a bushy tail tipped with white, and sharp, intelligent amber eyes are moving in the mall lobby\",\n \"30_D_loc1_171_t2n30_00ab_Hemi12_1\": \"a woman with long curly black hair, average build, a purple hoodie, black athletic shorts, and white running shoes and a bear with dark brown fur, small but fierce black eyes, and a broad and muscular build, radiating power are moving in the glacier\",\n \"31_D_loc1_121_t2n33_0079_Hemi12_1\": \"a man with short spiky blonde hair, slim build, a black trench coat, blue jeans, and brown hiking shoes and a fox with sleek russet fur, a bushy tail tipped with black, and bright green and cunning eyes are moving in the seaside street\",\n \"32_D_loc1_176_t2n34_00b0_Hemi12_1\": \"a man with short spiky brown hair, athletic build, a navy blue jacket, beige cargo pants, and black sneakers and a kangaroo with brown fur, powerful hind legs, and a muscular tail, showcasing its strength and agility are moving in the gymnastics room\",\n \"33_D_loc1_126_t2n36_007e_Hemi12_1\": \"a woman with long wavy blonde hair, petite figure, a red floral dress, white sandals, and a yellow shoulder bag and a polar bear with thick white fur, strong paws, and a black nose, embodying the essence of the Arctic are moving in the abandoned factory\",\n \"34_D_loc1_56_t2n38_0038_Hemi12_1\": \"a man with a shaved head, broad shoulders, a gray graphic t-shirt, dark jeans, and brown leather boots and a cheetah with a slender build, spotted golden fur, and sharp eyes, epitomizing speed and agility are moving in the autumn forest\",\n \"35_D_loc1_131_t2n41_0083_Hemi12_1\": \"a woman with shoulder-length straight auburn hair, a slender figure, a green button-up blouse, black leggings, and white sneakers and a dolphin with sleek grey skin, a curved dorsal fin, and intelligent, playful eyes, reflecting its nature are moving in the mountain village\",\n \"36_D_loc1_01_t2n1_0001_Hemi12_1\": \"a man with messy black hair, tall frame, a plaid red and black shirt, faded blue jeans, and tan hiking boots and a wolf with a body covered in thick silver fur, sharp ears, and piercing yellow eyes, showcasing its alertness are moving in the coastal harbor\",\n \"37_D_loc1_136_t2n2_0088_Hemi12_1\": \"a man with medium-length straight brown hair, tall and slender, a gray crew-neck t-shirt, beige trousers, and dark green sneakers and a leopard with a body covered in golden fur, dark rosettes, and a long muscular tail, emphasizing its strength are moving in the ancient ruins\",\n \"38_D_loc2_137_t2n3_0089_Hemi12_1\": \"a woman with short curly black hair, slender build, a pink hoodie, light gray joggers, and blue sneakers and a penguin with a body covered in smooth black-and-white feathers, short wings, and webbed feet are moving in the modern metropolis\",\n \"39_D_loc1_141_t2n6_008d_Hemi12_1\": \"a man with short black wavy hair, lean figure, a green and yellow plaid shirt, dark brown pants, and black suede shoes and a gazelle with a body covered in sleek tan fur, long legs, and elegant curved horns, showcasing its grace are moving in the desert\",\n \"40_D_loc2_142_t2n7_008e_Hemi12_1\": \"a man with curly black hair, muscular build, a dark green hoodie, gray joggers, and white running shoes and a rabbit with a body covered in soft fur, quick hops, and a playful demeanor, showcasing its energy are moving in the forest\",\n \"41_D_loc1_506_t2n9_01fa_Hemi12_1\": \"a woman with short curly black hair, slender build, a pink hoodie, light gray joggers, and blue sneakers and a koala with a body covered in soft grey fur, large round ears, and a black nose, radiating cuteness are moving in the city\",\n \"42_D_loc1_11_t2n10_000b_Hemi12_1\": \"a man with medium-length wavy brown hair, lean build, a black bomber jacket, olive green cargo pants, and brown hiking boots and a rhinoceros with a body covered in thick grey skin, a massive horn on its snout, and sturdy legs are moving in the snowy street\",\n \"43_D_loc1_81_t2n11_0051_Hemi12_1\": \"a man with buzz-cut blonde hair, stocky build, a gray zip-up sweater, black shorts, and red basketball shoes and a flamingo with a body covered in pink feathers, long slender legs, and a gracefully curved neck are moving in the park\",\n \"44_D_loc1_16_t2n13_0010_Hemi12_1\": \"a woman with long straight black hair, toned build, a blue denim jacket, light gray leggings, and black slip-on shoes and a parrot with bright red, blue, and yellow feathers, a curved beak, and sharp intelligent eyes are moving in the fjord\",\n \"45_D_loc1_86_t2n14_0056_Hemi12_1\": \"a man with short curly red hair, average build, a black leather jacket, dark blue cargo pants, and white sneakers and a hippopotamus with a body covered in thick grey-brown skin, massive jaws, and a large body are moving in the sunset beach\",\n \"46_D_loc1_1446_t2n15_05a6_Hemi12_1\": \"a woman with shoulder-length wavy brown hair, slim build, a green parka, black leggings, and gray hiking boots and a crocodile with a body covered in scaly green skin, a powerful tail, and sharp teeth are moving in the cave\",\n \"47_D_loc2_92_t2n17_005c_Hemi12_1\": \"a man with short straight black hair, tall and lean build, a navy blue sweater, khaki shorts, and brown sandals and a moose with a body covered in thick brown fur, massive antlers, and a bulky frame are moving in the snowy tundra\",\n \"48_D_loc1_156_t2n19_009c_Hemi12_1\": \"a woman with pixie-cut blonde hair, athletic build, a red windbreaker, blue ripped jeans, and black combat boots and a fluttering butterfly with intricate wing patterns, vivid colors, and graceful flight are moving in the prairie\",\n \"49_D_loc1_101_t2n21_0065_Hemi12_1\": \"a man with medium-length wavy gray hair, muscular build, a maroon t-shirt, beige chinos, and brown loafers and a chameleon with a body covered in vibrant green scales, bulging eyes, and a curled tail, showcasing its unique charm are moving in the asian town\",\n \"50_D_loc1_161_t2n22_00a1_Hemi12_1\": \"a woman with long curly black hair, average build, a purple hoodie, black athletic shorts, and white running shoes and a lemur with a body covered in soft grey fur, a ringed tail, and wide yellow eyes, and curious expression are moving in the rainforest\",\n \"51_D_loc1_36_t2n25_0024_Hemi12_1\": \"a man with short spiky blonde hair, slim build, a black trench coat, blue jeans, and brown hiking shoes and a squirrel with a body covered in bushy red fur, large eyes, and a fluffy tail are moving in the canyon\",\n \"52_D_loc1_166_t2n26_00a6_Hemi12_1\": \"a man with short spiky brown hair, athletic build, a navy blue jacket, beige cargo pants, and black sneakers and a panda with a body covered in fluffy black-and-white fur, a round face, and gentle eyes, radiating warmth are moving in the savanna\",\n \"53_D_loc1_221_t2n29_00dd_Hemi12_1\": \"a woman with long wavy blonde hair, petite figure, a red floral dress, white sandals, and a yellow shoulder bag and a porcupine with a body covered in spiky brown quills, a small nose, and curious eyes are moving in the urban rooftop garden\",\n \"54_D_loc1_171_t2n30_00ab_Hemi12_1\": \"a man with a shaved head, broad shoulders, a gray graphic t-shirt, dark jeans, and brown leather boots and a sedan with a sleek metallic silver body, long wheelbase, a low-profile hood, and a small rear spoiler are moving in the swamp\",\n \"55_D_loc1_121_t2n33_0079_Hemi12_1\": \"a woman with shoulder-length straight auburn hair, a slender figure, a green button-up blouse, black leggings, and white sneakers and an SUV with a matte black exterior, elevated suspension, a tall roofline, and a compact rear roof rack are moving in the riverbank\",\n \"56_D_loc1_176_t2n34_00b0_Hemi12_1\": \"a man with messy black hair, tall frame, a plaid red and black shirt, faded blue jeans, and tan hiking boots and a pickup truck with rugged dark green paint, extended cab, raised suspension, and a modest cargo bed cover are moving in the coral reef\",\n \"57_D_loc1_126_t2n36_007e_Hemi12_1\": \"a man with medium-length straight brown hair, tall and slender, a gray crew-neck t-shirt, beige trousers, and dark green sneakers and a vintage convertible with a body covered in shiny red paint, chrome bumpers, and a stylish design are moving in the volcanic landscape\",\n \"58_D_loc1_56_t2n38_0038_Hemi12_1\": \"a woman with short curly black hair, slender build, a pink hoodie, light gray joggers, and blue sneakers and a futuristic electric car with a minimalist silver design, slim LED lights, and smooth curves are moving in the wind farm\",\n \"59_D_loc1_131_t2n41_0083_Hemi12_1\": \"a man with short black wavy hair, lean figure, a green and yellow plaid shirt, dark brown pants, and black suede shoes and a compact electric vehicle with a silver finish, aerodynamic profile, and efficient battery are moving in the town street\",\n \"60_D_loc1_01_t2n1_0001_Hemi12_1\": \"a man with curly black hair, muscular build, a dark green hoodie, gray joggers, and white running shoes and a firefighting robot with a water cannon arm, heat sensors, and durable red-and-silver exterior are moving in the night city square\",\n \"61_D_loc1_136_t2n2_0088_Hemi12_1\": \"a woman with short blonde hair, slim athletic build, a red leather jacket, dark blue jeans, and white sneakers and an industrial welding robot with articulated arms, a laser precision welder, and heat-resistant shields are moving in the mall lobby\",\n \"62_D_loc2_137_t2n3_0089_Hemi12_1\": \"a woman with short blonde hair, slim athletic build, a red leather jacket, dark blue jeans, and white sneakers and a dog with a fluffy coat, wagging tail, and warm golden-brown fur, exuding a gentle and friendly charm are moving in the glacier\",\n \"63_D_loc1_141_t2n6_008d_Hemi12_1\": \"a man with buzz-cut blonde hair, stocky build, a gray zip-up sweater, black shorts, and red basketball shoes and a disaster rescue robot with reinforced limbs, advanced AI, and a rugged body designed to navigate are moving in the seaside street\",\n \"64_D_loc2_142_t2n7_008e_Hemi12_1\": \"a woman with long straight black hair, toned build, a blue denim jacket, light gray leggings, and black slip-on shoes and an exploration rover robot with solar panels, durable wheels, and advanced sensors for planetary exploration are moving in the gymnastics room\",\n \"65_D_loc1_506_t2n9_01fa_Hemi12_1\": \"a man with short curly red hair, average build, a black leather jacket, dark blue cargo pants, and white sneakers and a dog with a fluffy coat, wagging tail, and warm golden-brown fur, exuding a gentle and friendly charm are moving in the abandoned factory\",\n \"66_D_loc1_11_t2n10_000b_Hemi12_1\": \"a woman with shoulder-length wavy brown hair, slim build, a green parka, black leggings, and gray hiking boots and a tiger with vibrant orange and black stripes, piercing yellow eyes, and a powerful stance, exuding strength and grace are moving in the autumn forest\",\n \"67_D_loc1_81_t2n11_0051_Hemi12_1\": \"a man with short straight black hair, tall and lean build, a navy blue sweater, khaki shorts, and brown sandals and a giraffe with golden-yellow fur, long legs, a tall slender neck, and patches of brown spots, exuding elegance and calm are moving in the mountain village\",\n \"68_D_loc1_16_t2n13_0010_Hemi12_1\": \"a woman with pixie-cut blonde hair, athletic build, a red windbreaker, blue ripped jeans, and black combat boots and an alpaca with soft white wool, short legs, a thick neck, and a fluffy head of fur, radiating gentle charm are moving in the coastal harbor\",\n \"69_D_loc1_86_t2n14_0056_Hemi12_1\": \"a man with medium-length wavy gray hair, muscular build, a maroon t-shirt, beige chinos, and brown loafers and a zebra with black and white stripes, sturdy legs, a short neck, and a sleek mane running down its back are moving in the ancient ruins\",\n \"70_D_loc1_1446_t2n15_05a6_Hemi12_1\": \"a woman with long curly black hair, average build, a purple hoodie, black athletic shorts, and white running shoes and a deer with sleek tan fur, long slender legs, a graceful neck, and tiny antlers atop its head are moving in the modern metropolis\",\n \"71_D_loc2_92_t2n17_005c_Hemi12_1\": \"a man with short spiky blonde hair, slim build, a black trench coat, blue jeans, and brown hiking shoes and a gazelle with light golden fur, long slender legs, a thin neck, and short, sharp horns, embodying elegance and agility are moving in the desert\",\n \"72_D_loc1_156_t2n19_009c_Hemi12_1\": \"a man with short spiky brown hair, athletic build, a navy blue jacket, beige cargo pants, and black sneakers and a horse with chestnut brown fur, muscular legs, a slim neck, and a flowing mane, exuding strength and grace are moving in the forest\",\n \"73_D_loc1_101_t2n21_0065_Hemi12_1\": \"a man with short spiky blonde hair, slim build, a black trench coat, blue jeans, and brown hiking shoes and a sleek black panther with a smooth, glossy coat, emerald green eyes, and a powerful stance are moving in the city\",\n \"74_D_loc1_161_t2n22_00a1_Hemi12_1\": \"a man with a shaved head, broad shoulders, a gray graphic t-shirt, dark jeans, and brown leather boots and a cheetah with golden fur covered in black spots, intense amber eyes, and a slender, agile body are moving in the snowy street\",\n \"75_D_loc1_36_t2n25_0024_Hemi12_1\": \"a woman with shoulder-length straight auburn hair, a slender figure, a green button-up blouse, black leggings, and white sneakers and a regal lion with a thick, flowing golden mane, sharp brown eyes, and a powerful muscular frame are moving in the park\",\n \"76_D_loc1_166_t2n26_00a6_Hemi12_1\": \"a man with messy black hair, tall frame, a plaid red and black shirt, faded blue jeans, and tan hiking boots and a snow leopard with pale gray fur adorned with dark rosettes, icy blue eyes, and a stealthy, poised posture are moving in the fjord\",\n \"77_D_loc1_221_t2n29_00dd_Hemi12_1\": \"a man with medium-length straight brown hair, tall and slender, a gray crew-neck t-shirt, beige trousers, and dark green sneakers and a jaguar with a golden-yellow coat dotted with intricate black rosettes, deep green eyes, and a muscular build are moving in the sunset beach\",\n \"78_D_loc1_171_t2n30_00ab_Hemi12_1\": \"a woman with short curly black hair, slender build, a pink hoodie, light gray joggers, and blue sneakers and a wolf with thick silver-gray fur, alert golden eyes, and a lean yet strong body, exuding confidence and boldness are moving in the cave\",\n \"79_D_loc1_121_t2n33_0079_Hemi12_1\": \"a man with short black wavy hair, lean figure, a green and yellow plaid shirt, dark brown pants, and black suede shoes and a tiger with a pristine white coat marked by bold black stripes, bright blue eyes, and a graceful, poised form are moving in the snowy tundra\",\n \"80_D_loc1_176_t2n34_00b0_Hemi12_1\": \"a man with curly black hair, muscular build, a dark green hoodie, gray joggers, and white running shoes and a lynx with tufted ears, soft reddish-brown fur with faint spots, and intense yellow-green eyes are moving in the prairie\",\n \"81_D_loc1_126_t2n36_007e_Hemi12_1\": \"a woman with short blonde hair, slim athletic build, a red leather jacket, dark blue jeans, and white sneakers and a bear with dark brown fur, small but fierce black eyes, and a broad and muscular build, radiating power are moving in the asian town\",\n \"82_D_loc1_56_t2n38_0038_Hemi12_1\": \"a man with medium-length wavy brown hair, lean build, a black bomber jacket, olive green cargo pants, and brown hiking boots and a swift fox with reddish-orange fur, a bushy tail tipped with white, and sharp, intelligent amber eyes are moving in the rainforest\",\n \"83_D_loc1_131_t2n41_0083_Hemi12_1\": \"a man with buzz-cut blonde hair, stocky build, a gray zip-up sweater, black shorts, and red basketball shoes and a falcon with blue-gray feathers, sharp talons, and keen yellow eyes fixed on its prey below are moving in the canyon\",\n \"84_D_loc1_301_t3n3_012d_Hemi12_1\": \"a woman with long wavy blonde hair, petite figure, a red floral dress, white sandals, and a yellow shoulder bag and a squirrel with a body covered in bushy red fur, large eyes, and a fluffy tail and a penguin with a body covered in smooth black-and-white feathers, short wings, and webbed feet are moving in the sunset beach\",\n \"85_D_loc3_133_t3n1_0085_Hemi12_1\": \"a woman with short curly black hair, slender build, a pink hoodie, light gray joggers, and blue sneakers and a leopard with a body covered in golden fur, dark rosettes, and a long muscular tail, emphasizing its strength and a rhinoceros with a body covered in thick grey skin, a massive horn on its snout, and sturdy legs are moving in the rainforest\",\n \"86_D_loc1_46_t3n8_002e_Hemi12_1\": \"a woman with short blonde hair, slim athletic build, a red leather jacket, dark blue jeans, and white sneakers and a parrot with bright red, blue, and yellow feathers, a curved beak, and sharp intelligent eyes and a gazelle with light golden fur, long slender legs, a thin neck, and short, sharp horns, embodying elegance and agility are moving in the urban rooftop garden\",\n \"87_D_loc2_192_t3n9_00c0_Hemi12_1\": \"a woman with pixie-cut blonde hair, athletic build, a red windbreaker, blue ripped jeans, and black combat boots and a fluttering butterfly with intricate wing patterns, vivid colors, and graceful flight and a sleek black panther with a smooth, glossy coat, emerald green eyes, and a powerful stance are moving in the night city square\",\n \"88_D_loc3_93_t3n21_005d_Hemi12_1\": \"a man with medium-length wavy gray hair, muscular build, a maroon t-shirt, beige chinos, and brown loafers and a cheetah with golden fur covered in black spots, intense amber eyes, and a slender, agile body and a rhinoceros with a body covered in thick grey skin, a massive horn on its snout, and sturdy legs are moving in the mall lobby\",\n \"89_D_loc1_301_t3n3_012d_Hemi12_1\": \"a man with short spiky blonde hair, slim build, a black trench coat, blue jeans, and brown hiking shoes and a fluttering butterfly with intricate wing patterns, vivid colors, and graceful flight and a dolphin with sleek grey skin, a curved dorsal fin, and intelligent, playful eyes, reflecting its nature are moving in the seaside street\",\n \"90_D_loc1_46_t3n8_002e_Hemi12_1\": \"a woman with long wavy blonde hair, petite figure, a red floral dress, white sandals, and a yellow shoulder bag and a crocodile with a body covered in scaly green skin, a powerful tail, and sharp teeth and a dolphin with sleek grey skin, a curved dorsal fin, and intelligent, playful eyes, reflecting its nature are moving in the abandoned factory\",\n \"91_D_loc3_93_t3n21_005d_Hemi12_1\": \"a man with short straight black hair, tall and lean build, a navy blue sweater, khaki shorts, and brown sandals and a koala with a body covered in soft grey fur, large round ears, and a black nose, radiating cuteness and a parrot with bright red, blue, and yellow feathers, a curved beak, and sharp intelligent eyes are moving in the snowy tundra\",\n \"92_D_loc2_192_t3n9_00c0_Hemi12_1\": \"a woman with short curly black hair, slender build, a pink hoodie, light gray joggers, and blue sneakers and a koala with a body covered in soft grey fur, large round ears, and a black nose, radiating cuteness and a cheetah with golden fur covered in black spots, intense amber eyes, and a slender, agile body are moving in the wind farm\",\n \"93_D_loc3_133_t3n1_0085_Hemi12_1\": \"a man with curly black hair, muscular build, a dark green hoodie, gray joggers, and white running shoes and a sedan with a sleek metallic silver body, long wheelbase, a low-profile hood, and a small rear spoiler and a koala with a body covered in soft grey fur, large round ears, and a black nose, radiating cuteness are moving in the night city square\",\n \"94_D_loc1_301_t3n3_012d_Hemi12_1\": \"a woman with short blonde hair, slim athletic build, a red leather jacket, dark blue jeans, and white sneakers and a sleek black panther with a smooth, glossy coat, emerald green eyes, and a powerful stance and a parrot with bright red, blue, and yellow feathers, a curved beak, and sharp intelligent eyes are moving in the mall lobby\",\n \"95_D_loc3_93_t3n21_005d_Hemi12_1\": \"a man with short curly red hair, average build, a black leather jacket, dark blue cargo pants, and white sneakers and a koala with a body covered in soft grey fur, large round ears, and a black nose, radiating cuteness and a fox with sleek russet fur, a bushy tail tipped with black, and bright green and cunning eyes are moving in the abandoned factory\",\n \"96_D_loc1_46_t3n8_002e_Hemi12_1\": \"a woman with shoulder-length straight auburn hair, a slender figure, a green button-up blouse, black leggings, and white sneakers and a rhinoceros with a body covered in thick grey skin, a massive horn on its snout, and sturdy legs and a tiger with vibrant orange and black stripes, piercing yellow eyes, and a powerful stance, exuding strength and grace are moving in the park\",\n \"97_D_loc1_871_t3n7_0367_Hemi12_1\": \"a woman with shoulder-length wavy brown hair, slim build, a green parka, black leggings, and gray hiking boots and a polar bear with thick white fur, strong paws, and a black nose, embodying the essence of the Arctic and a vintage convertible with a body covered in shiny red paint, chrome bumpers, and a stylish design are moving in the swamp\",\n \"98_D_loc3_133_t3n1_0085_Hemi12_1\": \"a woman with long curly black hair, average build, a purple hoodie, black athletic shorts, and white running shoes and a fluttering butterfly with intricate wing patterns, vivid colors, and graceful flight and a sleek black panther with a smooth, glossy coat, emerald green eyes, and a powerful stance are moving in the wind farm\",\n \"99_D_loc3_133_t3n1_0085_Hemi12_1\": \"a woman with pixie-cut blonde hair, athletic build, a red windbreaker, blue ripped jeans, and black combat boots and a vintage convertible with a body covered in shiny red paint, chrome bumpers, and a stylish design and a falcon with blue-gray feathers, sharp talons, and keen yellow eyes fixed on its prey below are moving in the fjord\"\n}\n" }, { "path": "eval/common_metrics_on_video_quality/eval_visual.sh", "content": "basedir=eval_folder\nfolder1_path=${basedir}/base_t2v_eval_sets\nfolder2_path=${basedir}/eval_sets\n\n# calculate FVD\npython calculate_fvd_styleganv.py -v1_f ${folder1_path} -v2_f ${folder2_path}\n\n# calculate FID\npython -m pytorch_fid ${basedir}/eval_1 ${basedir}/eval_2\n\n# calculate CLIP-SIM\npython calculate_clip.py -v_f ${folder2_path}\n\nrm -rf ${basedir}/eval_1\nrm -rf ${basedir}/eval_2" } ]