[
{
"path": ".gitattributes",
"content": "*.jpg filter=lfs diff=lfs merge=lfs -text\n*.jpeg filter=lfs diff=lfs merge=lfs -text\n*.png filter=lfs diff=lfs merge=lfs -text\n*.gif filter=lfs diff=lfs merge=lfs -text\ntests/utils/car.png filter=lfs diff=lfs merge=lfs -text\n"
},
{
"path": ".github/workflows/pylint.yml",
"content": "name: Ruff\n\non: [push]\n\njobs:\n build:\n runs-on: ubuntu-latest\n strategy:\n matrix:\n python-version: [\"3.10\"]\n steps:\n - name: Checkout repository and submodules\n uses: actions/checkout@v4\n - name: Set up Python ${{ matrix.python-version }}\n uses: actions/setup-python@v5\n with:\n python-version: ${{ matrix.python-version }}\n - name: Install dependencies\n run: |\n python -m pip install --upgrade pip\n pip install ruff==0.2.2 black==24.2.0\n - name: Analyzing the code with ruff\n run: |\n ruff $(git ls-files '*.py')\n - name: Verify that no Black changes are required\n run: |\n black --check $(git ls-files '*.py')\n"
},
{
"path": ".gitignore",
"content": "# 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/\nlib/\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/latest/usage/project/#working-with-version-control\n.pdm.toml\n.pdm-python\n.pdm-build/\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# 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# From inference.py\noutputs/\n*.mp4\n*.png\n!tests/utils/car.png"
},
{
"path": ".pre-commit-config.yaml",
"content": "repos:\n - repo: https://github.com/astral-sh/ruff-pre-commit\n # Ruff version.\n rev: v0.2.2\n hooks:\n # Run the linter.\n - id: ruff\n args: [--fix] # Automatically fix issues if possible.\n types: [python] # Ensure it only runs on .py files.\n\n - repo: https://github.com/psf/black\n rev: 24.2.0 # Specify the version of Black you want\n hooks:\n - id: black\n name: Black code formatter\n language_version: python3 # Use the Python version you're targeting (e.g., 3.10)"
},
{
"path": "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 [yyyy] [name of copyright owner]\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": "README.md",
"content": "\n\n# LTX-Video\n\n[](https://ltx.video)\n[](https://huggingface.co/Lightricks/LTX-Video)\n[](https://app.ltx.studio/ltx-2-playground/t2v)\n[](https://arxiv.org/abs/2501.00103)\n[](https://github.com/Lightricks/LTX-Video-Trainer)\n[](https://discord.gg/ltxplatform)\n\nThis is the official repository for LTX-Video.\n\n
\n\n---\n\n## ๐ **New: LTX-2 is Now Available!**\n\n**We're excited to announce [LTX-2](https://github.com/Lightricks/LTX-2) - the next generation of LTX with synchronized audio+video generation!**\n\nLTX-2 is the first DiT-based audio-video foundation model that contains all core capabilities of modern video generation in one model. **LTX-2 is now the primary home for LTX development** and includes significant improvements:\n\n- ๐ต **Synchronized Audio+Video Generation** - Generate videos with perfectly synchronized audio\n- ๐ฌ **Latest Model** - LTX-2 with improved quality and capabilities\n- ๐ **ComfyUI Integration** - Built into ComfyUI core for seamless workflows\n- ๐ฏ **Advanced Features:**\n - Multiple keyframe support\n - IC-LoRA control models for precise generation\n - Standard LoRA support for style customization\n - Latent upsampler for multiscale pipelines\n- ๐ ๏ธ **Training Tools** - LoRA training capabilities\n- ๐ **Comprehensive Documentation** - Full documentation at [https://docs.ltx.video](https://docs.ltx.video)\n- ๐ **Active Development** - Ongoing improvements and community support\n\n**[๐ Check out LTX-2 here](https://github.com/Lightricks/LTX-2)**\n\n**[๐ View Documentation](https://docs.ltx.video)**\n\n---\n\n## Table of Contents\n\n- [Introduction](#introduction)\n- [What's New](#news)\n- [Models](#models)\n- [Quick Start Guide](#quick-start-guide)\n - [Online demo](#online-inference)\n - [Run locally](#run-locally)\n - [Installation](#installation)\n - [Inference](#inference)\n - [ComfyUI Integration](#comfyui-integration)\n - [Diffusers Integration](#diffusers-integration)\n- [Model User Guide](#model-user-guide)\n- [Community Contribution](#community-contribution)\n- [Training](#training)\n- [Control Models](#control-models)\n- [Join Us!](#join-us)\n- [Acknowledgement](#acknowledgement)\n\n# Introduction\n\nLTX-Video is the first DiT-based video generation model that contains all core capabilities of modern video generation in one model: synchronized audio and video, high fidelity, multiple performance modes, production-ready outputs, API access, and open access. It can generate up to 50 FPS videos at native 4K resolution with synchronized audio in one pass.\nThe model is trained on a large-scale dataset of diverse videos and can generate high-resolution videos with realistic and diverse content.\n\nThe model supports image-to-video, multi-keyframe conditioning, keyframe-based animation, video extension (both forward and backward), video-to-video transformations, and any combination of these features.\n\n### Image-to-video examples\n| | | |\n|:---:|:---:|:---:|\n|  |  |  |\n|  |  |  |\n|  |  |  |\n\n### Controlled video examples\n| | | |\n|:---:|:---:|:---:|\n|  |  |  |\n\n| | |\n|:---:|:---:|\n|  |  |\n\n# News\n\n## October 23, 2025: LTX-2 Announced\n\nToday we announced our newest foundation model, LTX-2. LTX-2 represents a major leap forward from our previous model, LTXV 0.9.8. Hereโs whatโs new:\n* **Audio + Video, Together**: Visuals and sound are generated in one coherent process, with motion, dialogue, ambience, and music flowing simultaneously.\n* **4K Fidelity**: Professional-grade precision with native 4K and up to 50 fps, sharp textures, clean motion, and synchronized audio.\n* **Longer Generations**: LTX-2 supports longer, continuous clips with synchronized audio up to 10 seconds.\n* **Low Cost & Efficiency**: Up to 50% lower compute cost than competing models, powered by a multi-GPU inference stack.\n* **Creative Control**: Multi-keyframe conditioning, 3D camera logic, and LoRA fine-tuning deliver frame-level precision and style consistency.\n\nFor more details, please see our [blog post](https://website.ltx.video/blog/introducing-ltx-2). LTX-2 model weights, code, and benchmarks will be released to the community later in 2025. \n\n\n## July, 16th, 2025: New Distilled models v0.9.8 with up to 60 seconds of video:\n- Long shot generation in LTXV-13B!\n * LTX-Video now supports up to 60 seconds of video.\n * Compatible also with the official IC-LoRAs.\n * Try now in [ComfyUI](https://github.com/Lightricks/ComfyUI-LTXVideo/tree/master/example_workflows/ltxv-13b-i2v-long-multi-prompt.json).\n- Release a new distilled models:\n * 13B distilled model [ltxv-13b-0.9.8-distilled](https://github.com/Lightricks/LTX-Video/blob/main/configs/ltxv-13b-0.9.8-distilled.yaml)\n * 2B distilled model [ltxv-2b-0.9.8-distilled](https://github.com/Lightricks/LTX-Video/blob/main/configs/ltxv-2b-0.9.8-distilled.yaml)\n * Both models are distilled from the same base model [ltxv-13b-0.9.8-dev](https://github.com/Lightricks/LTX-Video/blob/main/configs/ltxv-13b-0.9.8-dev.yaml) and are compatible for use together in the same multiscale pipeline.\n * Improved prompt understanding and detail generation\n * Includes corresponding FP8 weights and workflows.\n- Release a new detailer model [LTX-Video-ICLoRA-detailer-13B-0.9.8](https://huggingface.co/Lightricks/LTX-Video-ICLoRA-detailer-13b-0.9.8)\n * Available in [ComfyUI](https://github.com/Lightricks/ComfyUI-LTXVideo/tree/master/example_workflows/ltxv-13b-upscale.json).\n\n## July, 8th, 2025: New Control Models Released!\n- Released three new control models for LTX-Video on HuggingFace:\n * **Depth Control**: [LTX-Video-ICLoRA-depth-13b-0.9.7](https://huggingface.co/Lightricks/LTX-Video-ICLoRA-depth-13b-0.9.7)\n * **Pose Control**: [LTX-Video-ICLoRA-pose-13b-0.9.7](https://huggingface.co/Lightricks/LTX-Video-ICLoRA-pose-13b-0.9.7)\n * **Canny Control**: [LTX-Video-ICLoRA-canny-13b-0.9.7](https://huggingface.co/Lightricks/LTX-Video-ICLoRA-canny-13b-0.9.7)\n\n\n## May, 14th, 2025: New distilled model 13B v0.9.7:\n- Release a new 13B distilled model [ltxv-13b-0.9.7-distilled](https://huggingface.co/Lightricks/LTX-Video/blob/main/ltxv-13b-0.9.7-distilled.safetensors)\n * Amazing for iterative work - generates HD videos in 10 seconds, with low-res preview after just 3 seconds (on H100)!\n * Does not require classifier-free guidance and spatio-temporal guidance.\n * Supports sampling with 8 (recommended), or less diffusion steps.\n * Also released a LoRA version of the distilled model, [ltxv-13b-0.9.7-distilled-lora128](https://huggingface.co/Lightricks/LTX-Video/blob/main/ltxv-13b-0.9.7-distilled-lora128.safetensors)\n * Requires only 1GB of VRAM\n * Can be used with the full 13B model for fast inference\n- Release a new quantized distilled model [ltxv-13b-0.9.7-distilled-fp8](https://huggingface.co/Lightricks/LTX-Video/blob/main/ltxv-13b-0.9.7-distilled-fp8.safetensors) for *real-time* generation (on H100) with even less VRAM\n\n## May, 5th, 2025: New model 13B v0.9.7:\n- Release a new 13B model [ltxv-13b-0.9.7-dev](https://huggingface.co/Lightricks/LTX-Video/blob/main/ltxv-13b-0.9.7-dev.safetensors)\n- Release a new quantized model [ltxv-13b-0.9.7-dev-fp8](https://huggingface.co/Lightricks/LTX-Video/blob/main/ltxv-13b-0.9.7-dev-fp8.safetensors) for faster inference with less VRam\n- Release a new upscalers\n * [ltxv-temporal-upscaler-0.9.7](https://huggingface.co/Lightricks/LTX-Video/blob/main/ltxv-temporal-upscaler-0.9.7.safetensors)\n * [ltxv-spatial-upscaler-0.9.7](https://huggingface.co/Lightricks/LTX-Video/blob/main/ltxv-spatial-upscaler-0.9.7.safetensors)\n- Breakthrough prompt adherence and physical understanding.\n- New Pipeline for multi-scale video rendering for fast and high quality results\n\n\n## April, 15th, 2025: New checkpoints v0.9.6:\n- Release a new checkpoint [ltxv-2b-0.9.6-dev-04-25](https://huggingface.co/Lightricks/LTX-Video/blob/main/ltxv-2b-0.9.6-dev-04-25.safetensors) with improved quality\n- Release a new distilled model [ltxv-2b-0.9.6-distilled-04-25](https://huggingface.co/Lightricks/LTX-Video/blob/main/ltxv-2b-0.9.6-distilled-04-25.safetensors)\n * 15x faster inference than non-distilled model.\n * Does not require classifier-free guidance and spatio-temporal guidance.\n * Supports sampling with 8 (recommended), or less diffusion steps.\n- Improved prompt adherence, motion quality and fine details.\n- New default resolution and FPS: 1216 ร 704 pixels at 30 FPS\n * Still real time on H100 with the distilled model.\n * Other resolutions and FPS are still supported.\n- Support stochastic inference (can improve visual quality when using the distilled model)\n\n## March, 5th, 2025: New checkpoint v0.9.5\n- New license for commercial use ([OpenRail-M](https://huggingface.co/Lightricks/LTX-Video/ltx-video-2b-v0.9.5.license.txt))\n- Release a new checkpoint v0.9.5 with improved quality\n- Support keyframes and video extension\n- Support higher resolutions\n- Improved prompt understanding\n- Improved VAE\n- New online web app in [LTX-Studio](https://app.ltx.studio/ltx-video)\n- Automatic prompt enhancement\n\n## February, 20th, 2025: More inference options\n- Improve STG (Spatiotemporal Guidance) for LTX-Video\n- Support MPS on macOS with PyTorch 2.3.0\n- Add support for 8-bit model, LTX-VideoQ8\n- Add TeaCache for LTX-Video\n- Add [ComfyUI-LTXTricks](#comfyui-integration)\n- Add Diffusion-Pipe\n\n## December 31st, 2024: Research paper\n- Release the [research paper](https://arxiv.org/abs/2501.00103)\n\n## December 20th, 2024: New checkpoint v0.9.1\n- Release a new checkpoint v0.9.1 with improved quality\n- Support for STG / PAG\n- Support loading checkpoints of LTX-Video in Diffusers format (conversion is done on-the-fly)\n- Support offloading unused parts to CPU\n- Support the new timestep-conditioned VAE decoder\n- Reference contributions from the community in the readme file\n- Relax transformers dependency\n\n## November 21th, 2024: Initial release v0.9.0\n- Initial release of LTX-Video\n- Support text-to-video and image-to-video generation\n\n\n# Models\n\n| Name | Notes | inference.py config | ComfyUI workflow (Recommended) |\n|-------------------------|--------------------------------------------------------------------------------------------|----------------------------------------------------------------------------------------------------------------------------------------------------------|--------------------------------|\n| ltxv-13b-0.9.8-dev | Highest quality, requires more VRAM | [ltxv-13b-0.9.8-dev.yaml](https://github.com/Lightricks/LTX-Video/blob/main/configs/ltxv-13b-0.9.8-dev.yaml) | [ltxv-13b-i2v-base.json](https://github.com/Lightricks/ComfyUI-LTXVideo/blob/master/example_workflows/ltxv-13b-i2v-base.json) |\n| [ltxv-13b-0.9.8-mix](https://app.ltx.studio/motion-workspace?videoModel=ltxv-13b) | Mix ltxv-13b-dev and ltxv-13b-distilled in the same multi-scale rendering workflow for balanced speed-quality | N/A | [ltxv-13b-i2v-mixed-multiscale.json](https://github.com/Lightricks/ComfyUI-LTXVideo/blob/master/example_workflows/ltxv-13b-i2v-mixed-multiscale.json) |\n [ltxv-13b-0.9.8-distilled](https://app.ltx.studio/motion-workspace?videoModel=ltxv) | Faster, less VRAM usage, slight quality reduction compared to 13b. Ideal for rapid iterations | [ltxv-13b-0.9.8-distilled.yaml](https://github.com/Lightricks/LTX-Video/blob/main/configs/ltxv-13b-0.9.8-distilled.yaml) | [ltxv-13b-dist-i2v-base.json](https://github.com/Lightricks/ComfyUI-LTXVideo/blob/master/example_workflows/13b-distilled/ltxv-13b-dist-i2v-base.json) |\nltxv-2b-0.9.8-distilled | Smaller model, slight quality reduction compared to 13b distilled. Ideal for fast generation with light VRAM usage | [ltxv-2b-0.9.8-distilled.yaml](https://github.com/Lightricks/LTX-Video/blob/main/configs/ltxv-2b-0.9.8-distilled.yaml) | N/A |\n| ltxv-13b-0.9.8-dev-fp8 | Quantized version of ltxv-13b | [ltxv-13b-0.9.8-dev-fp8.yaml](https://github.com/Lightricks/LTX-Video/blob/main/configs/ltxv-13b-0.9.8-dev-fp8.yaml) | [ltxv-13b-i2v-base-fp8.json](https://github.com/Lightricks/ComfyUI-LTXVideo/blob/master/example_workflows/ltxv-13b-i2v-base-fp8.json) |\n| ltxv-13b-0.9.8-distilled-fp8 | Quantized version of ltxv-13b-distilled | [ltxv-13b-0.9.8-distilled-fp8.yaml](https://github.com/Lightricks/LTX-Video/blob/main/configs/ltxv-13b-0.9.8-distilled-fp8.yaml) | [ltxv-13b-dist-i2v-base-fp8.json](https://github.com/Lightricks/ComfyUI-LTXVideo/blob/master/example_workflows/13b-distilled/ltxv-13b-dist-i2v-base-fp8.json) |\n| ltxv-2b-0.9.8-distilled-fp8 | Quantized version of ltxv-2b-distilled | [ltxv-2b-0.9.8-distilled-fp8.yaml](https://github.com/Lightricks/LTX-Video/blob/main/configs/ltxv-2b-0.9.8-distilled-fp8.yaml) | N/A |\n| ltxv-2b-0.9.6 | Good quality, lower VRAM requirement than ltxv-13b | [ltxv-2b-0.9.6-dev.yaml](https://github.com/Lightricks/LTX-Video/blob/main/configs/ltxv-2b-0.9.6-dev.yaml) | [ltxvideo-i2v.json](https://github.com/Lightricks/ComfyUI-LTXVideo/blob/master/example_workflows/low_level/ltxvideo-i2v.json) |\n| ltxv-2b-0.9.6-distilled | 15ร faster, real-time capable, fewer steps needed, no STG/CFG required | [ltxv-2b-0.9.6-distilled.yaml](https://github.com/Lightricks/LTX-Video/blob/main/configs/ltxv-2b-0.9.6-distilled.yaml) | [ltxvideo-i2v-distilled.json](https://github.com/Lightricks/ComfyUI-LTXVideo/blob/master/example_workflows/low_level/ltxvideo-i2v-distilled.json) |\n\n\n# Quick Start Guide\n\n## Online inference\nThe model is accessible right away via the following links:\n- [LTX-Studio image-to-video (13B-mix)](https://app.ltx.studio/motion-workspace?videoModel=ltxv-13b)\n- [LTX-Studio image-to-video (13B distilled)](https://app.ltx.studio/motion-workspace?videoModel=ltxv)\n- [Fal.ai image-to-video (13B full)](https://fal.ai/models/fal-ai/ltx-video-13b-dev/image-to-video)\n- [Fal.ai image-to-video (13B distilled)](https://fal.ai/models/fal-ai/ltx-video-13b-distilled/image-to-video)\n- [Replicate image-to-video](https://replicate.com/lightricks/ltx-video)\n\n## Run locally\n\n### Installation\nThe codebase was tested with Python 3.10.5, CUDA version 12.2, and supports PyTorch >= 2.1.2.\nOn macOS, MPS was tested with PyTorch 2.3.0, and should support PyTorch == 2.3 or >= 2.6.\n\n```bash\ngit clone https://github.com/Lightricks/LTX-Video.git\ncd LTX-Video\n\n# create env\npython -m venv env\nsource env/bin/activate\npython -m pip install -e .\\[inference\\]\n```\n\n#### FP8 Kernels (optional)\n\n[FP8 kernels](https://github.com/Lightricks/LTXVideo-Q8-Kernels) developed for LTX-Video provide performance boost on supported graphics cards (Ada architecture and later). To install FP8 kernels, follow the instructions in that repository.\n\n### Inference\n\n๐ **Note:** For best results, we recommend using our [ComfyUI](#comfyui-integration) workflow. We're working on updating the inference.py script to match the high quality and output fidelity of ComfyUI.\n\nTo use our model, please follow the inference code in [inference.py](./inference.py):\n\n#### For image-to-video generation:\n\n```bash\npython inference.py --prompt \"PROMPT\" --conditioning_media_paths IMAGE_PATH --conditioning_start_frames 0 --height HEIGHT --width WIDTH --num_frames NUM_FRAMES --seed SEED --pipeline_config configs/ltxv-13b-0.9.8-distilled.yaml\n```\n\n#### Extending a video:\n\n๐ **Note:** Input video segments must contain a multiple of 8 frames plus 1 (e.g., 9, 17, 25, etc.), and the target frame number should be a multiple of 8.\n\n\n```bash\npython inference.py --prompt \"PROMPT\" --conditioning_media_paths VIDEO_PATH --conditioning_start_frames START_FRAME --height HEIGHT --width WIDTH --num_frames NUM_FRAMES --seed SEED --pipeline_config configs/ltxv-13b-0.9.8-distilled.yaml\n```\n\n#### For video generation with multiple conditions:\n\nYou can now generate a video conditioned on a set of images and/or short video segments.\nSimply provide a list of paths to the images or video segments you want to condition on, along with their target frame numbers in the generated video. You can also specify the conditioning strength for each item (default: 1.0).\n\n```bash\npython inference.py --prompt \"PROMPT\" --conditioning_media_paths IMAGE_OR_VIDEO_PATH_1 IMAGE_OR_VIDEO_PATH_2 --conditioning_start_frames TARGET_FRAME_1 TARGET_FRAME_2 --height HEIGHT --width WIDTH --num_frames NUM_FRAMES --seed SEED --pipeline_config configs/ltxv-13b-0.9.8-distilled.yaml\n```\n\n### Using as a library\n\n```python\nfrom ltx_video.inference import infer, InferenceConfig\n\ninfer(\n InferenceConfig(\n pipeline_config=\"configs/ltxv-13b-0.9.8-distilled.yaml\",\n prompt=PROMPT,\n height=HEIGHT,\n width=WIDTH,\n num_frames=NUM_FRAMES,\n output_path=\"output.mp4\",\n )\n)\n```\n\n## ComfyUI Integration\nTo use our model with ComfyUI, please follow the instructions at [https://github.com/Lightricks/ComfyUI-LTXVideo/](https://github.com/Lightricks/ComfyUI-LTXVideo/).\n\n## Diffusers Integration\nTo use our model with the Diffusers Python library, check out the [official documentation](https://huggingface.co/docs/diffusers/main/en/api/pipelines/ltx_video).\n\nDiffusers also support an 8-bit version of LTX-Video, [see details below](#ltx-videoq8)\n\n# Model User Guide\n\n## ๐ Prompt Engineering\n\nWhen writing prompts, focus on detailed, chronological descriptions of actions and scenes. Include specific movements, appearances, camera angles, and environmental details - all in a single flowing paragraph. Start directly with the action, and keep descriptions literal and precise. Think like a cinematographer describing a shot list. Keep within 200 words. For best results, build your prompts using this structure:\n\n* Start with main action in a single sentence\n* Add specific details about movements and gestures\n* Describe character/object appearances precisely\n* Include background and environment details\n* Specify camera angles and movements\n* Describe lighting and colors\n* Note any changes or sudden events\n* See [examples](#introduction) for more inspiration.\n\n### Automatic Prompt Enhancement\n\nWhen using `LTXVideoPipeline` directly, you can enable prompt enhancement by setting `enhance_prompt=True`.\n\n## ๐ฎ Parameter Guide\n\n* Resolution Preset: Higher resolutions for detailed scenes, lower for faster generation and simpler scenes. The model works on resolutions that are divisible by 32 and number of frames that are divisible by 8 + 1 (e.g. 257). In case the resolution or number of frames are not divisible by 32 or 8 + 1, the input will be padded with -1 and then cropped to the desired resolution and number of frames. The model works best on resolutions under 720 x 1280 and number of frames below 257\n* Seed: Save seed values to recreate specific styles or compositions you like\n* Guidance Scale: 3-3.5 are the recommended values\n* Inference Steps: More steps (40+) for quality, fewer steps (20-30) for speed\n\n๐ For advanced parameters usage, please see `python inference.py --help`\n\n## Community Contribution\n\n### ComfyUI-LTXTricks ๐ ๏ธ\n\nA community project providing additional nodes for enhanced control over the LTX Video model. It includes implementations of advanced techniques like RF-Inversion, RF-Edit, FlowEdit, and more. These nodes enable workflows such as Image and Video to Video (I+V2V), enhanced sampling via Spatiotemporal Skip Guidance (STG), and interpolation with precise frame settings.\n\n- **Repository:** [ComfyUI-LTXTricks](https://github.com/logtd/ComfyUI-LTXTricks)\n- **Features:**\n - ๐ **RF-Inversion:** Implements [RF-Inversion](https://rf-inversion.github.io/) with an [example workflow here](https://github.com/logtd/ComfyUI-LTXTricks/blob/main/example_workflows/example_ltx_inversion.json).\n - โ๏ธ **RF-Edit:** Implements [RF-Solver-Edit](https://github.com/wangjiangshan0725/RF-Solver-Edit) with an [example workflow here](https://github.com/logtd/ComfyUI-LTXTricks/blob/main/example_workflows/example_ltx_rf_edit.json).\n - ๐ **FlowEdit:** Implements [FlowEdit](https://github.com/fallenshock/FlowEdit) with an [example workflow here](https://github.com/logtd/ComfyUI-LTXTricks/blob/main/example_workflows/example_ltx_flow_edit.json).\n - ๐ฅ **I+V2V:** Enables Video to Video with a reference image. [Example workflow](https://github.com/logtd/ComfyUI-LTXTricks/blob/main/example_workflows/example_ltx_iv2v.json).\n - โจ **Enhance:** Partial implementation of [STGuidance](https://junhahyung.github.io/STGuidance/). [Example workflow](https://github.com/logtd/ComfyUI-LTXTricks/blob/main/example_workflows/example_ltxv_stg.json).\n - ๐ผ๏ธ **Interpolation and Frame Setting:** Nodes for precise control of latents per frame. [Example workflow](https://github.com/logtd/ComfyUI-LTXTricks/blob/main/example_workflows/example_ltx_interpolation.json).\n\n\n### LTX-VideoQ8 ๐ฑ \n\n**LTX-VideoQ8** is an 8-bit optimized version of [LTX-Video](https://github.com/Lightricks/LTX-Video), designed for faster performance on NVIDIA ADA GPUs.\n\n- **Repository:** [LTX-VideoQ8](https://github.com/KONAKONA666/LTX-Video)\n- **Features:**\n - ๐ Up to 3X speed-up with no accuracy loss\n - ๐ฅ Generate 720x480x121 videos in under a minute on RTX 4060 (8GB VRAM)\n - ๐ ๏ธ Fine-tune 2B transformer models with precalculated latents\n- **Community Discussion:** [Reddit Thread](https://www.reddit.com/r/StableDiffusion/comments/1h79ks2/fast_ltx_video_on_rtx_4060_and_other_ada_gpus/)\n- **Diffusers integration:** A diffusers integration for the 8-bit model is already out! [Details here](https://github.com/sayakpaul/q8-ltx-video)\n\n\n### TeaCache for LTX-Video ๐ต \n\n**TeaCache** is a training-free caching approach that leverages timestep differences across model outputs to accelerate LTX-Video inference by up to 2x without significant visual quality degradation.\n\n- **Repository:** [TeaCache4LTX-Video](https://github.com/ali-vilab/TeaCache/tree/main/TeaCache4LTX-Video)\n- **Features:**\n - ๐ Speeds up LTX-Video inference.\n - ๐ Adjustable trade-offs between speed (up to 2x) and visual quality using configurable parameters.\n - ๐ ๏ธ No retraining required: Works directly with existing models.\n\n### Your Contribution\n\n...is welcome! If you have a project or tool that integrates with LTX-Video,\nplease let us know by opening an issue or pull request.\n\n# Training\n\nWe provide an open-source repository for fine-tuning the LTX-Video model: [LTX-Video-Trainer](https://github.com/Lightricks/LTX-Video-Trainer).\nThis repository supports both the 2B and 13B model variants, enabling full fine-tuning as well as LoRA (Low-Rank Adaptation) fine-tuning for more efficient training. This includes:\n\n- **Control LoRAs**: Train custom control models like depth, pose, and canny control\n- **Effect LoRAs**: Create specialized effects and transformations for video generation\n\nExplore the repository to customize the model for your specific use cases!\nMore information and training instructions can be found in the [README](https://github.com/Lightricks/LTX-Video-Trainer/blob/main/README.md).\n\n# Control Models\n\n[ComfyUI-LTXVideo](https://github.com/Lightricks/ComfyUI-LTXVideo) repository now contains workflows and models for 3 specialized models that enable precise control over LTX-Video generation:\n\nPose Control, Depth Control and Canny Control\n\n**Example ComfyUI Workflow (for all control types):** [ic-lora.json](https://github.com/Lightricks/ComfyUI-LTXVideo/blob/master/example_workflows/ic_lora/ic-lora.json)\n\n# Join Us\n\nWant to work on cutting-edge AI research and make a real impact on millions of users worldwide?\n\nAt **Lightricks**, an AI-first company, we're revolutionizing how visual content is created.\n\nIf you are passionate about AI, computer vision, and video generation, we would love to hear from you!\n\nPlease visit our [careers page](https://careers.lightricks.com/careers?query=&office=all&department=R%26D) for more information.\n\n# Acknowledgement\n\nWe are grateful for the following awesome projects when implementing LTX-Video:\n* [DiT](https://github.com/facebookresearch/DiT) and [PixArt-alpha](https://github.com/PixArt-alpha/PixArt-alpha): vision transformers for image generation.\n\n\n## Citation\n\n๐ Our tech report is out! If you find our work helpful, please โญ๏ธ star the repository and cite our paper.\n\n```\n@article{HaCohen2024LTXVideo,\n title={LTX-Video: Realtime Video Latent Diffusion},\n author={HaCohen, Yoav and Chiprut, Nisan and Brazowski, Benny and Shalem, Daniel and Moshe, Dudu and Richardson, Eitan and Levin, Eran and Shiran, Guy and Zabari, Nir and Gordon, Ori and Panet, Poriya and Weissbuch, Sapir and Kulikov, Victor and Bitterman, Yaki and Melumian, Zeev and Bibi, Ofir},\n journal={arXiv preprint arXiv:2501.00103},\n year={2024}\n}\n```\n"
},
{
"path": "configs/ltxv-13b-0.9.8-dev-fp8.yaml",
"content": "pipeline_type: multi-scale\ncheckpoint_path: \"ltxv-13b-0.9.8-dev-fp8.safetensors\"\ndownscale_factor: 0.6666666\nspatial_upscaler_model_path: \"ltxv-spatial-upscaler-0.9.8.safetensors\"\nstg_mode: \"attention_values\" # options: \"attention_values\", \"attention_skip\", \"residual\", \"transformer_block\"\ndecode_timestep: 0.05\ndecode_noise_scale: 0.025\ntext_encoder_model_name_or_path: \"PixArt-alpha/PixArt-XL-2-1024-MS\"\nprecision: \"float8_e4m3fn\" # options: \"float8_e4m3fn\", \"bfloat16\", \"mixed_precision\"\nsampler: \"from_checkpoint\" # options: \"uniform\", \"linear-quadratic\", \"from_checkpoint\"\nprompt_enhancement_words_threshold: 120\nprompt_enhancer_image_caption_model_name_or_path: \"MiaoshouAI/Florence-2-large-PromptGen-v2.0\"\nprompt_enhancer_llm_model_name_or_path: \"unsloth/Llama-3.2-3B-Instruct\"\nstochastic_sampling: false\n\nfirst_pass:\n guidance_scale: [1, 1, 6, 8, 6, 1, 1]\n stg_scale: [0, 0, 4, 4, 4, 2, 1]\n rescaling_scale: [1, 1, 0.5, 0.5, 1, 1, 1]\n guidance_timesteps: [1.0, 0.996, 0.9933, 0.9850, 0.9767, 0.9008, 0.6180]\n skip_block_list: [[], [11, 25, 35, 39], [22, 35, 39], [28], [28], [28], [28]]\n num_inference_steps: 30\n skip_final_inference_steps: 3\n cfg_star_rescale: true\n\nsecond_pass:\n guidance_scale: [1]\n stg_scale: [1]\n rescaling_scale: [1]\n guidance_timesteps: [1.0]\n skip_block_list: [27]\n num_inference_steps: 30\n skip_initial_inference_steps: 17\n cfg_star_rescale: true\n"
},
{
"path": "configs/ltxv-13b-0.9.8-dev.yaml",
"content": "pipeline_type: multi-scale\ncheckpoint_path: \"ltxv-13b-0.9.8-dev.safetensors\"\ndownscale_factor: 0.6666666\nspatial_upscaler_model_path: \"ltxv-spatial-upscaler-0.9.8.safetensors\"\nstg_mode: \"attention_values\" # options: \"attention_values\", \"attention_skip\", \"residual\", \"transformer_block\"\ndecode_timestep: 0.05\ndecode_noise_scale: 0.025\ntext_encoder_model_name_or_path: \"PixArt-alpha/PixArt-XL-2-1024-MS\"\nprecision: \"bfloat16\"\nsampler: \"from_checkpoint\" # options: \"uniform\", \"linear-quadratic\", \"from_checkpoint\"\nprompt_enhancement_words_threshold: 120\nprompt_enhancer_image_caption_model_name_or_path: \"MiaoshouAI/Florence-2-large-PromptGen-v2.0\"\nprompt_enhancer_llm_model_name_or_path: \"unsloth/Llama-3.2-3B-Instruct\"\nstochastic_sampling: false\n\nfirst_pass:\n guidance_scale: [1, 1, 6, 8, 6, 1, 1]\n stg_scale: [0, 0, 4, 4, 4, 2, 1]\n rescaling_scale: [1, 1, 0.5, 0.5, 1, 1, 1]\n guidance_timesteps: [1.0, 0.996, 0.9933, 0.9850, 0.9767, 0.9008, 0.6180]\n skip_block_list: [[], [11, 25, 35, 39], [22, 35, 39], [28], [28], [28], [28]]\n num_inference_steps: 30\n skip_final_inference_steps: 3\n cfg_star_rescale: true\n\nsecond_pass:\n guidance_scale: [1]\n stg_scale: [1]\n rescaling_scale: [1]\n guidance_timesteps: [1.0]\n skip_block_list: [27]\n num_inference_steps: 30\n skip_initial_inference_steps: 17\n cfg_star_rescale: true"
},
{
"path": "configs/ltxv-13b-0.9.8-distilled-fp8.yaml",
"content": "pipeline_type: multi-scale\ncheckpoint_path: \"ltxv-13b-0.9.8-distilled-fp8.safetensors\"\ndownscale_factor: 0.6666666\nspatial_upscaler_model_path: \"ltxv-spatial-upscaler-0.9.8.safetensors\"\nstg_mode: \"attention_values\" # options: \"attention_values\", \"attention_skip\", \"residual\", \"transformer_block\"\ndecode_timestep: 0.05\ndecode_noise_scale: 0.025\ntext_encoder_model_name_or_path: \"PixArt-alpha/PixArt-XL-2-1024-MS\"\nprecision: \"float8_e4m3fn\" # options: \"float8_e4m3fn\", \"bfloat16\", \"mixed_precision\"\nsampler: \"from_checkpoint\" # options: \"uniform\", \"linear-quadratic\", \"from_checkpoint\"\nprompt_enhancement_words_threshold: 120\nprompt_enhancer_image_caption_model_name_or_path: \"MiaoshouAI/Florence-2-large-PromptGen-v2.0\"\nprompt_enhancer_llm_model_name_or_path: \"unsloth/Llama-3.2-3B-Instruct\"\nstochastic_sampling: false\n\nfirst_pass:\n timesteps: [1.0000, 0.9937, 0.9875, 0.9812, 0.9750, 0.9094, 0.7250]\n guidance_scale: 1\n stg_scale: 0\n rescaling_scale: 1\n skip_block_list: [42]\n\nsecond_pass:\n timesteps: [0.9094, 0.7250, 0.4219]\n guidance_scale: 1\n stg_scale: 0\n rescaling_scale: 1\n skip_block_list: [42]\n tone_map_compression_ratio: 0.6\n"
},
{
"path": "configs/ltxv-13b-0.9.8-distilled.yaml",
"content": "pipeline_type: multi-scale\ncheckpoint_path: \"ltxv-13b-0.9.8-distilled.safetensors\"\ndownscale_factor: 0.6666666\nspatial_upscaler_model_path: \"ltxv-spatial-upscaler-0.9.8.safetensors\"\nstg_mode: \"attention_values\" # options: \"attention_values\", \"attention_skip\", \"residual\", \"transformer_block\"\ndecode_timestep: 0.05\ndecode_noise_scale: 0.025\ntext_encoder_model_name_or_path: \"PixArt-alpha/PixArt-XL-2-1024-MS\"\nprecision: \"bfloat16\"\nsampler: \"from_checkpoint\" # options: \"uniform\", \"linear-quadratic\", \"from_checkpoint\"\nprompt_enhancement_words_threshold: 120\nprompt_enhancer_image_caption_model_name_or_path: \"MiaoshouAI/Florence-2-large-PromptGen-v2.0\"\nprompt_enhancer_llm_model_name_or_path: \"unsloth/Llama-3.2-3B-Instruct\"\nstochastic_sampling: false\n\nfirst_pass:\n timesteps: [1.0000, 0.9937, 0.9875, 0.9812, 0.9750, 0.9094, 0.7250]\n guidance_scale: 1\n stg_scale: 0\n rescaling_scale: 1\n skip_block_list: [42]\n\nsecond_pass:\n timesteps: [0.9094, 0.7250, 0.4219]\n guidance_scale: 1\n stg_scale: 0\n rescaling_scale: 1\n skip_block_list: [42]\n tone_map_compression_ratio: 0.6\n"
},
{
"path": "configs/ltxv-2b-0.9.1.yaml",
"content": "pipeline_type: base\ncheckpoint_path: \"ltx-video-2b-v0.9.1.safetensors\"\nguidance_scale: 3\nstg_scale: 1\nrescaling_scale: 0.7\nskip_block_list: [19]\nnum_inference_steps: 40\nstg_mode: \"attention_values\" # options: \"attention_values\", \"attention_skip\", \"residual\", \"transformer_block\"\ndecode_timestep: 0.05\ndecode_noise_scale: 0.025\ntext_encoder_model_name_or_path: \"PixArt-alpha/PixArt-XL-2-1024-MS\"\nprecision: \"bfloat16\"\nsampler: \"from_checkpoint\" # options: \"uniform\", \"linear-quadratic\", \"from_checkpoint\"\nprompt_enhancement_words_threshold: 120\nprompt_enhancer_image_caption_model_name_or_path: \"MiaoshouAI/Florence-2-large-PromptGen-v2.0\"\nprompt_enhancer_llm_model_name_or_path: \"unsloth/Llama-3.2-3B-Instruct\"\nstochastic_sampling: false"
},
{
"path": "configs/ltxv-2b-0.9.5.yaml",
"content": "pipeline_type: base\ncheckpoint_path: \"ltx-video-2b-v0.9.5.safetensors\"\nguidance_scale: 3\nstg_scale: 1\nrescaling_scale: 0.7\nskip_block_list: [19]\nnum_inference_steps: 40\nstg_mode: \"attention_values\" # options: \"attention_values\", \"attention_skip\", \"residual\", \"transformer_block\"\ndecode_timestep: 0.05\ndecode_noise_scale: 0.025\ntext_encoder_model_name_or_path: \"PixArt-alpha/PixArt-XL-2-1024-MS\"\nprecision: \"bfloat16\"\nsampler: \"from_checkpoint\" # options: \"uniform\", \"linear-quadratic\", \"from_checkpoint\"\nprompt_enhancement_words_threshold: 120\nprompt_enhancer_image_caption_model_name_or_path: \"MiaoshouAI/Florence-2-large-PromptGen-v2.0\"\nprompt_enhancer_llm_model_name_or_path: \"unsloth/Llama-3.2-3B-Instruct\"\nstochastic_sampling: false"
},
{
"path": "configs/ltxv-2b-0.9.6-dev.yaml",
"content": "pipeline_type: base\ncheckpoint_path: \"ltxv-2b-0.9.6-dev-04-25.safetensors\"\nguidance_scale: 3\nstg_scale: 1\nrescaling_scale: 0.7\nskip_block_list: [19]\nnum_inference_steps: 40\nstg_mode: \"attention_values\" # options: \"attention_values\", \"attention_skip\", \"residual\", \"transformer_block\"\ndecode_timestep: 0.05\ndecode_noise_scale: 0.025\ntext_encoder_model_name_or_path: \"PixArt-alpha/PixArt-XL-2-1024-MS\"\nprecision: \"bfloat16\"\nsampler: \"from_checkpoint\" # options: \"uniform\", \"linear-quadratic\", \"from_checkpoint\"\nprompt_enhancement_words_threshold: 120\nprompt_enhancer_image_caption_model_name_or_path: \"MiaoshouAI/Florence-2-large-PromptGen-v2.0\"\nprompt_enhancer_llm_model_name_or_path: \"unsloth/Llama-3.2-3B-Instruct\"\nstochastic_sampling: false"
},
{
"path": "configs/ltxv-2b-0.9.6-distilled.yaml",
"content": "pipeline_type: base\ncheckpoint_path: \"ltxv-2b-0.9.6-distilled-04-25.safetensors\"\nguidance_scale: 1\nstg_scale: 0\nrescaling_scale: 1\nnum_inference_steps: 8\nstg_mode: \"attention_values\" # options: \"attention_values\", \"attention_skip\", \"residual\", \"transformer_block\"\ndecode_timestep: 0.05\ndecode_noise_scale: 0.025\ntext_encoder_model_name_or_path: \"PixArt-alpha/PixArt-XL-2-1024-MS\"\nprecision: \"bfloat16\"\nsampler: \"from_checkpoint\" # options: \"uniform\", \"linear-quadratic\", \"from_checkpoint\"\nprompt_enhancement_words_threshold: 120\nprompt_enhancer_image_caption_model_name_or_path: \"MiaoshouAI/Florence-2-large-PromptGen-v2.0\"\nprompt_enhancer_llm_model_name_or_path: \"unsloth/Llama-3.2-3B-Instruct\"\nstochastic_sampling: true"
},
{
"path": "configs/ltxv-2b-0.9.8-distilled-fp8.yaml",
"content": "pipeline_type: multi-scale\ncheckpoint_path: \"ltxv-2b-0.9.8-distilled-fp8.safetensors\"\ndownscale_factor: 0.6666666\nspatial_upscaler_model_path: \"ltxv-spatial-upscaler-0.9.8.safetensors\"\nstg_mode: \"attention_values\" # options: \"attention_values\", \"attention_skip\", \"residual\", \"transformer_block\"\ndecode_timestep: 0.05\ndecode_noise_scale: 0.025\ntext_encoder_model_name_or_path: \"PixArt-alpha/PixArt-XL-2-1024-MS\"\nprecision: \"float8_e4m3fn\" # options: \"float8_e4m3fn\", \"bfloat16\", \"mixed_precision\"\nsampler: \"from_checkpoint\" # options: \"uniform\", \"linear-quadratic\", \"from_checkpoint\"\nprompt_enhancement_words_threshold: 120\nprompt_enhancer_image_caption_model_name_or_path: \"MiaoshouAI/Florence-2-large-PromptGen-v2.0\"\nprompt_enhancer_llm_model_name_or_path: \"unsloth/Llama-3.2-3B-Instruct\"\nstochastic_sampling: false\n\nfirst_pass:\n timesteps: [1.0000, 0.9937, 0.9875, 0.9812, 0.9750, 0.9094, 0.7250]\n guidance_scale: 1\n stg_scale: 0\n rescaling_scale: 1\n skip_block_list: [42]\n\nsecond_pass:\n timesteps: [0.9094, 0.7250, 0.4219]\n guidance_scale: 1\n stg_scale: 0\n rescaling_scale: 1\n skip_block_list: [42]\n"
},
{
"path": "configs/ltxv-2b-0.9.8-distilled.yaml",
"content": "pipeline_type: multi-scale\ncheckpoint_path: \"ltxv-2b-0.9.8-distilled.safetensors\"\ndownscale_factor: 0.6666666\nspatial_upscaler_model_path: \"ltxv-spatial-upscaler-0.9.8.safetensors\"\nstg_mode: \"attention_values\" # options: \"attention_values\", \"attention_skip\", \"residual\", \"transformer_block\"\ndecode_timestep: 0.05\ndecode_noise_scale: 0.025\ntext_encoder_model_name_or_path: \"PixArt-alpha/PixArt-XL-2-1024-MS\"\nprecision: \"bfloat16\"\nsampler: \"from_checkpoint\" # options: \"uniform\", \"linear-quadratic\", \"from_checkpoint\"\nprompt_enhancement_words_threshold: 120\nprompt_enhancer_image_caption_model_name_or_path: \"MiaoshouAI/Florence-2-large-PromptGen-v2.0\"\nprompt_enhancer_llm_model_name_or_path: \"unsloth/Llama-3.2-3B-Instruct\"\nstochastic_sampling: false\n\nfirst_pass:\n timesteps: [1.0000, 0.9937, 0.9875, 0.9812, 0.9750, 0.9094, 0.7250]\n guidance_scale: 1\n stg_scale: 0\n rescaling_scale: 1\n skip_block_list: [42]\n\nsecond_pass:\n timesteps: [0.9094, 0.7250, 0.4219]\n guidance_scale: 1\n stg_scale: 0\n rescaling_scale: 1\n skip_block_list: [42]\n"
},
{
"path": "configs/ltxv-2b-0.9.yaml",
"content": "pipeline_type: base\ncheckpoint_path: \"ltx-video-2b-v0.9.safetensors\"\nguidance_scale: 3\nstg_scale: 1\nrescaling_scale: 0.7\nskip_block_list: [19]\nnum_inference_steps: 40\nstg_mode: \"attention_values\" # options: \"attention_values\", \"attention_skip\", \"residual\", \"transformer_block\"\ndecode_timestep: 0.05\ndecode_noise_scale: 0.025\ntext_encoder_model_name_or_path: \"PixArt-alpha/PixArt-XL-2-1024-MS\"\nprecision: \"bfloat16\"\nsampler: \"from_checkpoint\" # options: \"uniform\", \"linear-quadratic\", \"from_checkpoint\"\nprompt_enhancement_words_threshold: 120\nprompt_enhancer_image_caption_model_name_or_path: \"MiaoshouAI/Florence-2-large-PromptGen-v2.0\"\nprompt_enhancer_llm_model_name_or_path: \"unsloth/Llama-3.2-3B-Instruct\"\nstochastic_sampling: false"
},
{
"path": "inference.py",
"content": "from transformers import HfArgumentParser\n\nfrom ltx_video.inference import infer, InferenceConfig\n\n\ndef main():\n parser = HfArgumentParser(InferenceConfig)\n config = parser.parse_args_into_dataclasses()[0]\n infer(config=config)\n\n\nif __name__ == \"__main__\":\n main()\n"
},
{
"path": "ltx_video/__init__.py",
"content": ""
},
{
"path": "ltx_video/inference.py",
"content": "import os\nimport random\nfrom datetime import datetime\nfrom pathlib import Path\nfrom diffusers.utils import logging\nfrom typing import Optional, List, Union\nimport yaml\n\nimport imageio\nimport json\nimport numpy as np\nimport torch\nfrom safetensors import safe_open\nfrom PIL import Image\nimport torchvision.transforms.functional as TVF\nfrom transformers import (\n T5EncoderModel,\n T5Tokenizer,\n AutoModelForCausalLM,\n AutoProcessor,\n AutoTokenizer,\n)\nfrom huggingface_hub import hf_hub_download\nfrom dataclasses import dataclass, field\n\nfrom ltx_video.models.autoencoders.causal_video_autoencoder import (\n CausalVideoAutoencoder,\n)\nfrom ltx_video.models.transformers.symmetric_patchifier import SymmetricPatchifier\nfrom ltx_video.models.transformers.transformer3d import Transformer3DModel\nfrom ltx_video.pipelines.pipeline_ltx_video import (\n ConditioningItem,\n LTXVideoPipeline,\n LTXMultiScalePipeline,\n)\nfrom ltx_video.schedulers.rf import RectifiedFlowScheduler\nfrom ltx_video.utils.skip_layer_strategy import SkipLayerStrategy\nfrom ltx_video.models.autoencoders.latent_upsampler import LatentUpsampler\nimport ltx_video.pipelines.crf_compressor as crf_compressor\n\nlogger = logging.get_logger(\"LTX-Video\")\n\n\ndef get_total_gpu_memory():\n if torch.cuda.is_available():\n total_memory = torch.cuda.get_device_properties(0).total_memory / (1024**3)\n return total_memory\n return 0\n\n\ndef get_device():\n if torch.cuda.is_available():\n return \"cuda\"\n elif torch.backends.mps.is_available():\n return \"mps\"\n return \"cpu\"\n\n\ndef load_image_to_tensor_with_resize_and_crop(\n image_input: Union[str, Image.Image],\n target_height: int = 512,\n target_width: int = 768,\n just_crop: bool = False,\n) -> torch.Tensor:\n \"\"\"Load and process an image into a tensor.\n\n Args:\n image_input: Either a file path (str) or a PIL Image object\n target_height: Desired height of output tensor\n target_width: Desired width of output tensor\n just_crop: If True, only crop the image to the target size without resizing\n \"\"\"\n if isinstance(image_input, str):\n image = Image.open(image_input).convert(\"RGB\")\n elif isinstance(image_input, Image.Image):\n image = image_input\n else:\n raise ValueError(\"image_input must be either a file path or a PIL Image object\")\n\n input_width, input_height = image.size\n aspect_ratio_target = target_width / target_height\n aspect_ratio_frame = input_width / input_height\n if aspect_ratio_frame > aspect_ratio_target:\n new_width = int(input_height * aspect_ratio_target)\n new_height = input_height\n x_start = (input_width - new_width) // 2\n y_start = 0\n else:\n new_width = input_width\n new_height = int(input_width / aspect_ratio_target)\n x_start = 0\n y_start = (input_height - new_height) // 2\n\n image = image.crop((x_start, y_start, x_start + new_width, y_start + new_height))\n if not just_crop:\n image = image.resize((target_width, target_height))\n\n frame_tensor = TVF.to_tensor(image) # PIL -> tensor (C, H, W), [0,1]\n frame_tensor = TVF.gaussian_blur(frame_tensor, kernel_size=3, sigma=1.0)\n frame_tensor_hwc = frame_tensor.permute(1, 2, 0) # (C, H, W) -> (H, W, C)\n frame_tensor_hwc = crf_compressor.compress(frame_tensor_hwc)\n frame_tensor = frame_tensor_hwc.permute(2, 0, 1) * 255.0 # (H, W, C) -> (C, H, W)\n frame_tensor = (frame_tensor / 127.5) - 1.0\n # Create 5D tensor: (batch_size=1, channels=3, num_frames=1, height, width)\n return frame_tensor.unsqueeze(0).unsqueeze(2)\n\n\ndef calculate_padding(\n source_height: int, source_width: int, target_height: int, target_width: int\n) -> tuple[int, int, int, int]:\n\n # Calculate total padding needed\n pad_height = target_height - source_height\n pad_width = target_width - source_width\n\n # Calculate padding for each side\n pad_top = pad_height // 2\n pad_bottom = pad_height - pad_top # Handles odd padding\n pad_left = pad_width // 2\n pad_right = pad_width - pad_left # Handles odd padding\n\n # Return padded tensor\n # Padding format is (left, right, top, bottom)\n padding = (pad_left, pad_right, pad_top, pad_bottom)\n return padding\n\n\ndef convert_prompt_to_filename(text: str, max_len: int = 20) -> str:\n # Remove non-letters and convert to lowercase\n clean_text = \"\".join(\n char.lower() for char in text if char.isalpha() or char.isspace()\n )\n\n # Split into words\n words = clean_text.split()\n\n # Build result string keeping track of length\n result = []\n current_length = 0\n\n for word in words:\n # Add word length plus 1 for underscore (except for first word)\n new_length = current_length + len(word)\n\n if new_length <= max_len:\n result.append(word)\n current_length += len(word)\n else:\n break\n\n return \"-\".join(result)\n\n\n# Generate output video name\ndef get_unique_filename(\n base: str,\n ext: str,\n prompt: str,\n seed: int,\n resolution: tuple[int, int, int],\n dir: Path,\n endswith=None,\n index_range=1000,\n) -> Path:\n base_filename = f\"{base}_{convert_prompt_to_filename(prompt, max_len=30)}_{seed}_{resolution[0]}x{resolution[1]}x{resolution[2]}\"\n for i in range(index_range):\n filename = dir / f\"{base_filename}_{i}{endswith if endswith else ''}{ext}\"\n if not os.path.exists(filename):\n return filename\n raise FileExistsError(\n f\"Could not find a unique filename after {index_range} attempts.\"\n )\n\n\ndef seed_everething(seed: int):\n random.seed(seed)\n np.random.seed(seed)\n torch.manual_seed(seed)\n if torch.cuda.is_available():\n torch.cuda.manual_seed(seed)\n if torch.backends.mps.is_available():\n torch.mps.manual_seed(seed)\n\n\ndef create_transformer(ckpt_path: str, precision: str) -> Transformer3DModel:\n if precision == \"float8_e4m3fn\":\n try:\n from q8_kernels.integration.patch_transformer import (\n patch_diffusers_transformer as patch_transformer_for_q8_kernels,\n )\n\n transformer = Transformer3DModel.from_pretrained(\n ckpt_path, dtype=torch.float8_e4m3fn\n )\n patch_transformer_for_q8_kernels(transformer)\n return transformer\n except ImportError:\n raise ValueError(\n \"Q8-Kernels not found. To use FP8 checkpoint, please install Q8 kernels from https://github.com/Lightricks/LTXVideo-Q8-Kernels\"\n )\n elif precision == \"bfloat16\":\n return Transformer3DModel.from_pretrained(ckpt_path).to(torch.bfloat16)\n else:\n return Transformer3DModel.from_pretrained(ckpt_path)\n\n\ndef create_ltx_video_pipeline(\n ckpt_path: str,\n precision: str,\n text_encoder_model_name_or_path: str,\n sampler: Optional[str] = None,\n device: Optional[str] = None,\n enhance_prompt: bool = False,\n prompt_enhancer_image_caption_model_name_or_path: Optional[str] = None,\n prompt_enhancer_llm_model_name_or_path: Optional[str] = None,\n) -> LTXVideoPipeline:\n ckpt_path = Path(ckpt_path)\n assert os.path.exists(\n ckpt_path\n ), f\"Ckpt path provided (--ckpt_path) {ckpt_path} does not exist\"\n\n with safe_open(ckpt_path, framework=\"pt\") as f:\n metadata = f.metadata()\n config_str = metadata.get(\"config\")\n configs = json.loads(config_str)\n allowed_inference_steps = configs.get(\"allowed_inference_steps\", None)\n\n vae = CausalVideoAutoencoder.from_pretrained(ckpt_path)\n transformer = create_transformer(ckpt_path, precision)\n\n # Use constructor if sampler is specified, otherwise use from_pretrained\n if sampler == \"from_checkpoint\" or not sampler:\n scheduler = RectifiedFlowScheduler.from_pretrained(ckpt_path)\n else:\n scheduler = RectifiedFlowScheduler(\n sampler=(\"Uniform\" if sampler.lower() == \"uniform\" else \"LinearQuadratic\")\n )\n\n text_encoder = T5EncoderModel.from_pretrained(\n text_encoder_model_name_or_path, subfolder=\"text_encoder\"\n )\n patchifier = SymmetricPatchifier(patch_size=1)\n tokenizer = T5Tokenizer.from_pretrained(\n text_encoder_model_name_or_path, subfolder=\"tokenizer\"\n )\n\n transformer = transformer.to(device)\n vae = vae.to(device)\n text_encoder = text_encoder.to(device)\n\n if enhance_prompt:\n prompt_enhancer_image_caption_model = AutoModelForCausalLM.from_pretrained(\n prompt_enhancer_image_caption_model_name_or_path, trust_remote_code=True\n )\n prompt_enhancer_image_caption_processor = AutoProcessor.from_pretrained(\n prompt_enhancer_image_caption_model_name_or_path, trust_remote_code=True\n )\n prompt_enhancer_llm_model = AutoModelForCausalLM.from_pretrained(\n prompt_enhancer_llm_model_name_or_path,\n torch_dtype=\"bfloat16\",\n )\n prompt_enhancer_llm_tokenizer = AutoTokenizer.from_pretrained(\n prompt_enhancer_llm_model_name_or_path,\n )\n else:\n prompt_enhancer_image_caption_model = None\n prompt_enhancer_image_caption_processor = None\n prompt_enhancer_llm_model = None\n prompt_enhancer_llm_tokenizer = None\n\n vae = vae.to(torch.bfloat16)\n text_encoder = text_encoder.to(torch.bfloat16)\n\n # Use submodels for the pipeline\n submodel_dict = {\n \"transformer\": transformer,\n \"patchifier\": patchifier,\n \"text_encoder\": text_encoder,\n \"tokenizer\": tokenizer,\n \"scheduler\": scheduler,\n \"vae\": vae,\n \"prompt_enhancer_image_caption_model\": prompt_enhancer_image_caption_model,\n \"prompt_enhancer_image_caption_processor\": prompt_enhancer_image_caption_processor,\n \"prompt_enhancer_llm_model\": prompt_enhancer_llm_model,\n \"prompt_enhancer_llm_tokenizer\": prompt_enhancer_llm_tokenizer,\n \"allowed_inference_steps\": allowed_inference_steps,\n }\n\n pipeline = LTXVideoPipeline(**submodel_dict)\n pipeline = pipeline.to(device)\n return pipeline\n\n\ndef create_latent_upsampler(latent_upsampler_model_path: str, device: str):\n latent_upsampler = LatentUpsampler.from_pretrained(latent_upsampler_model_path)\n latent_upsampler.to(device)\n latent_upsampler.eval()\n return latent_upsampler\n\n\ndef load_pipeline_config(pipeline_config: str):\n current_file = Path(__file__)\n\n path = None\n if os.path.isfile(current_file.parent / pipeline_config):\n path = current_file.parent / pipeline_config\n elif os.path.isfile(pipeline_config):\n path = pipeline_config\n else:\n raise ValueError(f\"Pipeline config file {pipeline_config} does not exist\")\n\n with open(path, \"r\") as f:\n return yaml.safe_load(f)\n\n\n@dataclass\nclass InferenceConfig:\n prompt: str = field(metadata={\"help\": \"Prompt for the generation\"})\n\n output_path: str = field(\n default_factory=lambda: Path(\n f\"outputs/{datetime.today().strftime('%Y-%m-%d')}\"\n ),\n metadata={\"help\": \"Path to the folder to save the output video\"},\n )\n\n # Pipeline settings\n pipeline_config: str = field(\n default=\"configs/ltxv-13b-0.9.7-dev.yaml\",\n metadata={\"help\": \"Path to the pipeline config file\"},\n )\n seed: int = field(\n default=171198, metadata={\"help\": \"Random seed for the inference\"}\n )\n height: int = field(\n default=704, metadata={\"help\": \"Height of the output video frames\"}\n )\n width: int = field(\n default=1216, metadata={\"help\": \"Width of the output video frames\"}\n )\n num_frames: int = field(\n default=121,\n metadata={\"help\": \"Number of frames to generate in the output video\"},\n )\n frame_rate: int = field(\n default=30, metadata={\"help\": \"Frame rate for the output video\"}\n )\n offload_to_cpu: bool = field(\n default=False, metadata={\"help\": \"Offloading unnecessary computations to CPU.\"}\n )\n negative_prompt: str = field(\n default=\"worst quality, inconsistent motion, blurry, jittery, distorted\",\n metadata={\"help\": \"Negative prompt for undesired features\"},\n )\n\n # Video-to-video arguments\n input_media_path: Optional[str] = field(\n default=None,\n metadata={\n \"help\": \"Path to the input video (or image) to be modified using the video-to-video pipeline\"\n },\n )\n\n # Conditioning\n image_cond_noise_scale: float = field(\n default=0.15,\n metadata={\"help\": \"Amount of noise to add to the conditioned image\"},\n )\n conditioning_media_paths: Optional[List[str]] = field(\n default=None,\n metadata={\n \"help\": \"List of paths to conditioning media (images or videos). Each path will be used as a conditioning item.\"\n },\n )\n conditioning_strengths: Optional[List[float]] = field(\n default=None,\n metadata={\n \"help\": \"List of conditioning strengths (between 0 and 1) for each conditioning item. Must match the number of conditioning items.\"\n },\n )\n conditioning_start_frames: Optional[List[int]] = field(\n default=None,\n metadata={\n \"help\": \"List of frame indices where each conditioning item should be applied. Must match the number of conditioning items.\"\n },\n )\n\n\ndef infer(config: InferenceConfig):\n pipeline_config = load_pipeline_config(config.pipeline_config)\n\n ltxv_model_name_or_path = pipeline_config[\"checkpoint_path\"]\n if not os.path.isfile(ltxv_model_name_or_path):\n ltxv_model_path = hf_hub_download(\n repo_id=\"Lightricks/LTX-Video\",\n filename=ltxv_model_name_or_path,\n repo_type=\"model\",\n )\n else:\n ltxv_model_path = ltxv_model_name_or_path\n\n spatial_upscaler_model_name_or_path = pipeline_config.get(\n \"spatial_upscaler_model_path\"\n )\n if spatial_upscaler_model_name_or_path and not os.path.isfile(\n spatial_upscaler_model_name_or_path\n ):\n spatial_upscaler_model_path = hf_hub_download(\n repo_id=\"Lightricks/LTX-Video\",\n filename=spatial_upscaler_model_name_or_path,\n repo_type=\"model\",\n )\n else:\n spatial_upscaler_model_path = spatial_upscaler_model_name_or_path\n\n conditioning_media_paths = config.conditioning_media_paths\n conditioning_strengths = config.conditioning_strengths\n conditioning_start_frames = config.conditioning_start_frames\n\n # Validate conditioning arguments\n if conditioning_media_paths:\n # Use default strengths of 1.0\n if not conditioning_strengths:\n conditioning_strengths = [1.0] * len(conditioning_media_paths)\n if not conditioning_start_frames:\n raise ValueError(\n \"If `conditioning_media_paths` is provided, \"\n \"`conditioning_start_frames` must also be provided\"\n )\n if len(conditioning_media_paths) != len(conditioning_strengths) or len(\n conditioning_media_paths\n ) != len(conditioning_start_frames):\n raise ValueError(\n \"`conditioning_media_paths`, `conditioning_strengths`, \"\n \"and `conditioning_start_frames` must have the same length\"\n )\n if any(s < 0 or s > 1 for s in conditioning_strengths):\n raise ValueError(\"All conditioning strengths must be between 0 and 1\")\n if any(f < 0 or f >= config.num_frames for f in conditioning_start_frames):\n raise ValueError(\n f\"All conditioning start frames must be between 0 and {config.num_frames-1}\"\n )\n\n seed_everething(config.seed)\n if config.offload_to_cpu and not torch.cuda.is_available():\n logger.warning(\n \"offload_to_cpu is set to True, but offloading will not occur since the model is already running on CPU.\"\n )\n offload_to_cpu = False\n else:\n offload_to_cpu = config.offload_to_cpu and get_total_gpu_memory() < 30\n\n output_dir = (\n Path(config.output_path)\n if config.output_path\n else Path(f\"outputs/{datetime.today().strftime('%Y-%m-%d')}\")\n )\n output_dir.mkdir(parents=True, exist_ok=True)\n\n # Adjust dimensions to be divisible by 32 and num_frames to be (N * 8 + 1)\n height_padded = ((config.height - 1) // 32 + 1) * 32\n width_padded = ((config.width - 1) // 32 + 1) * 32\n num_frames_padded = ((config.num_frames - 2) // 8 + 1) * 8 + 1\n\n padding = calculate_padding(\n config.height, config.width, height_padded, width_padded\n )\n\n logger.warning(\n f\"Padded dimensions: {height_padded}x{width_padded}x{num_frames_padded}\"\n )\n\n device = get_device()\n\n prompt_enhancement_words_threshold = pipeline_config[\n \"prompt_enhancement_words_threshold\"\n ]\n\n prompt_word_count = len(config.prompt.split())\n enhance_prompt = (\n prompt_enhancement_words_threshold > 0\n and prompt_word_count < prompt_enhancement_words_threshold\n )\n\n if prompt_enhancement_words_threshold > 0 and not enhance_prompt:\n logger.info(\n f\"Prompt has {prompt_word_count} words, which exceeds the threshold of {prompt_enhancement_words_threshold}. Prompt enhancement disabled.\"\n )\n\n precision = pipeline_config[\"precision\"]\n text_encoder_model_name_or_path = pipeline_config[\"text_encoder_model_name_or_path\"]\n sampler = pipeline_config.get(\"sampler\", None)\n prompt_enhancer_image_caption_model_name_or_path = pipeline_config[\n \"prompt_enhancer_image_caption_model_name_or_path\"\n ]\n prompt_enhancer_llm_model_name_or_path = pipeline_config[\n \"prompt_enhancer_llm_model_name_or_path\"\n ]\n\n pipeline = create_ltx_video_pipeline(\n ckpt_path=ltxv_model_path,\n precision=precision,\n text_encoder_model_name_or_path=text_encoder_model_name_or_path,\n sampler=sampler,\n device=device,\n enhance_prompt=enhance_prompt,\n prompt_enhancer_image_caption_model_name_or_path=prompt_enhancer_image_caption_model_name_or_path,\n prompt_enhancer_llm_model_name_or_path=prompt_enhancer_llm_model_name_or_path,\n )\n\n if pipeline_config.get(\"pipeline_type\", None) == \"multi-scale\":\n if not spatial_upscaler_model_path:\n raise ValueError(\n \"spatial upscaler model path is missing from pipeline config file and is required for multi-scale rendering\"\n )\n latent_upsampler = create_latent_upsampler(\n spatial_upscaler_model_path, pipeline.device\n )\n pipeline = LTXMultiScalePipeline(pipeline, latent_upsampler=latent_upsampler)\n\n media_item = None\n if config.input_media_path:\n media_item = load_media_file(\n media_path=config.input_media_path,\n height=config.height,\n width=config.width,\n max_frames=num_frames_padded,\n padding=padding,\n )\n\n conditioning_items = (\n prepare_conditioning(\n conditioning_media_paths=conditioning_media_paths,\n conditioning_strengths=conditioning_strengths,\n conditioning_start_frames=conditioning_start_frames,\n height=config.height,\n width=config.width,\n num_frames=config.num_frames,\n padding=padding,\n pipeline=pipeline,\n )\n if conditioning_media_paths\n else None\n )\n\n stg_mode = pipeline_config.get(\"stg_mode\", \"attention_values\")\n del pipeline_config[\"stg_mode\"]\n if stg_mode.lower() == \"stg_av\" or stg_mode.lower() == \"attention_values\":\n skip_layer_strategy = SkipLayerStrategy.AttentionValues\n elif stg_mode.lower() == \"stg_as\" or stg_mode.lower() == \"attention_skip\":\n skip_layer_strategy = SkipLayerStrategy.AttentionSkip\n elif stg_mode.lower() == \"stg_r\" or stg_mode.lower() == \"residual\":\n skip_layer_strategy = SkipLayerStrategy.Residual\n elif stg_mode.lower() == \"stg_t\" or stg_mode.lower() == \"transformer_block\":\n skip_layer_strategy = SkipLayerStrategy.TransformerBlock\n else:\n raise ValueError(f\"Invalid spatiotemporal guidance mode: {stg_mode}\")\n\n # Prepare input for the pipeline\n sample = {\n \"prompt\": config.prompt,\n \"prompt_attention_mask\": None,\n \"negative_prompt\": config.negative_prompt,\n \"negative_prompt_attention_mask\": None,\n }\n\n generator = torch.Generator(device=device).manual_seed(config.seed)\n\n images = pipeline(\n **pipeline_config,\n skip_layer_strategy=skip_layer_strategy,\n generator=generator,\n output_type=\"pt\",\n callback_on_step_end=None,\n height=height_padded,\n width=width_padded,\n num_frames=num_frames_padded,\n frame_rate=config.frame_rate,\n **sample,\n media_items=media_item,\n conditioning_items=conditioning_items,\n is_video=True,\n vae_per_channel_normalize=True,\n image_cond_noise_scale=config.image_cond_noise_scale,\n mixed_precision=(precision == \"mixed_precision\"),\n offload_to_cpu=offload_to_cpu,\n device=device,\n enhance_prompt=enhance_prompt,\n ).images\n\n # Crop the padded images to the desired resolution and number of frames\n (pad_left, pad_right, pad_top, pad_bottom) = padding\n pad_bottom = -pad_bottom\n pad_right = -pad_right\n if pad_bottom == 0:\n pad_bottom = images.shape[3]\n if pad_right == 0:\n pad_right = images.shape[4]\n images = images[:, :, : config.num_frames, pad_top:pad_bottom, pad_left:pad_right]\n\n for i in range(images.shape[0]):\n # Gathering from B, C, F, H, W to C, F, H, W and then permuting to F, H, W, C\n video_np = images[i].permute(1, 2, 3, 0).cpu().float().numpy()\n # Unnormalizing images to [0, 255] range\n video_np = (video_np * 255).astype(np.uint8)\n fps = config.frame_rate\n height, width = video_np.shape[1:3]\n # In case a single image is generated\n if video_np.shape[0] == 1:\n output_filename = get_unique_filename(\n f\"image_output_{i}\",\n \".png\",\n prompt=config.prompt,\n seed=config.seed,\n resolution=(height, width, config.num_frames),\n dir=output_dir,\n )\n imageio.imwrite(output_filename, video_np[0])\n else:\n output_filename = get_unique_filename(\n f\"video_output_{i}\",\n \".mp4\",\n prompt=config.prompt,\n seed=config.seed,\n resolution=(height, width, config.num_frames),\n dir=output_dir,\n )\n\n # Write video\n with imageio.get_writer(output_filename, fps=fps) as video:\n for frame in video_np:\n video.append_data(frame)\n\n logger.warning(f\"Output saved to {output_filename}\")\n\n\ndef prepare_conditioning(\n conditioning_media_paths: List[str],\n conditioning_strengths: List[float],\n conditioning_start_frames: List[int],\n height: int,\n width: int,\n num_frames: int,\n padding: tuple[int, int, int, int],\n pipeline: LTXVideoPipeline,\n) -> Optional[List[ConditioningItem]]:\n \"\"\"Prepare conditioning items based on input media paths and their parameters.\n\n Args:\n conditioning_media_paths: List of paths to conditioning media (images or videos)\n conditioning_strengths: List of conditioning strengths for each media item\n conditioning_start_frames: List of frame indices where each item should be applied\n height: Height of the output frames\n width: Width of the output frames\n num_frames: Number of frames in the output video\n padding: Padding to apply to the frames\n pipeline: LTXVideoPipeline object used for condition video trimming\n\n Returns:\n A list of ConditioningItem objects.\n \"\"\"\n conditioning_items = []\n for path, strength, start_frame in zip(\n conditioning_media_paths, conditioning_strengths, conditioning_start_frames\n ):\n num_input_frames = orig_num_input_frames = get_media_num_frames(path)\n if hasattr(pipeline, \"trim_conditioning_sequence\") and callable(\n getattr(pipeline, \"trim_conditioning_sequence\")\n ):\n num_input_frames = pipeline.trim_conditioning_sequence(\n start_frame, orig_num_input_frames, num_frames\n )\n if num_input_frames < orig_num_input_frames:\n logger.warning(\n f\"Trimming conditioning video {path} from {orig_num_input_frames} to {num_input_frames} frames.\"\n )\n\n media_tensor = load_media_file(\n media_path=path,\n height=height,\n width=width,\n max_frames=num_input_frames,\n padding=padding,\n just_crop=True,\n )\n conditioning_items.append(ConditioningItem(media_tensor, start_frame, strength))\n return conditioning_items\n\n\ndef get_media_num_frames(media_path: str) -> int:\n is_video = any(\n media_path.lower().endswith(ext) for ext in [\".mp4\", \".avi\", \".mov\", \".mkv\"]\n )\n num_frames = 1\n if is_video:\n reader = imageio.get_reader(media_path)\n num_frames = reader.count_frames()\n reader.close()\n return num_frames\n\n\ndef load_media_file(\n media_path: str,\n height: int,\n width: int,\n max_frames: int,\n padding: tuple[int, int, int, int],\n just_crop: bool = False,\n) -> torch.Tensor:\n is_video = any(\n media_path.lower().endswith(ext) for ext in [\".mp4\", \".avi\", \".mov\", \".mkv\"]\n )\n if is_video:\n reader = imageio.get_reader(media_path)\n num_input_frames = min(reader.count_frames(), max_frames)\n\n # Read and preprocess the relevant frames from the video file.\n frames = []\n for i in range(num_input_frames):\n frame = Image.fromarray(reader.get_data(i))\n frame_tensor = load_image_to_tensor_with_resize_and_crop(\n frame, height, width, just_crop=just_crop\n )\n frame_tensor = torch.nn.functional.pad(frame_tensor, padding)\n frames.append(frame_tensor)\n reader.close()\n\n # Stack frames along the temporal dimension\n media_tensor = torch.cat(frames, dim=2)\n else: # Input image\n media_tensor = load_image_to_tensor_with_resize_and_crop(\n media_path, height, width, just_crop=just_crop\n )\n media_tensor = torch.nn.functional.pad(media_tensor, padding)\n return media_tensor\n"
},
{
"path": "ltx_video/models/__init__.py",
"content": ""
},
{
"path": "ltx_video/models/autoencoders/__init__.py",
"content": ""
},
{
"path": "ltx_video/models/autoencoders/causal_conv3d.py",
"content": "from typing import Tuple, Union\n\nimport torch\nimport torch.nn as nn\n\n\nclass CausalConv3d(nn.Module):\n def __init__(\n self,\n in_channels,\n out_channels,\n kernel_size: int = 3,\n stride: Union[int, Tuple[int]] = 1,\n dilation: int = 1,\n groups: int = 1,\n spatial_padding_mode: str = \"zeros\",\n **kwargs,\n ):\n super().__init__()\n\n self.in_channels = in_channels\n self.out_channels = out_channels\n\n kernel_size = (kernel_size, kernel_size, kernel_size)\n self.time_kernel_size = kernel_size[0]\n\n dilation = (dilation, 1, 1)\n\n height_pad = kernel_size[1] // 2\n width_pad = kernel_size[2] // 2\n padding = (0, height_pad, width_pad)\n\n self.conv = nn.Conv3d(\n in_channels,\n out_channels,\n kernel_size,\n stride=stride,\n dilation=dilation,\n padding=padding,\n padding_mode=spatial_padding_mode,\n groups=groups,\n )\n\n def forward(self, x, causal: bool = True):\n if causal:\n first_frame_pad = x[:, :, :1, :, :].repeat(\n (1, 1, self.time_kernel_size - 1, 1, 1)\n )\n x = torch.concatenate((first_frame_pad, x), dim=2)\n else:\n first_frame_pad = x[:, :, :1, :, :].repeat(\n (1, 1, (self.time_kernel_size - 1) // 2, 1, 1)\n )\n last_frame_pad = x[:, :, -1:, :, :].repeat(\n (1, 1, (self.time_kernel_size - 1) // 2, 1, 1)\n )\n x = torch.concatenate((first_frame_pad, x, last_frame_pad), dim=2)\n x = self.conv(x)\n return x\n\n @property\n def weight(self):\n return self.conv.weight\n"
},
{
"path": "ltx_video/models/autoencoders/causal_video_autoencoder.py",
"content": "import json\nimport os\nfrom functools import partial\nfrom types import SimpleNamespace\nfrom typing import Any, Mapping, Optional, Tuple, Union, List\nfrom pathlib import Path\n\nimport torch\nimport numpy as np\nfrom einops import rearrange\nfrom torch import nn\nfrom diffusers.utils import logging\nimport torch.nn.functional as F\nfrom diffusers.models.embeddings import PixArtAlphaCombinedTimestepSizeEmbeddings\nfrom safetensors import safe_open\n\n\nfrom ltx_video.models.autoencoders.conv_nd_factory import make_conv_nd, make_linear_nd\nfrom ltx_video.models.autoencoders.pixel_norm import PixelNorm\nfrom ltx_video.models.autoencoders.pixel_shuffle import PixelShuffleND\nfrom ltx_video.models.autoencoders.vae import AutoencoderKLWrapper\nfrom ltx_video.models.transformers.attention import Attention\nfrom ltx_video.utils.diffusers_config_mapping import (\n diffusers_and_ours_config_mapping,\n make_hashable_key,\n VAE_KEYS_RENAME_DICT,\n)\n\nPER_CHANNEL_STATISTICS_PREFIX = \"per_channel_statistics.\"\nlogger = logging.get_logger(__name__) # pylint: disable=invalid-name\n\n\nclass CausalVideoAutoencoder(AutoencoderKLWrapper):\n @classmethod\n def from_pretrained(\n cls,\n pretrained_model_name_or_path: Optional[Union[str, os.PathLike]],\n *args,\n **kwargs,\n ):\n pretrained_model_name_or_path = Path(pretrained_model_name_or_path)\n if (\n pretrained_model_name_or_path.is_dir()\n and (pretrained_model_name_or_path / \"autoencoder.pth\").exists()\n ):\n config_local_path = pretrained_model_name_or_path / \"config.json\"\n config = cls.load_config(config_local_path, **kwargs)\n\n model_local_path = pretrained_model_name_or_path / \"autoencoder.pth\"\n state_dict = torch.load(model_local_path, map_location=torch.device(\"cpu\"))\n\n statistics_local_path = (\n pretrained_model_name_or_path / \"per_channel_statistics.json\"\n )\n if statistics_local_path.exists():\n with open(statistics_local_path, \"r\") as file:\n data = json.load(file)\n transposed_data = list(zip(*data[\"data\"]))\n data_dict = {\n col: torch.tensor(vals)\n for col, vals in zip(data[\"columns\"], transposed_data)\n }\n std_of_means = data_dict[\"std-of-means\"]\n mean_of_means = data_dict.get(\n \"mean-of-means\", torch.zeros_like(data_dict[\"std-of-means\"])\n )\n state_dict[f\"{PER_CHANNEL_STATISTICS_PREFIX}std-of-means\"] = (\n std_of_means\n )\n state_dict[f\"{PER_CHANNEL_STATISTICS_PREFIX}mean-of-means\"] = (\n mean_of_means\n )\n\n elif pretrained_model_name_or_path.is_dir():\n config_path = pretrained_model_name_or_path / \"vae\" / \"config.json\"\n with open(config_path, \"r\") as f:\n config = make_hashable_key(json.load(f))\n\n assert config in diffusers_and_ours_config_mapping, (\n \"Provided diffusers checkpoint config for VAE is not suppported. \"\n \"We only support diffusers configs found in Lightricks/LTX-Video.\"\n )\n\n config = diffusers_and_ours_config_mapping[config]\n\n state_dict_path = (\n pretrained_model_name_or_path\n / \"vae\"\n / \"diffusion_pytorch_model.safetensors\"\n )\n\n state_dict = {}\n with safe_open(state_dict_path, framework=\"pt\", device=\"cpu\") as f:\n for k in f.keys():\n state_dict[k] = f.get_tensor(k)\n for key in list(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\n state_dict[new_key] = state_dict.pop(key)\n\n elif pretrained_model_name_or_path.is_file() and str(\n pretrained_model_name_or_path\n ).endswith(\".safetensors\"):\n state_dict = {}\n with safe_open(\n pretrained_model_name_or_path, framework=\"pt\", device=\"cpu\"\n ) as f:\n metadata = f.metadata()\n for k in f.keys():\n state_dict[k] = f.get_tensor(k)\n configs = json.loads(metadata[\"config\"])\n config = configs[\"vae\"]\n\n video_vae = cls.from_config(config)\n if \"torch_dtype\" in kwargs:\n video_vae.to(kwargs[\"torch_dtype\"])\n video_vae.load_state_dict(state_dict)\n return video_vae\n\n @staticmethod\n def from_config(config):\n assert (\n config[\"_class_name\"] == \"CausalVideoAutoencoder\"\n ), \"config must have _class_name=CausalVideoAutoencoder\"\n if isinstance(config[\"dims\"], list):\n config[\"dims\"] = tuple(config[\"dims\"])\n\n assert config[\"dims\"] in [2, 3, (2, 1)], \"dims must be 2, 3 or (2, 1)\"\n\n double_z = config.get(\"double_z\", True)\n latent_log_var = config.get(\n \"latent_log_var\", \"per_channel\" if double_z else \"none\"\n )\n use_quant_conv = config.get(\"use_quant_conv\", True)\n normalize_latent_channels = config.get(\"normalize_latent_channels\", False)\n\n if use_quant_conv and latent_log_var in [\"uniform\", \"constant\"]:\n raise ValueError(\n f\"latent_log_var={latent_log_var} requires use_quant_conv=False\"\n )\n\n encoder = Encoder(\n dims=config[\"dims\"],\n in_channels=config.get(\"in_channels\", 3),\n out_channels=config[\"latent_channels\"],\n blocks=config.get(\"encoder_blocks\", config.get(\"blocks\")),\n patch_size=config.get(\"patch_size\", 1),\n latent_log_var=latent_log_var,\n norm_layer=config.get(\"norm_layer\", \"group_norm\"),\n base_channels=config.get(\"encoder_base_channels\", 128),\n spatial_padding_mode=config.get(\"spatial_padding_mode\", \"zeros\"),\n )\n\n decoder = Decoder(\n dims=config[\"dims\"],\n in_channels=config[\"latent_channels\"],\n out_channels=config.get(\"out_channels\", 3),\n blocks=config.get(\"decoder_blocks\", config.get(\"blocks\")),\n patch_size=config.get(\"patch_size\", 1),\n norm_layer=config.get(\"norm_layer\", \"group_norm\"),\n causal=config.get(\"causal_decoder\", False),\n timestep_conditioning=config.get(\"timestep_conditioning\", False),\n base_channels=config.get(\"decoder_base_channels\", 128),\n spatial_padding_mode=config.get(\"spatial_padding_mode\", \"zeros\"),\n )\n\n dims = config[\"dims\"]\n return CausalVideoAutoencoder(\n encoder=encoder,\n decoder=decoder,\n latent_channels=config[\"latent_channels\"],\n dims=dims,\n use_quant_conv=use_quant_conv,\n normalize_latent_channels=normalize_latent_channels,\n )\n\n @property\n def config(self):\n return SimpleNamespace(\n _class_name=\"CausalVideoAutoencoder\",\n dims=self.dims,\n in_channels=self.encoder.conv_in.in_channels // self.encoder.patch_size**2,\n out_channels=self.decoder.conv_out.out_channels\n // self.decoder.patch_size**2,\n latent_channels=self.decoder.conv_in.in_channels,\n encoder_blocks=self.encoder.blocks_desc,\n decoder_blocks=self.decoder.blocks_desc,\n scaling_factor=1.0,\n norm_layer=self.encoder.norm_layer,\n patch_size=self.encoder.patch_size,\n latent_log_var=self.encoder.latent_log_var,\n use_quant_conv=self.use_quant_conv,\n causal_decoder=self.decoder.causal,\n timestep_conditioning=self.decoder.timestep_conditioning,\n normalize_latent_channels=self.normalize_latent_channels,\n )\n\n @property\n def is_video_supported(self):\n \"\"\"\n Check if the model supports video inputs of shape (B, C, F, H, W). Otherwise, the model only supports 2D images.\n \"\"\"\n return self.dims != 2\n\n @property\n def spatial_downscale_factor(self):\n return (\n 2\n ** len(\n [\n block\n for block in self.encoder.blocks_desc\n if block[0]\n in [\n \"compress_space\",\n \"compress_all\",\n \"compress_all_res\",\n \"compress_space_res\",\n ]\n ]\n )\n * self.encoder.patch_size\n )\n\n @property\n def temporal_downscale_factor(self):\n return 2 ** len(\n [\n block\n for block in self.encoder.blocks_desc\n if block[0]\n in [\n \"compress_time\",\n \"compress_all\",\n \"compress_all_res\",\n \"compress_time_res\",\n ]\n ]\n )\n\n def to_json_string(self) -> str:\n import json\n\n return json.dumps(self.config.__dict__)\n\n def load_state_dict(self, state_dict: Mapping[str, Any], strict: bool = True):\n if any([key.startswith(\"vae.\") for key in state_dict.keys()]):\n state_dict = {\n key.replace(\"vae.\", \"\"): value\n for key, value in state_dict.items()\n if key.startswith(\"vae.\")\n }\n ckpt_state_dict = {\n key: value\n for key, value in state_dict.items()\n if not key.startswith(PER_CHANNEL_STATISTICS_PREFIX)\n }\n\n model_keys = set(name for name, _ in self.named_modules())\n\n key_mapping = {\n \".resnets.\": \".res_blocks.\",\n \"downsamplers.0\": \"downsample\",\n \"upsamplers.0\": \"upsample\",\n }\n converted_state_dict = {}\n for key, value in ckpt_state_dict.items():\n for k, v in key_mapping.items():\n key = key.replace(k, v)\n\n key_prefix = \".\".join(key.split(\".\")[:-1])\n if \"norm\" in key and key_prefix not in model_keys:\n logger.info(\n f\"Removing key {key} from state_dict as it is not present in the model\"\n )\n continue\n\n converted_state_dict[key] = value\n\n super().load_state_dict(converted_state_dict, strict=strict)\n\n data_dict = {\n key.removeprefix(PER_CHANNEL_STATISTICS_PREFIX): value\n for key, value in state_dict.items()\n if key.startswith(PER_CHANNEL_STATISTICS_PREFIX)\n }\n if len(data_dict) > 0:\n self.register_buffer(\"std_of_means\", data_dict[\"std-of-means\"])\n self.register_buffer(\n \"mean_of_means\",\n data_dict.get(\n \"mean-of-means\", torch.zeros_like(data_dict[\"std-of-means\"])\n ),\n )\n\n def last_layer(self):\n if hasattr(self.decoder, \"conv_out\"):\n if isinstance(self.decoder.conv_out, nn.Sequential):\n last_layer = self.decoder.conv_out[-1]\n else:\n last_layer = self.decoder.conv_out\n else:\n last_layer = self.decoder.layers[-1]\n return last_layer\n\n def set_use_tpu_flash_attention(self):\n for block in self.decoder.up_blocks:\n if isinstance(block, UNetMidBlock3D) and block.attention_blocks:\n for attention_block in block.attention_blocks:\n attention_block.set_use_tpu_flash_attention()\n\n\nclass Encoder(nn.Module):\n r\"\"\"\n The `Encoder` layer of a variational autoencoder that encodes its input into a latent representation.\n\n Args:\n dims (`int` or `Tuple[int, int]`, *optional*, defaults to 3):\n The number of dimensions to use in convolutions.\n in_channels (`int`, *optional*, defaults to 3):\n The number of input channels.\n out_channels (`int`, *optional*, defaults to 3):\n The number of output channels.\n blocks (`List[Tuple[str, int]]`, *optional*, defaults to `[(\"res_x\", 1)]`):\n The blocks to use. Each block is a tuple of the block name and the number of layers.\n base_channels (`int`, *optional*, defaults to 128):\n The number of output channels for the first convolutional layer.\n norm_num_groups (`int`, *optional*, defaults to 32):\n The number of groups for normalization.\n patch_size (`int`, *optional*, defaults to 1):\n The patch size to use. Should be a power of 2.\n norm_layer (`str`, *optional*, defaults to `group_norm`):\n The normalization layer to use. Can be either `group_norm` or `pixel_norm`.\n latent_log_var (`str`, *optional*, defaults to `per_channel`):\n The number of channels for the log variance. Can be either `per_channel`, `uniform`, `constant` or `none`.\n \"\"\"\n\n def __init__(\n self,\n dims: Union[int, Tuple[int, int]] = 3,\n in_channels: int = 3,\n out_channels: int = 3,\n blocks: List[Tuple[str, int | dict]] = [(\"res_x\", 1)],\n base_channels: int = 128,\n norm_num_groups: int = 32,\n patch_size: Union[int, Tuple[int]] = 1,\n norm_layer: str = \"group_norm\", # group_norm, pixel_norm\n latent_log_var: str = \"per_channel\",\n spatial_padding_mode: str = \"zeros\",\n ):\n super().__init__()\n self.patch_size = patch_size\n self.norm_layer = norm_layer\n self.latent_channels = out_channels\n self.latent_log_var = latent_log_var\n self.blocks_desc = blocks\n\n in_channels = in_channels * patch_size**2\n output_channel = base_channels\n\n self.conv_in = make_conv_nd(\n dims=dims,\n in_channels=in_channels,\n out_channels=output_channel,\n kernel_size=3,\n stride=1,\n padding=1,\n causal=True,\n spatial_padding_mode=spatial_padding_mode,\n )\n\n self.down_blocks = nn.ModuleList([])\n\n for block_name, block_params in blocks:\n input_channel = output_channel\n if isinstance(block_params, int):\n block_params = {\"num_layers\": block_params}\n\n if block_name == \"res_x\":\n block = UNetMidBlock3D(\n dims=dims,\n in_channels=input_channel,\n num_layers=block_params[\"num_layers\"],\n resnet_eps=1e-6,\n resnet_groups=norm_num_groups,\n norm_layer=norm_layer,\n spatial_padding_mode=spatial_padding_mode,\n )\n elif block_name == \"res_x_y\":\n output_channel = block_params.get(\"multiplier\", 2) * output_channel\n block = ResnetBlock3D(\n dims=dims,\n in_channels=input_channel,\n out_channels=output_channel,\n eps=1e-6,\n groups=norm_num_groups,\n norm_layer=norm_layer,\n spatial_padding_mode=spatial_padding_mode,\n )\n elif block_name == \"compress_time\":\n block = make_conv_nd(\n dims=dims,\n in_channels=input_channel,\n out_channels=output_channel,\n kernel_size=3,\n stride=(2, 1, 1),\n causal=True,\n spatial_padding_mode=spatial_padding_mode,\n )\n elif block_name == \"compress_space\":\n block = make_conv_nd(\n dims=dims,\n in_channels=input_channel,\n out_channels=output_channel,\n kernel_size=3,\n stride=(1, 2, 2),\n causal=True,\n spatial_padding_mode=spatial_padding_mode,\n )\n elif block_name == \"compress_all\":\n block = make_conv_nd(\n dims=dims,\n in_channels=input_channel,\n out_channels=output_channel,\n kernel_size=3,\n stride=(2, 2, 2),\n causal=True,\n spatial_padding_mode=spatial_padding_mode,\n )\n elif block_name == \"compress_all_x_y\":\n output_channel = block_params.get(\"multiplier\", 2) * output_channel\n block = make_conv_nd(\n dims=dims,\n in_channels=input_channel,\n out_channels=output_channel,\n kernel_size=3,\n stride=(2, 2, 2),\n causal=True,\n spatial_padding_mode=spatial_padding_mode,\n )\n elif block_name == \"compress_all_res\":\n output_channel = block_params.get(\"multiplier\", 2) * output_channel\n block = SpaceToDepthDownsample(\n dims=dims,\n in_channels=input_channel,\n out_channels=output_channel,\n stride=(2, 2, 2),\n spatial_padding_mode=spatial_padding_mode,\n )\n elif block_name == \"compress_space_res\":\n output_channel = block_params.get(\"multiplier\", 2) * output_channel\n block = SpaceToDepthDownsample(\n dims=dims,\n in_channels=input_channel,\n out_channels=output_channel,\n stride=(1, 2, 2),\n spatial_padding_mode=spatial_padding_mode,\n )\n elif block_name == \"compress_time_res\":\n output_channel = block_params.get(\"multiplier\", 2) * output_channel\n block = SpaceToDepthDownsample(\n dims=dims,\n in_channels=input_channel,\n out_channels=output_channel,\n stride=(2, 1, 1),\n spatial_padding_mode=spatial_padding_mode,\n )\n else:\n raise ValueError(f\"unknown block: {block_name}\")\n\n self.down_blocks.append(block)\n\n # out\n if norm_layer == \"group_norm\":\n self.conv_norm_out = nn.GroupNorm(\n num_channels=output_channel, num_groups=norm_num_groups, eps=1e-6\n )\n elif norm_layer == \"pixel_norm\":\n self.conv_norm_out = PixelNorm()\n elif norm_layer == \"layer_norm\":\n self.conv_norm_out = LayerNorm(output_channel, eps=1e-6)\n\n self.conv_act = nn.SiLU()\n\n conv_out_channels = out_channels\n if latent_log_var == \"per_channel\":\n conv_out_channels *= 2\n elif latent_log_var == \"uniform\":\n conv_out_channels += 1\n elif latent_log_var == \"constant\":\n conv_out_channels += 1\n elif latent_log_var != \"none\":\n raise ValueError(f\"Invalid latent_log_var: {latent_log_var}\")\n self.conv_out = make_conv_nd(\n dims,\n output_channel,\n conv_out_channels,\n 3,\n padding=1,\n causal=True,\n spatial_padding_mode=spatial_padding_mode,\n )\n\n self.gradient_checkpointing = False\n\n def forward(self, sample: torch.FloatTensor) -> torch.FloatTensor:\n r\"\"\"The forward method of the `Encoder` class.\"\"\"\n\n sample = patchify(sample, patch_size_hw=self.patch_size, patch_size_t=1)\n sample = self.conv_in(sample)\n\n checkpoint_fn = (\n partial(torch.utils.checkpoint.checkpoint, use_reentrant=False)\n if self.gradient_checkpointing and self.training\n else lambda x: x\n )\n\n for down_block in self.down_blocks:\n sample = checkpoint_fn(down_block)(sample)\n\n sample = self.conv_norm_out(sample)\n sample = self.conv_act(sample)\n sample = self.conv_out(sample)\n\n if self.latent_log_var == \"uniform\":\n last_channel = sample[:, -1:, ...]\n num_dims = sample.dim()\n\n if num_dims == 4:\n # For shape (B, C, H, W)\n repeated_last_channel = last_channel.repeat(\n 1, sample.shape[1] - 2, 1, 1\n )\n sample = torch.cat([sample, repeated_last_channel], dim=1)\n elif num_dims == 5:\n # For shape (B, C, F, H, W)\n repeated_last_channel = last_channel.repeat(\n 1, sample.shape[1] - 2, 1, 1, 1\n )\n sample = torch.cat([sample, repeated_last_channel], dim=1)\n else:\n raise ValueError(f\"Invalid input shape: {sample.shape}\")\n elif self.latent_log_var == \"constant\":\n sample = sample[:, :-1, ...]\n approx_ln_0 = (\n -30\n ) # this is the minimal clamp value in DiagonalGaussianDistribution objects\n sample = torch.cat(\n [sample, torch.ones_like(sample, device=sample.device) * approx_ln_0],\n dim=1,\n )\n\n return sample\n\n\nclass Decoder(nn.Module):\n r\"\"\"\n The `Decoder` layer of a variational autoencoder that decodes its latent representation into an output sample.\n\n Args:\n dims (`int` or `Tuple[int, int]`, *optional*, defaults to 3):\n The number of dimensions to use in convolutions.\n in_channels (`int`, *optional*, defaults to 3):\n The number of input channels.\n out_channels (`int`, *optional*, defaults to 3):\n The number of output channels.\n blocks (`List[Tuple[str, int]]`, *optional*, defaults to `[(\"res_x\", 1)]`):\n The blocks to use. Each block is a tuple of the block name and the number of layers.\n base_channels (`int`, *optional*, defaults to 128):\n The number of output channels for the first convolutional layer.\n norm_num_groups (`int`, *optional*, defaults to 32):\n The number of groups for normalization.\n patch_size (`int`, *optional*, defaults to 1):\n The patch size to use. Should be a power of 2.\n norm_layer (`str`, *optional*, defaults to `group_norm`):\n The normalization layer to use. Can be either `group_norm` or `pixel_norm`.\n causal (`bool`, *optional*, defaults to `True`):\n Whether to use causal convolutions or not.\n \"\"\"\n\n def __init__(\n self,\n dims,\n in_channels: int = 3,\n out_channels: int = 3,\n blocks: List[Tuple[str, int | dict]] = [(\"res_x\", 1)],\n base_channels: int = 128,\n layers_per_block: int = 2,\n norm_num_groups: int = 32,\n patch_size: int = 1,\n norm_layer: str = \"group_norm\",\n causal: bool = True,\n timestep_conditioning: bool = False,\n spatial_padding_mode: str = \"zeros\",\n ):\n super().__init__()\n self.patch_size = patch_size\n self.layers_per_block = layers_per_block\n out_channels = out_channels * patch_size**2\n self.causal = causal\n self.blocks_desc = blocks\n\n # Compute output channel to be product of all channel-multiplier blocks\n output_channel = base_channels\n for block_name, block_params in list(reversed(blocks)):\n block_params = block_params if isinstance(block_params, dict) else {}\n if block_name == \"res_x_y\":\n output_channel = output_channel * block_params.get(\"multiplier\", 2)\n if block_name.startswith(\"compress\"):\n output_channel = output_channel * block_params.get(\"multiplier\", 1)\n\n self.conv_in = make_conv_nd(\n dims,\n in_channels,\n output_channel,\n kernel_size=3,\n stride=1,\n padding=1,\n causal=True,\n spatial_padding_mode=spatial_padding_mode,\n )\n\n self.up_blocks = nn.ModuleList([])\n\n for block_name, block_params in list(reversed(blocks)):\n input_channel = output_channel\n if isinstance(block_params, int):\n block_params = {\"num_layers\": block_params}\n\n if block_name == \"res_x\":\n block = UNetMidBlock3D(\n dims=dims,\n in_channels=input_channel,\n num_layers=block_params[\"num_layers\"],\n resnet_eps=1e-6,\n resnet_groups=norm_num_groups,\n norm_layer=norm_layer,\n inject_noise=block_params.get(\"inject_noise\", False),\n timestep_conditioning=timestep_conditioning,\n spatial_padding_mode=spatial_padding_mode,\n )\n elif block_name == \"attn_res_x\":\n block = UNetMidBlock3D(\n dims=dims,\n in_channels=input_channel,\n num_layers=block_params[\"num_layers\"],\n resnet_groups=norm_num_groups,\n norm_layer=norm_layer,\n inject_noise=block_params.get(\"inject_noise\", False),\n timestep_conditioning=timestep_conditioning,\n attention_head_dim=block_params[\"attention_head_dim\"],\n spatial_padding_mode=spatial_padding_mode,\n )\n elif block_name == \"res_x_y\":\n output_channel = output_channel // block_params.get(\"multiplier\", 2)\n block = ResnetBlock3D(\n dims=dims,\n in_channels=input_channel,\n out_channels=output_channel,\n eps=1e-6,\n groups=norm_num_groups,\n norm_layer=norm_layer,\n inject_noise=block_params.get(\"inject_noise\", False),\n timestep_conditioning=False,\n spatial_padding_mode=spatial_padding_mode,\n )\n elif block_name == \"compress_time\":\n block = DepthToSpaceUpsample(\n dims=dims,\n in_channels=input_channel,\n stride=(2, 1, 1),\n spatial_padding_mode=spatial_padding_mode,\n )\n elif block_name == \"compress_space\":\n block = DepthToSpaceUpsample(\n dims=dims,\n in_channels=input_channel,\n stride=(1, 2, 2),\n spatial_padding_mode=spatial_padding_mode,\n )\n elif block_name == \"compress_all\":\n output_channel = output_channel // block_params.get(\"multiplier\", 1)\n block = DepthToSpaceUpsample(\n dims=dims,\n in_channels=input_channel,\n stride=(2, 2, 2),\n residual=block_params.get(\"residual\", False),\n out_channels_reduction_factor=block_params.get(\"multiplier\", 1),\n spatial_padding_mode=spatial_padding_mode,\n )\n else:\n raise ValueError(f\"unknown layer: {block_name}\")\n\n self.up_blocks.append(block)\n\n if norm_layer == \"group_norm\":\n self.conv_norm_out = nn.GroupNorm(\n num_channels=output_channel, num_groups=norm_num_groups, eps=1e-6\n )\n elif norm_layer == \"pixel_norm\":\n self.conv_norm_out = PixelNorm()\n elif norm_layer == \"layer_norm\":\n self.conv_norm_out = LayerNorm(output_channel, eps=1e-6)\n\n self.conv_act = nn.SiLU()\n self.conv_out = make_conv_nd(\n dims,\n output_channel,\n out_channels,\n 3,\n padding=1,\n causal=True,\n spatial_padding_mode=spatial_padding_mode,\n )\n\n self.gradient_checkpointing = False\n\n self.timestep_conditioning = timestep_conditioning\n\n if timestep_conditioning:\n self.timestep_scale_multiplier = nn.Parameter(\n torch.tensor(1000.0, dtype=torch.float32)\n )\n self.last_time_embedder = PixArtAlphaCombinedTimestepSizeEmbeddings(\n output_channel * 2, 0\n )\n self.last_scale_shift_table = nn.Parameter(\n torch.randn(2, output_channel) / output_channel**0.5\n )\n\n def forward(\n self,\n sample: torch.FloatTensor,\n target_shape,\n timestep: Optional[torch.Tensor] = None,\n ) -> torch.FloatTensor:\n r\"\"\"The forward method of the `Decoder` class.\"\"\"\n assert target_shape is not None, \"target_shape must be provided\"\n batch_size = sample.shape[0]\n\n sample = self.conv_in(sample, causal=self.causal)\n\n upscale_dtype = next(iter(self.up_blocks.parameters())).dtype\n\n checkpoint_fn = (\n partial(torch.utils.checkpoint.checkpoint, use_reentrant=False)\n if self.gradient_checkpointing and self.training\n else lambda x: x\n )\n\n sample = sample.to(upscale_dtype)\n\n if self.timestep_conditioning:\n assert (\n timestep is not None\n ), \"should pass timestep with timestep_conditioning=True\"\n scaled_timestep = timestep * self.timestep_scale_multiplier\n\n for up_block in self.up_blocks:\n if self.timestep_conditioning and isinstance(up_block, UNetMidBlock3D):\n sample = checkpoint_fn(up_block)(\n sample, causal=self.causal, timestep=scaled_timestep\n )\n else:\n sample = checkpoint_fn(up_block)(sample, causal=self.causal)\n\n sample = self.conv_norm_out(sample)\n\n if self.timestep_conditioning:\n embedded_timestep = self.last_time_embedder(\n timestep=scaled_timestep.flatten(),\n resolution=None,\n aspect_ratio=None,\n batch_size=sample.shape[0],\n hidden_dtype=sample.dtype,\n )\n embedded_timestep = embedded_timestep.view(\n batch_size, embedded_timestep.shape[-1], 1, 1, 1\n )\n ada_values = self.last_scale_shift_table[\n None, ..., None, None, None\n ] + embedded_timestep.reshape(\n batch_size,\n 2,\n -1,\n embedded_timestep.shape[-3],\n embedded_timestep.shape[-2],\n embedded_timestep.shape[-1],\n )\n shift, scale = ada_values.unbind(dim=1)\n sample = sample * (1 + scale) + shift\n\n sample = self.conv_act(sample)\n sample = self.conv_out(sample, causal=self.causal)\n\n sample = unpatchify(sample, patch_size_hw=self.patch_size, patch_size_t=1)\n\n return sample\n\n\nclass UNetMidBlock3D(nn.Module):\n \"\"\"\n A 3D UNet mid-block [`UNetMidBlock3D`] with multiple residual blocks.\n\n Args:\n in_channels (`int`): The number of input channels.\n dropout (`float`, *optional*, defaults to 0.0): The dropout rate.\n num_layers (`int`, *optional*, defaults to 1): The number of residual blocks.\n resnet_eps (`float`, *optional*, 1e-6 ): The epsilon value for the resnet blocks.\n resnet_groups (`int`, *optional*, defaults to 32):\n The number of groups to use in the group normalization layers of the resnet blocks.\n norm_layer (`str`, *optional*, defaults to `group_norm`):\n The normalization layer to use. Can be either `group_norm` or `pixel_norm`.\n inject_noise (`bool`, *optional*, defaults to `False`):\n Whether to inject noise into the hidden states.\n timestep_conditioning (`bool`, *optional*, defaults to `False`):\n Whether to condition the hidden states on the timestep.\n attention_head_dim (`int`, *optional*, defaults to -1):\n The dimension of the attention head. If -1, no attention is used.\n\n Returns:\n `torch.FloatTensor`: The output of the last residual block, which is a tensor of shape `(batch_size,\n in_channels, height, width)`.\n\n \"\"\"\n\n def __init__(\n self,\n dims: Union[int, Tuple[int, int]],\n in_channels: int,\n dropout: float = 0.0,\n num_layers: int = 1,\n resnet_eps: float = 1e-6,\n resnet_groups: int = 32,\n norm_layer: str = \"group_norm\",\n inject_noise: bool = False,\n timestep_conditioning: bool = False,\n attention_head_dim: int = -1,\n spatial_padding_mode: str = \"zeros\",\n ):\n super().__init__()\n resnet_groups = (\n resnet_groups if resnet_groups is not None else min(in_channels // 4, 32)\n )\n self.timestep_conditioning = timestep_conditioning\n\n if timestep_conditioning:\n self.time_embedder = PixArtAlphaCombinedTimestepSizeEmbeddings(\n in_channels * 4, 0\n )\n\n self.res_blocks = nn.ModuleList(\n [\n ResnetBlock3D(\n dims=dims,\n in_channels=in_channels,\n out_channels=in_channels,\n eps=resnet_eps,\n groups=resnet_groups,\n dropout=dropout,\n norm_layer=norm_layer,\n inject_noise=inject_noise,\n timestep_conditioning=timestep_conditioning,\n spatial_padding_mode=spatial_padding_mode,\n )\n for _ in range(num_layers)\n ]\n )\n\n self.attention_blocks = None\n\n if attention_head_dim > 0:\n if attention_head_dim > in_channels:\n raise ValueError(\n \"attention_head_dim must be less than or equal to in_channels\"\n )\n\n self.attention_blocks = nn.ModuleList(\n [\n Attention(\n query_dim=in_channels,\n heads=in_channels // attention_head_dim,\n dim_head=attention_head_dim,\n bias=True,\n out_bias=True,\n qk_norm=\"rms_norm\",\n residual_connection=True,\n )\n for _ in range(num_layers)\n ]\n )\n\n def forward(\n self,\n hidden_states: torch.FloatTensor,\n causal: bool = True,\n timestep: Optional[torch.Tensor] = None,\n ) -> torch.FloatTensor:\n timestep_embed = None\n if self.timestep_conditioning:\n assert (\n timestep is not None\n ), \"should pass timestep with timestep_conditioning=True\"\n batch_size = hidden_states.shape[0]\n timestep_embed = self.time_embedder(\n timestep=timestep.flatten(),\n resolution=None,\n aspect_ratio=None,\n batch_size=batch_size,\n hidden_dtype=hidden_states.dtype,\n )\n timestep_embed = timestep_embed.view(\n batch_size, timestep_embed.shape[-1], 1, 1, 1\n )\n\n if self.attention_blocks:\n for resnet, attention in zip(self.res_blocks, self.attention_blocks):\n hidden_states = resnet(\n hidden_states, causal=causal, timestep=timestep_embed\n )\n\n # Reshape the hidden states to be (batch_size, frames * height * width, channel)\n batch_size, channel, frames, height, width = hidden_states.shape\n hidden_states = hidden_states.view(\n batch_size, channel, frames * height * width\n ).transpose(1, 2)\n\n if attention.use_tpu_flash_attention:\n # Pad the second dimension to be divisible by block_k_major (block in flash attention)\n seq_len = hidden_states.shape[1]\n block_k_major = 512\n pad_len = (block_k_major - seq_len % block_k_major) % block_k_major\n if pad_len > 0:\n hidden_states = F.pad(\n hidden_states, (0, 0, 0, pad_len), \"constant\", 0\n )\n\n # Create a mask with ones for the original sequence length and zeros for the padded indexes\n mask = torch.ones(\n (hidden_states.shape[0], seq_len),\n device=hidden_states.device,\n dtype=hidden_states.dtype,\n )\n if pad_len > 0:\n mask = F.pad(mask, (0, pad_len), \"constant\", 0)\n\n hidden_states = attention(\n hidden_states,\n attention_mask=(\n None if not attention.use_tpu_flash_attention else mask\n ),\n )\n\n if attention.use_tpu_flash_attention:\n # Remove the padding\n if pad_len > 0:\n hidden_states = hidden_states[:, :-pad_len, :]\n\n # Reshape the hidden states back to (batch_size, channel, frames, height, width, channel)\n hidden_states = hidden_states.transpose(-1, -2).reshape(\n batch_size, channel, frames, height, width\n )\n else:\n for resnet in self.res_blocks:\n hidden_states = resnet(\n hidden_states, causal=causal, timestep=timestep_embed\n )\n\n return hidden_states\n\n\nclass SpaceToDepthDownsample(nn.Module):\n def __init__(self, dims, in_channels, out_channels, stride, spatial_padding_mode):\n super().__init__()\n self.stride = stride\n self.group_size = in_channels * np.prod(stride) // out_channels\n self.conv = make_conv_nd(\n dims=dims,\n in_channels=in_channels,\n out_channels=out_channels // np.prod(stride),\n kernel_size=3,\n stride=1,\n causal=True,\n spatial_padding_mode=spatial_padding_mode,\n )\n\n def forward(self, x, causal: bool = True):\n if self.stride[0] == 2:\n x = torch.cat(\n [x[:, :, :1, :, :], x], dim=2\n ) # duplicate first frames for padding\n\n # skip connection\n x_in = rearrange(\n x,\n \"b c (d p1) (h p2) (w p3) -> b (c p1 p2 p3) d h w\",\n p1=self.stride[0],\n p2=self.stride[1],\n p3=self.stride[2],\n )\n x_in = rearrange(x_in, \"b (c g) d h w -> b c g d h w\", g=self.group_size)\n x_in = x_in.mean(dim=2)\n\n # conv\n x = self.conv(x, causal=causal)\n x = rearrange(\n x,\n \"b c (d p1) (h p2) (w p3) -> b (c p1 p2 p3) d h w\",\n p1=self.stride[0],\n p2=self.stride[1],\n p3=self.stride[2],\n )\n\n x = x + x_in\n\n return x\n\n\nclass DepthToSpaceUpsample(nn.Module):\n def __init__(\n self,\n dims,\n in_channels,\n stride,\n residual=False,\n out_channels_reduction_factor=1,\n spatial_padding_mode=\"zeros\",\n ):\n super().__init__()\n self.stride = stride\n self.out_channels = (\n np.prod(stride) * in_channels // out_channels_reduction_factor\n )\n self.conv = make_conv_nd(\n dims=dims,\n in_channels=in_channels,\n out_channels=self.out_channels,\n kernel_size=3,\n stride=1,\n causal=True,\n spatial_padding_mode=spatial_padding_mode,\n )\n self.pixel_shuffle = PixelShuffleND(dims=dims, upscale_factors=stride)\n self.residual = residual\n self.out_channels_reduction_factor = out_channels_reduction_factor\n\n def forward(self, x, causal: bool = True):\n if self.residual:\n # Reshape and duplicate the input to match the output shape\n x_in = self.pixel_shuffle(x)\n num_repeat = np.prod(self.stride) // self.out_channels_reduction_factor\n x_in = x_in.repeat(1, num_repeat, 1, 1, 1)\n if self.stride[0] == 2:\n x_in = x_in[:, :, 1:, :, :]\n x = self.conv(x, causal=causal)\n x = self.pixel_shuffle(x)\n if self.stride[0] == 2:\n x = x[:, :, 1:, :, :]\n if self.residual:\n x = x + x_in\n return x\n\n\nclass LayerNorm(nn.Module):\n def __init__(self, dim, eps, elementwise_affine=True) -> None:\n super().__init__()\n self.norm = nn.LayerNorm(dim, eps=eps, elementwise_affine=elementwise_affine)\n\n def forward(self, x):\n x = rearrange(x, \"b c d h w -> b d h w c\")\n x = self.norm(x)\n x = rearrange(x, \"b d h w c -> b c d h w\")\n return x\n\n\nclass ResnetBlock3D(nn.Module):\n r\"\"\"\n A Resnet block.\n\n Parameters:\n in_channels (`int`): The number of channels in the input.\n out_channels (`int`, *optional*, default to be `None`):\n The number of output channels for the first conv layer. If None, same as `in_channels`.\n dropout (`float`, *optional*, defaults to `0.0`): The dropout probability to use.\n groups (`int`, *optional*, default to `32`): The number of groups to use for the first normalization layer.\n eps (`float`, *optional*, defaults to `1e-6`): The epsilon to use for the normalization.\n \"\"\"\n\n def __init__(\n self,\n dims: Union[int, Tuple[int, int]],\n in_channels: int,\n out_channels: Optional[int] = None,\n dropout: float = 0.0,\n groups: int = 32,\n eps: float = 1e-6,\n norm_layer: str = \"group_norm\",\n inject_noise: bool = False,\n timestep_conditioning: bool = False,\n spatial_padding_mode: str = \"zeros\",\n ):\n super().__init__()\n self.in_channels = in_channels\n out_channels = in_channels if out_channels is None else out_channels\n self.out_channels = out_channels\n self.inject_noise = inject_noise\n\n if norm_layer == \"group_norm\":\n self.norm1 = nn.GroupNorm(\n num_groups=groups, num_channels=in_channels, eps=eps, affine=True\n )\n elif norm_layer == \"pixel_norm\":\n self.norm1 = PixelNorm()\n elif norm_layer == \"layer_norm\":\n self.norm1 = LayerNorm(in_channels, eps=eps, elementwise_affine=True)\n\n self.non_linearity = nn.SiLU()\n\n self.conv1 = make_conv_nd(\n dims,\n in_channels,\n out_channels,\n kernel_size=3,\n stride=1,\n padding=1,\n causal=True,\n spatial_padding_mode=spatial_padding_mode,\n )\n\n if inject_noise:\n self.per_channel_scale1 = nn.Parameter(torch.zeros((in_channels, 1, 1)))\n\n if norm_layer == \"group_norm\":\n self.norm2 = nn.GroupNorm(\n num_groups=groups, num_channels=out_channels, eps=eps, affine=True\n )\n elif norm_layer == \"pixel_norm\":\n self.norm2 = PixelNorm()\n elif norm_layer == \"layer_norm\":\n self.norm2 = LayerNorm(out_channels, eps=eps, elementwise_affine=True)\n\n self.dropout = torch.nn.Dropout(dropout)\n\n self.conv2 = make_conv_nd(\n dims,\n out_channels,\n out_channels,\n kernel_size=3,\n stride=1,\n padding=1,\n causal=True,\n spatial_padding_mode=spatial_padding_mode,\n )\n\n if inject_noise:\n self.per_channel_scale2 = nn.Parameter(torch.zeros((in_channels, 1, 1)))\n\n self.conv_shortcut = (\n make_linear_nd(\n dims=dims, in_channels=in_channels, out_channels=out_channels\n )\n if in_channels != out_channels\n else nn.Identity()\n )\n\n self.norm3 = (\n LayerNorm(in_channels, eps=eps, elementwise_affine=True)\n if in_channels != out_channels\n else nn.Identity()\n )\n\n self.timestep_conditioning = timestep_conditioning\n\n if timestep_conditioning:\n self.scale_shift_table = nn.Parameter(\n torch.randn(4, in_channels) / in_channels**0.5\n )\n\n def _feed_spatial_noise(\n self, hidden_states: torch.FloatTensor, per_channel_scale: torch.FloatTensor\n ) -> torch.FloatTensor:\n spatial_shape = hidden_states.shape[-2:]\n device = hidden_states.device\n dtype = hidden_states.dtype\n\n # similar to the \"explicit noise inputs\" method in style-gan\n spatial_noise = torch.randn(spatial_shape, device=device, dtype=dtype)[None]\n scaled_noise = (spatial_noise * per_channel_scale)[None, :, None, ...]\n hidden_states = hidden_states + scaled_noise\n\n return hidden_states\n\n def forward(\n self,\n input_tensor: torch.FloatTensor,\n causal: bool = True,\n timestep: Optional[torch.Tensor] = None,\n ) -> torch.FloatTensor:\n hidden_states = input_tensor\n batch_size = hidden_states.shape[0]\n\n hidden_states = self.norm1(hidden_states)\n if self.timestep_conditioning:\n assert (\n timestep is not None\n ), \"should pass timestep with timestep_conditioning=True\"\n ada_values = self.scale_shift_table[\n None, ..., None, None, None\n ] + timestep.reshape(\n batch_size,\n 4,\n -1,\n timestep.shape[-3],\n timestep.shape[-2],\n timestep.shape[-1],\n )\n shift1, scale1, shift2, scale2 = ada_values.unbind(dim=1)\n\n hidden_states = hidden_states * (1 + scale1) + shift1\n\n hidden_states = self.non_linearity(hidden_states)\n\n hidden_states = self.conv1(hidden_states, causal=causal)\n\n if self.inject_noise:\n hidden_states = self._feed_spatial_noise(\n hidden_states, self.per_channel_scale1\n )\n\n hidden_states = self.norm2(hidden_states)\n\n if self.timestep_conditioning:\n hidden_states = hidden_states * (1 + scale2) + shift2\n\n hidden_states = self.non_linearity(hidden_states)\n\n hidden_states = self.dropout(hidden_states)\n\n hidden_states = self.conv2(hidden_states, causal=causal)\n\n if self.inject_noise:\n hidden_states = self._feed_spatial_noise(\n hidden_states, self.per_channel_scale2\n )\n\n input_tensor = self.norm3(input_tensor)\n\n batch_size = input_tensor.shape[0]\n\n input_tensor = self.conv_shortcut(input_tensor)\n\n output_tensor = input_tensor + hidden_states\n\n return output_tensor\n\n\ndef patchify(x, patch_size_hw, patch_size_t=1):\n if patch_size_hw == 1 and patch_size_t == 1:\n return x\n if x.dim() == 4:\n x = rearrange(\n x, \"b c (h q) (w r) -> b (c r q) h w\", q=patch_size_hw, r=patch_size_hw\n )\n elif x.dim() == 5:\n x = rearrange(\n x,\n \"b c (f p) (h q) (w r) -> b (c p r q) f h w\",\n p=patch_size_t,\n q=patch_size_hw,\n r=patch_size_hw,\n )\n else:\n raise ValueError(f\"Invalid input shape: {x.shape}\")\n\n return x\n\n\ndef unpatchify(x, patch_size_hw, patch_size_t=1):\n if patch_size_hw == 1 and patch_size_t == 1:\n return x\n\n if x.dim() == 4:\n x = rearrange(\n x, \"b (c r q) h w -> b c (h q) (w r)\", q=patch_size_hw, r=patch_size_hw\n )\n elif x.dim() == 5:\n x = rearrange(\n x,\n \"b (c p r q) f h w -> b c (f p) (h q) (w r)\",\n p=patch_size_t,\n q=patch_size_hw,\n r=patch_size_hw,\n )\n\n return x\n\n\ndef create_video_autoencoder_demo_config(\n latent_channels: int = 64,\n):\n encoder_blocks = [\n (\"res_x\", {\"num_layers\": 2}),\n (\"compress_space_res\", {\"multiplier\": 2}),\n (\"compress_time_res\", {\"multiplier\": 2}),\n (\"compress_all_res\", {\"multiplier\": 2}),\n (\"compress_all_res\", {\"multiplier\": 2}),\n (\"res_x\", {\"num_layers\": 1}),\n ]\n decoder_blocks = [\n (\"res_x\", {\"num_layers\": 2, \"inject_noise\": False}),\n (\"compress_all\", {\"residual\": True, \"multiplier\": 2}),\n (\"compress_all\", {\"residual\": True, \"multiplier\": 2}),\n (\"compress_all\", {\"residual\": True, \"multiplier\": 2}),\n (\"res_x\", {\"num_layers\": 2, \"inject_noise\": False}),\n ]\n return {\n \"_class_name\": \"CausalVideoAutoencoder\",\n \"dims\": 3,\n \"encoder_blocks\": encoder_blocks,\n \"decoder_blocks\": decoder_blocks,\n \"latent_channels\": latent_channels,\n \"norm_layer\": \"pixel_norm\",\n \"patch_size\": 4,\n \"latent_log_var\": \"uniform\",\n \"use_quant_conv\": False,\n \"causal_decoder\": False,\n \"timestep_conditioning\": True,\n \"spatial_padding_mode\": \"replicate\",\n }\n\n\ndef test_vae_patchify_unpatchify():\n import torch\n\n x = torch.randn(2, 3, 8, 64, 64)\n x_patched = patchify(x, patch_size_hw=4, patch_size_t=4)\n x_unpatched = unpatchify(x_patched, patch_size_hw=4, patch_size_t=4)\n assert torch.allclose(x, x_unpatched)\n\n\ndef demo_video_autoencoder_forward_backward():\n # Configuration for the VideoAutoencoder\n config = create_video_autoencoder_demo_config()\n\n # Instantiate the VideoAutoencoder with the specified configuration\n video_autoencoder = CausalVideoAutoencoder.from_config(config)\n\n print(video_autoencoder)\n video_autoencoder.eval()\n # Print the total number of parameters in the video autoencoder\n total_params = sum(p.numel() for p in video_autoencoder.parameters())\n print(f\"Total number of parameters in VideoAutoencoder: {total_params:,}\")\n\n # Create a mock input tensor simulating a batch of videos\n # Shape: (batch_size, channels, depth, height, width)\n # E.g., 4 videos, each with 3 color channels, 16 frames, and 64x64 pixels per frame\n input_videos = torch.randn(2, 3, 17, 64, 64)\n\n # Forward pass: encode and decode the input videos\n latent = video_autoencoder.encode(input_videos).latent_dist.mode()\n print(f\"input shape={input_videos.shape}\")\n print(f\"latent shape={latent.shape}\")\n\n timestep = torch.ones(input_videos.shape[0]) * 0.1\n reconstructed_videos = video_autoencoder.decode(\n latent, target_shape=input_videos.shape, timestep=timestep\n ).sample\n\n print(f\"reconstructed shape={reconstructed_videos.shape}\")\n\n # Validate that single image gets treated the same way as first frame\n input_image = input_videos[:, :, :1, :, :]\n image_latent = video_autoencoder.encode(input_image).latent_dist.mode()\n _ = video_autoencoder.decode(\n image_latent, target_shape=image_latent.shape, timestep=timestep\n ).sample\n\n first_frame_latent = latent[:, :, :1, :, :]\n\n assert torch.allclose(image_latent, first_frame_latent, atol=1e-6)\n # assert torch.allclose(reconstructed_image, reconstructed_videos[:, :, :1, :, :], atol=1e-6)\n # assert torch.allclose(image_latent, first_frame_latent, atol=1e-6)\n # assert (reconstructed_image == reconstructed_videos[:, :, :1, :, :]).all()\n\n # Calculate the loss (e.g., mean squared error)\n loss = torch.nn.functional.mse_loss(input_videos, reconstructed_videos)\n\n # Perform backward pass\n loss.backward()\n\n print(f\"Demo completed with loss: {loss.item()}\")\n\n\n# Ensure to call the demo function to execute the forward and backward pass\nif __name__ == \"__main__\":\n demo_video_autoencoder_forward_backward()\n"
},
{
"path": "ltx_video/models/autoencoders/conv_nd_factory.py",
"content": "from typing import Tuple, Union\n\nimport torch\n\nfrom ltx_video.models.autoencoders.dual_conv3d import DualConv3d\nfrom ltx_video.models.autoencoders.causal_conv3d import CausalConv3d\n\n\ndef make_conv_nd(\n dims: Union[int, Tuple[int, int]],\n in_channels: int,\n out_channels: int,\n kernel_size: int,\n stride=1,\n padding=0,\n dilation=1,\n groups=1,\n bias=True,\n causal=False,\n spatial_padding_mode=\"zeros\",\n temporal_padding_mode=\"zeros\",\n):\n if not (spatial_padding_mode == temporal_padding_mode or causal):\n raise NotImplementedError(\"spatial and temporal padding modes must be equal\")\n if dims == 2:\n return torch.nn.Conv2d(\n in_channels=in_channels,\n out_channels=out_channels,\n kernel_size=kernel_size,\n stride=stride,\n padding=padding,\n dilation=dilation,\n groups=groups,\n bias=bias,\n padding_mode=spatial_padding_mode,\n )\n elif dims == 3:\n if causal:\n return CausalConv3d(\n in_channels=in_channels,\n out_channels=out_channels,\n kernel_size=kernel_size,\n stride=stride,\n padding=padding,\n dilation=dilation,\n groups=groups,\n bias=bias,\n spatial_padding_mode=spatial_padding_mode,\n )\n return torch.nn.Conv3d(\n in_channels=in_channels,\n out_channels=out_channels,\n kernel_size=kernel_size,\n stride=stride,\n padding=padding,\n dilation=dilation,\n groups=groups,\n bias=bias,\n padding_mode=spatial_padding_mode,\n )\n elif dims == (2, 1):\n return DualConv3d(\n in_channels=in_channels,\n out_channels=out_channels,\n kernel_size=kernel_size,\n stride=stride,\n padding=padding,\n bias=bias,\n padding_mode=spatial_padding_mode,\n )\n else:\n raise ValueError(f\"unsupported dimensions: {dims}\")\n\n\ndef make_linear_nd(\n dims: int,\n in_channels: int,\n out_channels: int,\n bias=True,\n):\n if dims == 2:\n return torch.nn.Conv2d(\n in_channels=in_channels, out_channels=out_channels, kernel_size=1, bias=bias\n )\n elif dims == 3 or dims == (2, 1):\n return torch.nn.Conv3d(\n in_channels=in_channels, out_channels=out_channels, kernel_size=1, bias=bias\n )\n else:\n raise ValueError(f\"unsupported dimensions: {dims}\")\n"
},
{
"path": "ltx_video/models/autoencoders/dual_conv3d.py",
"content": "import math\nfrom typing import Tuple, Union\n\nimport torch\nimport torch.nn as nn\nimport torch.nn.functional as F\nfrom einops import rearrange\n\n\nclass DualConv3d(nn.Module):\n def __init__(\n self,\n in_channels,\n out_channels,\n kernel_size,\n stride: Union[int, Tuple[int, int, int]] = 1,\n padding: Union[int, Tuple[int, int, int]] = 0,\n dilation: Union[int, Tuple[int, int, int]] = 1,\n groups=1,\n bias=True,\n padding_mode=\"zeros\",\n ):\n super(DualConv3d, self).__init__()\n\n self.in_channels = in_channels\n self.out_channels = out_channels\n self.padding_mode = padding_mode\n # Ensure kernel_size, stride, padding, and dilation are tuples of length 3\n if isinstance(kernel_size, int):\n kernel_size = (kernel_size, kernel_size, kernel_size)\n if kernel_size == (1, 1, 1):\n raise ValueError(\n \"kernel_size must be greater than 1. Use make_linear_nd instead.\"\n )\n if isinstance(stride, int):\n stride = (stride, stride, stride)\n if isinstance(padding, int):\n padding = (padding, padding, padding)\n if isinstance(dilation, int):\n dilation = (dilation, dilation, dilation)\n\n # Set parameters for convolutions\n self.groups = groups\n self.bias = bias\n\n # Define the size of the channels after the first convolution\n intermediate_channels = (\n out_channels if in_channels < out_channels else in_channels\n )\n\n # Define parameters for the first convolution\n self.weight1 = nn.Parameter(\n torch.Tensor(\n intermediate_channels,\n in_channels // groups,\n 1,\n kernel_size[1],\n kernel_size[2],\n )\n )\n self.stride1 = (1, stride[1], stride[2])\n self.padding1 = (0, padding[1], padding[2])\n self.dilation1 = (1, dilation[1], dilation[2])\n if bias:\n self.bias1 = nn.Parameter(torch.Tensor(intermediate_channels))\n else:\n self.register_parameter(\"bias1\", None)\n\n # Define parameters for the second convolution\n self.weight2 = nn.Parameter(\n torch.Tensor(\n out_channels, intermediate_channels // groups, kernel_size[0], 1, 1\n )\n )\n self.stride2 = (stride[0], 1, 1)\n self.padding2 = (padding[0], 0, 0)\n self.dilation2 = (dilation[0], 1, 1)\n if bias:\n self.bias2 = nn.Parameter(torch.Tensor(out_channels))\n else:\n self.register_parameter(\"bias2\", None)\n\n # Initialize weights and biases\n self.reset_parameters()\n\n def reset_parameters(self):\n nn.init.kaiming_uniform_(self.weight1, a=math.sqrt(5))\n nn.init.kaiming_uniform_(self.weight2, a=math.sqrt(5))\n if self.bias:\n fan_in1, _ = nn.init._calculate_fan_in_and_fan_out(self.weight1)\n bound1 = 1 / math.sqrt(fan_in1)\n nn.init.uniform_(self.bias1, -bound1, bound1)\n fan_in2, _ = nn.init._calculate_fan_in_and_fan_out(self.weight2)\n bound2 = 1 / math.sqrt(fan_in2)\n nn.init.uniform_(self.bias2, -bound2, bound2)\n\n def forward(self, x, use_conv3d=False, skip_time_conv=False):\n if use_conv3d:\n return self.forward_with_3d(x=x, skip_time_conv=skip_time_conv)\n else:\n return self.forward_with_2d(x=x, skip_time_conv=skip_time_conv)\n\n def forward_with_3d(self, x, skip_time_conv):\n # First convolution\n x = F.conv3d(\n x,\n self.weight1,\n self.bias1,\n self.stride1,\n self.padding1,\n self.dilation1,\n self.groups,\n padding_mode=self.padding_mode,\n )\n\n if skip_time_conv:\n return x\n\n # Second convolution\n x = F.conv3d(\n x,\n self.weight2,\n self.bias2,\n self.stride2,\n self.padding2,\n self.dilation2,\n self.groups,\n padding_mode=self.padding_mode,\n )\n\n return x\n\n def forward_with_2d(self, x, skip_time_conv):\n b, c, d, h, w = x.shape\n\n # First 2D convolution\n x = rearrange(x, \"b c d h w -> (b d) c h w\")\n # Squeeze the depth dimension out of weight1 since it's 1\n weight1 = self.weight1.squeeze(2)\n # Select stride, padding, and dilation for the 2D convolution\n stride1 = (self.stride1[1], self.stride1[2])\n padding1 = (self.padding1[1], self.padding1[2])\n dilation1 = (self.dilation1[1], self.dilation1[2])\n x = F.conv2d(\n x,\n weight1,\n self.bias1,\n stride1,\n padding1,\n dilation1,\n self.groups,\n padding_mode=self.padding_mode,\n )\n\n _, _, h, w = x.shape\n\n if skip_time_conv:\n x = rearrange(x, \"(b d) c h w -> b c d h w\", b=b)\n return x\n\n # Second convolution which is essentially treated as a 1D convolution across the 'd' dimension\n x = rearrange(x, \"(b d) c h w -> (b h w) c d\", b=b)\n\n # Reshape weight2 to match the expected dimensions for conv1d\n weight2 = self.weight2.squeeze(-1).squeeze(-1)\n # Use only the relevant dimension for stride, padding, and dilation for the 1D convolution\n stride2 = self.stride2[0]\n padding2 = self.padding2[0]\n dilation2 = self.dilation2[0]\n x = F.conv1d(\n x,\n weight2,\n self.bias2,\n stride2,\n padding2,\n dilation2,\n self.groups,\n padding_mode=self.padding_mode,\n )\n x = rearrange(x, \"(b h w) c d -> b c d h w\", b=b, h=h, w=w)\n\n return x\n\n @property\n def weight(self):\n return self.weight2\n\n\ndef test_dual_conv3d_consistency():\n # Initialize parameters\n in_channels = 3\n out_channels = 5\n kernel_size = (3, 3, 3)\n stride = (2, 2, 2)\n padding = (1, 1, 1)\n\n # Create an instance of the DualConv3d class\n dual_conv3d = DualConv3d(\n in_channels=in_channels,\n out_channels=out_channels,\n kernel_size=kernel_size,\n stride=stride,\n padding=padding,\n bias=True,\n )\n\n # Example input tensor\n test_input = torch.randn(1, 3, 10, 10, 10)\n\n # Perform forward passes with both 3D and 2D settings\n output_conv3d = dual_conv3d(test_input, use_conv3d=True)\n output_2d = dual_conv3d(test_input, use_conv3d=False)\n\n # Assert that the outputs from both methods are sufficiently close\n assert torch.allclose(\n output_conv3d, output_2d, atol=1e-6\n ), \"Outputs are not consistent between 3D and 2D convolutions.\"\n"
},
{
"path": "ltx_video/models/autoencoders/latent_upsampler.py",
"content": "from typing import Optional, Union\nfrom pathlib import Path\nimport os\nimport json\n\nimport torch\nimport torch.nn as nn\nfrom einops import rearrange\nfrom diffusers import ConfigMixin, ModelMixin\nfrom safetensors.torch import safe_open\n\nfrom ltx_video.models.autoencoders.pixel_shuffle import PixelShuffleND\n\n\nclass ResBlock(nn.Module):\n def __init__(\n self, channels: int, mid_channels: Optional[int] = None, dims: int = 3\n ):\n super().__init__()\n if mid_channels is None:\n mid_channels = channels\n\n Conv = nn.Conv2d if dims == 2 else nn.Conv3d\n\n self.conv1 = Conv(channels, mid_channels, kernel_size=3, padding=1)\n self.norm1 = nn.GroupNorm(32, mid_channels)\n self.conv2 = Conv(mid_channels, channels, kernel_size=3, padding=1)\n self.norm2 = nn.GroupNorm(32, channels)\n self.activation = nn.SiLU()\n\n def forward(self, x: torch.Tensor) -> torch.Tensor:\n residual = x\n x = self.conv1(x)\n x = self.norm1(x)\n x = self.activation(x)\n x = self.conv2(x)\n x = self.norm2(x)\n x = self.activation(x + residual)\n return x\n\n\nclass LatentUpsampler(ModelMixin, ConfigMixin):\n \"\"\"\n Model to spatially upsample VAE latents.\n\n Args:\n in_channels (`int`): Number of channels in the input latent\n mid_channels (`int`): Number of channels in the middle layers\n num_blocks_per_stage (`int`): Number of ResBlocks to use in each stage (pre/post upsampling)\n dims (`int`): Number of dimensions for convolutions (2 or 3)\n spatial_upsample (`bool`): Whether to spatially upsample the latent\n temporal_upsample (`bool`): Whether to temporally upsample the latent\n \"\"\"\n\n def __init__(\n self,\n in_channels: int = 128,\n mid_channels: int = 512,\n num_blocks_per_stage: int = 4,\n dims: int = 3,\n spatial_upsample: bool = True,\n temporal_upsample: bool = False,\n ):\n super().__init__()\n\n self.in_channels = in_channels\n self.mid_channels = mid_channels\n self.num_blocks_per_stage = num_blocks_per_stage\n self.dims = dims\n self.spatial_upsample = spatial_upsample\n self.temporal_upsample = temporal_upsample\n\n Conv = nn.Conv2d if dims == 2 else nn.Conv3d\n\n self.initial_conv = Conv(in_channels, mid_channels, kernel_size=3, padding=1)\n self.initial_norm = nn.GroupNorm(32, mid_channels)\n self.initial_activation = nn.SiLU()\n\n self.res_blocks = nn.ModuleList(\n [ResBlock(mid_channels, dims=dims) for _ in range(num_blocks_per_stage)]\n )\n\n if spatial_upsample and temporal_upsample:\n self.upsampler = nn.Sequential(\n nn.Conv3d(mid_channels, 8 * mid_channels, kernel_size=3, padding=1),\n PixelShuffleND(3),\n )\n elif spatial_upsample:\n self.upsampler = nn.Sequential(\n nn.Conv2d(mid_channels, 4 * mid_channels, kernel_size=3, padding=1),\n PixelShuffleND(2),\n )\n elif temporal_upsample:\n self.upsampler = nn.Sequential(\n nn.Conv3d(mid_channels, 2 * mid_channels, kernel_size=3, padding=1),\n PixelShuffleND(1),\n )\n else:\n raise ValueError(\n \"Either spatial_upsample or temporal_upsample must be True\"\n )\n\n self.post_upsample_res_blocks = nn.ModuleList(\n [ResBlock(mid_channels, dims=dims) for _ in range(num_blocks_per_stage)]\n )\n\n self.final_conv = Conv(mid_channels, in_channels, kernel_size=3, padding=1)\n\n def forward(self, latent: torch.Tensor) -> torch.Tensor:\n b, c, f, h, w = latent.shape\n\n if self.dims == 2:\n x = rearrange(latent, \"b c f h w -> (b f) c h w\")\n x = self.initial_conv(x)\n x = self.initial_norm(x)\n x = self.initial_activation(x)\n\n for block in self.res_blocks:\n x = block(x)\n\n x = self.upsampler(x)\n\n for block in self.post_upsample_res_blocks:\n x = block(x)\n\n x = self.final_conv(x)\n x = rearrange(x, \"(b f) c h w -> b c f h w\", b=b, f=f)\n else:\n x = self.initial_conv(latent)\n x = self.initial_norm(x)\n x = self.initial_activation(x)\n\n for block in self.res_blocks:\n x = block(x)\n\n if self.temporal_upsample:\n x = self.upsampler(x)\n x = x[:, :, 1:, :, :]\n else:\n x = rearrange(x, \"b c f h w -> (b f) c h w\")\n x = self.upsampler(x)\n x = rearrange(x, \"(b f) c h w -> b c f h w\", b=b, f=f)\n\n for block in self.post_upsample_res_blocks:\n x = block(x)\n\n x = self.final_conv(x)\n\n return x\n\n @classmethod\n def from_config(cls, config):\n return cls(\n in_channels=config.get(\"in_channels\", 4),\n mid_channels=config.get(\"mid_channels\", 128),\n num_blocks_per_stage=config.get(\"num_blocks_per_stage\", 4),\n dims=config.get(\"dims\", 2),\n spatial_upsample=config.get(\"spatial_upsample\", True),\n temporal_upsample=config.get(\"temporal_upsample\", False),\n )\n\n def config(self):\n return {\n \"_class_name\": \"LatentUpsampler\",\n \"in_channels\": self.in_channels,\n \"mid_channels\": self.mid_channels,\n \"num_blocks_per_stage\": self.num_blocks_per_stage,\n \"dims\": self.dims,\n \"spatial_upsample\": self.spatial_upsample,\n \"temporal_upsample\": self.temporal_upsample,\n }\n\n @classmethod\n def from_pretrained(\n cls,\n pretrained_model_path: Optional[Union[str, os.PathLike]],\n *args,\n **kwargs,\n ):\n pretrained_model_path = Path(pretrained_model_path)\n if pretrained_model_path.is_file() and str(pretrained_model_path).endswith(\n \".safetensors\"\n ):\n state_dict = {}\n with safe_open(pretrained_model_path, framework=\"pt\", device=\"cpu\") as f:\n metadata = f.metadata()\n for k in f.keys():\n state_dict[k] = f.get_tensor(k)\n config = json.loads(metadata[\"config\"])\n with torch.device(\"meta\"):\n latent_upsampler = LatentUpsampler.from_config(config)\n latent_upsampler.load_state_dict(state_dict, assign=True)\n return latent_upsampler\n\n\nif __name__ == \"__main__\":\n latent_upsampler = LatentUpsampler(num_blocks_per_stage=4, dims=3)\n print(latent_upsampler)\n total_params = sum(p.numel() for p in latent_upsampler.parameters())\n print(f\"Total number of parameters: {total_params:,}\")\n latent = torch.randn(1, 128, 9, 16, 16)\n upsampled_latent = latent_upsampler(latent)\n print(f\"Upsampled latent shape: {upsampled_latent.shape}\")\n"
},
{
"path": "ltx_video/models/autoencoders/pixel_norm.py",
"content": "import torch\nfrom torch import nn\n\n\nclass PixelNorm(nn.Module):\n def __init__(self, dim=1, eps=1e-8):\n super(PixelNorm, self).__init__()\n self.dim = dim\n self.eps = eps\n\n def forward(self, x):\n return x / torch.sqrt(torch.mean(x**2, dim=self.dim, keepdim=True) + self.eps)\n"
},
{
"path": "ltx_video/models/autoencoders/pixel_shuffle.py",
"content": "import torch.nn as nn\nfrom einops import rearrange\n\n\nclass PixelShuffleND(nn.Module):\n def __init__(self, dims, upscale_factors=(2, 2, 2)):\n super().__init__()\n assert dims in [1, 2, 3], \"dims must be 1, 2, or 3\"\n self.dims = dims\n self.upscale_factors = upscale_factors\n\n def forward(self, x):\n if self.dims == 3:\n return rearrange(\n x,\n \"b (c p1 p2 p3) d h w -> b c (d p1) (h p2) (w p3)\",\n p1=self.upscale_factors[0],\n p2=self.upscale_factors[1],\n p3=self.upscale_factors[2],\n )\n elif self.dims == 2:\n return rearrange(\n x,\n \"b (c p1 p2) h w -> b c (h p1) (w p2)\",\n p1=self.upscale_factors[0],\n p2=self.upscale_factors[1],\n )\n elif self.dims == 1:\n return rearrange(\n x,\n \"b (c p1) f h w -> b c (f p1) h w\",\n p1=self.upscale_factors[0],\n )\n"
},
{
"path": "ltx_video/models/autoencoders/vae.py",
"content": "from typing import Optional, Union\n\nimport torch\nimport inspect\nimport math\nimport torch.nn as nn\nfrom diffusers import ConfigMixin, ModelMixin\nfrom diffusers.models.autoencoders.vae import (\n DecoderOutput,\n DiagonalGaussianDistribution,\n)\nfrom diffusers.models.modeling_outputs import AutoencoderKLOutput\nfrom ltx_video.models.autoencoders.conv_nd_factory import make_conv_nd\n\n\nclass AutoencoderKLWrapper(ModelMixin, ConfigMixin):\n \"\"\"Variational Autoencoder (VAE) model with KL loss.\n\n VAE from the paper Auto-Encoding Variational Bayes by Diederik P. Kingma and Max Welling.\n This model is a wrapper around an encoder and a decoder, and it adds a KL loss term to the reconstruction loss.\n\n Args:\n encoder (`nn.Module`):\n Encoder module.\n decoder (`nn.Module`):\n Decoder module.\n latent_channels (`int`, *optional*, defaults to 4):\n Number of latent channels.\n \"\"\"\n\n def __init__(\n self,\n encoder: nn.Module,\n decoder: nn.Module,\n latent_channels: int = 4,\n dims: int = 2,\n sample_size=512,\n use_quant_conv: bool = True,\n normalize_latent_channels: bool = False,\n ):\n super().__init__()\n\n # pass init params to Encoder\n self.encoder = encoder\n self.use_quant_conv = use_quant_conv\n self.normalize_latent_channels = normalize_latent_channels\n\n # pass init params to Decoder\n quant_dims = 2 if dims == 2 else 3\n self.decoder = decoder\n if use_quant_conv:\n self.quant_conv = make_conv_nd(\n quant_dims, 2 * latent_channels, 2 * latent_channels, 1\n )\n self.post_quant_conv = make_conv_nd(\n quant_dims, latent_channels, latent_channels, 1\n )\n else:\n self.quant_conv = nn.Identity()\n self.post_quant_conv = nn.Identity()\n\n if normalize_latent_channels:\n if dims == 2:\n self.latent_norm_out = nn.BatchNorm2d(latent_channels, affine=False)\n else:\n self.latent_norm_out = nn.BatchNorm3d(latent_channels, affine=False)\n else:\n self.latent_norm_out = nn.Identity()\n self.use_z_tiling = False\n self.use_hw_tiling = False\n self.dims = dims\n self.z_sample_size = 1\n\n self.decoder_params = inspect.signature(self.decoder.forward).parameters\n\n # only relevant if vae tiling is enabled\n self.set_tiling_params(sample_size=sample_size, overlap_factor=0.25)\n\n def set_tiling_params(self, sample_size: int = 512, overlap_factor: float = 0.25):\n self.tile_sample_min_size = sample_size\n num_blocks = len(self.encoder.down_blocks)\n self.tile_latent_min_size = int(sample_size / (2 ** (num_blocks - 1)))\n self.tile_overlap_factor = overlap_factor\n\n def enable_z_tiling(self, z_sample_size: int = 8):\n r\"\"\"\n Enable tiling during VAE decoding.\n\n When this option is enabled, the VAE will split the input tensor in tiles to compute decoding in several\n steps. This is useful to save some memory and allow larger batch sizes.\n \"\"\"\n self.use_z_tiling = z_sample_size > 1\n self.z_sample_size = z_sample_size\n assert (\n z_sample_size % 8 == 0 or z_sample_size == 1\n ), f\"z_sample_size must be a multiple of 8 or 1. Got {z_sample_size}.\"\n\n def disable_z_tiling(self):\n r\"\"\"\n Disable tiling during VAE decoding. If `use_tiling` was previously invoked, this method will go back to computing\n decoding in one step.\n \"\"\"\n self.use_z_tiling = False\n\n def enable_hw_tiling(self):\n r\"\"\"\n Enable tiling during VAE decoding along the height and width dimension.\n \"\"\"\n self.use_hw_tiling = True\n\n def disable_hw_tiling(self):\n r\"\"\"\n Disable tiling during VAE decoding along the height and width dimension.\n \"\"\"\n self.use_hw_tiling = False\n\n def _hw_tiled_encode(self, x: torch.FloatTensor, return_dict: bool = True):\n overlap_size = int(self.tile_sample_min_size * (1 - self.tile_overlap_factor))\n blend_extent = int(self.tile_latent_min_size * self.tile_overlap_factor)\n row_limit = self.tile_latent_min_size - blend_extent\n\n # Split the image into 512x512 tiles and encode them separately.\n rows = []\n for i in range(0, x.shape[3], overlap_size):\n row = []\n for j in range(0, x.shape[4], overlap_size):\n tile = x[\n :,\n :,\n :,\n i : i + self.tile_sample_min_size,\n j : j + self.tile_sample_min_size,\n ]\n tile = self.encoder(tile)\n tile = self.quant_conv(tile)\n row.append(tile)\n rows.append(row)\n result_rows = []\n for i, row in enumerate(rows):\n result_row = []\n for j, tile in enumerate(row):\n # blend the above tile and the left tile\n # to the current tile and add the current tile to the result row\n if i > 0:\n tile = self.blend_v(rows[i - 1][j], tile, blend_extent)\n if j > 0:\n tile = self.blend_h(row[j - 1], tile, blend_extent)\n result_row.append(tile[:, :, :, :row_limit, :row_limit])\n result_rows.append(torch.cat(result_row, dim=4))\n\n moments = torch.cat(result_rows, dim=3)\n return moments\n\n def blend_z(\n self, a: torch.Tensor, b: torch.Tensor, blend_extent: int\n ) -> torch.Tensor:\n blend_extent = min(a.shape[2], b.shape[2], blend_extent)\n for z in range(blend_extent):\n b[:, :, z, :, :] = a[:, :, -blend_extent + z, :, :] * (\n 1 - z / blend_extent\n ) + b[:, :, z, :, :] * (z / blend_extent)\n return b\n\n def blend_v(\n self, a: torch.Tensor, b: torch.Tensor, blend_extent: int\n ) -> torch.Tensor:\n blend_extent = min(a.shape[3], b.shape[3], blend_extent)\n for y in range(blend_extent):\n b[:, :, :, y, :] = a[:, :, :, -blend_extent + y, :] * (\n 1 - y / blend_extent\n ) + b[:, :, :, y, :] * (y / blend_extent)\n return b\n\n def blend_h(\n self, a: torch.Tensor, b: torch.Tensor, blend_extent: int\n ) -> torch.Tensor:\n blend_extent = min(a.shape[4], b.shape[4], blend_extent)\n for x in range(blend_extent):\n b[:, :, :, :, x] = a[:, :, :, :, -blend_extent + x] * (\n 1 - x / blend_extent\n ) + b[:, :, :, :, x] * (x / blend_extent)\n return b\n\n def _hw_tiled_decode(self, z: torch.FloatTensor, target_shape):\n overlap_size = int(self.tile_latent_min_size * (1 - self.tile_overlap_factor))\n blend_extent = int(self.tile_sample_min_size * self.tile_overlap_factor)\n row_limit = self.tile_sample_min_size - blend_extent\n tile_target_shape = (\n *target_shape[:3],\n self.tile_sample_min_size,\n self.tile_sample_min_size,\n )\n # Split z into overlapping 64x64 tiles and decode them separately.\n # The tiles have an overlap to avoid seams between tiles.\n rows = []\n for i in range(0, z.shape[3], overlap_size):\n row = []\n for j in range(0, z.shape[4], overlap_size):\n tile = z[\n :,\n :,\n :,\n i : i + self.tile_latent_min_size,\n j : j + self.tile_latent_min_size,\n ]\n tile = self.post_quant_conv(tile)\n decoded = self.decoder(tile, target_shape=tile_target_shape)\n row.append(decoded)\n rows.append(row)\n result_rows = []\n for i, row in enumerate(rows):\n result_row = []\n for j, tile in enumerate(row):\n # blend the above tile and the left tile\n # to the current tile and add the current tile to the result row\n if i > 0:\n tile = self.blend_v(rows[i - 1][j], tile, blend_extent)\n if j > 0:\n tile = self.blend_h(row[j - 1], tile, blend_extent)\n result_row.append(tile[:, :, :, :row_limit, :row_limit])\n result_rows.append(torch.cat(result_row, dim=4))\n\n dec = torch.cat(result_rows, dim=3)\n return dec\n\n def encode(\n self, z: torch.FloatTensor, return_dict: bool = True\n ) -> Union[DecoderOutput, torch.FloatTensor]:\n if self.use_z_tiling and z.shape[2] > self.z_sample_size > 1:\n num_splits = z.shape[2] // self.z_sample_size\n sizes = [self.z_sample_size] * num_splits\n sizes = (\n sizes + [z.shape[2] - sum(sizes)]\n if z.shape[2] - sum(sizes) > 0\n else sizes\n )\n tiles = z.split(sizes, dim=2)\n moments_tiles = [\n (\n self._hw_tiled_encode(z_tile, return_dict)\n if self.use_hw_tiling\n else self._encode(z_tile)\n )\n for z_tile in tiles\n ]\n moments = torch.cat(moments_tiles, dim=2)\n\n else:\n moments = (\n self._hw_tiled_encode(z, return_dict)\n if self.use_hw_tiling\n else self._encode(z)\n )\n\n posterior = DiagonalGaussianDistribution(moments)\n if not return_dict:\n return (posterior,)\n\n return AutoencoderKLOutput(latent_dist=posterior)\n\n def _normalize_latent_channels(self, z: torch.FloatTensor) -> torch.FloatTensor:\n if isinstance(self.latent_norm_out, nn.BatchNorm3d):\n _, c, _, _, _ = z.shape\n z = torch.cat(\n [\n self.latent_norm_out(z[:, : c // 2, :, :, :]),\n z[:, c // 2 :, :, :, :],\n ],\n dim=1,\n )\n elif isinstance(self.latent_norm_out, nn.BatchNorm2d):\n raise NotImplementedError(\"BatchNorm2d not supported\")\n return z\n\n def _unnormalize_latent_channels(self, z: torch.FloatTensor) -> torch.FloatTensor:\n if isinstance(self.latent_norm_out, nn.BatchNorm3d):\n running_mean = self.latent_norm_out.running_mean.view(1, -1, 1, 1, 1)\n running_var = self.latent_norm_out.running_var.view(1, -1, 1, 1, 1)\n eps = self.latent_norm_out.eps\n\n z = z * torch.sqrt(running_var + eps) + running_mean\n elif isinstance(self.latent_norm_out, nn.BatchNorm3d):\n raise NotImplementedError(\"BatchNorm2d not supported\")\n return z\n\n def _encode(self, x: torch.FloatTensor) -> AutoencoderKLOutput:\n h = self.encoder(x)\n moments = self.quant_conv(h)\n moments = self._normalize_latent_channels(moments)\n return moments\n\n def _decode(\n self,\n z: torch.FloatTensor,\n target_shape=None,\n timestep: Optional[torch.Tensor] = None,\n ) -> Union[DecoderOutput, torch.FloatTensor]:\n z = self._unnormalize_latent_channels(z)\n z = self.post_quant_conv(z)\n if \"timestep\" in self.decoder_params:\n dec = self.decoder(z, target_shape=target_shape, timestep=timestep)\n else:\n dec = self.decoder(z, target_shape=target_shape)\n return dec\n\n def decode(\n self,\n z: torch.FloatTensor,\n return_dict: bool = True,\n target_shape=None,\n timestep: Optional[torch.Tensor] = None,\n ) -> Union[DecoderOutput, torch.FloatTensor]:\n assert target_shape is not None, \"target_shape must be provided for decoding\"\n if self.use_z_tiling and z.shape[2] > self.z_sample_size > 1:\n reduction_factor = int(\n self.encoder.patch_size_t\n * 2\n ** (\n len(self.encoder.down_blocks)\n - 1\n - math.sqrt(self.encoder.patch_size)\n )\n )\n split_size = self.z_sample_size // reduction_factor\n num_splits = z.shape[2] // split_size\n\n # copy target shape, and divide frame dimension (=2) by the context size\n target_shape_split = list(target_shape)\n target_shape_split[2] = target_shape[2] // num_splits\n\n decoded_tiles = [\n (\n self._hw_tiled_decode(z_tile, target_shape_split)\n if self.use_hw_tiling\n else self._decode(z_tile, target_shape=target_shape_split)\n )\n for z_tile in torch.tensor_split(z, num_splits, dim=2)\n ]\n decoded = torch.cat(decoded_tiles, dim=2)\n else:\n decoded = (\n self._hw_tiled_decode(z, target_shape)\n if self.use_hw_tiling\n else self._decode(z, target_shape=target_shape, timestep=timestep)\n )\n\n if not return_dict:\n return (decoded,)\n\n return DecoderOutput(sample=decoded)\n\n def forward(\n self,\n sample: torch.FloatTensor,\n sample_posterior: bool = False,\n return_dict: bool = True,\n generator: Optional[torch.Generator] = None,\n ) -> Union[DecoderOutput, torch.FloatTensor]:\n r\"\"\"\n Args:\n sample (`torch.FloatTensor`): Input sample.\n sample_posterior (`bool`, *optional*, defaults to `False`):\n Whether to sample from the posterior.\n return_dict (`bool`, *optional*, defaults to `True`):\n Whether to return a [`DecoderOutput`] instead of a plain tuple.\n generator (`torch.Generator`, *optional*):\n Generator used to sample from the posterior.\n \"\"\"\n x = sample\n posterior = self.encode(x).latent_dist\n if sample_posterior:\n z = posterior.sample(generator=generator)\n else:\n z = posterior.mode()\n dec = self.decode(z, target_shape=sample.shape).sample\n\n if not return_dict:\n return (dec,)\n\n return DecoderOutput(sample=dec)\n"
},
{
"path": "ltx_video/models/autoencoders/vae_encode.py",
"content": "from typing import Tuple\nimport torch\nfrom diffusers import AutoencoderKL\nfrom einops import rearrange\nfrom torch import Tensor\n\n\nfrom ltx_video.models.autoencoders.causal_video_autoencoder import (\n CausalVideoAutoencoder,\n)\nfrom ltx_video.models.autoencoders.video_autoencoder import (\n Downsample3D,\n VideoAutoencoder,\n)\n\ntry:\n import torch_xla.core.xla_model as xm\nexcept ImportError:\n xm = None\n\n\ndef vae_encode(\n media_items: Tensor,\n vae: AutoencoderKL,\n split_size: int = 1,\n vae_per_channel_normalize=False,\n) -> Tensor:\n \"\"\"\n Encodes media items (images or videos) into latent representations using a specified VAE model.\n The function supports processing batches of images or video frames and can handle the processing\n in smaller sub-batches if needed.\n\n Args:\n media_items (Tensor): A torch Tensor containing the media items to encode. The expected\n shape is (batch_size, channels, height, width) for images or (batch_size, channels,\n frames, height, width) for videos.\n vae (AutoencoderKL): An instance of the `AutoencoderKL` class from the `diffusers` library,\n pre-configured and loaded with the appropriate model weights.\n split_size (int, optional): The number of sub-batches to split the input batch into for encoding.\n If set to more than 1, the input media items are processed in smaller batches according to\n this value. Defaults to 1, which processes all items in a single batch.\n\n Returns:\n Tensor: A torch Tensor of the encoded latent representations. The shape of the tensor is adjusted\n to match the input shape, scaled by the model's configuration.\n\n Examples:\n >>> import torch\n >>> from diffusers import AutoencoderKL\n >>> vae = AutoencoderKL.from_pretrained('your-model-name')\n >>> images = torch.rand(10, 3, 8 256, 256) # Example tensor with 10 videos of 8 frames.\n >>> latents = vae_encode(images, vae)\n >>> print(latents.shape) # Output shape will depend on the model's latent configuration.\n\n Note:\n In case of a video, the function encodes the media item frame-by frame.\n \"\"\"\n is_video_shaped = media_items.dim() == 5\n batch_size, channels = media_items.shape[0:2]\n\n if channels != 3:\n raise ValueError(f\"Expects tensors with 3 channels, got {channels}.\")\n\n if is_video_shaped and not isinstance(\n vae, (VideoAutoencoder, CausalVideoAutoencoder)\n ):\n media_items = rearrange(media_items, \"b c n h w -> (b n) c h w\")\n if split_size > 1:\n if len(media_items) % split_size != 0:\n raise ValueError(\n \"Error: The batch size must be divisible by 'train.vae_bs_split\"\n )\n encode_bs = len(media_items) // split_size\n # latents = [vae.encode(image_batch).latent_dist.sample() for image_batch in media_items.split(encode_bs)]\n latents = []\n if media_items.device.type == \"xla\":\n xm.mark_step()\n for image_batch in media_items.split(encode_bs):\n latents.append(vae.encode(image_batch).latent_dist.sample())\n if media_items.device.type == \"xla\":\n xm.mark_step()\n latents = torch.cat(latents, dim=0)\n else:\n latents = vae.encode(media_items).latent_dist.sample()\n\n latents = normalize_latents(latents, vae, vae_per_channel_normalize)\n if is_video_shaped and not isinstance(\n vae, (VideoAutoencoder, CausalVideoAutoencoder)\n ):\n latents = rearrange(latents, \"(b n) c h w -> b c n h w\", b=batch_size)\n return latents\n\n\ndef vae_decode(\n latents: Tensor,\n vae: AutoencoderKL,\n is_video: bool = True,\n split_size: int = 1,\n vae_per_channel_normalize=False,\n timestep=None,\n) -> Tensor:\n is_video_shaped = latents.dim() == 5\n batch_size = latents.shape[0]\n\n if is_video_shaped and not isinstance(\n vae, (VideoAutoencoder, CausalVideoAutoencoder)\n ):\n latents = rearrange(latents, \"b c n h w -> (b n) c h w\")\n if split_size > 1:\n if len(latents) % split_size != 0:\n raise ValueError(\n \"Error: The batch size must be divisible by 'train.vae_bs_split\"\n )\n encode_bs = len(latents) // split_size\n image_batch = [\n _run_decoder(\n latent_batch, vae, is_video, vae_per_channel_normalize, timestep\n )\n for latent_batch in latents.split(encode_bs)\n ]\n images = torch.cat(image_batch, dim=0)\n else:\n images = _run_decoder(\n latents, vae, is_video, vae_per_channel_normalize, timestep\n )\n\n if is_video_shaped and not isinstance(\n vae, (VideoAutoencoder, CausalVideoAutoencoder)\n ):\n images = rearrange(images, \"(b n) c h w -> b c n h w\", b=batch_size)\n return images\n\n\ndef _run_decoder(\n latents: Tensor,\n vae: AutoencoderKL,\n is_video: bool,\n vae_per_channel_normalize=False,\n timestep=None,\n) -> Tensor:\n if isinstance(vae, (VideoAutoencoder, CausalVideoAutoencoder)):\n *_, fl, hl, wl = latents.shape\n temporal_scale, spatial_scale, _ = get_vae_size_scale_factor(vae)\n latents = latents.to(vae.dtype)\n vae_decode_kwargs = {}\n if timestep is not None:\n vae_decode_kwargs[\"timestep\"] = timestep\n image = vae.decode(\n un_normalize_latents(latents, vae, vae_per_channel_normalize),\n return_dict=False,\n target_shape=(\n 1,\n 3,\n fl * temporal_scale if is_video else 1,\n hl * spatial_scale,\n wl * spatial_scale,\n ),\n **vae_decode_kwargs,\n )[0]\n else:\n image = vae.decode(\n un_normalize_latents(latents, vae, vae_per_channel_normalize),\n return_dict=False,\n )[0]\n return image\n\n\ndef get_vae_size_scale_factor(vae: AutoencoderKL) -> float:\n if isinstance(vae, CausalVideoAutoencoder):\n spatial = vae.spatial_downscale_factor\n temporal = vae.temporal_downscale_factor\n else:\n down_blocks = len(\n [\n block\n for block in vae.encoder.down_blocks\n if isinstance(block.downsample, Downsample3D)\n ]\n )\n spatial = vae.config.patch_size * 2**down_blocks\n temporal = (\n vae.config.patch_size_t * 2**down_blocks\n if isinstance(vae, VideoAutoencoder)\n else 1\n )\n\n return (temporal, spatial, spatial)\n\n\ndef latent_to_pixel_coords(\n latent_coords: Tensor, vae: AutoencoderKL, causal_fix: bool = False\n) -> Tensor:\n \"\"\"\n Converts latent coordinates to pixel coordinates by scaling them according to the VAE's\n configuration.\n\n Args:\n latent_coords (Tensor): A tensor of shape [batch_size, 3, num_latents]\n containing the latent corner coordinates of each token.\n vae (AutoencoderKL): The VAE model\n causal_fix (bool): Whether to take into account the different temporal scale\n of the first frame. Default = False for backwards compatibility.\n Returns:\n Tensor: A tensor of pixel coordinates corresponding to the input latent coordinates.\n \"\"\"\n\n scale_factors = get_vae_size_scale_factor(vae)\n causal_fix = isinstance(vae, CausalVideoAutoencoder) and causal_fix\n pixel_coords = latent_to_pixel_coords_from_factors(\n latent_coords, scale_factors, causal_fix\n )\n return pixel_coords\n\n\ndef latent_to_pixel_coords_from_factors(\n latent_coords: Tensor, scale_factors: Tuple, causal_fix: bool = False\n) -> Tensor:\n pixel_coords = (\n latent_coords\n * torch.tensor(scale_factors, device=latent_coords.device)[None, :, None]\n )\n if causal_fix:\n # Fix temporal scale for first frame to 1 due to causality\n pixel_coords[:, 0] = (pixel_coords[:, 0] + 1 - scale_factors[0]).clamp(min=0)\n return pixel_coords\n\n\ndef normalize_latents(\n latents: Tensor, vae: AutoencoderKL, vae_per_channel_normalize: bool = False\n) -> Tensor:\n return (\n (latents - vae.mean_of_means.to(latents.dtype).view(1, -1, 1, 1, 1))\n / vae.std_of_means.to(latents.dtype).view(1, -1, 1, 1, 1)\n if vae_per_channel_normalize\n else latents * vae.config.scaling_factor\n )\n\n\ndef un_normalize_latents(\n latents: Tensor, vae: AutoencoderKL, vae_per_channel_normalize: bool = False\n) -> Tensor:\n return (\n latents * vae.std_of_means.to(latents.dtype).view(1, -1, 1, 1, 1)\n + vae.mean_of_means.to(latents.dtype).view(1, -1, 1, 1, 1)\n if vae_per_channel_normalize\n else latents / vae.config.scaling_factor\n )\n"
},
{
"path": "ltx_video/models/autoencoders/video_autoencoder.py",
"content": "import json\nimport os\nfrom functools import partial\nfrom types import SimpleNamespace\nfrom typing import Any, Mapping, Optional, Tuple, Union\n\nimport torch\nfrom einops import rearrange\nfrom torch import nn\nfrom torch.nn import functional\n\nfrom diffusers.utils import logging\n\nfrom ltx_video.utils.torch_utils import Identity\nfrom ltx_video.models.autoencoders.conv_nd_factory import make_conv_nd, make_linear_nd\nfrom ltx_video.models.autoencoders.pixel_norm import PixelNorm\nfrom ltx_video.models.autoencoders.vae import AutoencoderKLWrapper\n\nlogger = logging.get_logger(__name__)\n\n\nclass VideoAutoencoder(AutoencoderKLWrapper):\n @classmethod\n def from_pretrained(\n cls,\n pretrained_model_name_or_path: Optional[Union[str, os.PathLike]],\n *args,\n **kwargs,\n ):\n config_local_path = pretrained_model_name_or_path / \"config.json\"\n config = cls.load_config(config_local_path, **kwargs)\n video_vae = cls.from_config(config)\n video_vae.to(kwargs[\"torch_dtype\"])\n\n model_local_path = pretrained_model_name_or_path / \"autoencoder.pth\"\n ckpt_state_dict = torch.load(model_local_path)\n video_vae.load_state_dict(ckpt_state_dict)\n\n statistics_local_path = (\n pretrained_model_name_or_path / \"per_channel_statistics.json\"\n )\n if statistics_local_path.exists():\n with open(statistics_local_path, \"r\") as file:\n data = json.load(file)\n transposed_data = list(zip(*data[\"data\"]))\n data_dict = {\n col: torch.tensor(vals)\n for col, vals in zip(data[\"columns\"], transposed_data)\n }\n video_vae.register_buffer(\"std_of_means\", data_dict[\"std-of-means\"])\n video_vae.register_buffer(\n \"mean_of_means\",\n data_dict.get(\n \"mean-of-means\", torch.zeros_like(data_dict[\"std-of-means\"])\n ),\n )\n\n return video_vae\n\n @staticmethod\n def from_config(config):\n assert (\n config[\"_class_name\"] == \"VideoAutoencoder\"\n ), \"config must have _class_name=VideoAutoencoder\"\n if isinstance(config[\"dims\"], list):\n config[\"dims\"] = tuple(config[\"dims\"])\n\n assert config[\"dims\"] in [2, 3, (2, 1)], \"dims must be 2, 3 or (2, 1)\"\n\n double_z = config.get(\"double_z\", True)\n latent_log_var = config.get(\n \"latent_log_var\", \"per_channel\" if double_z else \"none\"\n )\n use_quant_conv = config.get(\"use_quant_conv\", True)\n\n if use_quant_conv and latent_log_var == \"uniform\":\n raise ValueError(\"uniform latent_log_var requires use_quant_conv=False\")\n\n encoder = Encoder(\n dims=config[\"dims\"],\n in_channels=config.get(\"in_channels\", 3),\n out_channels=config[\"latent_channels\"],\n block_out_channels=config[\"block_out_channels\"],\n patch_size=config.get(\"patch_size\", 1),\n latent_log_var=latent_log_var,\n norm_layer=config.get(\"norm_layer\", \"group_norm\"),\n patch_size_t=config.get(\"patch_size_t\", config.get(\"patch_size\", 1)),\n add_channel_padding=config.get(\"add_channel_padding\", False),\n )\n\n decoder = Decoder(\n dims=config[\"dims\"],\n in_channels=config[\"latent_channels\"],\n out_channels=config.get(\"out_channels\", 3),\n block_out_channels=config[\"block_out_channels\"],\n patch_size=config.get(\"patch_size\", 1),\n norm_layer=config.get(\"norm_layer\", \"group_norm\"),\n patch_size_t=config.get(\"patch_size_t\", config.get(\"patch_size\", 1)),\n add_channel_padding=config.get(\"add_channel_padding\", False),\n )\n\n dims = config[\"dims\"]\n return VideoAutoencoder(\n encoder=encoder,\n decoder=decoder,\n latent_channels=config[\"latent_channels\"],\n dims=dims,\n use_quant_conv=use_quant_conv,\n )\n\n @property\n def config(self):\n return SimpleNamespace(\n _class_name=\"VideoAutoencoder\",\n dims=self.dims,\n in_channels=self.encoder.conv_in.in_channels\n // (self.encoder.patch_size_t * self.encoder.patch_size**2),\n out_channels=self.decoder.conv_out.out_channels\n // (self.decoder.patch_size_t * self.decoder.patch_size**2),\n latent_channels=self.decoder.conv_in.in_channels,\n block_out_channels=[\n self.encoder.down_blocks[i].res_blocks[-1].conv1.out_channels\n for i in range(len(self.encoder.down_blocks))\n ],\n scaling_factor=1.0,\n norm_layer=self.encoder.norm_layer,\n patch_size=self.encoder.patch_size,\n latent_log_var=self.encoder.latent_log_var,\n use_quant_conv=self.use_quant_conv,\n patch_size_t=self.encoder.patch_size_t,\n add_channel_padding=self.encoder.add_channel_padding,\n )\n\n @property\n def is_video_supported(self):\n \"\"\"\n Check if the model supports video inputs of shape (B, C, F, H, W). Otherwise, the model only supports 2D images.\n \"\"\"\n return self.dims != 2\n\n @property\n def downscale_factor(self):\n return self.encoder.downsample_factor\n\n def to_json_string(self) -> str:\n import json\n\n return json.dumps(self.config.__dict__)\n\n def load_state_dict(self, state_dict: Mapping[str, Any], strict: bool = True):\n model_keys = set(name for name, _ in self.named_parameters())\n\n key_mapping = {\n \".resnets.\": \".res_blocks.\",\n \"downsamplers.0\": \"downsample\",\n \"upsamplers.0\": \"upsample\",\n }\n\n converted_state_dict = {}\n for key, value in state_dict.items():\n for k, v in key_mapping.items():\n key = key.replace(k, v)\n\n if \"norm\" in key and key not in model_keys:\n logger.info(\n f\"Removing key {key} from state_dict as it is not present in the model\"\n )\n continue\n\n converted_state_dict[key] = value\n\n super().load_state_dict(converted_state_dict, strict=strict)\n\n def last_layer(self):\n if hasattr(self.decoder, \"conv_out\"):\n if isinstance(self.decoder.conv_out, nn.Sequential):\n last_layer = self.decoder.conv_out[-1]\n else:\n last_layer = self.decoder.conv_out\n else:\n last_layer = self.decoder.layers[-1]\n return last_layer\n\n\nclass Encoder(nn.Module):\n r\"\"\"\n The `Encoder` layer of a variational autoencoder that encodes its input into a latent representation.\n\n Args:\n in_channels (`int`, *optional*, defaults to 3):\n The number of input channels.\n out_channels (`int`, *optional*, defaults to 3):\n The number of output channels.\n block_out_channels (`Tuple[int, ...]`, *optional*, defaults to `(64,)`):\n The number of output channels for each block.\n layers_per_block (`int`, *optional*, defaults to 2):\n The number of layers per block.\n norm_num_groups (`int`, *optional*, defaults to 32):\n The number of groups for normalization.\n patch_size (`int`, *optional*, defaults to 1):\n The patch size to use. Should be a power of 2.\n norm_layer (`str`, *optional*, defaults to `group_norm`):\n The normalization layer to use. Can be either `group_norm` or `pixel_norm`.\n latent_log_var (`str`, *optional*, defaults to `per_channel`):\n The number of channels for the log variance. Can be either `per_channel`, `uniform`, or `none`.\n \"\"\"\n\n def __init__(\n self,\n dims: Union[int, Tuple[int, int]] = 3,\n in_channels: int = 3,\n out_channels: int = 3,\n block_out_channels: Tuple[int, ...] = (64,),\n layers_per_block: int = 2,\n norm_num_groups: int = 32,\n patch_size: Union[int, Tuple[int]] = 1,\n norm_layer: str = \"group_norm\", # group_norm, pixel_norm\n latent_log_var: str = \"per_channel\",\n patch_size_t: Optional[int] = None,\n add_channel_padding: Optional[bool] = False,\n ):\n super().__init__()\n self.patch_size = patch_size\n self.patch_size_t = patch_size_t if patch_size_t is not None else patch_size\n self.add_channel_padding = add_channel_padding\n self.layers_per_block = layers_per_block\n self.norm_layer = norm_layer\n self.latent_channels = out_channels\n self.latent_log_var = latent_log_var\n if add_channel_padding:\n in_channels = in_channels * self.patch_size**3\n else:\n in_channels = in_channels * self.patch_size_t * self.patch_size**2\n self.in_channels = in_channels\n output_channel = block_out_channels[0]\n\n self.conv_in = make_conv_nd(\n dims=dims,\n in_channels=in_channels,\n out_channels=output_channel,\n kernel_size=3,\n stride=1,\n padding=1,\n )\n\n self.down_blocks = nn.ModuleList([])\n\n for i in range(len(block_out_channels)):\n input_channel = output_channel\n output_channel = block_out_channels[i]\n is_final_block = i == len(block_out_channels) - 1\n\n down_block = DownEncoderBlock3D(\n dims=dims,\n in_channels=input_channel,\n out_channels=output_channel,\n num_layers=self.layers_per_block,\n add_downsample=not is_final_block and 2**i >= patch_size,\n resnet_eps=1e-6,\n downsample_padding=0,\n resnet_groups=norm_num_groups,\n norm_layer=norm_layer,\n )\n self.down_blocks.append(down_block)\n\n self.mid_block = UNetMidBlock3D(\n dims=dims,\n in_channels=block_out_channels[-1],\n num_layers=self.layers_per_block,\n resnet_eps=1e-6,\n resnet_groups=norm_num_groups,\n norm_layer=norm_layer,\n )\n\n # out\n if norm_layer == \"group_norm\":\n self.conv_norm_out = nn.GroupNorm(\n num_channels=block_out_channels[-1],\n num_groups=norm_num_groups,\n eps=1e-6,\n )\n elif norm_layer == \"pixel_norm\":\n self.conv_norm_out = PixelNorm()\n self.conv_act = nn.SiLU()\n\n conv_out_channels = out_channels\n if latent_log_var == \"per_channel\":\n conv_out_channels *= 2\n elif latent_log_var == \"uniform\":\n conv_out_channels += 1\n elif latent_log_var != \"none\":\n raise ValueError(f\"Invalid latent_log_var: {latent_log_var}\")\n self.conv_out = make_conv_nd(\n dims, block_out_channels[-1], conv_out_channels, 3, padding=1\n )\n\n self.gradient_checkpointing = False\n\n @property\n def downscale_factor(self):\n return (\n 2\n ** len(\n [\n block\n for block in self.down_blocks\n if isinstance(block.downsample, Downsample3D)\n ]\n )\n * self.patch_size\n )\n\n def forward(\n self, sample: torch.FloatTensor, return_features=False\n ) -> torch.FloatTensor:\n r\"\"\"The forward method of the `Encoder` class.\"\"\"\n\n downsample_in_time = sample.shape[2] != 1\n\n # patchify\n patch_size_t = self.patch_size_t if downsample_in_time else 1\n sample = patchify(\n sample,\n patch_size_hw=self.patch_size,\n patch_size_t=patch_size_t,\n add_channel_padding=self.add_channel_padding,\n )\n\n sample = self.conv_in(sample)\n\n checkpoint_fn = (\n partial(torch.utils.checkpoint.checkpoint, use_reentrant=False)\n if self.gradient_checkpointing and self.training\n else lambda x: x\n )\n\n if return_features:\n features = []\n for down_block in self.down_blocks:\n sample = checkpoint_fn(down_block)(\n sample, downsample_in_time=downsample_in_time\n )\n if return_features:\n features.append(sample)\n\n sample = checkpoint_fn(self.mid_block)(sample)\n\n # post-process\n sample = self.conv_norm_out(sample)\n sample = self.conv_act(sample)\n sample = self.conv_out(sample)\n\n if self.latent_log_var == \"uniform\":\n last_channel = sample[:, -1:, ...]\n num_dims = sample.dim()\n\n if num_dims == 4:\n # For shape (B, C, H, W)\n repeated_last_channel = last_channel.repeat(\n 1, sample.shape[1] - 2, 1, 1\n )\n sample = torch.cat([sample, repeated_last_channel], dim=1)\n elif num_dims == 5:\n # For shape (B, C, F, H, W)\n repeated_last_channel = last_channel.repeat(\n 1, sample.shape[1] - 2, 1, 1, 1\n )\n sample = torch.cat([sample, repeated_last_channel], dim=1)\n else:\n raise ValueError(f\"Invalid input shape: {sample.shape}\")\n\n if return_features:\n features.append(sample[:, : self.latent_channels, ...])\n return sample, features\n return sample\n\n\nclass Decoder(nn.Module):\n r\"\"\"\n The `Decoder` layer of a variational autoencoder that decodes its latent representation into an output sample.\n\n Args:\n in_channels (`int`, *optional*, defaults to 3):\n The number of input channels.\n out_channels (`int`, *optional*, defaults to 3):\n The number of output channels.\n block_out_channels (`Tuple[int, ...]`, *optional*, defaults to `(64,)`):\n The number of output channels for each block.\n layers_per_block (`int`, *optional*, defaults to 2):\n The number of layers per block.\n norm_num_groups (`int`, *optional*, defaults to 32):\n The number of groups for normalization.\n patch_size (`int`, *optional*, defaults to 1):\n The patch size to use. Should be a power of 2.\n norm_layer (`str`, *optional*, defaults to `group_norm`):\n The normalization layer to use. Can be either `group_norm` or `pixel_norm`.\n \"\"\"\n\n def __init__(\n self,\n dims,\n in_channels: int = 3,\n out_channels: int = 3,\n block_out_channels: Tuple[int, ...] = (64,),\n layers_per_block: int = 2,\n norm_num_groups: int = 32,\n patch_size: int = 1,\n norm_layer: str = \"group_norm\",\n patch_size_t: Optional[int] = None,\n add_channel_padding: Optional[bool] = False,\n ):\n super().__init__()\n self.patch_size = patch_size\n self.patch_size_t = patch_size_t if patch_size_t is not None else patch_size\n self.add_channel_padding = add_channel_padding\n self.layers_per_block = layers_per_block\n if add_channel_padding:\n out_channels = out_channels * self.patch_size**3\n else:\n out_channels = out_channels * self.patch_size_t * self.patch_size**2\n self.out_channels = out_channels\n\n self.conv_in = make_conv_nd(\n dims,\n in_channels,\n block_out_channels[-1],\n kernel_size=3,\n stride=1,\n padding=1,\n )\n\n self.mid_block = None\n self.up_blocks = nn.ModuleList([])\n\n self.mid_block = UNetMidBlock3D(\n dims=dims,\n in_channels=block_out_channels[-1],\n num_layers=self.layers_per_block,\n resnet_eps=1e-6,\n resnet_groups=norm_num_groups,\n norm_layer=norm_layer,\n )\n\n reversed_block_out_channels = list(reversed(block_out_channels))\n output_channel = reversed_block_out_channels[0]\n for i in range(len(reversed_block_out_channels)):\n prev_output_channel = output_channel\n output_channel = reversed_block_out_channels[i]\n\n is_final_block = i == len(block_out_channels) - 1\n\n up_block = UpDecoderBlock3D(\n dims=dims,\n num_layers=self.layers_per_block + 1,\n in_channels=prev_output_channel,\n out_channels=output_channel,\n add_upsample=not is_final_block\n and 2 ** (len(block_out_channels) - i - 1) > patch_size,\n resnet_eps=1e-6,\n resnet_groups=norm_num_groups,\n norm_layer=norm_layer,\n )\n self.up_blocks.append(up_block)\n\n if norm_layer == \"group_norm\":\n self.conv_norm_out = nn.GroupNorm(\n num_channels=block_out_channels[0], num_groups=norm_num_groups, eps=1e-6\n )\n elif norm_layer == \"pixel_norm\":\n self.conv_norm_out = PixelNorm()\n\n self.conv_act = nn.SiLU()\n self.conv_out = make_conv_nd(\n dims, block_out_channels[0], out_channels, 3, padding=1\n )\n\n self.gradient_checkpointing = False\n\n def forward(self, sample: torch.FloatTensor, target_shape) -> torch.FloatTensor:\n r\"\"\"The forward method of the `Decoder` class.\"\"\"\n assert target_shape is not None, \"target_shape must be provided\"\n upsample_in_time = sample.shape[2] < target_shape[2]\n\n sample = self.conv_in(sample)\n\n upscale_dtype = next(iter(self.up_blocks.parameters())).dtype\n\n checkpoint_fn = (\n partial(torch.utils.checkpoint.checkpoint, use_reentrant=False)\n if self.gradient_checkpointing and self.training\n else lambda x: x\n )\n\n sample = checkpoint_fn(self.mid_block)(sample)\n sample = sample.to(upscale_dtype)\n\n for up_block in self.up_blocks:\n sample = checkpoint_fn(up_block)(sample, upsample_in_time=upsample_in_time)\n\n # post-process\n sample = self.conv_norm_out(sample)\n sample = self.conv_act(sample)\n sample = self.conv_out(sample)\n\n # un-patchify\n patch_size_t = self.patch_size_t if upsample_in_time else 1\n sample = unpatchify(\n sample,\n patch_size_hw=self.patch_size,\n patch_size_t=patch_size_t,\n add_channel_padding=self.add_channel_padding,\n )\n\n return sample\n\n\nclass DownEncoderBlock3D(nn.Module):\n def __init__(\n self,\n dims: Union[int, Tuple[int, int]],\n in_channels: int,\n out_channels: int,\n dropout: float = 0.0,\n num_layers: int = 1,\n resnet_eps: float = 1e-6,\n resnet_groups: int = 32,\n add_downsample: bool = True,\n downsample_padding: int = 1,\n norm_layer: str = \"group_norm\",\n ):\n super().__init__()\n res_blocks = []\n\n for i in range(num_layers):\n in_channels = in_channels if i == 0 else out_channels\n res_blocks.append(\n ResnetBlock3D(\n dims=dims,\n in_channels=in_channels,\n out_channels=out_channels,\n eps=resnet_eps,\n groups=resnet_groups,\n dropout=dropout,\n norm_layer=norm_layer,\n )\n )\n\n self.res_blocks = nn.ModuleList(res_blocks)\n\n if add_downsample:\n self.downsample = Downsample3D(\n dims,\n out_channels,\n out_channels=out_channels,\n padding=downsample_padding,\n )\n else:\n self.downsample = Identity()\n\n def forward(\n self, hidden_states: torch.FloatTensor, downsample_in_time\n ) -> torch.FloatTensor:\n for resnet in self.res_blocks:\n hidden_states = resnet(hidden_states)\n\n hidden_states = self.downsample(\n hidden_states, downsample_in_time=downsample_in_time\n )\n\n return hidden_states\n\n\nclass UNetMidBlock3D(nn.Module):\n \"\"\"\n A 3D UNet mid-block [`UNetMidBlock3D`] with multiple residual blocks.\n\n Args:\n in_channels (`int`): The number of input channels.\n dropout (`float`, *optional*, defaults to 0.0): The dropout rate.\n num_layers (`int`, *optional*, defaults to 1): The number of residual blocks.\n resnet_eps (`float`, *optional*, 1e-6 ): The epsilon value for the resnet blocks.\n resnet_groups (`int`, *optional*, defaults to 32):\n The number of groups to use in the group normalization layers of the resnet blocks.\n\n Returns:\n `torch.FloatTensor`: The output of the last residual block, which is a tensor of shape `(batch_size,\n in_channels, height, width)`.\n\n \"\"\"\n\n def __init__(\n self,\n dims: Union[int, Tuple[int, int]],\n in_channels: int,\n dropout: float = 0.0,\n num_layers: int = 1,\n resnet_eps: float = 1e-6,\n resnet_groups: int = 32,\n norm_layer: str = \"group_norm\",\n ):\n super().__init__()\n resnet_groups = (\n resnet_groups if resnet_groups is not None else min(in_channels // 4, 32)\n )\n\n self.res_blocks = nn.ModuleList(\n [\n ResnetBlock3D(\n dims=dims,\n in_channels=in_channels,\n out_channels=in_channels,\n eps=resnet_eps,\n groups=resnet_groups,\n dropout=dropout,\n norm_layer=norm_layer,\n )\n for _ in range(num_layers)\n ]\n )\n\n def forward(self, hidden_states: torch.FloatTensor) -> torch.FloatTensor:\n for resnet in self.res_blocks:\n hidden_states = resnet(hidden_states)\n\n return hidden_states\n\n\nclass UpDecoderBlock3D(nn.Module):\n def __init__(\n self,\n dims: Union[int, Tuple[int, int]],\n in_channels: int,\n out_channels: int,\n resolution_idx: Optional[int] = None,\n dropout: float = 0.0,\n num_layers: int = 1,\n resnet_eps: float = 1e-6,\n resnet_groups: int = 32,\n add_upsample: bool = True,\n norm_layer: str = \"group_norm\",\n ):\n super().__init__()\n res_blocks = []\n\n for i in range(num_layers):\n input_channels = in_channels if i == 0 else out_channels\n\n res_blocks.append(\n ResnetBlock3D(\n dims=dims,\n in_channels=input_channels,\n out_channels=out_channels,\n eps=resnet_eps,\n groups=resnet_groups,\n dropout=dropout,\n norm_layer=norm_layer,\n )\n )\n\n self.res_blocks = nn.ModuleList(res_blocks)\n\n if add_upsample:\n self.upsample = Upsample3D(\n dims=dims, channels=out_channels, out_channels=out_channels\n )\n else:\n self.upsample = Identity()\n\n self.resolution_idx = resolution_idx\n\n def forward(\n self, hidden_states: torch.FloatTensor, upsample_in_time=True\n ) -> torch.FloatTensor:\n for resnet in self.res_blocks:\n hidden_states = resnet(hidden_states)\n\n hidden_states = self.upsample(hidden_states, upsample_in_time=upsample_in_time)\n\n return hidden_states\n\n\nclass ResnetBlock3D(nn.Module):\n r\"\"\"\n A Resnet block.\n\n Parameters:\n in_channels (`int`): The number of channels in the input.\n out_channels (`int`, *optional*, default to be `None`):\n The number of output channels for the first conv layer. If None, same as `in_channels`.\n dropout (`float`, *optional*, defaults to `0.0`): The dropout probability to use.\n groups (`int`, *optional*, default to `32`): The number of groups to use for the first normalization layer.\n eps (`float`, *optional*, defaults to `1e-6`): The epsilon to use for the normalization.\n \"\"\"\n\n def __init__(\n self,\n dims: Union[int, Tuple[int, int]],\n in_channels: int,\n out_channels: Optional[int] = None,\n conv_shortcut: bool = False,\n dropout: float = 0.0,\n groups: int = 32,\n eps: float = 1e-6,\n norm_layer: str = \"group_norm\",\n ):\n super().__init__()\n self.in_channels = in_channels\n out_channels = in_channels if out_channels is None else out_channels\n self.out_channels = out_channels\n self.use_conv_shortcut = conv_shortcut\n\n if norm_layer == \"group_norm\":\n self.norm1 = torch.nn.GroupNorm(\n num_groups=groups, num_channels=in_channels, eps=eps, affine=True\n )\n elif norm_layer == \"pixel_norm\":\n self.norm1 = PixelNorm()\n\n self.non_linearity = nn.SiLU()\n\n self.conv1 = make_conv_nd(\n dims, in_channels, out_channels, kernel_size=3, stride=1, padding=1\n )\n\n if norm_layer == \"group_norm\":\n self.norm2 = torch.nn.GroupNorm(\n num_groups=groups, num_channels=out_channels, eps=eps, affine=True\n )\n elif norm_layer == \"pixel_norm\":\n self.norm2 = PixelNorm()\n\n self.dropout = torch.nn.Dropout(dropout)\n\n self.conv2 = make_conv_nd(\n dims, out_channels, out_channels, kernel_size=3, stride=1, padding=1\n )\n\n self.conv_shortcut = (\n make_linear_nd(\n dims=dims, in_channels=in_channels, out_channels=out_channels\n )\n if in_channels != out_channels\n else nn.Identity()\n )\n\n def forward(\n self,\n input_tensor: torch.FloatTensor,\n ) -> torch.FloatTensor:\n hidden_states = input_tensor\n\n hidden_states = self.norm1(hidden_states)\n\n hidden_states = self.non_linearity(hidden_states)\n\n hidden_states = self.conv1(hidden_states)\n\n hidden_states = self.norm2(hidden_states)\n\n hidden_states = self.non_linearity(hidden_states)\n\n hidden_states = self.dropout(hidden_states)\n\n hidden_states = self.conv2(hidden_states)\n\n input_tensor = self.conv_shortcut(input_tensor)\n\n output_tensor = input_tensor + hidden_states\n\n return output_tensor\n\n\nclass Downsample3D(nn.Module):\n def __init__(\n self,\n dims,\n in_channels: int,\n out_channels: int,\n kernel_size: int = 3,\n padding: int = 1,\n ):\n super().__init__()\n stride: int = 2\n self.padding = padding\n self.in_channels = in_channels\n self.dims = dims\n self.conv = make_conv_nd(\n dims=dims,\n in_channels=in_channels,\n out_channels=out_channels,\n kernel_size=kernel_size,\n stride=stride,\n padding=padding,\n )\n\n def forward(self, x, downsample_in_time=True):\n conv = self.conv\n if self.padding == 0:\n if self.dims == 2:\n padding = (0, 1, 0, 1)\n else:\n padding = (0, 1, 0, 1, 0, 1 if downsample_in_time else 0)\n\n x = functional.pad(x, padding, mode=\"constant\", value=0)\n\n if self.dims == (2, 1) and not downsample_in_time:\n return conv(x, skip_time_conv=True)\n\n return conv(x)\n\n\nclass Upsample3D(nn.Module):\n \"\"\"\n An upsampling layer for 3D tensors of shape (B, C, D, H, W).\n\n :param channels: channels in the inputs and outputs.\n \"\"\"\n\n def __init__(self, dims, channels, out_channels=None):\n super().__init__()\n self.dims = dims\n self.channels = channels\n self.out_channels = out_channels or channels\n self.conv = make_conv_nd(\n dims, channels, out_channels, kernel_size=3, padding=1, bias=True\n )\n\n def forward(self, x, upsample_in_time):\n if self.dims == 2:\n x = functional.interpolate(\n x, (x.shape[2] * 2, x.shape[3] * 2), mode=\"nearest\"\n )\n else:\n time_scale_factor = 2 if upsample_in_time else 1\n # print(\"before:\", x.shape)\n b, c, d, h, w = x.shape\n x = rearrange(x, \"b c d h w -> (b d) c h w\")\n # height and width interpolate\n x = functional.interpolate(\n x, (x.shape[2] * 2, x.shape[3] * 2), mode=\"nearest\"\n )\n _, _, h, w = x.shape\n\n if not upsample_in_time and self.dims == (2, 1):\n x = rearrange(x, \"(b d) c h w -> b c d h w \", b=b, h=h, w=w)\n return self.conv(x, skip_time_conv=True)\n\n # Second ** upsampling ** which is essentially treated as a 1D convolution across the 'd' dimension\n x = rearrange(x, \"(b d) c h w -> (b h w) c 1 d\", b=b)\n\n # (b h w) c 1 d\n new_d = x.shape[-1] * time_scale_factor\n x = functional.interpolate(x, (1, new_d), mode=\"nearest\")\n # (b h w) c 1 new_d\n x = rearrange(\n x, \"(b h w) c 1 new_d -> b c new_d h w\", b=b, h=h, w=w, new_d=new_d\n )\n # b c d h w\n\n # x = functional.interpolate(\n # x, (x.shape[2] * time_scale_factor, x.shape[3] * 2, x.shape[4] * 2), mode=\"nearest\"\n # )\n # print(\"after:\", x.shape)\n\n return self.conv(x)\n\n\ndef patchify(x, patch_size_hw, patch_size_t=1, add_channel_padding=False):\n if patch_size_hw == 1 and patch_size_t == 1:\n return x\n if x.dim() == 4:\n x = rearrange(\n x, \"b c (h q) (w r) -> b (c r q) h w\", q=patch_size_hw, r=patch_size_hw\n )\n elif x.dim() == 5:\n x = rearrange(\n x,\n \"b c (f p) (h q) (w r) -> b (c p r q) f h w\",\n p=patch_size_t,\n q=patch_size_hw,\n r=patch_size_hw,\n )\n else:\n raise ValueError(f\"Invalid input shape: {x.shape}\")\n\n if (\n (x.dim() == 5)\n and (patch_size_hw > patch_size_t)\n and (patch_size_t > 1 or add_channel_padding)\n ):\n channels_to_pad = x.shape[1] * (patch_size_hw // patch_size_t) - x.shape[1]\n padding_zeros = torch.zeros(\n x.shape[0],\n channels_to_pad,\n x.shape[2],\n x.shape[3],\n x.shape[4],\n device=x.device,\n dtype=x.dtype,\n )\n x = torch.cat([padding_zeros, x], dim=1)\n\n return x\n\n\ndef unpatchify(x, patch_size_hw, patch_size_t=1, add_channel_padding=False):\n if patch_size_hw == 1 and patch_size_t == 1:\n return x\n\n if (\n (x.dim() == 5)\n and (patch_size_hw > patch_size_t)\n and (patch_size_t > 1 or add_channel_padding)\n ):\n channels_to_keep = int(x.shape[1] * (patch_size_t / patch_size_hw))\n x = x[:, :channels_to_keep, :, :, :]\n\n if x.dim() == 4:\n x = rearrange(\n x, \"b (c r q) h w -> b c (h q) (w r)\", q=patch_size_hw, r=patch_size_hw\n )\n elif x.dim() == 5:\n x = rearrange(\n x,\n \"b (c p r q) f h w -> b c (f p) (h q) (w r)\",\n p=patch_size_t,\n q=patch_size_hw,\n r=patch_size_hw,\n )\n\n return x\n\n\ndef create_video_autoencoder_config(\n latent_channels: int = 4,\n):\n config = {\n \"_class_name\": \"VideoAutoencoder\",\n \"dims\": (\n 2,\n 1,\n ), # 2 for Conv2, 3 for Conv3d, (2, 1) for Conv2d followed by Conv1d\n \"in_channels\": 3, # Number of input color channels (e.g., RGB)\n \"out_channels\": 3, # Number of output color channels\n \"latent_channels\": latent_channels, # Number of channels in the latent space representation\n \"block_out_channels\": [\n 128,\n 256,\n 512,\n 512,\n ], # Number of output channels of each encoder / decoder inner block\n \"patch_size\": 1,\n }\n\n return config\n\n\ndef create_video_autoencoder_pathify4x4x4_config(\n latent_channels: int = 4,\n):\n config = {\n \"_class_name\": \"VideoAutoencoder\",\n \"dims\": (\n 2,\n 1,\n ), # 2 for Conv2, 3 for Conv3d, (2, 1) for Conv2d followed by Conv1d\n \"in_channels\": 3, # Number of input color channels (e.g., RGB)\n \"out_channels\": 3, # Number of output color channels\n \"latent_channels\": latent_channels, # Number of channels in the latent space representation\n \"block_out_channels\": [512]\n * 4, # Number of output channels of each encoder / decoder inner block\n \"patch_size\": 4,\n \"latent_log_var\": \"uniform\",\n }\n\n return config\n\n\ndef create_video_autoencoder_pathify4x4_config(\n latent_channels: int = 4,\n):\n config = {\n \"_class_name\": \"VideoAutoencoder\",\n \"dims\": 2, # 2 for Conv2, 3 for Conv3d, (2, 1) for Conv2d followed by Conv1d\n \"in_channels\": 3, # Number of input color channels (e.g., RGB)\n \"out_channels\": 3, # Number of output color channels\n \"latent_channels\": latent_channels, # Number of channels in the latent space representation\n \"block_out_channels\": [512]\n * 4, # Number of output channels of each encoder / decoder inner block\n \"patch_size\": 4,\n \"norm_layer\": \"pixel_norm\",\n }\n\n return config\n\n\ndef test_vae_patchify_unpatchify():\n import torch\n\n x = torch.randn(2, 3, 8, 64, 64)\n x_patched = patchify(x, patch_size_hw=4, patch_size_t=4)\n x_unpatched = unpatchify(x_patched, patch_size_hw=4, patch_size_t=4)\n assert torch.allclose(x, x_unpatched)\n\n\ndef demo_video_autoencoder_forward_backward():\n # Configuration for the VideoAutoencoder\n config = create_video_autoencoder_pathify4x4x4_config()\n\n # Instantiate the VideoAutoencoder with the specified configuration\n video_autoencoder = VideoAutoencoder.from_config(config)\n\n print(video_autoencoder)\n\n # Print the total number of parameters in the video autoencoder\n total_params = sum(p.numel() for p in video_autoencoder.parameters())\n print(f\"Total number of parameters in VideoAutoencoder: {total_params:,}\")\n\n # Create a mock input tensor simulating a batch of videos\n # Shape: (batch_size, channels, depth, height, width)\n # E.g., 4 videos, each with 3 color channels, 16 frames, and 64x64 pixels per frame\n input_videos = torch.randn(2, 3, 8, 64, 64)\n\n # Forward pass: encode and decode the input videos\n latent = video_autoencoder.encode(input_videos).latent_dist.mode()\n print(f\"input shape={input_videos.shape}\")\n print(f\"latent shape={latent.shape}\")\n reconstructed_videos = video_autoencoder.decode(\n latent, target_shape=input_videos.shape\n ).sample\n\n print(f\"reconstructed shape={reconstructed_videos.shape}\")\n\n # Calculate the loss (e.g., mean squared error)\n loss = torch.nn.functional.mse_loss(input_videos, reconstructed_videos)\n\n # Perform backward pass\n loss.backward()\n\n print(f\"Demo completed with loss: {loss.item()}\")\n\n\n# Ensure to call the demo function to execute the forward and backward pass\nif __name__ == \"__main__\":\n demo_video_autoencoder_forward_backward()\n"
},
{
"path": "ltx_video/models/transformers/__init__.py",
"content": ""
},
{
"path": "ltx_video/models/transformers/attention.py",
"content": "import inspect\nfrom importlib import import_module\nfrom typing import Any, Dict, Optional, Tuple\n\nimport torch\nimport torch.nn.functional as F\nfrom diffusers.models.activations import GEGLU, GELU, ApproximateGELU\nfrom diffusers.models.attention import _chunked_feed_forward\nfrom diffusers.models.attention_processor import (\n LoRAAttnAddedKVProcessor,\n LoRAAttnProcessor,\n LoRAAttnProcessor2_0,\n LoRAXFormersAttnProcessor,\n SpatialNorm,\n)\nfrom diffusers.models.lora import LoRACompatibleLinear\nfrom diffusers.models.normalization import RMSNorm\nfrom diffusers.utils import deprecate, logging\nfrom diffusers.utils.torch_utils import maybe_allow_in_graph\nfrom einops import rearrange\nfrom torch import nn\n\nfrom ltx_video.utils.skip_layer_strategy import SkipLayerStrategy\n\ntry:\n from torch_xla.experimental.custom_kernel import flash_attention\nexcept ImportError:\n # workaround for automatic tests. Currently this function is manually patched\n # to the torch_xla lib on setup of container\n pass\n\n# code adapted from https://github.com/huggingface/diffusers/blob/main/src/diffusers/models/attention.py\n\nlogger = logging.get_logger(__name__)\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 qk_norm (`str`, *optional*, defaults to None):\n Set to 'layer_norm' or `rms_norm` to perform query and key normalization.\n adaptive_norm (`str`, *optional*, defaults to `\"single_scale_shift\"`):\n The type of adaptive norm to use. Can be `\"single_scale_shift\"`, `\"single_scale\"` or \"none\".\n standardization_norm (`str`, *optional*, defaults to `\"layer_norm\"`):\n The type of pre-normalization to use. Can be `\"layer_norm\"` or `\"rms_norm\"`.\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, # pylint: disable=unused-argument\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 adaptive_norm: str = \"single_scale_shift\", # 'single_scale_shift', 'single_scale' or 'none'\n standardization_norm: str = \"layer_norm\", # 'layer_norm' or 'rms_norm'\n norm_eps: float = 1e-5,\n qk_norm: Optional[str] = None,\n final_dropout: bool = False,\n attention_type: str = \"default\", # pylint: disable=unused-argument\n ff_inner_dim: Optional[int] = None,\n ff_bias: bool = True,\n attention_out_bias: bool = True,\n use_tpu_flash_attention: bool = False,\n use_rope: bool = False,\n ):\n super().__init__()\n self.only_cross_attention = only_cross_attention\n self.use_tpu_flash_attention = use_tpu_flash_attention\n self.adaptive_norm = adaptive_norm\n\n assert standardization_norm in [\"layer_norm\", \"rms_norm\"]\n assert adaptive_norm in [\"single_scale_shift\", \"single_scale\", \"none\"]\n\n make_norm_layer = (\n nn.LayerNorm if standardization_norm == \"layer_norm\" else RMSNorm\n )\n\n # Define 3 blocks. Each block has its own normalization layer.\n # 1. Self-Attn\n self.norm1 = make_norm_layer(\n dim, elementwise_affine=norm_elementwise_affine, eps=norm_eps\n )\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 use_tpu_flash_attention=use_tpu_flash_attention,\n qk_norm=qk_norm,\n use_rope=use_rope,\n )\n\n # 2. Cross-Attn\n if cross_attention_dim is not None or double_self_attention:\n self.attn2 = Attention(\n query_dim=dim,\n cross_attention_dim=(\n cross_attention_dim if not double_self_attention else None\n ),\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 use_tpu_flash_attention=use_tpu_flash_attention,\n qk_norm=qk_norm,\n use_rope=use_rope,\n ) # is self-attn if encoder_hidden_states is none\n\n if adaptive_norm == \"none\":\n self.attn2_norm = make_norm_layer(\n dim, norm_eps, norm_elementwise_affine\n )\n else:\n self.attn2 = None\n self.attn2_norm = None\n\n self.norm2 = make_norm_layer(dim, norm_eps, norm_elementwise_affine)\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 # 5. Scale-shift for PixArt-Alpha.\n if adaptive_norm != \"none\":\n num_ada_params = 4 if adaptive_norm == \"single_scale\" else 6\n self.scale_shift_table = nn.Parameter(\n torch.randn(num_ada_params, dim) / dim**0.5\n )\n\n # let chunk size default to None\n self._chunk_size = None\n self._chunk_dim = 0\n\n def set_use_tpu_flash_attention(self):\n r\"\"\"\n Function sets the flag in this object and propagates down the children. The flag will enforce the usage of TPU\n attention kernel.\n \"\"\"\n self.use_tpu_flash_attention = True\n self.attn1.set_use_tpu_flash_attention()\n self.attn2.set_use_tpu_flash_attention()\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.FloatTensor,\n freqs_cis: Optional[Tuple[torch.FloatTensor, torch.FloatTensor]] = None,\n attention_mask: Optional[torch.FloatTensor] = None,\n encoder_hidden_states: Optional[torch.FloatTensor] = None,\n encoder_attention_mask: Optional[torch.FloatTensor] = None,\n timestep: Optional[torch.LongTensor] = None,\n cross_attention_kwargs: Dict[str, Any] = None,\n class_labels: Optional[torch.LongTensor] = None,\n skip_layer_mask: Optional[torch.Tensor] = None,\n skip_layer_strategy: Optional[SkipLayerStrategy] = None,\n ) -> torch.FloatTensor:\n if cross_attention_kwargs is not None:\n if cross_attention_kwargs.get(\"scale\", None) is not None:\n logger.warning(\n \"Passing `scale` to `cross_attention_kwargs` is depcrecated. `scale` will be ignored.\"\n )\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 original_hidden_states = hidden_states\n\n norm_hidden_states = self.norm1(hidden_states)\n\n # Apply ada_norm_single\n if self.adaptive_norm in [\"single_scale_shift\", \"single_scale\"]:\n assert timestep.ndim == 3 # [batch, 1 or num_tokens, embedding_dim]\n num_ada_params = self.scale_shift_table.shape[0]\n ada_values = self.scale_shift_table[None, None] + timestep.reshape(\n batch_size, timestep.shape[1], num_ada_params, -1\n )\n if self.adaptive_norm == \"single_scale_shift\":\n shift_msa, scale_msa, gate_msa, shift_mlp, scale_mlp, gate_mlp = (\n ada_values.unbind(dim=2)\n )\n norm_hidden_states = norm_hidden_states * (1 + scale_msa) + shift_msa\n else:\n scale_msa, gate_msa, scale_mlp, gate_mlp = ada_values.unbind(dim=2)\n norm_hidden_states = norm_hidden_states * (1 + scale_msa)\n elif self.adaptive_norm == \"none\":\n scale_msa, gate_msa, scale_mlp, gate_mlp = None, None, None, None\n else:\n raise ValueError(f\"Unknown adaptive norm type: {self.adaptive_norm}\")\n\n norm_hidden_states = norm_hidden_states.squeeze(\n 1\n ) # TODO: Check if this is needed\n\n # 1. Prepare GLIGEN inputs\n cross_attention_kwargs = (\n cross_attention_kwargs.copy() if cross_attention_kwargs is not None else {}\n )\n\n attn_output = self.attn1(\n norm_hidden_states,\n freqs_cis=freqs_cis,\n encoder_hidden_states=(\n encoder_hidden_states if self.only_cross_attention else None\n ),\n attention_mask=attention_mask,\n skip_layer_mask=skip_layer_mask,\n skip_layer_strategy=skip_layer_strategy,\n **cross_attention_kwargs,\n )\n if gate_msa is not None:\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 # 3. Cross-Attention\n if self.attn2 is not None:\n if self.adaptive_norm == \"none\":\n attn_input = self.attn2_norm(hidden_states)\n else:\n attn_input = hidden_states\n attn_output = self.attn2(\n attn_input,\n freqs_cis=freqs_cis,\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 norm_hidden_states = self.norm2(hidden_states)\n if self.adaptive_norm == \"single_scale_shift\":\n norm_hidden_states = norm_hidden_states * (1 + scale_mlp) + shift_mlp\n elif self.adaptive_norm == \"single_scale\":\n norm_hidden_states = norm_hidden_states * (1 + scale_mlp)\n elif self.adaptive_norm == \"none\":\n pass\n else:\n raise ValueError(f\"Unknown adaptive norm type: {self.adaptive_norm}\")\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(\n self.ff, norm_hidden_states, self._chunk_dim, self._chunk_size\n )\n else:\n ff_output = self.ff(norm_hidden_states)\n if gate_mlp is not None:\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 if (\n skip_layer_mask is not None\n and skip_layer_strategy == SkipLayerStrategy.TransformerBlock\n ):\n skip_layer_mask = skip_layer_mask.view(-1, 1, 1)\n hidden_states = hidden_states * skip_layer_mask + original_hidden_states * (\n 1.0 - skip_layer_mask\n )\n\n return hidden_states\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 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 qk_norm (`str`, *optional*, defaults to None):\n Set to 'layer_norm' or `rms_norm` to perform query and key normalization.\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 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 added_kv_proj_dim: Optional[int] = None,\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 qk_norm: Optional[str] = None,\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 use_tpu_flash_attention: bool = False,\n use_rope: bool = False,\n ):\n super().__init__()\n self.inner_dim = out_dim if out_dim is not None else dim_head * 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 = (\n cross_attention_dim if cross_attention_dim is not None else query_dim\n )\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.use_tpu_flash_attention = use_tpu_flash_attention\n self.use_rope = use_rope\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 if qk_norm is None:\n self.q_norm = nn.Identity()\n self.k_norm = nn.Identity()\n elif qk_norm == \"rms_norm\":\n self.q_norm = RMSNorm(dim_head * heads, eps=1e-5)\n self.k_norm = RMSNorm(dim_head * heads, eps=1e-5)\n elif qk_norm == \"layer_norm\":\n self.q_norm = nn.LayerNorm(dim_head * heads, eps=1e-5)\n self.k_norm = nn.LayerNorm(dim_head * heads, eps=1e-5)\n else:\n raise ValueError(f\"Unsupported qk_norm method: {qk_norm}\")\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(\n num_channels=query_dim, num_groups=norm_num_groups, eps=eps, affine=True\n )\n else:\n self.group_norm = None\n\n if spatial_norm_dim is not None:\n self.spatial_norm = SpatialNorm(\n f_channels=query_dim, zq_channels=spatial_norm_dim\n )\n else:\n self.spatial_norm = None\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,\n num_groups=cross_attention_norm_num_groups,\n eps=1e-5,\n 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 linear_cls = nn.Linear\n\n self.linear_cls = linear_cls\n self.to_q = linear_cls(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 = linear_cls(self.cross_attention_dim, self.inner_dim, bias=bias)\n self.to_v = linear_cls(self.cross_attention_dim, self.inner_dim, bias=bias)\n else:\n self.to_k = None\n self.to_v = None\n\n if self.added_kv_proj_dim is not None:\n self.add_k_proj = linear_cls(added_kv_proj_dim, self.inner_dim)\n self.add_v_proj = linear_cls(added_kv_proj_dim, self.inner_dim)\n\n self.to_out = nn.ModuleList([])\n self.to_out.append(linear_cls(self.inner_dim, self.out_dim, bias=out_bias))\n self.to_out.append(nn.Dropout(dropout))\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 = AttnProcessor2_0()\n self.set_processor(processor)\n\n def set_use_tpu_flash_attention(self):\n r\"\"\"\n Function sets the flag in this object. The flag will enforce the usage of TPU attention kernel.\n \"\"\"\n self.use_tpu_flash_attention = True\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(\n f\"You are removing possibly trained weights of {self.processor} with {processor}\"\n )\n self._modules.pop(\"processor\")\n\n self.processor = processor\n\n def get_processor(\n self, return_deprecated_lora: bool = False\n ) -> \"AttentionProcessor\": # noqa: F821\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 # TODO(Sayak, Patrick). The rest of the function is needed to ensure backwards compatible\n # serialization format for LoRA Attention Processors. It should be deleted once the integration\n # with PEFT is completed.\n is_lora_activated = {\n name: module.lora_layer is not None\n for name, module in self.named_modules()\n if hasattr(module, \"lora_layer\")\n }\n\n # 1. if no layer has a LoRA activated we can return the processor as usual\n if not any(is_lora_activated.values()):\n return self.processor\n\n # If doesn't apply LoRA do `add_k_proj` or `add_v_proj`\n is_lora_activated.pop(\"add_k_proj\", None)\n is_lora_activated.pop(\"add_v_proj\", None)\n # 2. else it is not posssible that only some layers have LoRA activated\n if not all(is_lora_activated.values()):\n raise ValueError(\n f\"Make sure that either all layers or no layers have LoRA activated, but have {is_lora_activated}\"\n )\n\n # 3. And we need to merge the current LoRA layers into the corresponding LoRA attention processor\n non_lora_processor_cls_name = self.processor.__class__.__name__\n lora_processor_cls = getattr(\n import_module(__name__), \"LoRA\" + non_lora_processor_cls_name\n )\n\n hidden_size = self.inner_dim\n\n # now create a LoRA attention processor from the LoRA layers\n if lora_processor_cls in [\n LoRAAttnProcessor,\n LoRAAttnProcessor2_0,\n LoRAXFormersAttnProcessor,\n ]:\n kwargs = {\n \"cross_attention_dim\": self.cross_attention_dim,\n \"rank\": self.to_q.lora_layer.rank,\n \"network_alpha\": self.to_q.lora_layer.network_alpha,\n \"q_rank\": self.to_q.lora_layer.rank,\n \"q_hidden_size\": self.to_q.lora_layer.out_features,\n \"k_rank\": self.to_k.lora_layer.rank,\n \"k_hidden_size\": self.to_k.lora_layer.out_features,\n \"v_rank\": self.to_v.lora_layer.rank,\n \"v_hidden_size\": self.to_v.lora_layer.out_features,\n \"out_rank\": self.to_out[0].lora_layer.rank,\n \"out_hidden_size\": self.to_out[0].lora_layer.out_features,\n }\n\n if hasattr(self.processor, \"attention_op\"):\n kwargs[\"attention_op\"] = self.processor.attention_op\n\n lora_processor = lora_processor_cls(hidden_size, **kwargs)\n lora_processor.to_q_lora.load_state_dict(self.to_q.lora_layer.state_dict())\n lora_processor.to_k_lora.load_state_dict(self.to_k.lora_layer.state_dict())\n lora_processor.to_v_lora.load_state_dict(self.to_v.lora_layer.state_dict())\n lora_processor.to_out_lora.load_state_dict(\n self.to_out[0].lora_layer.state_dict()\n )\n elif lora_processor_cls == LoRAAttnAddedKVProcessor:\n lora_processor = lora_processor_cls(\n hidden_size,\n cross_attention_dim=self.add_k_proj.weight.shape[0],\n rank=self.to_q.lora_layer.rank,\n network_alpha=self.to_q.lora_layer.network_alpha,\n )\n lora_processor.to_q_lora.load_state_dict(self.to_q.lora_layer.state_dict())\n lora_processor.to_k_lora.load_state_dict(self.to_k.lora_layer.state_dict())\n lora_processor.to_v_lora.load_state_dict(self.to_v.lora_layer.state_dict())\n lora_processor.to_out_lora.load_state_dict(\n self.to_out[0].lora_layer.state_dict()\n )\n\n # only save if used\n if self.add_k_proj.lora_layer is not None:\n lora_processor.add_k_proj_lora.load_state_dict(\n self.add_k_proj.lora_layer.state_dict()\n )\n lora_processor.add_v_proj_lora.load_state_dict(\n self.add_v_proj.lora_layer.state_dict()\n )\n else:\n lora_processor.add_k_proj_lora = None\n lora_processor.add_v_proj_lora = None\n else:\n raise ValueError(f\"{lora_processor_cls} does not exist.\")\n\n return lora_processor\n\n def forward(\n self,\n hidden_states: torch.FloatTensor,\n freqs_cis: Optional[Tuple[torch.FloatTensor, torch.FloatTensor]] = None,\n encoder_hidden_states: Optional[torch.FloatTensor] = None,\n attention_mask: Optional[torch.FloatTensor] = None,\n skip_layer_mask: Optional[torch.Tensor] = None,\n skip_layer_strategy: Optional[SkipLayerStrategy] = 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 skip_layer_mask (`torch.Tensor`, *optional*):\n The skip layer mask to use. If `None`, no mask is applied.\n skip_layer_strategy (`SkipLayerStrategy`, *optional*, defaults to `None`):\n Controls which layers to skip for spatiotemporal guidance.\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(\n inspect.signature(self.processor.__call__).parameters.keys()\n )\n unused_kwargs = [\n k for k, _ in cross_attention_kwargs.items() if k not in attn_parameters\n ]\n if len(unused_kwargs) > 0:\n logger.warning(\n f\"cross_attention_kwargs {unused_kwargs} are not expected by\"\n f\" {self.processor.__class__.__name__} and will be ignored.\"\n )\n cross_attention_kwargs = {\n k: w for k, w in cross_attention_kwargs.items() if k in attn_parameters\n }\n\n return self.processor(\n self,\n hidden_states,\n freqs_cis=freqs_cis,\n encoder_hidden_states=encoder_hidden_states,\n attention_mask=attention_mask,\n skip_layer_mask=skip_layer_mask,\n skip_layer_strategy=skip_layer_strategy,\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(\n batch_size // head_size, seq_len, dim * head_size\n )\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\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(\n batch_size, seq_len * extra_dim, head_size, dim // head_size\n )\n tensor = tensor.permute(0, 2, 1, 3)\n\n if out_dim == 3:\n tensor = tensor.reshape(\n batch_size * head_size, seq_len * extra_dim, dim // head_size\n )\n\n return tensor\n\n def get_attention_scores(\n self,\n query: torch.Tensor,\n key: torch.Tensor,\n attention_mask: 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],\n query.shape[1],\n key.shape[1],\n dtype=query.dtype,\n 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,\n attention_mask: torch.Tensor,\n target_length: int,\n batch_size: int,\n 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 = (\n attention_mask.shape[0],\n attention_mask.shape[1],\n target_length,\n )\n padding = torch.zeros(\n padding_shape,\n dtype=attention_mask.dtype,\n device=attention_mask.device,\n )\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(\n self, encoder_hidden_states: torch.Tensor\n ) -> 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 (\n self.norm_cross is not None\n ), \"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 @staticmethod\n def apply_rotary_emb(\n input_tensor: torch.Tensor,\n freqs_cis: Tuple[torch.FloatTensor, torch.FloatTensor],\n ) -> Tuple[torch.Tensor, torch.Tensor]:\n cos_freqs = freqs_cis[0]\n sin_freqs = freqs_cis[1]\n\n t_dup = rearrange(input_tensor, \"... (d r) -> ... d r\", r=2)\n t1, t2 = t_dup.unbind(dim=-1)\n t_dup = torch.stack((-t2, t1), dim=-1)\n input_tensor_rot = rearrange(t_dup, \"... d r -> ... (d r)\")\n\n out = input_tensor * cos_freqs + input_tensor_rot * sin_freqs\n\n return out\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 pass\n\n def __call__(\n self,\n attn: Attention,\n hidden_states: torch.FloatTensor,\n freqs_cis: Tuple[torch.FloatTensor, torch.FloatTensor],\n encoder_hidden_states: Optional[torch.FloatTensor] = None,\n attention_mask: Optional[torch.FloatTensor] = None,\n temb: Optional[torch.FloatTensor] = None,\n skip_layer_mask: Optional[torch.FloatTensor] = None,\n skip_layer_strategy: Optional[SkipLayerStrategy] = None,\n *args,\n **kwargs,\n ) -> torch.FloatTensor:\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(\n batch_size, channel, height * width\n ).transpose(1, 2)\n\n batch_size, sequence_length, _ = (\n hidden_states.shape\n if encoder_hidden_states is None\n else encoder_hidden_states.shape\n )\n\n if skip_layer_mask is not None:\n skip_layer_mask = skip_layer_mask.reshape(batch_size, 1, 1)\n\n if (attention_mask is not None) and (not attn.use_tpu_flash_attention):\n attention_mask = attn.prepare_attention_mask(\n attention_mask, sequence_length, batch_size\n )\n # scaled_dot_product_attention expects attention_mask shape to be\n # (batch, heads, source_length, target_length)\n attention_mask = attention_mask.view(\n batch_size, attn.heads, -1, attention_mask.shape[-1]\n )\n\n if attn.group_norm is not None:\n hidden_states = attn.group_norm(hidden_states.transpose(1, 2)).transpose(\n 1, 2\n )\n\n query = attn.to_q(hidden_states)\n query = attn.q_norm(query)\n\n if encoder_hidden_states is not None:\n if attn.norm_cross:\n encoder_hidden_states = attn.norm_encoder_hidden_states(\n encoder_hidden_states\n )\n key = attn.to_k(encoder_hidden_states)\n key = attn.k_norm(key)\n else: # if no context provided do self-attention\n encoder_hidden_states = hidden_states\n key = attn.to_k(hidden_states)\n key = attn.k_norm(key)\n if attn.use_rope:\n key = attn.apply_rotary_emb(key, freqs_cis)\n query = attn.apply_rotary_emb(query, freqs_cis)\n\n value = attn.to_v(encoder_hidden_states)\n value_for_stg = value\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\n if attn.use_tpu_flash_attention: # use tpu attention offload 'flash attention'\n q_segment_indexes = None\n if (\n attention_mask is not None\n ): # if mask is required need to tune both segmenIds fields\n # attention_mask = torch.squeeze(attention_mask).to(torch.float32)\n attention_mask = attention_mask.to(torch.float32)\n q_segment_indexes = torch.ones(\n batch_size, query.shape[2], device=query.device, dtype=torch.float32\n )\n assert (\n attention_mask.shape[1] == key.shape[2]\n ), f\"ERROR: KEY SHAPE must be same as attention mask [{key.shape[2]}, {attention_mask.shape[1]}]\"\n\n assert (\n query.shape[2] % 128 == 0\n ), f\"ERROR: QUERY SHAPE must be divisible by 128 (TPU limitation) [{query.shape[2]}]\"\n assert (\n key.shape[2] % 128 == 0\n ), f\"ERROR: KEY SHAPE must be divisible by 128 (TPU limitation) [{key.shape[2]}]\"\n\n # run the TPU kernel implemented in jax with pallas\n hidden_states_a = flash_attention(\n q=query,\n k=key,\n v=value,\n q_segment_ids=q_segment_indexes,\n kv_segment_ids=attention_mask,\n sm_scale=attn.scale,\n )\n else:\n hidden_states_a = F.scaled_dot_product_attention(\n query,\n key,\n value,\n attn_mask=attention_mask,\n dropout_p=0.0,\n is_causal=False,\n )\n\n hidden_states_a = hidden_states_a.transpose(1, 2).reshape(\n batch_size, -1, attn.heads * head_dim\n )\n hidden_states_a = hidden_states_a.to(query.dtype)\n\n if (\n skip_layer_mask is not None\n and skip_layer_strategy == SkipLayerStrategy.AttentionSkip\n ):\n hidden_states = hidden_states_a * skip_layer_mask + hidden_states * (\n 1.0 - skip_layer_mask\n )\n elif (\n skip_layer_mask is not None\n and skip_layer_strategy == SkipLayerStrategy.AttentionValues\n ):\n hidden_states = hidden_states_a * skip_layer_mask + value_for_stg * (\n 1.0 - skip_layer_mask\n )\n else:\n hidden_states = hidden_states_a\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(\n batch_size, channel, height, width\n )\n if (\n skip_layer_mask is not None\n and skip_layer_strategy == SkipLayerStrategy.Residual\n ):\n skip_layer_mask = skip_layer_mask.reshape(batch_size, 1, 1, 1)\n\n if attn.residual_connection:\n if (\n skip_layer_mask is not None\n and skip_layer_strategy == SkipLayerStrategy.Residual\n ):\n hidden_states = hidden_states + residual * skip_layer_mask\n else:\n hidden_states = hidden_states + residual\n\n hidden_states = hidden_states / attn.rescale_output_factor\n\n return hidden_states\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.FloatTensor,\n encoder_hidden_states: Optional[torch.FloatTensor] = None,\n attention_mask: Optional[torch.FloatTensor] = None,\n temb: Optional[torch.FloatTensor] = 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(\n batch_size, channel, height * width\n ).transpose(1, 2)\n\n batch_size, sequence_length, _ = (\n hidden_states.shape\n if encoder_hidden_states is None\n else encoder_hidden_states.shape\n )\n attention_mask = attn.prepare_attention_mask(\n attention_mask, sequence_length, batch_size\n )\n\n if attn.group_norm is not None:\n hidden_states = attn.group_norm(hidden_states.transpose(1, 2)).transpose(\n 1, 2\n )\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(\n encoder_hidden_states\n )\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 query = attn.q_norm(query)\n key = attn.k_norm(key)\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(\n batch_size, channel, height, width\n )\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 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 linear_cls = nn.Linear\n\n if activation_fn == \"gelu\":\n act_fn = GELU(dim, inner_dim, bias=bias)\n elif 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 else:\n raise ValueError(f\"Unsupported activation function: {activation_fn}\")\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(linear_cls(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, scale: float = 1.0) -> torch.Tensor:\n compatible_cls = (GEGLU, LoRACompatibleLinear)\n for module in self.net:\n if isinstance(module, compatible_cls):\n hidden_states = module(hidden_states, scale)\n else:\n hidden_states = module(hidden_states)\n return hidden_states\n"
},
{
"path": "ltx_video/models/transformers/embeddings.py",
"content": "# Adapted from: https://github.com/huggingface/diffusers/blob/main/src/diffusers/models/embeddings.py\nimport math\n\nimport numpy as np\nimport torch\nfrom einops import rearrange\nfrom torch import nn\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 :param timesteps: a 1-D Tensor of N indices, one per batch element.\n These may be fractional.\n :param embedding_dim: the dimension of the output. :param max_period: controls the minimum frequency of the\n embeddings. :return: 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(embed_dim, grid, w, h, f):\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 grid = rearrange(grid, \"c (f h w) -> c f h w\", h=h, w=w)\n grid = rearrange(grid, \"c f h w -> c h w f\", h=h, w=w)\n grid = grid.reshape([3, 1, w, h, f])\n pos_embed = get_3d_sincos_pos_embed_from_grid(embed_dim, grid)\n pos_embed = pos_embed.transpose(1, 0, 2, 3)\n return rearrange(pos_embed, \"h w f c -> (f h w) c\")\n\n\ndef get_3d_sincos_pos_embed_from_grid(embed_dim, grid):\n if embed_dim % 3 != 0:\n raise ValueError(\"embed_dim must be divisible by 3\")\n\n # use half of dimensions to encode grid_h\n emb_f = get_1d_sincos_pos_embed_from_grid(embed_dim // 3, grid[0]) # (H*W*T, D/3)\n emb_h = get_1d_sincos_pos_embed_from_grid(embed_dim // 3, grid[1]) # (H*W*T, D/3)\n emb_w = get_1d_sincos_pos_embed_from_grid(embed_dim // 3, grid[2]) # (H*W*T, D/3)\n\n emb = np.concatenate([emb_h, emb_w, emb_f], axis=-1) # (H*W*T, 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_shape = pos.shape\n\n pos = pos.reshape(-1)\n out = np.einsum(\"m,d->md\", pos, omega) # (M, D/2), outer product\n out = out.reshape([*pos_shape, -1])[0]\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 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(\n torch.arange(0, embed_dim, 2) * (-math.log(10000.0) / embed_dim)\n )\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"
},
{
"path": "ltx_video/models/transformers/symmetric_patchifier.py",
"content": "from abc import ABC, abstractmethod\nfrom typing import Tuple\n\nimport torch\nfrom diffusers.configuration_utils import ConfigMixin\nfrom einops import rearrange\nfrom torch import Tensor\n\n\nclass Patchifier(ConfigMixin, ABC):\n def __init__(self, patch_size: int):\n super().__init__()\n self._patch_size = (1, patch_size, patch_size)\n\n @abstractmethod\n def patchify(self, latents: Tensor) -> Tuple[Tensor, Tensor]:\n raise NotImplementedError(\"Patchify method not implemented\")\n\n @abstractmethod\n def unpatchify(\n self,\n latents: Tensor,\n output_height: int,\n output_width: int,\n out_channels: int,\n ) -> Tuple[Tensor, Tensor]:\n pass\n\n @property\n def patch_size(self):\n return self._patch_size\n\n def get_latent_coords(\n self, latent_num_frames, latent_height, latent_width, batch_size, device\n ):\n \"\"\"\n Return a tensor of shape [batch_size, 3, num_patches] containing the\n top-left corner latent coordinates of each latent patch.\n The tensor is repeated for each batch element.\n \"\"\"\n latent_sample_coords = torch.meshgrid(\n torch.arange(0, latent_num_frames, self._patch_size[0], device=device),\n torch.arange(0, latent_height, self._patch_size[1], device=device),\n torch.arange(0, latent_width, self._patch_size[2], device=device),\n )\n latent_sample_coords = torch.stack(latent_sample_coords, dim=0)\n latent_coords = latent_sample_coords.unsqueeze(0).repeat(batch_size, 1, 1, 1, 1)\n latent_coords = rearrange(\n latent_coords, \"b c f h w -> b c (f h w)\", b=batch_size\n )\n return latent_coords\n\n\nclass SymmetricPatchifier(Patchifier):\n def patchify(self, latents: Tensor) -> Tuple[Tensor, Tensor]:\n b, _, f, h, w = latents.shape\n latent_coords = self.get_latent_coords(f, h, w, b, latents.device)\n latents = rearrange(\n latents,\n \"b c (f p1) (h p2) (w p3) -> b (f h w) (c p1 p2 p3)\",\n p1=self._patch_size[0],\n p2=self._patch_size[1],\n p3=self._patch_size[2],\n )\n return latents, latent_coords\n\n def unpatchify(\n self,\n latents: Tensor,\n output_height: int,\n output_width: int,\n out_channels: int,\n ) -> Tuple[Tensor, Tensor]:\n output_height = output_height // self._patch_size[1]\n output_width = output_width // self._patch_size[2]\n latents = rearrange(\n latents,\n \"b (f h w) (c p q) -> b c f (h p) (w q)\",\n h=output_height,\n w=output_width,\n p=self._patch_size[1],\n q=self._patch_size[2],\n )\n return latents\n"
},
{
"path": "ltx_video/models/transformers/transformer3d.py",
"content": "# Adapted from: https://github.com/huggingface/diffusers/blob/v0.26.3/src/diffusers/models/transformers/transformer_2d.py\nimport math\nfrom dataclasses import dataclass\nfrom typing import Any, Dict, List, Optional, Union\nimport os\nimport json\nimport glob\nfrom pathlib import Path\n\nimport torch\nfrom diffusers.configuration_utils import ConfigMixin, register_to_config\nfrom diffusers.models.embeddings import PixArtAlphaTextProjection\nfrom diffusers.models.modeling_utils import ModelMixin\nfrom diffusers.models.normalization import AdaLayerNormSingle\nfrom diffusers.utils import BaseOutput, is_torch_version\nfrom diffusers.utils import logging\nfrom torch import nn\nfrom safetensors import safe_open\n\n\nfrom ltx_video.models.transformers.attention import BasicTransformerBlock\nfrom ltx_video.utils.skip_layer_strategy import SkipLayerStrategy\n\nfrom ltx_video.utils.diffusers_config_mapping import (\n diffusers_and_ours_config_mapping,\n make_hashable_key,\n TRANSFORMER_KEYS_RENAME_DICT,\n)\n\n\nlogger = logging.get_logger(__name__)\n\n\n@dataclass\nclass Transformer3DModelOutput(BaseOutput):\n \"\"\"\n The output of [`Transformer2DModel`].\n\n Args:\n sample (`torch.FloatTensor` of shape `(batch_size, num_channels, height, width)` or `(batch size, num_vector_embeds - 1, num_latent_pixels)` if [`Transformer2DModel`] is discrete):\n The hidden states output conditioned on the `encoder_hidden_states` input. If discrete, returns probability\n distributions for the unnoised latent pixels.\n \"\"\"\n\n sample: torch.FloatTensor\n\n\nclass Transformer3DModel(ModelMixin, ConfigMixin):\n _supports_gradient_checkpointing = True\n\n @register_to_config\n def __init__(\n self,\n num_attention_heads: int = 16,\n attention_head_dim: int = 88,\n in_channels: Optional[int] = None,\n out_channels: Optional[int] = None,\n num_layers: int = 1,\n dropout: float = 0.0,\n norm_num_groups: int = 32,\n cross_attention_dim: Optional[int] = None,\n attention_bias: bool = False,\n num_vector_embeds: Optional[int] = None,\n activation_fn: str = \"geglu\",\n num_embeds_ada_norm: Optional[int] = None,\n use_linear_projection: bool = False,\n only_cross_attention: bool = False,\n double_self_attention: bool = False,\n upcast_attention: bool = False,\n adaptive_norm: str = \"single_scale_shift\", # 'single_scale_shift' or 'single_scale'\n standardization_norm: str = \"layer_norm\", # 'layer_norm' or 'rms_norm'\n norm_elementwise_affine: bool = True,\n norm_eps: float = 1e-5,\n attention_type: str = \"default\",\n caption_channels: int = None,\n use_tpu_flash_attention: bool = False, # if True uses the TPU attention offload ('flash attention')\n qk_norm: Optional[str] = None,\n positional_embedding_type: str = \"rope\",\n positional_embedding_theta: Optional[float] = None,\n positional_embedding_max_pos: Optional[List[int]] = None,\n timestep_scale_multiplier: Optional[float] = None,\n causal_temporal_positioning: bool = False, # For backward compatibility, will be deprecated\n ):\n super().__init__()\n self.use_tpu_flash_attention = (\n use_tpu_flash_attention # FIXME: push config down to the attention modules\n )\n self.use_linear_projection = use_linear_projection\n self.num_attention_heads = num_attention_heads\n self.attention_head_dim = attention_head_dim\n inner_dim = num_attention_heads * attention_head_dim\n self.inner_dim = inner_dim\n self.patchify_proj = nn.Linear(in_channels, inner_dim, bias=True)\n self.positional_embedding_type = positional_embedding_type\n self.positional_embedding_theta = positional_embedding_theta\n self.positional_embedding_max_pos = positional_embedding_max_pos\n self.use_rope = self.positional_embedding_type == \"rope\"\n self.timestep_scale_multiplier = timestep_scale_multiplier\n\n if self.positional_embedding_type == \"absolute\":\n raise ValueError(\"Absolute positional embedding is no longer supported\")\n elif self.positional_embedding_type == \"rope\":\n if positional_embedding_theta is None:\n raise ValueError(\n \"If `positional_embedding_type` type is rope, `positional_embedding_theta` must also be defined\"\n )\n if positional_embedding_max_pos is None:\n raise ValueError(\n \"If `positional_embedding_type` type is rope, `positional_embedding_max_pos` must also be defined\"\n )\n\n # 3. Define transformers blocks\n self.transformer_blocks = nn.ModuleList(\n [\n BasicTransformerBlock(\n inner_dim,\n num_attention_heads,\n attention_head_dim,\n dropout=dropout,\n cross_attention_dim=cross_attention_dim,\n activation_fn=activation_fn,\n num_embeds_ada_norm=num_embeds_ada_norm,\n attention_bias=attention_bias,\n only_cross_attention=only_cross_attention,\n double_self_attention=double_self_attention,\n upcast_attention=upcast_attention,\n adaptive_norm=adaptive_norm,\n standardization_norm=standardization_norm,\n norm_elementwise_affine=norm_elementwise_affine,\n norm_eps=norm_eps,\n attention_type=attention_type,\n use_tpu_flash_attention=use_tpu_flash_attention,\n qk_norm=qk_norm,\n use_rope=self.use_rope,\n )\n for d in range(num_layers)\n ]\n )\n\n # 4. Define output layers\n self.out_channels = in_channels if out_channels is None else out_channels\n self.norm_out = nn.LayerNorm(inner_dim, elementwise_affine=False, eps=1e-6)\n self.scale_shift_table = nn.Parameter(\n torch.randn(2, inner_dim) / inner_dim**0.5\n )\n self.proj_out = nn.Linear(inner_dim, self.out_channels)\n\n self.adaln_single = AdaLayerNormSingle(\n inner_dim, use_additional_conditions=False\n )\n if adaptive_norm == \"single_scale\":\n self.adaln_single.linear = nn.Linear(inner_dim, 4 * inner_dim, bias=True)\n\n self.caption_projection = None\n if caption_channels is not None:\n self.caption_projection = PixArtAlphaTextProjection(\n in_features=caption_channels, hidden_size=inner_dim\n )\n\n self.gradient_checkpointing = False\n\n def set_use_tpu_flash_attention(self):\n r\"\"\"\n Function sets the flag in this object and propagates down the children. The flag will enforce the usage of TPU\n attention kernel.\n \"\"\"\n logger.info(\"ENABLE TPU FLASH ATTENTION -> TRUE\")\n self.use_tpu_flash_attention = True\n # push config down to the attention modules\n for block in self.transformer_blocks:\n block.set_use_tpu_flash_attention()\n\n def create_skip_layer_mask(\n self,\n batch_size: int,\n num_conds: int,\n ptb_index: int,\n skip_block_list: Optional[List[int]] = None,\n ):\n if skip_block_list is None or len(skip_block_list) == 0:\n return None\n num_layers = len(self.transformer_blocks)\n mask = torch.ones(\n (num_layers, batch_size * num_conds), device=self.device, dtype=self.dtype\n )\n for block_idx in skip_block_list:\n mask[block_idx, ptb_index::num_conds] = 0\n return mask\n\n def _set_gradient_checkpointing(self, module, value=False):\n if hasattr(module, \"gradient_checkpointing\"):\n module.gradient_checkpointing = value\n\n def get_fractional_positions(self, indices_grid):\n fractional_positions = torch.stack(\n [\n indices_grid[:, i] / self.positional_embedding_max_pos[i]\n for i in range(3)\n ],\n dim=-1,\n )\n return fractional_positions\n\n def precompute_freqs_cis(self, indices_grid, spacing=\"exp\"):\n dtype = torch.float32 # We need full precision in the freqs_cis computation.\n dim = self.inner_dim\n theta = self.positional_embedding_theta\n\n fractional_positions = self.get_fractional_positions(indices_grid)\n\n start = 1\n end = theta\n device = fractional_positions.device\n if spacing == \"exp\":\n indices = theta ** (\n torch.linspace(\n math.log(start, theta),\n math.log(end, theta),\n dim // 6,\n device=device,\n dtype=dtype,\n )\n )\n indices = indices.to(dtype=dtype)\n elif spacing == \"exp_2\":\n indices = 1.0 / theta ** (torch.arange(0, dim, 6, device=device) / dim)\n indices = indices.to(dtype=dtype)\n elif spacing == \"linear\":\n indices = torch.linspace(start, end, dim // 6, device=device, dtype=dtype)\n elif spacing == \"sqrt\":\n indices = torch.linspace(\n start**2, end**2, dim // 6, device=device, dtype=dtype\n ).sqrt()\n\n indices = indices * math.pi / 2\n\n if spacing == \"exp_2\":\n freqs = (\n (indices * fractional_positions.unsqueeze(-1))\n .transpose(-1, -2)\n .flatten(2)\n )\n else:\n freqs = (\n (indices * (fractional_positions.unsqueeze(-1) * 2 - 1))\n .transpose(-1, -2)\n .flatten(2)\n )\n\n cos_freq = freqs.cos().repeat_interleave(2, dim=-1)\n sin_freq = freqs.sin().repeat_interleave(2, dim=-1)\n if dim % 6 != 0:\n cos_padding = torch.ones_like(cos_freq[:, :, : dim % 6])\n sin_padding = torch.zeros_like(cos_freq[:, :, : dim % 6])\n cos_freq = torch.cat([cos_padding, cos_freq], dim=-1)\n sin_freq = torch.cat([sin_padding, sin_freq], dim=-1)\n return cos_freq.to(self.dtype), sin_freq.to(self.dtype)\n\n def load_state_dict(\n self,\n state_dict: Dict,\n *args,\n **kwargs,\n ):\n if any([key.startswith(\"model.diffusion_model.\") for key in state_dict.keys()]):\n state_dict = {\n key.replace(\"model.diffusion_model.\", \"\"): value\n for key, value in state_dict.items()\n if key.startswith(\"model.diffusion_model.\")\n }\n super().load_state_dict(state_dict, *args, **kwargs)\n\n @classmethod\n def from_pretrained(\n cls,\n pretrained_model_path: Optional[Union[str, os.PathLike]],\n *args,\n **kwargs,\n ):\n pretrained_model_path = Path(pretrained_model_path)\n if pretrained_model_path.is_dir():\n config_path = pretrained_model_path / \"transformer\" / \"config.json\"\n with open(config_path, \"r\") as f:\n config = make_hashable_key(json.load(f))\n\n assert config in diffusers_and_ours_config_mapping, (\n \"Provided diffusers checkpoint config for transformer is not suppported. \"\n \"We only support diffusers configs found in Lightricks/LTX-Video.\"\n )\n\n config = diffusers_and_ours_config_mapping[config]\n state_dict = {}\n ckpt_paths = (\n pretrained_model_path\n / \"transformer\"\n / \"diffusion_pytorch_model*.safetensors\"\n )\n dict_list = glob.glob(str(ckpt_paths))\n for dict_path in dict_list:\n part_dict = {}\n with safe_open(dict_path, framework=\"pt\", device=\"cpu\") as f:\n for k in f.keys():\n part_dict[k] = f.get_tensor(k)\n state_dict.update(part_dict)\n\n for key in list(state_dict.keys()):\n new_key = key\n for replace_key, rename_key in TRANSFORMER_KEYS_RENAME_DICT.items():\n new_key = new_key.replace(replace_key, rename_key)\n state_dict[new_key] = state_dict.pop(key)\n\n with torch.device(\"meta\"):\n transformer = cls.from_config(config)\n transformer.load_state_dict(state_dict, assign=True, strict=True)\n elif pretrained_model_path.is_file() and str(pretrained_model_path).endswith(\n \".safetensors\"\n ):\n comfy_single_file_state_dict = {}\n with safe_open(pretrained_model_path, framework=\"pt\", device=\"cpu\") as f:\n metadata = f.metadata()\n for k in f.keys():\n comfy_single_file_state_dict[k] = f.get_tensor(k)\n configs = json.loads(metadata[\"config\"])\n transformer_config = configs[\"transformer\"]\n with torch.device(\"meta\"):\n transformer = Transformer3DModel.from_config(transformer_config)\n transformer.load_state_dict(comfy_single_file_state_dict, assign=True)\n return transformer\n\n def forward(\n self,\n hidden_states: torch.Tensor,\n indices_grid: torch.Tensor,\n encoder_hidden_states: Optional[torch.Tensor] = None,\n timestep: Optional[torch.LongTensor] = None,\n class_labels: Optional[torch.LongTensor] = None,\n cross_attention_kwargs: Dict[str, Any] = None,\n attention_mask: Optional[torch.Tensor] = None,\n encoder_attention_mask: Optional[torch.Tensor] = None,\n skip_layer_mask: Optional[torch.Tensor] = None,\n skip_layer_strategy: Optional[SkipLayerStrategy] = None,\n return_dict: bool = True,\n ):\n \"\"\"\n The [`Transformer2DModel`] forward method.\n\n Args:\n hidden_states (`torch.LongTensor` of shape `(batch size, num latent pixels)` if discrete, `torch.FloatTensor` of shape `(batch size, channel, height, width)` if continuous):\n Input `hidden_states`.\n indices_grid (`torch.LongTensor` of shape `(batch size, 3, num latent pixels)`):\n encoder_hidden_states ( `torch.FloatTensor` of shape `(batch size, sequence len, embed dims)`, *optional*):\n Conditional embeddings for cross attention layer. If not given, cross-attention defaults to\n self-attention.\n timestep ( `torch.LongTensor`, *optional*):\n Used to indicate denoising step. Optional timestep to be applied as an embedding in `AdaLayerNorm`.\n class_labels ( `torch.LongTensor` of shape `(batch size, num classes)`, *optional*):\n Used to indicate class labels conditioning. Optional class labels to be applied as an embedding in\n `AdaLayerZeroNorm`.\n cross_attention_kwargs ( `Dict[str, Any]`, *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 attention_mask ( `torch.Tensor`, *optional*):\n An attention mask of shape `(batch, key_tokens)` is applied to `encoder_hidden_states`. If `1` the mask\n is kept, otherwise if `0` it is discarded. Mask will be converted into a bias, which adds large\n negative values to the attention scores corresponding to \"discard\" tokens.\n encoder_attention_mask ( `torch.Tensor`, *optional*):\n Cross-attention mask applied to `encoder_hidden_states`. Two formats supported:\n\n * Mask `(batch, sequence_length)` True = keep, False = discard.\n * Bias `(batch, 1, sequence_length)` 0 = keep, -10000 = discard.\n\n If `ndim == 2`: will be interpreted as a mask, then converted into a bias consistent with the format\n above. This bias will be added to the cross-attention scores.\n skip_layer_mask ( `torch.Tensor`, *optional*):\n A mask of shape `(num_layers, batch)` that indicates which layers to skip. `0` at position\n `layer, batch_idx` indicates that the layer should be skipped for the corresponding batch index.\n skip_layer_strategy ( `SkipLayerStrategy`, *optional*, defaults to `None`):\n Controls which layers are skipped when calculating a perturbed latent for spatiotemporal guidance.\n return_dict (`bool`, *optional*, defaults to `True`):\n Whether or not to return a [`~models.unets.unet_2d_condition.UNet2DConditionOutput`] instead of a plain\n tuple.\n\n Returns:\n If `return_dict` is True, an [`~models.transformer_2d.Transformer2DModelOutput`] is returned, otherwise a\n `tuple` where the first element is the sample tensor.\n \"\"\"\n # for tpu attention offload 2d token masks are used. No need to transform.\n if not self.use_tpu_flash_attention:\n # ensure attention_mask is a bias, and give it a singleton query_tokens dimension.\n # we may have done this conversion already, e.g. if we came here via UNet2DConditionModel#forward.\n # we can tell by counting dims; if ndim == 2: it's a mask rather than a bias.\n # expects mask of shape:\n # [batch, key_tokens]\n # adds singleton query_tokens dimension:\n # [batch, 1, key_tokens]\n # this helps to broadcast it as a bias over attention scores, which will be in one of the following shapes:\n # [batch, heads, query_tokens, key_tokens] (e.g. torch sdp attn)\n # [batch * heads, query_tokens, key_tokens] (e.g. xformers or classic attn)\n if attention_mask is not None and attention_mask.ndim == 2:\n # assume that mask is expressed as:\n # (1 = keep, 0 = discard)\n # convert mask into a bias that can be added to attention scores:\n # (keep = +0, discard = -10000.0)\n attention_mask = (1 - attention_mask.to(hidden_states.dtype)) * -10000.0\n attention_mask = attention_mask.unsqueeze(1)\n\n # convert encoder_attention_mask to a bias the same way we do for attention_mask\n if encoder_attention_mask is not None and encoder_attention_mask.ndim == 2:\n encoder_attention_mask = (\n 1 - encoder_attention_mask.to(hidden_states.dtype)\n ) * -10000.0\n encoder_attention_mask = encoder_attention_mask.unsqueeze(1)\n\n # 1. Input\n hidden_states = self.patchify_proj(hidden_states)\n\n if self.timestep_scale_multiplier:\n timestep = self.timestep_scale_multiplier * timestep\n\n freqs_cis = self.precompute_freqs_cis(indices_grid)\n\n batch_size = hidden_states.shape[0]\n timestep, embedded_timestep = self.adaln_single(\n timestep.flatten(),\n {\"resolution\": None, \"aspect_ratio\": None},\n batch_size=batch_size,\n hidden_dtype=hidden_states.dtype,\n )\n # Second dimension is 1 or number of tokens (if timestep_per_token)\n timestep = timestep.view(batch_size, -1, timestep.shape[-1])\n embedded_timestep = embedded_timestep.view(\n batch_size, -1, embedded_timestep.shape[-1]\n )\n\n # 2. Blocks\n if self.caption_projection is not None:\n batch_size = hidden_states.shape[0]\n encoder_hidden_states = self.caption_projection(encoder_hidden_states)\n encoder_hidden_states = encoder_hidden_states.view(\n batch_size, -1, hidden_states.shape[-1]\n )\n\n for block_idx, block in enumerate(self.transformer_blocks):\n if self.training and self.gradient_checkpointing:\n\n def create_custom_forward(module, return_dict=None):\n def custom_forward(*inputs):\n if return_dict is not None:\n return module(*inputs, return_dict=return_dict)\n else:\n return module(*inputs)\n\n return custom_forward\n\n ckpt_kwargs: Dict[str, Any] = (\n {\"use_reentrant\": False} if is_torch_version(\">=\", \"1.11.0\") else {}\n )\n hidden_states = torch.utils.checkpoint.checkpoint(\n create_custom_forward(block),\n hidden_states,\n freqs_cis,\n attention_mask,\n encoder_hidden_states,\n encoder_attention_mask,\n timestep,\n cross_attention_kwargs,\n class_labels,\n (\n skip_layer_mask[block_idx]\n if skip_layer_mask is not None\n else None\n ),\n skip_layer_strategy,\n **ckpt_kwargs,\n )\n else:\n hidden_states = block(\n hidden_states,\n freqs_cis=freqs_cis,\n attention_mask=attention_mask,\n encoder_hidden_states=encoder_hidden_states,\n encoder_attention_mask=encoder_attention_mask,\n timestep=timestep,\n cross_attention_kwargs=cross_attention_kwargs,\n class_labels=class_labels,\n skip_layer_mask=(\n skip_layer_mask[block_idx]\n if skip_layer_mask is not None\n else None\n ),\n skip_layer_strategy=skip_layer_strategy,\n )\n\n # 3. Output\n scale_shift_values = (\n self.scale_shift_table[None, None] + embedded_timestep[:, :, None]\n )\n shift, scale = scale_shift_values[:, :, 0], scale_shift_values[:, :, 1]\n hidden_states = self.norm_out(hidden_states)\n # Modulation\n hidden_states = hidden_states * (1 + scale) + shift\n hidden_states = self.proj_out(hidden_states)\n if not return_dict:\n return (hidden_states,)\n\n return Transformer3DModelOutput(sample=hidden_states)\n"
},
{
"path": "ltx_video/pipelines/__init__.py",
"content": ""
},
{
"path": "ltx_video/pipelines/crf_compressor.py",
"content": "import av\nimport torch\nimport io\nimport numpy as np\n\n\ndef _encode_single_frame(output_file, image_array: np.ndarray, crf):\n container = av.open(output_file, \"w\", format=\"mp4\")\n try:\n stream = container.add_stream(\n \"libx264\", rate=1, options={\"crf\": str(crf), \"preset\": \"veryfast\"}\n )\n stream.height = image_array.shape[0]\n stream.width = image_array.shape[1]\n av_frame = av.VideoFrame.from_ndarray(image_array, format=\"rgb24\").reformat(\n format=\"yuv420p\"\n )\n container.mux(stream.encode(av_frame))\n container.mux(stream.encode())\n finally:\n container.close()\n\n\ndef _decode_single_frame(video_file):\n container = av.open(video_file)\n try:\n stream = next(s for s in container.streams if s.type == \"video\")\n frame = next(container.decode(stream))\n finally:\n container.close()\n return frame.to_ndarray(format=\"rgb24\")\n\n\ndef compress(image: torch.Tensor, crf=29):\n if crf == 0:\n return image\n\n image_array = (\n (image[: (image.shape[0] // 2) * 2, : (image.shape[1] // 2) * 2] * 255.0)\n .byte()\n .cpu()\n .numpy()\n )\n with io.BytesIO() as output_file:\n _encode_single_frame(output_file, image_array, crf)\n video_bytes = output_file.getvalue()\n with io.BytesIO(video_bytes) as video_file:\n image_array = _decode_single_frame(video_file)\n tensor = torch.tensor(image_array, dtype=image.dtype, device=image.device) / 255.0\n return tensor\n"
},
{
"path": "ltx_video/pipelines/pipeline_ltx_video.py",
"content": "# Adapted from: https://github.com/huggingface/diffusers/blob/main/src/diffusers/pipelines/pixart_alpha/pipeline_pixart_alpha.py\nimport copy\nimport inspect\nimport math\nimport re\nfrom contextlib import nullcontext\nfrom dataclasses import dataclass\nfrom typing import Any, Callable, Dict, List, Optional, Tuple, Union\n\nimport torch\nimport torch.nn.functional as F\nfrom diffusers.image_processor import VaeImageProcessor\nfrom diffusers.models import AutoencoderKL\nfrom diffusers.pipelines.pipeline_utils import DiffusionPipeline, ImagePipelineOutput\nfrom diffusers.schedulers import DPMSolverMultistepScheduler\nfrom diffusers.utils import deprecate, logging\nfrom diffusers.utils.torch_utils import randn_tensor\nfrom einops import rearrange\nfrom transformers import (\n T5EncoderModel,\n T5Tokenizer,\n AutoModelForCausalLM,\n AutoProcessor,\n AutoTokenizer,\n)\n\nfrom ltx_video.models.autoencoders.causal_video_autoencoder import (\n CausalVideoAutoencoder,\n)\nfrom ltx_video.models.autoencoders.vae_encode import (\n get_vae_size_scale_factor,\n latent_to_pixel_coords,\n vae_decode,\n vae_encode,\n)\nfrom ltx_video.models.transformers.symmetric_patchifier import Patchifier\nfrom ltx_video.models.transformers.transformer3d import Transformer3DModel\nfrom ltx_video.schedulers.rf import TimestepShifter\nfrom ltx_video.utils.skip_layer_strategy import SkipLayerStrategy\nfrom ltx_video.utils.prompt_enhance_utils import generate_cinematic_prompt\nfrom ltx_video.models.autoencoders.latent_upsampler import LatentUpsampler\nfrom ltx_video.models.autoencoders.vae_encode import (\n un_normalize_latents,\n normalize_latents,\n)\n\n\nlogger = logging.get_logger(__name__) # pylint: disable=invalid-name\n\n\nASPECT_RATIO_1024_BIN = {\n \"0.25\": [512.0, 2048.0],\n \"0.28\": [512.0, 1856.0],\n \"0.32\": [576.0, 1792.0],\n \"0.33\": [576.0, 1728.0],\n \"0.35\": [576.0, 1664.0],\n \"0.4\": [640.0, 1600.0],\n \"0.42\": [640.0, 1536.0],\n \"0.48\": [704.0, 1472.0],\n \"0.5\": [704.0, 1408.0],\n \"0.52\": [704.0, 1344.0],\n \"0.57\": [768.0, 1344.0],\n \"0.6\": [768.0, 1280.0],\n \"0.68\": [832.0, 1216.0],\n \"0.72\": [832.0, 1152.0],\n \"0.78\": [896.0, 1152.0],\n \"0.82\": [896.0, 1088.0],\n \"0.88\": [960.0, 1088.0],\n \"0.94\": [960.0, 1024.0],\n \"1.0\": [1024.0, 1024.0],\n \"1.07\": [1024.0, 960.0],\n \"1.13\": [1088.0, 960.0],\n \"1.21\": [1088.0, 896.0],\n \"1.29\": [1152.0, 896.0],\n \"1.38\": [1152.0, 832.0],\n \"1.46\": [1216.0, 832.0],\n \"1.67\": [1280.0, 768.0],\n \"1.75\": [1344.0, 768.0],\n \"2.0\": [1408.0, 704.0],\n \"2.09\": [1472.0, 704.0],\n \"2.4\": [1536.0, 640.0],\n \"2.5\": [1600.0, 640.0],\n \"3.0\": [1728.0, 576.0],\n \"4.0\": [2048.0, 512.0],\n}\n\nASPECT_RATIO_512_BIN = {\n \"0.25\": [256.0, 1024.0],\n \"0.28\": [256.0, 928.0],\n \"0.32\": [288.0, 896.0],\n \"0.33\": [288.0, 864.0],\n \"0.35\": [288.0, 832.0],\n \"0.4\": [320.0, 800.0],\n \"0.42\": [320.0, 768.0],\n \"0.48\": [352.0, 736.0],\n \"0.5\": [352.0, 704.0],\n \"0.52\": [352.0, 672.0],\n \"0.57\": [384.0, 672.0],\n \"0.6\": [384.0, 640.0],\n \"0.68\": [416.0, 608.0],\n \"0.72\": [416.0, 576.0],\n \"0.78\": [448.0, 576.0],\n \"0.82\": [448.0, 544.0],\n \"0.88\": [480.0, 544.0],\n \"0.94\": [480.0, 512.0],\n \"1.0\": [512.0, 512.0],\n \"1.07\": [512.0, 480.0],\n \"1.13\": [544.0, 480.0],\n \"1.21\": [544.0, 448.0],\n \"1.29\": [576.0, 448.0],\n \"1.38\": [576.0, 416.0],\n \"1.46\": [608.0, 416.0],\n \"1.67\": [640.0, 384.0],\n \"1.75\": [672.0, 384.0],\n \"2.0\": [704.0, 352.0],\n \"2.09\": [736.0, 352.0],\n \"2.4\": [768.0, 320.0],\n \"2.5\": [800.0, 320.0],\n \"3.0\": [864.0, 288.0],\n \"4.0\": [1024.0, 256.0],\n}\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 skip_initial_inference_steps: int = 0,\n skip_final_inference_steps: int = 0,\n **kwargs,\n):\n \"\"\"\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,\n `timesteps` 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 support arbitrary spacing between timesteps. If `None`, then the default\n timestep spacing strategy of the scheduler is used. If `timesteps` is passed, `num_inference_steps`\n must be `None`.\n max_timestep ('float', *optional*, defaults to 1.0):\n The initial noising level for image-to-image/video-to-video. The list if timestamps will be\n truncated to start with a timestamp greater or equal to this.\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:\n accepts_timesteps = \"timesteps\" in set(\n inspect.signature(scheduler.set_timesteps).parameters.keys()\n )\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 else:\n scheduler.set_timesteps(num_inference_steps, device=device, **kwargs)\n timesteps = scheduler.timesteps\n\n if (\n skip_initial_inference_steps < 0\n or skip_final_inference_steps < 0\n or skip_initial_inference_steps + skip_final_inference_steps\n >= num_inference_steps\n ):\n raise ValueError(\n \"invalid skip inference step values: must be non-negative and the sum of skip_initial_inference_steps and skip_final_inference_steps must be less than the number of inference steps\"\n )\n\n timesteps = timesteps[\n skip_initial_inference_steps : len(timesteps) - skip_final_inference_steps\n ]\n scheduler.set_timesteps(timesteps=timesteps, device=device, **kwargs)\n num_inference_steps = len(timesteps)\n\n return timesteps, num_inference_steps\n\n\n@dataclass\nclass ConditioningItem:\n \"\"\"\n Defines a single frame-conditioning item - a single frame or a sequence of frames.\n\n Attributes:\n media_item (torch.Tensor): shape=(b, 3, f, h, w). The media item to condition on.\n media_frame_number (int): The start-frame number of the media item in the generated video.\n conditioning_strength (float): The strength of the conditioning (1.0 = full conditioning).\n media_x (Optional[int]): Optional left x coordinate of the media item in the generated frame.\n media_y (Optional[int]): Optional top y coordinate of the media item in the generated frame.\n \"\"\"\n\n media_item: torch.Tensor\n media_frame_number: int\n conditioning_strength: float\n media_x: Optional[int] = None\n media_y: Optional[int] = None\n\n\nclass LTXVideoPipeline(DiffusionPipeline):\n r\"\"\"\n Pipeline for text-to-image generation using LTX-Video.\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 images to and from latent representations.\n text_encoder ([`T5EncoderModel`]):\n Frozen text-encoder. This 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 ([`Transformer2DModel`]):\n A text conditioned `Transformer2DModel` to denoise the encoded image latents.\n scheduler ([`SchedulerMixin`]):\n A scheduler to be used in combination with `transformer` to denoise the encoded image latents.\n \"\"\"\n\n bad_punct_regex = re.compile(\n r\"[\"\n + \"#ยฎโขยฉโข&@ยทยบยฝยพยฟยกยง~\"\n + r\"\\)\"\n + r\"\\(\"\n + r\"\\]\"\n + r\"\\[\"\n + r\"\\}\"\n + r\"\\{\"\n + r\"\\|\"\n + \"\\\\\"\n + r\"\\/\"\n + r\"\\*\"\n + r\"]{1,}\"\n ) # noqa\n\n _optional_components = [\n \"tokenizer\",\n \"text_encoder\",\n \"prompt_enhancer_image_caption_model\",\n \"prompt_enhancer_image_caption_processor\",\n \"prompt_enhancer_llm_model\",\n \"prompt_enhancer_llm_tokenizer\",\n ]\n model_cpu_offload_seq = \"prompt_enhancer_image_caption_model->prompt_enhancer_llm_model->text_encoder->transformer->vae\"\n\n def __init__(\n self,\n tokenizer: T5Tokenizer,\n text_encoder: T5EncoderModel,\n vae: AutoencoderKL,\n transformer: Transformer3DModel,\n scheduler: DPMSolverMultistepScheduler,\n patchifier: Patchifier,\n prompt_enhancer_image_caption_model: AutoModelForCausalLM,\n prompt_enhancer_image_caption_processor: AutoProcessor,\n prompt_enhancer_llm_model: AutoModelForCausalLM,\n prompt_enhancer_llm_tokenizer: AutoTokenizer,\n allowed_inference_steps: Optional[List[float]] = None,\n ):\n super().__init__()\n\n self.register_modules(\n tokenizer=tokenizer,\n text_encoder=text_encoder,\n vae=vae,\n transformer=transformer,\n scheduler=scheduler,\n patchifier=patchifier,\n prompt_enhancer_image_caption_model=prompt_enhancer_image_caption_model,\n prompt_enhancer_image_caption_processor=prompt_enhancer_image_caption_processor,\n prompt_enhancer_llm_model=prompt_enhancer_llm_model,\n prompt_enhancer_llm_tokenizer=prompt_enhancer_llm_tokenizer,\n )\n\n self.video_scale_factor, self.vae_scale_factor, _ = get_vae_size_scale_factor(\n self.vae\n )\n self.image_processor = VaeImageProcessor(vae_scale_factor=self.vae_scale_factor)\n\n self.allowed_inference_steps = allowed_inference_steps\n\n def mask_text_embeddings(self, emb, mask):\n if emb.shape[0] == 1:\n keep_index = mask.sum().item()\n return emb[:, :, :keep_index, :], keep_index\n else:\n masked_feature = emb * mask[:, None, :, None]\n return masked_feature, emb.shape[2]\n\n # Adapted from diffusers.pipelines.deepfloyd_if.pipeline_if.encode_prompt\n def encode_prompt(\n self,\n prompt: Union[str, List[str]],\n do_classifier_free_guidance: bool = True,\n negative_prompt: str = \"\",\n num_images_per_prompt: int = 1,\n device: Optional[torch.device] = None,\n prompt_embeds: Optional[torch.FloatTensor] = None,\n negative_prompt_embeds: Optional[torch.FloatTensor] = None,\n prompt_attention_mask: Optional[torch.FloatTensor] = None,\n negative_prompt_attention_mask: Optional[torch.FloatTensor] = None,\n text_encoder_max_tokens: int = 256,\n **kwargs,\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 not to guide the image generation. If not defined, one has to pass `negative_prompt_embeds`\n instead. Ignored when not using guidance (i.e., ignored if `guidance_scale` is less than `1`). For\n This should be \"\".\n do_classifier_free_guidance (`bool`, *optional*, defaults to `True`):\n whether to use classifier free guidance or not\n num_images_per_prompt (`int`, *optional*, defaults to 1):\n number of images that should be generated per prompt\n device: (`torch.device`, *optional*):\n torch device to place the resulting embeddings on\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.\n \"\"\"\n\n if \"mask_feature\" in kwargs:\n deprecation_message = \"The use of `mask_feature` is deprecated. It is no longer used in any computation and that doesn't affect the end results. It will be removed in a future version.\"\n deprecate(\"mask_feature\", \"1.0.0\", deprecation_message, standard_warn=False)\n\n if device is None:\n device = self._execution_device\n\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 # See Section 3.1. of the paper.\n max_length = (\n text_encoder_max_tokens # TPU supports only lengths multiple of 128\n )\n if prompt_embeds is None:\n assert (\n self.text_encoder is not None\n ), \"You should provide either prompt_embeds or self.text_encoder should not be None,\"\n text_enc_device = next(self.text_encoder.parameters()).device\n prompt = self._text_preprocessing(prompt)\n text_inputs = self.tokenizer(\n prompt,\n padding=\"max_length\",\n max_length=max_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(\n prompt, padding=\"longest\", return_tensors=\"pt\"\n ).input_ids\n\n if untruncated_ids.shape[-1] >= text_input_ids.shape[\n -1\n ] and not torch.equal(text_input_ids, untruncated_ids):\n removed_text = self.tokenizer.batch_decode(\n untruncated_ids[:, max_length - 1 : -1]\n )\n logger.warning(\n \"The following part of your input was truncated because CLIP can only handle sequences up to\"\n f\" {max_length} tokens: {removed_text}\"\n )\n\n prompt_attention_mask = text_inputs.attention_mask\n prompt_attention_mask = prompt_attention_mask.to(text_enc_device)\n prompt_attention_mask = prompt_attention_mask.to(device)\n\n prompt_embeds = self.text_encoder(\n text_input_ids.to(text_enc_device), attention_mask=prompt_attention_mask\n )\n prompt_embeds = prompt_embeds[0]\n\n if self.text_encoder is not None:\n dtype = self.text_encoder.dtype\n elif self.transformer is not None:\n dtype = self.transformer.dtype\n else:\n dtype = None\n\n prompt_embeds = prompt_embeds.to(dtype=dtype, device=device)\n\n bs_embed, seq_len, _ = prompt_embeds.shape\n # duplicate text embeddings and attention mask for each generation per prompt, using mps friendly method\n prompt_embeds = prompt_embeds.repeat(1, num_images_per_prompt, 1)\n prompt_embeds = prompt_embeds.view(\n bs_embed * num_images_per_prompt, seq_len, -1\n )\n prompt_attention_mask = prompt_attention_mask.repeat(1, num_images_per_prompt)\n prompt_attention_mask = prompt_attention_mask.view(\n bs_embed * num_images_per_prompt, -1\n )\n\n # get unconditional embeddings for classifier free guidance\n if do_classifier_free_guidance and negative_prompt_embeds is None:\n uncond_tokens = self._text_preprocessing(negative_prompt)\n uncond_tokens = uncond_tokens * batch_size\n max_length = prompt_embeds.shape[1]\n uncond_input = self.tokenizer(\n uncond_tokens,\n padding=\"max_length\",\n max_length=max_length,\n truncation=True,\n return_attention_mask=True,\n add_special_tokens=True,\n return_tensors=\"pt\",\n )\n negative_prompt_attention_mask = uncond_input.attention_mask\n negative_prompt_attention_mask = negative_prompt_attention_mask.to(\n text_enc_device\n )\n\n negative_prompt_embeds = self.text_encoder(\n uncond_input.input_ids.to(text_enc_device),\n attention_mask=negative_prompt_attention_mask,\n )\n negative_prompt_embeds = negative_prompt_embeds[0]\n\n if do_classifier_free_guidance:\n # duplicate unconditional embeddings for each generation per prompt, using mps friendly method\n seq_len = negative_prompt_embeds.shape[1]\n\n negative_prompt_embeds = negative_prompt_embeds.to(\n dtype=dtype, device=device\n )\n\n negative_prompt_embeds = negative_prompt_embeds.repeat(\n 1, num_images_per_prompt, 1\n )\n negative_prompt_embeds = negative_prompt_embeds.view(\n batch_size * num_images_per_prompt, seq_len, -1\n )\n\n negative_prompt_attention_mask = negative_prompt_attention_mask.repeat(\n 1, num_images_per_prompt\n )\n negative_prompt_attention_mask = negative_prompt_attention_mask.view(\n bs_embed * num_images_per_prompt, -1\n )\n else:\n negative_prompt_embeds = None\n negative_prompt_attention_mask = None\n\n return (\n prompt_embeds,\n prompt_attention_mask,\n negative_prompt_embeds,\n negative_prompt_attention_mask,\n )\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(\n inspect.signature(self.scheduler.step).parameters.keys()\n )\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(\n inspect.signature(self.scheduler.step).parameters.keys()\n )\n if accepts_generator:\n extra_step_kwargs[\"generator\"] = generator\n return extra_step_kwargs\n\n def check_inputs(\n self,\n prompt,\n height,\n width,\n negative_prompt,\n prompt_embeds=None,\n negative_prompt_embeds=None,\n prompt_attention_mask=None,\n negative_prompt_attention_mask=None,\n enhance_prompt=False,\n ):\n if height % 8 != 0 or width % 8 != 0:\n raise ValueError(\n f\"`height` and `width` have to be divisible by 8 but are {height} and {width}.\"\n )\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 (\n not isinstance(prompt, str) and not isinstance(prompt, list)\n ):\n raise ValueError(\n f\"`prompt` has to be of type `str` or `list` but is {type(prompt)}\"\n )\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 prompt_attention_mask is None:\n raise ValueError(\n \"Must provide `prompt_attention_mask` when specifying `prompt_embeds`.\"\n )\n\n if (\n negative_prompt_embeds is not None\n and negative_prompt_attention_mask is None\n ):\n raise ValueError(\n \"Must provide `negative_prompt_attention_mask` when specifying `negative_prompt_embeds`.\"\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 if prompt_attention_mask.shape != negative_prompt_attention_mask.shape:\n raise ValueError(\n \"`prompt_attention_mask` and `negative_prompt_attention_mask` must have the same shape when passed directly, but\"\n f\" got: `prompt_attention_mask` {prompt_attention_mask.shape} != `negative_prompt_attention_mask`\"\n f\" {negative_prompt_attention_mask.shape}.\"\n )\n\n if enhance_prompt:\n assert (\n self.prompt_enhancer_image_caption_model is not None\n ), \"Image caption model must be initialized if enhance_prompt is True\"\n assert (\n self.prompt_enhancer_image_caption_processor is not None\n ), \"Image caption processor must be initialized if enhance_prompt is True\"\n assert (\n self.prompt_enhancer_llm_model is not None\n ), \"Text prompt enhancer model must be initialized if enhance_prompt is True\"\n assert (\n self.prompt_enhancer_llm_tokenizer is not None\n ), \"Text prompt enhancer tokenizer must be initialized if enhance_prompt is True\"\n\n def _text_preprocessing(self, text):\n if not isinstance(text, (tuple, list)):\n text = [text]\n\n def process(text: str):\n text = text.strip()\n return text\n\n return [process(t) for t in text]\n\n @staticmethod\n def add_noise_to_image_conditioning_latents(\n t: float,\n init_latents: torch.Tensor,\n latents: torch.Tensor,\n noise_scale: float,\n conditioning_mask: torch.Tensor,\n generator,\n eps=1e-6,\n ):\n \"\"\"\n Add timestep-dependent noise to the hard-conditioning latents.\n This helps with motion continuity, especially when conditioned on a single frame.\n \"\"\"\n noise = randn_tensor(\n latents.shape,\n generator=generator,\n device=latents.device,\n dtype=latents.dtype,\n )\n # Add noise only to hard-conditioning latents (conditioning_mask = 1.0)\n need_to_noise = (conditioning_mask > 1.0 - eps).unsqueeze(-1)\n noised_latents = init_latents + noise_scale * noise * (t**2)\n latents = torch.where(need_to_noise, noised_latents, latents)\n return latents\n\n # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.StableDiffusionPipeline.prepare_latents\n def prepare_latents(\n self,\n latents: torch.Tensor | None,\n media_items: torch.Tensor | None,\n timestep: float,\n latent_shape: torch.Size | Tuple[Any, ...],\n dtype: torch.dtype,\n device: torch.device,\n generator: torch.Generator | List[torch.Generator],\n vae_per_channel_normalize: bool = True,\n ):\n \"\"\"\n Prepare the initial latent tensor to be denoised.\n The latents are either pure noise or a noised version of the encoded media items.\n Args:\n latents (`torch.FloatTensor` or `None`):\n The latents to use (provided by the user) or `None` to create new latents.\n media_items (`torch.FloatTensor` or `None`):\n An image or video to be updated using img2img or vid2vid. The media item is encoded and noised.\n timestep (`float`):\n The timestep to noise the encoded media_items to.\n latent_shape (`torch.Size`):\n The target latent shape.\n dtype (`torch.dtype`):\n The target dtype.\n device (`torch.device`):\n The target device.\n generator (`torch.Generator` or `List[torch.Generator]`):\n Generator(s) to be used for the noising process.\n vae_per_channel_normalize ('bool'):\n When encoding the media_items, whether to normalize the latents per-channel.\n Returns:\n `torch.FloatTensor`: The latents to be used for the denoising process. This is a tensor of shape\n (batch_size, num_channels, height, width).\n \"\"\"\n if isinstance(generator, list) and len(generator) != latent_shape[0]:\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 {latent_shape[0]}. Make sure the batch size matches the length of the generators.\"\n )\n\n # Initialize the latents with the given latents or encoded media item, if provided\n assert (\n latents is None or media_items is None\n ), \"Cannot provide both latents and media_items. Please provide only one of the two.\"\n\n assert (\n latents is None and media_items is None or timestep < 1.0\n ), \"Input media_item or latents are provided, but they will be replaced with noise.\"\n\n if media_items is not None:\n latents = vae_encode(\n media_items.to(dtype=self.vae.dtype, device=self.vae.device),\n self.vae,\n vae_per_channel_normalize=vae_per_channel_normalize,\n )\n if latents is not None:\n assert (\n latents.shape == latent_shape\n ), f\"Latents have to be of shape {latent_shape} but are {latents.shape}.\"\n latents = latents.to(device=device, dtype=dtype)\n\n # For backward compatibility, generate in the \"patchified\" shape and rearrange\n b, c, f, h, w = latent_shape\n noise = randn_tensor(\n (b, f * h * w, c), generator=generator, device=device, dtype=dtype\n )\n noise = rearrange(noise, \"b (f h w) c -> b c f h w\", f=f, h=h, w=w)\n\n # scale the initial noise by the standard deviation required by the scheduler\n noise = noise * self.scheduler.init_noise_sigma\n\n if latents is None:\n latents = noise\n else:\n # Noise the latents to the required (first) timestep\n latents = timestep * noise + (1 - timestep) * latents\n\n return latents\n\n @staticmethod\n def classify_height_width_bin(\n height: int, width: int, ratios: dict\n ) -> Tuple[int, int]:\n \"\"\"Returns binned height and width.\"\"\"\n ar = float(height / width)\n closest_ratio = min(ratios.keys(), key=lambda ratio: abs(float(ratio) - ar))\n default_hw = ratios[closest_ratio]\n return int(default_hw[0]), int(default_hw[1])\n\n @staticmethod\n def resize_and_crop_tensor(\n samples: torch.Tensor, new_width: int, new_height: int\n ) -> torch.Tensor:\n n_frames, orig_height, orig_width = samples.shape[-3:]\n\n # Check if resizing is needed\n if orig_height != new_height or orig_width != new_width:\n ratio = max(new_height / orig_height, new_width / orig_width)\n resized_width = int(orig_width * ratio)\n resized_height = int(orig_height * ratio)\n\n # Resize\n samples = LTXVideoPipeline.resize_tensor(\n samples, resized_height, resized_width\n )\n\n # Center Crop\n start_x = (resized_width - new_width) // 2\n end_x = start_x + new_width\n start_y = (resized_height - new_height) // 2\n end_y = start_y + new_height\n samples = samples[..., start_y:end_y, start_x:end_x]\n\n return samples\n\n @staticmethod\n def resize_tensor(media_items, height, width):\n n_frames = media_items.shape[2]\n if media_items.shape[-2:] != (height, width):\n media_items = rearrange(media_items, \"b c n h w -> (b n) c h w\")\n media_items = F.interpolate(\n media_items,\n size=(height, width),\n mode=\"bilinear\",\n align_corners=False,\n )\n media_items = rearrange(media_items, \"(b n) c h w -> b c n h w\", n=n_frames)\n return media_items\n\n @torch.no_grad()\n def __call__(\n self,\n height: int,\n width: int,\n num_frames: int,\n frame_rate: float,\n prompt: Union[str, List[str]] = None,\n negative_prompt: str = \"\",\n num_inference_steps: int = 20,\n skip_initial_inference_steps: int = 0,\n skip_final_inference_steps: int = 0,\n timesteps: List[int] = None,\n guidance_scale: Union[float, List[float]] = 4.5,\n cfg_star_rescale: bool = False,\n skip_layer_strategy: Optional[SkipLayerStrategy] = None,\n skip_block_list: Optional[Union[List[List[int]], List[int]]] = None,\n stg_scale: Union[float, List[float]] = 1.0,\n rescaling_scale: Union[float, List[float]] = 0.7,\n guidance_timesteps: Optional[List[int]] = None,\n num_images_per_prompt: Optional[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 prompt_attention_mask: Optional[torch.FloatTensor] = None,\n negative_prompt_embeds: Optional[torch.FloatTensor] = None,\n negative_prompt_attention_mask: Optional[torch.FloatTensor] = None,\n output_type: Optional[str] = \"pil\",\n return_dict: bool = True,\n callback_on_step_end: Optional[Callable[[int, int, Dict], None]] = None,\n conditioning_items: Optional[List[ConditioningItem]] = None,\n decode_timestep: Union[List[float], float] = 0.0,\n decode_noise_scale: Optional[List[float]] = None,\n mixed_precision: bool = False,\n offload_to_cpu: bool = False,\n enhance_prompt: bool = False,\n text_encoder_max_tokens: int = 256,\n stochastic_sampling: bool = False,\n media_items: Optional[torch.Tensor] = None,\n tone_map_compression_ratio: float = 0.0,\n **kwargs,\n ) -> Union[ImagePipelineOutput, 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 num_inference_steps (`int`, *optional*, defaults to 100):\n The number of denoising steps. More denoising steps usually lead to a higher quality image at the\n expense of slower inference. If `timesteps` is provided, this parameter is ignored.\n skip_initial_inference_steps (`int`, *optional*, defaults to 0):\n The number of initial timesteps to skip. After calculating the timesteps, this number of timesteps will\n be removed from the beginning of the timesteps list. Meaning the highest-timesteps values will not run.\n skip_final_inference_steps (`int`, *optional*, defaults to 0):\n The number of final timesteps to skip. After calculating the timesteps, this number of timesteps will\n be removed from the end of the timesteps list. Meaning the lowest-timesteps values will not run.\n timesteps (`List[int]`, *optional*):\n Custom timesteps to use for the denoising process. If not defined, equal spaced `num_inference_steps`\n timesteps are used. Must be in descending order.\n guidance_scale (`float`, *optional*, defaults to 4.5):\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 cfg_star_rescale (`bool`, *optional*, defaults to `False`):\n If set to `True`, applies the CFG star rescale. Scales the negative prediction according to dot\n product between positive and negative.\n num_images_per_prompt (`int`, *optional*, defaults to 1):\n The number of images to generate per prompt.\n height (`int`, *optional*, defaults to self.unet.config.sample_size):\n The height in pixels of the generated image.\n width (`int`, *optional*, defaults to self.unet.config.sample_size):\n The width in pixels of the generated image.\n eta (`float`, *optional*, defaults to 0.0):\n Corresponds to parameter eta (ฮท) in the DDIM paper: https://arxiv.org/abs/2010.02502. Only applies to\n [`schedulers.DDIMScheduler`], will be ignored for others.\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 prompt_attention_mask (`torch.FloatTensor`, *optional*): Pre-generated attention mask for text embeddings.\n negative_prompt_embeds (`torch.FloatTensor`, *optional*):\n Pre-generated negative text embeddings. This negative prompt should be \"\". If not\n provided, negative_prompt_embeds will be generated from `negative_prompt` input argument.\n negative_prompt_attention_mask (`torch.FloatTensor`, *optional*):\n Pre-generated attention mask for negative text embeddings.\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 to return a [`~pipelines.stable_diffusion.IFPipelineOutput`] instead of a plain tuple.\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 use_resolution_binning (`bool` defaults to `True`):\n If set to `True`, the requested height and width are first mapped to the closest resolutions using\n `ASPECT_RATIO_1024_BIN`. After the produced latents are decoded into images, they are resized back to\n the requested resolution. Useful for generating non-square images.\n enhance_prompt (`bool`, *optional*, defaults to `False`):\n If set to `True`, the prompt is enhanced using a LLM model.\n text_encoder_max_tokens (`int`, *optional*, defaults to `256`):\n The maximum number of tokens to use for the text encoder.\n stochastic_sampling (`bool`, *optional*, defaults to `False`):\n If set to `True`, the sampling is stochastic. If set to `False`, the sampling is deterministic.\n media_items ('torch.Tensor', *optional*):\n The input media item used for image-to-image / video-to-video.\n tone_map_compression_ratio: compression ratio for tone mapping, defaults to 0.0.\n If set to 0.0, no tone mapping is applied. If set to 1.0 - full compression is applied.\n Examples:\n\n Returns:\n [`~pipelines.ImagePipelineOutput`] or `tuple`:\n If `return_dict` is `True`, [`~pipelines.ImagePipelineOutput`] is returned, otherwise a `tuple` is\n returned where the first element is a list with the generated images\n \"\"\"\n if \"mask_feature\" in kwargs:\n deprecation_message = \"The use of `mask_feature` is deprecated. It is no longer used in any computation and that doesn't affect the end results. It will be removed in a future version.\"\n deprecate(\"mask_feature\", \"1.0.0\", deprecation_message, standard_warn=False)\n\n is_video = kwargs.get(\"is_video\", False)\n self.check_inputs(\n prompt,\n height,\n width,\n negative_prompt,\n prompt_embeds,\n negative_prompt_embeds,\n prompt_attention_mask,\n negative_prompt_attention_mask,\n )\n\n # 2. Default height and width to transformer\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 self.video_scale_factor = self.video_scale_factor if is_video else 1\n vae_per_channel_normalize = kwargs.get(\"vae_per_channel_normalize\", True)\n image_cond_noise_scale = kwargs.get(\"image_cond_noise_scale\", 0.0)\n\n latent_height = height // self.vae_scale_factor\n latent_width = width // self.vae_scale_factor\n latent_num_frames = num_frames // self.video_scale_factor\n if isinstance(self.vae, CausalVideoAutoencoder) and is_video:\n latent_num_frames += 1\n latent_shape = (\n batch_size * num_images_per_prompt,\n self.transformer.config.in_channels,\n latent_num_frames,\n latent_height,\n latent_width,\n )\n\n # Prepare the list of denoising time-steps\n\n retrieve_timesteps_kwargs = {}\n if isinstance(self.scheduler, TimestepShifter):\n retrieve_timesteps_kwargs[\"samples_shape\"] = latent_shape\n\n assert (\n skip_initial_inference_steps == 0\n or latents is not None\n or media_items is not None\n ), (\n f\"skip_initial_inference_steps ({skip_initial_inference_steps}) is used for image-to-image/video-to-video - \"\n \"media_item or latents should be provided.\"\n )\n\n timesteps, num_inference_steps = retrieve_timesteps(\n self.scheduler,\n num_inference_steps,\n device,\n timesteps,\n skip_initial_inference_steps=skip_initial_inference_steps,\n skip_final_inference_steps=skip_final_inference_steps,\n **retrieve_timesteps_kwargs,\n )\n\n if self.allowed_inference_steps is not None:\n for timestep in [round(x, 4) for x in timesteps.tolist()]:\n assert (\n timestep in self.allowed_inference_steps\n ), f\"Invalid inference timestep {timestep}. Allowed timesteps are {self.allowed_inference_steps}.\"\n\n if guidance_timesteps:\n guidance_mapping = []\n for timestep in timesteps:\n indices = [\n i for i, val in enumerate(guidance_timesteps) if val <= timestep\n ]\n # assert len(indices) > 0, f\"No guidance timestep found for {timestep}\"\n guidance_mapping.append(\n indices[0] if len(indices) > 0 else (len(guidance_timesteps) - 1)\n )\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 if not isinstance(guidance_scale, List):\n guidance_scale = [guidance_scale] * len(timesteps)\n else:\n guidance_scale = [\n guidance_scale[guidance_mapping[i]] for i in range(len(timesteps))\n ]\n\n if not isinstance(stg_scale, List):\n stg_scale = [stg_scale] * len(timesteps)\n else:\n stg_scale = [stg_scale[guidance_mapping[i]] for i in range(len(timesteps))]\n\n if not isinstance(rescaling_scale, List):\n rescaling_scale = [rescaling_scale] * len(timesteps)\n else:\n rescaling_scale = [\n rescaling_scale[guidance_mapping[i]] for i in range(len(timesteps))\n ]\n\n # Normalize skip_block_list to always be None or a list of lists matching timesteps\n if skip_block_list is not None:\n # Convert single list to list of lists if needed\n if len(skip_block_list) == 0 or not isinstance(skip_block_list[0], list):\n skip_block_list = [skip_block_list] * len(timesteps)\n else:\n new_skip_block_list = []\n for i, timestep in enumerate(timesteps):\n new_skip_block_list.append(skip_block_list[guidance_mapping[i]])\n skip_block_list = new_skip_block_list\n\n if enhance_prompt:\n self.prompt_enhancer_image_caption_model = (\n self.prompt_enhancer_image_caption_model.to(self._execution_device)\n )\n self.prompt_enhancer_llm_model = self.prompt_enhancer_llm_model.to(\n self._execution_device\n )\n\n prompt = generate_cinematic_prompt(\n self.prompt_enhancer_image_caption_model,\n self.prompt_enhancer_image_caption_processor,\n self.prompt_enhancer_llm_model,\n self.prompt_enhancer_llm_tokenizer,\n prompt,\n conditioning_items,\n max_new_tokens=text_encoder_max_tokens,\n )\n\n # 3. Encode input prompt\n if self.text_encoder is not None:\n self.text_encoder = self.text_encoder.to(self._execution_device)\n\n (\n prompt_embeds,\n prompt_attention_mask,\n negative_prompt_embeds,\n negative_prompt_attention_mask,\n ) = self.encode_prompt(\n prompt,\n True,\n negative_prompt=negative_prompt,\n num_images_per_prompt=num_images_per_prompt,\n device=device,\n prompt_embeds=prompt_embeds,\n negative_prompt_embeds=negative_prompt_embeds,\n prompt_attention_mask=prompt_attention_mask,\n negative_prompt_attention_mask=negative_prompt_attention_mask,\n text_encoder_max_tokens=text_encoder_max_tokens,\n )\n\n if offload_to_cpu and self.text_encoder is not None:\n self.text_encoder = self.text_encoder.cpu()\n\n self.transformer = self.transformer.to(self._execution_device)\n\n prompt_embeds_batch = prompt_embeds\n prompt_attention_mask_batch = prompt_attention_mask\n negative_prompt_embeds = (\n torch.zeros_like(prompt_embeds)\n if negative_prompt_embeds is None\n else negative_prompt_embeds\n )\n negative_prompt_attention_mask = (\n torch.zeros_like(prompt_attention_mask)\n if negative_prompt_attention_mask is None\n else negative_prompt_attention_mask\n )\n\n prompt_embeds_batch = torch.cat(\n [negative_prompt_embeds, prompt_embeds, prompt_embeds], dim=0\n )\n prompt_attention_mask_batch = torch.cat(\n [\n negative_prompt_attention_mask,\n prompt_attention_mask,\n prompt_attention_mask,\n ],\n dim=0,\n )\n # 4. Prepare the initial latents using the provided media and conditioning items\n\n # Prepare the initial latents tensor, shape = (b, c, f, h, w)\n latents = self.prepare_latents(\n latents=latents,\n media_items=media_items,\n timestep=timesteps[0],\n latent_shape=latent_shape,\n dtype=prompt_embeds.dtype,\n device=device,\n generator=generator,\n vae_per_channel_normalize=vae_per_channel_normalize,\n )\n\n # Update the latents with the conditioning items and patchify them into (b, n, c)\n latents, pixel_coords, conditioning_mask, num_cond_latents = (\n self.prepare_conditioning(\n conditioning_items=conditioning_items,\n init_latents=latents,\n num_frames=num_frames,\n height=height,\n width=width,\n vae_per_channel_normalize=vae_per_channel_normalize,\n generator=generator,\n )\n )\n init_latents = latents.clone() # Used for image_cond_noise_update\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. Denoising loop\n num_warmup_steps = max(\n len(timesteps) - num_inference_steps * self.scheduler.order, 0\n )\n\n orig_conditioning_mask = conditioning_mask\n\n # Befor compiling this code please be aware:\n # This code might generate different input shapes if some timesteps have no STG or CFG.\n # This means that the codes might need to be compiled mutliple times.\n # To avoid that, use the same STG and CFG values for all timesteps.\n\n with self.progress_bar(total=num_inference_steps) as progress_bar:\n for i, t in enumerate(timesteps):\n do_classifier_free_guidance = guidance_scale[i] > 1.0\n do_spatio_temporal_guidance = stg_scale[i] > 0\n do_rescaling = rescaling_scale[i] != 1.0\n\n num_conds = 1\n if do_classifier_free_guidance:\n num_conds += 1\n if do_spatio_temporal_guidance:\n num_conds += 1\n\n if do_classifier_free_guidance and do_spatio_temporal_guidance:\n indices = slice(batch_size * 0, batch_size * 3)\n elif do_classifier_free_guidance:\n indices = slice(batch_size * 0, batch_size * 2)\n elif do_spatio_temporal_guidance:\n indices = slice(batch_size * 1, batch_size * 3)\n else:\n indices = slice(batch_size * 1, batch_size * 2)\n\n # Prepare skip layer masks\n skip_layer_mask: Optional[torch.Tensor] = None\n if do_spatio_temporal_guidance:\n if skip_block_list is not None:\n skip_layer_mask = self.transformer.create_skip_layer_mask(\n batch_size, num_conds, num_conds - 1, skip_block_list[i]\n )\n\n batch_pixel_coords = torch.cat([pixel_coords] * num_conds)\n conditioning_mask = orig_conditioning_mask\n if conditioning_mask is not None and is_video:\n assert num_images_per_prompt == 1\n conditioning_mask = torch.cat([conditioning_mask] * num_conds)\n fractional_coords = batch_pixel_coords.to(torch.float32)\n fractional_coords[:, 0] = fractional_coords[:, 0] * (1.0 / frame_rate)\n\n if conditioning_mask is not None and image_cond_noise_scale > 0.0:\n latents = self.add_noise_to_image_conditioning_latents(\n t,\n init_latents,\n latents,\n image_cond_noise_scale,\n orig_conditioning_mask,\n generator,\n )\n\n latent_model_input = (\n torch.cat([latents] * num_conds) if num_conds > 1 else latents\n )\n latent_model_input = self.scheduler.scale_model_input(\n latent_model_input, t\n )\n\n current_timestep = t\n if not torch.is_tensor(current_timestep):\n # TODO: this requires sync between CPU and GPU. So try to pass timesteps as tensors if you can\n # This would be a good case for the `match` statement (Python 3.10+)\n is_mps = latent_model_input.device.type == \"mps\"\n if isinstance(current_timestep, float):\n dtype = torch.float32 if is_mps else torch.float64\n else:\n dtype = torch.int32 if is_mps else torch.int64\n current_timestep = torch.tensor(\n [current_timestep],\n dtype=dtype,\n device=latent_model_input.device,\n )\n elif len(current_timestep.shape) == 0:\n current_timestep = current_timestep[None].to(\n latent_model_input.device\n )\n # broadcast to batch dimension in a way that's compatible with ONNX/Core ML\n current_timestep = current_timestep.expand(\n latent_model_input.shape[0]\n ).unsqueeze(-1)\n\n if conditioning_mask is not None:\n # Conditioning latents have an initial timestep and noising level of (1.0 - conditioning_mask)\n # and will start to be denoised when the current timestep is lower than their conditioning timestep.\n current_timestep = torch.min(\n current_timestep, 1.0 - conditioning_mask\n )\n\n # Choose the appropriate context manager based on `mixed_precision`\n if mixed_precision:\n context_manager = torch.autocast(device.type, dtype=torch.bfloat16)\n else:\n context_manager = nullcontext() # Dummy context manager\n\n # predict noise model_output\n with context_manager:\n noise_pred = self.transformer(\n latent_model_input.to(self.transformer.dtype),\n indices_grid=fractional_coords,\n encoder_hidden_states=prompt_embeds_batch[indices].to(\n self.transformer.dtype\n ),\n encoder_attention_mask=prompt_attention_mask_batch[indices],\n timestep=current_timestep,\n skip_layer_mask=skip_layer_mask,\n skip_layer_strategy=skip_layer_strategy,\n return_dict=False,\n )[0]\n\n # perform guidance\n if do_spatio_temporal_guidance:\n noise_pred_text, noise_pred_text_perturb = noise_pred.chunk(\n num_conds\n )[-2:]\n if do_classifier_free_guidance:\n noise_pred_uncond, noise_pred_text = noise_pred.chunk(num_conds)[:2]\n\n if cfg_star_rescale:\n # Rescales the unconditional noise prediction using the projection of the conditional prediction onto it:\n # ฮฑ = (โจฮต_text, ฮต_uncondโฉ / ||ฮต_uncond||ยฒ), then ฮต_uncond โ ฮฑ * ฮต_uncond\n # where ฮต_text is the conditional noise prediction and ฮต_uncond is the unconditional one.\n positive_flat = noise_pred_text.view(batch_size, -1)\n negative_flat = noise_pred_uncond.view(batch_size, -1)\n dot_product = torch.sum(\n positive_flat * negative_flat, dim=1, keepdim=True\n )\n squared_norm = (\n torch.sum(negative_flat**2, dim=1, keepdim=True) + 1e-8\n )\n alpha = dot_product / squared_norm\n noise_pred_uncond = alpha * noise_pred_uncond\n\n noise_pred = noise_pred_uncond + guidance_scale[i] * (\n noise_pred_text - noise_pred_uncond\n )\n elif do_spatio_temporal_guidance:\n noise_pred = noise_pred_text\n if do_spatio_temporal_guidance:\n noise_pred = noise_pred + stg_scale[i] * (\n noise_pred_text - noise_pred_text_perturb\n )\n if do_rescaling and stg_scale[i] > 0.0:\n noise_pred_text_std = noise_pred_text.view(batch_size, -1).std(\n dim=1, keepdim=True\n )\n noise_pred_std = noise_pred.view(batch_size, -1).std(\n dim=1, keepdim=True\n )\n\n factor = noise_pred_text_std / noise_pred_std\n factor = rescaling_scale[i] * factor + (1 - rescaling_scale[i])\n\n noise_pred = noise_pred * factor.view(batch_size, 1, 1)\n\n current_timestep = current_timestep[:1]\n # learned sigma\n if (\n self.transformer.config.out_channels // 2\n == self.transformer.config.in_channels\n ):\n noise_pred = noise_pred.chunk(2, dim=1)[0]\n\n # compute previous image: x_t -> x_t-1\n latents = self.denoising_step(\n latents,\n noise_pred,\n current_timestep,\n orig_conditioning_mask,\n t,\n extra_step_kwargs,\n stochastic_sampling=stochastic_sampling,\n )\n\n # call the callback, if provided\n if i == len(timesteps) - 1 or (\n (i + 1) > num_warmup_steps and (i + 1) % self.scheduler.order == 0\n ):\n progress_bar.update()\n\n if callback_on_step_end is not None:\n callback_on_step_end(self, i, t, {})\n\n if offload_to_cpu:\n self.transformer = self.transformer.cpu()\n if self._execution_device == \"cuda\":\n torch.cuda.empty_cache()\n\n # Remove the added conditioning latents\n latents = latents[:, num_cond_latents:]\n\n latents = self.patchifier.unpatchify(\n latents=latents,\n output_height=latent_height,\n output_width=latent_width,\n out_channels=self.transformer.in_channels\n // math.prod(self.patchifier.patch_size),\n )\n if output_type != \"latent\":\n if self.vae.decoder.timestep_conditioning:\n noise = torch.randn_like(latents)\n if not isinstance(decode_timestep, list):\n decode_timestep = [decode_timestep] * latents.shape[0]\n if decode_noise_scale is None:\n decode_noise_scale = decode_timestep\n elif not isinstance(decode_noise_scale, list):\n decode_noise_scale = [decode_noise_scale] * latents.shape[0]\n\n decode_timestep = torch.tensor(decode_timestep).to(latents.device)\n decode_noise_scale = torch.tensor(decode_noise_scale).to(\n latents.device\n )[:, None, None, None, None]\n latents = (\n latents * (1 - decode_noise_scale) + noise * decode_noise_scale\n )\n else:\n decode_timestep = None\n latents = self.tone_map_latents(latents, tone_map_compression_ratio)\n image = vae_decode(\n latents,\n self.vae,\n is_video,\n vae_per_channel_normalize=kwargs[\"vae_per_channel_normalize\"],\n timestep=decode_timestep,\n )\n\n image = self.image_processor.postprocess(image, output_type=output_type)\n\n else:\n image = latents\n\n # Offload all models\n self.maybe_free_model_hooks()\n\n if not return_dict:\n return (image,)\n\n return ImagePipelineOutput(images=image)\n\n def denoising_step(\n self,\n latents: torch.Tensor,\n noise_pred: torch.Tensor,\n current_timestep: torch.Tensor,\n conditioning_mask: torch.Tensor,\n t: float,\n extra_step_kwargs,\n t_eps=1e-6,\n stochastic_sampling=False,\n ):\n \"\"\"\n Perform the denoising step for the required tokens, based on the current timestep and\n conditioning mask:\n Conditioning latents have an initial timestep and noising level of (1.0 - conditioning_mask)\n and will start to be denoised when the current timestep is equal or lower than their\n conditioning timestep.\n (hard-conditioning latents with conditioning_mask = 1.0 are never denoised)\n \"\"\"\n # Denoise the latents using the scheduler\n denoised_latents = self.scheduler.step(\n noise_pred,\n t if current_timestep is None else current_timestep,\n latents,\n **extra_step_kwargs,\n return_dict=False,\n stochastic_sampling=stochastic_sampling,\n )[0]\n\n if conditioning_mask is None:\n return denoised_latents\n\n tokens_to_denoise_mask = (t - t_eps < (1.0 - conditioning_mask)).unsqueeze(-1)\n return torch.where(tokens_to_denoise_mask, denoised_latents, latents)\n\n def prepare_conditioning(\n self,\n conditioning_items: Optional[List[ConditioningItem]],\n init_latents: torch.Tensor,\n num_frames: int,\n height: int,\n width: int,\n vae_per_channel_normalize: bool = False,\n generator=None,\n ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, int]:\n \"\"\"\n Prepare conditioning tokens based on the provided conditioning items.\n\n This method encodes provided conditioning items (video frames or single frames) into latents\n and integrates them with the initial latent tensor. It also calculates corresponding pixel\n coordinates, a mask indicating the influence of conditioning latents, and the total number of\n conditioning latents.\n\n Args:\n conditioning_items (Optional[List[ConditioningItem]]): A list of ConditioningItem objects.\n init_latents (torch.Tensor): The initial latent tensor of shape (b, c, f_l, h_l, w_l), where\n `f_l` is the number of latent frames, and `h_l` and `w_l` are latent spatial dimensions.\n num_frames, height, width: The dimensions of the generated video.\n vae_per_channel_normalize (bool, optional): Whether to normalize channels during VAE encoding.\n Defaults to `False`.\n generator: The random generator\n\n Returns:\n Tuple[torch.Tensor, torch.Tensor, torch.Tensor, int]:\n - `init_latents` (torch.Tensor): The updated latent tensor including conditioning latents,\n patchified into (b, n, c) shape.\n - `init_pixel_coords` (torch.Tensor): The pixel coordinates corresponding to the updated\n latent tensor.\n - `conditioning_mask` (torch.Tensor): A mask indicating the conditioning-strength of each\n latent token.\n - `num_cond_latents` (int): The total number of latent tokens added from conditioning items.\n\n Raises:\n AssertionError: If input shapes, dimensions, or conditions for applying conditioning are invalid.\n \"\"\"\n assert isinstance(self.vae, CausalVideoAutoencoder)\n\n if conditioning_items:\n batch_size, _, num_latent_frames = init_latents.shape[:3]\n\n init_conditioning_mask = torch.zeros(\n init_latents[:, 0, :, :, :].shape,\n dtype=torch.float32,\n device=init_latents.device,\n )\n\n extra_conditioning_latents = []\n extra_conditioning_pixel_coords = []\n extra_conditioning_mask = []\n extra_conditioning_num_latents = 0 # Number of extra conditioning latents added (should be removed before decoding)\n\n # Process each conditioning item\n for conditioning_item in conditioning_items:\n conditioning_item = self._resize_conditioning_item(\n conditioning_item, height, width\n )\n media_item = conditioning_item.media_item\n media_frame_number = conditioning_item.media_frame_number\n strength = conditioning_item.conditioning_strength\n assert media_item.ndim == 5 # (b, c, f, h, w)\n b, c, n_frames, h, w = media_item.shape\n assert (\n height == h and width == w\n ) or media_frame_number == 0, f\"Dimensions do not match: {height}x{width} != {h}x{w} - allowed only when media_frame_number == 0\"\n assert n_frames % 8 == 1\n assert (\n media_frame_number >= 0\n and media_frame_number + n_frames <= num_frames\n )\n\n # Encode the provided conditioning media item\n media_item_latents = vae_encode(\n media_item.to(dtype=self.vae.dtype, device=self.vae.device),\n self.vae,\n vae_per_channel_normalize=vae_per_channel_normalize,\n ).to(dtype=init_latents.dtype)\n\n # Handle the different conditioning cases\n if media_frame_number == 0:\n # Get the target spatial position of the latent conditioning item\n media_item_latents, l_x, l_y = self._get_latent_spatial_position(\n media_item_latents,\n conditioning_item,\n height,\n width,\n strip_latent_border=True,\n )\n b, c_l, f_l, h_l, w_l = media_item_latents.shape\n\n # First frame or sequence - just update the initial noise latents and the mask\n init_latents[:, :, :f_l, l_y : l_y + h_l, l_x : l_x + w_l] = (\n torch.lerp(\n init_latents[:, :, :f_l, l_y : l_y + h_l, l_x : l_x + w_l],\n media_item_latents,\n strength,\n )\n )\n init_conditioning_mask[\n :, :f_l, l_y : l_y + h_l, l_x : l_x + w_l\n ] = strength\n else:\n # Non-first frame or sequence\n if n_frames > 1:\n # Handle non-first sequence.\n # Encoded latents are either fully consumed, or the prefix is handled separately below.\n (\n init_latents,\n init_conditioning_mask,\n media_item_latents,\n ) = self._handle_non_first_conditioning_sequence(\n init_latents,\n init_conditioning_mask,\n media_item_latents,\n media_frame_number,\n strength,\n )\n\n # Single frame or sequence-prefix latents\n if media_item_latents is not None:\n noise = randn_tensor(\n media_item_latents.shape,\n generator=generator,\n device=media_item_latents.device,\n dtype=media_item_latents.dtype,\n )\n\n media_item_latents = torch.lerp(\n noise, media_item_latents, strength\n )\n\n # Patchify the extra conditioning latents and calculate their pixel coordinates\n media_item_latents, latent_coords = self.patchifier.patchify(\n latents=media_item_latents\n )\n pixel_coords = latent_to_pixel_coords(\n latent_coords,\n self.vae,\n causal_fix=self.transformer.config.causal_temporal_positioning,\n )\n\n # Update the frame numbers to match the target frame number\n pixel_coords[:, 0] += media_frame_number\n extra_conditioning_num_latents += media_item_latents.shape[1]\n\n conditioning_mask = torch.full(\n media_item_latents.shape[:2],\n strength,\n dtype=torch.float32,\n device=init_latents.device,\n )\n\n extra_conditioning_latents.append(media_item_latents)\n extra_conditioning_pixel_coords.append(pixel_coords)\n extra_conditioning_mask.append(conditioning_mask)\n\n # Patchify the updated latents and calculate their pixel coordinates\n init_latents, init_latent_coords = self.patchifier.patchify(\n latents=init_latents\n )\n init_pixel_coords = latent_to_pixel_coords(\n init_latent_coords,\n self.vae,\n causal_fix=self.transformer.config.causal_temporal_positioning,\n )\n\n if not conditioning_items:\n return init_latents, init_pixel_coords, None, 0\n\n init_conditioning_mask, _ = self.patchifier.patchify(\n latents=init_conditioning_mask.unsqueeze(1)\n )\n init_conditioning_mask = init_conditioning_mask.squeeze(-1)\n\n if extra_conditioning_latents:\n # Stack the extra conditioning latents, pixel coordinates and mask\n init_latents = torch.cat([*extra_conditioning_latents, init_latents], dim=1)\n init_pixel_coords = torch.cat(\n [*extra_conditioning_pixel_coords, init_pixel_coords], dim=2\n )\n init_conditioning_mask = torch.cat(\n [*extra_conditioning_mask, init_conditioning_mask], dim=1\n )\n\n if self.transformer.use_tpu_flash_attention:\n # When flash attention is used, keep the original number of tokens by removing\n # tokens from the end.\n init_latents = init_latents[:, :-extra_conditioning_num_latents]\n init_pixel_coords = init_pixel_coords[\n :, :, :-extra_conditioning_num_latents\n ]\n init_conditioning_mask = init_conditioning_mask[\n :, :-extra_conditioning_num_latents\n ]\n\n return (\n init_latents,\n init_pixel_coords,\n init_conditioning_mask,\n extra_conditioning_num_latents,\n )\n\n @staticmethod\n def _resize_conditioning_item(\n conditioning_item: ConditioningItem,\n height: int,\n width: int,\n ):\n if conditioning_item.media_x or conditioning_item.media_y:\n raise ValueError(\n \"Provide media_item in the target size for spatial conditioning.\"\n )\n new_conditioning_item = copy.copy(conditioning_item)\n new_conditioning_item.media_item = LTXVideoPipeline.resize_tensor(\n conditioning_item.media_item, height, width\n )\n return new_conditioning_item\n\n def _get_latent_spatial_position(\n self,\n latents: torch.Tensor,\n conditioning_item: ConditioningItem,\n height: int,\n width: int,\n strip_latent_border,\n ):\n \"\"\"\n Get the spatial position of the conditioning item in the latent space.\n If requested, strip the conditioning latent borders that do not align with target borders.\n (border latents look different then other latents and might confuse the model)\n \"\"\"\n scale = self.vae_scale_factor\n h, w = conditioning_item.media_item.shape[-2:]\n assert (\n h <= height and w <= width\n ), f\"Conditioning item size {h}x{w} is larger than target size {height}x{width}\"\n assert h % scale == 0 and w % scale == 0\n\n # Compute the start and end spatial positions of the media item\n x_start, y_start = conditioning_item.media_x, conditioning_item.media_y\n x_start = (width - w) // 2 if x_start is None else x_start\n y_start = (height - h) // 2 if y_start is None else y_start\n x_end, y_end = x_start + w, y_start + h\n assert (\n x_end <= width and y_end <= height\n ), f\"Conditioning item {x_start}:{x_end}x{y_start}:{y_end} is out of bounds for target size {width}x{height}\"\n\n if strip_latent_border:\n # Strip one latent from left/right and/or top/bottom, update x, y accordingly\n if x_start > 0:\n x_start += scale\n latents = latents[:, :, :, :, 1:]\n\n if y_start > 0:\n y_start += scale\n latents = latents[:, :, :, 1:, :]\n\n if x_end < width:\n latents = latents[:, :, :, :, :-1]\n\n if y_end < height:\n latents = latents[:, :, :, :-1, :]\n\n return latents, x_start // scale, y_start // scale\n\n @staticmethod\n def _handle_non_first_conditioning_sequence(\n init_latents: torch.Tensor,\n init_conditioning_mask: torch.Tensor,\n latents: torch.Tensor,\n media_frame_number: int,\n strength: float,\n num_prefix_latent_frames: int = 2,\n prefix_latents_mode: str = \"concat\",\n prefix_soft_conditioning_strength: float = 0.15,\n ):\n \"\"\"\n Special handling for a conditioning sequence that does not start on the first frame.\n The special handling is required to allow a short encoded video to be used as middle\n (or last) sequence in a longer video.\n Args:\n init_latents (torch.Tensor): The initial noise latents to be updated.\n init_conditioning_mask (torch.Tensor): The initial conditioning mask to be updated.\n latents (torch.Tensor): The encoded conditioning item.\n media_frame_number (int): The target frame number of the first frame in the conditioning sequence.\n strength (float): The conditioning strength for the conditioning latents.\n num_prefix_latent_frames (int, optional): The length of the sequence prefix, to be handled\n separately. Defaults to 2.\n prefix_latents_mode (str, optional): Special treatment for prefix (boundary) latents.\n - \"drop\": Drop the prefix latents.\n - \"soft\": Use the prefix latents, but with soft-conditioning\n - \"concat\": Add the prefix latents as extra tokens (like single frames)\n prefix_soft_conditioning_strength (float, optional): The strength of the soft-conditioning for\n the prefix latents, relevant if `prefix_latents_mode` is \"soft\". Defaults to 0.1.\n\n \"\"\"\n f_l = latents.shape[2]\n f_l_p = num_prefix_latent_frames\n assert f_l >= f_l_p\n assert media_frame_number % 8 == 0\n if f_l > f_l_p:\n # Insert the conditioning latents **excluding the prefix** into the sequence\n f_l_start = media_frame_number // 8 + f_l_p\n f_l_end = f_l_start + f_l - f_l_p\n init_latents[:, :, f_l_start:f_l_end] = torch.lerp(\n init_latents[:, :, f_l_start:f_l_end],\n latents[:, :, f_l_p:],\n strength,\n )\n # Mark these latent frames as conditioning latents\n init_conditioning_mask[:, f_l_start:f_l_end] = strength\n\n # Handle the prefix-latents\n if prefix_latents_mode == \"soft\":\n if f_l_p > 1:\n # Drop the first (single-frame) latent and soft-condition the remaining prefix\n f_l_start = media_frame_number // 8 + 1\n f_l_end = f_l_start + f_l_p - 1\n strength = min(prefix_soft_conditioning_strength, strength)\n init_latents[:, :, f_l_start:f_l_end] = torch.lerp(\n init_latents[:, :, f_l_start:f_l_end],\n latents[:, :, 1:f_l_p],\n strength,\n )\n # Mark these latent frames as conditioning latents\n init_conditioning_mask[:, f_l_start:f_l_end] = strength\n latents = None # No more latents to handle\n elif prefix_latents_mode == \"drop\":\n # Drop the prefix latents\n latents = None\n elif prefix_latents_mode == \"concat\":\n # Pass-on the prefix latents to be handled as extra conditioning frames\n latents = latents[:, :, :f_l_p]\n else:\n raise ValueError(f\"Invalid prefix_latents_mode: {prefix_latents_mode}\")\n return (\n init_latents,\n init_conditioning_mask,\n latents,\n )\n\n def trim_conditioning_sequence(\n self, start_frame: int, sequence_num_frames: int, target_num_frames: int\n ):\n \"\"\"\n Trim a conditioning sequence to the allowed number of frames.\n\n Args:\n start_frame (int): The target frame number of the first frame in the sequence.\n sequence_num_frames (int): The number of frames in the sequence.\n target_num_frames (int): The target number of frames in the generated video.\n\n Returns:\n int: updated sequence length\n \"\"\"\n scale_factor = self.video_scale_factor\n num_frames = min(sequence_num_frames, target_num_frames - start_frame)\n # Trim down to a multiple of temporal_scale_factor frames plus 1\n num_frames = (num_frames - 1) // scale_factor * scale_factor + 1\n return num_frames\n\n @staticmethod\n def tone_map_latents(\n latents: torch.Tensor,\n compression: float,\n ) -> torch.Tensor:\n \"\"\"\n Applies a non-linear tone-mapping function to latent values to reduce their dynamic range\n in a perceptually smooth way using a sigmoid-based compression.\n\n This is useful for regularizing high-variance latents or for conditioning outputs\n during generation, especially when controlling dynamic behavior with a `compression` factor.\n\n Parameters:\n ----------\n latents : torch.Tensor\n Input latent tensor with arbitrary shape. Expected to be roughly in [-1, 1] or [0, 1] range.\n compression : float\n Compression strength in the range [0, 1].\n - 0.0: No tone-mapping (identity transform)\n - 1.0: Full compression effect\n\n Returns:\n -------\n torch.Tensor\n The tone-mapped latent tensor of the same shape as input.\n \"\"\"\n if not (0 <= compression <= 1):\n raise ValueError(\"Compression must be in the range [0, 1]\")\n\n # Remap [0-1] to [0-0.75] and apply sigmoid compression in one shot\n scale_factor = compression * 0.75\n abs_latents = torch.abs(latents)\n\n # Sigmoid compression: sigmoid shifts large values toward 0.2, small values stay ~1.0\n # When scale_factor=0, sigmoid term vanishes, when scale_factor=0.75, full effect\n sigmoid_term = torch.sigmoid(4.0 * scale_factor * (abs_latents - 1.0))\n scales = 1.0 - 0.8 * scale_factor * sigmoid_term\n\n filtered = latents * scales\n return filtered\n\n\ndef adain_filter_latent(\n latents: torch.Tensor, reference_latents: torch.Tensor, factor=1.0\n):\n \"\"\"\n Applies Adaptive Instance Normalization (AdaIN) to a latent tensor based on\n statistics from a reference latent tensor.\n\n Args:\n latent (torch.Tensor): Input latents to normalize\n reference_latent (torch.Tensor): The reference latents providing style statistics.\n factor (float): Blending factor between original and transformed latent.\n Range: -10.0 to 10.0, Default: 1.0\n\n Returns:\n torch.Tensor: The transformed latent tensor\n \"\"\"\n result = latents.clone()\n\n for i in range(latents.size(0)):\n for c in range(latents.size(1)):\n r_sd, r_mean = torch.std_mean(\n reference_latents[i, c], dim=None\n ) # index by original dim order\n i_sd, i_mean = torch.std_mean(result[i, c], dim=None)\n\n result[i, c] = ((result[i, c] - i_mean) / i_sd) * r_sd + r_mean\n\n result = torch.lerp(latents, result, factor)\n return result\n\n\nclass LTXMultiScalePipeline:\n def _upsample_latents(\n self, latest_upsampler: LatentUpsampler, latents: torch.Tensor\n ):\n assert latents.device == latest_upsampler.device\n\n latents = un_normalize_latents(\n latents, self.vae, vae_per_channel_normalize=True\n )\n upsampled_latents = latest_upsampler(latents)\n upsampled_latents = normalize_latents(\n upsampled_latents, self.vae, vae_per_channel_normalize=True\n )\n return upsampled_latents\n\n def __init__(\n self, video_pipeline: LTXVideoPipeline, latent_upsampler: LatentUpsampler\n ):\n self.video_pipeline = video_pipeline\n self.vae = video_pipeline.vae\n self.latent_upsampler = latent_upsampler\n\n def __call__(\n self,\n downscale_factor: float,\n first_pass: dict,\n second_pass: dict,\n *args: Any,\n **kwargs: Any,\n ) -> Any:\n original_kwargs = kwargs.copy()\n original_output_type = kwargs[\"output_type\"]\n original_width = kwargs[\"width\"]\n original_height = kwargs[\"height\"]\n\n x_width = int(kwargs[\"width\"] * downscale_factor)\n downscaled_width = x_width - (x_width % self.video_pipeline.vae_scale_factor)\n x_height = int(kwargs[\"height\"] * downscale_factor)\n downscaled_height = x_height - (x_height % self.video_pipeline.vae_scale_factor)\n\n kwargs[\"output_type\"] = \"latent\"\n kwargs[\"width\"] = downscaled_width\n kwargs[\"height\"] = downscaled_height\n kwargs.update(**first_pass)\n result = self.video_pipeline(*args, **kwargs)\n latents = result.images\n\n upsampled_latents = self._upsample_latents(self.latent_upsampler, latents)\n upsampled_latents = adain_filter_latent(\n latents=upsampled_latents, reference_latents=latents\n )\n\n kwargs = original_kwargs\n\n kwargs[\"latents\"] = upsampled_latents\n kwargs[\"output_type\"] = original_output_type\n kwargs[\"width\"] = downscaled_width * 2\n kwargs[\"height\"] = downscaled_height * 2\n kwargs.update(**second_pass)\n\n result = self.video_pipeline(*args, **kwargs)\n if original_output_type != \"latent\":\n num_frames = result.images.shape[2]\n videos = rearrange(result.images, \"b c f h w -> (b f) c h w\")\n\n videos = F.interpolate(\n videos,\n size=(original_height, original_width),\n mode=\"bilinear\",\n align_corners=False,\n )\n videos = rearrange(videos, \"(b f) c h w -> b c f h w\", f=num_frames)\n result.images = videos\n\n return result\n"
},
{
"path": "ltx_video/schedulers/__init__.py",
"content": ""
},
{
"path": "ltx_video/schedulers/rf.py",
"content": "import math\nfrom abc import ABC, abstractmethod\nfrom dataclasses import dataclass\nfrom typing import Callable, Optional, Tuple, Union\nimport json\nimport os\nfrom pathlib import Path\n\nimport torch\nfrom diffusers.configuration_utils import ConfigMixin, register_to_config\nfrom diffusers.schedulers.scheduling_utils import SchedulerMixin\nfrom diffusers.utils import BaseOutput\nfrom torch import Tensor\nfrom safetensors import safe_open\n\n\nfrom ltx_video.utils.torch_utils import append_dims\n\nfrom ltx_video.utils.diffusers_config_mapping import (\n diffusers_and_ours_config_mapping,\n make_hashable_key,\n)\n\n\ndef linear_quadratic_schedule(num_steps, threshold_noise=0.025, linear_steps=None):\n if num_steps == 1:\n return torch.tensor([1.0])\n if linear_steps is None:\n linear_steps = num_steps // 2\n linear_sigma_schedule = [\n i * threshold_noise / linear_steps for i in range(linear_steps)\n ]\n threshold_noise_step_diff = linear_steps - threshold_noise * num_steps\n quadratic_steps = num_steps - linear_steps\n quadratic_coef = threshold_noise_step_diff / (linear_steps * quadratic_steps**2)\n linear_coef = threshold_noise / linear_steps - 2 * threshold_noise_step_diff / (\n quadratic_steps**2\n )\n const = quadratic_coef * (linear_steps**2)\n quadratic_sigma_schedule = [\n quadratic_coef * (i**2) + linear_coef * i + const\n for i in range(linear_steps, num_steps)\n ]\n sigma_schedule = linear_sigma_schedule + quadratic_sigma_schedule + [1.0]\n sigma_schedule = [1.0 - x for x in sigma_schedule]\n return torch.tensor(sigma_schedule[:-1])\n\n\ndef simple_diffusion_resolution_dependent_timestep_shift(\n samples_shape: torch.Size,\n timesteps: Tensor,\n n: int = 32 * 32,\n) -> Tensor:\n if len(samples_shape) == 3:\n _, m, _ = samples_shape\n elif len(samples_shape) in [4, 5]:\n m = math.prod(samples_shape[2:])\n else:\n raise ValueError(\n \"Samples must have shape (b, t, c), (b, c, h, w) or (b, c, f, h, w)\"\n )\n snr = (timesteps / (1 - timesteps)) ** 2\n shift_snr = torch.log(snr) + 2 * math.log(m / n)\n shifted_timesteps = torch.sigmoid(0.5 * shift_snr)\n\n return shifted_timesteps\n\n\ndef time_shift(mu: float, sigma: float, t: Tensor):\n return math.exp(mu) / (math.exp(mu) + (1 / t - 1) ** sigma)\n\n\ndef get_normal_shift(\n n_tokens: int,\n min_tokens: int = 1024,\n max_tokens: int = 4096,\n min_shift: float = 0.95,\n max_shift: float = 2.05,\n) -> Callable[[float], float]:\n m = (max_shift - min_shift) / (max_tokens - min_tokens)\n b = min_shift - m * min_tokens\n return m * n_tokens + b\n\n\ndef strech_shifts_to_terminal(shifts: Tensor, terminal=0.1):\n \"\"\"\n Stretch a function (given as sampled shifts) so that its final value matches the given terminal value\n using the provided formula.\n\n Parameters:\n - shifts (Tensor): The samples of the function to be stretched (PyTorch Tensor).\n - terminal (float): The desired terminal value (value at the last sample).\n\n Returns:\n - Tensor: The stretched shifts such that the final value equals `terminal`.\n \"\"\"\n if shifts.numel() == 0:\n raise ValueError(\"The 'shifts' tensor must not be empty.\")\n\n # Ensure terminal value is valid\n if terminal <= 0 or terminal >= 1:\n raise ValueError(\"The terminal value must be between 0 and 1 (exclusive).\")\n\n # Transform the shifts using the given formula\n one_minus_z = 1 - shifts\n scale_factor = one_minus_z[-1] / (1 - terminal)\n stretched_shifts = 1 - (one_minus_z / scale_factor)\n\n return stretched_shifts\n\n\ndef sd3_resolution_dependent_timestep_shift(\n samples_shape: torch.Size,\n timesteps: Tensor,\n target_shift_terminal: Optional[float] = None,\n) -> Tensor:\n \"\"\"\n Shifts the timestep schedule as a function of the generated resolution.\n\n In the SD3 paper, the authors empirically how to shift the timesteps based on the resolution of the target images.\n For more details: https://arxiv.org/pdf/2403.03206\n\n In Flux they later propose a more dynamic resolution dependent timestep shift, see:\n https://github.com/black-forest-labs/flux/blob/87f6fff727a377ea1c378af692afb41ae84cbe04/src/flux/sampling.py#L66\n\n\n Args:\n samples_shape (torch.Size): The samples batch shape (batch_size, channels, height, width) or\n (batch_size, channels, frame, height, width).\n timesteps (Tensor): A batch of timesteps with shape (batch_size,).\n target_shift_terminal (float): The target terminal value for the shifted timesteps.\n\n Returns:\n Tensor: The shifted timesteps.\n \"\"\"\n if len(samples_shape) == 3:\n _, m, _ = samples_shape\n elif len(samples_shape) in [4, 5]:\n m = math.prod(samples_shape[2:])\n else:\n raise ValueError(\n \"Samples must have shape (b, t, c), (b, c, h, w) or (b, c, f, h, w)\"\n )\n\n shift = get_normal_shift(m)\n time_shifts = time_shift(shift, 1, timesteps)\n if target_shift_terminal is not None: # Stretch the shifts to the target terminal\n time_shifts = strech_shifts_to_terminal(time_shifts, target_shift_terminal)\n return time_shifts\n\n\nclass TimestepShifter(ABC):\n @abstractmethod\n def shift_timesteps(self, samples_shape: torch.Size, timesteps: Tensor) -> Tensor:\n pass\n\n\n@dataclass\nclass RectifiedFlowSchedulerOutput(BaseOutput):\n \"\"\"\n Output class for the scheduler's step function output.\n\n Args:\n prev_sample (`torch.FloatTensor` of shape `(batch_size, num_channels, height, width)` for images):\n Computed sample (x_{t-1}) of previous timestep. `prev_sample` should be used as next model input in the\n denoising loop.\n pred_original_sample (`torch.FloatTensor` of shape `(batch_size, num_channels, height, width)` for images):\n The predicted denoised sample (x_{0}) based on the model output from the current timestep.\n `pred_original_sample` can be used to preview progress or for guidance.\n \"\"\"\n\n prev_sample: torch.FloatTensor\n pred_original_sample: Optional[torch.FloatTensor] = None\n\n\nclass RectifiedFlowScheduler(SchedulerMixin, ConfigMixin, TimestepShifter):\n order = 1\n\n @register_to_config\n def __init__(\n self,\n num_train_timesteps=1000,\n shifting: Optional[str] = None,\n base_resolution: int = 32**2,\n target_shift_terminal: Optional[float] = None,\n sampler: Optional[str] = \"Uniform\",\n shift: Optional[float] = None,\n ):\n super().__init__()\n self.init_noise_sigma = 1.0\n self.num_inference_steps = None\n self.sampler = sampler\n self.shifting = shifting\n self.base_resolution = base_resolution\n self.target_shift_terminal = target_shift_terminal\n self.timesteps = self.sigmas = self.get_initial_timesteps(\n num_train_timesteps, shift=shift\n )\n self.shift = shift\n\n def get_initial_timesteps(\n self, num_timesteps: int, shift: Optional[float] = None\n ) -> Tensor:\n if self.sampler == \"Uniform\":\n return torch.linspace(1, 1 / num_timesteps, num_timesteps)\n elif self.sampler == \"LinearQuadratic\":\n return linear_quadratic_schedule(num_timesteps)\n elif self.sampler == \"Constant\":\n assert (\n shift is not None\n ), \"Shift must be provided for constant time shift sampler.\"\n return time_shift(\n shift, 1, torch.linspace(1, 1 / num_timesteps, num_timesteps)\n )\n\n def shift_timesteps(self, samples_shape: torch.Size, timesteps: Tensor) -> Tensor:\n if self.shifting == \"SD3\":\n return sd3_resolution_dependent_timestep_shift(\n samples_shape, timesteps, self.target_shift_terminal\n )\n elif self.shifting == \"SimpleDiffusion\":\n return simple_diffusion_resolution_dependent_timestep_shift(\n samples_shape, timesteps, self.base_resolution\n )\n return timesteps\n\n def set_timesteps(\n self,\n num_inference_steps: Optional[int] = None,\n samples_shape: Optional[torch.Size] = None,\n timesteps: Optional[Tensor] = None,\n device: Union[str, torch.device] = None,\n ):\n \"\"\"\n Sets the discrete timesteps used for the diffusion chain. Supporting function to be run before inference.\n If `timesteps` are provided, they will be used instead of the scheduled timesteps.\n\n Args:\n num_inference_steps (`int` *optional*): The number of diffusion steps used when generating samples.\n samples_shape (`torch.Size` *optional*): The samples batch shape, used for shifting.\n timesteps ('torch.Tensor' *optional*): Specific timesteps to use instead of scheduled timesteps.\n device (`Union[str, torch.device]`, *optional*): The device to which the timesteps tensor will be moved.\n \"\"\"\n if timesteps is not None and num_inference_steps is not None:\n raise ValueError(\n \"You cannot provide both `timesteps` and `num_inference_steps`.\"\n )\n if timesteps is None:\n num_inference_steps = min(\n self.config.num_train_timesteps, num_inference_steps\n )\n timesteps = self.get_initial_timesteps(\n num_inference_steps, shift=self.shift\n ).to(device)\n timesteps = self.shift_timesteps(samples_shape, timesteps)\n else:\n timesteps = torch.Tensor(timesteps).to(device)\n num_inference_steps = len(timesteps)\n self.timesteps = timesteps\n self.num_inference_steps = num_inference_steps\n self.sigmas = self.timesteps\n\n @staticmethod\n def from_pretrained(pretrained_model_path: Union[str, os.PathLike]):\n pretrained_model_path = Path(pretrained_model_path)\n if pretrained_model_path.is_file():\n comfy_single_file_state_dict = {}\n with safe_open(pretrained_model_path, framework=\"pt\", device=\"cpu\") as f:\n metadata = f.metadata()\n for k in f.keys():\n comfy_single_file_state_dict[k] = f.get_tensor(k)\n configs = json.loads(metadata[\"config\"])\n config = configs[\"scheduler\"]\n del comfy_single_file_state_dict\n\n elif pretrained_model_path.is_dir():\n diffusers_noise_scheduler_config_path = (\n pretrained_model_path / \"scheduler\" / \"scheduler_config.json\"\n )\n\n with open(diffusers_noise_scheduler_config_path, \"r\") as f:\n scheduler_config = json.load(f)\n hashable_config = make_hashable_key(scheduler_config)\n if hashable_config in diffusers_and_ours_config_mapping:\n config = diffusers_and_ours_config_mapping[hashable_config]\n return RectifiedFlowScheduler.from_config(config)\n\n def scale_model_input(\n self, sample: torch.FloatTensor, timestep: Optional[int] = None\n ) -> torch.FloatTensor:\n # pylint: disable=unused-argument\n \"\"\"\n Ensures interchangeability with schedulers that need to scale the denoising model input depending on the\n current timestep.\n\n Args:\n sample (`torch.FloatTensor`): input sample\n timestep (`int`, optional): current timestep\n\n Returns:\n `torch.FloatTensor`: scaled input sample\n \"\"\"\n return sample\n\n def step(\n self,\n model_output: torch.FloatTensor,\n timestep: torch.FloatTensor,\n sample: torch.FloatTensor,\n return_dict: bool = True,\n stochastic_sampling: Optional[bool] = False,\n **kwargs,\n ) -> Union[RectifiedFlowSchedulerOutput, Tuple]:\n \"\"\"\n Predict the sample from the previous timestep by reversing the SDE. This function propagates the diffusion\n process from the learned model outputs (most often the predicted noise).\n z_{t_1} = z_t - Delta_t * v\n The method finds the next timestep that is lower than the input timestep(s) and denoises the latents\n to that level. The input timestep(s) are not required to be one of the predefined timesteps.\n\n Args:\n model_output (`torch.FloatTensor`):\n The direct output from learned diffusion model - the velocity,\n timestep (`float`):\n The current discrete timestep in the diffusion chain (global or per-token).\n sample (`torch.FloatTensor`):\n A current latent tokens to be de-noised.\n return_dict (`bool`, *optional*, defaults to `True`):\n Whether or not to return a [`~schedulers.scheduling_ddim.DDIMSchedulerOutput`] or `tuple`.\n stochastic_sampling (`bool`, *optional*, defaults to `False`):\n Whether to use stochastic sampling for the sampling process.\n\n Returns:\n [`~schedulers.scheduling_utils.RectifiedFlowSchedulerOutput`] or `tuple`:\n If return_dict is `True`, [`~schedulers.rf_scheduler.RectifiedFlowSchedulerOutput`] is returned,\n otherwise a tuple is returned where the first element is the sample tensor.\n \"\"\"\n if self.num_inference_steps is None:\n raise ValueError(\n \"Number of inference steps is 'None', you need to run 'set_timesteps' after creating the scheduler\"\n )\n t_eps = 1e-6 # Small epsilon to avoid numerical issues in timestep values\n\n timesteps_padded = torch.cat(\n [self.timesteps, torch.zeros(1, device=self.timesteps.device)]\n )\n\n # Find the next lower timestep(s) and compute the dt from the current timestep(s)\n if timestep.ndim == 0:\n # Global timestep case\n lower_mask = timesteps_padded < timestep - t_eps\n lower_timestep = timesteps_padded[lower_mask][0] # Closest lower timestep\n dt = timestep - lower_timestep\n\n else:\n # Per-token case\n assert timestep.ndim == 2\n lower_mask = timesteps_padded[:, None, None] < timestep[None] - t_eps\n lower_timestep = lower_mask * timesteps_padded[:, None, None]\n lower_timestep, _ = lower_timestep.max(dim=0)\n dt = (timestep - lower_timestep)[..., None]\n\n # Compute previous sample\n if stochastic_sampling:\n x0 = sample - timestep[..., None] * model_output\n next_timestep = timestep[..., None] - dt\n prev_sample = self.add_noise(x0, torch.randn_like(sample), next_timestep)\n else:\n prev_sample = sample - dt * model_output\n\n if not return_dict:\n return (prev_sample,)\n\n return RectifiedFlowSchedulerOutput(prev_sample=prev_sample)\n\n def add_noise(\n self,\n original_samples: torch.FloatTensor,\n noise: torch.FloatTensor,\n timesteps: torch.FloatTensor,\n ) -> torch.FloatTensor:\n sigmas = timesteps\n sigmas = append_dims(sigmas, original_samples.ndim)\n alphas = 1 - sigmas\n noisy_samples = alphas * original_samples + sigmas * noise\n return noisy_samples\n"
},
{
"path": "ltx_video/utils/__init__.py",
"content": ""
},
{
"path": "ltx_video/utils/diffusers_config_mapping.py",
"content": "def make_hashable_key(dict_key):\n def convert_value(value):\n if isinstance(value, list):\n return tuple(value)\n elif isinstance(value, dict):\n return tuple(sorted((k, convert_value(v)) for k, v in value.items()))\n else:\n return value\n\n return tuple(sorted((k, convert_value(v)) for k, v in dict_key.items()))\n\n\nDIFFUSERS_SCHEDULER_CONFIG = {\n \"_class_name\": \"FlowMatchEulerDiscreteScheduler\",\n \"_diffusers_version\": \"0.32.0.dev0\",\n \"base_image_seq_len\": 1024,\n \"base_shift\": 0.95,\n \"invert_sigmas\": False,\n \"max_image_seq_len\": 4096,\n \"max_shift\": 2.05,\n \"num_train_timesteps\": 1000,\n \"shift\": 1.0,\n \"shift_terminal\": 0.1,\n \"use_beta_sigmas\": False,\n \"use_dynamic_shifting\": True,\n \"use_exponential_sigmas\": False,\n \"use_karras_sigmas\": False,\n}\nDIFFUSERS_TRANSFORMER_CONFIG = {\n \"_class_name\": \"LTXVideoTransformer3DModel\",\n \"_diffusers_version\": \"0.32.0.dev0\",\n \"activation_fn\": \"gelu-approximate\",\n \"attention_bias\": True,\n \"attention_head_dim\": 64,\n \"attention_out_bias\": True,\n \"caption_channels\": 4096,\n \"cross_attention_dim\": 2048,\n \"in_channels\": 128,\n \"norm_elementwise_affine\": False,\n \"norm_eps\": 1e-06,\n \"num_attention_heads\": 32,\n \"num_layers\": 28,\n \"out_channels\": 128,\n \"patch_size\": 1,\n \"patch_size_t\": 1,\n \"qk_norm\": \"rms_norm_across_heads\",\n}\nDIFFUSERS_VAE_CONFIG = {\n \"_class_name\": \"AutoencoderKLLTXVideo\",\n \"_diffusers_version\": \"0.32.0.dev0\",\n \"block_out_channels\": [128, 256, 512, 512],\n \"decoder_causal\": False,\n \"encoder_causal\": True,\n \"in_channels\": 3,\n \"latent_channels\": 128,\n \"layers_per_block\": [4, 3, 3, 3, 4],\n \"out_channels\": 3,\n \"patch_size\": 4,\n \"patch_size_t\": 1,\n \"resnet_norm_eps\": 1e-06,\n \"scaling_factor\": 1.0,\n \"spatio_temporal_scaling\": [True, True, True, False],\n}\n\nOURS_SCHEDULER_CONFIG = {\n \"_class_name\": \"RectifiedFlowScheduler\",\n \"_diffusers_version\": \"0.25.1\",\n \"num_train_timesteps\": 1000,\n \"shifting\": \"SD3\",\n \"base_resolution\": None,\n \"target_shift_terminal\": 0.1,\n}\n\nOURS_TRANSFORMER_CONFIG = {\n \"_class_name\": \"Transformer3DModel\",\n \"_diffusers_version\": \"0.25.1\",\n \"_name_or_path\": \"PixArt-alpha/PixArt-XL-2-256x256\",\n \"activation_fn\": \"gelu-approximate\",\n \"attention_bias\": True,\n \"attention_head_dim\": 64,\n \"attention_type\": \"default\",\n \"caption_channels\": 4096,\n \"cross_attention_dim\": 2048,\n \"double_self_attention\": False,\n \"dropout\": 0.0,\n \"in_channels\": 128,\n \"norm_elementwise_affine\": False,\n \"norm_eps\": 1e-06,\n \"norm_num_groups\": 32,\n \"num_attention_heads\": 32,\n \"num_embeds_ada_norm\": 1000,\n \"num_layers\": 28,\n \"num_vector_embeds\": None,\n \"only_cross_attention\": False,\n \"out_channels\": 128,\n \"project_to_2d_pos\": True,\n \"upcast_attention\": False,\n \"use_linear_projection\": False,\n \"qk_norm\": \"rms_norm\",\n \"standardization_norm\": \"rms_norm\",\n \"positional_embedding_type\": \"rope\",\n \"positional_embedding_theta\": 10000.0,\n \"positional_embedding_max_pos\": [20, 2048, 2048],\n \"timestep_scale_multiplier\": 1000,\n}\nOURS_VAE_CONFIG = {\n \"_class_name\": \"CausalVideoAutoencoder\",\n \"dims\": 3,\n \"in_channels\": 3,\n \"out_channels\": 3,\n \"latent_channels\": 128,\n \"blocks\": [\n [\"res_x\", 4],\n [\"compress_all\", 1],\n [\"res_x_y\", 1],\n [\"res_x\", 3],\n [\"compress_all\", 1],\n [\"res_x_y\", 1],\n [\"res_x\", 3],\n [\"compress_all\", 1],\n [\"res_x\", 3],\n [\"res_x\", 4],\n ],\n \"scaling_factor\": 1.0,\n \"norm_layer\": \"pixel_norm\",\n \"patch_size\": 4,\n \"latent_log_var\": \"uniform\",\n \"use_quant_conv\": False,\n \"causal_decoder\": False,\n}\n\n\ndiffusers_and_ours_config_mapping = {\n make_hashable_key(DIFFUSERS_SCHEDULER_CONFIG): OURS_SCHEDULER_CONFIG,\n make_hashable_key(DIFFUSERS_TRANSFORMER_CONFIG): OURS_TRANSFORMER_CONFIG,\n make_hashable_key(DIFFUSERS_VAE_CONFIG): OURS_VAE_CONFIG,\n}\n\n\nTRANSFORMER_KEYS_RENAME_DICT = {\n \"proj_in\": \"patchify_proj\",\n \"time_embed\": \"adaln_single\",\n \"norm_q\": \"q_norm\",\n \"norm_k\": \"k_norm\",\n}\n\n\nVAE_KEYS_RENAME_DICT = {\n \"decoder.up_blocks.3.conv_in\": \"decoder.up_blocks.7\",\n \"decoder.up_blocks.3.upsamplers.0\": \"decoder.up_blocks.8\",\n \"decoder.up_blocks.3\": \"decoder.up_blocks.9\",\n \"decoder.up_blocks.2.upsamplers.0\": \"decoder.up_blocks.5\",\n \"decoder.up_blocks.2.conv_in\": \"decoder.up_blocks.4\",\n \"decoder.up_blocks.2\": \"decoder.up_blocks.6\",\n \"decoder.up_blocks.1.upsamplers.0\": \"decoder.up_blocks.2\",\n \"decoder.up_blocks.1\": \"decoder.up_blocks.3\",\n \"decoder.up_blocks.0\": \"decoder.up_blocks.1\",\n \"decoder.mid_block\": \"decoder.up_blocks.0\",\n \"encoder.down_blocks.3\": \"encoder.down_blocks.8\",\n \"encoder.down_blocks.2.downsamplers.0\": \"encoder.down_blocks.7\",\n \"encoder.down_blocks.2\": \"encoder.down_blocks.6\",\n \"encoder.down_blocks.1.downsamplers.0\": \"encoder.down_blocks.4\",\n \"encoder.down_blocks.1.conv_out\": \"encoder.down_blocks.5\",\n \"encoder.down_blocks.1\": \"encoder.down_blocks.3\",\n \"encoder.down_blocks.0.conv_out\": \"encoder.down_blocks.2\",\n \"encoder.down_blocks.0.downsamplers.0\": \"encoder.down_blocks.1\",\n \"encoder.down_blocks.0\": \"encoder.down_blocks.0\",\n \"encoder.mid_block\": \"encoder.down_blocks.9\",\n \"conv_shortcut.conv\": \"conv_shortcut\",\n \"resnets\": \"res_blocks\",\n \"norm3\": \"norm3.norm\",\n \"latents_mean\": \"per_channel_statistics.mean-of-means\",\n \"latents_std\": \"per_channel_statistics.std-of-means\",\n}\n"
},
{
"path": "ltx_video/utils/prompt_enhance_utils.py",
"content": "import logging\nfrom typing import Union, List, Optional\n\nimport torch\nfrom PIL import Image\n\nlogger = logging.getLogger(__name__) # pylint: disable=invalid-name\n\nT2V_CINEMATIC_PROMPT = \"\"\"You are an expert cinematic director with many award winning movies, When writing prompts based on the user input, focus on detailed, chronological descriptions of actions and scenes.\nInclude specific movements, appearances, camera angles, and environmental details - all in a single flowing paragraph.\nStart directly with the action, and keep descriptions literal and precise.\nThink like a cinematographer describing a shot list.\nDo not change the user input intent, just enhance it.\nKeep within 150 words.\nFor best results, build your prompts using this structure:\nStart with main action in a single sentence\nAdd specific details about movements and gestures\nDescribe character/object appearances precisely\nInclude background and environment details\nSpecify camera angles and movements\nDescribe lighting and colors\nNote any changes or sudden events\nDo not exceed the 150 word limit!\nOutput the enhanced prompt only.\n\"\"\"\n\nI2V_CINEMATIC_PROMPT = \"\"\"You are an expert cinematic director with many award winning movies, When writing prompts based on the user input, focus on detailed, chronological descriptions of actions and scenes.\nInclude specific movements, appearances, camera angles, and environmental details - all in a single flowing paragraph.\nStart directly with the action, and keep descriptions literal and precise.\nThink like a cinematographer describing a shot list.\nKeep within 150 words.\nFor best results, build your prompts using this structure:\nDescribe the image first and then add the user input. Image description should be in first priority! Align to the image caption if it contradicts the user text input.\nStart with main action in a single sentence\nAdd specific details about movements and gestures\nDescribe character/object appearances precisely\nInclude background and environment details\nSpecify camera angles and movements\nDescribe lighting and colors\nNote any changes or sudden events\nAlign to the image caption if it contradicts the user text input.\nDo not exceed the 150 word limit!\nOutput the enhanced prompt only.\n\"\"\"\n\n\ndef tensor_to_pil(tensor):\n # Ensure tensor is in range [-1, 1]\n assert tensor.min() >= -1 and tensor.max() <= 1\n\n # Convert from [-1, 1] to [0, 1]\n tensor = (tensor + 1) / 2\n\n # Rearrange from [C, H, W] to [H, W, C]\n tensor = tensor.permute(1, 2, 0)\n\n # Convert to numpy array and then to uint8 range [0, 255]\n numpy_image = (tensor.cpu().numpy() * 255).astype(\"uint8\")\n\n # Convert to PIL Image\n return Image.fromarray(numpy_image)\n\n\ndef generate_cinematic_prompt(\n image_caption_model,\n image_caption_processor,\n prompt_enhancer_model,\n prompt_enhancer_tokenizer,\n prompt: Union[str, List[str]],\n conditioning_items: Optional[List] = None,\n max_new_tokens: int = 256,\n) -> List[str]:\n prompts = [prompt] if isinstance(prompt, str) else prompt\n\n if conditioning_items is None:\n prompts = _generate_t2v_prompt(\n prompt_enhancer_model,\n prompt_enhancer_tokenizer,\n prompts,\n max_new_tokens,\n T2V_CINEMATIC_PROMPT,\n )\n else:\n if len(conditioning_items) > 1 or conditioning_items[0].media_frame_number != 0:\n logger.warning(\n \"prompt enhancement does only support unconditional or first frame of conditioning items, returning original prompts\"\n )\n return prompts\n\n first_frame_conditioning_item = conditioning_items[0]\n first_frames = _get_first_frames_from_conditioning_item(\n first_frame_conditioning_item\n )\n\n assert len(first_frames) == len(\n prompts\n ), \"Number of conditioning frames must match number of prompts\"\n\n prompts = _generate_i2v_prompt(\n image_caption_model,\n image_caption_processor,\n prompt_enhancer_model,\n prompt_enhancer_tokenizer,\n prompts,\n first_frames,\n max_new_tokens,\n I2V_CINEMATIC_PROMPT,\n )\n\n return prompts\n\n\ndef _get_first_frames_from_conditioning_item(conditioning_item) -> List[Image.Image]:\n frames_tensor = conditioning_item.media_item\n return [\n tensor_to_pil(frames_tensor[i, :, 0, :, :])\n for i in range(frames_tensor.shape[0])\n ]\n\n\ndef _generate_t2v_prompt(\n prompt_enhancer_model,\n prompt_enhancer_tokenizer,\n prompts: List[str],\n max_new_tokens: int,\n system_prompt: str,\n) -> List[str]:\n messages = [\n [\n {\"role\": \"system\", \"content\": system_prompt},\n {\"role\": \"user\", \"content\": f\"user_prompt: {p}\"},\n ]\n for p in prompts\n ]\n\n texts = [\n prompt_enhancer_tokenizer.apply_chat_template(\n m, tokenize=False, add_generation_prompt=True\n )\n for m in messages\n ]\n model_inputs = prompt_enhancer_tokenizer(texts, return_tensors=\"pt\").to(\n prompt_enhancer_model.device\n )\n\n return _generate_and_decode_prompts(\n prompt_enhancer_model, prompt_enhancer_tokenizer, model_inputs, max_new_tokens\n )\n\n\ndef _generate_i2v_prompt(\n image_caption_model,\n image_caption_processor,\n prompt_enhancer_model,\n prompt_enhancer_tokenizer,\n prompts: List[str],\n first_frames: List[Image.Image],\n max_new_tokens: int,\n system_prompt: str,\n) -> List[str]:\n image_captions = _generate_image_captions(\n image_caption_model, image_caption_processor, first_frames\n )\n\n messages = [\n [\n {\"role\": \"system\", \"content\": system_prompt},\n {\"role\": \"user\", \"content\": f\"user_prompt: {p}\\nimage_caption: {c}\"},\n ]\n for p, c in zip(prompts, image_captions)\n ]\n\n texts = [\n prompt_enhancer_tokenizer.apply_chat_template(\n m, tokenize=False, add_generation_prompt=True\n )\n for m in messages\n ]\n model_inputs = prompt_enhancer_tokenizer(texts, return_tensors=\"pt\").to(\n prompt_enhancer_model.device\n )\n\n return _generate_and_decode_prompts(\n prompt_enhancer_model, prompt_enhancer_tokenizer, model_inputs, max_new_tokens\n )\n\n\ndef _generate_image_captions(\n image_caption_model,\n image_caption_processor,\n images: List[Image.Image],\n system_prompt: str = \"\",\n) -> List[str]:\n image_caption_prompts = [system_prompt] * len(images)\n inputs = image_caption_processor(\n image_caption_prompts, images, return_tensors=\"pt\"\n ).to(image_caption_model.device)\n\n with torch.inference_mode():\n generated_ids = image_caption_model.generate(\n input_ids=inputs[\"input_ids\"],\n pixel_values=inputs[\"pixel_values\"],\n max_new_tokens=1024,\n do_sample=False,\n num_beams=3,\n )\n\n return image_caption_processor.batch_decode(generated_ids, skip_special_tokens=True)\n\n\ndef _generate_and_decode_prompts(\n prompt_enhancer_model, prompt_enhancer_tokenizer, model_inputs, max_new_tokens: int\n) -> List[str]:\n with torch.inference_mode():\n outputs = prompt_enhancer_model.generate(\n **model_inputs, max_new_tokens=max_new_tokens\n )\n generated_ids = [\n output_ids[len(input_ids) :]\n for input_ids, output_ids in zip(model_inputs.input_ids, outputs)\n ]\n decoded_prompts = prompt_enhancer_tokenizer.batch_decode(\n generated_ids, skip_special_tokens=True\n )\n\n return decoded_prompts\n"
},
{
"path": "ltx_video/utils/skip_layer_strategy.py",
"content": "from enum import Enum, auto\n\n\nclass SkipLayerStrategy(Enum):\n AttentionSkip = auto()\n AttentionValues = auto()\n Residual = auto()\n TransformerBlock = auto()\n"
},
{
"path": "ltx_video/utils/torch_utils.py",
"content": "import torch\nfrom torch import nn\n\n\ndef append_dims(x: torch.Tensor, target_dims: int) -> torch.Tensor:\n \"\"\"Appends dimensions to the end of a tensor until it has target_dims dimensions.\"\"\"\n dims_to_append = target_dims - x.ndim\n if dims_to_append < 0:\n raise ValueError(\n f\"input has {x.ndim} dims but target_dims is {target_dims}, which is less\"\n )\n elif dims_to_append == 0:\n return x\n return x[(...,) + (None,) * dims_to_append]\n\n\nclass Identity(nn.Module):\n \"\"\"A placeholder identity operator that is argument-insensitive.\"\"\"\n\n def __init__(self, *args, **kwargs) -> None: # pylint: disable=unused-argument\n super().__init__()\n\n # pylint: disable=unused-argument\n def forward(self, x: torch.Tensor, *args, **kwargs) -> torch.Tensor:\n return x\n"
},
{
"path": "pyproject.toml",
"content": "[build-system]\nrequires = [\"setuptools>=42\", \"wheel\"]\nbuild-backend = \"setuptools.build_meta\"\n\n[project]\nname = \"ltx-video\"\nversion = \"0.1.2\"\ndescription = \"A package for LTX-Video model\"\nauthors = [\n { name = \"LTX-Video Team\", email = \"ltx-video@lightricks.com\" }\n]\nrequires-python = \">=3.10\"\nreadme = \"README.md\"\nclassifiers = [\n \"Programming Language :: Python :: 3\",\n \"Operating System :: OS Independent\"\n]\ndependencies = [\n \"torch>=2.1.0\",\n \"diffusers>=0.28.2\",\n \"transformers>=4.47.2,<4.52.0\",\n \"sentencepiece>=0.1.96\",\n \"huggingface-hub~=0.30\",\n \"einops\",\n \"timm\"\n]\n\n[project.optional-dependencies]\ninference = [\n \"imageio[ffmpeg]\",\n \"av\",\n \"torchvision\"\n]\ntest = [\n \"pytest\",\n]\n\n[tool.setuptools.packages.find]\ninclude = [\"ltx_video*\"]\n\n[tool.setuptools.package-data]\nltx_video = [\"configs/*.yaml\"]\n"
},
{
"path": "tests/conftest.py",
"content": "import json\nimport pytest\nimport safetensors.torch\nimport torch\n\nfrom ltx_video.models.autoencoders.causal_video_autoencoder import (\n CausalVideoAutoencoder,\n create_video_autoencoder_demo_config,\n PER_CHANNEL_STATISTICS_PREFIX,\n)\nfrom ltx_video.models.transformers.transformer3d import Transformer3DModel\n\n\ndef pytest_make_parametrize_id(config, val, argname):\n if isinstance(val, str):\n return f\"{argname}-{val}\"\n return f\"{argname}-{repr(val)}\"\n\n\n@pytest.fixture\ndef num_latent_channels():\n return 16\n\n\n@pytest.fixture\ndef video_autoencoder(num_latent_channels):\n config = create_video_autoencoder_demo_config(latent_channels=num_latent_channels)\n model = CausalVideoAutoencoder.from_config(config)\n model.eval().to(torch.bfloat16)\n return model\n\n\n@pytest.fixture\ndef transformer_config(num_latent_channels):\n transformer_config = {\n \"activation_fn\": \"gelu-approximate\",\n \"attention_bias\": True,\n \"attention_head_dim\": 12,\n \"attention_type\": \"default\",\n \"caption_channels\": 4096,\n \"cross_attention_dim\": 192,\n \"double_self_attention\": False,\n \"dropout\": 0.0,\n \"in_channels\": num_latent_channels,\n \"norm_elementwise_affine\": False,\n \"norm_eps\": 1e-06,\n \"norm_num_groups\": 32,\n \"num_attention_heads\": 16,\n \"num_embeds_ada_norm\": 1000,\n \"num_layers\": 2,\n \"num_vector_embeds\": None,\n \"only_cross_attention\": False,\n \"out_channels\": num_latent_channels,\n \"upcast_attention\": False,\n \"use_linear_projection\": False,\n \"qk_norm\": \"rms_norm\",\n \"standardization_norm\": \"rms_norm\",\n \"positional_embedding_type\": \"rope\",\n \"positional_embedding_theta\": 10000.0,\n \"positional_embedding_max_pos\": [120, 1, 1],\n \"timestep_scale_multiplier\": 1000,\n }\n return transformer_config\n\n\n@pytest.fixture\ndef synthetic_ckpt_path(\n tmp_path, video_autoencoder, num_latent_channels, transformer_config\n):\n # Create transformer\n transformer = Transformer3DModel.from_config(transformer_config)\n transformer.to(torch.bfloat16)\n\n # Prepare configs and state dicts\n configs = {\"transformer\": transformer_config, \"vae\": vars(video_autoencoder.config)}\n transformer_sd = transformer.state_dict()\n transformer_sd = {\n \"model.diffusion_model.\" + key: value for key, value in transformer_sd.items()\n }\n\n # Prepare VAE state dict with per-channel statistics\n vae_sd = video_autoencoder.state_dict()\n vae_sd[f\"{PER_CHANNEL_STATISTICS_PREFIX}std-of-means\"] = torch.rand(\n num_latent_channels,\n )\n vae_sd[f\"{PER_CHANNEL_STATISTICS_PREFIX}mean-of-means\"] = torch.rand(\n num_latent_channels,\n )\n vae_sd = {\"vae.\" + key: value for key, value in vae_sd.items()}\n\n out_file_path = f\"{tmp_path}/test_ckpt.safetensors\"\n safetensors.torch.save_file(\n {**transformer_sd, **vae_sd},\n out_file_path,\n metadata={\"config\": json.dumps(configs)},\n )\n return out_file_path\n"
},
{
"path": "tests/test_configs.py",
"content": "import pytest\nfrom pathlib import Path\n\nfrom ltx_video.inference import infer, InferenceConfig\n\nCONFIGS_DIR = Path(__file__).parents[1] / \"configs\"\n\n\n@pytest.fixture\ndef prompt():\n return \"A video of a cat playing with a ball.\"\n\n\n# mark as slow to avoid running these tests by default\n@pytest.mark.slow\n@pytest.mark.parametrize(\n \"pipeline_config\",\n [pytest.param(config, id=config.stem) for config in CONFIGS_DIR.glob(\"*.yaml\")],\n)\ndef test_run_config(tmp_path, prompt, pipeline_config):\n if \"fp8\" in pipeline_config.stem:\n pytest.skip(\"Skipping fp8 configs as they require specific hardware support.\")\n\n inference_config = InferenceConfig(prompt=prompt)\n inference_config.pipeline_config = CONFIGS_DIR / pipeline_config\n inference_config.output_path = tmp_path / f\"{pipeline_config.stem}\"\n inference_config.height = 256\n inference_config.width = 320\n inference_config.num_frames = 33\n infer(config=inference_config)\n"
},
{
"path": "tests/test_inference.py",
"content": "from dataclasses import asdict\nimport pytest\nimport torch\nimport yaml\n\nfrom ltx_video.inference import (\n create_ltx_video_pipeline,\n get_device,\n infer,\n InferenceConfig,\n)\nfrom ltx_video.utils.skip_layer_strategy import SkipLayerStrategy\n\n\n@pytest.fixture\ndef input_image_path():\n return \"tests/utils/woman.jpeg\"\n\n\n@pytest.fixture\ndef input_video_path():\n return \"tests/utils/woman.mp4\"\n\n\ndef base_inference_config(tmp_path, pipeline_config):\n temp_config_path = tmp_path / \"config.yaml\"\n with open(temp_config_path, \"w\") as f:\n yaml.dump(pipeline_config, f)\n\n return InferenceConfig(\n seed=42,\n height=256,\n width=320,\n num_frames=49,\n frame_rate=25,\n prompt=\"A young woman with wavy, shoulder-length light brown hair stands outdoors on a foggy day. She wears a cozy pink turtleneck sweater, with a serene expression and piercing blue eyes. A wooden fence and a misty, grassy field fade into the background, evoking a calm and introspective mood.\",\n negative_prompt=\"worst quality, inconsistent motion, blurry, jittery, distorted\",\n output_path=tmp_path,\n pipeline_config=temp_config_path,\n )\n\n\n@pytest.fixture\ndef base_pipeline_config(synthetic_ckpt_path):\n return {\n \"num_inference_steps\": 1,\n \"stg_mode\": \"attention_values\",\n \"skip_block_list\": [1],\n \"precision\": \"bfloat16\",\n \"decode_timestep\": 0.05,\n \"decode_noise_scale\": 0.025,\n \"checkpoint_path\": synthetic_ckpt_path,\n \"text_encoder_model_name_or_path\": \"PixArt-alpha/PixArt-XL-2-1024-MS\",\n \"prompt_enhancer_image_caption_model_name_or_path\": \"MiaoshouAI/Florence-2-large-PromptGen-v2.0\",\n \"prompt_enhancer_llm_model_name_or_path\": \"unsloth/Llama-3.2-3B-Instruct\",\n \"prompt_enhancement_words_threshold\": 120,\n \"sampler\": \"LinearQuadratic\",\n }\n\n\n@pytest.mark.parametrize(\n \"conditioning_test_mode\",\n [\"unconditional\", \"first-frame\", \"first-sequence\", \"sequence-and-frame\"],\n ids=lambda x: f\"conditioning_test_mode={x}\",\n)\ndef test_condition_modes(\n tmp_path,\n conditioning_test_mode,\n input_image_path,\n input_video_path,\n base_pipeline_config,\n):\n inference_config = base_inference_config(tmp_path, base_pipeline_config)\n\n if conditioning_test_mode == \"unconditional\":\n pass\n elif conditioning_test_mode == \"first-frame\":\n inference_config.conditioning_media_paths = [input_image_path]\n inference_config.conditioning_start_frames = [0]\n elif conditioning_test_mode == \"first-sequence\":\n inference_config.conditioning_media_paths = [input_video_path]\n inference_config.conditioning_start_frames = [0]\n elif conditioning_test_mode == \"sequence-and-frame\":\n inference_config.conditioning_media_paths = [input_video_path, input_image_path]\n inference_config.conditioning_start_frames = [16, 43]\n else:\n raise ValueError(f\"Unknown conditioning mode: {conditioning_test_mode}\")\n\n # Test that the infer function runs without errors\n infer(inference_config)\n\n\ndef test_vid2vid(tmp_path, input_video_path, base_pipeline_config):\n pipeline_config = base_pipeline_config\n pipeline_config[\"num_inference_steps\"] = 3\n pipeline_config[\"skip_initial_inference_steps\"] = 1\n\n inference_config = base_inference_config(tmp_path, pipeline_config)\n inference_config.num_frames = 25\n inference_config.input_media_path = input_video_path\n\n # Test that the infer function runs without errors\n infer(inference_config)\n\n\ndef test_pipeline_on_batch(tmp_path, base_pipeline_config):\n pipeline_config = base_pipeline_config\n inference_config = base_inference_config(tmp_path, pipeline_config)\n inference_config.num_frames = 1 # For faster test, we use a single frame\n\n device = get_device()\n pipeline = create_ltx_video_pipeline(\n ckpt_path=pipeline_config[\"checkpoint_path\"],\n device=device,\n precision=pipeline_config[\"precision\"],\n text_encoder_model_name_or_path=pipeline_config[\n \"text_encoder_model_name_or_path\"\n ],\n enhance_prompt=False,\n prompt_enhancer_image_caption_model_name_or_path=pipeline_config[\n \"prompt_enhancer_image_caption_model_name_or_path\"\n ],\n prompt_enhancer_llm_model_name_or_path=pipeline_config[\n \"prompt_enhancer_llm_model_name_or_path\"\n ],\n sampler=\"LinearQuadratic\",\n )\n\n first_prompt = \"A vintage yellow car drives along a wet mountain road, its rear wheels kicking up a light spray as it moves. The camera follows close behind, capturing the curvature of the road as it winds through rocky cliffs and lush green hills. The sunlight pierces through scattered clouds, reflecting off the car's rain-speckled surface, creating a dynamic, cinematic moment. The scene conveys a sense of freedom and exploration as the car disappears into the distance.\"\n second_prompt = \"A woman with blonde hair styled up, wearing a black dress with sequins and pearl earrings, looks down with a sad expression on her face. The camera remains stationary, focused on the woman's face. The lighting is dim, casting soft shadows on her face. The scene appears to be from a movie or TV show.\"\n\n def get_images(prompts):\n generators = [\n torch.Generator(device=device).manual_seed(inference_config.seed)\n for _ in range(2)\n ]\n torch.manual_seed(inference_config.seed)\n\n params = asdict(inference_config)\n params[\"prompt\"] = prompts\n\n pipeline_result = pipeline(\n generator=generators,\n output_type=\"pt\",\n vae_per_channel_normalize=True,\n **params,\n )\n return pipeline_result.images\n\n # Run the pipeline on two different batches of prompts\n batch_diff_images = get_images([first_prompt, second_prompt])\n batch_same_images = get_images([second_prompt, second_prompt])\n\n # Take the second image from both runs, which should be equal\n image2_not_same = batch_diff_images[1, :, 0, :, :]\n image2_same = batch_same_images[1, :, 0, :, :]\n\n assert torch.allclose(image2_not_same, image2_same)\n\n\ndef test_prompt_enhancement(tmp_path, base_pipeline_config):\n pipeline_config = base_pipeline_config\n inference_config = base_inference_config(tmp_path, pipeline_config)\n inference_config.num_frames = 1 # For faster test, we use a single frame\n\n device = get_device()\n pipeline = create_ltx_video_pipeline(\n ckpt_path=pipeline_config[\"checkpoint_path\"],\n device=device,\n precision=pipeline_config[\"precision\"],\n text_encoder_model_name_or_path=pipeline_config[\n \"text_encoder_model_name_or_path\"\n ],\n enhance_prompt=True,\n prompt_enhancer_image_caption_model_name_or_path=pipeline_config[\n \"prompt_enhancer_image_caption_model_name_or_path\"\n ],\n prompt_enhancer_llm_model_name_or_path=pipeline_config[\n \"prompt_enhancer_llm_model_name_or_path\"\n ],\n sampler=\"LinearQuadratic\",\n )\n # Mock the pipeline's _encode_prompt method to verify the prompt being used\n original_encode_prompt = pipeline.encode_prompt\n\n def mock_encode_prompt(prompt, *args, **kwargs):\n prompts_used.append(prompt[0] if isinstance(prompt, list) else prompt)\n return original_encode_prompt(prompt, *args, **kwargs)\n\n pipeline.encode_prompt = mock_encode_prompt\n\n original_prompt = \"A cat sitting on a windowsill\"\n inference_config.prompt = original_prompt\n\n def run_pipeline(enhance_prompt):\n params = asdict(inference_config)\n pipeline(\n enhance_prompt=enhance_prompt,\n **params,\n skip_layer_strategy=SkipLayerStrategy.AttentionValues,\n vae_per_channel_normalize=True,\n output_type=\"pt\",\n )\n assert (\n len(prompts_used) > 0\n ), f\"No prompts were used in the pipeline run with enhance_prompt={enhance_prompt}\"\n\n if enhance_prompt:\n # Verify that the enhanced prompt was used\n assert (\n prompts_used[0] != original_prompt\n ), f\"Expected enhanced prompt to be different from original prompt, but got: {original_prompt}\"\n else:\n # Verify that the original prompt was used\n assert (\n prompts_used[0] == original_prompt\n ), f\"Expected original prompt to be used, but got: {prompts_used[0]}\"\n\n # Run pipeline with prompt enhancement enabled\n prompts_used = []\n run_pipeline(enhance_prompt=True)\n # Run pipeline with prompt enhancement disabled\n prompts_used = []\n run_pipeline(enhance_prompt=False)\n"
},
{
"path": "tests/test_scheduler.py",
"content": "import pytest\nimport torch\nfrom ltx_video.schedulers.rf import RectifiedFlowScheduler\n\n\ndef init_latents_and_scheduler(sampler):\n batch_size, n_tokens, n_channels = 2, 4096, 128\n num_steps = 20\n scheduler = RectifiedFlowScheduler(\n sampler=(\"Uniform\" if sampler.lower() == \"uniform\" else \"LinearQuadratic\")\n )\n latents = torch.randn(size=(batch_size, n_tokens, n_channels))\n scheduler.set_timesteps(num_inference_steps=num_steps, samples_shape=latents.shape)\n return scheduler, latents\n\n\n@pytest.mark.parametrize(\"sampler\", [\"LinearQuadratic\", \"Uniform\"])\ndef test_scheduler_default_behavior(sampler):\n \"\"\"\n Test the case of a single timestep from the list of timesteps.\n \"\"\"\n scheduler, latents = init_latents_and_scheduler(sampler)\n\n for i, t in enumerate(scheduler.timesteps):\n noise_pred = torch.randn_like(latents)\n denoised_latents = scheduler.step(\n noise_pred,\n t,\n latents,\n return_dict=False,\n )[0]\n\n # Verify the denoising\n next_t = scheduler.timesteps[i + 1] if i < len(scheduler.timesteps) - 1 else 0.0\n dt = t - next_t\n expected_denoised_latents = latents - dt * noise_pred\n assert torch.allclose(denoised_latents, expected_denoised_latents, atol=1e-06)\n\n\n@pytest.mark.parametrize(\"sampler\", [\"LinearQuadratic\", \"Uniform\"])\ndef test_scheduler_per_token(sampler):\n \"\"\"\n Test the case of a timestep per token (from the list of timesteps).\n Some tokens are set with timestep of 0.\n \"\"\"\n scheduler, latents = init_latents_and_scheduler(sampler)\n batch_size, n_tokens = latents.shape[:2]\n for i, t in enumerate(scheduler.timesteps):\n timesteps = torch.full((batch_size, n_tokens), t)\n timesteps[:, 0] = 0.0\n noise_pred = torch.randn_like(latents)\n denoised_latents = scheduler.step(\n noise_pred,\n timesteps,\n latents,\n return_dict=False,\n )[0]\n\n # Verify the denoising\n next_t = scheduler.timesteps[i + 1] if i < len(scheduler.timesteps) - 1 else 0.0\n next_timesteps = torch.full((batch_size, n_tokens), next_t)\n dt = timesteps - next_timesteps\n expected_denoised_latents = latents - dt.unsqueeze(-1) * noise_pred\n assert torch.allclose(\n denoised_latents[:, 1:], expected_denoised_latents[:, 1:], atol=1e-06\n )\n assert torch.allclose(denoised_latents[:, 0], latents[:, 0], atol=1e-06)\n\n\n@pytest.mark.parametrize(\"sampler\", [\"LinearQuadratic\", \"Uniform\"])\ndef test_scheduler_t_not_in_list(sampler):\n \"\"\"\n Test the case of a timestep per token NOT from the list of timesteps.\n \"\"\"\n scheduler, latents = init_latents_and_scheduler(sampler)\n batch_size, n_tokens = latents.shape[:2]\n for i in range(len(scheduler.timesteps)):\n if i < len(scheduler.timesteps) - 1:\n t = (scheduler.timesteps[i] + scheduler.timesteps[i + 1]) / 2\n else:\n t = scheduler.timesteps[i] / 2\n timesteps = torch.full((batch_size, n_tokens), t)\n noise_pred = torch.randn_like(latents)\n denoised_latents = scheduler.step(\n noise_pred,\n timesteps,\n latents,\n return_dict=False,\n )[0]\n\n # Verify the denoising\n next_t = scheduler.timesteps[i + 1] if i < len(scheduler.timesteps) - 1 else 0.0\n next_timesteps = torch.full((batch_size, n_tokens), next_t)\n dt = timesteps - next_timesteps\n expected_denoised_latents = latents - dt.unsqueeze(-1) * noise_pred\n assert torch.allclose(denoised_latents, expected_denoised_latents, atol=1e-06)\n"
},
{
"path": "tests/test_vae.py",
"content": "import pytest\nimport torch\nfrom ltx_video.models.autoencoders.causal_video_autoencoder import (\n CausalVideoAutoencoder,\n)\n\n\ndef test_encode_decode_shape(video_autoencoder, num_latent_channels):\n spatial_factor = video_autoencoder.spatial_downscale_factor\n temporal_factor = video_autoencoder.temporal_downscale_factor\n input_videos = torch.randn(2, 3, 17, 64, 64, dtype=torch.bfloat16)\n\n # Encode\n latent = video_autoencoder.encode(input_videos).latent_dist.mode()\n expected_shape = (\n input_videos.shape[0],\n num_latent_channels,\n (input_videos.shape[2] + 7) // temporal_factor,\n input_videos.shape[3] // spatial_factor,\n input_videos.shape[4] // spatial_factor,\n )\n assert latent.shape == expected_shape\n\n # Decode\n timestep = torch.ones(input_videos.shape[0]) * 0.1\n reconstructed_videos = video_autoencoder.decode(\n latent, target_shape=input_videos.shape, timestep=timestep\n ).sample\n assert input_videos.shape == reconstructed_videos.shape\n\n\ndef test_temporal_causality(video_autoencoder):\n # validate temporal causality in encoder\n input_videos = torch.randn(2, 3, 17, 64, 64, dtype=torch.bfloat16)\n latent = video_autoencoder.encode(input_videos).latent_dist.mode()\n\n # Check that encoding a single frame matches the corresponding slice in the full latent\n input_image = input_videos[:, :, :1, :, :]\n image_latent = video_autoencoder.encode(input_image).latent_dist.mode()\n assert torch.allclose(image_latent, latent[:, :, :1, :, :], atol=1e-6)\n\n # Check that encoding a sequence of frames matches the corresponding slice in the full latent\n input_sequence = input_videos[:, :, :9, :, :]\n sequence_latent = video_autoencoder.encode(input_sequence).latent_dist.mode()\n assert torch.allclose(sequence_latent, latent[:, :, :2, :, :], atol=1e-6)\n\n\n@pytest.mark.parametrize(\n \"layer_name,expected_temporal_factor,expected_spatial_factor\",\n [\n (\"compress_space_res\", 1, 2),\n (\"compress_space\", 1, 2),\n (\"compress_time_res\", 2, 1),\n (\"compress_time\", 2, 1),\n (\"compress_all_res\", 2, 2),\n (\"compress_all\", 2, 2),\n ],\n)\ndef test_downscale_factors(\n num_latent_channels, layer_name, expected_temporal_factor, expected_spatial_factor\n):\n patch_size = 4\n encoder_blocks = [\n (layer_name, {\"multiplier\": 2}),\n ]\n decoder_blocks = [\n (\"compress_all\", {\"residual\": True, \"multiplier\": 2}),\n ]\n config = {\n \"_class_name\": \"CausalVideoAutoencoder\",\n \"dims\": 3,\n \"encoder_blocks\": encoder_blocks,\n \"decoder_blocks\": decoder_blocks,\n \"latent_channels\": num_latent_channels,\n \"norm_layer\": \"pixel_norm\",\n \"patch_size\": patch_size,\n \"latent_log_var\": \"uniform\",\n \"use_quant_conv\": False,\n \"causal_decoder\": False,\n \"timestep_conditioning\": True,\n \"spatial_padding_mode\": \"replicate\",\n }\n model = CausalVideoAutoencoder.from_config(config)\n assert model.temporal_downscale_factor == expected_temporal_factor\n assert model.spatial_downscale_factor == expected_spatial_factor * patch_size\n"
},
{
"path": "tests/utils/.gitattributes",
"content": "*.mp4 filter=lfs diff=lfs merge=lfs -text\n"
}
]