indigo/LTX-2
1
0
Fork 0
Code Issues Pull Requests 1 Actions Packages Projects Releases Wiki Activity

Automated PR - 2026-01-05

Browse Source
This commit is contained in:
sync-bot
2026-01-05 20:10:38 +00:00
parent fc3b319d34
commit 9ce438b353
153 changed files with 28100 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files
.gitattributes
+9
View File
@@ -0,0 +1,9 @@
*.gif filter=lfs diff=lfs merge=lfs -text
*.safetensors filter=lfs diff=lfs merge=lfs -text
*.sft filter=lfs diff=lfs merge=lfs -text
*.pt filter=lfs diff=lfs merge=lfs -text
*.mp4 filter=lfs diff=lfs merge=lfs -text
*.png filter=lfs diff=lfs merge=lfs -text
*.jpeg filter=lfs diff=lfs merge=lfs -text
*.jpg filter=lfs diff=lfs merge=lfs -text
*.webp filter=lfs diff=lfs merge=lfs -text
.gitignore
+41
View File
@@ -0,0 +1,41 @@
# Python-generated files
__pycache__/
*.py[oc]
build/
dist/
wheels/
checkpoints/
*.egg-info
# Virtual environments
.venv
.python-version
# IDE settings
.idea/
.vscode/
# Other files
.DS_Store
tmp
.wandb
# Model checkpoints
*.ckpt
*.pt
*.safetensors
*.sft
# Media files
*.gif
*.heic
*.heif
*.jpg
*.jpeg
*.json
*.m4a
*.mov
*.mp4
*.png
*.wav
*.webp
LICENSE
+381
View File
@@ -0,0 +1,381 @@
LTX-2 Community License Agreement
License date: January 5, 2026
By using or distributing any portion or element of LTX-2, you agree
to be bound by this Agreement.
1. Definitions.
"Agreement" means the terms and conditions for the license, use,
reproduction, and distribution of LTX-2 and the Complementary
Materials, as specified in this document.
"Control" means the direct or indirect ownership of more than
fifty percent (50%) of the voting securities or other ownership
interests, or the power to direct the management and policies of
such Entity through voting rights, contract, or otherwise.
"Data" means a collection of information and/or content extracted
from the dataset used with LTX-2, including to train, pretrain,
or otherwise evaluate LTX-2. The Data is not licensed under this
Agreement.
"Derivatives of LTX-2" means all modifications to LTX-2, works
based on LTX-2, or any other model which is created or initialized
by transfer of patterns of the weights, parameters, activations or
output of LTX-2, to the other model, in order to cause the other
model to perform similarly to LTX-2, including – but not limited
to - distillation methods entailing the use of intermediate data
representations or methods based on the generation of synthetic
data by LTX-2 for training the other model. For clarity, Derivatives
of LTX-2 include: (i) any fine-tuned or adapted weights, parameters,
or checkpoints derived from LTX-2; (ii) derivative model architectures
that incorporate or are based upon LTX-2's architecture; and
(iii) any modified or extended versions of the Complementary
Materials. All intellectual property rights in Derivatives of LTX-2
shall be subject to the terms of this Agreement, and you may not
claim exclusive ownership rights in any Derivatives of LTX-2 that
would restrict the rights granted herein.
"Entity" means any individual, corporation, partnership, limited
liability company, or other legal entity. For purposes of this
Agreement, an Entity shall be deemed to include, on an aggregative
basis, all subsidiaries, affiliates, and other companies under
common Control with such Entity. When determining whether an Entity
meets any threshold under this Agreement (including revenue
thresholds), all subsidiaries, affiliates, and companies under
common Control shall be considered collectively.
"Harm" includes but is not limited to physical, mental,
psychological, financial and reputational damage, pain, or loss.
"Licensor" or "Lightricks" means the owner that is granting the
license under this Agreement. For the purposes of this Agreement,
the Licensor is Lightricks Ltd.
"LTX-2" means the large language models, text/image/video/audio/3D
generation models, and multimodal large language models and their
software and algorithms, including trained model weights, parameters
(including optimizer states), machine-learning model code,
inference-enabling code, training-enabling code, fine-tuning
enabling code, accompanying source code, scripts, documentation,
tutorials, examples, and all other elements of the foregoing
distributed and made publicly available by Lightricks (including,
for example, at https://github.com/Lightricks/LTX-2) for the LTX-2
model released on January 5, 2026. This license is applicable to
all LTX-2 versions released since January 5, 2026, and all future
releases of LTX-2 under this license.
"Output" means the results of operating LTX-2 as embodied in
informational content resulting therefrom.
"you" (or "your") means an individual or legal Entity licensing
LTX-2 in accordance with this Agreement and/or making use of LTX-2
for whichever purpose and in any field of use, including usage of
LTX-2 in an end-use application - e.g. chatbot, translator, image
generator.
2. Grant of License. Subject to the terms and conditions of this
Agreement, you are granted a non-exclusive, worldwide,
non-transferable and royalty-free limited license under Licensor's
intellectual property or other rights owned by Licensor embodied
in LTX-2 to use, reproduce, prepare, distribute, publicly display,
publicly perform, sublicense, copy, create derivative works of,
and make modifications to LTX-2, for any purpose, subject to the
restrictions set forth in Attachment A; provided however, that
Entities with annual revenues of at least $10,000,000 (the
"Commercial Entities") are required to obtain a paid commercial
use license in order to use LTX-2 and Derivatives of LTX-2,
subject to the terms and provisions of a different license (the
"Commercial Use Agreement"), as will be provided by the Licensor.
Commercial Entities interested in such a commercial license are
required to contact Licensor. Any commercial use of LTX-2 or
Derivatives of LTX-2 by the Commercial Entities not in accordance
with this Agreement and/or the Commercial Use Agreement is strictly
prohibited and shall be deemed a material breach of this Agreement.
Such material breach will be subject, in addition to any license
fees owed to Licensor for the period such Commercial Entity used
LTX-2 (as will be determined by Licensor), to liquidated damages,
which will be paid to Licensor immediately upon demand, in an
amount equal to double the amount that would otherwise have been
paid by you for the relevant period of time. Such amount reflects
a reasonable estimation of the losses and administrative costs
incurred due to such breach. You agree and understand that this
remedy does not limit the Licensor's right to pursue other remedies
available at law or equity.
3. Distribution and Redistribution. You may host for third parties
remote access purposes (e.g. software-as-a-service), reproduce
and distribute copies of LTX-2 or Derivatives of LTX-2 thereof in
any medium, with or without modifications, provided that you meet
the following conditions:
(a) Use-based restrictions as referenced in paragraph 4 and all
provisions of Attachment A MUST be included as an enforceable
provision by you in any type of legal agreement (e.g. a
license) governing the use and/or distribution of LTX-2 or
Derivatives of LTX-2, and you shall give notice to subsequent
users you distribute to, that LTX-2 or Derivatives of LTX-2
are subject to paragraph 4 and Attachment A in their entirety,
including all use restrictions and acceptable use policies;
(b) You must provide any third party recipients of LTX-2 or
Derivatives of LTX-2 a copy of this Agreement, including all
attachments and use policies. Any Derivative of LTX-2 (as
defined in Section 1, including but not limited to fine-tuned
weights, modified training code, models trained on Outputs, or
any other derivative) must be distributed exclusively under
the terms of this Agreement with a complete copy of this
license included;
(c) You must cause any modified files to carry prominent notices
stating that you changed the files;
(d) You must retain all copyright, patent, trademark, and
attribution notices excluding those notices that do not
pertain to any part of LTX-2, Derivatives of LTX-2.
You may add your own copyright statement to your modifications and
may provide additional or different license terms and conditions -
respecting paragraph 3(a) - for use, reproduction, or distribution
of your modifications, or for any such Derivatives of LTX-2 as a
whole, provided your use, reproduction, and distribution of LTX-2
otherwise complies with the conditions stated in this Agreement,
and you provide a complete copy of this Agreement with any such
use, reproduction and distribution of LTX-2 and any Derivatives
thereof.
4. Use-based restrictions. The restrictions set forth in Attachment A
are considered Use-based restrictions. Therefore, you cannot use
LTX-2 and the Derivatives of LTX-2 in violation of the specified
restricted uses. You may use LTX-2 subject to this Agreement,
including only for lawful purposes and in accordance with the
Agreement. "Use" may include creating any content with, fine-tuning,
updating, running, training, evaluating and/or re-parametrizing
LTX-2. You shall require all of your users who use LTX-2 or a
Derivative of LTX-2 to comply with the terms of this paragraph 4.
5. The Output You Generate. Except as set forth herein, Licensor
claims no rights in the Output you generate using LTX-2. You are
accountable for input you insert into LTX-2, the Output you
generate and its subsequent uses. No use of the Output can
contravene any provision as stated in the Agreement.
6. Updates and Runtime Restrictions. To the maximum extent permitted
by law, Licensor reserves the right to restrict (remotely or
otherwise) usage of LTX-2 in violation of this Agreement, update
LTX-2 through electronic means, or modify the Output of LTX-2
based on updates. You shall undertake reasonable efforts to use
the latest version of LTX-2. Any use of the non-current version
of LTX-2 is done solely at your risk.
7. Export Controls and Sanctions Compliance. You acknowledge that
LTX-2, Derivatives of LTX-2 may be subject to export control laws
and regulations, including but not limited to the U.S. Export
Administration Regulations and sanctions programs administered by
the Office of Foreign Assets Control (OFAC). You represent and
warrant that you and any users of LTX-2 are not (i) located in,
organized under the laws of, or ordinarily resident in any country
or territory subject to comprehensive sanctions; (ii) identified
on any U.S. government restricted party list, including the
Specially Designated Nationals and Blocked Persons List; or
(iii) otherwise prohibited from receiving LTX-2 under applicable
law. You shall not export, re-export, or transfer LTX-2, directly
or indirectly, in violation of any applicable export control or
sanctions laws or regulations. You agree to comply with all
applicable trade control laws and shall indemnify and hold
Licensor harmless from any claims arising from your failure to
comply with such laws.
8. Trademarks and related. Nothing in this Agreement permits you to
make use of Licensor's trademarks, trade names, logos or to
otherwise suggest endorsement or misrepresent the relationship
between the parties; and any rights not expressly granted herein
are reserved by the Licensor.
9. Disclaimer of Warranty. Unless required by applicable law or
agreed to in writing, Licensor provides LTX-2 on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or
implied, including, without limitation, any warranties or
conditions of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS
FOR A PARTICULAR PURPOSE. You are solely responsible for
determining the appropriateness of using or redistributing LTX-2
and Derivatives of LTX-2 and assume any risks associated with
your exercise of permissions under this Agreement.
10. Limitation of Liability. In no event and under no legal theory,
whether in tort (including negligence), contract, or otherwise,
unless required by applicable law (such as deliberate and grossly
negligent acts) or agreed to in writing, shall Licensor be liable
to you for damages, including any direct, indirect, special,
incidental, or consequential damages of any character arising as
a result of this Agreement or out of the use or inability to use
LTX-2 (including but not limited to damages for loss of goodwill,
work stoppage, computer failure or malfunction, or any and all
other commercial damages or losses), even if Licensor has been
advised of the possibility of such damages.
11. Accepting Warranty or Additional Liability. While redistributing
LTX-2 and Derivatives of LTX-2, you may, provided you do not
violate the terms of this Agreement, choose to offer and charge
a fee for, acceptance of support, warranty, indemnity, or other
liability obligations. However, in accepting such obligations,
you may act only on your own behalf and on your sole
responsibility, not on behalf of Licensor, and only if you agree
to indemnify, defend, and hold Licensor harmless for any liability
incurred by, or claims asserted against Licensor, by reason of
your accepting any such warranty or additional liability.
12. Governing Law. This Agreement and all relations, disputes, claims
and other matters arising hereunder (including non-contractual
disputes or claims) will be governed exclusively by, and construed
exclusively in accordance with, the laws of the State of New York.
To the extent permitted by law, choice of laws rules and the
United Nations Convention on Contracts for the International Sale
of Goods will not apply. For the purposes of adjudicating any
action or proceeding to enforce the terms of this Agreement, you
hereby irrevocably consent to the exclusive jurisdiction of, and
venue in, the federal and state courts located in the County of
New York within the State of New York. The prevailing party in
any claim or dispute between the parties under this Agreement
will be entitled to reimbursement of its reasonable attorneys'
fees and costs. You hereby waive the right to a trial by jury,
to participate in a class or representative action (including in
arbitration), or to combine individual proceedings in court or
in arbitration without the consent of all parties.
13. Term and Termination. This Agreement is effective upon your
acceptance and continues until terminated. Licensor may terminate
this Agreement immediately upon written notice to you if you
breach any provision of this Agreement, including but not limited
to violations of the use restrictions in Attachment A or
unauthorized commercial use. Upon termination: (a) all rights
granted to you under this Agreement will immediately cease;
(b) you must immediately cease all use of LTX-2 and Derivatives
of LTX-2; (c) you must delete or destroy all copies of LTX-2
and Derivatives of LTX-2 in your possession or control; and
(d) you must notify any third parties to whom you distributed
LTX-2 or Derivatives of LTX-2 of the termination. Sections 8-13,
and Section 15 shall survive termination of this Agreement.
Termination does not relieve you of any obligations incurred
prior to termination, including payment obligations under
Section 2. In addition, if You commence a lawsuit or other
proceedings (including a cross-claim or counterclaim in a lawsuit)
against Licensor or any person or entity alleging that LTX-2 or
any Output, or any portion of any of the foregoing, infringe any
intellectual property or other right owned or licensable by you,
then all licenses granted to you under this Agreement shall
terminate as of the date such lawsuit or other proceeding is filed.
14. Disputes and Arbitration. All disputes arising in connection with
this Agreement shall be finally settled by arbitration under the
Rules of Arbitration of the International Chamber of Commerce
("ICC Rules"), by one (1) arbitrator appointed in accordance with
the ICC Rules. The seat of arbitration shall be New York, NY, USA,
and the proceedings shall be conducted in English. The arbitrator
shall be empowered to grant any relief that a court could grant.
Judgment on the arbitration award may be entered by any court
having jurisdiction thereof. Each party waives its right to a
trial by jury and to participate in any class or representative
action.
15. If any provision of this Agreement is held to be
invalid, illegal
or unenforceable, the remaining provisions shall be unaffected
thereby and remain valid as if such provision had not been set
forth herein.
END OF TERMS AND CONDITIONS
ATTACHMENT A: Use Restrictions
When using the Outputs, LTX-2 and any Derivatives thereof, you
will comply with the Acceptable Use Policy. In addition, you
agree not to use the Outputs, LTX-2 or its Derivatives in any
of the following ways:
1. In any way that violates any applicable national, federal,
state, local or international law or regulation;
2. For the purpose of exploiting, Harming or attempting to
exploit or Harm minors in any way;
3. To generate or disseminate false information and/or content
with the purpose of Harming others;
4. To generate or disseminate personal identifiable information
that can be used to Harm an individual;
5. To generate or disseminate information and/or content (e.g.
images, code, posts, articles), and place the information
and/or content in any context (e.g. bot generating tweets)
without expressly and intelligibly disclaiming that the
information and/or content is machine generated;
6. To defame, disparage or otherwise harass others;
7. To impersonate or attempt to impersonate (e.g. deepfakes)
others without their consent;
8. For fully automated decision making that adversely impacts an
individual's legal rights or otherwise creates or modifies a
binding, enforceable obligation;
9. For any use intended to or which has the effect of
discriminating against or Harming individuals or groups based
on online or offline social behavior or known or predicted
personal or personality characteristics;
10. To exploit any of the vulnerabilities of a specific group of
persons based on their age, social, physical or mental
characteristics, in order to materially distort the behavior
of a person pertaining to that group in a manner that causes
or is likely to cause that person or another person physical
or psychological Harm;
11. For any use intended to or which has the effect of
discriminating against individuals or groups based on legally
protected characteristics or categories;
12. To provide medical advice and medical results interpretation;
13. To generate or disseminate information for the purpose to be
used for administration of justice, law enforcement,
immigration or asylum processes, such as predicting an
individual will commit fraud/crime commitment (e.g. by text
profiling, drawing causal relationships between assertions
made in documents, indiscriminate and arbitrarily-targeted use);
14. To generate and/or disseminate malware (including – but not
limited to – ransomware) or any other content to be used for
the purpose of harming electronic systems;
15. To engage in, promote, incite, or facilitate discrimination
or other unlawful or harmful conduct in the provision of
employment, employment benefits, credit, housing, or other
essential goods and services;
16. To engage in, promote, incite, or facilitate the harassment,
abuse, threatening, or bullying of individuals or groups of
individuals;
17. For military, warfare, nuclear industries or applications,
weapons development, or any use in connection with activities
that may cause death, personal injury, or severe physical or
environmental damage;
18. For commercial use only: To train, improve, or fine-tune any
other machine learning model, artificial intelligence system,
or competing model, except for Derivatives of LTX-2 as
expressly permitted under this Agreement;
19. To circumvent, disable, or interfere with any technical
limitations, safety features, content filters, or use
restrictions implemented in LTX-2 by Licensor;
20. To use LTX-2 or Derivatives of LTX-2 in any product, service,
or application that directly competes with Licensor's
commercial products or services, or is designed to replace or
substitute Licensor's offerings in the market, without
obtaining a separate commercial license from Licensor.
README.md
+115
View File
@@ -1 +1,116 @@
# LTX-2
[![Website](https://img.shields.io/badge/Website-LTX-181717?logo=google-chrome)](https://ltx.io)
[![Model](https://img.shields.io/badge/HuggingFace-Model-orange?logo=huggingface)](https://huggingface.co/Lightricks/LTX-2)
[![Demo](https://img.shields.io/badge/Demo-Try%20Now-brightgreen?logo=vercel)](https://app.ltx.studio/ltx-2-playground/i2v)
[![Paper](https://img.shields.io/badge/Paper-arXiv-B31B1B?logo=arxiv)](https://arxiv.org/abs/2501.00103)
[![Discord](https://img.shields.io/badge/Join-Discord-5865F2?logo=discord)](https://discord.gg/ltxplatform)
**LTX-2** is the first DiT-based audio-video foundation 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.
https://github.com/user-attachments/assets/63f33254-b11b-4f73-81f4-b02cffcad2ed
## 🚀 Quick Start
```bash
# Clone the repository
git clone https://github.com/Lightricks/LTX-2.git
cd LTX-2
# Set up the environment
uv sync --frozen
source .venv/bin/activate
```
### Required Models
Download the following models from the [LTX-2 HuggingFace repository](https://huggingface.co/Lightricks/LTX-2):
**LTX-2 Model Checkpoint** (choose and download one of the following)
* [`ltx-2-19b-dev-fp8.safetensors`](https://huggingface.co/Lightricks/LTX-2/blob/main/ltx-2-19b-dev-fp8.safetensors) - [Download](https://huggingface.co/Lightricks/LTX-2/resolve/main/ltx-2-19b-dev-fp8.safetensors)
* [`ltx-2-19b-dev.safetensors`](https://huggingface.co/Lightricks/LTX-2/blob/main/ltx-2-19b-dev.safetensors) - [Download](https://huggingface.co/Lightricks/LTX-2/resolve/main/ltx-2-19b-dev.safetensors)
* [`ltx-2-19b-distilled.safetensors`](https://huggingface.co/Lightricks/LTX-2/blob/main/ltx-2-19b-distilled.safetensors) - [Download](https://huggingface.co/Lightricks/LTX-2/resolve/main/ltx-2-19b-distilled.safetensors)
* [`ltx-2-19b-distilled-fp8.safetensors`](https://huggingface.co/Lightricks/LTX-2/blob/main/ltx-2-19b-distilled-fp8.safetensors) - [Download](https://huggingface.co/Lightricks/LTX-2/resolve/main/ltx-2-19b-distilled-fp8.safetensors)
**Spatial Upscaler** - Required for current two-stage pipeline implementations in this repository
* [`ltx-2-spatial-upscaler-x2-1.0.safetensors`](https://huggingface.co/Lightricks/LTX-2/blob/main/ltx-2-spatial-upscaler-x2-1.0.safetensors) - [Download](https://huggingface.co/Lightricks/LTX-2/resolve/main/ltx-2-spatial-upscaler-x2-1.0.safetensors)
**Temporal Upscaler** - Required for current two-stage pipeline implementations in this repository
* [`ltx-2-temporal-upscaler-x2-1.0.safetensors`](https://huggingface.co/Lightricks/LTX-2/blob/main/ltx-2-temporal-upscaler-x2-1.0.safetensors) - [Download](https://huggingface.co/Lightricks/LTX-2/resolve/main/ltx-2-temporal-upscaler-x2-1.0.safetensors)
**Distilled LoRA** - Required for current two-stage pipeline implementations in this repository (except DistilledPipeline and ICLoraPipeline)
* [`ltx-2-19b-distilled-lora-384.safetensors`](https://huggingface.co/Lightricks/LTX-2/blob/main/ltx-2-19b-distilled-lora-384.safetensors) - [Download](https://huggingface.co/Lightricks/LTX-2/resolve/main/ltx-2-19b-distilled-lora-384.safetensors)
**Gemma Text Encoder** (download all assets from the repository)
* [`Gemma 3`](https://huggingface.co/google/gemma-3-12b-it-qat-q4_0-unquantized/tree/main)
**LoRAs**
* [`LTX-2-19b-IC-LoRA-Canny-Control`](https://huggingface.co/Lightricks/LTX-2-19b-IC-LoRA-Canny-Control) - [Download](https://huggingface.co/Lightricks/LTX-2-19b-IC-LoRA-Canny-Control/resolve/main/ltx-2-19b-ic-lora-canny-control.safetensors)
* [`LTX-2-19b-IC-LoRA-Depth-Control`](https://huggingface.co/Lightricks/LTX-2-19b-IC-LoRA-Depth-Control) - [Download](https://huggingface.co/Lightricks/LTX-2-19b-IC-LoRA-Depth-Control/resolve/main/ltx-2-19b-ic-lora-depth-control.safetensors)
* [`LTX-2-19b-IC-LoRA-Detailer`](https://huggingface.co/Lightricks/LTX-2-19b-IC-LoRA-Detailer) - [Download](https://huggingface.co/Lightricks/LTX-2-19b-IC-LoRA-Detailer/resolve/main/ltx-2-19b-ic-lora-detailer.safetensors)
* [`LTX-2-19b-IC-LoRA-Pose-Control`](https://huggingface.co/Lightricks/LTX-2-19b-IC-LoRA-Pose-Control) - [Download](https://huggingface.co/Lightricks/LTX-2-19b-IC-LoRA-Pose-Control/resolve/main/ltx-2-19b-ic-lora-pose-control.safetensors)
* [`LTX-2-19b-LoRA-Camera-Control-Dolly-In`](https://huggingface.co/Lightricks/LTX-2-19b-LoRA-Camera-Control-Dolly-In) - [Download](https://huggingface.co/Lightricks/LTX-2-19b-LoRA-Camera-Control-Dolly-In/resolve/main/ltx-2-19b-lora-camera-control-dolly-in.safetensors)
* [`LTX-2-19b-LoRA-Camera-Control-Dolly-Left`](https://huggingface.co/Lightricks/LTX-2-19b-LoRA-Camera-Control-Dolly-Left) - [Download](https://huggingface.co/Lightricks/LTX-2-19b-LoRA-Camera-Control-Dolly-Left/resolve/main/ltx-2-19b-lora-camera-control-dolly-left.safetensors)
* [`LTX-2-19b-LoRA-Camera-Control-Dolly-Out`](https://huggingface.co/Lightricks/LTX-2-19b-LoRA-Camera-Control-Dolly-Out) - [Download](https://huggingface.co/Lightricks/LTX-2-19b-LoRA-Camera-Control-Dolly-Out/resolve/main/ltx-2-19b-lora-camera-control-dolly-out.safetensors)
* [`LTX-2-19b-LoRA-Camera-Control-Dolly-Right`](https://huggingface.co/Lightricks/LTX-2-19b-LoRA-Camera-Control-Dolly-Right) - [Download](https://huggingface.co/Lightricks/LTX-2-19b-LoRA-Camera-Control-Dolly-Right/resolve/main/ltx-2-19b-lora-camera-control-dolly-right.safetensors)
* [`LTX-2-19b-LoRA-Camera-Control-Jib-Down`](https://huggingface.co/Lightricks/LTX-2-19b-LoRA-Camera-Control-Jib-Down) - [Download](https://huggingface.co/Lightricks/LTX-2-19b-LoRA-Camera-Control-Jib-Down/resolve/main/ltx-2-19b-lora-camera-control-jib-down.safetensors)
* [`LTX-2-19b-LoRA-Camera-Control-Jib-Up`](https://huggingface.co/Lightricks/LTX-2-19b-LoRA-Camera-Control-Jib-Up) - [Download](https://huggingface.co/Lightricks/LTX-2-19b-LoRA-Camera-Control-Jib-Up/resolve/main/ltx-2-19b-lora-camera-control-jib-up.safetensors)
* [`LTX-2-19b-LoRA-Camera-Control-Static`](https://huggingface.co/Lightricks/LTX-2-19b-LoRA-Camera-Control-Static) - [Download](https://huggingface.co/Lightricks/LTX-2-19b-LoRA-Camera-Control-Static/resolve/main/ltx-2-19b-lora-camera-control-static.safetensors)
### Available Pipelines
* **[TI2VidTwoStagesPipeline](packages/ltx-pipelines/src/ltx_pipelines/ti2vid_two_stages.py)** - Production-quality text-to-video with 2x upsampling (recommended)
* **[TI2VidOneStagePipeline](packages/ltx-pipelines/src/ltx_pipelines/ti2vid_one_stage.py)** - Single-stage generation for quick prototyping
* **[DistilledPipeline](packages/ltx-pipelines/src/ltx_pipelines/distilled.py)** - Fastest inference with 8 predefined sigmas
* **[ICLoraPipeline](packages/ltx-pipelines/src/ltx_pipelines/ic_lora.py)** - Video-to-video and image-to-video transformations
* **[KeyframeInterpolationPipeline](packages/ltx-pipelines/src/ltx_pipelines/keyframe_interpolation.py)** - Interpolate between keyframe images
### ⚡ Optimization Tips
* **Use DistilledPipeline** - Fastest inference with only 8 predefined sigmas (8 steps stage 1, 4 steps stage 2)
* **Enable FP8 transformer** - Enables lower memory footprint: `--enable-fp8` (CLI) or `fp8transformer=True` (Python)
* **Install attention optimizations** - Use xFormers (`uv sync --extra xformers`) or [Flash Attention 3](https://github.com/Dao-AILab/flash-attention) for Hopper GPUs
* **Use gradient estimation** - Reduce inference steps from 40 to 20-30 while maintaining quality (see [pipeline documentation](packages/ltx-pipelines/README.md#denoising-loop-optimization))
* **Skip memory cleanup** - If you have sufficient VRAM, disable automatic memory cleanup between stages for faster processing
* **Choose single-stage pipeline** - Use `TI2VidOneStagePipeline` for faster generation when high resolution isn't required
## ✍️ Prompting for LTX-2
When 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:
- Start with main action in a single sentence
- Add specific details about movements and gestures
- Describe character/object appearances precisely
- Include background and environment details
- Specify camera angles and movements
- Describe lighting and colors
- Note any changes or sudden events
For additional guidance on writing a prompt please refer to <https://ltx.video/blog/how-to-prompt-for-ltx-2>
### Automatic Prompt Enhancement
LTX-2 pipelines support automatic prompt enhancement via an `enhance_prompt` parameter.
## 🔌 ComfyUI Integration
To use our model with ComfyUI, please follow the instructions at <https://github.com/Lightricks/ComfyUI-LTXVideo/>.
## 📦 Packages
This repository is organized as a monorepo with three main packages:
* **[ltx-core](packages/ltx-core/)** - Core model implementation, inference stack, and utilities
* **[ltx-pipelines](packages/ltx-pipelines/)** - High-level pipeline implementations for text-to-video, image-to-video, and other generation modes
* **[ltx-trainer](packages/ltx-trainer/)** - Training and fine-tuning tools for LoRA, full fine-tuning, and IC-LoRA
Each package has its own README and documentation. See the [Documentation](#-documentation) section below.
## 📚 Documentation
Each package includes comprehensive documentation:
* **[LTX-Core README](packages/ltx-core/README.md)** - Core model implementation, inference stack, and utilities
* **[LTX-Pipelines README](packages/ltx-pipelines/README.md)** - High-level pipeline implementations and usage guides
* **[LTX-Trainer README](packages/ltx-trainer/README.md)** - Training and fine-tuning documentation with detailed guides
packages/ltx-core/README.md
+280
View File
@@ -0,0 +1,280 @@
# LTX-Core
The foundational library for the LTX-2 Audio-Video generation model. This package contains the raw model definitions, component implementations, and loading logic used by `ltx-pipelines` and `ltx-trainer`.
## 📦 What's Inside?
- **`components/`**: Modular diffusion components (Schedulers, Guiders, Noisers, Patchifiers) following standard protocols
- **`conditioning/`**: Tools for preparing latent states and applying conditioning (image, video, keyframes)
- **`guidance/`**: Perturbation system for fine-grained control over attention mechanisms
- **`loader/`**: Utilities for loading weights from `.safetensors`, fusing LoRAs, and managing memory
- **`model/`**: PyTorch implementations of the LTX-2 Transformer, Video VAE, Audio VAE, Vocoder and Upscaler
- **`text_encoders/gemma`**: Gemma text encoder implementation with tokenizers, feature extractors, and separate encoders for audio-video and video-only generation
## 🚀 Quick Start
`ltx-core` provides the building blocks (models, components, and utilities) needed to construct inference flows. For ready-made inference pipelines use [`ltx-pipelines`](../ltx-pipelines/) or [`ltx-trainer`](../ltx-trainer/) for training.
## 🔧 Installation
```bash
# From the repository root
uv sync --frozen
# Or install as a package
pip install -e packages/ltx-core
```
## Building Blocks Overview
`ltx-core` provides modular components that can be combined to build custom inference flows:
### Core Models
- **Transformer** ([`model/transformer/`](src/ltx_core/model/transformer/)): The 48-layer LTX-2 transformer with cross-modal attention for joint audio-video processing. Expects inputs in [`Modality`](src/ltx_core/model/transformer/modality.py) format
- **Video VAE** ([`model/video_vae/`](src/ltx_core/model/video_vae/)): Encodes/decodes video pixels to/from latent space with temporal and spatial compression
- **Audio VAE** ([`model/audio_vae/`](src/ltx_core/model/audio_vae/)): Encodes/decodes audio spectrograms to/from latent space
- **Vocoder** ([`model/audio_vae/`](src/ltx_core/model/audio_vae/)): Neural vocoder that converts mel spectrograms to audio waveforms
- **Text Encoder** ([`text_encoders/`](src/ltx_core/text_encoders/)): Gemma-based encoder that produces separate embeddings for video and audio conditioning
- **Spatial Upscaler** ([`model/upsampler/`](src/ltx_core/model/upsampler/)): Upsamples latent representations for higher-resolution generation
### Diffusion Components
- **Schedulers** ([`components/schedulers.py`](src/ltx_core/components/schedulers.py)): Noise schedules (LTX2Scheduler, LinearQuadratic, Beta) that control the denoising process
- **Guiders** ([`components/guiders.py`](src/ltx_core/components/guiders.py)): Guidance strategies (CFG, STG, APG) for controlling generation quality and adherence to prompts
- **Noisers** ([`components/noisers.py`](src/ltx_core/components/noisers.py)): Add noise to latents according to the diffusion schedule
- **Patchifiers** ([`components/patchifiers.py`](src/ltx_core/components/patchifiers.py)): Convert between spatial latents `[B, C, F, H, W]` and sequence format `[B, seq_len, dim]` for transformer processing
### Conditioning & Control
- **Conditioning** ([`conditioning/`](src/ltx_core/conditioning/)): Tools for preparing and applying various conditioning types (image, video, keyframes)
- **Guidance** ([`guidance/`](src/ltx_core/guidance/)): Perturbation system for fine-grained control over attention mechanisms (e.g., skipping specific attention layers)
### Utilities
- **Loader** ([`loader/`](src/ltx_core/loader/)): Model loading from `.safetensors`, LoRA fusion, weight remapping, and memory management
For complete, production-ready pipeline implementations that combine these building blocks, see the [`ltx-pipelines`](../ltx-pipelines/) package.
---
# Architecture Overview
This section provides a deep dive into the internal architecture of the LTX-2 Audio-Video generation model.
## Table of Contents
1. [High-Level Architecture](#high-level-architecture)
2. [The Transformer](#the-transformer)
3. [Video VAE](#video-vae)
4. [Audio VAE](#audio-vae)
5. [Text Encoding (Gemma)](#text-encoding-gemma)
6. [Spatial Upscaler](#spatial-upsampler)
7. [Data Flow](#data-flow)
---
## High-Level Architecture
LTX-2 is a **joint Audio-Video diffusion transformer** that processes both modalities simultaneously in a unified architecture. Unlike traditional models that handle video and audio separately, LTX-2 uses cross-modal attention to enable natural synchronization.
```text
┌─────────────────────────────────────────────────────────────┐
│ INPUT PREPARATION │
│ │
│ Video Pixels → Video VAE Encoder → Video Latents │
│ Audio Waveform → Audio VAE Encoder → Audio Latents │
│ Text Prompt → Gemma Encoder → Text Embeddings │
└─────────────────────────────────────────────────────────────┘
↓
┌─────────────────────────────────────────────────────────────┐
│ LTX-2 TRANSFORMER (48 Blocks) │
│ │
│ ┌──────────────┐ ┌──────────────┐ │
│ │ Video Stream │ │ Audio Stream │ │
│ │ │ │ │ │
│ │ Self-Attn │ │ Self-Attn │ │
│ │ Cross-Attn │ │ Cross-Attn │ │
│ │ │◄────────────►│ │ │
│ │ A↔V Cross │ │ A↔V Cross │ │
│ │ Feed-Forward │ │ Feed-Forward │ │
│ └──────────────┘ └──────────────┘ │
└─────────────────────────────────────────────────────────────┘
↓
┌─────────────────────────────────────────────────────────────┐
│ OUTPUT DECODING │
│ │
│ Video Latents → Video VAE Decoder → Video Pixels │
│ Audio Latents → Audio VAE Decoder → Mel Spectrogram │
│ Mel Spectrogram → Vocoder → Audio Waveform │
└─────────────────────────────────────────────────────────────┘
```
---
## The Transformer
The core of LTX-2 is a 48-layer transformer that processes both video and audio tokens simultaneously.
### Model Structure
**Source**: [`src/ltx_core/model/transformer/model.py`](src/ltx_core/model/transformer/model.py)
The `LTXModel` class implements the transformer. It supports both video-only and audio-video generation modes. For actual usage, see the [`ltx-pipelines`](../ltx-pipelines/) package which handles model loading and initialization.
### Transformer Block Architecture
**Source**: [`src/ltx_core/model/transformer/transformer.py`](src/ltx_core/model/transformer/transformer.py)
```text
┌─────────────────────────────────────────────────────────────┐
│ TRANSFORMER BLOCK │
│ │
│ VIDEO PATH: │
│ Input → RMSNorm → AdaLN → Self-Attn (attn1) │
│ → RMSNorm → Cross-Attn (attn2, text) │
│ → RMSNorm → AdaLN → A↔V Cross-Attn │
│ → RMSNorm → AdaLN → Feed-Forward (ff) → Output │
│ │
│ AUDIO PATH: │
│ Input → RMSNorm → AdaLN → Self-Attn (audio_attn1) │
│ → RMSNorm → Cross-Attn (audio_attn2, text) │
│ → RMSNorm → AdaLN → A↔V Cross-Attn │
│ → RMSNorm → AdaLN → Feed-Forward (audio_ff) │
│ │
│ AdaLN (Adaptive Layer Normalization): │
│ - Uses scale_shift_table (6 params) for video/audio │
│ - Uses scale_shift_table_a2v_ca (5 params) for A↔V CA │
│ - Conditioned on per-token timestep embeddings │
└─────────────────────────────────────────────────────────────┘
```
### Perturbations
The transformer supports [**perturbations**](src/ltx_core/guidance/perturbations.py) that selectively skip attention operations.
Perturbations allow you to disable specific attention mechanisms during inference, which is useful for guidance techniques like STG (Spatio-Temporal Guidance).
**Supported Perturbation Types**:
- `SKIP_VIDEO_SELF_ATTN`: Skip video self-attention
- `SKIP_AUDIO_SELF_ATTN`: Skip audio self-attention
- `SKIP_A2V_CROSS_ATTN`: Skip audio-to-video cross-attention
- `SKIP_V2A_CROSS_ATTN`: Skip video-to-audio cross-attention
Perturbations are used internally by guidance mechanisms like STG (Spatio-Temporal Guidance). For usage examples, see the [`ltx-pipelines`](../ltx-pipelines/) package.
---
## Video VAE
The Video VAE ([`src/ltx_core/model/video_vae/`](src/ltx_core/model/video_vae/)) encodes video pixels into latent representations and decodes them back.
### Architecture
- **Encoder**: Compresses `[B, 3, F, H, W]` pixels → `[B, 128, F', H/32, W/32]` latents
- Where `F' = 1 + (F-1)/8` (frame count must satisfy `(F-1) % 8 == 0`)
- Example: `[B, 3, 33, 512, 512]` → `[B, 128, 5, 16, 16]`
- **Decoder**: Expands `[B, 128, F, H, W]` latents → `[B, 3, F', H*32, W*32]` pixels
- Where `F' = 1 + (F-1)*8`
- Example: `[B, 128, 5, 16, 16]` → `[B, 3, 33, 512, 512]`
The Video VAE is used internally by pipelines for encoding video pixels to latents and decoding latents back to pixels. For usage examples, see the [`ltx-pipelines`](../ltx-pipelines/) package.
---
## Audio VAE
The Audio VAE ([`src/ltx_core/model/audio_vae/`](src/ltx_core/model/audio_vae/)) processes audio spectrograms.
### Audio VAE Architecture
- **Encoder**: Compresses mel spectrogram `[B, mel_bins, T]` → `[B, 8, T/4, 16]` latents
- Temporal downsampling: 4× (`LATENT_DOWNSAMPLE_FACTOR = 4`)
- Frequency bins: Fixed 16 mel bins in latent space
- Latent channels: 8
- **Decoder**: Expands `[B, 8, T, 16]` latents → mel spectrogram `[B, mel_bins, T*4]`
- **Vocoder**: Converts mel spectrogram → audio waveform
**Downsampling**:
- Temporal: 4× (time steps)
- Frequency: Variable (input mel_bins → fixed 16 in latent space)
The Audio VAE is used internally by pipelines for encoding mel spectrograms to latents and decoding latents back to mel spectrograms. The vocoder converts mel spectrograms to audio waveforms. For usage examples, see the [`ltx-pipelines`](../ltx-pipelines/) package.
---
## Text Encoding (Gemma)
LTX-2 uses **Gemma** (Google's open LLM) as the text encoder, located in [`src/ltx_core/text_encoders/gemma/`](src/ltx_core/text_encoders/gemma/).
### Text Encoder Architecture
- **Tokenizer**: Converts text → token IDs
- **Gemma Model**: Processes tokens → embeddings
- **Text Projection**: Uses `PixArtAlphaTextProjection` to project caption embeddings
- Two-layer MLP with GELU (tanh approximation) or SiLU activation
- Projects from caption channels (3840) to model dimensions
- **Feature Extractor**: Extracts video/audio-specific embeddings
- **Separate Encoders**:
- `AVEncoder`: For audio-video generation (outputs separate video and audio contexts)
- `VideoOnlyEncoder`: For video-only generation
### System Prompts
System prompts are also used to enhance user's prompts.
- **Text-to-Video**: [`gemma_t2v_system_prompt.txt`](src/ltx_core/text_encoders/gemma/encoders/prompts/gemma_t2v_system_prompt.txt)
- **Image-to-Video**: [`gemma_i2v_system_prompt.txt`](src/ltx_core/text_encoders/gemma/encoders/prompts/gemma_i2v_system_prompt.txt)
**Important**: Video and audio receive **different** context embeddings, even from the same prompt. This allows better modality-specific conditioning.
**Output Format**:
- Video context: `[B, seq_len, 4096]` - Video-specific text embeddings
- Audio context: `[B, seq_len, 2048]` - Audio-specific text embeddings
The text encoder is used internally by pipelines. For usage examples, see the [`ltx-pipelines`](../ltx-pipelines/) package.
---
## Upscaler
The Upscaler ([`src/ltx_core/model/upsampler/`](src/ltx_core/model/upsampler/)) upsamples latent representations for higher-resolution output.
The spatial upsampler is used internally by two-stage pipelines (e.g., [`TI2VidTwoStagesPipeline`](../ltx-pipelines/src/ltx_pipelines/ti2vid_two_stages.py), [`ICLoraPipeline`](../ltx-pipelines/src/ltx_pipelines/ic_lora.py)) to upsample low-resolution latents before final VAE decoding. For usage examples, see the [`ltx-pipelines`](../ltx-pipelines/) package.
---
## Data Flow
### Complete Generation Pipeline
Here's how all the components work together conceptually ([`src/ltx_core/components/`](src/ltx_core/components/)):
**Pipeline Steps**:
1. **Text Encoding**: Text prompt → Gemma encoder → separate video/audio embeddings
2. **Latent Initialization**: Initialize noise latents in spatial format `[B, C, F, H, W]`
3. **Patchification**: Convert spatial latents to sequence format `[B, seq_len, dim]` for transformer
4. **Sigma Schedule**: Generate noise schedule (adapts to token count)
5. **Denoising Loop**: Iteratively denoise using transformer predictions
- Create Modality inputs with per-token timesteps and RoPE positions
- Forward pass through transformer (conditional and unconditional for CFG)
- Apply guidance (CFG, STG, etc.)
- Update latents using diffusion step (Euler, etc.)
6. **Unpatchification**: Convert sequence back to spatial format
7. **VAE Decoding**: Decode latents to pixel space (with optional upsampling for two-stage)
- [`TI2VidTwoStagesPipeline`](../ltx-pipelines/src/ltx_pipelines/ti2vid_two_stages.py) - Two-stage text-to-video (recommended)
- [`ICLoraPipeline`](../ltx-pipelines/src/ltx_pipelines/ic_lora.py) - Video-to-video with IC-LoRA control
- [`DistilledPipeline`](../ltx-pipelines/src/ltx_pipelines/distilled.py) - Fast inference with distilled model
- [`KeyframeInterpolationPipeline`](../ltx-pipelines/src/ltx_pipelines/keyframe_interpolation.py) - Keyframe-based interpolation
See the [ltx-pipelines README](../ltx-pipelines/README.md) for usage examples.
## 🔗 Related Projects
- **[ltx-pipelines](../ltx-pipelines/)** - High-level pipeline implementations for text-to-video, image-to-video, and video-to-video
- **[ltx-trainer](../ltx-trainer/)** - Training and fine-tuning tools
packages/ltx-core/pyproject.toml
+37
View File
@@ -0,0 +1,37 @@
[project]
name = "ltx-core"
version = "1.0.0"
description = "Core implementation of Lightricks' LTX-2 model"
readme = "README.md"
requires-python = ">=3.10"
dependencies = [
"torch~=2.7",
"torchaudio",
"einops",
"numpy",
"transformers",
"safetensors",
"accelerate",
"scipy>=1.14",
]
[project.optional-dependencies]
xformers = ["xformers"]
[tool.uv.sources]
xformers = { index = "pytorch" }
[[tool.uv.index]]
name = "pytorch"
url = "https://download.pytorch.org/whl/cu129"
explicit = true
[build-system]
requires = ["uv_build>=0.9.8,<0.10.0"]
build-backend = "uv_build"
[dependency-groups]
dev = [
"scikit-image>=0.25.2",
]
packages/ltx-core/src/ltx_core/__init__.py
View File
packages/ltx-core/src/ltx_core/components/__init__.py
+10
View File
@@ -0,0 +1,10 @@
"""
Diffusion pipeline components.
Submodules:
diffusion_steps - Diffusion stepping algorithms (EulerDiffusionStep)
guiders - Guidance strategies (CFGGuider, STGGuider, APG variants)
noisers - Noise samplers (GaussianNoiser)
patchifiers - Latent patchification (VideoLatentPatchifier, AudioPatchifier)
protocols - Protocol definitions (Patchifier, etc.)
schedulers - Sigma schedulers (LTX2Scheduler, LinearQuadraticScheduler)
"""
packages/ltx-core/src/ltx_core/components/diffusion_steps.py
+22
View File
@@ -0,0 +1,22 @@
import torch
from ltx_core.components.protocols import DiffusionStepProtocol
from ltx_core.utils import to_velocity
class EulerDiffusionStep(DiffusionStepProtocol):
"""
First-order Euler method for diffusion sampling.
Takes a single step from the current noise level (sigma) to the next by
computing velocity from the denoised prediction and applying: sample + velocity * dt.
"""
def step(
self, sample: torch.Tensor, denoised_sample: torch.Tensor, sigmas: torch.Tensor, step_index: int
) -> torch.Tensor:
sigma = sigmas[step_index]
sigma_next = sigmas[step_index + 1]
dt = sigma_next - sigma
velocity = to_velocity(sample, sigma, denoised_sample)
return (sample.to(torch.float32) + velocity.to(torch.float32) * dt).to(sample.dtype)
packages/ltx-core/src/ltx_core/components/guiders.py
+198
View File
@@ -0,0 +1,198 @@
from dataclasses import dataclass
import torch
from ltx_core.components.protocols import GuiderProtocol
@dataclass(frozen=True)
class CFGGuider(GuiderProtocol):
"""
Classifier-free guidance (CFG) guider.
Computes the guidance delta as (scale - 1) * (cond - uncond), steering the
denoising process toward the conditioned prediction.
Attributes:
scale: Guidance strength. 1.0 means no guidance, higher values increase
adherence to the conditioning.
"""
scale: float
def delta(self, cond: torch.Tensor, uncond: torch.Tensor) -> torch.Tensor:
return (self.scale - 1) * (cond - uncond)
def enabled(self) -> bool:
return self.scale != 1.0
@dataclass(frozen=True)
class CFGStarRescalingGuider(GuiderProtocol):
"""
Calculates the CFG delta between conditioned and unconditioned samples.
To minimize offset in the denoising direction and move mostly along the
conditioning axis within the distribution, the unconditioned sample is
rescaled in accordance with the norm of the conditioned sample.
Attributes:
scale (float):
Global guidance strength. A value of 1.0 corresponds to no extra
guidance beyond the base model prediction. Values > 1.0 increase
the influence of the conditioned sample relative to the
unconditioned one.
"""
scale: float
def delta(self, cond: torch.Tensor, uncond: torch.Tensor) -> torch.Tensor:
rescaled_neg = projection_coef(cond, uncond) * uncond
return (self.scale - 1) * (cond - rescaled_neg)
def enabled(self) -> bool:
return self.scale != 1.0
@dataclass(frozen=True)
class STGGuider(GuiderProtocol):
"""
Calculates the STG delta between conditioned and perturbed denoised samples.
Perturbed samples are the result of the denoising process with perturbations,
e.g. attentions acting as passthrough for certain layers and modalities.
Attributes:
scale (float):
Global strength of the STG guidance. A value of 0.0 disables the
guidance. Larger values increase the correction applied in the
direction of (pos_denoised - perturbed_denoised).
"""
scale: float
def delta(self, pos_denoised: torch.Tensor, perturbed_denoised: torch.Tensor) -> torch.Tensor:
return self.scale * (pos_denoised - perturbed_denoised)
def enabled(self) -> bool:
return self.scale != 0.0
@dataclass(frozen=True)
class LtxAPGGuider(GuiderProtocol):
"""
Calculates the APG (adaptive projected guidance) delta between conditioned
and unconditioned samples.
To minimize offset in the denoising direction and move mostly along the
conditioning axis within the distribution, the (cond - uncond) delta is
decomposed into components parallel and orthogonal to the conditioned
sample. The `eta` parameter weights the parallel component, while `scale`
is applied to the orthogonal component. Optionally, a norm threshold can
be used to suppress guidance when the magnitude of the correction is small.
Attributes:
scale (float):
Strength applied to the component of the guidance that is orthogonal
to the conditioned sample. Controls how aggressively we move in
directions that change semantics but stay consistent with the
conditioning manifold.
eta (float):
Weight of the component of the guidance that is parallel to the
conditioned sample. A value of 1.0 keeps the full parallel
component; values in [0, 1] attenuate it, and values > 1.0 amplify
motion along the conditioning direction.
norm_threshold (float):
Minimum L2 norm of the guidance delta below which the guidance
can be reduced or ignored (depending on implementation).
This is useful for avoiding noisy or unstable updates when the
guidance signal is very small.
"""
scale: float
eta: float = 1.0
norm_threshold: float = 0.0
def delta(self, cond: torch.Tensor, uncond: torch.Tensor) -> torch.Tensor:
guidance = cond - uncond
if self.norm_threshold > 0:
ones = torch.ones_like(guidance)
guidance_norm = guidance.norm(p=2, dim=[-1, -2, -3], keepdim=True)
scale_factor = torch.minimum(ones, self.norm_threshold / guidance_norm)
guidance = guidance * scale_factor
proj_coeff = projection_coef(guidance, cond)
g_parallel = proj_coeff * cond
g_orth = guidance - g_parallel
g_apg = g_parallel * self.eta + g_orth
return g_apg * (self.scale - 1)
def enabled(self) -> bool:
return self.scale != 1.0
@dataclass(frozen=False)
class LegacyStatefulAPGGuider(GuiderProtocol):
"""
Calculates the APG (adaptive projected guidance) delta between conditioned
and unconditioned samples.
To minimize offset in the denoising direction and move mostly along the
conditioning axis within the distribution, the (cond - uncond) delta is
decomposed into components parallel and orthogonal to the conditioned
sample. The `eta` parameter weights the parallel component, while `scale`
is applied to the orthogonal component. Optionally, a norm threshold can
be used to suppress guidance when the magnitude of the correction is small.
Attributes:
scale (float):
Strength applied to the component of the guidance that is orthogonal
to the conditioned sample. Controls how aggressively we move in
directions that change semantics but stay consistent with the
conditioning manifold.
eta (float):
Weight of the component of the guidance that is parallel to the
conditioned sample. A value of 1.0 keeps the full parallel
component; values in [0, 1] attenuate it, and values > 1.0 amplify
motion along the conditioning direction.
norm_threshold (float):
Minimum L2 norm of the guidance delta below which the guidance
can be reduced or ignored (depending on implementation).
This is useful for avoiding noisy or unstable updates when the
guidance signal is very small.
momentum (float):
Exponential moving-average coefficient for accumulating guidance
over time. running_avg = momentum * running_avg + guidance
"""
scale: float
eta: float
norm_threshold: float = 5.0
momentum: float = 0.0
# it is user's responsibility not to use same APGGuider for several denoisings or different modalities
# in order not to share accumulated average across different denoisings or modalities
running_avg: torch.Tensor | None = None
def delta(self, cond: torch.Tensor, uncond: torch.Tensor) -> torch.Tensor:
guidance = cond - uncond
if self.momentum != 0:
if self.running_avg is None:
self.running_avg = guidance.clone()
else:
self.running_avg = self.momentum * self.running_avg + guidance
guidance = self.running_avg
if self.norm_threshold > 0:
ones = torch.ones_like(guidance)
guidance_norm = guidance.norm(p=2, dim=[-1, -2, -3], keepdim=True)
scale_factor = torch.minimum(ones, self.norm_threshold / guidance_norm)
guidance = guidance * scale_factor
proj_coeff = projection_coef(guidance, cond)
g_parallel = proj_coeff * cond
g_orth = guidance - g_parallel
g_apg = g_parallel * self.eta + g_orth
return g_apg * self.scale
def enabled(self) -> bool:
return self.scale != 0.0
def projection_coef(to_project: torch.Tensor, project_onto: torch.Tensor) -> torch.Tensor:
batch_size = to_project.shape[0]
positive_flat = to_project.reshape(batch_size, -1)
negative_flat = project_onto.reshape(batch_size, -1)
dot_product = torch.sum(positive_flat * negative_flat, dim=1, keepdim=True)
squared_norm = torch.sum(negative_flat**2, dim=1, keepdim=True) + 1e-8
return dot_product / squared_norm
packages/ltx-core/src/ltx_core/components/noisers.py
+35
View File
@@ -0,0 +1,35 @@
from dataclasses import replace
from typing import Protocol
import torch
from ltx_core.types import LatentState
class Noiser(Protocol):
"""Protocol for adding noise to a latent state during diffusion."""
def __call__(self, latent_state: LatentState, noise_scale: float) -> LatentState: ...
class GaussianNoiser(Noiser):
"""Adds Gaussian noise to a latent state, scaled by the denoise mask."""
def __init__(self, generator: torch.Generator):
super().__init__()
self.generator = generator
def __call__(self, latent_state: LatentState, noise_scale: float = 1.0) -> LatentState:
noise = torch.randn(
*latent_state.latent.shape,
device=latent_state.latent.device,
dtype=latent_state.latent.dtype,
generator=self.generator,
)
scaled_mask = latent_state.denoise_mask * noise_scale
latent = noise * scaled_mask + latent_state.latent * (1 - scaled_mask)
return replace(
latent_state,
latent=latent.to(latent_state.latent.dtype),
)
packages/ltx-core/src/ltx_core/components/patchifiers.py
+348
View File
@@ -0,0 +1,348 @@
import math
from typing import Optional, Tuple
import einops
import torch
from ltx_core.components.protocols import Patchifier
from ltx_core.types import AudioLatentShape, SpatioTemporalScaleFactors, VideoLatentShape
class VideoLatentPatchifier(Patchifier):
def __init__(self, patch_size: int):
# Patch sizes for video latents.
self._patch_size = (
1, # temporal dimension
patch_size, # height dimension
patch_size, # width dimension
)
@property
def patch_size(self) -> Tuple[int, int, int]:
return self._patch_size
def get_token_count(self, tgt_shape: VideoLatentShape) -> int:
return math.prod(tgt_shape.to_torch_shape()[2:]) // math.prod(self._patch_size)
def patchify(
self,
latents: torch.Tensor,
) -> torch.Tensor:
latents = einops.rearrange(
latents,
"b c (f p1) (h p2) (w p3) -> b (f h w) (c p1 p2 p3)",
p1=self._patch_size[0],
p2=self._patch_size[1],
p3=self._patch_size[2],
)
return latents
def unpatchify(
self,
latents: torch.Tensor,
output_shape: VideoLatentShape,
) -> torch.Tensor:
assert self._patch_size[0] == 1, "Temporal patch size must be 1 for symmetric patchifier"
patch_grid_frames = output_shape.frames // self._patch_size[0]
patch_grid_height = output_shape.height // self._patch_size[1]
patch_grid_width = output_shape.width // self._patch_size[2]
latents = einops.rearrange(
latents,
"b (f h w) (c p q) -> b c f (h p) (w q)",
f=patch_grid_frames,
h=patch_grid_height,
w=patch_grid_width,
p=self._patch_size[1],
q=self._patch_size[2],
)
return latents
def get_patch_grid_bounds(
self,
output_shape: AudioLatentShape | VideoLatentShape,
device: Optional[torch.device] = None,
) -> torch.Tensor:
"""
Return the per-dimension bounds [inclusive start, exclusive end) for every
patch produced by `patchify`. The bounds are expressed in the original
video grid coordinates: frame/time, height, and width.
The resulting tensor is shaped `[batch_size, 3, num_patches, 2]`, where:
- axis 1 (size 3) enumerates (frame/time, height, width) dimensions
- axis 3 (size 2) stores `[start, end)` indices within each dimension
Args:
output_shape: Video grid description containing frames, height, and width.
device: Device of the latent tensor.
"""
if not isinstance(output_shape, VideoLatentShape):
raise ValueError("VideoLatentPatchifier expects VideoLatentShape when computing coordinates")
frames = output_shape.frames
height = output_shape.height
width = output_shape.width
batch_size = output_shape.batch
# Validate inputs to ensure positive dimensions
assert frames > 0, f"frames must be positive, got {frames}"
assert height > 0, f"height must be positive, got {height}"
assert width > 0, f"width must be positive, got {width}"
assert batch_size > 0, f"batch_size must be positive, got {batch_size}"
# Generate grid coordinates for each dimension (frame, height, width)
# We use torch.arange to create the starting coordinates for each patch.
# indexing='ij' ensures the dimensions are in the order (frame, height, width).
grid_coords = torch.meshgrid(
torch.arange(start=0, end=frames, step=self._patch_size[0], device=device),
torch.arange(start=0, end=height, step=self._patch_size[1], device=device),
torch.arange(start=0, end=width, step=self._patch_size[2], device=device),
indexing="ij",
)
# Stack the grid coordinates to create the start coordinates tensor.
# Shape becomes (3, grid_f, grid_h, grid_w)
patch_starts = torch.stack(grid_coords, dim=0)
# Create a tensor containing the size of a single patch:
# (frame_patch_size, height_patch_size, width_patch_size).
# Reshape to (3, 1, 1, 1) to enable broadcasting when adding to the start coordinates.
patch_size_delta = torch.tensor(
self._patch_size,
device=patch_starts.device,
dtype=patch_starts.dtype,
).view(3, 1, 1, 1)
# Calculate end coordinates: start + patch_size
# Shape becomes (3, grid_f, grid_h, grid_w)
patch_ends = patch_starts + patch_size_delta
# Stack start and end coordinates together along the last dimension
# Shape becomes (3, grid_f, grid_h, grid_w, 2), where the last dimension is [start, end]
latent_coords = torch.stack((patch_starts, patch_ends), dim=-1)
# Broadcast to batch size and flatten all spatial/temporal dimensions into one sequence.
# Final Shape: (batch_size, 3, num_patches, 2)
latent_coords = einops.repeat(
latent_coords,
"c f h w bounds -> b c (f h w) bounds",
b=batch_size,
bounds=2,
)
return latent_coords
def get_pixel_coords(
latent_coords: torch.Tensor,
scale_factors: SpatioTemporalScaleFactors,
causal_fix: bool = False,
) -> torch.Tensor:
"""
Map latent-space `[start, end)` coordinates to their pixel-space equivalents by scaling
each axis (frame/time, height, width) with the corresponding VAE downsampling factors.
Optionally compensate for causal encoding that keeps the first frame at unit temporal scale.
Args:
latent_coords: Tensor of latent bounds shaped `(batch, 3, num_patches, 2)`.
scale_factors: SpatioTemporalScaleFactors tuple `(temporal, height, width)` with integer scale factors applied
per axis.
causal_fix: When True, rewrites the temporal axis of the first frame so causal VAEs
that treat frame zero differently still yield non-negative timestamps.
"""
# Broadcast the VAE scale factors so they align with the `(batch, axis, patch, bound)` layout.
broadcast_shape = [1] * latent_coords.ndim
broadcast_shape[1] = -1 # axis dimension corresponds to (frame/time, height, width)
scale_tensor = torch.tensor(scale_factors, device=latent_coords.device).view(*broadcast_shape)
# Apply per-axis scaling to convert latent bounds into pixel-space coordinates.
pixel_coords = latent_coords * scale_tensor
if causal_fix:
# VAE temporal stride for the very first frame is 1 instead of `scale_factors[0]`.
# Shift and clamp to keep the first-frame timestamps causal and non-negative.
pixel_coords[:, 0, ...] = (pixel_coords[:, 0, ...] + 1 - scale_factors[0]).clamp(min=0)
return pixel_coords
class AudioPatchifier(Patchifier):
def __init__(
self,
patch_size: int,
sample_rate: int = 16000,
hop_length: int = 160,
audio_latent_downsample_factor: int = 4,
is_causal: bool = True,
shift: int = 0,
):
"""
Patchifier tailored for spectrogram/audio latents.
Args:
patch_size: Number of mel bins combined into a single patch. This
controls the resolution along the frequency axis.
sample_rate: Original waveform sampling rate. Used to map latent
indices back to seconds so downstream consumers can align audio
and video cues.
hop_length: Window hop length used for the spectrogram. Determines
how many real-time samples separate two consecutive latent frames.
audio_latent_downsample_factor: Ratio between spectrogram frames and
latent frames; compensates for additional downsampling inside the
VAE encoder.
is_causal: When True, timing is shifted to account for causal
receptive fields so timestamps do not peek into the future.
shift: Integer offset applied to the latent indices. Enables
constructing overlapping windows from the same latent sequence.
"""
self.hop_length = hop_length
self.sample_rate = sample_rate
self.audio_latent_downsample_factor = audio_latent_downsample_factor
self.is_causal = is_causal
self.shift = shift
self._patch_size = (1, patch_size, patch_size)
@property
def patch_size(self) -> Tuple[int, int, int]:
return self._patch_size
def get_token_count(self, tgt_shape: AudioLatentShape) -> int:
return tgt_shape.frames
def _get_audio_latent_time_in_sec(
self,
start_latent: int,
end_latent: int,
dtype: torch.dtype,
device: Optional[torch.device] = None,
) -> torch.Tensor:
"""
Converts latent indices into real-time seconds while honoring causal
offsets and the configured hop length.
Args:
start_latent: Inclusive start index inside the latent sequence. This
sets the first timestamp returned.
end_latent: Exclusive end index. Determines how many timestamps get
generated.
dtype: Floating-point dtype used for the returned tensor, allowing
callers to control precision.
device: Target device for the timestamp tensor. When omitted the
computation occurs on CPU to avoid surprising GPU allocations.
"""
if device is None:
device = torch.device("cpu")
audio_latent_frame = torch.arange(start_latent, end_latent, dtype=dtype, device=device)
audio_mel_frame = audio_latent_frame * self.audio_latent_downsample_factor
if self.is_causal:
# Frame offset for causal alignment.
# The "+1" ensures the timestamp corresponds to the first sample that is fully available.
causal_offset = 1
audio_mel_frame = (audio_mel_frame + causal_offset - self.audio_latent_downsample_factor).clip(min=0)
return audio_mel_frame * self.hop_length / self.sample_rate
def _compute_audio_timings(
self,
batch_size: int,
num_steps: int,
device: Optional[torch.device] = None,
) -> torch.Tensor:
"""
Builds a `(B, 1, T, 2)` tensor containing timestamps for each latent frame.
This helper method underpins `get_patch_grid_bounds` for the audio patchifier.
Args:
batch_size: Number of sequences to broadcast the timings over.
num_steps: Number of latent frames (time steps) to convert into timestamps.
device: Device on which the resulting tensor should reside.
"""
resolved_device = device
if resolved_device is None:
resolved_device = torch.device("cpu")
start_timings = self._get_audio_latent_time_in_sec(
self.shift,
num_steps + self.shift,
torch.float32,
resolved_device,
)
start_timings = start_timings.unsqueeze(0).expand(batch_size, -1).unsqueeze(1)
end_timings = self._get_audio_latent_time_in_sec(
self.shift + 1,
num_steps + self.shift + 1,
torch.float32,
resolved_device,
)
end_timings = end_timings.unsqueeze(0).expand(batch_size, -1).unsqueeze(1)
return torch.stack([start_timings, end_timings], dim=-1)
def patchify(
self,
audio_latents: torch.Tensor,
) -> torch.Tensor:
"""
Flattens the audio latent tensor along time. Use `get_patch_grid_bounds`
to derive timestamps for each latent frame based on the configured hop
length and downsampling.
Args:
audio_latents: Latent tensor to patchify.
Returns:
Flattened patch tokens tensor. Use `get_patch_grid_bounds` to compute the
corresponding timing metadata when needed.
"""
audio_latents = einops.rearrange(
audio_latents,
"b c t f -> b t (c f)",
)
return audio_latents
def unpatchify(
self,
audio_latents: torch.Tensor,
output_shape: AudioLatentShape,
) -> torch.Tensor:
"""
Restores the `(B, C, T, F)` spectrogram tensor from flattened patches.
Use `get_patch_grid_bounds` to recompute the timestamps that describe each
frame's position in real time.
Args:
audio_latents: Latent tensor to unpatchify.
output_shape: Shape of the unpatched output tensor.
Returns:
Unpatched latent tensor. Use `get_patch_grid_bounds` to compute the timing
metadata associated with the restored latents.
"""
# audio_latents shape: (batch, time, freq * channels)
audio_latents = einops.rearrange(
audio_latents,
"b t (c f) -> b c t f",
c=output_shape.channels,
f=output_shape.mel_bins,
)
return audio_latents
def get_patch_grid_bounds(
self,
output_shape: AudioLatentShape | VideoLatentShape,
device: Optional[torch.device] = None,
) -> torch.Tensor:
"""
Return the temporal bounds `[inclusive start, exclusive end)` for every
patch emitted by `patchify`. For audio this corresponds to timestamps in
seconds aligned with the original spectrogram grid.
The returned tensor has shape `[batch_size, 1, time_steps, 2]`, where:
- axis 1 (size 1) represents the temporal dimension
- axis 3 (size 2) stores the `[start, end)` timestamps per patch
Args:
output_shape: Audio grid specification describing the number of time steps.
device: Target device for the returned tensor.
"""
if not isinstance(output_shape, AudioLatentShape):
raise ValueError("AudioPatchifier expects AudioLatentShape when computing coordinates")
return self._compute_audio_timings(output_shape.batch, output_shape.frames, device)
packages/ltx-core/src/ltx_core/components/protocols.py
+101
View File
@@ -0,0 +1,101 @@
from typing import Protocol, Tuple
import torch
from ltx_core.types import AudioLatentShape, VideoLatentShape
class Patchifier(Protocol):
"""
Protocol for patchifiers that convert latent tensors into patches and assemble them back.
"""
def patchify(
self,
latents: torch.Tensor,
) -> torch.Tensor:
...
"""
Convert latent tensors into flattened patch tokens.
Args:
latents: Latent tensor to patchify.
Returns:
Flattened patch tokens tensor.
"""
def unpatchify(
self,
latents: torch.Tensor,
output_shape: AudioLatentShape | VideoLatentShape,
) -> torch.Tensor:
"""
Converts latent tensors between spatio-temporal formats and flattened sequence representations.
Args:
latents: Patch tokens that must be rearranged back into the latent grid constructed by `patchify`.
output_shape: Shape of the output tensor. Note that output_shape is either AudioLatentShape or
VideoLatentShape.
Returns:
Dense latent tensor restored from the flattened representation.
"""
@property
def patch_size(self) -> Tuple[int, int, int]:
...
"""
Returns the patch size as a tuple of (temporal, height, width) dimensions
"""
def get_patch_grid_bounds(
self,
output_shape: AudioLatentShape | VideoLatentShape,
device: torch.device | None = None,
) -> torch.Tensor:
...
"""
Compute metadata describing where each latent patch resides within the
grid specified by `output_shape`.
Args:
output_shape: Target grid layout for the patches.
device: Target device for the returned tensor.
Returns:
Tensor containing patch coordinate metadata such as spatial or temporal intervals.
"""
class SchedulerProtocol(Protocol):
"""
Protocol for schedulers that provide a sigmas schedule tensor for a
given number of steps. Device is cpu.
"""
def execute(self, steps: int, **kwargs) -> torch.FloatTensor: ...
class GuiderProtocol(Protocol):
"""
Protocol for guiders that compute a delta tensor given conditioning inputs.
The returned delta should be added to the conditional output (cond), enabling
multiple guiders to be chained together by accumulating their deltas.
"""
scale: float
def delta(self, cond: torch.Tensor, uncond: torch.Tensor) -> torch.Tensor: ...
def enabled(self) -> bool:
"""
Returns whether the corresponding perturbation is enabled. E.g. for CFG, this should return False if the scale
is 1.0.
"""
...
class DiffusionStepProtocol(Protocol):
"""
Protocol for diffusion steps that provide a next sample tensor for a given current sample tensor,
current denoised sample tensor, and sigmas tensor.
"""
def step(
self, sample: torch.Tensor, denoised_sample: torch.Tensor, sigmas: torch.Tensor, step_index: int
) -> torch.Tensor: ...
packages/ltx-core/src/ltx_core/components/schedulers.py
+129
View File
@@ -0,0 +1,129 @@
import math
from functools import lru_cache
import numpy
import scipy
import torch
from ltx_core.components.protocols import SchedulerProtocol
BASE_SHIFT_ANCHOR = 1024
MAX_SHIFT_ANCHOR = 4096
class LTX2Scheduler(SchedulerProtocol):
"""
Default scheduler for LTX-2 diffusion sampling.
Generates a sigma schedule with token-count-dependent shifting and optional
stretching to a terminal value.
"""
def execute(
self,
steps: int,
latent: torch.Tensor | None = None,
max_shift: float = 2.05,
base_shift: float = 0.95,
stretch: bool = True,
terminal: float = 0.1,
**_kwargs,
) -> torch.FloatTensor:
tokens = math.prod(latent.shape[2:]) if latent is not None else MAX_SHIFT_ANCHOR
sigmas = torch.linspace(1.0, 0.0, steps + 1)
x1 = BASE_SHIFT_ANCHOR
x2 = MAX_SHIFT_ANCHOR
mm = (max_shift - base_shift) / (x2 - x1)
b = base_shift - mm * x1
sigma_shift = (tokens) * mm + b
power = 1
sigmas = torch.where(
sigmas != 0,
math.exp(sigma_shift) / (math.exp(sigma_shift) + (1 / sigmas - 1) ** power),
0,
)
# Stretch sigmas so that its final value matches the given terminal value.
if stretch:
non_zero_mask = sigmas != 0
non_zero_sigmas = sigmas[non_zero_mask]
one_minus_z = 1.0 - non_zero_sigmas
scale_factor = one_minus_z[-1] / (1.0 - terminal)
stretched = 1.0 - (one_minus_z / scale_factor)
sigmas[non_zero_mask] = stretched
return sigmas.to(torch.float32)
class LinearQuadraticScheduler(SchedulerProtocol):
"""
Scheduler with linear steps followed by quadratic steps.
Produces a sigma schedule that transitions linearly up to a threshold,
then follows a quadratic curve for the remaining steps.
"""
def execute(
self, steps: int, threshold_noise: float = 0.025, linear_steps: int | None = None, **_kwargs
) -> torch.FloatTensor:
if steps == 1:
return torch.FloatTensor([1.0, 0.0])
if linear_steps is None:
linear_steps = steps // 2
linear_sigma_schedule = [i * threshold_noise / linear_steps for i in range(linear_steps)]
threshold_noise_step_diff = linear_steps - threshold_noise * steps
quadratic_steps = steps - linear_steps
quadratic_sigma_schedule = []
if quadratic_steps > 0:
quadratic_coef = threshold_noise_step_diff / (linear_steps * quadratic_steps**2)
linear_coef = threshold_noise / linear_steps - 2 * threshold_noise_step_diff / (quadratic_steps**2)
const = quadratic_coef * (linear_steps**2)
quadratic_sigma_schedule = [
quadratic_coef * (i**2) + linear_coef * i + const for i in range(linear_steps, steps)
]
sigma_schedule = linear_sigma_schedule + quadratic_sigma_schedule + [1.0]
sigma_schedule = [1.0 - x for x in sigma_schedule]
return torch.FloatTensor(sigma_schedule)
class BetaScheduler(SchedulerProtocol):
"""
Scheduler using a beta distribution to sample timesteps.
Based on: https://arxiv.org/abs/2407.12173
"""
shift = 2.37
timesteps_length = 10000
def execute(self, steps: int, alpha: float = 0.6, beta: float = 0.6) -> torch.FloatTensor:
"""
Execute the beta scheduler.
Args:
steps: The number of steps to execute the scheduler for.
alpha: The alpha parameter for the beta distribution.
beta: The beta parameter for the beta distribution.
Warnings:
The number of steps within `sigmas` theoretically might be less than `steps+1`,
because of the deduplication of the identical timesteps
Returns:
A tensor of sigmas.
"""
model_sampling_sigmas = _precalculate_model_sampling_sigmas(self.shift, self.timesteps_length)
total_timesteps = len(model_sampling_sigmas) - 1
ts = 1 - numpy.linspace(0, 1, steps, endpoint=False)
ts = numpy.rint(scipy.stats.beta.ppf(ts, alpha, beta) * total_timesteps).tolist()
ts = list(dict.fromkeys(ts))
sigmas = [float(model_sampling_sigmas[int(t)]) for t in ts] + [0.0]
return torch.FloatTensor(sigmas)
@lru_cache(maxsize=5)
def _precalculate_model_sampling_sigmas(shift: float, timesteps_length: int) -> torch.Tensor:
timesteps = torch.arange(1, timesteps_length + 1, 1) / timesteps_length
return torch.Tensor([flux_time_shift(shift, 1.0, t) for t in timesteps])
def flux_time_shift(mu: float, sigma: float, t: float) -> float:
return math.exp(mu) / (math.exp(mu) + (1 / t - 1) ** sigma)
packages/ltx-core/src/ltx_core/conditioning/__init__.py
+12
View File
@@ -0,0 +1,12 @@
"""Conditioning utilities: latent state, tools, and conditioning types."""
from ltx_core.conditioning.exceptions import ConditioningError
from ltx_core.conditioning.item import ConditioningItem
from ltx_core.conditioning.types import VideoConditionByKeyframeIndex, VideoConditionByLatentIndex
__all__ = [
"ConditioningError",
"ConditioningItem",
"VideoConditionByKeyframeIndex",
"VideoConditionByLatentIndex",
]
packages/ltx-core/src/ltx_core/conditioning/exceptions.py
+4
View File
@@ -0,0 +1,4 @@
class ConditioningError(Exception):
"""
Class for conditioning-related errors.
"""
packages/ltx-core/src/ltx_core/conditioning/item.py
+20
View File
@@ -0,0 +1,20 @@
from typing import Protocol
from ltx_core.tools import LatentTools
from ltx_core.types import LatentState
class ConditioningItem(Protocol):
"""Protocol for conditioning items that modify latent state during diffusion."""
def apply_to(self, latent_state: LatentState, latent_tools: LatentTools) -> LatentState:
"""
Apply the conditioning to the latent state.
Args:
latent_state: The latent state to apply the conditioning to. This is state always patchified.
Returns:
The latent state after the conditioning has been applied.
IMPORTANT: If the conditioning needs to add extra tokens to the latent, it should add them to the end of the
latent.
"""
...
packages/ltx-core/src/ltx_core/conditioning/types/__init__.py
+9
View File
@@ -0,0 +1,9 @@
"""Conditioning type implementations."""
from ltx_core.conditioning.types.keyframe_cond import VideoConditionByKeyframeIndex
from ltx_core.conditioning.types.latent_cond import VideoConditionByLatentIndex
__all__ = [
"VideoConditionByKeyframeIndex",
"VideoConditionByLatentIndex",
]
packages/ltx-core/src/ltx_core/conditioning/types/keyframe_cond.py
+53
View File
@@ -0,0 +1,53 @@
import torch
from ltx_core.components.patchifiers import get_pixel_coords
from ltx_core.conditioning.item import ConditioningItem
from ltx_core.tools import VideoLatentTools
from ltx_core.types import LatentState, VideoLatentShape
class VideoConditionByKeyframeIndex(ConditioningItem):
"""
Conditions video generation on keyframe latents at a specific frame index.
Appends keyframe tokens to the latent state with positions offset by frame_idx,
and sets denoise strength according to the strength parameter.
"""
def __init__(self, keyframes: torch.Tensor, frame_idx: int, strength: float):
self.keyframes = keyframes
self.frame_idx = frame_idx
self.strength = strength
def apply_to(
self,
latent_state: LatentState,
latent_tools: VideoLatentTools,
) -> LatentState:
tokens = latent_tools.patchifier.patchify(self.keyframes)
latent_coords = latent_tools.patchifier.get_patch_grid_bounds(
output_shape=VideoLatentShape.from_torch_shape(self.keyframes.shape),
device=self.keyframes.device,
)
positions = get_pixel_coords(
latent_coords=latent_coords,
scale_factors=latent_tools.scale_factors,
causal_fix=latent_tools.causal_fix if self.frame_idx == 0 else False,
)
positions[:, 0, ...] += self.frame_idx
positions = positions.to(dtype=torch.float32)
positions[:, 0, ...] /= latent_tools.fps
denoise_mask = torch.full(
size=(*tokens.shape[:2], 1),
fill_value=1.0 - self.strength,
device=self.keyframes.device,
dtype=self.keyframes.dtype,
)
return LatentState(
latent=torch.cat([latent_state.latent, tokens], dim=1),
denoise_mask=torch.cat([latent_state.denoise_mask, denoise_mask], dim=1),
positions=torch.cat([latent_state.positions, positions], dim=2),
clean_latent=torch.cat([latent_state.clean_latent, tokens], dim=1),
)
packages/ltx-core/src/ltx_core/conditioning/types/latent_cond.py
+44
View File
@@ -0,0 +1,44 @@
import torch
from ltx_core.conditioning.exceptions import ConditioningError
from ltx_core.conditioning.item import ConditioningItem
from ltx_core.tools import LatentTools
from ltx_core.types import LatentState
class VideoConditionByLatentIndex(ConditioningItem):
"""
Conditions video generation by injecting latents at a specific latent frame index.
Replaces tokens in the latent state at positions corresponding to latent_idx,
and sets denoise strength according to the strength parameter.
"""
def __init__(self, latent: torch.Tensor, strength: float, latent_idx: int):
self.latent = latent
self.strength = strength
self.latent_idx = latent_idx
def apply_to(self, latent_state: LatentState, latent_tools: LatentTools) -> LatentState:
cond_batch, cond_channels, _, cond_height, cond_width = self.latent.shape
tgt_batch, tgt_channels, tgt_frames, tgt_height, tgt_width = latent_tools.target_shape.to_torch_shape()
if (cond_batch, cond_channels, cond_height, cond_width) != (tgt_batch, tgt_channels, tgt_height, tgt_width):
raise ConditioningError(
f"Can't apply image conditioning item to latent with shape {latent_tools.target_shape}, expected "
f"shape is ({tgt_batch}, {tgt_channels}, {tgt_frames}, {tgt_height}, {tgt_width}). Make sure "
"the image and latent have the same spatial shape."
)
tokens = latent_tools.patchifier.patchify(self.latent)
start_token = latent_tools.patchifier.get_token_count(
latent_tools.target_shape._replace(frames=self.latent_idx)
)
stop_token = start_token + tokens.shape[1]
latent_state = latent_state.clone()
latent_state.latent[:, start_token:stop_token] = tokens
latent_state.clean_latent[:, start_token:stop_token] = tokens
latent_state.denoise_mask[:, start_token:stop_token] = 1.0 - self.strength
return latent_state
packages/ltx-core/src/ltx_core/guidance/__init__.py
+15
View File
@@ -0,0 +1,15 @@
"""Guidance and perturbation utilities for attention manipulation."""
from ltx_core.guidance.perturbations import (
BatchedPerturbationConfig,
Perturbation,
PerturbationConfig,
PerturbationType,
)
__all__ = [
"BatchedPerturbationConfig",
"Perturbation",
"PerturbationConfig",
"PerturbationType",
]
packages/ltx-core/src/ltx_core/guidance/perturbations.py
+79
View File
@@ -0,0 +1,79 @@
from dataclasses import dataclass
from enum import Enum
import torch
from torch._prims_common import DeviceLikeType
class PerturbationType(Enum):
"""Types of attention perturbations for STG (Spatio-Temporal Guidance)."""
SKIP_A2V_CROSS_ATTN = "skip_a2v_cross_attn"
SKIP_V2A_CROSS_ATTN = "skip_v2a_cross_attn"
SKIP_VIDEO_SELF_ATTN = "skip_video_self_attn"
SKIP_AUDIO_SELF_ATTN = "skip_audio_self_attn"
@dataclass(frozen=True)
class Perturbation:
"""A single perturbation specifying which attention type to skip and in which blocks."""
type: PerturbationType
blocks: list[int] | None # None means all blocks
def is_perturbed(self, perturbation_type: PerturbationType, block: int) -> bool:
if self.type != perturbation_type:
return False
if self.blocks is None:
return True
return block in self.blocks
@dataclass(frozen=True)
class PerturbationConfig:
"""Configuration holding a list of perturbations for a single sample."""
perturbations: list[Perturbation] | None
def is_perturbed(self, perturbation_type: PerturbationType, block: int) -> bool:
if self.perturbations is None:
return False
return any(perturbation.is_perturbed(perturbation_type, block) for perturbation in self.perturbations)
@staticmethod
def empty() -> "PerturbationConfig":
return PerturbationConfig([])
@dataclass(frozen=True)
class BatchedPerturbationConfig:
"""Perturbation configurations for a batch, with utilities for generating attention masks."""
perturbations: list[PerturbationConfig]
def mask(
self, perturbation_type: PerturbationType, block: int, device: DeviceLikeType, dtype: torch.dtype
) -> torch.Tensor:
mask = torch.ones((len(self.perturbations),), device=device, dtype=dtype)
for batch_idx, perturbation in enumerate(self.perturbations):
if perturbation.is_perturbed(perturbation_type, block):
mask[batch_idx] = 0
return mask
def mask_like(self, perturbation_type: PerturbationType, block: int, values: torch.Tensor) -> torch.Tensor:
mask = self.mask(perturbation_type, block, values.device, values.dtype)
return mask.view(mask.numel(), *([1] * len(values.shape[1:])))
def any_in_batch(self, perturbation_type: PerturbationType, block: int) -> bool:
return any(perturbation.is_perturbed(perturbation_type, block) for perturbation in self.perturbations)
def all_in_batch(self, perturbation_type: PerturbationType, block: int) -> bool:
return all(perturbation.is_perturbed(perturbation_type, block) for perturbation in self.perturbations)
@staticmethod
def empty(batch_size: int) -> "BatchedPerturbationConfig":
return BatchedPerturbationConfig([PerturbationConfig.empty() for _ in range(batch_size)])
packages/ltx-core/src/ltx_core/loader/__init__.py
+48
View File
@@ -0,0 +1,48 @@
"""Loader utilities for model weights, LoRAs, and safetensor operations."""
from ltx_core.loader.fuse_loras import apply_loras
from ltx_core.loader.module_ops import ModuleOps
from ltx_core.loader.primitives import (
LoRAAdaptableProtocol,
LoraPathStrengthAndSDOps,
LoraStateDictWithStrength,
ModelBuilderProtocol,
StateDict,
StateDictLoader,
)
from ltx_core.loader.registry import DummyRegistry, Registry, StateDictRegistry
from ltx_core.loader.sd_ops import (
LTXV_LORA_COMFY_RENAMING_MAP,
ContentMatching,
ContentReplacement,
KeyValueOperation,
KeyValueOperationResult,
SDKeyValueOperation,
SDOps,
)
from ltx_core.loader.sft_loader import SafetensorsModelStateDictLoader, SafetensorsStateDictLoader
from ltx_core.loader.single_gpu_model_builder import SingleGPUModelBuilder
__all__ = [
"LTXV_LORA_COMFY_RENAMING_MAP",
"ContentMatching",
"ContentReplacement",
"DummyRegistry",
"KeyValueOperation",
"KeyValueOperationResult",
"LoRAAdaptableProtocol",
"LoraPathStrengthAndSDOps",
"LoraStateDictWithStrength",
"ModelBuilderProtocol",
"ModuleOps",
"Registry",
"SDKeyValueOperation",
"SDOps",
"SafetensorsModelStateDictLoader",
"SafetensorsStateDictLoader",
"SingleGPUModelBuilder",
"StateDict",
"StateDictLoader",
"StateDictRegistry",
"apply_loras",
]
packages/ltx-core/src/ltx_core/loader/fuse_loras.py
+100
View File
@@ -0,0 +1,100 @@
import torch
import triton
from ltx_core.loader.kernels import fused_add_round_kernel
from ltx_core.loader.primitives import LoraStateDictWithStrength, StateDict
BLOCK_SIZE = 1024
def fused_add_round_launch(target_weight: torch.Tensor, original_weight: torch.Tensor, seed: int) -> torch.Tensor:
if original_weight.dtype == torch.float8_e4m3fn:
exponent_bits, mantissa_bits, exponent_bias = 4, 3, 7
elif original_weight.dtype == torch.float8_e5m2:
exponent_bits, mantissa_bits, exponent_bias = 5, 2, 15 # noqa: F841
else:
raise ValueError("Unsupported dtype")
if target_weight.dtype != torch.bfloat16:
raise ValueError("target_weight dtype must be bfloat16")
# Calculate grid and block sizes
n_elements = original_weight.numel()
grid = (triton.cdiv(n_elements, BLOCK_SIZE),)
# Launch kernel
fused_add_round_kernel[grid](
original_weight,
target_weight,
seed,
n_elements,
exponent_bias,
mantissa_bits,
BLOCK_SIZE,
)
return target_weight
def calculate_weight_float8_(target_weights: torch.Tensor, original_weights: torch.Tensor) -> torch.Tensor:
result = fused_add_round_launch(target_weights, original_weights, seed=0).to(target_weights.dtype)
target_weights.copy_(result, non_blocking=True)
return target_weights
def _prepare_deltas(
lora_sd_and_strengths: list[LoraStateDictWithStrength], key: str, dtype: torch.dtype, device: torch.device
) -> torch.Tensor | None:
deltas = []
prefix = key[: -len(".weight")]
key_a = f"{prefix}.lora_A.weight"
key_b = f"{prefix}.lora_B.weight"
for lsd, coef in lora_sd_and_strengths:
if key_a not in lsd.sd or key_b not in lsd.sd:
continue
product = torch.matmul(lsd.sd[key_b] * coef, lsd.sd[key_a])
deltas.append(product.to(dtype=dtype, device=device))
if len(deltas) == 0:
return None
elif len(deltas) == 1:
return deltas[0]
return torch.sum(torch.stack(deltas, dim=0), dim=0)
def apply_loras(
model_sd: StateDict,
lora_sd_and_strengths: list[LoraStateDictWithStrength],
dtype: torch.dtype,
destination_sd: StateDict | None = None,
) -> StateDict:
sd = {}
if destination_sd is not None:
sd = destination_sd.sd
size = 0
device = torch.device("meta")
inner_dtypes = set()
for key, weight in model_sd.sd.items():
if weight is None:
continue
device = weight.device
target_dtype = dtype if dtype is not None else weight.dtype
deltas_dtype = target_dtype if target_dtype not in [torch.float8_e4m3fn, torch.float8_e5m2] else torch.bfloat16
deltas = _prepare_deltas(lora_sd_and_strengths, key, deltas_dtype, device)
if deltas is None:
if key in sd:
continue
deltas = weight.clone().to(dtype=target_dtype, device=device)
elif weight.dtype == torch.float8_e4m3fn:
if str(device).startswith("cuda"):
deltas = calculate_weight_float8_(deltas, weight)
else:
deltas.add_(weight.to(dtype=deltas.dtype, device=device))
elif weight.dtype == torch.bfloat16:
deltas.add_(weight)
else:
raise ValueError(f"Unsupported dtype: {weight.dtype}")
sd[key] = deltas.to(dtype=target_dtype)
inner_dtypes.add(target_dtype)
size += deltas.nbytes
if destination_sd is not None:
return destination_sd
return StateDict(sd, device, size, inner_dtypes)
packages/ltx-core/src/ltx_core/loader/kernels.py
+72
View File
@@ -0,0 +1,72 @@
# ruff: noqa: ANN001, ANN201, ERA001, N803, N806
import triton
import triton.language as tl
@triton.jit
def fused_add_round_kernel(
x_ptr,
output_ptr, # contents will be added to the output
seed,
n_elements,
EXPONENT_BIAS,
MANTISSA_BITS,
BLOCK_SIZE: tl.constexpr,
):
"""
A kernel to upcast 8bit quantized weights to bfloat16 with stochastic rounding
and add them to bfloat16 output weights. Might be used to upcast original model weights
and to further add them to precalculated deltas coming from LoRAs.
"""
# Get program ID and compute offsets
pid = tl.program_id(axis=0)
block_start = pid * BLOCK_SIZE
offsets = block_start + tl.arange(0, BLOCK_SIZE)
mask = offsets < n_elements
# Load data
x = tl.load(x_ptr + offsets, mask=mask)
rand_vals = tl.rand(seed, offsets) - 0.5
x = tl.cast(x, tl.float16)
delta = tl.load(output_ptr + offsets, mask=mask)
delta = tl.cast(delta, tl.float16)
x = x + delta
x_bits = tl.cast(x, tl.int16, bitcast=True)
# Calculate the exponent. Unbiased fp16 exponent is ((x_bits & 0x7C00) >> 10) - 15 for
# normal numbers and -14 for subnormals.
fp16_exponent_bits = (x_bits & 0x7C00) >> 10
fp16_normals = fp16_exponent_bits > 0
fp16_exponent = tl.where(fp16_normals, fp16_exponent_bits - 15, -14)
# Add the target dtype's exponent bias and clamp to the target dtype's exponent range.
exponent = fp16_exponent + EXPONENT_BIAS
MAX_EXPONENT = 2 * EXPONENT_BIAS + 1
exponent = tl.where(exponent > MAX_EXPONENT, MAX_EXPONENT, exponent)
exponent = tl.where(exponent < 0, 0, exponent)
# Normal ULP exponent, expressed as an fp16 exponent field:
# (exponent - EXPONENT_BIAS - MANTISSA_BITS) + 15
# Simplifies to: fp16_exponent - MANTISSA_BITS + 15
# See https://en.wikipedia.org/wiki/Unit_in_the_last_place
eps_exp = tl.maximum(0, tl.minimum(31, exponent - EXPONENT_BIAS - MANTISSA_BITS + 15))
# Calculate epsilon in the target dtype
eps_normal = tl.cast(tl.cast(eps_exp << 10, tl.int16), tl.float16, bitcast=True)
# Subnormal ULP: 2^(1 - EXPONENT_BIAS - MANTISSA_BITS) ->
# fp16 exponent bits: (1 - EXPONENT_BIAS - MANTISSA_BITS) + 15 =
# 16 - EXPONENT_BIAS - MANTISSA_BITS
eps_subnormal = tl.cast(tl.cast((16 - EXPONENT_BIAS - MANTISSA_BITS) << 10, tl.int16), tl.float16, bitcast=True)
eps = tl.where(exponent > 0, eps_normal, eps_subnormal)
# Apply zero mask to epsilon
eps = tl.where(x == 0, 0.0, eps)
# Apply stochastic rounding
output = tl.cast(x + rand_vals * eps, tl.bfloat16)
# Store the result
tl.store(output_ptr + offsets, output, mask=mask)
packages/ltx-core/src/ltx_core/loader/module_ops.py
+14
View File
@@ -0,0 +1,14 @@
from typing import Callable, NamedTuple
import torch
class ModuleOps(NamedTuple):
"""
Defines a named operation for matching and mutating PyTorch modules.
Used to selectively transform modules in a model (e.g., replacing layers with quantized versions).
"""
name: str
matcher: Callable[[torch.nn.Module], bool]
mutator: Callable[[torch.nn.Module], torch.nn.Module]
packages/ltx-core/src/ltx_core/loader/primitives.py
+109
View File
@@ -0,0 +1,109 @@
from dataclasses import dataclass
from typing import NamedTuple, Protocol
import torch
from ltx_core.loader.module_ops import ModuleOps
from ltx_core.loader.sd_ops import SDOps
from ltx_core.model.model_protocol import ModelType
@dataclass(frozen=True)
class StateDict:
"""
Immutable container for a PyTorch state dictionary.
Contains:
- sd: Dictionary of tensors (weights, buffers, etc.)
- device: Device where tensors are stored
- size: Total memory footprint in bytes
- dtype: Set of tensor dtypes present
"""
sd: dict
device: torch.device
size: int
dtype: set[torch.dtype]
def footprint(self) -> tuple[int, torch.device]:
return self.size, self.device
class StateDictLoader(Protocol):
"""
Protocol for loading state dictionaries from various sources.
Implementations must provide:
- metadata: Extract model metadata from a single path
- load: Load state dict from path(s) and apply SDOps transformations
"""
def metadata(self, path: str) -> dict:
"""
Load metadata from path
"""
def load(self, path: str | list[str], sd_ops: SDOps | None = None, device: torch.device | None = None) -> StateDict:
"""
Load state dict from path or paths (for sharded model storage) and apply sd_ops
"""
class ModelBuilderProtocol(Protocol[ModelType]):
"""
Protocol for building PyTorch models from configuration dictionaries.
Implementations must provide:
- meta_model: Create a model from configuration dictionary and apply module operations
- build: Create and initialize a model from state dictionary and apply dtype transformations
"""
def meta_model(self, config: dict, module_ops: list[ModuleOps] | None = None) -> ModelType:
"""
Create a model on the meta device from a configuration dictionary.
This decouples model creation from weight loading, allowing the model
architecture to be instantiated without allocating memory for parameters.
Args:
config: Model configuration dictionary.
module_ops: Optional list of module operations to apply (e.g., quantization).
Returns:
Model instance on meta device (no actual memory allocated for parameters).
"""
...
def build(self, dtype: torch.dtype | None = None) -> ModelType:
"""
Build the model
Args:
dtype: Target dtype for the model, if None, uses the dtype of the model_path model
Returns:
Model instance
"""
...
class LoRAAdaptableProtocol(Protocol):
"""
Protocol for models that can be adapted with LoRAs.
Implementations must provide:
- lora: Add a LoRA to the model
"""
def lora(self, lora_path: str, strength: float) -> "LoRAAdaptableProtocol":
pass
class LoraPathStrengthAndSDOps(NamedTuple):
"""
Tuple containing a LoRA path, strength, and SDOps for applying to the LoRA state dict.
"""
path: str
strength: float
sd_ops: SDOps
class LoraStateDictWithStrength(NamedTuple):
"""
Tuple containing a LoRA state dict and strength for applying to the model.
"""
state_dict: StateDict
strength: float
packages/ltx-core/src/ltx_core/loader/registry.py
+84
View File
@@ -0,0 +1,84 @@
import hashlib
import threading
from dataclasses import dataclass, field
from pathlib import Path
from typing import Protocol
from ltx_core.loader.primitives import StateDict
from ltx_core.loader.sd_ops import SDOps
class Registry(Protocol):
"""
Protocol for managing state dictionaries in a registry.
It is used to store state dictionaries and reuse them later without loading them again.
Implementations must provide:
- add: Add a state dictionary to the registry
- pop: Remove a state dictionary from the registry
- get: Retrieve a state dictionary from the registry
- clear: Clear all state dictionaries from the registry
"""
def add(self, paths: list[str], sd_ops: SDOps | None, state_dict: StateDict) -> None: ...
def pop(self, paths: list[str], sd_ops: SDOps | None) -> StateDict | None: ...
def get(self, paths: list[str], sd_ops: SDOps | None) -> StateDict | None: ...
def clear(self) -> None: ...
class DummyRegistry(Registry):
"""
Dummy registry that does not store state dictionaries.
"""
def add(self, paths: list[str], sd_ops: SDOps | None, state_dict: StateDict) -> None:
pass
def pop(self, paths: list[str], sd_ops: SDOps | None) -> StateDict | None:
pass
def get(self, paths: list[str], sd_ops: SDOps | None) -> StateDict | None:
pass
def clear(self) -> None:
pass
@dataclass
class StateDictRegistry(Registry):
"""
Registry that stores state dictionaries in a dictionary.
"""
_state_dicts: dict[str, StateDict] = field(default_factory=dict)
_lock: threading.Lock = field(default_factory=threading.Lock)
def _generate_id(self, paths: list[str], sd_ops: SDOps) -> str:
m = hashlib.sha256()
parts = [str(Path(p).resolve()) for p in paths]
if sd_ops is not None:
parts.append(sd_ops.name)
m.update("\0".join(parts).encode("utf-8"))
return m.hexdigest()
def add(self, paths: list[str], sd_ops: SDOps | None, state_dict: StateDict) -> str:
sd_id = self._generate_id(paths, sd_ops)
with self._lock:
if sd_id in self._state_dicts:
raise ValueError(f"State dict retrieved from {paths} with {sd_ops} already added, check with get first")
self._state_dicts[sd_id] = state_dict
return sd_id
def pop(self, paths: list[str], sd_ops: SDOps | None) -> StateDict | None:
with self._lock:
return self._state_dicts.pop(self._generate_id(paths, sd_ops), None)
def get(self, paths: list[str], sd_ops: SDOps | None) -> StateDict | None:
with self._lock:
return self._state_dicts.get(self._generate_id(paths, sd_ops), None)
def clear(self) -> None:
with self._lock:
self._state_dicts.clear()
packages/ltx-core/src/ltx_core/loader/sd_ops.py
+127
View File
@@ -0,0 +1,127 @@
from dataclasses import dataclass, replace
from typing import NamedTuple, Protocol
import torch
@dataclass(frozen=True, slots=True)
class ContentReplacement:
"""
Represents a content replacement operation.
Used to replace a specific content with a replacement in a state dict key.
"""
content: str
replacement: str
@dataclass(frozen=True, slots=True)
class ContentMatching:
"""
Represents a content matching operation.
Used to match a specific prefix and suffix in a state dict key.
"""
prefix: str = ""
suffix: str = ""
class KeyValueOperationResult(NamedTuple):
"""
Represents the result of a key-value operation.
Contains the new key and value after the operation has been applied.
"""
new_key: str
new_value: torch.Tensor
class KeyValueOperation(Protocol):
"""
Protocol for key-value operations.
Used to apply operations to a specific key and value in a state dict.
"""
def __call__(self, tensor_key: str, tensor_value: torch.Tensor) -> list[KeyValueOperationResult]: ...
@dataclass(frozen=True, slots=True)
class SDKeyValueOperation:
"""
Represents a key-value operation.
Used to apply operations to a specific key and value in a state dict.
"""
key_matcher: ContentMatching
kv_operation: KeyValueOperation
@dataclass(frozen=True, slots=True)
class SDOps:
"""Immutable class representing state dict key operations."""
name: str
mapping: tuple[
ContentReplacement | ContentMatching | SDKeyValueOperation, ...
] = () # Immutable tuple of (key, value) pairs
def with_replacement(self, content: str, replacement: str) -> "SDOps":
"""Create a new SDOps instance with the specified replacement added to the mapping."""
new_mapping = (*self.mapping, ContentReplacement(content, replacement))
return replace(self, mapping=new_mapping)
def with_matching(self, prefix: str = "", suffix: str = "") -> "SDOps":
"""Create a new SDOps instance with the specified prefix and suffix matching added to the mapping."""
new_mapping = (*self.mapping, ContentMatching(prefix, suffix))
return replace(self, mapping=new_mapping)
def with_kv_operation(
self,
operation: KeyValueOperation,
key_prefix: str = "",
key_suffix: str = "",
) -> "SDOps":
"""Create a new SDOps instance with the specified value operation added to the mapping."""
key_matcher = ContentMatching(key_prefix, key_suffix)
sd_kv_operation = SDKeyValueOperation(key_matcher, operation)
new_mapping = (*self.mapping, sd_kv_operation)
return replace(self, mapping=new_mapping)
def apply_to_key(self, key: str) -> str | None:
"""Apply the mapping to the given name."""
matchers = [content for content in self.mapping if isinstance(content, ContentMatching)]
valid = any(key.startswith(f.prefix) and key.endswith(f.suffix) for f in matchers)
if not valid:
return None
for replacement in self.mapping:
if not isinstance(replacement, ContentReplacement):
continue
if replacement.content in key:
key = key.replace(replacement.content, replacement.replacement)
return key
def apply_to_key_value(self, key: str, value: torch.Tensor) -> list[KeyValueOperationResult]:
"""Apply the value operation to the given name and associated value."""
for operation in self.mapping:
if not isinstance(operation, SDKeyValueOperation):
continue
if key.startswith(operation.key_matcher.prefix) and key.endswith(operation.key_matcher.suffix):
return operation.kv_operation(key, value)
return [KeyValueOperationResult(key, value)]
# Predefined SDOps instances
LTXV_LORA_COMFY_RENAMING_MAP = (
SDOps("LTXV_LORA_COMFY_PREFIX_MAP").with_matching().with_replacement("diffusion_model.", "")
)
LTXV_LORA_COMFY_TARGET_MAP = (
SDOps("LTXV_LORA_COMFY_TARGET_MAP")
.with_matching()
.with_replacement("diffusion_model.", "")
.with_replacement(".lora_A.weight", ".weight")
.with_replacement(".lora_B.weight", ".weight")
)
packages/ltx-core/src/ltx_core/loader/sft_loader.py
+63
View File
@@ -0,0 +1,63 @@
import json
import safetensors
import torch
from ltx_core.loader.primitives import StateDict, StateDictLoader
from ltx_core.loader.sd_ops import SDOps
class SafetensorsStateDictLoader(StateDictLoader):
"""
Loads weights from safetensors files without metadata support.
Use this for loading raw weight files. For model files that include
configuration metadata, use SafetensorsModelStateDictLoader instead.
"""
def metadata(self, path: str) -> dict:
raise NotImplementedError("Not implemented")
def load(self, path: str | list[str], sd_ops: SDOps, device: torch.device | None = None) -> StateDict:
"""
Load state dict from path or paths (for sharded model storage) and apply sd_ops
"""
sd = {}
size = 0
dtype = set()
device = device or torch.device("cpu")
model_paths = path if isinstance(path, list) else [path]
for shard_path in model_paths:
with safetensors.safe_open(shard_path, framework="pt", device=str(device)) as f:
safetensor_keys = f.keys()
for name in safetensor_keys:
expected_name = name if sd_ops is None else sd_ops.apply_to_key(name)
if expected_name is None:
continue
value = f.get_tensor(name).to(device=device, non_blocking=True, copy=False)
key_value_pairs = ((expected_name, value),)
if sd_ops is not None:
key_value_pairs = sd_ops.apply_to_key_value(expected_name, value)
for key, value in key_value_pairs:
size += value.nbytes
dtype.add(value.dtype)
sd[key] = value
return StateDict(sd=sd, device=device, size=size, dtype=dtype)
class SafetensorsModelStateDictLoader(StateDictLoader):
"""
Loads weights and configuration metadata from safetensors model files.
Unlike SafetensorsStateDictLoader, this loader can read model configuration
from the safetensors file metadata via the metadata() method.
"""
def __init__(self, weight_loader: SafetensorsStateDictLoader | None = None):
self.weight_loader = weight_loader if weight_loader is not None else SafetensorsStateDictLoader()
def metadata(self, path: str) -> dict:
with safetensors.safe_open(path, framework="pt") as f:
return json.loads(f.metadata()["config"])
def load(self, path: str | list[str], sd_ops: SDOps | None = None, device: torch.device | None = None) -> StateDict:
return self.weight_loader.load(path, sd_ops, device)
packages/ltx-core/src/ltx_core/loader/single_gpu_model_builder.py
+101
View File
@@ -0,0 +1,101 @@
import logging
from dataclasses import dataclass, field, replace
from typing import Generic
import torch
from ltx_core.loader.fuse_loras import apply_loras
from ltx_core.loader.module_ops import ModuleOps
from ltx_core.loader.primitives import (
LoRAAdaptableProtocol,
LoraPathStrengthAndSDOps,
LoraStateDictWithStrength,
ModelBuilderProtocol,
StateDict,
StateDictLoader,
)
from ltx_core.loader.registry import DummyRegistry, Registry
from ltx_core.loader.sd_ops import SDOps
from ltx_core.loader.sft_loader import SafetensorsModelStateDictLoader
from ltx_core.model.model_protocol import ModelConfigurator, ModelType
logger: logging.Logger = logging.getLogger(__name__)
@dataclass(frozen=True)
class SingleGPUModelBuilder(Generic[ModelType], ModelBuilderProtocol[ModelType], LoRAAdaptableProtocol):
"""
Builder for PyTorch models residing on a single GPU.
"""
model_class_configurator: type[ModelConfigurator[ModelType]]
model_path: str | tuple[str, ...]
model_sd_ops: SDOps | None = None
module_ops: tuple[ModuleOps, ...] = field(default_factory=tuple)
loras: tuple[LoraPathStrengthAndSDOps, ...] = field(default_factory=tuple)
model_loader: StateDictLoader = field(default_factory=SafetensorsModelStateDictLoader)
registry: Registry = field(default_factory=DummyRegistry)
def lora(self, lora_path: str, strength: float = 1.0, sd_ops: SDOps | None = None) -> "SingleGPUModelBuilder":
return replace(self, loras=(*self.loras, LoraPathStrengthAndSDOps(lora_path, strength, sd_ops)))
def model_config(self) -> dict:
first_shard_path = self.model_path[0] if isinstance(self.model_path, tuple) else self.model_path
return self.model_loader.metadata(first_shard_path)
def meta_model(self, config: dict, module_ops: tuple[ModuleOps, ...]) -> ModelType:
with torch.device("meta"):
model = self.model_class_configurator.from_config(config)
for module_op in module_ops:
if module_op.matcher(model):
model = module_op.mutator(model)
return model
def load_sd(
self, paths: list[str], registry: Registry, device: torch.device | None, sd_ops: SDOps | None = None
) -> StateDict:
state_dict = registry.get(paths, sd_ops)
if state_dict is None:
state_dict = self.model_loader.load(paths, sd_ops=sd_ops, device=device)
registry.add(paths, sd_ops=sd_ops, state_dict=state_dict)
return state_dict
def _return_model(self, meta_model: ModelType, device: torch.device) -> ModelType:
uninitialized_params = [name for name, param in meta_model.named_parameters() if str(param.device) == "meta"]
uninitialized_buffers = [name for name, buffer in meta_model.named_buffers() if str(buffer.device) == "meta"]
if uninitialized_params or uninitialized_buffers:
logger.warning(f"Uninitialized parameters or buffers: {uninitialized_params + uninitialized_buffers}")
return meta_model
retval = meta_model.to(device)
return retval
def build(self, device: torch.device | None = None, dtype: torch.dtype | None = None) -> ModelType:
device = torch.device("cuda") if device is None else device
config = self.model_config()
meta_model = self.meta_model(config, self.module_ops)
model_paths = self.model_path if isinstance(self.model_path, tuple) else [self.model_path]
model_state_dict = self.load_sd(model_paths, sd_ops=self.model_sd_ops, registry=self.registry, device=device)
lora_strengths = [lora.strength for lora in self.loras]
if not lora_strengths or (min(lora_strengths) == 0 and max(lora_strengths) == 0):
sd = model_state_dict.sd
if dtype is not None:
sd = {key: value.to(dtype=dtype) for key, value in model_state_dict.sd.items()}
meta_model.load_state_dict(sd, strict=False, assign=True)
return self._return_model(meta_model, device)
lora_state_dicts = [
self.load_sd([lora.path], sd_ops=lora.sd_ops, registry=self.registry, device=device) for lora in self.loras
]
lora_sd_and_strengths = [
LoraStateDictWithStrength(sd, strength)
for sd, strength in zip(lora_state_dicts, lora_strengths, strict=True)
]
final_sd = apply_loras(
model_sd=model_state_dict,
lora_sd_and_strengths=lora_sd_and_strengths,
dtype=dtype,
destination_sd=model_state_dict if isinstance(self.registry, DummyRegistry) else None,
)
meta_model.load_state_dict(final_sd.sd, strict=False, assign=True)
return self._return_model(meta_model, device)
packages/ltx-core/src/ltx_core/model/__init__.py
+8
View File
@@ -0,0 +1,8 @@
"""Model definitions for LTX-2."""
from ltx_core.model.model_protocol import ModelConfigurator, ModelType
__all__ = [
"ModelConfigurator",
"ModelType",
]
packages/ltx-core/src/ltx_core/model/audio_vae/__init__.py
+27
View File
@@ -0,0 +1,27 @@
"""Audio VAE model components."""
from ltx_core.model.audio_vae.audio_vae import AudioDecoder, AudioEncoder, decode_audio
from ltx_core.model.audio_vae.model_configurator import (
AUDIO_VAE_DECODER_COMFY_KEYS_FILTER,
AUDIO_VAE_ENCODER_COMFY_KEYS_FILTER,
VOCODER_COMFY_KEYS_FILTER,
AudioDecoderConfigurator,
AudioEncoderConfigurator,
VocoderConfigurator,
)
from ltx_core.model.audio_vae.ops import AudioProcessor
from ltx_core.model.audio_vae.vocoder import Vocoder
__all__ = [
"AUDIO_VAE_DECODER_COMFY_KEYS_FILTER",
"AUDIO_VAE_ENCODER_COMFY_KEYS_FILTER",
"VOCODER_COMFY_KEYS_FILTER",
"AudioDecoder",
"AudioDecoderConfigurator",
"AudioEncoder",
"AudioEncoderConfigurator",
"AudioProcessor",
"Vocoder",
"VocoderConfigurator",
"decode_audio",
]
packages/ltx-core/src/ltx_core/model/audio_vae/attention.py
+71
View File
@@ -0,0 +1,71 @@
from enum import Enum
import torch
from ltx_core.model.common.normalization import NormType, build_normalization_layer
class AttentionType(Enum):
"""Enum for specifying the attention mechanism type."""
VANILLA = "vanilla"
LINEAR = "linear"
NONE = "none"
class AttnBlock(torch.nn.Module):
def __init__(
self,
in_channels: int,
norm_type: NormType = NormType.GROUP,
) -> None:
super().__init__()
self.in_channels = in_channels
self.norm = build_normalization_layer(in_channels, normtype=norm_type)
self.q = torch.nn.Conv2d(in_channels, in_channels, kernel_size=1, stride=1, padding=0)
self.k = torch.nn.Conv2d(in_channels, in_channels, kernel_size=1, stride=1, padding=0)
self.v = torch.nn.Conv2d(in_channels, in_channels, kernel_size=1, stride=1, padding=0)
self.proj_out = torch.nn.Conv2d(in_channels, in_channels, kernel_size=1, stride=1, padding=0)
def forward(self, x: torch.Tensor) -> torch.Tensor:
h_ = x
h_ = self.norm(h_)
q = self.q(h_)
k = self.k(h_)
v = self.v(h_)
# compute attention
b, c, h, w = q.shape
q = q.reshape(b, c, h * w).contiguous()
q = q.permute(0, 2, 1).contiguous() # b,hw,c
k = k.reshape(b, c, h * w).contiguous() # b,c,hw
w_ = torch.bmm(q, k).contiguous() # b,hw,hw w[b,i,j]=sum_c q[b,i,c]k[b,c,j]
w_ = w_ * (int(c) ** (-0.5))
w_ = torch.nn.functional.softmax(w_, dim=2)
# attend to values
v = v.reshape(b, c, h * w).contiguous()
w_ = w_.permute(0, 2, 1).contiguous() # b,hw,hw (first hw of k, second of q)
h_ = torch.bmm(v, w_).contiguous() # b, c,hw (hw of q) h_[b,c,j] = sum_i v[b,c,i] w_[b,i,j]
h_ = h_.reshape(b, c, h, w).contiguous()
h_ = self.proj_out(h_)
return x + h_
def make_attn(
in_channels: int,
attn_type: AttentionType = AttentionType.VANILLA,
norm_type: NormType = NormType.GROUP,
) -> torch.nn.Module:
match attn_type:
case AttentionType.VANILLA:
return AttnBlock(in_channels, norm_type=norm_type)
case AttentionType.NONE:
return torch.nn.Identity()
case AttentionType.LINEAR:
raise NotImplementedError(f"Attention type {attn_type.value} is not supported yet.")
case _:
raise ValueError(f"Unknown attention type: {attn_type}")
packages/ltx-core/src/ltx_core/model/audio_vae/audio_vae.py
+480
View File
@@ -0,0 +1,480 @@
from typing import Set, Tuple
import torch
import torch.nn.functional as F
from ltx_core.components.patchifiers import AudioPatchifier
from ltx_core.model.audio_vae.attention import AttentionType, make_attn
from ltx_core.model.audio_vae.causal_conv_2d import make_conv2d
from ltx_core.model.audio_vae.causality_axis import CausalityAxis
from ltx_core.model.audio_vae.downsample import build_downsampling_path
from ltx_core.model.audio_vae.ops import PerChannelStatistics
from ltx_core.model.audio_vae.resnet import ResnetBlock
from ltx_core.model.audio_vae.upsample import build_upsampling_path
from ltx_core.model.audio_vae.vocoder import Vocoder
from ltx_core.model.common.normalization import NormType, build_normalization_layer
from ltx_core.types import AudioLatentShape
LATENT_DOWNSAMPLE_FACTOR = 4
def build_mid_block(
channels: int,
temb_channels: int,
dropout: float,
norm_type: NormType,
causality_axis: CausalityAxis,
attn_type: AttentionType,
add_attention: bool,
) -> torch.nn.Module:
"""Build the middle block with two ResNet blocks and optional attention."""
mid = torch.nn.Module()
mid.block_1 = ResnetBlock(
in_channels=channels,
out_channels=channels,
temb_channels=temb_channels,
dropout=dropout,
norm_type=norm_type,
causality_axis=causality_axis,
)
mid.attn_1 = make_attn(channels, attn_type=attn_type, norm_type=norm_type) if add_attention else torch.nn.Identity()
mid.block_2 = ResnetBlock(
in_channels=channels,
out_channels=channels,
temb_channels=temb_channels,
dropout=dropout,
norm_type=norm_type,
causality_axis=causality_axis,
)
return mid
def run_mid_block(mid: torch.nn.Module, features: torch.Tensor) -> torch.Tensor:
"""Run features through the middle block."""
features = mid.block_1(features, temb=None)
features = mid.attn_1(features)
return mid.block_2(features, temb=None)
class AudioEncoder(torch.nn.Module):
"""
Encoder that compresses audio spectrograms into latent representations.
The encoder uses a series of downsampling blocks with residual connections,
attention mechanisms, and configurable causal convolutions.
"""
def __init__( # noqa: PLR0913
self,
*,
ch: int,
ch_mult: Tuple[int, ...] = (1, 2, 4, 8),
num_res_blocks: int,
attn_resolutions: Set[int],
dropout: float = 0.0,
resamp_with_conv: bool = True,
in_channels: int,
resolution: int,
z_channels: int,
double_z: bool = True,
attn_type: AttentionType = AttentionType.VANILLA,
mid_block_add_attention: bool = True,
norm_type: NormType = NormType.GROUP,
causality_axis: CausalityAxis = CausalityAxis.WIDTH,
sample_rate: int = 16000,
mel_hop_length: int = 160,
n_fft: int = 1024,
is_causal: bool = True,
mel_bins: int = 64,
**_ignore_kwargs,
) -> None:
"""
Initialize the Encoder.
Args:
Arguments are configuration parameters, loaded from the audio VAE checkpoint config
(audio_vae.model.params.ddconfig):
ch: Base number of feature channels used in the first convolution layer.
ch_mult: Multiplicative factors for the number of channels at each resolution level.
num_res_blocks: Number of residual blocks to use at each resolution level.
attn_resolutions: Spatial resolutions (e.g., in time/frequency) at which to apply attention.
resolution: Input spatial resolution of the spectrogram (height, width).
z_channels: Number of channels in the latent representation.
norm_type: Normalization layer type to use within the network (e.g., group, batch).
causality_axis: Axis along which convolutions should be causal (e.g., time axis).
sample_rate: Audio sample rate in Hz for the input signals.
mel_hop_length: Hop length used when computing the mel spectrogram.
n_fft: FFT size used to compute the spectrogram.
mel_bins: Number of mel-frequency bins in the input spectrogram.
in_channels: Number of channels in the input spectrogram tensor.
double_z: If True, predict both mean and log-variance (doubling latent channels).
is_causal: If True, use causal convolutions suitable for streaming setups.
dropout: Dropout probability used in residual and mid blocks.
attn_type: Type of attention mechanism to use in attention blocks.
resamp_with_conv: If True, perform resolution changes using strided convolutions.
mid_block_add_attention: If True, add an attention block in the mid-level of the encoder.
"""
super().__init__()
self.per_channel_statistics = PerChannelStatistics(latent_channels=ch)
self.sample_rate = sample_rate
self.mel_hop_length = mel_hop_length
self.n_fft = n_fft
self.is_causal = is_causal
self.mel_bins = mel_bins
self.patchifier = AudioPatchifier(
patch_size=1,
audio_latent_downsample_factor=LATENT_DOWNSAMPLE_FACTOR,
sample_rate=sample_rate,
hop_length=mel_hop_length,
is_causal=is_causal,
)
self.ch = ch
self.temb_ch = 0
self.num_resolutions = len(ch_mult)
self.num_res_blocks = num_res_blocks
self.resolution = resolution
self.in_channels = in_channels
self.z_channels = z_channels
self.double_z = double_z
self.norm_type = norm_type
self.causality_axis = causality_axis
self.attn_type = attn_type
# downsampling
self.conv_in = make_conv2d(
in_channels,
self.ch,
kernel_size=3,
stride=1,
causality_axis=self.causality_axis,
)
self.non_linearity = torch.nn.SiLU()
self.down, block_in = build_downsampling_path(
ch=ch,
ch_mult=ch_mult,
num_resolutions=self.num_resolutions,
num_res_blocks=num_res_blocks,
resolution=resolution,
temb_channels=self.temb_ch,
dropout=dropout,
norm_type=self.norm_type,
causality_axis=self.causality_axis,
attn_type=self.attn_type,
attn_resolutions=attn_resolutions,
resamp_with_conv=resamp_with_conv,
)
self.mid = build_mid_block(
channels=block_in,
temb_channels=self.temb_ch,
dropout=dropout,
norm_type=self.norm_type,
causality_axis=self.causality_axis,
attn_type=self.attn_type,
add_attention=mid_block_add_attention,
)
self.norm_out = build_normalization_layer(block_in, normtype=self.norm_type)
self.conv_out = make_conv2d(
block_in,
2 * z_channels if double_z else z_channels,
kernel_size=3,
stride=1,
causality_axis=self.causality_axis,
)
def forward(self, spectrogram: torch.Tensor) -> torch.Tensor:
"""
Encode audio spectrogram into latent representations.
Args:
spectrogram: Input spectrogram of shape (batch, channels, time, frequency)
Returns:
Encoded latent representation of shape (batch, channels, frames, mel_bins)
"""
h = self.conv_in(spectrogram)
h = self._run_downsampling_path(h)
h = run_mid_block(self.mid, h)
h = self._finalize_output(h)
return self._normalize_latents(h)
def _run_downsampling_path(self, h: torch.Tensor) -> torch.Tensor:
for level in range(self.num_resolutions):
stage = self.down[level]
for block_idx in range(self.num_res_blocks):
h = stage.block[block_idx](h, temb=None)
if stage.attn:
h = stage.attn[block_idx](h)
if level != self.num_resolutions - 1:
h = stage.downsample(h)
return h
def _finalize_output(self, h: torch.Tensor) -> torch.Tensor:
h = self.norm_out(h)
h = self.non_linearity(h)
return self.conv_out(h)
def _normalize_latents(self, latent_output: torch.Tensor) -> torch.Tensor:
"""
Normalize encoder latents using per-channel statistics.
When the encoder is configured with ``double_z=True``, the final
convolution produces twice the number of latent channels, typically
interpreted as two concatenated tensors along the channel dimension
(e.g., mean and variance or other auxiliary parameters).
This method intentionally uses only the first half of the channels
(the "mean" component) as input to the patchifier and normalization
logic. The remaining channels are left unchanged by this method and
are expected to be consumed elsewhere in the VAE pipeline.
If ``double_z=False``, the encoder output already contains only the
mean latents and the chunking operation simply returns that tensor.
"""
means = torch.chunk(latent_output, 2, dim=1)[0]
latent_shape = AudioLatentShape(
batch=means.shape[0],
channels=means.shape[1],
frames=means.shape[2],
mel_bins=means.shape[3],
)
latent_patched = self.patchifier.patchify(means)
latent_normalized = self.per_channel_statistics.normalize(latent_patched)
return self.patchifier.unpatchify(latent_normalized, latent_shape)
class AudioDecoder(torch.nn.Module):
"""
Symmetric decoder that reconstructs audio spectrograms from latent features.
The decoder mirrors the encoder structure with configurable channel multipliers,
attention resolutions, and causal convolutions.
"""
def __init__( # noqa: PLR0913
self,
*,
ch: int,
out_ch: int,
ch_mult: Tuple[int, ...] = (1, 2, 4, 8),
num_res_blocks: int,
attn_resolutions: Set[int],
resolution: int,
z_channels: int,
norm_type: NormType = NormType.GROUP,
causality_axis: CausalityAxis = CausalityAxis.WIDTH,
dropout: float = 0.0,
mid_block_add_attention: bool = True,
sample_rate: int = 16000,
mel_hop_length: int = 160,
is_causal: bool = True,
mel_bins: int | None = None,
) -> None:
"""
Initialize the Decoder.
Args:
Arguments are configuration parameters, loaded from the audio VAE checkpoint config
(audio_vae.model.params.ddconfig):
- ch, out_ch, ch_mult, num_res_blocks, attn_resolutions
- resolution, z_channels
- norm_type, causality_axis
"""
super().__init__()
# Internal behavioural defaults that are not driven by the checkpoint.
resamp_with_conv = True
attn_type = AttentionType.VANILLA
# Per-channel statistics for denormalizing latents
self.per_channel_statistics = PerChannelStatistics(latent_channels=ch)
self.sample_rate = sample_rate
self.mel_hop_length = mel_hop_length
self.is_causal = is_causal
self.mel_bins = mel_bins
self.patchifier = AudioPatchifier(
patch_size=1,
audio_latent_downsample_factor=LATENT_DOWNSAMPLE_FACTOR,
sample_rate=sample_rate,
hop_length=mel_hop_length,
is_causal=is_causal,
)
self.ch = ch
self.temb_ch = 0
self.num_resolutions = len(ch_mult)
self.num_res_blocks = num_res_blocks
self.resolution = resolution
self.out_ch = out_ch
self.give_pre_end = False
self.tanh_out = False
self.norm_type = norm_type
self.z_channels = z_channels
self.channel_multipliers = ch_mult
self.attn_resolutions = attn_resolutions
self.causality_axis = causality_axis
self.attn_type = attn_type
base_block_channels = ch * self.channel_multipliers[-1]
base_resolution = resolution // (2 ** (self.num_resolutions - 1))
self.z_shape = (1, z_channels, base_resolution, base_resolution)
self.conv_in = make_conv2d(
z_channels, base_block_channels, kernel_size=3, stride=1, causality_axis=self.causality_axis
)
self.non_linearity = torch.nn.SiLU()
self.mid = build_mid_block(
channels=base_block_channels,
temb_channels=self.temb_ch,
dropout=dropout,
norm_type=self.norm_type,
causality_axis=self.causality_axis,
attn_type=self.attn_type,
add_attention=mid_block_add_attention,
)
self.up, final_block_channels = build_upsampling_path(
ch=ch,
ch_mult=ch_mult,
num_resolutions=self.num_resolutions,
num_res_blocks=num_res_blocks,
resolution=resolution,
temb_channels=self.temb_ch,
dropout=dropout,
norm_type=self.norm_type,
causality_axis=self.causality_axis,
attn_type=self.attn_type,
attn_resolutions=attn_resolutions,
resamp_with_conv=resamp_with_conv,
initial_block_channels=base_block_channels,
)
self.norm_out = build_normalization_layer(final_block_channels, normtype=self.norm_type)
self.conv_out = make_conv2d(
final_block_channels, out_ch, kernel_size=3, stride=1, causality_axis=self.causality_axis
)
def forward(self, sample: torch.Tensor) -> torch.Tensor:
"""
Decode latent features back to audio spectrograms.
Args:
sample: Encoded latent representation of shape (batch, channels, frames, mel_bins)
Returns:
Reconstructed audio spectrogram of shape (batch, channels, time, frequency)
"""
sample, target_shape = self._denormalize_latents(sample)
h = self.conv_in(sample)
h = run_mid_block(self.mid, h)
h = self._run_upsampling_path(h)
h = self._finalize_output(h)
return self._adjust_output_shape(h, target_shape)
def _denormalize_latents(self, sample: torch.Tensor) -> tuple[torch.Tensor, AudioLatentShape]:
latent_shape = AudioLatentShape(
batch=sample.shape[0],
channels=sample.shape[1],
frames=sample.shape[2],
mel_bins=sample.shape[3],
)
sample_patched = self.patchifier.patchify(sample)
sample_denormalized = self.per_channel_statistics.un_normalize(sample_patched)
sample = self.patchifier.unpatchify(sample_denormalized, latent_shape)
target_frames = latent_shape.frames * LATENT_DOWNSAMPLE_FACTOR
if self.causality_axis != CausalityAxis.NONE:
target_frames = max(target_frames - (LATENT_DOWNSAMPLE_FACTOR - 1), 1)
target_shape = AudioLatentShape(
batch=latent_shape.batch,
channels=self.out_ch,
frames=target_frames,
mel_bins=self.mel_bins if self.mel_bins is not None else latent_shape.mel_bins,
)
return sample, target_shape
def _adjust_output_shape(
self,
decoded_output: torch.Tensor,
target_shape: AudioLatentShape,
) -> torch.Tensor:
"""
Adjust output shape to match target dimensions for variable-length audio.
This function handles the common case where decoded audio spectrograms need to be
resized to match a specific target shape.
Args:
decoded_output: Tensor of shape (batch, channels, time, frequency)
target_shape: AudioLatentShape describing (batch, channels, time, mel bins)
Returns:
Tensor adjusted to match target_shape exactly
"""
# Current output shape: (batch, channels, time, frequency)
_, _, current_time, current_freq = decoded_output.shape
target_channels = target_shape.channels
target_time = target_shape.frames
target_freq = target_shape.mel_bins
# Step 1: Crop first to avoid exceeding target dimensions
decoded_output = decoded_output[
:, :target_channels, : min(current_time, target_time), : min(current_freq, target_freq)
]
# Step 2: Calculate padding needed for time and frequency dimensions
time_padding_needed = target_time - decoded_output.shape[2]
freq_padding_needed = target_freq - decoded_output.shape[3]
# Step 3: Apply padding if needed
if time_padding_needed > 0 or freq_padding_needed > 0:
# PyTorch padding format: (pad_left, pad_right, pad_top, pad_bottom)
# For audio: pad_left/right = frequency, pad_top/bottom = time
padding = (
0,
max(freq_padding_needed, 0), # frequency padding (left, right)
0,
max(time_padding_needed, 0), # time padding (top, bottom)
)
decoded_output = F.pad(decoded_output, padding)
# Step 4: Final safety crop to ensure exact target shape
decoded_output = decoded_output[:, :target_channels, :target_time, :target_freq]
return decoded_output
def _run_upsampling_path(self, h: torch.Tensor) -> torch.Tensor:
for level in reversed(range(self.num_resolutions)):
stage = self.up[level]
for block_idx, block in enumerate(stage.block):
h = block(h, temb=None)
if stage.attn:
h = stage.attn[block_idx](h)
if level != 0 and hasattr(stage, "upsample"):
h = stage.upsample(h)
return h
def _finalize_output(self, h: torch.Tensor) -> torch.Tensor:
if self.give_pre_end:
return h
h = self.norm_out(h)
h = self.non_linearity(h)
h = self.conv_out(h)
return torch.tanh(h) if self.tanh_out else h
def decode_audio(latent: torch.Tensor, audio_decoder: "AudioDecoder", vocoder: "Vocoder") -> torch.Tensor:
"""
Decode an audio latent representation using the provided audio decoder and vocoder.
Args:
latent: Input audio latent tensor.
audio_decoder: Model to decode the latent to waveform features.
vocoder: Model to convert decoded features to audio waveform.
Returns:
Decoded audio as a float tensor.
"""
decoded_audio = audio_decoder(latent)
decoded_audio = vocoder(decoded_audio).squeeze(0).float()
return decoded_audio
packages/ltx-core/src/ltx_core/model/audio_vae/causal_conv_2d.py
+110
View File
@@ -0,0 +1,110 @@
import torch
import torch.nn.functional as F
from ltx_core.model.audio_vae.causality_axis import CausalityAxis
class CausalConv2d(torch.nn.Module):
"""
A causal 2D convolution.
This layer ensures that the output at time `t` only depends on inputs
at time `t` and earlier. It achieves this by applying asymmetric padding
to the time dimension (width) before the convolution.
"""
def __init__(
self,
in_channels: int,
out_channels: int,
kernel_size: int | tuple[int, int],
stride: int = 1,
dilation: int | tuple[int, int] = 1,
groups: int = 1,
bias: bool = True,
causality_axis: CausalityAxis = CausalityAxis.HEIGHT,
) -> None:
super().__init__()
self.causality_axis = causality_axis
# Ensure kernel_size and dilation are tuples
kernel_size = torch.nn.modules.utils._pair(kernel_size)
dilation = torch.nn.modules.utils._pair(dilation)
# Calculate padding dimensions
pad_h = (kernel_size[0] - 1) * dilation[0]
pad_w = (kernel_size[1] - 1) * dilation[1]
# The padding tuple for F.pad is (pad_left, pad_right, pad_top, pad_bottom)
match self.causality_axis:
case CausalityAxis.NONE:
self.padding = (pad_w // 2, pad_w - pad_w // 2, pad_h // 2, pad_h - pad_h // 2)
case CausalityAxis.WIDTH | CausalityAxis.WIDTH_COMPATIBILITY:
self.padding = (pad_w, 0, pad_h // 2, pad_h - pad_h // 2)
case CausalityAxis.HEIGHT:
self.padding = (pad_w // 2, pad_w - pad_w // 2, pad_h, 0)
case _:
raise ValueError(f"Invalid causality_axis: {causality_axis}")
# The internal convolution layer uses no padding, as we handle it manually
self.conv = torch.nn.Conv2d(
in_channels,
out_channels,
kernel_size,
stride=stride,
padding=0,
dilation=dilation,
groups=groups,
bias=bias,
)
def forward(self, x: torch.Tensor) -> torch.Tensor:
# Apply causal padding before convolution
x = F.pad(x, self.padding)
return self.conv(x)
def make_conv2d(
in_channels: int,
out_channels: int,
kernel_size: int | tuple[int, int],
stride: int = 1,
padding: tuple[int, int, int, int] | None = None,
dilation: int = 1,
groups: int = 1,
bias: bool = True,
causality_axis: CausalityAxis | None = None,
) -> torch.nn.Module:
"""
Create a 2D convolution layer that can be either causal or non-causal.
Args:
in_channels: Number of input channels
out_channels: Number of output channels
kernel_size: Size of the convolution kernel
stride: Convolution stride
padding: Padding (if None, will be calculated based on causal flag)
dilation: Dilation rate
groups: Number of groups for grouped convolution
bias: Whether to use bias
causality_axis: Dimension along which to apply causality.
Returns:
Either a regular Conv2d or CausalConv2d layer
"""
if causality_axis is not None:
# For causal convolution, padding is handled internally by CausalConv2d
return CausalConv2d(in_channels, out_channels, kernel_size, stride, dilation, groups, bias, causality_axis)
else:
# For non-causal convolution, use symmetric padding if not specified
if padding is None:
padding = kernel_size // 2 if isinstance(kernel_size, int) else tuple(k // 2 for k in kernel_size)
return torch.nn.Conv2d(
in_channels,
out_channels,
kernel_size,
stride,
padding,
dilation,
groups,
bias,
)
packages/ltx-core/src/ltx_core/model/audio_vae/causality_axis.py
+10
View File
@@ -0,0 +1,10 @@
from enum import Enum
class CausalityAxis(Enum):
"""Enum for specifying the causality axis in causal convolutions."""
NONE = None
WIDTH = "width"
HEIGHT = "height"
WIDTH_COMPATIBILITY = "width-compatibility"
packages/ltx-core/src/ltx_core/model/audio_vae/downsample.py
+110
View File
@@ -0,0 +1,110 @@
from typing import Set, Tuple
import torch
from ltx_core.model.audio_vae.attention import AttentionType, make_attn
from ltx_core.model.audio_vae.causality_axis import CausalityAxis
from ltx_core.model.audio_vae.resnet import ResnetBlock
from ltx_core.model.common.normalization import NormType
class Downsample(torch.nn.Module):
"""
A downsampling layer that can use either a strided convolution
or average pooling. Supports standard and causal padding for the
convolutional mode.
"""
def __init__(
self,
in_channels: int,
with_conv: bool,
causality_axis: CausalityAxis = CausalityAxis.WIDTH,
) -> None:
super().__init__()
self.with_conv = with_conv
self.causality_axis = causality_axis
if self.causality_axis != CausalityAxis.NONE and not self.with_conv:
raise ValueError("causality is only supported when `with_conv=True`.")
if self.with_conv:
# Do time downsampling here
# no asymmetric padding in torch conv, must do it ourselves
self.conv = torch.nn.Conv2d(in_channels, in_channels, kernel_size=3, stride=2, padding=0)
def forward(self, x: torch.Tensor) -> torch.Tensor:
if self.with_conv:
# Padding tuple is in the order: (left, right, top, bottom).
match self.causality_axis:
case CausalityAxis.NONE:
pad = (0, 1, 0, 1)
case CausalityAxis.WIDTH:
pad = (2, 0, 0, 1)
case CausalityAxis.HEIGHT:
pad = (0, 1, 2, 0)
case CausalityAxis.WIDTH_COMPATIBILITY:
pad = (1, 0, 0, 1)
case _:
raise ValueError(f"Invalid causality_axis: {self.causality_axis}")
x = torch.nn.functional.pad(x, pad, mode="constant", value=0)
x = self.conv(x)
else:
# This branch is only taken if with_conv=False, which implies causality_axis is NONE.
x = torch.nn.functional.avg_pool2d(x, kernel_size=2, stride=2)
return x
def build_downsampling_path( # noqa: PLR0913
*,
ch: int,
ch_mult: Tuple[int, ...],
num_resolutions: int,
num_res_blocks: int,
resolution: int,
temb_channels: int,
dropout: float,
norm_type: NormType,
causality_axis: CausalityAxis,
attn_type: AttentionType,
attn_resolutions: Set[int],
resamp_with_conv: bool,
) -> tuple[torch.nn.ModuleList, int]:
"""Build the downsampling path with residual blocks, attention, and downsampling layers."""
down_modules = torch.nn.ModuleList()
curr_res = resolution
in_ch_mult = (1, *tuple(ch_mult))
block_in = ch
for i_level in range(num_resolutions):
block = torch.nn.ModuleList()
attn = torch.nn.ModuleList()
block_in = ch * in_ch_mult[i_level]
block_out = ch * ch_mult[i_level]
for _ in range(num_res_blocks):
block.append(
ResnetBlock(
in_channels=block_in,
out_channels=block_out,
temb_channels=temb_channels,
dropout=dropout,
norm_type=norm_type,
causality_axis=causality_axis,
)
)
block_in = block_out
if curr_res in attn_resolutions:
attn.append(make_attn(block_in, attn_type=attn_type, norm_type=norm_type))
down = torch.nn.Module()
down.block = block
down.attn = attn
if i_level != num_resolutions - 1:
down.downsample = Downsample(block_in, resamp_with_conv, causality_axis=causality_axis)
curr_res = curr_res // 2
down_modules.append(down)
return down_modules, block_in
packages/ltx-core/src/ltx_core/model/audio_vae/model_configurator.py
+123
View File
@@ -0,0 +1,123 @@
from ltx_core.loader.sd_ops import SDOps
from ltx_core.model.audio_vae.attention import AttentionType
from ltx_core.model.audio_vae.audio_vae import AudioDecoder, AudioEncoder
from ltx_core.model.audio_vae.causality_axis import CausalityAxis
from ltx_core.model.audio_vae.vocoder import Vocoder
from ltx_core.model.common.normalization import NormType
from ltx_core.model.model_protocol import ModelConfigurator
class VocoderConfigurator(ModelConfigurator[Vocoder]):
@classmethod
def from_config(cls: type[Vocoder], config: dict) -> Vocoder:
config = config.get("vocoder", {})
return Vocoder(
resblock_kernel_sizes=config.get("resblock_kernel_sizes", [3, 7, 11]),
upsample_rates=config.get("upsample_rates", [6, 5, 2, 2, 2]),
upsample_kernel_sizes=config.get("upsample_kernel_sizes", [16, 15, 8, 4, 4]),
resblock_dilation_sizes=config.get("resblock_dilation_sizes", [[1, 3, 5], [1, 3, 5], [1, 3, 5]]),
upsample_initial_channel=config.get("upsample_initial_channel", 1024),
stereo=config.get("stereo", True),
resblock=config.get("resblock", "1"),
output_sample_rate=config.get("output_sample_rate", 24000),
)
VOCODER_COMFY_KEYS_FILTER = (
SDOps("VOCODER_COMFY_KEYS_FILTER").with_matching(prefix="vocoder.").with_replacement("vocoder.", "")
)
class AudioDecoderConfigurator(ModelConfigurator[AudioDecoder]):
@classmethod
def from_config(cls: type[AudioDecoder], config: dict) -> AudioDecoder:
audio_vae_cfg = config.get("audio_vae", {})
model_cfg = audio_vae_cfg.get("model", {})
model_params = model_cfg.get("params", {})
ddconfig = model_params.get("ddconfig", {})
preprocessing_cfg = audio_vae_cfg.get("preprocessing", {})
stft_cfg = preprocessing_cfg.get("stft", {})
mel_cfg = preprocessing_cfg.get("mel", {})
variables_cfg = audio_vae_cfg.get("variables", {})
sample_rate = model_params.get("sampling_rate", 16000)
mel_hop_length = stft_cfg.get("hop_length", 160)
is_causal = stft_cfg.get("causal", True)
mel_bins = ddconfig.get("mel_bins") or mel_cfg.get("n_mel_channels") or variables_cfg.get("mel_bins")
return AudioDecoder(
ch=ddconfig.get("ch", 128),
out_ch=ddconfig.get("out_ch", 2),
ch_mult=tuple(ddconfig.get("ch_mult", (1, 2, 4))),
num_res_blocks=ddconfig.get("num_res_blocks", 2),
attn_resolutions=ddconfig.get("attn_resolutions", {8, 16, 32}),
resolution=ddconfig.get("resolution", 256),
z_channels=ddconfig.get("z_channels", 8),
norm_type=NormType(ddconfig.get("norm_type", "pixel")),
causality_axis=CausalityAxis(ddconfig.get("causality_axis", "height")),
dropout=ddconfig.get("dropout", 0.0),
mid_block_add_attention=ddconfig.get("mid_block_add_attention", True),
sample_rate=sample_rate,
mel_hop_length=mel_hop_length,
is_causal=is_causal,
mel_bins=mel_bins,
)
class AudioEncoderConfigurator(ModelConfigurator[AudioEncoder]):
@classmethod
def from_config(cls: type[AudioEncoder], config: dict) -> AudioEncoder:
audio_vae_cfg = config.get("audio_vae", {})
model_cfg = audio_vae_cfg.get("model", {})
model_params = model_cfg.get("params", {})
ddconfig = model_params.get("ddconfig", {})
preprocessing_cfg = audio_vae_cfg.get("preprocessing", {})
stft_cfg = preprocessing_cfg.get("stft", {})
mel_cfg = preprocessing_cfg.get("mel", {})
variables_cfg = audio_vae_cfg.get("variables", {})
sample_rate = model_params.get("sampling_rate", 16000)
mel_hop_length = stft_cfg.get("hop_length", 160)
n_fft = stft_cfg.get("filter_length", 1024)
is_causal = stft_cfg.get("causal", True)
mel_bins = ddconfig.get("mel_bins") or mel_cfg.get("n_mel_channels") or variables_cfg.get("mel_bins")
return AudioEncoder(
ch=ddconfig.get("ch", 128),
ch_mult=tuple(ddconfig.get("ch_mult", (1, 2, 4))),
num_res_blocks=ddconfig.get("num_res_blocks", 2),
attn_resolutions=ddconfig.get("attn_resolutions", {8, 16, 32}),
resolution=ddconfig.get("resolution", 256),
z_channels=ddconfig.get("z_channels", 8),
double_z=ddconfig.get("double_z", True),
dropout=ddconfig.get("dropout", 0.0),
resamp_with_conv=ddconfig.get("resamp_with_conv", True),
in_channels=ddconfig.get("in_channels", 2),
attn_type=AttentionType(ddconfig.get("attn_type", "vanilla")),
mid_block_add_attention=ddconfig.get("mid_block_add_attention", True),
norm_type=NormType(ddconfig.get("norm_type", "pixel")),
causality_axis=CausalityAxis(ddconfig.get("causality_axis", "height")),
sample_rate=sample_rate,
mel_hop_length=mel_hop_length,
n_fft=n_fft,
is_causal=is_causal,
mel_bins=mel_bins,
)
AUDIO_VAE_DECODER_COMFY_KEYS_FILTER = (
SDOps("AUDIO_VAE_DECODER_COMFY_KEYS_FILTER")
.with_matching(prefix="audio_vae.decoder.")
.with_matching(prefix="audio_vae.per_channel_statistics.")
.with_replacement("audio_vae.decoder.", "")
.with_replacement("audio_vae.per_channel_statistics.", "per_channel_statistics.")
)
AUDIO_VAE_ENCODER_COMFY_KEYS_FILTER = (
SDOps("AUDIO_VAE_ENCODER_COMFY_KEYS_FILTER")
.with_matching(prefix="audio_vae.encoder.")
.with_matching(prefix="audio_vae.per_channel_statistics.")
.with_replacement("audio_vae.encoder.", "")
.with_replacement("audio_vae.per_channel_statistics.", "per_channel_statistics.")
)
packages/ltx-core/src/ltx_core/model/audio_vae/ops.py
+76
View File
@@ -0,0 +1,76 @@
import torch
import torchaudio
from torch import nn
class AudioProcessor(nn.Module):
"""Converts audio waveforms to log-mel spectrograms with optional resampling."""
def __init__(
self,
sample_rate: int,
mel_bins: int,
mel_hop_length: int,
n_fft: int,
) -> None:
super().__init__()
self.sample_rate = sample_rate
self.mel_transform = torchaudio.transforms.MelSpectrogram(
sample_rate=sample_rate,
n_fft=n_fft,
win_length=n_fft,
hop_length=mel_hop_length,
f_min=0.0,
f_max=sample_rate / 2.0,
n_mels=mel_bins,
window_fn=torch.hann_window,
center=True,
pad_mode="reflect",
power=1.0,
mel_scale="slaney",
norm="slaney",
)
def resample_waveform(
self,
waveform: torch.Tensor,
source_rate: int,
target_rate: int,
) -> torch.Tensor:
"""Resample waveform to target sample rate if needed."""
if source_rate == target_rate:
return waveform
resampled = torchaudio.functional.resample(waveform, source_rate, target_rate)
return resampled.to(device=waveform.device, dtype=waveform.dtype)
def waveform_to_mel(
self,
waveform: torch.Tensor,
waveform_sample_rate: int,
) -> torch.Tensor:
"""Convert waveform to log-mel spectrogram [batch, channels, time, n_mels]."""
waveform = self.resample_waveform(waveform, waveform_sample_rate, self.sample_rate)
mel = self.mel_transform(waveform)
mel = torch.log(torch.clamp(mel, min=1e-5))
mel = mel.to(device=waveform.device, dtype=waveform.dtype)
return mel.permute(0, 1, 3, 2).contiguous()
class PerChannelStatistics(nn.Module):
"""
Per-channel statistics for normalizing and denormalizing the latent representation.
This statics is computed over the entire dataset and stored in model's checkpoint under AudioVAE state_dict.
"""
def __init__(self, latent_channels: int = 128) -> None:
super().__init__()
self.register_buffer("std-of-means", torch.empty(latent_channels))
self.register_buffer("mean-of-means", torch.empty(latent_channels))
def un_normalize(self, x: torch.Tensor) -> torch.Tensor:
return (x * self.get_buffer("std-of-means").to(x)) + self.get_buffer("mean-of-means").to(x)
def normalize(self, x: torch.Tensor) -> torch.Tensor:
return (x - self.get_buffer("mean-of-means").to(x)) / self.get_buffer("std-of-means").to(x)
packages/ltx-core/src/ltx_core/model/audio_vae/resnet.py
+176
View File
@@ -0,0 +1,176 @@
from typing import Tuple
import torch
from ltx_core.model.audio_vae.causal_conv_2d import make_conv2d
from ltx_core.model.audio_vae.causality_axis import CausalityAxis
from ltx_core.model.common.normalization import NormType, build_normalization_layer
LRELU_SLOPE = 0.1
class ResBlock1(torch.nn.Module):
def __init__(self, channels: int, kernel_size: int = 3, dilation: Tuple[int, int, int] = (1, 3, 5)):
super(ResBlock1, self).__init__()
self.convs1 = torch.nn.ModuleList(
[
torch.nn.Conv1d(
channels,
channels,
kernel_size,
1,
dilation=dilation[0],
padding="same",
),
torch.nn.Conv1d(
channels,
channels,
kernel_size,
1,
dilation=dilation[1],
padding="same",
),
torch.nn.Conv1d(
channels,
channels,
kernel_size,
1,
dilation=dilation[2],
padding="same",
),
]
)
self.convs2 = torch.nn.ModuleList(
[
torch.nn.Conv1d(
channels,
channels,
kernel_size,
1,
dilation=1,
padding="same",
),
torch.nn.Conv1d(
channels,
channels,
kernel_size,
1,
dilation=1,
padding="same",
),
torch.nn.Conv1d(
channels,
channels,
kernel_size,
1,
dilation=1,
padding="same",
),
]
)
def forward(self, x: torch.Tensor) -> torch.Tensor:
for conv1, conv2 in zip(self.convs1, self.convs2, strict=True):
xt = torch.nn.functional.leaky_relu(x, LRELU_SLOPE)
xt = conv1(xt)
xt = torch.nn.functional.leaky_relu(xt, LRELU_SLOPE)
xt = conv2(xt)
x = xt + x
return x
class ResBlock2(torch.nn.Module):
def __init__(self, channels: int, kernel_size: int = 3, dilation: Tuple[int, int] = (1, 3)):
super(ResBlock2, self).__init__()
self.convs = torch.nn.ModuleList(
[
torch.nn.Conv1d(
channels,
channels,
kernel_size,
1,
dilation=dilation[0],
padding="same",
),
torch.nn.Conv1d(
channels,
channels,
kernel_size,
1,
dilation=dilation[1],
padding="same",
),
]
)
def forward(self, x: torch.Tensor) -> torch.Tensor:
for conv in self.convs:
xt = torch.nn.functional.leaky_relu(x, LRELU_SLOPE)
xt = conv(xt)
x = xt + x
return x
class ResnetBlock(torch.nn.Module):
def __init__(
self,
*,
in_channels: int,
out_channels: int | None = None,
conv_shortcut: bool = False,
dropout: float = 0.0,
temb_channels: int = 512,
norm_type: NormType = NormType.GROUP,
causality_axis: CausalityAxis = CausalityAxis.HEIGHT,
) -> None:
super().__init__()
self.causality_axis = causality_axis
if self.causality_axis != CausalityAxis.NONE and norm_type == NormType.GROUP:
raise ValueError("Causal ResnetBlock with GroupNorm is not supported.")
self.in_channels = in_channels
out_channels = in_channels if out_channels is None else out_channels
self.out_channels = out_channels
self.use_conv_shortcut = conv_shortcut
self.norm1 = build_normalization_layer(in_channels, normtype=norm_type)
self.non_linearity = torch.nn.SiLU()
self.conv1 = make_conv2d(in_channels, out_channels, kernel_size=3, stride=1, causality_axis=causality_axis)
if temb_channels > 0:
self.temb_proj = torch.nn.Linear(temb_channels, out_channels)
self.norm2 = build_normalization_layer(out_channels, normtype=norm_type)
self.dropout = torch.nn.Dropout(dropout)
self.conv2 = make_conv2d(out_channels, out_channels, kernel_size=3, stride=1, causality_axis=causality_axis)
if self.in_channels != self.out_channels:
if self.use_conv_shortcut:
self.conv_shortcut = make_conv2d(
in_channels, out_channels, kernel_size=3, stride=1, causality_axis=causality_axis
)
else:
self.nin_shortcut = make_conv2d(
in_channels, out_channels, kernel_size=1, stride=1, causality_axis=causality_axis
)
def forward(
self,
x: torch.Tensor,
temb: torch.Tensor | None = None,
) -> torch.Tensor:
h = x
h = self.norm1(h)
h = self.non_linearity(h)
h = self.conv1(h)
if temb is not None:
h = h + self.temb_proj(self.non_linearity(temb))[:, :, None, None]
h = self.norm2(h)
h = self.non_linearity(h)
h = self.dropout(h)
h = self.conv2(h)
if self.in_channels != self.out_channels:
x = self.conv_shortcut(x) if self.use_conv_shortcut else self.nin_shortcut(x)
return x + h
packages/ltx-core/src/ltx_core/model/audio_vae/upsample.py
+106
View File
@@ -0,0 +1,106 @@
from typing import Set, Tuple
import torch
from ltx_core.model.audio_vae.attention import AttentionType, make_attn
from ltx_core.model.audio_vae.causal_conv_2d import make_conv2d
from ltx_core.model.audio_vae.causality_axis import CausalityAxis
from ltx_core.model.audio_vae.resnet import ResnetBlock
from ltx_core.model.common.normalization import NormType
class Upsample(torch.nn.Module):
def __init__(
self,
in_channels: int,
with_conv: bool,
causality_axis: CausalityAxis = CausalityAxis.HEIGHT,
) -> None:
super().__init__()
self.with_conv = with_conv
self.causality_axis = causality_axis
if self.with_conv:
self.conv = make_conv2d(in_channels, in_channels, kernel_size=3, stride=1, causality_axis=causality_axis)
def forward(self, x: torch.Tensor) -> torch.Tensor:
x = torch.nn.functional.interpolate(x, scale_factor=2.0, mode="nearest")
if self.with_conv:
x = self.conv(x)
# Drop FIRST element in the causal axis to undo encoder's padding, while keeping the length 1 + 2 * n.
# For example, if the input is [0, 1, 2], after interpolation, the output is [0, 0, 1, 1, 2, 2].
# The causal convolution will pad the first element as [-, -, 0, 0, 1, 1, 2, 2],
# So the output elements rely on the following windows:
# 0: [-,-,0]
# 1: [-,0,0]
# 2: [0,0,1]
# 3: [0,1,1]
# 4: [1,1,2]
# 5: [1,2,2]
# Notice that the first and second elements in the output rely only on the first element in the input,
# while all other elements rely on two elements in the input.
# So we can drop the first element to undo the padding (rather than the last element).
# This is a no-op for non-causal convolutions.
match self.causality_axis:
case CausalityAxis.NONE:
pass # x remains unchanged
case CausalityAxis.HEIGHT:
x = x[:, :, 1:, :]
case CausalityAxis.WIDTH:
x = x[:, :, :, 1:]
case CausalityAxis.WIDTH_COMPATIBILITY:
pass # x remains unchanged
case _:
raise ValueError(f"Invalid causality_axis: {self.causality_axis}")
return x
def build_upsampling_path( # noqa: PLR0913
*,
ch: int,
ch_mult: Tuple[int, ...],
num_resolutions: int,
num_res_blocks: int,
resolution: int,
temb_channels: int,
dropout: float,
norm_type: NormType,
causality_axis: CausalityAxis,
attn_type: AttentionType,
attn_resolutions: Set[int],
resamp_with_conv: bool,
initial_block_channels: int,
) -> tuple[torch.nn.ModuleList, int]:
"""Build the upsampling path with residual blocks, attention, and upsampling layers."""
up_modules = torch.nn.ModuleList()
block_in = initial_block_channels
curr_res = resolution // (2 ** (num_resolutions - 1))
for level in reversed(range(num_resolutions)):
stage = torch.nn.Module()
stage.block = torch.nn.ModuleList()
stage.attn = torch.nn.ModuleList()
block_out = ch * ch_mult[level]
for _ in range(num_res_blocks + 1):
stage.block.append(
ResnetBlock(
in_channels=block_in,
out_channels=block_out,
temb_channels=temb_channels,
dropout=dropout,
norm_type=norm_type,
causality_axis=causality_axis,
)
)
block_in = block_out
if curr_res in attn_resolutions:
stage.attn.append(make_attn(block_in, attn_type=attn_type, norm_type=norm_type))
if level != 0:
stage.upsample = Upsample(block_in, resamp_with_conv, causality_axis=causality_axis)
curr_res *= 2
up_modules.insert(0, stage)
return up_modules, block_in
packages/ltx-core/src/ltx_core/model/audio_vae/vocoder.py
+123
View File
@@ -0,0 +1,123 @@
import math
from typing import List
import einops
import torch
import torch.nn.functional as F
from torch import nn
from ltx_core.model.audio_vae.resnet import LRELU_SLOPE, ResBlock1, ResBlock2
class Vocoder(torch.nn.Module):
"""
Vocoder model for synthesizing audio from Mel spectrograms.
Args:
resblock_kernel_sizes: List of kernel sizes for the residual blocks.
This value is read from the checkpoint at `config.vocoder.resblock_kernel_sizes`.
upsample_rates: List of upsampling rates.
This value is read from the checkpoint at `config.vocoder.upsample_rates`.
upsample_kernel_sizes: List of kernel sizes for the upsampling layers.
This value is read from the checkpoint at `config.vocoder.upsample_kernel_sizes`.
resblock_dilation_sizes: List of dilation sizes for the residual blocks.
This value is read from the checkpoint at `config.vocoder.resblock_dilation_sizes`.
upsample_initial_channel: Initial number of channels for the upsampling layers.
This value is read from the checkpoint at `config.vocoder.upsample_initial_channel`.
stereo: Whether to use stereo output.
This value is read from the checkpoint at `config.vocoder.stereo`.
resblock: Type of residual block to use.
This value is read from the checkpoint at `config.vocoder.resblock`.
output_sample_rate: Waveform sample rate.
This value is read from the checkpoint at `config.vocoder.output_sample_rate`.
"""
def __init__(
self,
resblock_kernel_sizes: List[int] | None = None,
upsample_rates: List[int] | None = None,
upsample_kernel_sizes: List[int] | None = None,
resblock_dilation_sizes: List[List[int]] | None = None,
upsample_initial_channel: int = 1024,
stereo: bool = True,
resblock: str = "1",
output_sample_rate: int = 24000,
):
super().__init__()
# Initialize default values if not provided. Note that mutable default values are not supported.
if resblock_kernel_sizes is None:
resblock_kernel_sizes = [3, 7, 11]
if upsample_rates is None:
upsample_rates = [6, 5, 2, 2, 2]
if upsample_kernel_sizes is None:
upsample_kernel_sizes = [16, 15, 8, 4, 4]
if resblock_dilation_sizes is None:
resblock_dilation_sizes = [[1, 3, 5], [1, 3, 5], [1, 3, 5]]
self.output_sample_rate = output_sample_rate
self.num_kernels = len(resblock_kernel_sizes)
self.num_upsamples = len(upsample_rates)
in_channels = 128 if stereo else 64
self.conv_pre = nn.Conv1d(in_channels, upsample_initial_channel, 7, 1, padding=3)
resblock_class = ResBlock1 if resblock == "1" else ResBlock2
self.ups = nn.ModuleList()
for i, (stride, kernel_size) in enumerate(zip(upsample_rates, upsample_kernel_sizes, strict=True)):
self.ups.append(
nn.ConvTranspose1d(
upsample_initial_channel // (2**i),
upsample_initial_channel // (2 ** (i + 1)),
kernel_size,
stride,
padding=(kernel_size - stride) // 2,
)
)
self.resblocks = nn.ModuleList()
for i, _ in enumerate(self.ups):
ch = upsample_initial_channel // (2 ** (i + 1))
for kernel_size, dilations in zip(resblock_kernel_sizes, resblock_dilation_sizes, strict=True):
self.resblocks.append(resblock_class(ch, kernel_size, dilations))
out_channels = 2 if stereo else 1
final_channels = upsample_initial_channel // (2**self.num_upsamples)
self.conv_post = nn.Conv1d(final_channels, out_channels, 7, 1, padding=3)
self.upsample_factor = math.prod(layer.stride[0] for layer in self.ups)
def forward(self, x: torch.Tensor) -> torch.Tensor:
"""
Forward pass of the vocoder.
Args:
x: Input Mel spectrogram tensor. Can be either:
- 3D: (batch_size, time, mel_bins) for mono
- 4D: (batch_size, 2, time, mel_bins) for stereo
Returns:
Audio waveform tensor of shape (batch_size, out_channels, audio_length)
"""
x = x.transpose(2, 3) # (batch, channels, time, mel_bins) -> (batch, channels, mel_bins, time)
if x.dim() == 4: # stereo
assert x.shape[1] == 2, "Input must have 2 channels for stereo"
x = einops.rearrange(x, "b s c t -> b (s c) t")
x = self.conv_pre(x)
for i in range(self.num_upsamples):
x = F.leaky_relu(x, LRELU_SLOPE)
x = self.ups[i](x)
start = i * self.num_kernels
end = start + self.num_kernels
# Evaluate all resblocks with the same input tensor so they can run
# independently (and thus in parallel on accelerator hardware) before
# aggregating their outputs via mean.
block_outputs = torch.stack(
[self.resblocks[idx](x) for idx in range(start, end)],
dim=0,
)
x = block_outputs.mean(dim=0)
x = self.conv_post(F.leaky_relu(x))
return torch.tanh(x)
packages/ltx-core/src/ltx_core/model/common/__init__.py
+9
View File
@@ -0,0 +1,9 @@
"""Common model utilities."""
from ltx_core.model.common.normalization import NormType, PixelNorm, build_normalization_layer
__all__ = [
"NormType",
"PixelNorm",
"build_normalization_layer",
]
packages/ltx-core/src/ltx_core/model/common/normalization.py
+59
View File
@@ -0,0 +1,59 @@
from enum import Enum
import torch
from torch import nn
class NormType(Enum):
"""Normalization layer types: GROUP (GroupNorm) or PIXEL (per-location RMS norm)."""
GROUP = "group"
PIXEL = "pixel"
class PixelNorm(nn.Module):
"""
Per-pixel (per-location) RMS normalization layer.
For each element along the chosen dimension, this layer normalizes the tensor
by the root-mean-square of its values across that dimension:
y = x / sqrt(mean(x^2, dim=dim, keepdim=True) + eps)
"""
def __init__(self, dim: int = 1, eps: float = 1e-8) -> None:
"""
Args:
dim: Dimension along which to compute the RMS (typically channels).
eps: Small constant added for numerical stability.
"""
super().__init__()
self.dim = dim
self.eps = eps
def forward(self, x: torch.Tensor) -> torch.Tensor:
"""
Apply RMS normalization along the configured dimension.
"""
# Compute mean of squared values along `dim`, keep dimensions for broadcasting.
mean_sq = torch.mean(x**2, dim=self.dim, keepdim=True)
# Normalize by the root-mean-square (RMS).
rms = torch.sqrt(mean_sq + self.eps)
return x / rms
def build_normalization_layer(
in_channels: int, *, num_groups: int = 32, normtype: NormType = NormType.GROUP
) -> nn.Module:
"""
Create a normalization layer based on the normalization type.
Args:
in_channels: Number of input channels
num_groups: Number of groups for group normalization
normtype: Type of normalization: "group" or "pixel"
Returns:
A normalization layer
"""
if normtype == NormType.GROUP:
return torch.nn.GroupNorm(num_groups=num_groups, num_channels=in_channels, eps=1e-6, affine=True)
if normtype == NormType.PIXEL:
return PixelNorm(dim=1, eps=1e-6)
raise ValueError(f"Invalid normalization type: {normtype}")
packages/ltx-core/src/ltx_core/model/model_protocol.py
+10
View File
@@ -0,0 +1,10 @@
from typing import Protocol, TypeVar
ModelType = TypeVar("ModelType")
class ModelConfigurator(Protocol[ModelType]):
"""Protocol for model loader classes that instantiates models from a configuration dictionary."""
@classmethod
def from_config(cls, config: dict) -> ModelType: ...
packages/ltx-core/src/ltx_core/model/transformer/__init__.py
+24
View File
@@ -0,0 +1,24 @@
"""Transformer model components."""
from ltx_core.model.transformer.modality import Modality
from ltx_core.model.transformer.model import LTXModel, X0Model
from ltx_core.model.transformer.model_configurator import (
LTXV_MODEL_COMFY_RENAMING_MAP,
LTXV_MODEL_COMFY_RENAMING_WITH_TRANSFORMER_LINEAR_DOWNCAST_MAP,
UPCAST_DURING_INFERENCE,
LTXModelConfigurator,
LTXVideoOnlyModelConfigurator,
UpcastWithStochasticRounding,
)
__all__ = [
"LTXV_MODEL_COMFY_RENAMING_MAP",
"LTXV_MODEL_COMFY_RENAMING_WITH_TRANSFORMER_LINEAR_DOWNCAST_MAP",
"UPCAST_DURING_INFERENCE",
"LTXModel",
"LTXModelConfigurator",
"LTXVideoOnlyModelConfigurator",
"Modality",
"UpcastWithStochasticRounding",
"X0Model",
]
packages/ltx-core/src/ltx_core/model/transformer/adaln.py
+34
View File
@@ -0,0 +1,34 @@
from typing import Optional, Tuple
import torch
from ltx_core.model.transformer.timestep_embedding import PixArtAlphaCombinedTimestepSizeEmbeddings
class AdaLayerNormSingle(torch.nn.Module):
r"""
Norm layer adaptive layer norm single (adaLN-single).
As proposed in PixArt-Alpha (see: https://arxiv.org/abs/2310.00426; Section 2.3).
Parameters:
embedding_dim (`int`): The size of each embedding vector.
use_additional_conditions (`bool`): To use additional conditions for normalization or not.
"""
def __init__(self, embedding_dim: int, embedding_coefficient: int = 6):
super().__init__()
self.emb = PixArtAlphaCombinedTimestepSizeEmbeddings(
embedding_dim,
size_emb_dim=embedding_dim // 3,
)
self.silu = torch.nn.SiLU()
self.linear = torch.nn.Linear(embedding_dim, embedding_coefficient * embedding_dim, bias=True)
def forward(
self,
timestep: torch.Tensor,
hidden_dtype: Optional[torch.dtype] = None,
) -> Tuple[torch.Tensor, torch.Tensor]:
embedded_timestep = self.emb(timestep, hidden_dtype=hidden_dtype)
return self.linear(self.silu(embedded_timestep)), embedded_timestep
packages/ltx-core/src/ltx_core/model/transformer/attention.py
+195
View File
@@ -0,0 +1,195 @@
from enum import Enum
from typing import Protocol
import torch
from ltx_core.model.transformer.rope import LTXRopeType, apply_rotary_emb
memory_efficient_attention = None
flash_attn_interface = None
try:
from xformers.ops import memory_efficient_attention
except ImportError:
memory_efficient_attention = None
try:
# FlashAttention3 and XFormersAttention cannot be used together
if memory_efficient_attention is None:
import flash_attn_interface
except ImportError:
flash_attn_interface = None
class AttentionCallable(Protocol):
def __call__(
self, q: torch.Tensor, k: torch.Tensor, v: torch.Tensor, heads: int, mask: torch.Tensor | None = None
) -> torch.Tensor: ...
class PytorchAttention(AttentionCallable):
def __call__(
self, q: torch.Tensor, k: torch.Tensor, v: torch.Tensor, heads: int, mask: torch.Tensor | None = None
) -> torch.Tensor:
b, _, dim_head = q.shape
dim_head //= heads
q, k, v = (t.view(b, -1, heads, dim_head).transpose(1, 2) for t in (q, k, v))
if mask is not None:
# add a batch dimension if there isn't already one
if mask.ndim == 2:
mask = mask.unsqueeze(0)
# add a heads dimension if there isn't already one
if mask.ndim == 3:
mask = mask.unsqueeze(1)
out = torch.nn.functional.scaled_dot_product_attention(q, k, v, attn_mask=mask, dropout_p=0.0, is_causal=False)
out = out.transpose(1, 2).reshape(b, -1, heads * dim_head)
return out
class XFormersAttention(AttentionCallable):
def __call__(
self,
q: torch.Tensor,
k: torch.Tensor,
v: torch.Tensor,
heads: int,
mask: torch.Tensor | None = None,
) -> torch.Tensor:
if memory_efficient_attention is None:
raise RuntimeError("XFormersAttention was selected but `xformers` is not installed.")
b, _, dim_head = q.shape
dim_head //= heads
# xformers expects [B, M, H, K]
q, k, v = (t.view(b, -1, heads, dim_head) for t in (q, k, v))
if mask is not None:
# add a singleton batch dimension
if mask.ndim == 2:
mask = mask.unsqueeze(0)
# add a singleton heads dimension
if mask.ndim == 3:
mask = mask.unsqueeze(1)
# pad to a multiple of 8
pad = 8 - mask.shape[-1] % 8
# the xformers docs says that it's allowed to have a mask of shape (1, Nq, Nk)
# but when using separated heads, the shape has to be (B, H, Nq, Nk)
# in flux, this matrix ends up being over 1GB
# here, we create a mask with the same batch/head size as the input mask (potentially singleton or full)
mask_out = torch.empty(
[mask.shape[0], mask.shape[1], q.shape[1], mask.shape[-1] + pad], dtype=q.dtype, device=q.device
)
mask_out[..., : mask.shape[-1]] = mask
# doesn't this remove the padding again??
mask = mask_out[..., : mask.shape[-1]]
mask = mask.expand(b, heads, -1, -1)
out = memory_efficient_attention(q.to(v.dtype), k.to(v.dtype), v, attn_bias=mask, p=0.0)
out = out.reshape(b, -1, heads * dim_head)
return out
class FlashAttention3(AttentionCallable):
def __call__(
self,
q: torch.Tensor,
k: torch.Tensor,
v: torch.Tensor,
heads: int,
mask: torch.Tensor | None = None,
) -> torch.Tensor:
if flash_attn_interface is None:
raise RuntimeError("FlashAttention3 was selected but `FlashAttention3` is not installed.")
b, _, dim_head = q.shape
dim_head //= heads
q, k, v = (t.view(b, -1, heads, dim_head) for t in (q, k, v))
if mask is not None:
raise NotImplementedError("Mask is not supported for FlashAttention3")
out = flash_attn_interface.flash_attn_func(q.to(v.dtype), k.to(v.dtype), v)
out = out.reshape(b, -1, heads * dim_head)
return out
class AttentionFunction(Enum):
PYTORCH = "pytorch"
XFORMERS = "xformers"
FLASH_ATTENTION_3 = "flash_attention_3"
DEFAULT = "default"
def __call__(
self, q: torch.Tensor, k: torch.Tensor, v: torch.Tensor, heads: int, mask: torch.Tensor | None = None
) -> torch.Tensor:
if self is AttentionFunction.PYTORCH:
return PytorchAttention()(q, k, v, heads, mask)
elif self is AttentionFunction.XFORMERS:
return XFormersAttention()(q, k, v, heads, mask)
elif self is AttentionFunction.FLASH_ATTENTION_3:
return FlashAttention3()(q, k, v, heads, mask)
else:
# Default behavior: XFormers if installed else - PyTorch
return (
XFormersAttention()(q, k, v, heads, mask)
if memory_efficient_attention is not None
else PytorchAttention()(q, k, v, heads, mask)
)
class Attention(torch.nn.Module):
def __init__(
self,
query_dim: int,
context_dim: int | None = None,
heads: int = 8,
dim_head: int = 64,
norm_eps: float = 1e-6,
rope_type: LTXRopeType = LTXRopeType.INTERLEAVED,
attention_function: AttentionCallable | AttentionFunction = AttentionFunction.DEFAULT,
) -> None:
super().__init__()
self.rope_type = rope_type
self.attention_function = attention_function
inner_dim = dim_head * heads
context_dim = query_dim if context_dim is None else context_dim
self.heads = heads
self.dim_head = dim_head
self.q_norm = torch.nn.RMSNorm(inner_dim, eps=norm_eps)
self.k_norm = torch.nn.RMSNorm(inner_dim, eps=norm_eps)
self.to_q = torch.nn.Linear(query_dim, inner_dim, bias=True)
self.to_k = torch.nn.Linear(context_dim, inner_dim, bias=True)
self.to_v = torch.nn.Linear(context_dim, inner_dim, bias=True)
self.to_out = torch.nn.Sequential(torch.nn.Linear(inner_dim, query_dim, bias=True), torch.nn.Identity())
def forward(
self,
x: torch.Tensor,
context: torch.Tensor | None = None,
mask: torch.Tensor | None = None,
pe: torch.Tensor | None = None,
k_pe: torch.Tensor | None = None,
) -> torch.Tensor:
q = self.to_q(x)
context = x if context is None else context
k = self.to_k(context)
v = self.to_v(context)
q = self.q_norm(q)
k = self.k_norm(k)
if pe is not None:
q = apply_rotary_emb(q, pe, self.rope_type)
k = apply_rotary_emb(k, pe if k_pe is None else k_pe, self.rope_type)
# attention_function can be an enum *or* a custom callable
out = self.attention_function(q, k, v, self.heads, mask)
return self.to_out(out)
packages/ltx-core/src/ltx_core/model/transformer/feed_forward.py
+15
View File
@@ -0,0 +1,15 @@
import torch
from ltx_core.model.transformer.gelu_approx import GELUApprox
class FeedForward(torch.nn.Module):
def __init__(self, dim: int, dim_out: int, mult: int = 4) -> None:
super().__init__()
inner_dim = int(dim * mult)
project_in = GELUApprox(dim, inner_dim)
self.net = torch.nn.Sequential(project_in, torch.nn.Identity(), torch.nn.Linear(inner_dim, dim_out))
def forward(self, x: torch.Tensor) -> torch.Tensor:
return self.net(x)
packages/ltx-core/src/ltx_core/model/transformer/gelu_approx.py
+10
View File
@@ -0,0 +1,10 @@
import torch
class GELUApprox(torch.nn.Module):
def __init__(self, dim_in: int, dim_out: int) -> None:
super().__init__()
self.proj = torch.nn.Linear(dim_in, dim_out)
def forward(self, x: torch.Tensor) -> torch.Tensor:
return torch.nn.functional.gelu(self.proj(x), approximate="tanh")
packages/ltx-core/src/ltx_core/model/transformer/modality.py
+23
View File
@@ -0,0 +1,23 @@
from dataclasses import dataclass
import torch
@dataclass(frozen=True)
class Modality:
"""
Input data for a single modality (video or audio) in the transformer.
Bundles the latent tokens, timestep embeddings, positional information,
and text conditioning context for processing by the diffusion transformer.
"""
latent: (
torch.Tensor
) # Shape: (B, T, D) where B is the batch size, T is the number of tokens, and D is input dimension
timesteps: torch.Tensor # Shape: (B, T) where T is the number of timesteps
positions: (
torch.Tensor
) # Shape: (B, 3, T) for video, where 3 is the number of dimensions and T is the number of tokens
context: torch.Tensor
enabled: bool = True
context_mask: torch.Tensor | None = None
packages/ltx-core/src/ltx_core/model/transformer/model.py
+468
View File
@@ -0,0 +1,468 @@
from enum import Enum
import torch
from ltx_core.guidance.perturbations import BatchedPerturbationConfig
from ltx_core.model.transformer.adaln import AdaLayerNormSingle
from ltx_core.model.transformer.attention import AttentionCallable, AttentionFunction
from ltx_core.model.transformer.modality import Modality
from ltx_core.model.transformer.rope import LTXRopeType
from ltx_core.model.transformer.text_projection import PixArtAlphaTextProjection
from ltx_core.model.transformer.transformer import BasicAVTransformerBlock, TransformerConfig
from ltx_core.model.transformer.transformer_args import (
MultiModalTransformerArgsPreprocessor,
TransformerArgs,
TransformerArgsPreprocessor,
)
from ltx_core.utils import to_denoised
class LTXModelType(Enum):
AudioVideo = "ltx av model"
VideoOnly = "ltx video only model"
AudioOnly = "ltx audio only model"
def is_video_enabled(self) -> bool:
return self in (LTXModelType.AudioVideo, LTXModelType.VideoOnly)
def is_audio_enabled(self) -> bool:
return self in (LTXModelType.AudioVideo, LTXModelType.AudioOnly)
class LTXModel(torch.nn.Module):
"""
LTX model transformer implementation.
This class implements the transformer blocks for the LTX model.
"""
def __init__( # noqa: PLR0913
self,
*,
model_type: LTXModelType = LTXModelType.AudioVideo,
num_attention_heads: int = 32,
attention_head_dim: int = 128,
in_channels: int = 128,
out_channels: int = 128,
num_layers: int = 48,
cross_attention_dim: int = 4096,
norm_eps: float = 1e-06,
attention_type: AttentionFunction | AttentionCallable = AttentionFunction.DEFAULT,
caption_channels: int = 3840,
positional_embedding_theta: float = 10000.0,
positional_embedding_max_pos: list[int] | None = None,
timestep_scale_multiplier: int = 1000,
use_middle_indices_grid: bool = True,
audio_num_attention_heads: int = 32,
audio_attention_head_dim: int = 64,
audio_in_channels: int = 128,
audio_out_channels: int = 128,
audio_cross_attention_dim: int = 2048,
audio_positional_embedding_max_pos: list[int] | None = None,
av_ca_timestep_scale_multiplier: int = 1,
rope_type: LTXRopeType = LTXRopeType.INTERLEAVED,
double_precision_rope: bool = False,
):
super().__init__()
self._enable_gradient_checkpointing = False
self.use_middle_indices_grid = use_middle_indices_grid
self.rope_type = rope_type
self.double_precision_rope = double_precision_rope
self.timestep_scale_multiplier = timestep_scale_multiplier
self.positional_embedding_theta = positional_embedding_theta
self.model_type = model_type
cross_pe_max_pos = None
if model_type.is_video_enabled():
if positional_embedding_max_pos is None:
positional_embedding_max_pos = [20, 2048, 2048]
self.positional_embedding_max_pos = positional_embedding_max_pos
self.num_attention_heads = num_attention_heads
self.inner_dim = num_attention_heads * attention_head_dim
self._init_video(
in_channels=in_channels,
out_channels=out_channels,
caption_channels=caption_channels,
norm_eps=norm_eps,
)
if model_type.is_audio_enabled():
if audio_positional_embedding_max_pos is None:
audio_positional_embedding_max_pos = [20]
self.audio_positional_embedding_max_pos = audio_positional_embedding_max_pos
self.audio_num_attention_heads = audio_num_attention_heads
self.audio_inner_dim = self.audio_num_attention_heads * audio_attention_head_dim
self._init_audio(
in_channels=audio_in_channels,
out_channels=audio_out_channels,
caption_channels=caption_channels,
norm_eps=norm_eps,
)
if model_type.is_video_enabled() and model_type.is_audio_enabled():
cross_pe_max_pos = max(self.positional_embedding_max_pos[0], self.audio_positional_embedding_max_pos[0])
self.av_ca_timestep_scale_multiplier = av_ca_timestep_scale_multiplier
self.audio_cross_attention_dim = audio_cross_attention_dim
self._init_audio_video(num_scale_shift_values=4)
self._init_preprocessors(cross_pe_max_pos)
# Initialize transformer blocks
self._init_transformer_blocks(
num_layers=num_layers,
attention_head_dim=attention_head_dim if model_type.is_video_enabled() else 0,
cross_attention_dim=cross_attention_dim,
audio_attention_head_dim=audio_attention_head_dim if model_type.is_audio_enabled() else 0,
audio_cross_attention_dim=audio_cross_attention_dim,
norm_eps=norm_eps,
attention_type=attention_type,
)
def _init_video(
self,
in_channels: int,
out_channels: int,
caption_channels: int,
norm_eps: float,
) -> None:
"""Initialize video-specific components."""
# Video input components
self.patchify_proj = torch.nn.Linear(in_channels, self.inner_dim, bias=True)
self.adaln_single = AdaLayerNormSingle(self.inner_dim)
# Video caption projection
self.caption_projection = PixArtAlphaTextProjection(
in_features=caption_channels,
hidden_size=self.inner_dim,
)
# Video output components
self.scale_shift_table = torch.nn.Parameter(torch.empty(2, self.inner_dim))
self.norm_out = torch.nn.LayerNorm(self.inner_dim, elementwise_affine=False, eps=norm_eps)
self.proj_out = torch.nn.Linear(self.inner_dim, out_channels)
def _init_audio(
self,
in_channels: int,
out_channels: int,
caption_channels: int,
norm_eps: float,
) -> None:
"""Initialize audio-specific components."""
# Audio input components
self.audio_patchify_proj = torch.nn.Linear(in_channels, self.audio_inner_dim, bias=True)
self.audio_adaln_single = AdaLayerNormSingle(
self.audio_inner_dim,
)
# Audio caption projection
self.audio_caption_projection = PixArtAlphaTextProjection(
in_features=caption_channels,
hidden_size=self.audio_inner_dim,
)
# Audio output components
self.audio_scale_shift_table = torch.nn.Parameter(torch.empty(2, self.audio_inner_dim))
self.audio_norm_out = torch.nn.LayerNorm(self.audio_inner_dim, elementwise_affine=False, eps=norm_eps)
self.audio_proj_out = torch.nn.Linear(self.audio_inner_dim, out_channels)
def _init_audio_video(
self,
num_scale_shift_values: int,
) -> None:
"""Initialize audio-video cross-attention components."""
self.av_ca_video_scale_shift_adaln_single = AdaLayerNormSingle(
self.inner_dim,
embedding_coefficient=num_scale_shift_values,
)
self.av_ca_audio_scale_shift_adaln_single = AdaLayerNormSingle(
self.audio_inner_dim,
embedding_coefficient=num_scale_shift_values,
)
self.av_ca_a2v_gate_adaln_single = AdaLayerNormSingle(
self.inner_dim,
embedding_coefficient=1,
)
self.av_ca_v2a_gate_adaln_single = AdaLayerNormSingle(
self.audio_inner_dim,
embedding_coefficient=1,
)
def _init_preprocessors(
self,
cross_pe_max_pos: int | None = None,
) -> None:
"""Initialize preprocessors for LTX."""
if self.model_type.is_video_enabled() and self.model_type.is_audio_enabled():
self.video_args_preprocessor = MultiModalTransformerArgsPreprocessor(
patchify_proj=self.patchify_proj,
adaln=self.adaln_single,
caption_projection=self.caption_projection,
cross_scale_shift_adaln=self.av_ca_video_scale_shift_adaln_single,
cross_gate_adaln=self.av_ca_a2v_gate_adaln_single,
inner_dim=self.inner_dim,
max_pos=self.positional_embedding_max_pos,
num_attention_heads=self.num_attention_heads,
cross_pe_max_pos=cross_pe_max_pos,
use_middle_indices_grid=self.use_middle_indices_grid,
audio_cross_attention_dim=self.audio_cross_attention_dim,
timestep_scale_multiplier=self.timestep_scale_multiplier,
double_precision_rope=self.double_precision_rope,
positional_embedding_theta=self.positional_embedding_theta,
rope_type=self.rope_type,
av_ca_timestep_scale_multiplier=self.av_ca_timestep_scale_multiplier,
)
self.audio_args_preprocessor = MultiModalTransformerArgsPreprocessor(
patchify_proj=self.audio_patchify_proj,
adaln=self.audio_adaln_single,
caption_projection=self.audio_caption_projection,
cross_scale_shift_adaln=self.av_ca_audio_scale_shift_adaln_single,
cross_gate_adaln=self.av_ca_v2a_gate_adaln_single,
inner_dim=self.audio_inner_dim,
max_pos=self.audio_positional_embedding_max_pos,
num_attention_heads=self.audio_num_attention_heads,
cross_pe_max_pos=cross_pe_max_pos,
use_middle_indices_grid=self.use_middle_indices_grid,
audio_cross_attention_dim=self.audio_cross_attention_dim,
timestep_scale_multiplier=self.timestep_scale_multiplier,
double_precision_rope=self.double_precision_rope,
positional_embedding_theta=self.positional_embedding_theta,
rope_type=self.rope_type,
av_ca_timestep_scale_multiplier=self.av_ca_timestep_scale_multiplier,
)
elif self.model_type.is_video_enabled():
self.video_args_preprocessor = TransformerArgsPreprocessor(
patchify_proj=self.patchify_proj,
adaln=self.adaln_single,
caption_projection=self.caption_projection,
inner_dim=self.inner_dim,
max_pos=self.positional_embedding_max_pos,
num_attention_heads=self.num_attention_heads,
use_middle_indices_grid=self.use_middle_indices_grid,
timestep_scale_multiplier=self.timestep_scale_multiplier,
double_precision_rope=self.double_precision_rope,
positional_embedding_theta=self.positional_embedding_theta,
rope_type=self.rope_type,
)
elif self.model_type.is_audio_enabled():
self.audio_args_preprocessor = TransformerArgsPreprocessor(
patchify_proj=self.audio_patchify_proj,
adaln=self.audio_adaln_single,
caption_projection=self.audio_caption_projection,
inner_dim=self.audio_inner_dim,
max_pos=self.audio_positional_embedding_max_pos,
num_attention_heads=self.audio_num_attention_heads,
use_middle_indices_grid=self.use_middle_indices_grid,
timestep_scale_multiplier=self.timestep_scale_multiplier,
double_precision_rope=self.double_precision_rope,
positional_embedding_theta=self.positional_embedding_theta,
rope_type=self.rope_type,
)
def _init_transformer_blocks(
self,
num_layers: int,
attention_head_dim: int,
cross_attention_dim: int,
audio_attention_head_dim: int,
audio_cross_attention_dim: int,
norm_eps: float,
attention_type: AttentionFunction | AttentionCallable,
) -> None:
"""Initialize transformer blocks for LTX."""
video_config = (
TransformerConfig(
dim=self.inner_dim,
heads=self.num_attention_heads,
d_head=attention_head_dim,
context_dim=cross_attention_dim,
)
if self.model_type.is_video_enabled()
else None
)
audio_config = (
TransformerConfig(
dim=self.audio_inner_dim,
heads=self.audio_num_attention_heads,
d_head=audio_attention_head_dim,
context_dim=audio_cross_attention_dim,
)
if self.model_type.is_audio_enabled()
else None
)
self.transformer_blocks = torch.nn.ModuleList(
[
BasicAVTransformerBlock(
idx=idx,
video=video_config,
audio=audio_config,
rope_type=self.rope_type,
norm_eps=norm_eps,
attention_function=attention_type,
)
for idx in range(num_layers)
]
)
def set_gradient_checkpointing(self, enable: bool) -> None:
"""Enable or disable gradient checkpointing for transformer blocks.
Gradient checkpointing trades compute for memory by recomputing activations
during the backward pass instead of storing them. This can significantly
reduce memory usage at the cost of ~20-30% slower training.
Args:
enable: Whether to enable gradient checkpointing
"""
self._enable_gradient_checkpointing = enable
def _process_transformer_blocks(
self,
video: TransformerArgs | None,
audio: TransformerArgs | None,
perturbations: BatchedPerturbationConfig,
) -> tuple[TransformerArgs, TransformerArgs]:
"""Process transformer blocks for LTXAV."""
# Process transformer blocks
for block in self.transformer_blocks:
if self._enable_gradient_checkpointing and self.training:
# Use gradient checkpointing to save memory during training.
# With use_reentrant=False, we can pass dataclasses directly -
# PyTorch will track all tensor leaves in the computation graph.
video, audio = torch.utils.checkpoint.checkpoint(
block,
video,
audio,
perturbations,
use_reentrant=False,
)
else:
video, audio = block(
video=video,
audio=audio,
perturbations=perturbations,
)
return video, audio
def _process_output(
self,
scale_shift_table: torch.Tensor,
norm_out: torch.nn.LayerNorm,
proj_out: torch.nn.Linear,
x: torch.Tensor,
embedded_timestep: torch.Tensor,
) -> torch.Tensor:
"""Process output for LTXV."""
# Apply scale-shift modulation
scale_shift_values = (
scale_shift_table[None, None].to(device=x.device, dtype=x.dtype) + embedded_timestep[:, :, None]
)
shift, scale = scale_shift_values[:, :, 0], scale_shift_values[:, :, 1]
x = norm_out(x)
x = x * (1 + scale) + shift
x = proj_out(x)
return x
def forward(
self, video: Modality | None, audio: Modality | None, perturbations: BatchedPerturbationConfig
) -> tuple[torch.Tensor, torch.Tensor]:
"""
Forward pass for LTX models.
Returns:
Processed output tensors
"""
if not self.model_type.is_video_enabled() and video is not None:
raise ValueError("Video is not enabled for this model")
if not self.model_type.is_audio_enabled() and audio is not None:
raise ValueError("Audio is not enabled for this model")
video_args = self.video_args_preprocessor.prepare(video) if video is not None else None
audio_args = self.audio_args_preprocessor.prepare(audio) if audio is not None else None
# Process transformer blocks
video_out, audio_out = self._process_transformer_blocks(
video=video_args,
audio=audio_args,
perturbations=perturbations,
)
# Process output
vx = (
self._process_output(
self.scale_shift_table, self.norm_out, self.proj_out, video_out.x, video_out.embedded_timestep
)
if video_out is not None
else None
)
ax = (
self._process_output(
self.audio_scale_shift_table,
self.audio_norm_out,
self.audio_proj_out,
audio_out.x,
audio_out.embedded_timestep,
)
if audio_out is not None
else None
)
return vx, ax
class LegacyX0Model(torch.nn.Module):
"""
Legacy X0 model implementation.
Returns fully denoised output based on the velocities produced by the base model.
"""
def __init__(self, velocity_model: LTXModel):
super().__init__()
self.velocity_model = velocity_model
def forward(
self,
video: Modality | None,
audio: Modality | None,
perturbations: BatchedPerturbationConfig,
sigma: float,
) -> tuple[torch.Tensor | None, torch.Tensor | None]:
"""
Denoise the video and audio according to the sigma.
Returns:
Denoised video and audio
"""
vx, ax = self.velocity_model(video, audio, perturbations)
denoised_video = to_denoised(video.latent, vx, sigma) if vx is not None else None
denoised_audio = to_denoised(audio.latent, ax, sigma) if ax is not None else None
return denoised_video, denoised_audio
class X0Model(torch.nn.Module):
"""
X0 model implementation.
Returns fully denoised outputs based on the velocities produced by the base model.
Applies scaled denoising to the video and audio according to the timesteps = sigma * denoising_mask.
"""
def __init__(self, velocity_model: LTXModel):
super().__init__()
self.velocity_model = velocity_model
def forward(
self,
video: Modality | None,
audio: Modality | None,
perturbations: BatchedPerturbationConfig,
) -> tuple[torch.Tensor | None, torch.Tensor | None]:
"""
Denoise the video and audio according to the sigma.
Returns:
Denoised video and audio
"""
vx, ax = self.velocity_model(video, audio, perturbations)
denoised_video = to_denoised(video.latent, vx, video.timesteps) if vx is not None else None
denoised_audio = to_denoised(audio.latent, ax, audio.timesteps) if ax is not None else None
return denoised_video, denoised_audio
packages/ltx-core/src/ltx_core/model/transformer/model_configurator.py
+237
View File
@@ -0,0 +1,237 @@
import torch
from ltx_core.loader.fuse_loras import fused_add_round_launch
from ltx_core.loader.module_ops import ModuleOps
from ltx_core.loader.sd_ops import KeyValueOperationResult, SDOps
from ltx_core.model.model_protocol import ModelConfigurator
from ltx_core.model.transformer.attention import AttentionFunction
from ltx_core.model.transformer.model import LTXModel, LTXModelType
from ltx_core.model.transformer.rope import LTXRopeType
from ltx_core.utils import check_config_value
class LTXModelConfigurator(ModelConfigurator[LTXModel]):
"""
Configurator for LTX model.
Used to create an LTX model from a configuration dictionary.
"""
@classmethod
def from_config(cls: type[LTXModel], config: dict) -> LTXModel:
config = config.get("transformer", {})
check_config_value(config, "dropout", 0.0)
check_config_value(config, "attention_bias", True)
check_config_value(config, "num_vector_embeds", None)
check_config_value(config, "activation_fn", "gelu-approximate")
check_config_value(config, "num_embeds_ada_norm", 1000)
check_config_value(config, "use_linear_projection", False)
check_config_value(config, "only_cross_attention", False)
check_config_value(config, "cross_attention_norm", True)
check_config_value(config, "double_self_attention", False)
check_config_value(config, "upcast_attention", False)
check_config_value(config, "standardization_norm", "rms_norm")
check_config_value(config, "norm_elementwise_affine", False)
check_config_value(config, "qk_norm", "rms_norm")
check_config_value(config, "positional_embedding_type", "rope")
check_config_value(config, "use_audio_video_cross_attention", True)
check_config_value(config, "share_ff", False)
check_config_value(config, "av_cross_ada_norm", True)
check_config_value(config, "use_middle_indices_grid", True)
return LTXModel(
model_type=LTXModelType.AudioVideo,
num_attention_heads=config.get("num_attention_heads", 32),
attention_head_dim=config.get("attention_head_dim", 128),
in_channels=config.get("in_channels", 128),
out_channels=config.get("out_channels", 128),
num_layers=config.get("num_layers", 48),
cross_attention_dim=config.get("cross_attention_dim", 4096),
norm_eps=config.get("norm_eps", 1e-06),
attention_type=AttentionFunction(config.get("attention_type", "default")),
caption_channels=config.get("caption_channels", 3840),
positional_embedding_theta=config.get("positional_embedding_theta", 10000.0),
positional_embedding_max_pos=config.get("positional_embedding_max_pos", [20, 2048, 2048]),
timestep_scale_multiplier=config.get("timestep_scale_multiplier", 1000),
use_middle_indices_grid=config.get("use_middle_indices_grid", True),
audio_num_attention_heads=config.get("audio_num_attention_heads", 32),
audio_attention_head_dim=config.get("audio_attention_head_dim", 64),
audio_in_channels=config.get("audio_in_channels", 128),
audio_out_channels=config.get("audio_out_channels", 128),
audio_cross_attention_dim=config.get("audio_cross_attention_dim", 2048),
audio_positional_embedding_max_pos=config.get("audio_positional_embedding_max_pos", [20]),
av_ca_timestep_scale_multiplier=config.get("av_ca_timestep_scale_multiplier", 1),
rope_type=LTXRopeType(config.get("rope_type", "interleaved")),
double_precision_rope=config.get("frequencies_precision", False) == "float64",
)
class LTXVideoOnlyModelConfigurator(ModelConfigurator[LTXModel]):
"""
Configurator for LTX video only model.
Used to create an LTX video only model from a configuration dictionary.
"""
@classmethod
def from_config(cls: type[LTXModel], config: dict) -> LTXModel:
config = config.get("transformer", {})
check_config_value(config, "dropout", 0.0)
check_config_value(config, "attention_bias", True)
check_config_value(config, "num_vector_embeds", None)
check_config_value(config, "activation_fn", "gelu-approximate")
check_config_value(config, "num_embeds_ada_norm", 1000)
check_config_value(config, "use_linear_projection", False)
check_config_value(config, "only_cross_attention", False)
check_config_value(config, "cross_attention_norm", True)
check_config_value(config, "double_self_attention", False)
check_config_value(config, "upcast_attention", False)
check_config_value(config, "standardization_norm", "rms_norm")
check_config_value(config, "norm_elementwise_affine", False)
check_config_value(config, "qk_norm", "rms_norm")
check_config_value(config, "positional_embedding_type", "rope")
check_config_value(config, "use_middle_indices_grid", True)
return LTXModel(
model_type=LTXModelType.VideoOnly,
num_attention_heads=config.get("num_attention_heads", 32),
attention_head_dim=config.get("attention_head_dim", 128),
in_channels=config.get("in_channels", 128),
out_channels=config.get("out_channels", 128),
num_layers=config.get("num_layers", 48),
cross_attention_dim=config.get("cross_attention_dim", 4096),
norm_eps=config.get("norm_eps", 1e-06),
attention_type=AttentionFunction(config.get("attention_type", "default")),
caption_channels=config.get("caption_channels", 3840),
positional_embedding_theta=config.get("positional_embedding_theta", 10000.0),
positional_embedding_max_pos=config.get("positional_embedding_max_pos", [20, 2048, 2048]),
timestep_scale_multiplier=config.get("timestep_scale_multiplier", 1000),
use_middle_indices_grid=config.get("use_middle_indices_grid", True),
rope_type=LTXRopeType(config.get("rope_type", "interleaved")),
double_precision_rope=config.get("frequencies_precision", False) == "float64",
)
def _naive_weight_or_bias_downcast(key: str, value: torch.Tensor) -> list[KeyValueOperationResult]:
"""
Downcast the weight or bias to the float8_e4m3fn dtype.
"""
return [KeyValueOperationResult(key, value.to(dtype=torch.float8_e4m3fn))]
def _upcast_and_round(
weight: torch.Tensor, dtype: torch.dtype, with_stochastic_rounding: bool = False, seed: int = 0
) -> torch.Tensor:
"""
Upcast the weight to the given dtype and optionally apply stochastic rounding.
Input weight needs to have float8_e4m3fn or float8_e5m2 dtype.
"""
if not with_stochastic_rounding:
return weight.to(dtype)
return fused_add_round_launch(torch.zeros_like(weight, dtype=dtype), weight, seed)
def replace_fwd_with_upcast(layer: torch.nn.Linear, with_stochastic_rounding: bool = False, seed: int = 0) -> None:
"""
Replace linear.forward and rms_norm.forward with a version that:
- upcasts weight and bias to input's dtype
- returns F.linear or F.rms_norm calculated in that dtype
"""
layer.original_forward = layer.forward
def new_linear_forward(*args, **_kwargs) -> torch.Tensor:
# assume first arg is the input tensor
x = args[0]
w_up = _upcast_and_round(layer.weight, x.dtype, with_stochastic_rounding, seed)
b_up = None
if layer.bias is not None:
b_up = _upcast_and_round(layer.bias, x.dtype, with_stochastic_rounding, seed)
return torch.nn.functional.linear(x, w_up, b_up)
layer.forward = new_linear_forward
def amend_forward_with_upcast(
model: torch.nn.Module, with_stochastic_rounding: bool = False, seed: int = 0
) -> torch.nn.Module:
"""
Replace the forward method of the model's Linear and RMSNorm layers to forward
with upcast and optional stochastic rounding.
"""
for m in model.modules():
if isinstance(m, (torch.nn.Linear)):
replace_fwd_with_upcast(m, with_stochastic_rounding, seed)
return model
LTXV_MODEL_COMFY_RENAMING_MAP = (
SDOps("LTXV_MODEL_COMFY_PREFIX_MAP")
.with_matching(prefix="model.diffusion_model.")
.with_replacement("model.diffusion_model.", "")
)
LTXV_MODEL_COMFY_RENAMING_WITH_TRANSFORMER_LINEAR_DOWNCAST_MAP = (
SDOps("LTXV_MODEL_COMFY_PREFIX_MAP")
.with_matching(prefix="model.diffusion_model.")
.with_replacement("model.diffusion_model.", "")
.with_kv_operation(
key_prefix="transformer_blocks.", key_suffix=".to_q.weight", operation=_naive_weight_or_bias_downcast
)
.with_kv_operation(
key_prefix="transformer_blocks.", key_suffix=".to_q.bias", operation=_naive_weight_or_bias_downcast
)
.with_kv_operation(
key_prefix="transformer_blocks.", key_suffix=".to_k.weight", operation=_naive_weight_or_bias_downcast
)
.with_kv_operation(
key_prefix="transformer_blocks.", key_suffix=".to_k.bias", operation=_naive_weight_or_bias_downcast
)
.with_kv_operation(
key_prefix="transformer_blocks.", key_suffix=".to_v.weight", operation=_naive_weight_or_bias_downcast
)
.with_kv_operation(
key_prefix="transformer_blocks.", key_suffix=".to_v.bias", operation=_naive_weight_or_bias_downcast
)
.with_kv_operation(
key_prefix="transformer_blocks.", key_suffix=".to_out.0.weight", operation=_naive_weight_or_bias_downcast
)
.with_kv_operation(
key_prefix="transformer_blocks.", key_suffix=".to_out.0.bias", operation=_naive_weight_or_bias_downcast
)
.with_kv_operation(
key_prefix="transformer_blocks.", key_suffix=".ff.net.0.proj.weight", operation=_naive_weight_or_bias_downcast
)
.with_kv_operation(
key_prefix="transformer_blocks.", key_suffix=".ff.net.0.proj.bias", operation=_naive_weight_or_bias_downcast
)
.with_kv_operation(
key_prefix="transformer_blocks.", key_suffix=".ff.net.2.weight", operation=_naive_weight_or_bias_downcast
)
.with_kv_operation(
key_prefix="transformer_blocks.", key_suffix=".ff.net.2.bias", operation=_naive_weight_or_bias_downcast
)
)
UPCAST_DURING_INFERENCE = ModuleOps(
name="upcast_fp8_during_linear_forward",
matcher=lambda model: isinstance(model, LTXModel),
mutator=lambda model: amend_forward_with_upcast(model, False),
)
class UpcastWithStochasticRounding(ModuleOps):
"""
ModuleOps for upcasting the model's float8_e4m3fn weights and biases to the bfloat16 dtype
and applying stochastic rounding during linear forward.
"""
def __new__(cls, seed: int = 0):
return super().__new__(
cls,
name="upcast_fp8_during_linear_forward_with_stochastic_rounding",
matcher=lambda model: isinstance(model, LTXModel),
mutator=lambda model: amend_forward_with_upcast(model, True, seed),
)
packages/ltx-core/src/ltx_core/model/transformer/rope.py
+204
View File
@@ -0,0 +1,204 @@
import functools
import math
from enum import Enum
from typing import Callable, Tuple
import numpy as np
import torch
from einops import rearrange
class LTXRopeType(Enum):
INTERLEAVED = "interleaved"
SPLIT = "split"
def apply_rotary_emb(
input_tensor: torch.Tensor,
freqs_cis: Tuple[torch.Tensor, torch.Tensor],
rope_type: LTXRopeType = LTXRopeType.INTERLEAVED,
) -> torch.Tensor:
if rope_type == LTXRopeType.INTERLEAVED:
return apply_interleaved_rotary_emb(input_tensor, *freqs_cis)
elif rope_type == LTXRopeType.SPLIT:
return apply_split_rotary_emb(input_tensor, *freqs_cis)
else:
raise ValueError(f"Invalid rope type: {rope_type}")
def apply_interleaved_rotary_emb(
input_tensor: torch.Tensor, cos_freqs: torch.Tensor, sin_freqs: torch.Tensor
) -> torch.Tensor:
t_dup = rearrange(input_tensor, "... (d r) -> ... d r", r=2)
t1, t2 = t_dup.unbind(dim=-1)
t_dup = torch.stack((-t2, t1), dim=-1)
input_tensor_rot = rearrange(t_dup, "... d r -> ... (d r)")
out = input_tensor * cos_freqs + input_tensor_rot * sin_freqs
return out
def apply_split_rotary_emb(
input_tensor: torch.Tensor, cos_freqs: torch.Tensor, sin_freqs: torch.Tensor
) -> torch.Tensor:
needs_reshape = False
if input_tensor.ndim != 4 and cos_freqs.ndim == 4:
b, h, t, _ = cos_freqs.shape
input_tensor = input_tensor.reshape(b, t, h, -1).swapaxes(1, 2)
needs_reshape = True
split_input = rearrange(input_tensor, "... (d r) -> ... d r", d=2)
first_half_input = split_input[..., :1, :]
second_half_input = split_input[..., 1:, :]
output = split_input * cos_freqs.unsqueeze(-2)
first_half_output = output[..., :1, :]
second_half_output = output[..., 1:, :]
first_half_output.addcmul_(-sin_freqs.unsqueeze(-2), second_half_input)
second_half_output.addcmul_(sin_freqs.unsqueeze(-2), first_half_input)
output = rearrange(output, "... d r -> ... (d r)")
if needs_reshape:
output = output.swapaxes(1, 2).reshape(b, t, -1)
return output
@functools.lru_cache(maxsize=5)
def generate_freq_grid_np(
positional_embedding_theta: float, positional_embedding_max_pos_count: int, inner_dim: int
) -> torch.Tensor:
theta = positional_embedding_theta
start = 1
end = theta
n_elem = 2 * positional_embedding_max_pos_count
pow_indices = np.power(
theta,
np.linspace(
np.log(start) / np.log(theta),
np.log(end) / np.log(theta),
inner_dim // n_elem,
dtype=np.float64,
),
)
return torch.tensor(pow_indices * math.pi / 2, dtype=torch.float32)
@functools.lru_cache(maxsize=5)
def generate_freq_grid_pytorch(
positional_embedding_theta: float, positional_embedding_max_pos_count: int, inner_dim: int
) -> torch.Tensor:
theta = positional_embedding_theta
start = 1
end = theta
n_elem = 2 * positional_embedding_max_pos_count
indices = theta ** (
torch.linspace(
math.log(start, theta),
math.log(end, theta),
inner_dim // n_elem,
dtype=torch.float32,
)
)
indices = indices.to(dtype=torch.float32)
indices = indices * math.pi / 2
return indices
def get_fractional_positions(indices_grid: torch.Tensor, max_pos: list[int]) -> torch.Tensor:
n_pos_dims = indices_grid.shape[1]
assert n_pos_dims == len(max_pos), (
f"Number of position dimensions ({n_pos_dims}) must match max_pos length ({len(max_pos)})"
)
fractional_positions = torch.stack(
[indices_grid[:, i] / max_pos[i] for i in range(n_pos_dims)],
dim=-1,
)
return fractional_positions
def generate_freqs(
indices: torch.Tensor, indices_grid: torch.Tensor, max_pos: list[int], use_middle_indices_grid: bool
) -> torch.Tensor:
if use_middle_indices_grid:
assert len(indices_grid.shape) == 4
assert indices_grid.shape[-1] == 2
indices_grid_start, indices_grid_end = indices_grid[..., 0], indices_grid[..., 1]
indices_grid = (indices_grid_start + indices_grid_end) / 2.0
elif len(indices_grid.shape) == 4:
indices_grid = indices_grid[..., 0]
fractional_positions = get_fractional_positions(indices_grid, max_pos)
indices = indices.to(device=fractional_positions.device)
freqs = (indices * (fractional_positions.unsqueeze(-1) * 2 - 1)).transpose(-1, -2).flatten(2)
return freqs
def split_freqs_cis(freqs: torch.Tensor, pad_size: int, num_attention_heads: int) -> tuple[torch.Tensor, torch.Tensor]:
cos_freq = freqs.cos()
sin_freq = freqs.sin()
if pad_size != 0:
cos_padding = torch.ones_like(cos_freq[:, :, :pad_size])
sin_padding = torch.zeros_like(sin_freq[:, :, :pad_size])
cos_freq = torch.concatenate([cos_padding, cos_freq], axis=-1)
sin_freq = torch.concatenate([sin_padding, sin_freq], axis=-1)
# Reshape freqs to be compatible with multi-head attention
b = cos_freq.shape[0]
t = cos_freq.shape[1]
cos_freq = cos_freq.reshape(b, t, num_attention_heads, -1)
sin_freq = sin_freq.reshape(b, t, num_attention_heads, -1)
cos_freq = torch.swapaxes(cos_freq, 1, 2) # (B,H,T,D//2)
sin_freq = torch.swapaxes(sin_freq, 1, 2) # (B,H,T,D//2)
return cos_freq, sin_freq
def interleaved_freqs_cis(freqs: torch.Tensor, pad_size: int) -> tuple[torch.Tensor, torch.Tensor]:
cos_freq = freqs.cos().repeat_interleave(2, dim=-1)
sin_freq = freqs.sin().repeat_interleave(2, dim=-1)
if pad_size != 0:
cos_padding = torch.ones_like(cos_freq[:, :, :pad_size])
sin_padding = torch.zeros_like(cos_freq[:, :, :pad_size])
cos_freq = torch.cat([cos_padding, cos_freq], dim=-1)
sin_freq = torch.cat([sin_padding, sin_freq], dim=-1)
return cos_freq, sin_freq
def precompute_freqs_cis(
indices_grid: torch.Tensor,
dim: int,
out_dtype: torch.dtype,
theta: float = 10000.0,
max_pos: list[int] | None = None,
use_middle_indices_grid: bool = False,
num_attention_heads: int = 32,
rope_type: LTXRopeType = LTXRopeType.INTERLEAVED,
freq_grid_generator: Callable[[float, int, int, torch.device], torch.Tensor] = generate_freq_grid_pytorch,
) -> tuple[torch.Tensor, torch.Tensor]:
if max_pos is None:
max_pos = [20, 2048, 2048]
indices = freq_grid_generator(theta, indices_grid.shape[1], dim)
freqs = generate_freqs(indices, indices_grid, max_pos, use_middle_indices_grid)
if rope_type == LTXRopeType.SPLIT:
expected_freqs = dim // 2
current_freqs = freqs.shape[-1]
pad_size = expected_freqs - current_freqs
cos_freq, sin_freq = split_freqs_cis(freqs, pad_size, num_attention_heads)
else:
# 2 because of cos and sin by 3 for (t, x, y), 1 for temporal only
n_elem = 2 * indices_grid.shape[1]
cos_freq, sin_freq = interleaved_freqs_cis(freqs, dim % n_elem)
return cos_freq.to(out_dtype), sin_freq.to(out_dtype)
packages/ltx-core/src/ltx_core/model/transformer/text_projection.py
+27
View File
@@ -0,0 +1,27 @@
import torch
class PixArtAlphaTextProjection(torch.nn.Module):
"""
Projects caption embeddings. Also handles dropout for classifier-free guidance.
Adapted from https://github.com/PixArt-alpha/PixArt-alpha/blob/master/diffusion/model/nets/PixArt_blocks.py
"""
def __init__(self, in_features: int, hidden_size: int, out_features: int | None = None, act_fn: str = "gelu_tanh"):
super().__init__()
if out_features is None:
out_features = hidden_size
self.linear_1 = torch.nn.Linear(in_features=in_features, out_features=hidden_size, bias=True)
if act_fn == "gelu_tanh":
self.act_1 = torch.nn.GELU(approximate="tanh")
elif act_fn == "silu":
self.act_1 = torch.nn.SiLU()
else:
raise ValueError(f"Unknown activation function: {act_fn}")
self.linear_2 = torch.nn.Linear(in_features=hidden_size, out_features=out_features, bias=True)
def forward(self, caption: torch.Tensor) -> torch.Tensor:
hidden_states = self.linear_1(caption)
hidden_states = self.act_1(hidden_states)
hidden_states = self.linear_2(hidden_states)
return hidden_states
packages/ltx-core/src/ltx_core/model/transformer/timestep_embedding.py
+143
View File
@@ -0,0 +1,143 @@
import math
import torch
def get_timestep_embedding(
timesteps: torch.Tensor,
embedding_dim: int,
flip_sin_to_cos: bool = False,
downscale_freq_shift: float = 1,
scale: float = 1,
max_period: int = 10000,
) -> torch.Tensor:
"""
This matches the implementation in Denoising Diffusion Probabilistic Models: Create sinusoidal timestep embeddings.
Args
timesteps (torch.Tensor):
a 1-D Tensor of N indices, one per batch element. These may be fractional.
embedding_dim (int):
the dimension of the output.
flip_sin_to_cos (bool):
Whether the embedding order should be `cos, sin` (if True) or `sin, cos` (if False)
downscale_freq_shift (float):
Controls the delta between frequencies between dimensions
scale (float):
Scaling factor applied to the embeddings.
max_period (int):
Controls the maximum frequency of the embeddings
Returns
torch.Tensor: an [N x dim] Tensor of positional embeddings.
"""
assert len(timesteps.shape) == 1, "Timesteps should be a 1d-array"
half_dim = embedding_dim // 2
exponent = -math.log(max_period) * torch.arange(start=0, end=half_dim, dtype=torch.float32, device=timesteps.device)
exponent = exponent / (half_dim - downscale_freq_shift)
emb = torch.exp(exponent)
emb = timesteps[:, None].float() * emb[None, :]
# scale embeddings
emb = scale * emb
# concat sine and cosine embeddings
emb = torch.cat([torch.sin(emb), torch.cos(emb)], dim=-1)
# flip sine and cosine embeddings
if flip_sin_to_cos:
emb = torch.cat([emb[:, half_dim:], emb[:, :half_dim]], dim=-1)
# zero pad
if embedding_dim % 2 == 1:
emb = torch.nn.functional.pad(emb, (0, 1, 0, 0))
return emb
class TimestepEmbedding(torch.nn.Module):
def __init__(
self,
in_channels: int,
time_embed_dim: int,
out_dim: int | None = None,
post_act_fn: str | None = None,
cond_proj_dim: int | None = None,
sample_proj_bias: bool = True,
):
super().__init__()
self.linear_1 = torch.nn.Linear(in_channels, time_embed_dim, sample_proj_bias)
if cond_proj_dim is not None:
self.cond_proj = torch.nn.Linear(cond_proj_dim, in_channels, bias=False)
else:
self.cond_proj = None
self.act = torch.nn.SiLU()
time_embed_dim_out = out_dim if out_dim is not None else time_embed_dim
self.linear_2 = torch.nn.Linear(time_embed_dim, time_embed_dim_out, sample_proj_bias)
if post_act_fn is None:
self.post_act = None
def forward(self, sample: torch.Tensor, condition: torch.Tensor | None = None) -> torch.Tensor:
if condition is not None:
sample = sample + self.cond_proj(condition)
sample = self.linear_1(sample)
if self.act is not None:
sample = self.act(sample)
sample = self.linear_2(sample)
if self.post_act is not None:
sample = self.post_act(sample)
return sample
class Timesteps(torch.nn.Module):
def __init__(self, num_channels: int, flip_sin_to_cos: bool, downscale_freq_shift: float, scale: int = 1):
super().__init__()
self.num_channels = num_channels
self.flip_sin_to_cos = flip_sin_to_cos
self.downscale_freq_shift = downscale_freq_shift
self.scale = scale
def forward(self, timesteps: torch.Tensor) -> torch.Tensor:
t_emb = get_timestep_embedding(
timesteps,
self.num_channels,
flip_sin_to_cos=self.flip_sin_to_cos,
downscale_freq_shift=self.downscale_freq_shift,
scale=self.scale,
)
return t_emb
class PixArtAlphaCombinedTimestepSizeEmbeddings(torch.nn.Module):
"""
For PixArt-Alpha.
Reference:
https://github.com/PixArt-alpha/PixArt-alpha/blob/0f55e922376d8b797edd44d25d0e7464b260dcab/diffusion/model/nets/PixArtMS.py#L164C9-L168C29
"""
def __init__(
self,
embedding_dim: int,
size_emb_dim: int,
):
super().__init__()
self.outdim = size_emb_dim
self.time_proj = Timesteps(num_channels=256, flip_sin_to_cos=True, downscale_freq_shift=0)
self.timestep_embedder = TimestepEmbedding(in_channels=256, time_embed_dim=embedding_dim)
def forward(
self,
timestep: torch.Tensor,
hidden_dtype: torch.dtype,
) -> torch.Tensor:
timesteps_proj = self.time_proj(timestep)
timesteps_emb = self.timestep_embedder(timesteps_proj.to(dtype=hidden_dtype)) # (N, D)
return timesteps_emb
packages/ltx-core/src/ltx_core/model/transformer/transformer.py
+274
View File
@@ -0,0 +1,274 @@
from dataclasses import dataclass, replace
import torch
from ltx_core.guidance.perturbations import BatchedPerturbationConfig, PerturbationType
from ltx_core.model.transformer.attention import Attention, AttentionCallable, AttentionFunction
from ltx_core.model.transformer.feed_forward import FeedForward
from ltx_core.model.transformer.rope import LTXRopeType
from ltx_core.model.transformer.transformer_args import TransformerArgs
from ltx_core.utils import rms_norm
@dataclass
class TransformerConfig:
dim: int
heads: int
d_head: int
context_dim: int
class BasicAVTransformerBlock(torch.nn.Module):
def __init__(
self,
idx: int,
video: TransformerConfig | None = None,
audio: TransformerConfig | None = None,
rope_type: LTXRopeType = LTXRopeType.INTERLEAVED,
norm_eps: float = 1e-6,
attention_function: AttentionFunction | AttentionCallable = AttentionFunction.DEFAULT,
):
super().__init__()
self.idx = idx
if video is not None:
self.attn1 = Attention(
query_dim=video.dim,
heads=video.heads,
dim_head=video.d_head,
context_dim=None,
rope_type=rope_type,
norm_eps=norm_eps,
attention_function=attention_function,
)
self.attn2 = Attention(
query_dim=video.dim,
context_dim=video.context_dim,
heads=video.heads,
dim_head=video.d_head,
rope_type=rope_type,
norm_eps=norm_eps,
attention_function=attention_function,
)
self.ff = FeedForward(video.dim, dim_out=video.dim)
self.scale_shift_table = torch.nn.Parameter(torch.empty(6, video.dim))
if audio is not None:
self.audio_attn1 = Attention(
query_dim=audio.dim,
heads=audio.heads,
dim_head=audio.d_head,
context_dim=None,
rope_type=rope_type,
norm_eps=norm_eps,
attention_function=attention_function,
)
self.audio_attn2 = Attention(
query_dim=audio.dim,
context_dim=audio.context_dim,
heads=audio.heads,
dim_head=audio.d_head,
rope_type=rope_type,
norm_eps=norm_eps,
attention_function=attention_function,
)
self.audio_ff = FeedForward(audio.dim, dim_out=audio.dim)
self.audio_scale_shift_table = torch.nn.Parameter(torch.empty(6, audio.dim))
if audio is not None and video is not None:
# Q: Video, K,V: Audio
self.audio_to_video_attn = Attention(
query_dim=video.dim,
context_dim=audio.dim,
heads=audio.heads,
dim_head=audio.d_head,
rope_type=rope_type,
norm_eps=norm_eps,
attention_function=attention_function,
)
# Q: Audio, K,V: Video
self.video_to_audio_attn = Attention(
query_dim=audio.dim,
context_dim=video.dim,
heads=audio.heads,
dim_head=audio.d_head,
rope_type=rope_type,
norm_eps=norm_eps,
attention_function=attention_function,
)
self.scale_shift_table_a2v_ca_audio = torch.nn.Parameter(torch.empty(5, audio.dim))
self.scale_shift_table_a2v_ca_video = torch.nn.Parameter(torch.empty(5, video.dim))
self.norm_eps = norm_eps
def get_ada_values(
self, scale_shift_table: torch.Tensor, batch_size: int, timestep: torch.Tensor, indices: slice
) -> tuple[torch.Tensor, ...]:
num_ada_params = scale_shift_table.shape[0]
ada_values = (
scale_shift_table[indices].unsqueeze(0).unsqueeze(0).to(device=timestep.device, dtype=timestep.dtype)
+ timestep.reshape(batch_size, timestep.shape[1], num_ada_params, -1)[:, :, indices, :]
).unbind(dim=2)
return ada_values
def get_av_ca_ada_values(
self,
scale_shift_table: torch.Tensor,
batch_size: int,
scale_shift_timestep: torch.Tensor,
gate_timestep: torch.Tensor,
num_scale_shift_values: int = 4,
) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
scale_shift_ada_values = self.get_ada_values(
scale_shift_table[:num_scale_shift_values, :], batch_size, scale_shift_timestep, slice(None, None)
)
gate_ada_values = self.get_ada_values(
scale_shift_table[num_scale_shift_values:, :], batch_size, gate_timestep, slice(None, None)
)
scale_shift_chunks = [t.squeeze(2) for t in scale_shift_ada_values]
gate_ada_values = [t.squeeze(2) for t in gate_ada_values]
return (*scale_shift_chunks, *gate_ada_values)
def forward( # noqa: PLR0915
self,
video: TransformerArgs | None,
audio: TransformerArgs | None,
perturbations: BatchedPerturbationConfig | None = None,
) -> tuple[TransformerArgs | None, TransformerArgs | None]:
batch_size = video.x.shape[0]
if perturbations is None:
perturbations = BatchedPerturbationConfig.empty(batch_size)
vx = video.x if video is not None else None
ax = audio.x if audio is not None else None
run_vx = video is not None and video.enabled and vx.numel() > 0
run_ax = audio is not None and audio.enabled and ax.numel() > 0
run_a2v = run_vx and (audio is not None and audio.enabled and ax.numel() > 0)
run_v2a = run_ax and (video is not None and video.enabled and vx.numel() > 0)
if run_vx:
vshift_msa, vscale_msa, vgate_msa = self.get_ada_values(
self.scale_shift_table, vx.shape[0], video.timesteps, slice(0, 3)
)
if not perturbations.all_in_batch(PerturbationType.SKIP_VIDEO_SELF_ATTN, self.idx):
norm_vx = rms_norm(vx, eps=self.norm_eps) * (1 + vscale_msa) + vshift_msa
v_mask = perturbations.mask_like(PerturbationType.SKIP_VIDEO_SELF_ATTN, self.idx, vx)
vx = vx + self.attn1(norm_vx, pe=video.positional_embeddings) * vgate_msa * v_mask
vx = vx + self.attn2(rms_norm(vx, eps=self.norm_eps), context=video.context, mask=video.context_mask)
del vshift_msa, vscale_msa, vgate_msa
if run_ax:
ashift_msa, ascale_msa, agate_msa = self.get_ada_values(
self.audio_scale_shift_table, ax.shape[0], audio.timesteps, slice(0, 3)
)
if not perturbations.all_in_batch(PerturbationType.SKIP_AUDIO_SELF_ATTN, self.idx):
norm_ax = rms_norm(ax, eps=self.norm_eps) * (1 + ascale_msa) + ashift_msa
a_mask = perturbations.mask_like(PerturbationType.SKIP_AUDIO_SELF_ATTN, self.idx, ax)
ax = ax + self.audio_attn1(norm_ax, pe=audio.positional_embeddings) * agate_msa * a_mask
ax = ax + self.audio_attn2(rms_norm(ax, eps=self.norm_eps), context=audio.context, mask=audio.context_mask)
del ashift_msa, ascale_msa, agate_msa
# Audio - Video cross attention.
if run_a2v or run_v2a:
vx_norm3 = rms_norm(vx, eps=self.norm_eps)
ax_norm3 = rms_norm(ax, eps=self.norm_eps)
(
scale_ca_audio_hidden_states_a2v,
shift_ca_audio_hidden_states_a2v,
scale_ca_audio_hidden_states_v2a,
shift_ca_audio_hidden_states_v2a,
gate_out_v2a,
) = self.get_av_ca_ada_values(
self.scale_shift_table_a2v_ca_audio,
ax.shape[0],
audio.cross_scale_shift_timestep,
audio.cross_gate_timestep,
)
(
scale_ca_video_hidden_states_a2v,
shift_ca_video_hidden_states_a2v,
scale_ca_video_hidden_states_v2a,
shift_ca_video_hidden_states_v2a,
gate_out_a2v,
) = self.get_av_ca_ada_values(
self.scale_shift_table_a2v_ca_video,
vx.shape[0],
video.cross_scale_shift_timestep,
video.cross_gate_timestep,
)
if run_a2v:
vx_scaled = vx_norm3 * (1 + scale_ca_video_hidden_states_a2v) + shift_ca_video_hidden_states_a2v
ax_scaled = ax_norm3 * (1 + scale_ca_audio_hidden_states_a2v) + shift_ca_audio_hidden_states_a2v
a2v_mask = perturbations.mask_like(PerturbationType.SKIP_A2V_CROSS_ATTN, self.idx, vx)
vx = vx + (
self.audio_to_video_attn(
vx_scaled,
context=ax_scaled,
pe=video.cross_positional_embeddings,
k_pe=audio.cross_positional_embeddings,
)
* gate_out_a2v
* a2v_mask
)
if run_v2a:
ax_scaled = ax_norm3 * (1 + scale_ca_audio_hidden_states_v2a) + shift_ca_audio_hidden_states_v2a
vx_scaled = vx_norm3 * (1 + scale_ca_video_hidden_states_v2a) + shift_ca_video_hidden_states_v2a
v2a_mask = perturbations.mask_like(PerturbationType.SKIP_V2A_CROSS_ATTN, self.idx, ax)
ax = ax + (
self.video_to_audio_attn(
ax_scaled,
context=vx_scaled,
pe=audio.cross_positional_embeddings,
k_pe=video.cross_positional_embeddings,
)
* gate_out_v2a
* v2a_mask
)
del gate_out_a2v, gate_out_v2a
del (
scale_ca_video_hidden_states_a2v,
shift_ca_video_hidden_states_a2v,
scale_ca_audio_hidden_states_a2v,
shift_ca_audio_hidden_states_a2v,
scale_ca_video_hidden_states_v2a,
shift_ca_video_hidden_states_v2a,
scale_ca_audio_hidden_states_v2a,
shift_ca_audio_hidden_states_v2a,
)
if run_vx:
vshift_mlp, vscale_mlp, vgate_mlp = self.get_ada_values(
self.scale_shift_table, vx.shape[0], video.timesteps, slice(3, None)
)
vx_scaled = rms_norm(vx, eps=self.norm_eps) * (1 + vscale_mlp) + vshift_mlp
vx = vx + self.ff(vx_scaled) * vgate_mlp
del vshift_mlp, vscale_mlp, vgate_mlp
if run_ax:
ashift_mlp, ascale_mlp, agate_mlp = self.get_ada_values(
self.audio_scale_shift_table, ax.shape[0], audio.timesteps, slice(3, None)
)
ax_scaled = rms_norm(ax, eps=self.norm_eps) * (1 + ascale_mlp) + ashift_mlp
ax = ax + self.audio_ff(ax_scaled) * agate_mlp
del ashift_mlp, ascale_mlp, agate_mlp
return replace(video, x=vx) if video is not None else None, replace(audio, x=ax) if audio is not None else None
packages/ltx-core/src/ltx_core/model/transformer/transformer_args.py
+239
View File
@@ -0,0 +1,239 @@
from dataclasses import dataclass, replace
import torch
from ltx_core.model.transformer.adaln import AdaLayerNormSingle
from ltx_core.model.transformer.modality import Modality
from ltx_core.model.transformer.rope import (
LTXRopeType,
generate_freq_grid_np,
generate_freq_grid_pytorch,
precompute_freqs_cis,
)
from ltx_core.model.transformer.text_projection import PixArtAlphaTextProjection
@dataclass(frozen=True)
class TransformerArgs:
x: torch.Tensor
context: torch.Tensor
context_mask: torch.Tensor
timesteps: torch.Tensor
embedded_timestep: torch.Tensor
positional_embeddings: torch.Tensor
cross_positional_embeddings: torch.Tensor | None
cross_scale_shift_timestep: torch.Tensor | None
cross_gate_timestep: torch.Tensor | None
enabled: bool
class TransformerArgsPreprocessor:
def __init__( # noqa: PLR0913
self,
patchify_proj: torch.nn.Linear,
adaln: AdaLayerNormSingle,
caption_projection: PixArtAlphaTextProjection,
inner_dim: int,
max_pos: list[int],
num_attention_heads: int,
use_middle_indices_grid: bool,
timestep_scale_multiplier: int,
double_precision_rope: bool,
positional_embedding_theta: float,
rope_type: LTXRopeType,
) -> None:
self.patchify_proj = patchify_proj
self.adaln = adaln
self.caption_projection = caption_projection
self.inner_dim = inner_dim
self.max_pos = max_pos
self.num_attention_heads = num_attention_heads
self.use_middle_indices_grid = use_middle_indices_grid
self.timestep_scale_multiplier = timestep_scale_multiplier
self.double_precision_rope = double_precision_rope
self.positional_embedding_theta = positional_embedding_theta
self.rope_type = rope_type
def _prepare_timestep(
self, timestep: torch.Tensor, batch_size: int, hidden_dtype: torch.dtype
) -> tuple[torch.Tensor, torch.Tensor]:
"""Prepare timestep embeddings."""
timestep = timestep * self.timestep_scale_multiplier
timestep, embedded_timestep = self.adaln(
timestep.flatten(),
hidden_dtype=hidden_dtype,
)
# Second dimension is 1 or number of tokens (if timestep_per_token)
timestep = timestep.view(batch_size, -1, timestep.shape[-1])
embedded_timestep = embedded_timestep.view(batch_size, -1, embedded_timestep.shape[-1])
return timestep, embedded_timestep
def _prepare_context(
self,
context: torch.Tensor,
x: torch.Tensor,
attention_mask: torch.Tensor | None = None,
) -> tuple[torch.Tensor, torch.Tensor | None]:
"""Prepare context for transformer blocks."""
batch_size = x.shape[0]
context = self.caption_projection(context)
context = context.view(batch_size, -1, x.shape[-1])
return context, attention_mask
def _prepare_attention_mask(self, attention_mask: torch.Tensor | None, x_dtype: torch.dtype) -> torch.Tensor | None:
"""Prepare attention mask."""
if attention_mask is None or torch.is_floating_point(attention_mask):
return attention_mask
return (attention_mask - 1).to(x_dtype).reshape(
(attention_mask.shape[0], 1, -1, attention_mask.shape[-1])
) * torch.finfo(x_dtype).max
def _prepare_positional_embeddings(
self,
positions: torch.Tensor,
inner_dim: int,
max_pos: list[int],
use_middle_indices_grid: bool,
num_attention_heads: int,
x_dtype: torch.dtype,
) -> torch.Tensor:
"""Prepare positional embeddings."""
freq_grid_generator = generate_freq_grid_np if self.double_precision_rope else generate_freq_grid_pytorch
pe = precompute_freqs_cis(
positions,
dim=inner_dim,
out_dtype=x_dtype,
theta=self.positional_embedding_theta,
max_pos=max_pos,
use_middle_indices_grid=use_middle_indices_grid,
num_attention_heads=num_attention_heads,
rope_type=self.rope_type,
freq_grid_generator=freq_grid_generator,
)
return pe
def prepare(
self,
modality: Modality,
) -> TransformerArgs:
x = self.patchify_proj(modality.latent)
timestep, embedded_timestep = self._prepare_timestep(modality.timesteps, x.shape[0], modality.latent.dtype)
context, attention_mask = self._prepare_context(modality.context, x, modality.context_mask)
attention_mask = self._prepare_attention_mask(attention_mask, modality.latent.dtype)
pe = self._prepare_positional_embeddings(
positions=modality.positions,
inner_dim=self.inner_dim,
max_pos=self.max_pos,
use_middle_indices_grid=self.use_middle_indices_grid,
num_attention_heads=self.num_attention_heads,
x_dtype=modality.latent.dtype,
)
return TransformerArgs(
x=x,
context=context,
context_mask=attention_mask,
timesteps=timestep,
embedded_timestep=embedded_timestep,
positional_embeddings=pe,
cross_positional_embeddings=None,
cross_scale_shift_timestep=None,
cross_gate_timestep=None,
enabled=modality.enabled,
)
class MultiModalTransformerArgsPreprocessor:
def __init__( # noqa: PLR0913
self,
patchify_proj: torch.nn.Linear,
adaln: AdaLayerNormSingle,
caption_projection: PixArtAlphaTextProjection,
cross_scale_shift_adaln: AdaLayerNormSingle,
cross_gate_adaln: AdaLayerNormSingle,
inner_dim: int,
max_pos: list[int],
num_attention_heads: int,
cross_pe_max_pos: int,
use_middle_indices_grid: bool,
audio_cross_attention_dim: int,
timestep_scale_multiplier: int,
double_precision_rope: bool,
positional_embedding_theta: float,
rope_type: LTXRopeType,
av_ca_timestep_scale_multiplier: int,
) -> None:
self.simple_preprocessor = TransformerArgsPreprocessor(
patchify_proj=patchify_proj,
adaln=adaln,
caption_projection=caption_projection,
inner_dim=inner_dim,
max_pos=max_pos,
num_attention_heads=num_attention_heads,
use_middle_indices_grid=use_middle_indices_grid,
timestep_scale_multiplier=timestep_scale_multiplier,
double_precision_rope=double_precision_rope,
positional_embedding_theta=positional_embedding_theta,
rope_type=rope_type,
)
self.cross_scale_shift_adaln = cross_scale_shift_adaln
self.cross_gate_adaln = cross_gate_adaln
self.cross_pe_max_pos = cross_pe_max_pos
self.audio_cross_attention_dim = audio_cross_attention_dim
self.av_ca_timestep_scale_multiplier = av_ca_timestep_scale_multiplier
def prepare(
self,
modality: Modality,
) -> TransformerArgs:
transformer_args = self.simple_preprocessor.prepare(modality)
cross_pe = self.simple_preprocessor._prepare_positional_embeddings(
positions=modality.positions[:, 0:1, :],
inner_dim=self.audio_cross_attention_dim,
max_pos=[self.cross_pe_max_pos],
use_middle_indices_grid=True,
num_attention_heads=self.simple_preprocessor.num_attention_heads,
x_dtype=modality.latent.dtype,
)
cross_scale_shift_timestep, cross_gate_timestep = self._prepare_cross_attention_timestep(
timestep=modality.timesteps,
timestep_scale_multiplier=self.simple_preprocessor.timestep_scale_multiplier,
batch_size=transformer_args.x.shape[0],
hidden_dtype=modality.latent.dtype,
)
return replace(
transformer_args,
cross_positional_embeddings=cross_pe,
cross_scale_shift_timestep=cross_scale_shift_timestep,
cross_gate_timestep=cross_gate_timestep,
)
def _prepare_cross_attention_timestep(
self,
timestep: torch.Tensor,
timestep_scale_multiplier: int,
batch_size: int,
hidden_dtype: torch.dtype,
) -> tuple[torch.Tensor, torch.Tensor]:
"""Prepare cross attention timestep embeddings."""
timestep = timestep * timestep_scale_multiplier
av_ca_factor = self.av_ca_timestep_scale_multiplier / timestep_scale_multiplier
scale_shift_timestep, _ = self.cross_scale_shift_adaln(
timestep.flatten(),
hidden_dtype=hidden_dtype,
)
scale_shift_timestep = scale_shift_timestep.view(batch_size, -1, scale_shift_timestep.shape[-1])
gate_noise_timestep, _ = self.cross_gate_adaln(
timestep.flatten() * av_ca_factor,
hidden_dtype=hidden_dtype,
)
gate_noise_timestep = gate_noise_timestep.view(batch_size, -1, gate_noise_timestep.shape[-1])
return scale_shift_timestep, gate_noise_timestep
packages/ltx-core/src/ltx_core/model/upsampler/__init__.py
+10
View File
@@ -0,0 +1,10 @@
"""Latent upsampler model components."""
from ltx_core.model.upsampler.model import LatentUpsampler, upsample_video
from ltx_core.model.upsampler.model_configurator import LatentUpsamplerConfigurator
__all__ = [
"LatentUpsampler",
"LatentUpsamplerConfigurator",
"upsample_video",
]
packages/ltx-core/src/ltx_core/model/upsampler/blur_downsample.py
+53
View File
@@ -0,0 +1,53 @@
import math
import torch
import torch.nn.functional as F
from einops import rearrange
class BlurDownsample(torch.nn.Module):
"""
Anti-aliased spatial downsampling by integer stride using a fixed separable binomial kernel.
Applies only on H,W. Works for dims=2 or dims=3 (per-frame).
"""
def __init__(self, dims: int, stride: int, kernel_size: int = 5) -> None:
super().__init__()
assert dims in (2, 3)
assert isinstance(stride, int)
assert stride >= 1
assert kernel_size >= 3
assert kernel_size % 2 == 1
self.dims = dims
self.stride = stride
self.kernel_size = kernel_size
# 5x5 separable binomial kernel using binomial coefficients [1, 4, 6, 4, 1] from
# the 4th row of Pascal's triangle. This kernel is used for anti-aliasing and
# provides a smooth approximation of a Gaussian filter (often called a "binomial filter").
# The 2D kernel is constructed as the outer product and normalized.
k = torch.tensor([math.comb(kernel_size - 1, k) for k in range(kernel_size)])
k2d = k[:, None] @ k[None, :]
k2d = (k2d / k2d.sum()).float() # shape (kernel_size, kernel_size)
self.register_buffer("kernel", k2d[None, None, :, :]) # (1, 1, kernel_size, kernel_size)
def forward(self, x: torch.Tensor) -> torch.Tensor:
if self.stride == 1:
return x
if self.dims == 2:
return self._apply_2d(x)
else:
# dims == 3: apply per-frame on H,W
b, _, f, _, _ = x.shape
x = rearrange(x, "b c f h w -> (b f) c h w")
x = self._apply_2d(x)
h2, w2 = x.shape[-2:]
x = rearrange(x, "(b f) c h w -> b c f h w", b=b, f=f, h=h2, w=w2)
return x
def _apply_2d(self, x2d: torch.Tensor) -> torch.Tensor:
c = x2d.shape[1]
weight = self.kernel.expand(c, 1, self.kernel_size, self.kernel_size) # depthwise
x2d = F.conv2d(x2d, weight=weight, bias=None, stride=self.stride, padding=self.kernel_size // 2, groups=c)
return x2d
packages/ltx-core/src/ltx_core/model/upsampler/model.py
+142
View File
@@ -0,0 +1,142 @@
import torch
from einops import rearrange
from ltx_core.model.upsampler.pixel_shuffle import PixelShuffleND
from ltx_core.model.upsampler.res_block import ResBlock
from ltx_core.model.upsampler.spatial_rational_resampler import SpatialRationalResampler
from ltx_core.model.video_vae import VideoEncoder
class LatentUpsampler(torch.nn.Module):
"""
Model to upsample VAE latents spatially and/or temporally.
Args:
in_channels (`int`): Number of channels in the input latent
mid_channels (`int`): Number of channels in the middle layers
num_blocks_per_stage (`int`): Number of ResBlocks to use in each stage (pre/post upsampling)
dims (`int`): Number of dimensions for convolutions (2 or 3)
spatial_upsample (`bool`): Whether to spatially upsample the latent
temporal_upsample (`bool`): Whether to temporally upsample the latent
spatial_scale (`float`): Scale factor for spatial upsampling
rational_resampler (`bool`): Whether to use a rational resampler for spatial upsampling
"""
def __init__(
self,
in_channels: int = 128,
mid_channels: int = 512,
num_blocks_per_stage: int = 4,
dims: int = 3,
spatial_upsample: bool = True,
temporal_upsample: bool = False,
spatial_scale: float = 2.0,
rational_resampler: bool = False,
):
super().__init__()
self.in_channels = in_channels
self.mid_channels = mid_channels
self.num_blocks_per_stage = num_blocks_per_stage
self.dims = dims
self.spatial_upsample = spatial_upsample
self.temporal_upsample = temporal_upsample
self.spatial_scale = float(spatial_scale)
self.rational_resampler = rational_resampler
conv = torch.nn.Conv2d if dims == 2 else torch.nn.Conv3d
self.initial_conv = conv(in_channels, mid_channels, kernel_size=3, padding=1)
self.initial_norm = torch.nn.GroupNorm(32, mid_channels)
self.initial_activation = torch.nn.SiLU()
self.res_blocks = torch.nn.ModuleList([ResBlock(mid_channels, dims=dims) for _ in range(num_blocks_per_stage)])
if spatial_upsample and temporal_upsample:
self.upsampler = torch.nn.Sequential(
torch.nn.Conv3d(mid_channels, 8 * mid_channels, kernel_size=3, padding=1),
PixelShuffleND(3),
)
elif spatial_upsample:
if rational_resampler:
self.upsampler = SpatialRationalResampler(mid_channels=mid_channels, scale=self.spatial_scale)
else:
self.upsampler = torch.nn.Sequential(
torch.nn.Conv2d(mid_channels, 4 * mid_channels, kernel_size=3, padding=1),
PixelShuffleND(2),
)
elif temporal_upsample:
self.upsampler = torch.nn.Sequential(
torch.nn.Conv3d(mid_channels, 2 * mid_channels, kernel_size=3, padding=1),
PixelShuffleND(1),
)
else:
raise ValueError("Either spatial_upsample or temporal_upsample must be True")
self.post_upsample_res_blocks = torch.nn.ModuleList(
[ResBlock(mid_channels, dims=dims) for _ in range(num_blocks_per_stage)]
)
self.final_conv = conv(mid_channels, in_channels, kernel_size=3, padding=1)
def forward(self, latent: torch.Tensor) -> torch.Tensor:
b, _, f, _, _ = latent.shape
if self.dims == 2:
x = rearrange(latent, "b c f h w -> (b f) c h w")
x = self.initial_conv(x)
x = self.initial_norm(x)
x = self.initial_activation(x)
for block in self.res_blocks:
x = block(x)
x = self.upsampler(x)
for block in self.post_upsample_res_blocks:
x = block(x)
x = self.final_conv(x)
x = rearrange(x, "(b f) c h w -> b c f h w", b=b, f=f)
else:
x = self.initial_conv(latent)
x = self.initial_norm(x)
x = self.initial_activation(x)
for block in self.res_blocks:
x = block(x)
if self.temporal_upsample:
x = self.upsampler(x)
# remove the first frame after upsampling.
# This is done because the first frame encodes one pixel frame.
x = x[:, :, 1:, :, :]
elif isinstance(self.upsampler, SpatialRationalResampler):
x = self.upsampler(x)
else:
x = rearrange(x, "b c f h w -> (b f) c h w")
x = self.upsampler(x)
x = rearrange(x, "(b f) c h w -> b c f h w", b=b, f=f)
for block in self.post_upsample_res_blocks:
x = block(x)
x = self.final_conv(x)
return x
def upsample_video(latent: torch.Tensor, video_encoder: VideoEncoder, upsampler: "LatentUpsampler") -> torch.Tensor:
"""
Apply upsampling to the latent representation using the provided upsampler,
with normalization and un-normalization based on the video encoder's per-channel statistics.
Args:
latent: Input latent tensor of shape [B, C, F, H, W].
video_encoder: VideoEncoder with per_channel_statistics for normalization.
upsampler: LatentUpsampler module to perform upsampling.
Returns:
torch.Tensor: Upsampled and re-normalized latent tensor.
"""
latent = video_encoder.per_channel_statistics.un_normalize(latent)
latent = upsampler(latent)
latent = video_encoder.per_channel_statistics.normalize(latent)
return latent
packages/ltx-core/src/ltx_core/model/upsampler/model_configurator.py
+30
View File
@@ -0,0 +1,30 @@
from ltx_core.model.model_protocol import ModelConfigurator
from ltx_core.model.upsampler.model import LatentUpsampler
class LatentUpsamplerConfigurator(ModelConfigurator[LatentUpsampler]):
"""
Configurator for LatentUpsampler model.
Used to create a LatentUpsampler model from a configuration dictionary.
"""
@classmethod
def from_config(cls: type[LatentUpsampler], config: dict) -> LatentUpsampler:
in_channels = config.get("in_channels", 128)
mid_channels = config.get("mid_channels", 512)
num_blocks_per_stage = config.get("num_blocks_per_stage", 4)
dims = config.get("dims", 3)
spatial_upsample = config.get("spatial_upsample", True)
temporal_upsample = config.get("temporal_upsample", False)
spatial_scale = config.get("spatial_scale", 2.0)
rational_resampler = config.get("rational_resampler", False)
return LatentUpsampler(
in_channels=in_channels,
mid_channels=mid_channels,
num_blocks_per_stage=num_blocks_per_stage,
dims=dims,
spatial_upsample=spatial_upsample,
temporal_upsample=temporal_upsample,
spatial_scale=spatial_scale,
rational_resampler=rational_resampler,
)
packages/ltx-core/src/ltx_core/model/upsampler/pixel_shuffle.py
+54
View File
@@ -0,0 +1,54 @@
import torch
from einops import rearrange
class PixelShuffleND(torch.nn.Module):
"""
N-dimensional pixel shuffle operation for upsampling tensors.
Args:
dims (int): Number of dimensions to apply pixel shuffle to.
- 1: Temporal (e.g., frames)
- 2: Spatial (e.g., height and width)
- 3: Spatiotemporal (e.g., depth, height, width)
upscale_factors (tuple[int, int, int], optional): Upscaling factors for each dimension.
For dims=1, only the first value is used.
For dims=2, the first two values are used.
For dims=3, all three values are used.
The input tensor is rearranged so that the channel dimension is split into
smaller channels and upscaling factors, and the upscaling factors are moved
into the corresponding spatial/temporal dimensions.
Note:
This operation is equivalent to the patchifier operation in for the models. Consider
using this class instead.
"""
def __init__(self, dims: int, upscale_factors: tuple[int, int, int] = (2, 2, 2)):
super().__init__()
assert dims in [1, 2, 3], "dims must be 1, 2, or 3"
self.dims = dims
self.upscale_factors = upscale_factors
def forward(self, x: torch.Tensor) -> torch.Tensor:
if self.dims == 3:
return rearrange(
x,
"b (c p1 p2 p3) d h w -> b c (d p1) (h p2) (w p3)",
p1=self.upscale_factors[0],
p2=self.upscale_factors[1],
p3=self.upscale_factors[2],
)
elif self.dims == 2:
return rearrange(
x,
"b (c p1 p2) h w -> b c (h p1) (w p2)",
p1=self.upscale_factors[0],
p2=self.upscale_factors[1],
)
elif self.dims == 1:
return rearrange(
x,
"b (c p1) f h w -> b c (f p1) h w",
p1=self.upscale_factors[0],
)
else:
raise ValueError(f"Unsupported dims: {self.dims}")
packages/ltx-core/src/ltx_core/model/upsampler/res_block.py
+37
View File
@@ -0,0 +1,37 @@
from typing import Optional
import torch
class ResBlock(torch.nn.Module):
"""
Residual block with two convolutional layers, group normalization, and SiLU activation.
Args:
channels (int): Number of input and output channels.
mid_channels (Optional[int]): Number of channels in the intermediate convolution layer. Defaults to `channels`
if not specified.
dims (int): Dimensionality of the convolution (2 for Conv2d, 3 for Conv3d). Defaults to 3.
"""
def __init__(self, channels: int, mid_channels: Optional[int] = None, dims: int = 3):
super().__init__()
if mid_channels is None:
mid_channels = channels
conv = torch.nn.Conv2d if dims == 2 else torch.nn.Conv3d
self.conv1 = conv(channels, mid_channels, kernel_size=3, padding=1)
self.norm1 = torch.nn.GroupNorm(32, mid_channels)
self.conv2 = conv(mid_channels, channels, kernel_size=3, padding=1)
self.norm2 = torch.nn.GroupNorm(32, channels)
self.activation = torch.nn.SiLU()
def forward(self, x: torch.Tensor) -> torch.Tensor:
residual = x
x = self.conv1(x)
x = self.norm1(x)
x = self.activation(x)
x = self.conv2(x)
x = self.norm2(x)
x = self.activation(x + residual)
return x
packages/ltx-core/src/ltx_core/model/upsampler/spatial_rational_resampler.py
+47
View File
@@ -0,0 +1,47 @@
from typing import Tuple
import torch
from einops import rearrange
from ltx_core.model.upsampler.blur_downsample import BlurDownsample
from ltx_core.model.upsampler.pixel_shuffle import PixelShuffleND
def _rational_for_scale(scale: float) -> Tuple[int, int]:
mapping = {0.75: (3, 4), 1.5: (3, 2), 2.0: (2, 1), 4.0: (4, 1)}
if float(scale) not in mapping:
raise ValueError(f"Unsupported scale {scale}. Choose from {list(mapping.keys())}")
return mapping[float(scale)]
class SpatialRationalResampler(torch.nn.Module):
"""
Fully-learned rational spatial scaling: up by 'num' via PixelShuffle, then anti-aliased
downsample by 'den' using fixed blur + stride. Operates on H,W only.
For dims==3, work per-frame for spatial scaling (temporal axis untouched).
Args:
mid_channels (`int`): Number of intermediate channels for the convolution layer
scale (`float`): Spatial scaling factor. Supported values are:
- 0.75: Downsample by 3/4 (reduce spatial size)
- 1.5: Upsample by 3/2 (increase spatial size)
- 2.0: Upsample by 2x (double spatial size)
- 4.0: Upsample by 4x (quadruple spatial size)
Any other value will raise a ValueError.
"""
def __init__(self, mid_channels: int, scale: float):
super().__init__()
self.scale = float(scale)
self.num, self.den = _rational_for_scale(self.scale)
self.conv = torch.nn.Conv2d(mid_channels, (self.num**2) * mid_channels, kernel_size=3, padding=1)
self.pixel_shuffle = PixelShuffleND(2, upscale_factors=(self.num, self.num))
self.blur_down = BlurDownsample(dims=2, stride=self.den)
def forward(self, x: torch.Tensor) -> torch.Tensor:
b, _, f, _, _ = x.shape
x = rearrange(x, "b c f h w -> (b f) c h w")
x = self.conv(x)
x = self.pixel_shuffle(x)
x = self.blur_down(x)
x = rearrange(x, "(b f) c h w -> b c f h w", b=b, f=f)
return x
packages/ltx-core/src/ltx_core/model/video_vae/__init__.py
+24
View File
@@ -0,0 +1,24 @@
"""Video VAE package."""
from ltx_core.model.video_vae.model_configurator import (
VAE_DECODER_COMFY_KEYS_FILTER,
VAE_ENCODER_COMFY_KEYS_FILTER,
VideoDecoderConfigurator,
VideoEncoderConfigurator,
)
from ltx_core.model.video_vae.tiling import SpatialTilingConfig, TemporalTilingConfig, TilingConfig
from ltx_core.model.video_vae.video_vae import VideoDecoder, VideoEncoder, decode_video, get_video_chunks_number
__all__ = [
"VAE_DECODER_COMFY_KEYS_FILTER",
"VAE_ENCODER_COMFY_KEYS_FILTER",
"SpatialTilingConfig",
"TemporalTilingConfig",
"TilingConfig",
"VideoDecoder",
"VideoDecoderConfigurator",
"VideoEncoder",
"VideoEncoderConfigurator",
"decode_video",
"get_video_chunks_number",
]
packages/ltx-core/src/ltx_core/model/video_vae/convolution.py
+317
View File
@@ -0,0 +1,317 @@
from typing import Tuple, Union
import torch
from einops import rearrange
from torch import nn
from torch.nn import functional as F
from ltx_core.model.video_vae.enums import PaddingModeType
def make_conv_nd( # noqa: PLR0913
dims: Union[int, Tuple[int, int]],
in_channels: int,
out_channels: int,
kernel_size: int,
stride: int = 1,
padding: int = 0,
dilation: int = 1,
groups: int = 1,
bias: bool = True,
causal: bool = False,
spatial_padding_mode: PaddingModeType = PaddingModeType.ZEROS,
temporal_padding_mode: PaddingModeType = PaddingModeType.ZEROS,
) -> nn.Module:
if not (spatial_padding_mode == temporal_padding_mode or causal):
raise NotImplementedError("spatial and temporal padding modes must be equal")
if dims == 2:
return nn.Conv2d(
in_channels=in_channels,
out_channels=out_channels,
kernel_size=kernel_size,
stride=stride,
padding=padding,
dilation=dilation,
groups=groups,
bias=bias,
padding_mode=spatial_padding_mode.value,
)
elif dims == 3:
if causal:
return CausalConv3d(
in_channels=in_channels,
out_channels=out_channels,
kernel_size=kernel_size,
stride=stride,
dilation=dilation,
groups=groups,
bias=bias,
spatial_padding_mode=spatial_padding_mode,
)
return nn.Conv3d(
in_channels=in_channels,
out_channels=out_channels,
kernel_size=kernel_size,
stride=stride,
padding=padding,
dilation=dilation,
groups=groups,
bias=bias,
padding_mode=spatial_padding_mode.value,
)
elif dims == (2, 1):
return DualConv3d(
in_channels=in_channels,
out_channels=out_channels,
kernel_size=kernel_size,
stride=stride,
padding=padding,
bias=bias,
padding_mode=spatial_padding_mode.value,
)
else:
raise ValueError(f"unsupported dimensions: {dims}")
def make_linear_nd(
dims: int,
in_channels: int,
out_channels: int,
bias: bool = True,
) -> nn.Module:
if dims == 2:
return nn.Conv2d(in_channels=in_channels, out_channels=out_channels, kernel_size=1, bias=bias)
elif dims in (3, (2, 1)):
return nn.Conv3d(in_channels=in_channels, out_channels=out_channels, kernel_size=1, bias=bias)
else:
raise ValueError(f"unsupported dimensions: {dims}")
class DualConv3d(nn.Module):
def __init__(
self,
in_channels: int,
out_channels: int,
kernel_size: int,
stride: Union[int, Tuple[int, int, int]] = 1,
padding: Union[int, Tuple[int, int, int]] = 0,
dilation: Union[int, Tuple[int, int, int]] = 1,
groups: int = 1,
bias: bool = True,
padding_mode: str = "zeros",
) -> None:
super(DualConv3d, self).__init__()
self.in_channels = in_channels
self.out_channels = out_channels
self.padding_mode = padding_mode
# Ensure kernel_size, stride, padding, and dilation are tuples of length 3
if isinstance(kernel_size, int):
kernel_size = (kernel_size, kernel_size, kernel_size)
if kernel_size == (1, 1, 1):
raise ValueError("kernel_size must be greater than 1. Use make_linear_nd instead.")
if isinstance(stride, int):
stride = (stride, stride, stride)
if isinstance(padding, int):
padding = (padding, padding, padding)
if isinstance(dilation, int):
dilation = (dilation, dilation, dilation)
# Set parameters for convolutions
self.groups = groups
self.bias = bias
# Define the size of the channels after the first convolution
intermediate_channels = out_channels if in_channels < out_channels else in_channels
# Define parameters for the first convolution
self.weight1 = nn.Parameter(
torch.Tensor(
intermediate_channels,
in_channels // groups,
1,
kernel_size[1],
kernel_size[2],
)
)
self.stride1 = (1, stride[1], stride[2])
self.padding1 = (0, padding[1], padding[2])
self.dilation1 = (1, dilation[1], dilation[2])
if bias:
self.bias1 = nn.Parameter(torch.Tensor(intermediate_channels))
else:
self.register_parameter("bias1", None)
# Define parameters for the second convolution
self.weight2 = nn.Parameter(torch.Tensor(out_channels, intermediate_channels // groups, kernel_size[0], 1, 1))
self.stride2 = (stride[0], 1, 1)
self.padding2 = (padding[0], 0, 0)
self.dilation2 = (dilation[0], 1, 1)
if bias:
self.bias2 = nn.Parameter(torch.Tensor(out_channels))
else:
self.register_parameter("bias2", None)
# Initialize weights and biases
self.reset_parameters()
def reset_parameters(self) -> None:
nn.init.kaiming_uniform_(self.weight1, a=torch.sqrt(5))
nn.init.kaiming_uniform_(self.weight2, a=torch.sqrt(5))
if self.bias:
fan_in1, _ = nn.init._calculate_fan_in_and_fan_out(self.weight1)
bound1 = 1 / torch.sqrt(fan_in1)
nn.init.uniform_(self.bias1, -bound1, bound1)
fan_in2, _ = nn.init._calculate_fan_in_and_fan_out(self.weight2)
bound2 = 1 / torch.sqrt(fan_in2)
nn.init.uniform_(self.bias2, -bound2, bound2)
def forward(
self,
x: torch.Tensor,
use_conv3d: bool = False,
skip_time_conv: bool = False,
) -> torch.Tensor:
if use_conv3d:
return self.forward_with_3d(x=x, skip_time_conv=skip_time_conv)
else:
return self.forward_with_2d(x=x, skip_time_conv=skip_time_conv)
def forward_with_3d(self, x: torch.Tensor, skip_time_conv: bool = False) -> torch.Tensor:
# First convolution
x = F.conv3d(
x,
self.weight1,
self.bias1,
self.stride1,
self.padding1,
self.dilation1,
self.groups,
padding_mode=self.padding_mode,
)
if skip_time_conv:
return x
# Second convolution
x = F.conv3d(
x,
self.weight2,
self.bias2,
self.stride2,
self.padding2,
self.dilation2,
self.groups,
padding_mode=self.padding_mode,
)
return x
def forward_with_2d(self, x: torch.Tensor, skip_time_conv: bool = False) -> torch.Tensor:
b, _, _, h, w = x.shape
# First 2D convolution
x = rearrange(x, "b c d h w -> (b d) c h w")
# Squeeze the depth dimension out of weight1 since it's 1
weight1 = self.weight1.squeeze(2)
# Select stride, padding, and dilation for the 2D convolution
stride1 = (self.stride1[1], self.stride1[2])
padding1 = (self.padding1[1], self.padding1[2])
dilation1 = (self.dilation1[1], self.dilation1[2])
x = F.conv2d(
x,
weight1,
self.bias1,
stride1,
padding1,
dilation1,
self.groups,
padding_mode=self.padding_mode,
)
_, _, h, w = x.shape
if skip_time_conv:
x = rearrange(x, "(b d) c h w -> b c d h w", b=b)
return x
# Second convolution which is essentially treated as a 1D convolution across the 'd' dimension
x = rearrange(x, "(b d) c h w -> (b h w) c d", b=b)
# Reshape weight2 to match the expected dimensions for conv1d
weight2 = self.weight2.squeeze(-1).squeeze(-1)
# Use only the relevant dimension for stride, padding, and dilation for the 1D convolution
stride2 = self.stride2[0]
padding2 = self.padding2[0]
dilation2 = self.dilation2[0]
x = F.conv1d(
x,
weight2,
self.bias2,
stride2,
padding2,
dilation2,
self.groups,
padding_mode=self.padding_mode,
)
x = rearrange(x, "(b h w) c d -> b c d h w", b=b, h=h, w=w)
return x
@property
def weight(self) -> torch.Tensor:
return self.weight2
class CausalConv3d(nn.Module):
def __init__(
self,
in_channels: int,
out_channels: int,
kernel_size: int = 3,
stride: Union[int, Tuple[int]] = 1,
dilation: int = 1,
groups: int = 1,
bias: bool = True,
spatial_padding_mode: PaddingModeType = PaddingModeType.ZEROS,
) -> None:
super().__init__()
self.in_channels = in_channels
self.out_channels = out_channels
kernel_size = (kernel_size, kernel_size, kernel_size)
self.time_kernel_size = kernel_size[0]
dilation = (dilation, 1, 1)
height_pad = kernel_size[1] // 2
width_pad = kernel_size[2] // 2
padding = (0, height_pad, width_pad)
self.conv = nn.Conv3d(
in_channels,
out_channels,
kernel_size,
stride=stride,
dilation=dilation,
padding=padding,
padding_mode=spatial_padding_mode.value,
groups=groups,
bias=bias,
)
def forward(self, x: torch.Tensor, causal: bool = True) -> torch.Tensor:
if causal:
first_frame_pad = x[:, :, :1, :, :].repeat((1, 1, self.time_kernel_size - 1, 1, 1))
x = torch.concatenate((first_frame_pad, x), dim=2)
else:
first_frame_pad = x[:, :, :1, :, :].repeat((1, 1, (self.time_kernel_size - 1) // 2, 1, 1))
last_frame_pad = x[:, :, -1:, :, :].repeat((1, 1, (self.time_kernel_size - 1) // 2, 1, 1))
x = torch.concatenate((first_frame_pad, x, last_frame_pad), dim=2)
x = self.conv(x)
return x
@property
def weight(self) -> torch.Tensor:
return self.conv.weight
packages/ltx-core/src/ltx_core/model/video_vae/enums.py
+20
View File
@@ -0,0 +1,20 @@
from enum import Enum
class NormLayerType(Enum):
GROUP_NORM = "group_norm"
PIXEL_NORM = "pixel_norm"
class LogVarianceType(Enum):
PER_CHANNEL = "per_channel"
UNIFORM = "uniform"
CONSTANT = "constant"
NONE = "none"
class PaddingModeType(Enum):
ZEROS = "zeros"
REFLECT = "reflect"
REPLICATE = "replicate"
CIRCULAR = "circular"
packages/ltx-core/src/ltx_core/model/video_vae/model_configurator.py
+77
View File
@@ -0,0 +1,77 @@
from ltx_core.loader.sd_ops import SDOps
from ltx_core.model.model_protocol import ModelConfigurator
from ltx_core.model.video_vae.enums import LogVarianceType, NormLayerType, PaddingModeType
from ltx_core.model.video_vae.video_vae import VideoDecoder, VideoEncoder
class VideoEncoderConfigurator(ModelConfigurator[VideoEncoder]):
"""Configurator for creating a video VAE Encoder from a configuration dictionary."""
@classmethod
def from_config(cls: type[VideoEncoder], config: dict) -> VideoEncoder:
config = config.get("vae", {})
convolution_dimensions = config.get("dims", 3)
in_channels = config.get("in_channels", 3)
latent_channels = config.get("latent_channels", 128)
encoder_spatial_padding_mode = PaddingModeType(config.get("encoder_spatial_padding_mode", "zeros"))
encoder_blocks = config.get("encoder_blocks", [])
patch_size = config.get("patch_size", 4)
norm_layer_str = config.get("norm_layer", "pixel_norm")
latent_log_var_str = config.get("latent_log_var", "uniform")
return VideoEncoder(
convolution_dimensions=convolution_dimensions,
in_channels=in_channels,
out_channels=latent_channels,
encoder_blocks=encoder_blocks,
patch_size=patch_size,
norm_layer=NormLayerType(norm_layer_str),
latent_log_var=LogVarianceType(latent_log_var_str),
encoder_spatial_padding_mode=encoder_spatial_padding_mode,
)
class VideoDecoderConfigurator(ModelConfigurator[VideoDecoder]):
"""Configurator for creating a video VAE Decoder from a configuration dictionary."""
@classmethod
def from_config(cls: type[VideoDecoder], config: dict) -> VideoDecoder:
config = config.get("vae", {})
convolution_dimensions = config.get("dims", 3)
latent_channels = config.get("latent_channels", 128)
decoder_spatial_padding_mode = PaddingModeType(config.get("decoder_spatial_padding_mode", "reflect"))
out_channels = config.get("out_channels", 3)
decoder_blocks = config.get("decoder_blocks", [])
patch_size = config.get("patch_size", 4)
norm_layer_str = config.get("norm_layer", "pixel_norm")
causal = config.get("causal_decoder", False)
timestep_conditioning = config.get("timestep_conditioning", True)
return VideoDecoder(
convolution_dimensions=convolution_dimensions,
in_channels=latent_channels,
out_channels=out_channels,
decoder_blocks=decoder_blocks,
patch_size=patch_size,
norm_layer=NormLayerType(norm_layer_str),
causal=causal,
timestep_conditioning=timestep_conditioning,
decoder_spatial_padding_mode=decoder_spatial_padding_mode,
)
VAE_DECODER_COMFY_KEYS_FILTER = (
SDOps("VAE_DECODER_COMFY_KEYS_FILTER")
.with_matching(prefix="vae.decoder.")
.with_matching(prefix="vae.per_channel_statistics.")
.with_replacement("vae.decoder.", "")
.with_replacement("vae.per_channel_statistics.", "per_channel_statistics.")
)
VAE_ENCODER_COMFY_KEYS_FILTER = (
SDOps("VAE_ENCODER_COMFY_KEYS_FILTER")
.with_matching(prefix="vae.encoder.")
.with_matching(prefix="vae.per_channel_statistics.")
.with_replacement("vae.encoder.", "")
.with_replacement("vae.per_channel_statistics.", "per_channel_statistics.")
)
packages/ltx-core/src/ltx_core/model/video_vae/normalization.py
+3
View File
@@ -0,0 +1,3 @@
from ltx_core.model.common.normalization import PixelNorm, build_normalization_layer
__all__ = ["PixelNorm", "build_normalization_layer"]
packages/ltx-core/src/ltx_core/model/video_vae/ops.py
+85
View File
@@ -0,0 +1,85 @@
import torch
from einops import rearrange
from torch import nn
def patchify(x: torch.Tensor, patch_size_hw: int, patch_size_t: int = 1) -> torch.Tensor:
"""
Rearrange spatial dimensions into channels. Divides image into patch_size x patch_size blocks
and moves pixels from each block into separate channels (space-to-depth).
Args:
x: Input tensor (4D or 5D)
patch_size_hw: Spatial patch size for height and width. With patch_size_hw=4, divides HxW into 4x4 blocks.
patch_size_t: Temporal patch size for frames. Default=1 (no temporal patching).
For 5D: (B, C, F, H, W) -> (B, Cx(patch_size_hw^2)x(patch_size_t), F/patch_size_t, H/patch_size_hw, W/patch_size_hw)
Example: (B, 3, 33, 512, 512) with patch_size_hw=4, patch_size_t=1 -> (B, 48, 33, 128, 128)
"""
if patch_size_hw == 1 and patch_size_t == 1:
return x
if x.dim() == 4:
x = rearrange(x, "b c (h q) (w r) -> b (c r q) h w", q=patch_size_hw, r=patch_size_hw)
elif x.dim() == 5:
x = rearrange(
x,
"b c (f p) (h q) (w r) -> b (c p r q) f h w",
p=patch_size_t,
q=patch_size_hw,
r=patch_size_hw,
)
else:
raise ValueError(f"Invalid input shape: {x.shape}")
return x
def unpatchify(x: torch.Tensor, patch_size_hw: int, patch_size_t: int = 1) -> torch.Tensor:
"""
Rearrange channels back into spatial dimensions. Inverse of patchify - moves pixels from
channels back into patch_size x patch_size blocks (depth-to-space).
Args:
x: Input tensor (4D or 5D)
patch_size_hw: Spatial patch size for height and width. With patch_size_hw=4, expands HxW by 4x.
patch_size_t: Temporal patch size for frames. Default=1 (no temporal expansion).
For 5D: (B, Cx(patch_size_hw^2)x(patch_size_t), F, H, W) -> (B, C, Fxpatch_size_t, Hxpatch_size_hw, Wxpatch_size_hw)
Example: (B, 48, 33, 128, 128) with patch_size_hw=4, patch_size_t=1 -> (B, 3, 33, 512, 512)
"""
if patch_size_hw == 1 and patch_size_t == 1:
return x
if x.dim() == 4:
x = rearrange(x, "b (c r q) h w -> b c (h q) (w r)", q=patch_size_hw, r=patch_size_hw)
elif x.dim() == 5:
x = rearrange(
x,
"b (c p r q) f h w -> b c (f p) (h q) (w r)",
p=patch_size_t,
q=patch_size_hw,
r=patch_size_hw,
)
return x
class PerChannelStatistics(nn.Module):
"""
Per-channel statistics for normalizing and denormalizing the latent representation.
This statics is computed over the entire dataset and stored in model's checkpoint under VAE state_dict.
"""
def __init__(self, latent_channels: int = 128):
super().__init__()
self.register_buffer("std-of-means", torch.empty(latent_channels))
self.register_buffer("mean-of-means", torch.empty(latent_channels))
self.register_buffer("mean-of-stds", torch.empty(latent_channels))
self.register_buffer("mean-of-stds_over_std-of-means", torch.empty(latent_channels))
self.register_buffer("channel", torch.empty(latent_channels))
def un_normalize(self, x: torch.Tensor) -> torch.Tensor:
return (x * self.get_buffer("std-of-means").view(1, -1, 1, 1, 1).to(x)) + self.get_buffer("mean-of-means").view(
1, -1, 1, 1, 1
).to(x)
def normalize(self, x: torch.Tensor) -> torch.Tensor:
return (x - self.get_buffer("mean-of-means").view(1, -1, 1, 1, 1).to(x)) / self.get_buffer("std-of-means").view(
1, -1, 1, 1, 1
).to(x)
packages/ltx-core/src/ltx_core/model/video_vae/resnet.py
+277
View File
@@ -0,0 +1,277 @@
from typing import Optional, Tuple, Union
import torch
from torch import nn
from ltx_core.model.common.normalization import PixelNorm
from ltx_core.model.transformer.timestep_embedding import PixArtAlphaCombinedTimestepSizeEmbeddings
from ltx_core.model.video_vae.convolution import make_conv_nd, make_linear_nd
from ltx_core.model.video_vae.enums import NormLayerType, PaddingModeType
class ResnetBlock3D(nn.Module):
r"""
A Resnet block.
Parameters:
in_channels (`int`): The number of channels in the input.
out_channels (`int`, *optional*, default to be `None`):
The number of output channels for the first conv layer. If None, same as `in_channels`.
dropout (`float`, *optional*, defaults to `0.0`): The dropout probability to use.
groups (`int`, *optional*, default to `32`): The number of groups to use for the first normalization layer.
eps (`float`, *optional*, defaults to `1e-6`): The epsilon to use for the normalization.
"""
def __init__(
self,
dims: Union[int, Tuple[int, int]],
in_channels: int,
out_channels: Optional[int] = None,
dropout: float = 0.0,
groups: int = 32,
eps: float = 1e-6,
norm_layer: NormLayerType = NormLayerType.PIXEL_NORM,
inject_noise: bool = False,
timestep_conditioning: bool = False,
spatial_padding_mode: PaddingModeType = PaddingModeType.ZEROS,
):
super().__init__()
self.in_channels = in_channels
out_channels = in_channels if out_channels is None else out_channels
self.out_channels = out_channels
self.inject_noise = inject_noise
if norm_layer == NormLayerType.GROUP_NORM:
self.norm1 = nn.GroupNorm(num_groups=groups, num_channels=in_channels, eps=eps, affine=True)
elif norm_layer == NormLayerType.PIXEL_NORM:
self.norm1 = PixelNorm()
self.non_linearity = nn.SiLU()
self.conv1 = make_conv_nd(
dims,
in_channels,
out_channels,
kernel_size=3,
stride=1,
padding=1,
causal=True,
spatial_padding_mode=spatial_padding_mode,
)
if inject_noise:
self.per_channel_scale1 = nn.Parameter(torch.zeros((in_channels, 1, 1)))
if norm_layer == NormLayerType.GROUP_NORM:
self.norm2 = nn.GroupNorm(num_groups=groups, num_channels=out_channels, eps=eps, affine=True)
elif norm_layer == NormLayerType.PIXEL_NORM:
self.norm2 = PixelNorm()
self.dropout = torch.nn.Dropout(dropout)
self.conv2 = make_conv_nd(
dims,
out_channels,
out_channels,
kernel_size=3,
stride=1,
padding=1,
causal=True,
spatial_padding_mode=spatial_padding_mode,
)
if inject_noise:
self.per_channel_scale2 = nn.Parameter(torch.zeros((in_channels, 1, 1)))
self.conv_shortcut = (
make_linear_nd(dims=dims, in_channels=in_channels, out_channels=out_channels)
if in_channels != out_channels
else nn.Identity()
)
# Using GroupNorm with 1 group is equivalent to LayerNorm but works with (B, C, ...) layout
# avoiding the need for dimension rearrangement used in standard nn.LayerNorm
self.norm3 = (
nn.GroupNorm(num_groups=1, num_channels=in_channels, eps=eps, affine=True)
if in_channels != out_channels
else nn.Identity()
)
self.timestep_conditioning = timestep_conditioning
if timestep_conditioning:
self.scale_shift_table = nn.Parameter(torch.randn(4, in_channels) / in_channels**0.5)
def _feed_spatial_noise(
self,
hidden_states: torch.Tensor,
per_channel_scale: torch.Tensor,
generator: Optional[torch.Generator] = None,
) -> torch.Tensor:
spatial_shape = hidden_states.shape[-2:]
device = hidden_states.device
dtype = hidden_states.dtype
# similar to the "explicit noise inputs" method in style-gan
spatial_noise = torch.randn(spatial_shape, device=device, dtype=dtype, generator=generator)[None]
scaled_noise = (spatial_noise * per_channel_scale)[None, :, None, ...]
hidden_states = hidden_states + scaled_noise
return hidden_states
def forward(
self,
input_tensor: torch.Tensor,
causal: bool = True,
timestep: Optional[torch.Tensor] = None,
generator: Optional[torch.Generator] = None,
) -> torch.Tensor:
hidden_states = input_tensor
batch_size = hidden_states.shape[0]
hidden_states = self.norm1(hidden_states)
if self.timestep_conditioning:
if timestep is None:
raise ValueError("'timestep' parameter must be provided when 'timestep_conditioning' is True")
ada_values = self.scale_shift_table[None, ..., None, None, None].to(
device=hidden_states.device, dtype=hidden_states.dtype
) + timestep.reshape(
batch_size,
4,
-1,
timestep.shape[-3],
timestep.shape[-2],
timestep.shape[-1],
)
shift1, scale1, shift2, scale2 = ada_values.unbind(dim=1)
hidden_states = hidden_states * (1 + scale1) + shift1
hidden_states = self.non_linearity(hidden_states)
hidden_states = self.conv1(hidden_states, causal=causal)
if self.inject_noise:
hidden_states = self._feed_spatial_noise(
hidden_states,
self.per_channel_scale1.to(device=hidden_states.device, dtype=hidden_states.dtype),
generator=generator,
)
hidden_states = self.norm2(hidden_states)
if self.timestep_conditioning:
hidden_states = hidden_states * (1 + scale2) + shift2
hidden_states = self.non_linearity(hidden_states)
hidden_states = self.dropout(hidden_states)
hidden_states = self.conv2(hidden_states, causal=causal)
if self.inject_noise:
hidden_states = self._feed_spatial_noise(
hidden_states,
self.per_channel_scale2.to(device=hidden_states.device, dtype=hidden_states.dtype),
generator=generator,
)
input_tensor = self.norm3(input_tensor)
batch_size = input_tensor.shape[0]
input_tensor = self.conv_shortcut(input_tensor)
output_tensor = input_tensor + hidden_states
return output_tensor
class UNetMidBlock3D(nn.Module):
"""
A 3D UNet mid-block [`UNetMidBlock3D`] with multiple residual blocks.
Args:
in_channels (`int`): The number of input channels.
dropout (`float`, *optional*, defaults to 0.0): The dropout rate.
num_layers (`int`, *optional*, defaults to 1): The number of residual blocks.
resnet_eps (`float`, *optional*, 1e-6 ): The epsilon value for the resnet blocks.
resnet_groups (`int`, *optional*, defaults to 32):
The number of groups to use in the group normalization layers of the resnet blocks.
norm_layer (`str`, *optional*, defaults to `group_norm`):
The normalization layer to use. Can be either `group_norm` or `pixel_norm`.
inject_noise (`bool`, *optional*, defaults to `False`):
Whether to inject noise into the hidden states.
timestep_conditioning (`bool`, *optional*, defaults to `False`):
Whether to condition the hidden states on the timestep.
Returns:
`torch.Tensor`: The output of the last residual block, which is a tensor of shape `(batch_size,
in_channels, height, width)`.
"""
def __init__(
self,
dims: Union[int, Tuple[int, int]],
in_channels: int,
dropout: float = 0.0,
num_layers: int = 1,
resnet_eps: float = 1e-6,
resnet_groups: int = 32,
norm_layer: NormLayerType = NormLayerType.GROUP_NORM,
inject_noise: bool = False,
timestep_conditioning: bool = False,
spatial_padding_mode: PaddingModeType = PaddingModeType.ZEROS,
):
super().__init__()
resnet_groups = resnet_groups if resnet_groups is not None else min(in_channels // 4, 32)
self.timestep_conditioning = timestep_conditioning
if timestep_conditioning:
self.time_embedder = PixArtAlphaCombinedTimestepSizeEmbeddings(
embedding_dim=in_channels * 4, size_emb_dim=0
)
self.res_blocks = nn.ModuleList(
[
ResnetBlock3D(
dims=dims,
in_channels=in_channels,
out_channels=in_channels,
eps=resnet_eps,
groups=resnet_groups,
dropout=dropout,
norm_layer=norm_layer,
inject_noise=inject_noise,
timestep_conditioning=timestep_conditioning,
spatial_padding_mode=spatial_padding_mode,
)
for _ in range(num_layers)
]
)
def forward(
self,
hidden_states: torch.Tensor,
causal: bool = True,
timestep: Optional[torch.Tensor] = None,
generator: Optional[torch.Generator] = None,
) -> torch.Tensor:
timestep_embed = None
if self.timestep_conditioning:
if timestep is None:
raise ValueError("'timestep' parameter must be provided when 'timestep_conditioning' is True")
batch_size = hidden_states.shape[0]
timestep_embed = self.time_embedder(
timestep=timestep.flatten(),
hidden_dtype=hidden_states.dtype,
)
timestep_embed = timestep_embed.view(batch_size, timestep_embed.shape[-1], 1, 1, 1)
for resnet in self.res_blocks:
hidden_states = resnet(
hidden_states,
causal=causal,
timestep=timestep_embed,
generator=generator,
)
return hidden_states
packages/ltx-core/src/ltx_core/model/video_vae/sampling.py
+123
View File
@@ -0,0 +1,123 @@
import math
from typing import Tuple, Union
import torch
from einops import rearrange
from torch import nn
from .convolution import make_conv_nd
from .enums import PaddingModeType
class SpaceToDepthDownsample(nn.Module):
def __init__(
self,
dims: Union[int, Tuple[int, int]],
in_channels: int,
out_channels: int,
stride: Tuple[int, int, int],
spatial_padding_mode: PaddingModeType = PaddingModeType.ZEROS,
):
super().__init__()
self.stride = stride
self.group_size = in_channels * math.prod(stride) // out_channels
self.conv = make_conv_nd(
dims=dims,
in_channels=in_channels,
out_channels=out_channels // math.prod(stride),
kernel_size=3,
stride=1,
causal=True,
spatial_padding_mode=spatial_padding_mode,
)
def forward(
self,
x: torch.Tensor,
causal: bool = True,
) -> torch.Tensor:
if self.stride[0] == 2:
x = torch.cat([x[:, :, :1, :, :], x], dim=2) # duplicate first frames for padding
# skip connection
x_in = rearrange(
x,
"b c (d p1) (h p2) (w p3) -> b (c p1 p2 p3) d h w",
p1=self.stride[0],
p2=self.stride[1],
p3=self.stride[2],
)
x_in = rearrange(x_in, "b (c g) d h w -> b c g d h w", g=self.group_size)
x_in = x_in.mean(dim=2)
# conv
x = self.conv(x, causal=causal)
x = rearrange(
x,
"b c (d p1) (h p2) (w p3) -> b (c p1 p2 p3) d h w",
p1=self.stride[0],
p2=self.stride[1],
p3=self.stride[2],
)
x = x + x_in
return x
class DepthToSpaceUpsample(nn.Module):
def __init__(
self,
dims: int | Tuple[int, int],
in_channels: int,
stride: Tuple[int, int, int],
residual: bool = False,
out_channels_reduction_factor: int = 1,
spatial_padding_mode: PaddingModeType = PaddingModeType.ZEROS,
):
super().__init__()
self.stride = stride
self.out_channels = math.prod(stride) * in_channels // out_channels_reduction_factor
self.conv = make_conv_nd(
dims=dims,
in_channels=in_channels,
out_channels=self.out_channels,
kernel_size=3,
stride=1,
causal=True,
spatial_padding_mode=spatial_padding_mode,
)
self.residual = residual
self.out_channels_reduction_factor = out_channels_reduction_factor
def forward(
self,
x: torch.Tensor,
causal: bool = True,
) -> torch.Tensor:
if self.residual:
# Reshape and duplicate the input to match the output shape
x_in = rearrange(
x,
"b (c p1 p2 p3) d h w -> b c (d p1) (h p2) (w p3)",
p1=self.stride[0],
p2=self.stride[1],
p3=self.stride[2],
)
num_repeat = math.prod(self.stride) // self.out_channels_reduction_factor
x_in = x_in.repeat(1, num_repeat, 1, 1, 1)
if self.stride[0] == 2:
x_in = x_in[:, :, 1:, :, :]
x = self.conv(x, causal=causal)
x = rearrange(
x,
"b (c p1 p2 p3) d h w -> b c (d p1) (h p2) (w p3)",
p1=self.stride[0],
p2=self.stride[1],
p3=self.stride[2],
)
if self.stride[0] == 2:
x = x[:, :, 1:, :, :]
if self.residual:
x = x + x_in
return x
packages/ltx-core/src/ltx_core/model/video_vae/tiling.py
+260
View File
@@ -0,0 +1,260 @@
import itertools
from dataclasses import dataclass
from typing import Callable, List, NamedTuple, Tuple
import torch
def compute_trapezoidal_mask_1d(
length: int,
ramp_left: int,
ramp_right: int,
left_starts_from_0: bool = False,
) -> torch.Tensor:
"""
Generate a 1D trapezoidal blending mask with linear ramps.
Args:
length: Output length of the mask.
ramp_left: Fade-in length on the left.
ramp_right: Fade-out length on the right.
left_starts_from_0: Whether the ramp starts from 0 or first non-zero value.
Useful for temporal tiles where the first tile is causal.
Returns:
A 1D tensor of shape `(length,)` with values in [0, 1].
"""
if length <= 0:
raise ValueError("Mask length must be positive.")
ramp_left = max(0, min(ramp_left, length))
ramp_right = max(0, min(ramp_right, length))
mask = torch.ones(length)
if ramp_left > 0:
interval_length = ramp_left + 1 if left_starts_from_0 else ramp_left + 2
fade_in = torch.linspace(0.0, 1.0, interval_length)[:-1]
if not left_starts_from_0:
fade_in = fade_in[1:]
mask[:ramp_left] *= fade_in
if ramp_right > 0:
fade_out = torch.linspace(1.0, 0.0, steps=ramp_right + 2)[1:-1]
mask[-ramp_right:] *= fade_out
return mask.clamp_(0, 1)
@dataclass(frozen=True)
class SpatialTilingConfig:
"""Configuration for dividing each frame into spatial tiles with optional overlap.
Args:
tile_size_in_pixels (int): Size of each tile in pixels. Must be at least 64 and divisible by 32.
tile_overlap_in_pixels (int, optional): Overlap between tiles in pixels. Must be divisible by 32. Defaults to 0.
"""
tile_size_in_pixels: int
tile_overlap_in_pixels: int = 0
def __post_init__(self) -> None:
if self.tile_size_in_pixels < 64:
raise ValueError(f"tile_size_in_pixels must be at least 64, got {self.tile_size_in_pixels}")
if self.tile_size_in_pixels % 32 != 0:
raise ValueError(f"tile_size_in_pixels must be divisible by 32, got {self.tile_size_in_pixels}")
if self.tile_overlap_in_pixels % 32 != 0:
raise ValueError(f"tile_overlap_in_pixels must be divisible by 32, got {self.tile_overlap_in_pixels}")
if self.tile_overlap_in_pixels >= self.tile_size_in_pixels:
raise ValueError(
f"Overlap must be less than tile size, got {self.tile_overlap_in_pixels} and {self.tile_size_in_pixels}"
)
@dataclass(frozen=True)
class TemporalTilingConfig:
"""Configuration for dividing a video into temporal tiles (chunks of frames) with optional overlap.
Args:
tile_size_in_frames (int): Number of frames in each tile. Must be at least 16 and divisible by 8.
tile_overlap_in_frames (int, optional): Number of overlapping frames between consecutive tiles.
Must be divisible by 8. Defaults to 0.
"""
tile_size_in_frames: int
tile_overlap_in_frames: int = 0
def __post_init__(self) -> None:
if self.tile_size_in_frames < 16:
raise ValueError(f"tile_size_in_frames must be at least 16, got {self.tile_size_in_frames}")
if self.tile_size_in_frames % 8 != 0:
raise ValueError(f"tile_size_in_frames must be divisible by 8, got {self.tile_size_in_frames}")
if self.tile_overlap_in_frames % 8 != 0:
raise ValueError(f"tile_overlap_in_frames must be divisible by 8, got {self.tile_overlap_in_frames}")
if self.tile_overlap_in_frames >= self.tile_size_in_frames:
raise ValueError(
f"Overlap must be less than tile size, got {self.tile_overlap_in_frames} and {self.tile_size_in_frames}"
)
@dataclass(frozen=True)
class TilingConfig:
"""Configuration for splitting video into tiles with optional overlap.
Attributes:
spatial_config: Configuration for splitting spatial dimensions into tiles.
temporal_config: Configuration for splitting temporal dimension into tiles.
"""
spatial_config: SpatialTilingConfig | None = None
temporal_config: TemporalTilingConfig | None = None
@classmethod
def default(cls) -> "TilingConfig":
return cls(
spatial_config=SpatialTilingConfig(tile_size_in_pixels=512, tile_overlap_in_pixels=64),
temporal_config=TemporalTilingConfig(tile_size_in_frames=64, tile_overlap_in_frames=24),
)
@dataclass(frozen=True)
class DimensionIntervals:
"""Intervals which a single dimension of the latent space is split into.
Each interval is defined by its start, end, left ramp, and right ramp.
The start and end are the indices of the first and last element (exclusive) in the interval.
Ramps are regions of the interval where the value of the mask tensor is
interpolated between 0 and 1 for blending with neighboring intervals.
The left ramp and right ramp values are the lengths of the left and right ramps.
"""
starts: List[int]
ends: List[int]
left_ramps: List[int]
right_ramps: List[int]
@dataclass(frozen=True)
class LatentIntervals:
"""Intervals which the latent tensor of given shape is split into.
Each dimension of the latent space is split into intervals based on the length along said dimension.
"""
original_shape: torch.Size
dimension_intervals: Tuple[DimensionIntervals, ...]
# Operation to split a single dimension of the tensor into intervals based on the length along the dimension.
SplitOperation = Callable[[int], DimensionIntervals]
# Operation to map the intervals in input dimension to slices and masks along a corresponding output dimension.
MappingOperation = Callable[[DimensionIntervals], tuple[list[slice], list[torch.Tensor | None]]]
def default_split_operation(length: int) -> DimensionIntervals:
return DimensionIntervals(starts=[0], ends=[length], left_ramps=[0], right_ramps=[0])
DEFAULT_SPLIT_OPERATION: SplitOperation = default_split_operation
def default_mapping_operation(
_intervals: DimensionIntervals,
) -> tuple[list[slice], list[torch.Tensor | None]]:
return [slice(0, None)], [None]
DEFAULT_MAPPING_OPERATION: MappingOperation = default_mapping_operation
class Tile(NamedTuple):
"""
Represents a single tile.
Attributes:
in_coords:
Tuple of slices specifying where to cut the tile from the INPUT tensor.
out_coords:
Tuple of slices specifying where this tile's OUTPUT should be placed in the reconstructed OUTPUT tensor.
masks_1d:
Per-dimension masks in OUTPUT units.
These are used to create all-dimensional blending mask.
Methods:
blend_mask:
Create a single N-D mask from the per-dimension masks.
"""
in_coords: Tuple[slice, ...]
out_coords: Tuple[slice, ...]
masks_1d: Tuple[Tuple[torch.Tensor, ...]]
@property
def blend_mask(self) -> torch.Tensor:
num_dims = len(self.out_coords)
per_dimension_masks: List[torch.Tensor] = []
for dim_idx in range(num_dims):
mask_1d = self.masks_1d[dim_idx]
view_shape = [1] * num_dims
if mask_1d is None:
# Broadcast mask along this dimension (length 1).
one = torch.ones(1)
view_shape[dim_idx] = 1
per_dimension_masks.append(one.view(*view_shape))
continue
# Reshape (L,) -> (1, ..., L, ..., 1) so masks across dimensions broadcast-multiply.
view_shape[dim_idx] = mask_1d.shape[0]
per_dimension_masks.append(mask_1d.view(*view_shape))
# Multiply per-dimension masks to form the full N-D mask (separable blending window).
combined_mask = per_dimension_masks[0]
for mask in per_dimension_masks[1:]:
combined_mask = combined_mask * mask
return combined_mask
def create_tiles_from_intervals_and_mappers(
intervals: LatentIntervals,
mappers: List[MappingOperation],
) -> List[Tile]:
full_dim_input_slices = []
full_dim_output_slices = []
full_dim_masks_1d = []
for axis_index in range(len(intervals.original_shape)):
dimension_intervals = intervals.dimension_intervals[axis_index]
starts = dimension_intervals.starts
ends = dimension_intervals.ends
input_slices = [slice(s, e) for s, e in zip(starts, ends, strict=True)]
output_slices, masks_1d = mappers[axis_index](dimension_intervals)
full_dim_input_slices.append(input_slices)
full_dim_output_slices.append(output_slices)
full_dim_masks_1d.append(masks_1d)
tiles = []
tile_in_coords = list(itertools.product(*full_dim_input_slices))
tile_out_coords = list(itertools.product(*full_dim_output_slices))
tile_mask_1ds = list(itertools.product(*full_dim_masks_1d))
for in_coord, out_coord, mask_1d in zip(tile_in_coords, tile_out_coords, tile_mask_1ds, strict=True):
tiles.append(
Tile(
in_coords=in_coord,
out_coords=out_coord,
masks_1d=mask_1d,
)
)
return tiles
def create_tiles(
latent_shape: torch.Size,
splitters: List[SplitOperation],
mappers: List[MappingOperation],
) -> List[Tile]:
if len(splitters) != len(latent_shape):
raise ValueError(
f"Number of splitters must be equal to number of dimensions in latent shape, "
f"got {len(splitters)} and {len(latent_shape)}"
)
if len(mappers) != len(latent_shape):
raise ValueError(
f"Number of mappers must be equal to number of dimensions in latent shape, "
f"got {len(mappers)} and {len(latent_shape)}"
)
intervals = [splitter(length) for splitter, length in zip(splitters, latent_shape, strict=True)]
latent_intervals = LatentIntervals(original_shape=latent_shape, dimension_intervals=tuple(intervals))
return create_tiles_from_intervals_and_mappers(latent_intervals, mappers)
packages/ltx-core/src/ltx_core/model/video_vae/video_vae.py
+925
View File
@@ -0,0 +1,925 @@
from dataclasses import replace
from typing import Any, Callable, Iterator, List, Optional, Tuple
import torch
from einops import rearrange
from torch import nn
from ltx_core.model.common.normalization import PixelNorm
from ltx_core.model.transformer.timestep_embedding import PixArtAlphaCombinedTimestepSizeEmbeddings
from ltx_core.model.video_vae.convolution import make_conv_nd
from ltx_core.model.video_vae.enums import LogVarianceType, NormLayerType, PaddingModeType
from ltx_core.model.video_vae.ops import PerChannelStatistics, patchify, unpatchify
from ltx_core.model.video_vae.resnet import ResnetBlock3D, UNetMidBlock3D
from ltx_core.model.video_vae.sampling import DepthToSpaceUpsample, SpaceToDepthDownsample
from ltx_core.model.video_vae.tiling import (
DEFAULT_MAPPING_OPERATION,
DEFAULT_SPLIT_OPERATION,
DimensionIntervals,
MappingOperation,
SplitOperation,
Tile,
TilingConfig,
compute_trapezoidal_mask_1d,
create_tiles,
)
from ltx_core.types import SpatioTemporalScaleFactors, VideoLatentShape
def _make_encoder_block(
block_name: str,
block_config: dict[str, Any],
in_channels: int,
convolution_dimensions: int,
norm_layer: NormLayerType,
norm_num_groups: int,
spatial_padding_mode: PaddingModeType,
) -> Tuple[nn.Module, int]:
out_channels = in_channels
if block_name == "res_x":
block = UNetMidBlock3D(
dims=convolution_dimensions,
in_channels=in_channels,
num_layers=block_config["num_layers"],
resnet_eps=1e-6,
resnet_groups=norm_num_groups,
norm_layer=norm_layer,
spatial_padding_mode=spatial_padding_mode,
)
elif block_name == "res_x_y":
out_channels = in_channels * block_config.get("multiplier", 2)
block = ResnetBlock3D(
dims=convolution_dimensions,
in_channels=in_channels,
out_channels=out_channels,
eps=1e-6,
groups=norm_num_groups,
norm_layer=norm_layer,
spatial_padding_mode=spatial_padding_mode,
)
elif block_name == "compress_time":
block = make_conv_nd(
dims=convolution_dimensions,
in_channels=in_channels,
out_channels=out_channels,
kernel_size=3,
stride=(2, 1, 1),
causal=True,
spatial_padding_mode=spatial_padding_mode,
)
elif block_name == "compress_space":
block = make_conv_nd(
dims=convolution_dimensions,
in_channels=in_channels,
out_channels=out_channels,
kernel_size=3,
stride=(1, 2, 2),
causal=True,
spatial_padding_mode=spatial_padding_mode,
)
elif block_name == "compress_all":
block = make_conv_nd(
dims=convolution_dimensions,
in_channels=in_channels,
out_channels=out_channels,
kernel_size=3,
stride=(2, 2, 2),
causal=True,
spatial_padding_mode=spatial_padding_mode,
)
elif block_name == "compress_all_x_y":
out_channels = in_channels * block_config.get("multiplier", 2)
block = make_conv_nd(
dims=convolution_dimensions,
in_channels=in_channels,
out_channels=out_channels,
kernel_size=3,
stride=(2, 2, 2),
causal=True,
spatial_padding_mode=spatial_padding_mode,
)
elif block_name == "compress_all_res":
out_channels = in_channels * block_config.get("multiplier", 2)
block = SpaceToDepthDownsample(
dims=convolution_dimensions,
in_channels=in_channels,
out_channels=out_channels,
stride=(2, 2, 2),
spatial_padding_mode=spatial_padding_mode,
)
elif block_name == "compress_space_res":
out_channels = in_channels * block_config.get("multiplier", 2)
block = SpaceToDepthDownsample(
dims=convolution_dimensions,
in_channels=in_channels,
out_channels=out_channels,
stride=(1, 2, 2),
spatial_padding_mode=spatial_padding_mode,
)
elif block_name == "compress_time_res":
out_channels = in_channels * block_config.get("multiplier", 2)
block = SpaceToDepthDownsample(
dims=convolution_dimensions,
in_channels=in_channels,
out_channels=out_channels,
stride=(2, 1, 1),
spatial_padding_mode=spatial_padding_mode,
)
else:
raise ValueError(f"unknown block: {block_name}")
return block, out_channels
class VideoEncoder(nn.Module):
_DEFAULT_NORM_NUM_GROUPS = 32
"""
Variational Autoencoder Encoder. Encodes video frames into a latent representation.
The encoder compresses the input video through a series of downsampling operations controlled by
patch_size and encoder_blocks. The output is a normalized latent tensor with shape (B, 128, F', H', W').
Compression Behavior:
The total compression is determined by:
1. Initial spatial compression via patchify: H -> H/4, W -> W/4 (patch_size=4)
2. Sequential compression through encoder_blocks based on their stride patterns
Compression blocks apply 2x compression in specified dimensions:
- "compress_time" / "compress_time_res": temporal only
- "compress_space" / "compress_space_res": spatial only (H and W)
- "compress_all" / "compress_all_res": all dimensions (F, H, W)
- "res_x" / "res_x_y": no compression
Standard LTX Video configuration:
- patch_size=4
- encoder_blocks: 1x compress_space_res, 1x compress_time_res, 2x compress_all_res
- Final dimensions: F' = 1 + (F-1)/8, H' = H/32, W' = W/32
- Example: (B, 3, 33, 512, 512) -> (B, 128, 5, 16, 16)
- Note: Input must have 1 + 8*k frames (e.g., 1, 9, 17, 25, 33...)
Args:
convolution_dimensions: The number of dimensions to use in convolutions (2D or 3D).
in_channels: The number of input channels. For RGB images, this is 3.
out_channels: The number of output channels (latent channels). For latent channels, this is 128.
encoder_blocks: The list of blocks to construct the encoder. Each block is a tuple of (block_name, params)
where params is either an int (num_layers) or a dict with configuration.
patch_size: The patch size for initial spatial compression. Should be a power of 2.
norm_layer: The normalization layer to use. Can be either `group_norm` or `pixel_norm`.
latent_log_var: The log variance mode. Can be either `per_channel`, `uniform`, `constant` or `none`.
"""
def __init__(
self,
convolution_dimensions: int = 3,
in_channels: int = 3,
out_channels: int = 128,
encoder_blocks: List[Tuple[str, int]] | List[Tuple[str, dict[str, Any]]] = [], # noqa: B006
patch_size: int = 4,
norm_layer: NormLayerType = NormLayerType.PIXEL_NORM,
latent_log_var: LogVarianceType = LogVarianceType.UNIFORM,
encoder_spatial_padding_mode: PaddingModeType = PaddingModeType.ZEROS,
):
super().__init__()
self.patch_size = patch_size
self.norm_layer = norm_layer
self.latent_channels = out_channels
self.latent_log_var = latent_log_var
self._norm_num_groups = self._DEFAULT_NORM_NUM_GROUPS
# Per-channel statistics for normalizing latents
self.per_channel_statistics = PerChannelStatistics(latent_channels=out_channels)
in_channels = in_channels * patch_size**2
feature_channels = out_channels
self.conv_in = make_conv_nd(
dims=convolution_dimensions,
in_channels=in_channels,
out_channels=feature_channels,
kernel_size=3,
stride=1,
padding=1,
causal=True,
spatial_padding_mode=encoder_spatial_padding_mode,
)
self.down_blocks = nn.ModuleList([])
for block_name, block_params in encoder_blocks:
# Convert int to dict format for uniform handling
block_config = {"num_layers": block_params} if isinstance(block_params, int) else block_params
block, feature_channels = _make_encoder_block(
block_name=block_name,
block_config=block_config,
in_channels=feature_channels,
convolution_dimensions=convolution_dimensions,
norm_layer=norm_layer,
norm_num_groups=self._norm_num_groups,
spatial_padding_mode=encoder_spatial_padding_mode,
)
self.down_blocks.append(block)
# out
if norm_layer == NormLayerType.GROUP_NORM:
self.conv_norm_out = nn.GroupNorm(num_channels=feature_channels, num_groups=self._norm_num_groups, eps=1e-6)
elif norm_layer == NormLayerType.PIXEL_NORM:
self.conv_norm_out = PixelNorm()
self.conv_act = nn.SiLU()
conv_out_channels = out_channels
if latent_log_var == LogVarianceType.PER_CHANNEL:
conv_out_channels *= 2
elif latent_log_var in {LogVarianceType.UNIFORM, LogVarianceType.CONSTANT}:
conv_out_channels += 1
elif latent_log_var != LogVarianceType.NONE:
raise ValueError(f"Invalid latent_log_var: {latent_log_var}")
self.conv_out = make_conv_nd(
dims=convolution_dimensions,
in_channels=feature_channels,
out_channels=conv_out_channels,
kernel_size=3,
padding=1,
causal=True,
spatial_padding_mode=encoder_spatial_padding_mode,
)
def forward(self, sample: torch.Tensor) -> torch.Tensor:
r"""
Encode video frames into normalized latent representation.
Args:
sample: Input video (B, C, F, H, W). F must be 1 + 8*k (e.g., 1, 9, 17, 25, 33...).
Returns:
Normalized latent means (B, 128, F', H', W') where F' = 1+(F-1)/8, H' = H/32, W' = W/32.
Example: (B, 3, 33, 512, 512) -> (B, 128, 5, 16, 16).
"""
# Validate frame count
frames_count = sample.shape[2]
if ((frames_count - 1) % 8) != 0:
raise ValueError(
"Invalid number of frames: Encode input must have 1 + 8 * x frames "
"(e.g., 1, 9, 17, ...). Please check your input."
)
# Initial spatial compression: trade spatial resolution for channel depth
# This reduces H,W by patch_size and increases channels, making convolutions more efficient
# Example: (B, 3, F, 512, 512) -> (B, 48, F, 128, 128) with patch_size=4
sample = patchify(sample, patch_size_hw=self.patch_size, patch_size_t=1)
sample = self.conv_in(sample)
for down_block in self.down_blocks:
sample = down_block(sample)
sample = self.conv_norm_out(sample)
sample = self.conv_act(sample)
sample = self.conv_out(sample)
if self.latent_log_var == LogVarianceType.UNIFORM:
# Uniform Variance: model outputs N means and 1 shared log-variance channel.
# We need to expand the single logvar to match the number of means channels
# to create a format compatible with PER_CHANNEL (means + logvar, each with N channels).
# Sample shape: (B, N+1, ...) where N = latent_channels (e.g., 128 means + 1 logvar = 129)
# Target shape: (B, 2*N, ...) where first N are means, last N are logvar
if sample.shape[1] < 2:
raise ValueError(
f"Invalid channel count for UNIFORM mode: expected at least 2 channels "
f"(N means + 1 logvar), got {sample.shape[1]}"
)
# Extract means (first N channels) and logvar (last 1 channel)
means = sample[:, :-1, ...] # (B, N, ...)
logvar = sample[:, -1:, ...] # (B, 1, ...)
# Repeat logvar N times to match means channels
# Use expand/repeat pattern that works for both 4D and 5D tensors
num_channels = means.shape[1]
repeat_shape = [1, num_channels] + [1] * (sample.ndim - 2)
repeated_logvar = logvar.repeat(*repeat_shape) # (B, N, ...)
# Concatenate to create (B, 2*N, ...) format: [means, repeated_logvar]
sample = torch.cat([means, repeated_logvar], dim=1)
elif self.latent_log_var == LogVarianceType.CONSTANT:
sample = sample[:, :-1, ...]
approx_ln_0 = -30 # this is the minimal clamp value in DiagonalGaussianDistribution objects
sample = torch.cat(
[sample, torch.ones_like(sample, device=sample.device) * approx_ln_0],
dim=1,
)
# Split into means and logvar, then normalize means
means, _ = torch.chunk(sample, 2, dim=1)
return self.per_channel_statistics.normalize(means)
def _make_decoder_block(
block_name: str,
block_config: dict[str, Any],
in_channels: int,
convolution_dimensions: int,
norm_layer: NormLayerType,
timestep_conditioning: bool,
norm_num_groups: int,
spatial_padding_mode: PaddingModeType,
) -> Tuple[nn.Module, int]:
out_channels = in_channels
if block_name == "res_x":
block = UNetMidBlock3D(
dims=convolution_dimensions,
in_channels=in_channels,
num_layers=block_config["num_layers"],
resnet_eps=1e-6,
resnet_groups=norm_num_groups,
norm_layer=norm_layer,
inject_noise=block_config.get("inject_noise", False),
timestep_conditioning=timestep_conditioning,
spatial_padding_mode=spatial_padding_mode,
)
elif block_name == "attn_res_x":
block = UNetMidBlock3D(
dims=convolution_dimensions,
in_channels=in_channels,
num_layers=block_config["num_layers"],
resnet_groups=norm_num_groups,
norm_layer=norm_layer,
inject_noise=block_config.get("inject_noise", False),
timestep_conditioning=timestep_conditioning,
attention_head_dim=block_config["attention_head_dim"],
spatial_padding_mode=spatial_padding_mode,
)
elif block_name == "res_x_y":
out_channels = in_channels // block_config.get("multiplier", 2)
block = ResnetBlock3D(
dims=convolution_dimensions,
in_channels=in_channels,
out_channels=out_channels,
eps=1e-6,
groups=norm_num_groups,
norm_layer=norm_layer,
inject_noise=block_config.get("inject_noise", False),
timestep_conditioning=False,
spatial_padding_mode=spatial_padding_mode,
)
elif block_name == "compress_time":
block = DepthToSpaceUpsample(
dims=convolution_dimensions,
in_channels=in_channels,
stride=(2, 1, 1),
spatial_padding_mode=spatial_padding_mode,
)
elif block_name == "compress_space":
block = DepthToSpaceUpsample(
dims=convolution_dimensions,
in_channels=in_channels,
stride=(1, 2, 2),
spatial_padding_mode=spatial_padding_mode,
)
elif block_name == "compress_all":
out_channels = in_channels // block_config.get("multiplier", 1)
block = DepthToSpaceUpsample(
dims=convolution_dimensions,
in_channels=in_channels,
stride=(2, 2, 2),
residual=block_config.get("residual", False),
out_channels_reduction_factor=block_config.get("multiplier", 1),
spatial_padding_mode=spatial_padding_mode,
)
else:
raise ValueError(f"unknown layer: {block_name}")
return block, out_channels
class VideoDecoder(nn.Module):
_DEFAULT_NORM_NUM_GROUPS = 32
"""
Variational Autoencoder Decoder. Decodes latent representation into video frames.
The decoder upsamples latents through a series of upsampling operations (inverse of encoder).
Output dimensions: F = 8x(F'-1) + 1, H = 32xH', W = 32xW' for standard LTX Video configuration.
Upsampling blocks expand dimensions by 2x in specified dimensions:
- "compress_time": temporal only
- "compress_space": spatial only (H and W)
- "compress_all": all dimensions (F, H, W)
- "res_x" / "res_x_y" / "attn_res_x": no upsampling
Causal Mode:
causal=False (standard): Symmetric padding, allows future frame dependencies.
causal=True: Causal padding, each frame depends only on past/current frames.
First frame removed after temporal upsampling in both modes. Output shape unchanged.
Example: (B, 128, 5, 16, 16) -> (B, 3, 33, 512, 512) for both modes.
Args:
convolution_dimensions: The number of dimensions to use in convolutions (2D or 3D).
in_channels: The number of input channels (latent channels). Default is 128.
out_channels: The number of output channels. For RGB images, this is 3.
decoder_blocks: The list of blocks to construct the decoder. Each block is a tuple of (block_name, params)
where params is either an int (num_layers) or a dict with configuration.
patch_size: Final spatial expansion factor. For standard LTX Video, use 4 for 4x spatial expansion:
H -> Hx4, W -> Wx4. Should be a power of 2.
norm_layer: The normalization layer to use. Can be either `group_norm` or `pixel_norm`.
causal: Whether to use causal convolutions. For standard LTX Video, use False for symmetric padding.
When True, uses causal padding (past/current frames only).
timestep_conditioning: Whether to condition the decoder on timestep for denoising.
"""
def __init__(
self,
convolution_dimensions: int = 3,
in_channels: int = 128,
out_channels: int = 3,
decoder_blocks: List[Tuple[str, int | dict]] = [], # noqa: B006
patch_size: int = 4,
norm_layer: NormLayerType = NormLayerType.PIXEL_NORM,
causal: bool = False,
timestep_conditioning: bool = False,
decoder_spatial_padding_mode: PaddingModeType = PaddingModeType.REFLECT,
):
super().__init__()
# Spatiotemporal downscaling between decoded video space and VAE latents.
# According to the LTXV paper, the standard configuration downsamples
# video inputs by a factor of 8 in the temporal dimension and 32 in
# each spatial dimension (height and width). This parameter determines how
# many video frames and pixels correspond to a single latent cell.
self.video_downscale_factors = SpatioTemporalScaleFactors(
time=8,
width=32,
height=32,
)
self.patch_size = patch_size
out_channels = out_channels * patch_size**2
self.causal = causal
self.timestep_conditioning = timestep_conditioning
self._norm_num_groups = self._DEFAULT_NORM_NUM_GROUPS
# Per-channel statistics for denormalizing latents
self.per_channel_statistics = PerChannelStatistics(latent_channels=in_channels)
# Noise and timestep parameters for decoder conditioning
self.decode_noise_scale = 0.025
self.decode_timestep = 0.05
# Compute initial feature_channels by going through blocks in reverse
# This determines the channel width at the start of the decoder
feature_channels = in_channels
for block_name, block_params in list(reversed(decoder_blocks)):
block_config = block_params if isinstance(block_params, dict) else {}
if block_name == "res_x_y":
feature_channels = feature_channels * block_config.get("multiplier", 2)
if block_name == "compress_all":
feature_channels = feature_channels * block_config.get("multiplier", 1)
self.conv_in = make_conv_nd(
dims=convolution_dimensions,
in_channels=in_channels,
out_channels=feature_channels,
kernel_size=3,
stride=1,
padding=1,
causal=True,
spatial_padding_mode=decoder_spatial_padding_mode,
)
self.up_blocks = nn.ModuleList([])
for block_name, block_params in list(reversed(decoder_blocks)):
# Convert int to dict format for uniform handling
block_config = {"num_layers": block_params} if isinstance(block_params, int) else block_params
block, feature_channels = _make_decoder_block(
block_name=block_name,
block_config=block_config,
in_channels=feature_channels,
convolution_dimensions=convolution_dimensions,
norm_layer=norm_layer,
timestep_conditioning=timestep_conditioning,
norm_num_groups=self._norm_num_groups,
spatial_padding_mode=decoder_spatial_padding_mode,
)
self.up_blocks.append(block)
if norm_layer == NormLayerType.GROUP_NORM:
self.conv_norm_out = nn.GroupNorm(num_channels=feature_channels, num_groups=self._norm_num_groups, eps=1e-6)
elif norm_layer == NormLayerType.PIXEL_NORM:
self.conv_norm_out = PixelNorm()
self.conv_act = nn.SiLU()
self.conv_out = make_conv_nd(
dims=convolution_dimensions,
in_channels=feature_channels,
out_channels=out_channels,
kernel_size=3,
padding=1,
causal=True,
spatial_padding_mode=decoder_spatial_padding_mode,
)
if timestep_conditioning:
self.timestep_scale_multiplier = nn.Parameter(torch.tensor(1000.0))
self.last_time_embedder = PixArtAlphaCombinedTimestepSizeEmbeddings(
embedding_dim=feature_channels * 2, size_emb_dim=0
)
self.last_scale_shift_table = nn.Parameter(torch.empty(2, feature_channels))
# def forward(self, sample: torch.Tensor, target_shape) -> torch.Tensor:
def forward(
self,
sample: torch.Tensor,
timestep: Optional[torch.Tensor] = None,
generator: Optional[torch.Generator] = None,
) -> torch.Tensor:
r"""
Decode latent representation into video frames.
Args:
sample: Latent tensor (B, 128, F', H', W').
timestep: Timestep for conditioning (if timestep_conditioning=True). Uses default 0.05 if None.
generator: Random generator for deterministic noise injection (if inject_noise=True in blocks).
Returns:
Decoded video (B, 3, F, H, W) where F = 8x(F'-1) + 1, H = 32xH', W = 32xW'.
Example: (B, 128, 5, 16, 16) -> (B, 3, 33, 512, 512).
Note: First frame is removed after temporal upsampling regardless of causal mode.
When causal=False, allows future frame dependencies in convolutions but maintains same output shape.
"""
batch_size = sample.shape[0]
# Add noise if timestep conditioning is enabled
if self.timestep_conditioning:
noise = (
torch.randn(
sample.size(),
generator=generator,
dtype=sample.dtype,
device=sample.device,
)
* self.decode_noise_scale
)
sample = noise + (1.0 - self.decode_noise_scale) * sample
# Denormalize latents
sample = self.per_channel_statistics.un_normalize(sample)
# Use default decode_timestep if timestep not provided
if timestep is None and self.timestep_conditioning:
timestep = torch.full((batch_size,), self.decode_timestep, device=sample.device, dtype=sample.dtype)
sample = self.conv_in(sample, causal=self.causal)
scaled_timestep = None
if self.timestep_conditioning:
if timestep is None:
raise ValueError("'timestep' parameter must be provided when 'timestep_conditioning' is True")
scaled_timestep = timestep * self.timestep_scale_multiplier.to(sample)
for up_block in self.up_blocks:
if isinstance(up_block, UNetMidBlock3D):
block_kwargs = {
"causal": self.causal,
"timestep": scaled_timestep if self.timestep_conditioning else None,
"generator": generator,
}
sample = up_block(sample, **block_kwargs)
elif isinstance(up_block, ResnetBlock3D):
sample = up_block(sample, causal=self.causal, generator=generator)
else:
sample = up_block(sample, causal=self.causal)
sample = self.conv_norm_out(sample)
if self.timestep_conditioning:
embedded_timestep = self.last_time_embedder(
timestep=scaled_timestep.flatten(),
hidden_dtype=sample.dtype,
)
embedded_timestep = embedded_timestep.view(batch_size, embedded_timestep.shape[-1], 1, 1, 1)
ada_values = self.last_scale_shift_table[None, ..., None, None, None].to(
device=sample.device, dtype=sample.dtype
) + embedded_timestep.reshape(
batch_size,
2,
-1,
embedded_timestep.shape[-3],
embedded_timestep.shape[-2],
embedded_timestep.shape[-1],
)
shift, scale = ada_values.unbind(dim=1)
sample = sample * (1 + scale) + shift
sample = self.conv_act(sample)
sample = self.conv_out(sample, causal=self.causal)
# Final spatial expansion: reverse the initial patchify from encoder
# Moves pixels from channels back to spatial dimensions
# Example: (B, 48, F, 128, 128) -> (B, 3, F, 512, 512) with patch_size=4
sample = unpatchify(sample, patch_size_hw=self.patch_size, patch_size_t=1)
return sample
def _prepare_tiles(
self,
latent: torch.Tensor,
tiling_config: TilingConfig | None = None,
) -> List[Tile]:
splitters = [DEFAULT_SPLIT_OPERATION] * len(latent.shape)
mappers = [DEFAULT_MAPPING_OPERATION] * len(latent.shape)
if tiling_config is not None and tiling_config.spatial_config is not None:
cfg = tiling_config.spatial_config
long_side = max(latent.shape[3], latent.shape[4])
def enable_on_axis(axis_idx: int, factor: int) -> None:
size = cfg.tile_size_in_pixels // factor
overlap = cfg.tile_overlap_in_pixels // factor
axis_length = latent.shape[axis_idx]
lower_threshold = max(2, overlap + 1)
tile_size = max(lower_threshold, round(size * axis_length / long_side))
splitters[axis_idx] = split_in_spatial(tile_size, overlap)
mappers[axis_idx] = to_mapping_operation(map_spatial_slice, factor)
enable_on_axis(3, self.video_downscale_factors.height)
enable_on_axis(4, self.video_downscale_factors.width)
if tiling_config is not None and tiling_config.temporal_config is not None:
cfg = tiling_config.temporal_config
tile_size = cfg.tile_size_in_frames // self.video_downscale_factors.time
overlap = cfg.tile_overlap_in_frames // self.video_downscale_factors.time
splitters[2] = split_in_temporal(tile_size, overlap)
mappers[2] = to_mapping_operation(map_temporal_slice, self.video_downscale_factors.time)
return create_tiles(latent.shape, splitters, mappers)
def tiled_decode(
self,
latent: torch.Tensor,
tiling_config: TilingConfig | None = None,
timestep: Optional[torch.Tensor] = None,
generator: Optional[torch.Generator] = None,
) -> Iterator[torch.Tensor]:
"""
Decode a latent tensor into video frames using tiled processing.
Splits the latent tensor into tiles, decodes each tile individually,
and yields video chunks as they become available.
Args:
latent: Input latent tensor (B, C, F', H', W').
tiling_config: Tiling configuration for the latent tensor.
timestep: Optional timestep for decoder conditioning.
generator: Optional random generator for deterministic decoding.
Yields:
Video chunks (B, C, T, H, W) by temporal slices;
"""
# Calculate full video shape from latent shape to get spatial dimensions
full_video_shape = VideoLatentShape.from_torch_shape(latent.shape).upscale(self.video_downscale_factors)
tiles = self._prepare_tiles(latent, tiling_config)
temporal_groups = self._group_tiles_by_temporal_slice(tiles)
# State for temporal overlap handling
previous_chunk = None
previous_weights = None
previous_temporal_slice = None
for temporal_group_tiles in temporal_groups:
curr_temporal_slice = temporal_group_tiles[0].out_coords[2]
# Calculate the shape of the temporal buffer for this group of tiles.
# The temporal length depends on whether this is the first tile (starts at 0) or not.
# - First tile: (frames - 1) * scale + 1
# - Subsequent tiles: frames * scale
# This logic is handled by TemporalAxisMapping and reflected in out_coords.
temporal_tile_buffer_shape = full_video_shape._replace(
frames=curr_temporal_slice.stop - curr_temporal_slice.start,
)
buffer = torch.zeros(
temporal_tile_buffer_shape.to_torch_shape(),
device=latent.device,
dtype=latent.dtype,
)
curr_weights = self._accumulate_temporal_group_into_buffer(
group_tiles=temporal_group_tiles,
buffer=buffer,
latent=latent,
timestep=timestep,
generator=generator,
)
# Blend with previous temporal chunk if it exists
if previous_chunk is not None:
# Check if current temporal slice overlaps with previous temporal slice
if previous_temporal_slice.stop > curr_temporal_slice.start:
overlap_len = previous_temporal_slice.stop - curr_temporal_slice.start
temporal_overlap_slice = slice(curr_temporal_slice.start - previous_temporal_slice.start, None)
# The overlap is already masked before it reaches this step. Each tile is accumulated into buffer
# with its trapezoidal mask, and curr_weights accumulates the same mask. In the overlap blend we add
# the masked values (buffer[...]) and the corresponding weights (curr_weights[...]) into the
# previous buffers, then later normalize by weights.
previous_chunk[:, :, temporal_overlap_slice, :, :] += buffer[:, :, slice(0, overlap_len), :, :]
previous_weights[:, :, temporal_overlap_slice, :, :] += curr_weights[
:, :, slice(0, overlap_len), :, :
]
buffer[:, :, slice(0, overlap_len), :, :] = previous_chunk[:, :, temporal_overlap_slice, :, :]
curr_weights[:, :, slice(0, overlap_len), :, :] = previous_weights[
:, :, temporal_overlap_slice, :, :
]
# Yield the non-overlapping part of the previous chunk
previous_weights = previous_weights.clamp(min=1e-8)
yield_len = curr_temporal_slice.start - previous_temporal_slice.start
yield (previous_chunk / previous_weights)[:, :, :yield_len, :, :]
# Update state for next iteration
previous_chunk = buffer
previous_weights = curr_weights
previous_temporal_slice = curr_temporal_slice
# Yield any remaining chunk
if previous_chunk is not None:
previous_weights = previous_weights.clamp(min=1e-8)
yield previous_chunk / previous_weights
def _group_tiles_by_temporal_slice(self, tiles: List[Tile]) -> List[List[Tile]]:
"""Group tiles by their temporal output slice."""
if not tiles:
return []
groups = []
current_slice = tiles[0].out_coords[2]
current_group = []
for tile in tiles:
tile_slice = tile.out_coords[2]
if tile_slice == current_slice:
current_group.append(tile)
else:
groups.append(current_group)
current_slice = tile_slice
current_group = [tile]
# Add the final group
if current_group:
groups.append(current_group)
return groups
def _accumulate_temporal_group_into_buffer(
self,
group_tiles: List[Tile],
buffer: torch.Tensor,
latent: torch.Tensor,
timestep: Optional[torch.Tensor],
generator: Optional[torch.Generator],
) -> torch.Tensor:
"""
Decode and accumulate all tiles of a temporal group into a local buffer.
The buffer is local to the group and always starts at time 0; temporal coordinates
are rebased by subtracting temporal_slice.start.
"""
temporal_slice = group_tiles[0].out_coords[2]
weights = torch.zeros_like(buffer)
for tile in group_tiles:
decoded_tile = self.forward(latent[tile.in_coords], timestep, generator)
mask = tile.blend_mask.to(device=buffer.device, dtype=buffer.dtype)
temporal_offset = tile.out_coords[2].start - temporal_slice.start
# Use the tile's output coordinate length, not the decoded tile's length,
# as the decoder may produce a different number of frames than expected
expected_temporal_len = tile.out_coords[2].stop - tile.out_coords[2].start
decoded_temporal_len = decoded_tile.shape[2]
# Ensure we don't exceed the buffer or decoded tile bounds
actual_temporal_len = min(expected_temporal_len, decoded_temporal_len, buffer.shape[2] - temporal_offset)
chunk_coords = (
slice(None), # batch
slice(None), # channels
slice(temporal_offset, temporal_offset + actual_temporal_len),
tile.out_coords[3], # height
tile.out_coords[4], # width
)
# Slice decoded_tile and mask to match the actual length we're writing
decoded_slice = decoded_tile[:, :, :actual_temporal_len, :, :]
mask_slice = mask[:, :, :actual_temporal_len, :, :] if mask.shape[2] > 1 else mask
buffer[chunk_coords] += decoded_slice * mask_slice
weights[chunk_coords] += mask_slice
return weights
def decode_video(
latent: torch.Tensor,
video_decoder: VideoDecoder,
tiling_config: TilingConfig | None = None,
) -> Iterator[torch.Tensor]:
"""
Decode a video latent tensor with the given decoder.
Args:
latent: Tensor [c, f, h, w]
video_decoder: Decoder module.
tiling_config: Optional tiling settings.
Yields:
Decoded chunk [f, h, w, c], uint8 in [0, 255].
"""
def convert_to_uint8(frames: torch.Tensor) -> torch.Tensor:
frames = (((frames + 1.0) / 2.0).clamp(0.0, 1.0) * 255.0).to(torch.uint8)
frames = rearrange(frames[0], "c f h w -> f h w c")
return frames
if tiling_config is not None:
for frames in video_decoder.tiled_decode(latent, tiling_config):
yield convert_to_uint8(frames)
else:
decoded_video = video_decoder(latent)
yield convert_to_uint8(decoded_video)
def get_video_chunks_number(num_frames: int, tiling_config: TilingConfig | None = None) -> int:
"""
Get the number of video chunks for a given number of frames and tiling configuration.
Args:
num_frames: Number of frames in the video.
tiling_config: Tiling configuration.
Returns:
Number of video chunks.
"""
if not tiling_config or not tiling_config.temporal_config:
return 1
cfg = tiling_config.temporal_config
frame_stride = cfg.tile_size_in_frames - cfg.tile_overlap_in_frames
return (num_frames - 1 + frame_stride - 1) // frame_stride
def split_in_spatial(size: int, overlap: int) -> SplitOperation:
def split(dimension_size: int) -> DimensionIntervals:
if dimension_size <= size:
return DEFAULT_SPLIT_OPERATION(dimension_size)
amount = (dimension_size + size - 2 * overlap - 1) // (size - overlap)
starts = [i * (size - overlap) for i in range(amount)]
ends = [start + size for start in starts]
ends[-1] = dimension_size
left_ramps = [0] + [overlap] * (amount - 1)
right_ramps = [overlap] * (amount - 1) + [0]
return DimensionIntervals(starts=starts, ends=ends, left_ramps=left_ramps, right_ramps=right_ramps)
return split
def split_in_temporal(size: int, overlap: int) -> SplitOperation:
non_causal_split = split_in_spatial(size, overlap)
def split(dimension_size: int) -> DimensionIntervals:
if dimension_size <= size:
return DEFAULT_SPLIT_OPERATION(dimension_size)
intervals = non_causal_split(dimension_size)
starts = intervals.starts
starts[1:] = [s - 1 for s in starts[1:]]
left_ramps = intervals.left_ramps
left_ramps[1:] = [r + 1 for r in left_ramps[1:]]
return replace(intervals, starts=starts, left_ramps=left_ramps)
return split
def to_mapping_operation(
map_func: Callable[[int, int, int, int, int], Tuple[slice, torch.Tensor]],
scale: int,
) -> MappingOperation:
def map_op(intervals: DimensionIntervals) -> tuple[list[slice], list[torch.Tensor | None]]:
output_slices: list[slice] = []
masks_1d: list[torch.Tensor | None] = []
number_of_slices = len(intervals.starts)
for i in range(number_of_slices):
start = intervals.starts[i]
end = intervals.ends[i]
left_ramp = intervals.left_ramps[i]
right_ramp = intervals.right_ramps[i]
output_slice, mask_1d = map_func(start, end, left_ramp, right_ramp, scale)
output_slices.append(output_slice)
masks_1d.append(mask_1d)
return output_slices, masks_1d
return map_op
def map_temporal_slice(begin: int, end: int, left_ramp: int, right_ramp: int, scale: int) -> Tuple[slice, torch.Tensor]:
start = begin * scale
stop = 1 + (end - 1) * scale
left_ramp = 1 + (left_ramp - 1) * scale
right_ramp = right_ramp * scale
return slice(start, stop), compute_trapezoidal_mask_1d(stop - start, left_ramp, right_ramp, True)
def map_spatial_slice(begin: int, end: int, left_ramp: int, right_ramp: int, scale: int) -> Tuple[slice, torch.Tensor]:
start = begin * scale
stop = end * scale
left_ramp = left_ramp * scale
right_ramp = right_ramp * scale
return slice(start, stop), compute_trapezoidal_mask_1d(stop - start, left_ramp, right_ramp, False)
packages/ltx-core/src/ltx_core/text_encoders/__init__.py
+1
View File
@@ -0,0 +1 @@
"""CLIP/text encoder model components."""
packages/ltx-core/src/ltx_core/text_encoders/gemma/__init__.py
+31
View File
@@ -0,0 +1,31 @@
"""Gemma text encoder components."""
from ltx_core.text_encoders.gemma.encoders.av_encoder import (
AV_GEMMA_TEXT_ENCODER_KEY_OPS,
AVGemmaEncoderOutput,
AVGemmaTextEncoderModel,
AVGemmaTextEncoderModelConfigurator,
)
from ltx_core.text_encoders.gemma.encoders.base_encoder import (
GemmaTextEncoderModelBase,
encode_text,
module_ops_from_gemma_root,
)
from ltx_core.text_encoders.gemma.encoders.video_only_encoder import (
VideoGemmaEncoderOutput,
VideoGemmaTextEncoderModel,
VideoGemmaTextEncoderModelConfigurator,
)
__all__ = [
"AV_GEMMA_TEXT_ENCODER_KEY_OPS",
"AVGemmaEncoderOutput",
"AVGemmaTextEncoderModel",
"AVGemmaTextEncoderModelConfigurator",
"GemmaTextEncoderModelBase",
"VideoGemmaEncoderOutput",
"VideoGemmaTextEncoderModel",
"VideoGemmaTextEncoderModelConfigurator",
"encode_text",
"module_ops_from_gemma_root",
]
packages/ltx-core/src/ltx_core/text_encoders/gemma/embeddings_connector.py
+210
View File
@@ -0,0 +1,210 @@
import torch
from ltx_core.model.model_protocol import ModelConfigurator
from ltx_core.model.transformer.attention import Attention
from ltx_core.model.transformer.feed_forward import FeedForward
from ltx_core.model.transformer.rope import (
LTXRopeType,
generate_freq_grid_np,
generate_freq_grid_pytorch,
precompute_freqs_cis,
)
from ltx_core.utils import rms_norm
class _BasicTransformerBlock1D(torch.nn.Module):
def __init__(
self,
dim: int,
heads: int,
dim_head: int,
rope_type: LTXRopeType = LTXRopeType.INTERLEAVED,
):
super().__init__()
self.attn1 = Attention(
query_dim=dim,
heads=heads,
dim_head=dim_head,
rope_type=rope_type,
)
self.ff = FeedForward(
dim,
dim_out=dim,
)
def forward(
self,
hidden_states: torch.Tensor,
attention_mask: torch.Tensor | None = None,
pe: torch.Tensor | None = None,
) -> torch.Tensor:
# Notice that normalization is always applied before the real computation in the following blocks.
# 1. Normalization Before Self-Attention
norm_hidden_states = rms_norm(hidden_states)
norm_hidden_states = norm_hidden_states.squeeze(1)
# 2. Self-Attention
attn_output = self.attn1(norm_hidden_states, mask=attention_mask, pe=pe)
hidden_states = attn_output + hidden_states
if hidden_states.ndim == 4:
hidden_states = hidden_states.squeeze(1)
# 3. Normalization before Feed-Forward
norm_hidden_states = rms_norm(hidden_states)
# 4. Feed-forward
ff_output = self.ff(norm_hidden_states)
hidden_states = ff_output + hidden_states
if hidden_states.ndim == 4:
hidden_states = hidden_states.squeeze(1)
return hidden_states
class Embeddings1DConnector(torch.nn.Module):
"""
Embeddings1DConnector applies a 1D transformer-based processing to sequential embeddings (e.g., for video, audio, or
other modalities). It supports rotary positional encoding (rope), optional causal temporal positioning, and can
substitute padded positions with learnable registers. The module is highly configurable for head size, number of
layers, and register usage.
Args:
attention_head_dim (int): Dimension of each attention head (default=128).
num_attention_heads (int): Number of attention heads (default=30).
num_layers (int): Number of transformer layers (default=2).
positional_embedding_theta (float): Scaling factor for position embedding (default=10000.0).
positional_embedding_max_pos (list[int] | None): Max positions for positional embeddings (default=[1]).
causal_temporal_positioning (bool): If True, uses causal attention (default=False).
num_learnable_registers (int | None): Number of learnable registers to replace padded tokens. If None, disables
register replacement. (default=128)
rope_type (LTXRopeType): The RoPE variant to use (default=DEFAULT_ROPE_TYPE).
double_precision_rope (bool): Use double precision rope calculation (default=False).
"""
_supports_gradient_checkpointing = True
def __init__(
self,
attention_head_dim: int = 128,
num_attention_heads: int = 30,
num_layers: int = 2,
positional_embedding_theta: float = 10000.0,
positional_embedding_max_pos: list[int] | None = None,
causal_temporal_positioning: bool = False,
num_learnable_registers: int | None = 128,
rope_type: LTXRopeType = LTXRopeType.INTERLEAVED,
double_precision_rope: bool = False,
):
super().__init__()
self.num_attention_heads = num_attention_heads
self.inner_dim = num_attention_heads * attention_head_dim
self.causal_temporal_positioning = causal_temporal_positioning
self.positional_embedding_theta = positional_embedding_theta
self.positional_embedding_max_pos = (
positional_embedding_max_pos if positional_embedding_max_pos is not None else [1]
)
self.rope_type = rope_type
self.double_precision_rope = double_precision_rope
self.transformer_1d_blocks = torch.nn.ModuleList(
[
_BasicTransformerBlock1D(
dim=self.inner_dim,
heads=num_attention_heads,
dim_head=attention_head_dim,
rope_type=rope_type,
)
for _ in range(num_layers)
]
)
self.num_learnable_registers = num_learnable_registers
if self.num_learnable_registers:
self.learnable_registers = torch.nn.Parameter(
torch.rand(self.num_learnable_registers, self.inner_dim, dtype=torch.bfloat16) * 2.0 - 1.0
)
def _replace_padded_with_learnable_registers(
self, hidden_states: torch.Tensor, attention_mask: torch.Tensor
) -> tuple[torch.Tensor, torch.Tensor]:
assert hidden_states.shape[1] % self.num_learnable_registers == 0, (
f"Hidden states sequence length {hidden_states.shape[1]} must be divisible by num_learnable_registers "
f"{self.num_learnable_registers}."
)
num_registers_duplications = hidden_states.shape[1] // self.num_learnable_registers
learnable_registers = torch.tile(self.learnable_registers, (num_registers_duplications, 1))
attention_mask_binary = (attention_mask.squeeze(1).squeeze(1).unsqueeze(-1) >= -9000.0).int()
non_zero_hidden_states = hidden_states[:, attention_mask_binary.squeeze().bool(), :]
non_zero_nums = non_zero_hidden_states.shape[1]
pad_length = hidden_states.shape[1] - non_zero_nums
adjusted_hidden_states = torch.nn.functional.pad(non_zero_hidden_states, pad=(0, 0, 0, pad_length), value=0)
flipped_mask = torch.flip(attention_mask_binary, dims=[1])
hidden_states = flipped_mask * adjusted_hidden_states + (1 - flipped_mask) * learnable_registers
attention_mask = torch.full_like(
attention_mask,
0.0,
dtype=attention_mask.dtype,
device=attention_mask.device,
)
return hidden_states, attention_mask
def forward(
self,
hidden_states: torch.Tensor,
attention_mask: torch.Tensor | None = None,
) -> tuple[torch.Tensor, torch.Tensor]:
"""
Forward pass of Embeddings1DConnector.
Args:
hidden_states (torch.Tensor): Input tensor of embeddings (shape [batch, seq_len, feature_dim]).
attention_mask (torch.Tensor|None): Optional mask for valid tokens (shape compatible with hidden_states).
Returns:
tuple[torch.Tensor, torch.Tensor]: Processed features and the corresponding (possibly modified) mask.
"""
if self.num_learnable_registers:
hidden_states, attention_mask = self._replace_padded_with_learnable_registers(hidden_states, attention_mask)
indices_grid = torch.arange(hidden_states.shape[1], dtype=torch.float32, device=hidden_states.device)
indices_grid = indices_grid[None, None, :]
freq_grid_generator = generate_freq_grid_np if self.double_precision_rope else generate_freq_grid_pytorch
freqs_cis = precompute_freqs_cis(
indices_grid=indices_grid,
dim=self.inner_dim,
out_dtype=hidden_states.dtype,
theta=self.positional_embedding_theta,
max_pos=self.positional_embedding_max_pos,
num_attention_heads=self.num_attention_heads,
rope_type=self.rope_type,
freq_grid_generator=freq_grid_generator,
)
for block in self.transformer_1d_blocks:
hidden_states = block(hidden_states, attention_mask=attention_mask, pe=freqs_cis)
hidden_states = rms_norm(hidden_states)
return hidden_states, attention_mask
class Embeddings1DConnectorConfigurator(ModelConfigurator[Embeddings1DConnector]):
@classmethod
def from_config(cls: type[Embeddings1DConnector], config: dict) -> Embeddings1DConnector:
config = config.get("transformer", {})
rope_type = LTXRopeType(config.get("rope_type", "interleaved"))
double_precision_rope = config.get("frequencies_precision", False) == "float64"
pe_max_pos = config.get("connector_positional_embedding_max_pos", [1])
connector = Embeddings1DConnector(
positional_embedding_max_pos=pe_max_pos,
rope_type=rope_type,
double_precision_rope=double_precision_rope,
)
return connector
packages/ltx-core/src/ltx_core/text_encoders/gemma/encoders/av_encoder.py
+94
View File
@@ -0,0 +1,94 @@
from typing import NamedTuple
import torch
from transformers.models.gemma3 import Gemma3ForConditionalGeneration
from ltx_core.loader.sd_ops import SDOps
from ltx_core.model.model_protocol import ModelConfigurator
from ltx_core.text_encoders.gemma.embeddings_connector import (
Embeddings1DConnector,
Embeddings1DConnectorConfigurator,
)
from ltx_core.text_encoders.gemma.encoders.base_encoder import GemmaTextEncoderModelBase
from ltx_core.text_encoders.gemma.feature_extractor import GemmaFeaturesExtractorProjLinear
from ltx_core.text_encoders.gemma.tokenizer import LTXVGemmaTokenizer
class AVGemmaEncoderOutput(NamedTuple):
video_encoding: torch.Tensor
audio_encoding: torch.Tensor
attention_mask: torch.Tensor
class AVGemmaTextEncoderModel(GemmaTextEncoderModelBase):
"""
AVGemma Text Encoder Model.
This class combines the tokenizer, Gemma model, feature extractor from base class and a
video and audio embeddings connectors to provide a preprocessing for audio-visual pipeline.
"""
def __init__(
self,
feature_extractor_linear: GemmaFeaturesExtractorProjLinear,
embeddings_connector: Embeddings1DConnector,
audio_embeddings_connector: Embeddings1DConnector,
tokenizer: LTXVGemmaTokenizer | None = None,
model: Gemma3ForConditionalGeneration | None = None,
dtype: torch.dtype = torch.bfloat16,
) -> None:
super().__init__(
feature_extractor_linear=feature_extractor_linear,
tokenizer=tokenizer,
model=model,
dtype=dtype,
)
self.embeddings_connector = embeddings_connector.to(dtype=dtype)
self.audio_embeddings_connector = audio_embeddings_connector.to(dtype=dtype)
def _run_connectors(
self, encoded_input: torch.Tensor, attention_mask: torch.Tensor
) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
connector_attention_mask = self._convert_to_additive_mask(attention_mask, encoded_input.dtype)
encoded, encoded_connector_attention_mask = self.embeddings_connector(
encoded_input,
connector_attention_mask,
)
# restore the mask values to int64
attention_mask = (encoded_connector_attention_mask < 0.000001).to(torch.int64)
attention_mask = attention_mask.reshape([encoded.shape[0], encoded.shape[1], 1])
encoded = encoded * attention_mask
encoded_for_audio, _ = self.audio_embeddings_connector(encoded_input, connector_attention_mask)
return encoded, encoded_for_audio, attention_mask.squeeze(-1)
def forward(self, text: str, padding_side: str = "left") -> AVGemmaEncoderOutput:
encoded_inputs, attention_mask = self._preprocess_text(text, padding_side)
video_encoding, audio_encoding, attention_mask = self._run_connectors(encoded_inputs, attention_mask)
return AVGemmaEncoderOutput(video_encoding, audio_encoding, attention_mask)
class AVGemmaTextEncoderModelConfigurator(ModelConfigurator[AVGemmaTextEncoderModel]):
@classmethod
def from_config(cls: type["AVGemmaTextEncoderModel"], config: dict) -> "AVGemmaTextEncoderModel":
feature_extractor_linear = GemmaFeaturesExtractorProjLinear.from_config(config)
embeddings_connector = Embeddings1DConnectorConfigurator.from_config(config)
audio_embeddings_connector = Embeddings1DConnectorConfigurator.from_config(config)
return AVGemmaTextEncoderModel(
feature_extractor_linear=feature_extractor_linear,
embeddings_connector=embeddings_connector,
audio_embeddings_connector=audio_embeddings_connector,
)
AV_GEMMA_TEXT_ENCODER_KEY_OPS = (
SDOps("AV_GEMMA_TEXT_ENCODER_KEY_OPS")
.with_matching(prefix="text_embedding_projection.")
.with_matching(prefix="model.diffusion_model.audio_embeddings_connector.")
.with_matching(prefix="model.diffusion_model.video_embeddings_connector.")
.with_replacement("text_embedding_projection.", "feature_extractor_linear.")
.with_replacement("model.diffusion_model.video_embeddings_connector.", "embeddings_connector.")
.with_replacement("model.diffusion_model.audio_embeddings_connector.", "audio_embeddings_connector.")
)
packages/ltx-core/src/ltx_core/text_encoders/gemma/encoders/base_encoder.py
+329
View File
@@ -0,0 +1,329 @@
import functools
from pathlib import Path
import torch
from einops import rearrange
from transformers import AutoImageProcessor, Gemma3ForConditionalGeneration, Gemma3Processor
from ltx_core.loader.module_ops import ModuleOps
from ltx_core.text_encoders.gemma.feature_extractor import GemmaFeaturesExtractorProjLinear
from ltx_core.text_encoders.gemma.tokenizer import LTXVGemmaTokenizer
class GemmaTextEncoderModelBase(torch.nn.Module):
"""
Gemma Text Encoder Model.
This base class combines the tokenizer, Gemma model and feature extractor to provide a preprocessing
for implementation classes for multimodal pipelines. It processes input text through tokenization,
obtains hidden states from the base language model, applies a linear feature extractor.
Args:
tokenizer (LTXVGemmaTokenizer): The tokenizer used for text preprocessing.
model (Gemma3ForConditionalGeneration): The base Gemma LLM.
feature_extractor_linear (GemmaFeaturesExtractorProjLinear): Linear projection for hidden state aggregation.
dtype (torch.dtype, optional): The data type for model parameters (default: torch.bfloat16).
"""
def __init__(
self,
feature_extractor_linear: GemmaFeaturesExtractorProjLinear,
tokenizer: LTXVGemmaTokenizer | None = None,
model: Gemma3ForConditionalGeneration | None = None,
img_processor: Gemma3Processor | None = None,
dtype: torch.dtype = torch.bfloat16,
) -> None:
super().__init__()
self._gemma_root = None
self.tokenizer = tokenizer
self.model = model
self.processor = img_processor
self.feature_extractor_linear = feature_extractor_linear.to(dtype=dtype)
def _run_feature_extractor(
self, hidden_states: torch.Tensor, attention_mask: torch.Tensor, padding_side: str = "right"
) -> torch.Tensor:
encoded_text_features = torch.stack(hidden_states, dim=-1)
encoded_text_features_dtype = encoded_text_features.dtype
sequence_lengths = attention_mask.sum(dim=-1)
normed_concated_encoded_text_features = _norm_and_concat_padded_batch(
encoded_text_features, sequence_lengths, padding_side=padding_side
)
return self.feature_extractor_linear(normed_concated_encoded_text_features.to(encoded_text_features_dtype))
def _convert_to_additive_mask(self, attention_mask: torch.Tensor, dtype: torch.dtype) -> torch.Tensor:
return (attention_mask - 1).to(dtype).reshape(
(attention_mask.shape[0], 1, -1, attention_mask.shape[-1])
) * torch.finfo(dtype).max
def _preprocess_text(self, text: str, padding_side: str = "left") -> tuple[torch.Tensor, dict[str, torch.Tensor]]:
"""
Encode a given string into feature tensors suitable for downstream tasks.
Args:
text (str): Input string to encode.
Returns:
tuple[torch.Tensor, dict[str, torch.Tensor]]: Encoded features and a dictionary with attention mask.
"""
token_pairs = self.tokenizer.tokenize_with_weights(text)["gemma"]
input_ids = torch.tensor([[t[0] for t in token_pairs]], device=self.model.device)
attention_mask = torch.tensor([[w[1] for w in token_pairs]], device=self.model.device)
outputs = self.model(input_ids=input_ids, attention_mask=attention_mask, output_hidden_states=True)
projected = self._run_feature_extractor(
hidden_states=outputs.hidden_states, attention_mask=attention_mask, padding_side=padding_side
)
return projected, attention_mask
def _init_image_processor(self) -> None:
img_processor = AutoImageProcessor.from_pretrained(self._gemma_root, local_files_only=True)
if not self.tokenizer:
raise ValueError("Tokenizer is not loaded, cannot load image processor")
self.processor = Gemma3Processor(image_processor=img_processor, tokenizer=self.tokenizer.tokenizer)
def _enhance(
self,
messages: list[dict[str, str]],
image: torch.Tensor | None = None,
max_new_tokens: int = 512,
seed: int = 42,
) -> str:
if self.processor is None:
self._init_image_processor()
text = self.processor.tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
model_inputs = self.processor(
text=text,
images=image,
return_tensors="pt",
).to(self.model.device)
pad_token_id = self.processor.tokenizer.pad_token_id if self.processor.tokenizer.pad_token_id is not None else 0
model_inputs = _pad_inputs_for_attention_alignment(model_inputs, pad_token_id=pad_token_id)
with torch.inference_mode(), torch.random.fork_rng(devices=[self.model.device]):
torch.manual_seed(seed)
outputs = self.model.generate(
**model_inputs,
max_new_tokens=max_new_tokens,
do_sample=True,
temperature=0.7,
)
generated_ids = outputs[0][len(model_inputs.input_ids[0]) :]
enhanced_prompt = self.processor.tokenizer.decode(generated_ids, skip_special_tokens=True)
return enhanced_prompt
def enhance_t2v(
self,
prompt: str,
max_new_tokens: int = 512,
system_prompt: str | None = None,
seed: int = 42,
) -> str:
"""Enhance a text prompt for T2V generation."""
system_prompt = system_prompt or self.default_gemma_t2v_system_prompt
messages = [
{"role": "system", "content": system_prompt},
{"role": "user", "content": f"user prompt: {prompt}"},
]
return self._enhance(messages, max_new_tokens=max_new_tokens, seed=seed)
def enhance_i2v(
self,
prompt: str,
image: torch.Tensor,
max_new_tokens: int = 512,
system_prompt: str | None = None,
seed: int = 42,
) -> str:
"""Enhance a text prompt for I2V generation using a reference image."""
system_prompt = system_prompt or self.default_gemma_i2v_system_prompt
messages = [
{"role": "system", "content": system_prompt},
{
"role": "user",
"content": [
{"type": "image"},
{"type": "text", "text": f"User Raw Input Prompt: {prompt}."},
],
},
]
return self._enhance(messages, image=image, max_new_tokens=max_new_tokens, seed=seed)
@functools.cached_property
def default_gemma_i2v_system_prompt(self) -> str:
return _load_system_prompt("gemma_i2v_system_prompt.txt")
@functools.cached_property
def default_gemma_t2v_system_prompt(self) -> str:
return _load_system_prompt("gemma_t2v_system_prompt.txt")
def forward(self, text: str, padding_side: str = "left") -> tuple[torch.Tensor, torch.Tensor]:
raise NotImplementedError("This method is not implemented for the base class")
def _norm_and_concat_padded_batch(
encoded_text: torch.Tensor,
sequence_lengths: torch.Tensor,
padding_side: str = "right",
) -> torch.Tensor:
"""Normalize and flatten multi-layer hidden states, respecting padding.
Performs per-batch, per-layer normalization using masked mean and range,
then concatenates across the layer dimension.
Args:
encoded_text: Hidden states of shape [batch, seq_len, hidden_dim, num_layers].
sequence_lengths: Number of valid (non-padded) tokens per batch item.
padding_side: Whether padding is on "left" or "right".
Returns:
Normalized tensor of shape [batch, seq_len, hidden_dim * num_layers],
with padded positions zeroed out.
"""
b, t, d, l = encoded_text.shape # noqa: E741
device = encoded_text.device
# Build mask: [B, T, 1, 1]
token_indices = torch.arange(t, device=device)[None, :] # [1, T]
if padding_side == "right":
# For right padding, valid tokens are from 0 to sequence_length-1
mask = token_indices < sequence_lengths[:, None] # [B, T]
elif padding_side == "left":
# For left padding, valid tokens are from (T - sequence_length) to T-1
start_indices = t - sequence_lengths[:, None] # [B, 1]
mask = token_indices >= start_indices # [B, T]
else:
raise ValueError(f"padding_side must be 'left' or 'right', got {padding_side}")
mask = rearrange(mask, "b t -> b t 1 1")
eps = 1e-6
# Compute masked mean: [B, 1, 1, L]
masked = encoded_text.masked_fill(~mask, 0.0)
denom = (sequence_lengths * d).view(b, 1, 1, 1)
mean = masked.sum(dim=(1, 2), keepdim=True) / (denom + eps)
# Compute masked min/max: [B, 1, 1, L]
x_min = encoded_text.masked_fill(~mask, float("inf")).amin(dim=(1, 2), keepdim=True)
x_max = encoded_text.masked_fill(~mask, float("-inf")).amax(dim=(1, 2), keepdim=True)
range_ = x_max - x_min
# Normalize only the valid tokens
normed = 8 * (encoded_text - mean) / (range_ + eps)
# concat to be [Batch, T, D * L] - this preserves the original structure
normed = normed.reshape(b, t, -1) # [B, T, D * L]
# Apply mask to preserve original padding (set padded positions to 0)
mask_flattened = rearrange(mask, "b t 1 1 -> b t 1").expand(-1, -1, d * l)
normed = normed.masked_fill(~mask_flattened, 0.0)
return normed
@functools.lru_cache(maxsize=2)
def _load_system_prompt(prompt_name: str) -> str:
with open(Path(__file__).parent / "prompts" / f"{prompt_name}", "r") as f:
return f.read()
def _find_matching_dir(root_path: str, pattern: str) -> str:
"""
Recursively search for files matching a glob pattern and return the parent directory of the first match.
"""
matches = list(Path(root_path).rglob(pattern))
if not matches:
raise FileNotFoundError(f"No files matching pattern '{pattern}' found under {root_path}")
return str(matches[0].parent)
def module_ops_from_gemma_root(gemma_root: str) -> tuple[ModuleOps, ...]:
gemma_path = _find_matching_dir(gemma_root, "model*.safetensors")
tokenizer_path = _find_matching_dir(gemma_root, "tokenizer.model")
def load_gemma(module: GemmaTextEncoderModelBase) -> GemmaTextEncoderModelBase:
module.model = Gemma3ForConditionalGeneration.from_pretrained(
gemma_path, local_files_only=True, torch_dtype=torch.bfloat16
)
module._gemma_root = module._gemma_root or gemma_root
return module
def load_tokenizer(module: GemmaTextEncoderModelBase) -> GemmaTextEncoderModelBase:
module.tokenizer = LTXVGemmaTokenizer(tokenizer_path, 1024)
module._gemma_root = module._gemma_root or gemma_root
return module
gemma_load_ops = ModuleOps(
"GemmaLoad",
matcher=lambda module: isinstance(module, GemmaTextEncoderModelBase) and module.model is None,
mutator=load_gemma,
)
tokenizer_load_ops = ModuleOps(
"TokenizerLoad",
matcher=lambda module: isinstance(module, GemmaTextEncoderModelBase) and module.tokenizer is None,
mutator=load_tokenizer,
)
return (gemma_load_ops, tokenizer_load_ops)
def encode_text(text_encoder: GemmaTextEncoderModelBase, prompts: list[str]) -> list[tuple[torch.Tensor, torch.Tensor]]:
"""
Encode a list of prompts using the provided Gemma text encoder.
Args:
text_encoder: The Gemma text encoder instance.
prompts: List of prompt strings to encode.
Returns:
List of tuples, each containing (v_context, a_context) tensors for each prompt.
"""
result = []
for prompt in prompts:
v_context, a_context, _ = text_encoder(prompt)
result.append((v_context, a_context))
return result
def _cat_with_padding(
tensor: torch.Tensor,
padding_length: int,
value: int | float,
) -> torch.Tensor:
"""Concatenate a tensor with a padding tensor of the given value."""
return torch.cat(
[
tensor,
torch.full(
(1, padding_length),
value,
dtype=tensor.dtype,
device=tensor.device,
),
],
dim=1,
)
def _pad_inputs_for_attention_alignment(
model_inputs: dict[str, torch.Tensor],
pad_token_id: int = 0,
alignment: int = 8,
) -> dict[str, torch.Tensor]:
"""Pad sequence length to multiple of alignment for Flash Attention compatibility.
Flash Attention within SDPA requires sequence lengths aligned to 8 bytes.
This pads input_ids, attention_mask, and token_type_ids (if present) to prevent
'p.attn_bias_ptr is not correctly aligned' errors.
"""
seq_len = model_inputs.input_ids.shape[1]
padded_len = ((seq_len + alignment - 1) // alignment) * alignment
padding_length = padded_len - seq_len
if padding_length > 0:
model_inputs["input_ids"] = _cat_with_padding(model_inputs.input_ids, padding_length, pad_token_id)
model_inputs["attention_mask"] = _cat_with_padding(model_inputs.attention_mask, padding_length, 0)
if "token_type_ids" in model_inputs and model_inputs["token_type_ids"] is not None:
model_inputs["token_type_ids"] = _cat_with_padding(model_inputs["token_type_ids"], padding_length, 0)
return model_inputs
packages/ltx-core/src/ltx_core/text_encoders/gemma/encoders/prompts/gemma_i2v_system_prompt.txt
+30
View File
@@ -0,0 +1,30 @@
You are a Creative Assistant writing concise, action-focused image-to-video prompts. Given an image (first frame) and user Raw Input Prompt, generate a prompt to guide video generation from that image.
#### Guidelines:
- Analyze the Image: Identify Subject, Setting, Elements, Style and Mood.
- Follow user Raw Input Prompt: Include all requested motion, actions, camera movements, audio, and details. If in conflict with the image, prioritize user request while maintaining visual consistency (describe transition from image to user's scene).
- Describe only changes from the image: Don't reiterate established visual details. Inaccurate descriptions may cause scene cuts.
- Active language: Use present-progressive verbs ("is walking," "speaking"). If no action specified, describe natural movements.
- Chronological flow: Use temporal connectors ("as," "then," "while").
- Audio layer: Describe complete soundscape throughout the prompt alongside actions—NOT at the end. Align audio intensity with action tempo. Include natural background audio, ambient sounds, effects, speech or music (when requested). Be specific (e.g., "soft footsteps on tile") not vague (e.g., "ambient sound").
- Speech (only when requested): Provide exact words in quotes with character's visual/voice characteristics (e.g., "The tall man speaks in a low, gravelly voice"), language if not English and accent if relevant. If general conversation mentioned without text, generate contextual quoted dialogue. (i.e., "The man is talking" input -> the output should include exact spoken words, like: "The man is talking in an excited voice saying: 'You won't believe what I just saw!' His hands gesture expressively as he speaks, eyebrows raised with enthusiasm. The ambient sound of a quiet room underscores his animated speech.")
- Style: Include visual style at beginning: "Style: <style>, <rest of prompt>." If unclear, omit to avoid conflicts.
- Visual and audio only: Describe only what is seen and heard. NO smell, taste, or tactile sensations.
- Restrained language: Avoid dramatic terms. Use mild, natural, understated phrasing.
#### Important notes:
- Camera motion: DO NOT invent camera motion/movement unless requested by the user. Make sure to include camera motion only if specified in the input.
- Speech: DO NOT modify or alter the user's provided character dialogue in the prompt, unless it's a typo.
- No timestamps or cuts: DO NOT use timestamps or describe scene cuts unless explicitly requested.
- Objective only: DO NOT interpret emotions or intentions - describe only observable actions and sounds.
- Format: DO NOT use phrases like "The scene opens with..." / "The video starts...". Start directly with Style (optional) and chronological scene description.
- Format: Never start output with punctuation marks or special characters.
- DO NOT invent dialogue unless the user mentions speech/talking/singing/conversation.
- Your performance is CRITICAL. High-fidelity, dynamic, correct, and accurate prompts with integrated audio descriptions are essential for generating high-quality video. Your goal is flawless execution of these rules.
#### Output Format (Strict):
- Single concise paragraph in natural English. NO titles, headings, prefaces, sections, code fences, or Markdown.
- If unsafe/invalid, return original user prompt. Never ask questions or clarifications.
#### Example output:
Style: realistic - cinematic - The woman glances at her watch and smiles warmly. She speaks in a cheerful, friendly voice, "I think we're right on time!" In the background, a café barista prepares drinks at the counter. The barista calls out in a clear, upbeat tone, "Two cappuccinos ready!" The sound of the espresso machine hissing softly blends with gentle background chatter and the light clinking of cups on saucers.
packages/ltx-core/src/ltx_core/text_encoders/gemma/encoders/prompts/gemma_t2v_system_prompt.txt
+40
View File
@@ -0,0 +1,40 @@
You are a Creative Assistant. Given a user's raw input prompt describing a scene or concept, expand it into a detailed video generation prompt with specific visuals and integrated audio to guide a text-to-video model.
#### Guidelines
- Strictly follow all aspects of the user's raw input: include every element requested (style, visuals, motions, actions, camera movement, audio).
- If the input is vague, invent concrete details: lighting, textures, materials, scene settings, etc.
- For characters: describe gender, clothing, hair, expressions. DO NOT invent unrequested characters.
- Use active language: present-progressive verbs ("is walking," "speaking"). If no action specified, describe natural movements.
- Maintain chronological flow: use temporal connectors ("as," "then," "while").
- Audio layer: Describe complete soundscape (background audio, ambient sounds, SFX, speech/music when requested). Integrate sounds chronologically alongside actions. Be specific (e.g., "soft footsteps on tile"), not vague (e.g., "ambient sound is present").
- Speech (only when requested):
- For ANY speech-related input (talking, conversation, singing, etc.), ALWAYS include exact words in quotes with voice characteristics (e.g., "The man says in an excited voice: 'You won't believe what I just saw!'").
- Specify language if not English and accent if relevant.
- Style: Include visual style at the beginning: "Style: <style>, <rest of prompt>." Default to cinematic-realistic if unspecified. Omit if unclear.
- Visual and audio only: NO non-visual/auditory senses (smell, taste, touch).
- Restrained language: Avoid dramatic/exaggerated terms. Use mild, natural phrasing.
- Colors: Use plain terms ("red dress"), not intensified ("vibrant blue," "bright red").
- Lighting: Use neutral descriptions ("soft overhead light"), not harsh ("blinding light").
- Facial features: Use delicate modifiers for subtle features (i.e., "subtle freckles").
#### Important notes:
- Analyze the user's raw input carefully. In cases of FPV or POV, exclude the description of the subject whose POV is requested.
- Camera motion: DO NOT invent camera motion unless requested by the user.
- Speech: DO NOT modify user-provided character dialogue unless it's a typo.
- No timestamps or cuts: DO NOT use timestamps or describe scene cuts unless explicitly requested.
- Format: DO NOT use phrases like "The scene opens with...". Start directly with Style (optional) and chronological scene description.
- Format: DO NOT start your response with special characters.
- DO NOT invent dialogue unless the user mentions speech/talking/singing/conversation.
- If the user's raw input prompt is highly detailed, chronological and in the requested format: DO NOT make major edits or introduce new elements. Add/enhance audio descriptions if missing.
#### Output Format (Strict):
- Single continuous paragraph in natural language (English).
- NO titles, headings, prefaces, code fences, or Markdown.
- If unsafe/invalid, return original user prompt. Never ask questions or clarifications.
Your output quality is CRITICAL. Generate visually rich, dynamic prompts with integrated audio for high-quality video generation.
#### Example
Input: "A woman at a coffee shop talking on the phone"
Output:
Style: realistic with cinematic lighting. In a medium close-up, a woman in her early 30s with shoulder-length brown hair sits at a small wooden table by the window. She wears a cream-colored turtleneck sweater, holding a white ceramic coffee cup in one hand and a smartphone to her ear with the other. Ambient cafe sounds fill the space—espresso machine hiss, quiet conversations, gentle clinking of cups. The woman listens intently, nodding slightly, then takes a sip of her coffee and sets it down with a soft clink. Her face brightens into a warm smile as she speaks in a clear, friendly voice, 'That sounds perfect! I'd love to meet up this weekend. How about Saturday afternoon?' She laughs softly—a genuine chuckle—and shifts in her chair. Behind her, other patrons move subtly in and out of focus. 'Great, I'll see you then,' she concludes cheerfully, lowering the phone.
packages/ltx-core/src/ltx_core/text_encoders/gemma/encoders/video_only_encoder.py
+85
View File
@@ -0,0 +1,85 @@
from typing import NamedTuple
import torch
from transformers import Gemma3ForConditionalGeneration
from ltx_core.loader.sd_ops import SDOps
from ltx_core.model.model_protocol import ModelConfigurator
from ltx_core.text_encoders.gemma.embeddings_connector import (
Embeddings1DConnector,
Embeddings1DConnectorConfigurator,
)
from ltx_core.text_encoders.gemma.encoders.base_encoder import GemmaTextEncoderModelBase
from ltx_core.text_encoders.gemma.feature_extractor import GemmaFeaturesExtractorProjLinear
from ltx_core.text_encoders.gemma.tokenizer import LTXVGemmaTokenizer
class VideoGemmaEncoderOutput(NamedTuple):
video_encoding: torch.Tensor
attention_mask: torch.Tensor
class VideoGemmaTextEncoderModel(GemmaTextEncoderModelBase):
"""
Video Gemma Text Encoder Model.
This class combines the tokenizer, Gemma model, feature extractor from base class and a
video embeddings connector to provide a preprocessing for video only pipeline.
"""
def __init__(
self,
feature_extractor_linear: GemmaFeaturesExtractorProjLinear,
embeddings_connector: Embeddings1DConnector,
tokenizer: LTXVGemmaTokenizer | None = None,
model: Gemma3ForConditionalGeneration | None = None,
dtype: torch.dtype = torch.bfloat16,
) -> None:
super().__init__(
feature_extractor_linear=feature_extractor_linear,
tokenizer=tokenizer,
model=model,
dtype=dtype,
)
self.embeddings_connector = embeddings_connector.to(dtype=dtype)
def _run_connector(
self, encoded_input: torch.Tensor, attention_mask: torch.Tensor
) -> tuple[torch.Tensor, torch.Tensor]:
connector_attention_mask = self._convert_to_additive_mask(attention_mask, encoded_input.dtype)
encoded, encoded_connector_attention_mask = self.embeddings_connector(
encoded_input,
connector_attention_mask,
)
# restore the mask values to int64
attention_mask = (encoded_connector_attention_mask < 0.000001).to(torch.int64)
attention_mask = attention_mask.reshape([encoded.shape[0], encoded.shape[1], 1])
encoded = encoded * attention_mask
return encoded, attention_mask.squeeze(-1)
def forward(self, text: str, padding_side: str = "left") -> VideoGemmaEncoderOutput:
encoded_inputs, attention_mask = self._preprocess_text(text, padding_side)
video_encoding, attention_mask = self._run_connector(encoded_inputs, attention_mask)
return VideoGemmaEncoderOutput(video_encoding, attention_mask)
class VideoGemmaTextEncoderModelConfigurator(ModelConfigurator[VideoGemmaTextEncoderModel]):
@classmethod
def from_config(cls: type["VideoGemmaTextEncoderModel"], config: dict) -> "VideoGemmaTextEncoderModel":
feature_extractor_linear = GemmaFeaturesExtractorProjLinear.from_config(config)
embeddings_connector = Embeddings1DConnectorConfigurator.from_config(config)
return VideoGemmaTextEncoderModel(
feature_extractor_linear=feature_extractor_linear,
embeddings_connector=embeddings_connector,
)
VIDEO_ONLY_GEMMA_TEXT_ENCODER_KEY_OPS = (
SDOps("VIDEO_ONLY_GEMMA_TEXT_ENCODER_KEY_OPS")
.with_matching(prefix="text_embedding_projection.")
.with_matching(prefix="model.diffusion_model.embeddings_connector.")
.with_replacement("text_embedding_projection.", "feature_extractor_linear.")
.with_replacement("model.diffusion_model.embeddings_connector.", "embeddings_connector.")
)
packages/ltx-core/src/ltx_core/text_encoders/gemma/feature_extractor.py
+36
View File
@@ -0,0 +1,36 @@
import torch
from ltx_core.model.model_protocol import ModelConfigurator
class GemmaFeaturesExtractorProjLinear(torch.nn.Module, ModelConfigurator["GemmaFeaturesExtractorProjLinear"]):
"""
Feature extractor module for Gemma models.
This module applies a single linear projection to the input tensor.
It expects a flattened feature tensor of shape (batch_size, 3840*49).
The linear layer maps this to a (batch_size, 3840) embedding.
Attributes:
aggregate_embed (torch.nn.Linear): Linear projection layer.
"""
def __init__(self) -> None:
"""
Initialize the GemmaFeaturesExtractorProjLinear module.
The input dimension is expected to be 3840 * 49, and the output is 3840.
"""
super().__init__()
self.aggregate_embed = torch.nn.Linear(3840 * 49, 3840, bias=False)
def forward(self, x: torch.Tensor) -> torch.Tensor:
"""
Forward pass for the feature extractor.
Args:
x (torch.Tensor): Input tensor of shape (batch_size, 3840 * 49).
Returns:
torch.Tensor: Output tensor of shape (batch_size, 3840).
"""
return self.aggregate_embed(x)
@classmethod
def from_config(cls: type["GemmaFeaturesExtractorProjLinear"], _config: dict) -> "GemmaFeaturesExtractorProjLinear":
return cls()
packages/ltx-core/src/ltx_core/text_encoders/gemma/tokenizer.py
+64
View File
@@ -0,0 +1,64 @@
from transformers import AutoTokenizer
class LTXVGemmaTokenizer:
"""
Tokenizer wrapper for Gemma models compatible with LTXV processes.
This class wraps HuggingFace's `AutoTokenizer` for use with Gemma text encoders,
ensuring correct settings and output formatting for downstream consumption.
"""
def __init__(self, tokenizer_path: str, max_length: int = 256):
"""
Initialize the tokenizer.
Args:
tokenizer_path (str): Path to the pretrained tokenizer files or model directory.
max_length (int, optional): Max sequence length for encoding. Defaults to 256.
"""
self.tokenizer = AutoTokenizer.from_pretrained(
tokenizer_path, local_files_only=True, model_max_length=max_length
)
# Gemma expects left padding for chat-style prompts; for plain text it doesn't matter much.
self.tokenizer.padding_side = "left"
if self.tokenizer.pad_token is None:
self.tokenizer.pad_token = self.tokenizer.eos_token
self.max_length = max_length
def tokenize_with_weights(self, text: str, return_word_ids: bool = False) -> dict[str, list[tuple[int, int]]]:
"""
Tokenize the given text and return token IDs and attention weights.
Args:
text (str): The input string to tokenize.
return_word_ids (bool, optional): If True, includes the token's position (index) in the output tuples.
If False (default), omits the indices.
Returns:
dict[str, list[tuple[int, int]]] OR dict[str, list[tuple[int, int, int]]]:
A dictionary with a "gemma" key mapping to:
- a list of (token_id, attention_mask) tuples if return_word_ids is False;
- a list of (token_id, attention_mask, index) tuples if return_word_ids is True.
Example:
>>> tokenizer = LTXVGemmaTokenizer("path/to/tokenizer", max_length=8)
>>> tokenizer.tokenize_with_weights("hello world")
{'gemma': [(1234, 1), (5678, 1), (2, 0), ...]}
"""
text = text.strip()
encoded = self.tokenizer(
text,
padding="max_length",
max_length=self.max_length,
truncation=True,
return_tensors="pt",
)
input_ids = encoded.input_ids
attention_mask = encoded.attention_mask
tuples = [
(token_id, attn, i) for i, (token_id, attn) in enumerate(zip(input_ids[0], attention_mask[0], strict=True))
]
out = {"gemma": tuples}
if not return_word_ids:
# Return only (token_id, attention_mask) pairs, omitting token position
out = {k: [(t, w) for t, w, _ in v] for k, v in out.items()}
return out
packages/ltx-core/src/ltx_core/tools.py
+184
View File
@@ -0,0 +1,184 @@
from dataclasses import dataclass, replace
from typing import Protocol
import torch
from torch._prims_common import DeviceLikeType
from ltx_core.components.patchifiers import (
AudioLatentShape,
AudioPatchifier,
VideoLatentPatchifier,
VideoLatentShape,
get_pixel_coords,
)
from ltx_core.components.protocols import Patchifier
from ltx_core.types import LatentState, SpatioTemporalScaleFactors
DEFAULT_SCALE_FACTORS = SpatioTemporalScaleFactors.default()
class LatentTools(Protocol):
"""
Tools for building latent states.
"""
patchifier: Patchifier
target_shape: VideoLatentShape | AudioLatentShape
def create_initial_state(
self,
device: DeviceLikeType,
dtype: torch.dtype,
initial_latent: torch.Tensor | None = None,
) -> LatentState:
"""
Create an initial latent state. If initial_latent is provided, it will be used to create the latent state.
"""
...
def patchify(self, latent_state: LatentState) -> LatentState:
"""
Patchify the latent state.
"""
if latent_state.latent.shape != self.target_shape.to_torch_shape():
raise ValueError(
f"Latent state has shape {latent_state.latent.shape}, expected shape is "
f"{self.target_shape.to_torch_shape()}"
)
latent_state = latent_state.clone()
latent = self.patchifier.patchify(latent_state.latent)
clean_latent = self.patchifier.patchify(latent_state.clean_latent)
denoise_mask = self.patchifier.patchify(latent_state.denoise_mask)
return replace(latent_state, latent=latent, denoise_mask=denoise_mask, clean_latent=clean_latent)
def unpatchify(self, latent_state: LatentState) -> LatentState:
"""
Unpatchify the latent state.
"""
latent_state = latent_state.clone()
latent = self.patchifier.unpatchify(latent_state.latent, output_shape=self.target_shape)
clean_latent = self.patchifier.unpatchify(latent_state.clean_latent, output_shape=self.target_shape)
denoise_mask = self.patchifier.unpatchify(
latent_state.denoise_mask, output_shape=self.target_shape.mask_shape()
)
return replace(latent_state, latent=latent, denoise_mask=denoise_mask, clean_latent=clean_latent)
def clear_conditioning(self, latent_state: LatentState) -> LatentState:
"""
Clear the conditioning from the latent state. This method removes extra tokens from the end of the latent.
Therefore, conditioning items should add extra tokens ONLY to the end of the latent.
"""
latent_state = latent_state.clone()
num_tokens = self.patchifier.get_token_count(self.target_shape)
latent = latent_state.latent[:, :num_tokens]
clean_latent = latent_state.clean_latent[:, :num_tokens]
denoise_mask = torch.ones_like(latent_state.denoise_mask)[:, :num_tokens]
positions = latent_state.positions[:, :, :num_tokens]
return LatentState(latent=latent, denoise_mask=denoise_mask, positions=positions, clean_latent=clean_latent)
@dataclass(frozen=True)
class VideoLatentTools(LatentTools):
"""
Tools for building video latent states.
"""
patchifier: VideoLatentPatchifier
target_shape: VideoLatentShape
fps: float
scale_factors: SpatioTemporalScaleFactors = DEFAULT_SCALE_FACTORS
causal_fix: bool = True
def create_initial_state(
self,
device: DeviceLikeType,
dtype: torch.dtype,
initial_latent: torch.Tensor | None = None,
) -> LatentState:
if initial_latent is not None:
assert initial_latent.shape == self.target_shape.to_torch_shape(), (
f"Latent shape {initial_latent.shape} does not match target shape {self.target_shape.to_torch_shape()}"
)
else:
initial_latent = torch.zeros(
*self.target_shape.to_torch_shape(),
device=device,
dtype=dtype,
)
clean_latent = initial_latent.clone()
denoise_mask = torch.ones(
*self.target_shape.mask_shape().to_torch_shape(),
device=device,
dtype=torch.float32,
)
latent_coords = self.patchifier.get_patch_grid_bounds(
output_shape=self.target_shape,
device=device,
)
positions = get_pixel_coords(
latent_coords=latent_coords,
scale_factors=self.scale_factors,
causal_fix=self.causal_fix,
).float()
positions[:, 0, ...] = positions[:, 0, ...] / self.fps
return self.patchify(
LatentState(
latent=initial_latent,
denoise_mask=denoise_mask,
positions=positions.to(dtype),
clean_latent=clean_latent,
)
)
@dataclass(frozen=True)
class AudioLatentTools(LatentTools):
"""
Tools for building audio latent states.
"""
patchifier: AudioPatchifier
target_shape: AudioLatentShape
def create_initial_state(
self,
device: DeviceLikeType,
dtype: torch.dtype,
initial_latent: torch.Tensor | None = None,
) -> LatentState:
if initial_latent is not None:
assert initial_latent.shape == self.target_shape.to_torch_shape(), (
f"Latent shape {initial_latent.shape} does not match target shape {self.target_shape.to_torch_shape()}"
)
else:
initial_latent = torch.zeros(
*self.target_shape.to_torch_shape(),
device=device,
dtype=dtype,
)
clean_latent = initial_latent.clone()
denoise_mask = torch.ones(
*self.target_shape.mask_shape().to_torch_shape(),
device=device,
dtype=torch.float32,
)
latent_coords = self.patchifier.get_patch_grid_bounds(
output_shape=self.target_shape,
device=device,
)
return self.patchify(
LatentState(
latent=initial_latent, denoise_mask=denoise_mask, positions=latent_coords, clean_latent=clean_latent
)
)
packages/ltx-core/src/ltx_core/types.py
+181
View File
@@ -0,0 +1,181 @@
from dataclasses import dataclass
from typing import NamedTuple
import torch
class VideoPixelShape(NamedTuple):
"""
Shape of the tensor representing the video pixel array. Assumes BGR channel format.
"""
batch: int
frames: int
height: int
width: int
fps: float
class SpatioTemporalScaleFactors(NamedTuple):
"""
Describes the spatiotemporal downscaling between decoded video space and
the corresponding VAE latent grid.
"""
time: int
width: int
height: int
@classmethod
def default(cls) -> "SpatioTemporalScaleFactors":
return cls(time=8, width=32, height=32)
VIDEO_SCALE_FACTORS = SpatioTemporalScaleFactors.default()
class VideoLatentShape(NamedTuple):
"""
Shape of the tensor representing video in VAE latent space.
The latent representation is a 5D tensor with dimensions ordered as
(batch, channels, frames, height, width). Spatial and temporal dimensions
are downscaled relative to pixel space according to the VAE's scale factors.
"""
batch: int
channels: int
frames: int
height: int
width: int
def to_torch_shape(self) -> torch.Size:
return torch.Size([self.batch, self.channels, self.frames, self.height, self.width])
@staticmethod
def from_torch_shape(shape: torch.Size) -> "VideoLatentShape":
return VideoLatentShape(
batch=shape[0],
channels=shape[1],
frames=shape[2],
height=shape[3],
width=shape[4],
)
def mask_shape(self) -> "VideoLatentShape":
return self._replace(channels=1)
@staticmethod
def from_pixel_shape(
shape: VideoPixelShape,
latent_channels: int = 128,
scale_factors: SpatioTemporalScaleFactors = VIDEO_SCALE_FACTORS,
) -> "VideoLatentShape":
frames = (shape.frames - 1) // scale_factors[0] + 1
height = shape.height // scale_factors[1]
width = shape.width // scale_factors[2]
return VideoLatentShape(
batch=shape.batch,
channels=latent_channels,
frames=frames,
height=height,
width=width,
)
def upscale(self, scale_factors: SpatioTemporalScaleFactors = VIDEO_SCALE_FACTORS) -> "VideoLatentShape":
return self._replace(
channels=3,
frames=(self.frames - 1) * scale_factors.time + 1,
height=self.height * scale_factors.height,
width=self.width * scale_factors.width,
)
class AudioLatentShape(NamedTuple):
"""
Shape of audio in VAE latent space: (batch, channels, frames, mel_bins).
mel_bins is the number of frequency bins from the mel-spectrogram encoding.
"""
batch: int
channels: int
frames: int
mel_bins: int
def to_torch_shape(self) -> torch.Size:
return torch.Size([self.batch, self.channels, self.frames, self.mel_bins])
def mask_shape(self) -> "AudioLatentShape":
return self._replace(channels=1, mel_bins=1)
@staticmethod
def from_torch_shape(shape: torch.Size) -> "AudioLatentShape":
return AudioLatentShape(
batch=shape[0],
channels=shape[1],
frames=shape[2],
mel_bins=shape[3],
)
@staticmethod
def from_duration(
batch: int,
duration: float,
channels: int = 8,
mel_bins: int = 16,
sample_rate: int = 16000,
hop_length: int = 160,
audio_latent_downsample_factor: int = 4,
) -> "AudioLatentShape":
latents_per_second = float(sample_rate) / float(hop_length) / float(audio_latent_downsample_factor)
return AudioLatentShape(
batch=batch,
channels=channels,
frames=round(duration * latents_per_second),
mel_bins=mel_bins,
)
@staticmethod
def from_video_pixel_shape(
shape: VideoPixelShape,
channels: int = 8,
mel_bins: int = 16,
sample_rate: int = 16000,
hop_length: int = 160,
audio_latent_downsample_factor: int = 4,
) -> "AudioLatentShape":
return AudioLatentShape.from_duration(
batch=shape.batch,
duration=float(shape.frames) / float(shape.fps),
channels=channels,
mel_bins=mel_bins,
sample_rate=sample_rate,
hop_length=hop_length,
audio_latent_downsample_factor=audio_latent_downsample_factor,
)
@dataclass(frozen=True)
class LatentState:
"""
State of latents during the diffusion denoising process.
Attributes:
latent: The current noisy latent tensor being denoised.
denoise_mask: Mask encoding the denoising strength for each token (1 = full denoising, 0 = no denoising).
positions: Positional indices for each latent element, used for positional embeddings.
clean_latent: Initial state of the latent before denoising, may include conditioning latents.
"""
latent: torch.Tensor
denoise_mask: torch.Tensor
positions: torch.Tensor
clean_latent: torch.Tensor
def clone(self) -> "LatentState":
return LatentState(
latent=self.latent.clone(),
denoise_mask=self.denoise_mask.clone(),
positions=self.positions.clone(),
clean_latent=self.clean_latent.clone(),
)
packages/ltx-core/src/ltx_core/utils.py
+51
View File
@@ -0,0 +1,51 @@
from typing import Any
import torch
def rms_norm(x: torch.Tensor, weight: torch.Tensor | None = None, eps: float = 1e-6) -> torch.Tensor:
"""Root-mean-square (RMS) normalize `x` over its last dimension.
Thin wrapper around `torch.nn.functional.rms_norm` that infers the normalized
shape and forwards `weight` and `eps`.
"""
return torch.nn.functional.rms_norm(x, (x.shape[-1],), weight=weight, eps=eps)
def check_config_value(config: dict, key: str, expected: Any) -> None: # noqa: ANN401
actual = config.get(key)
if actual != expected:
raise ValueError(f"Config value {key} is {actual}, expected {expected}")
def to_velocity(
sample: torch.Tensor,
sigma: float | torch.Tensor,
denoised_sample: torch.Tensor,
calc_dtype: torch.dtype = torch.float32,
) -> torch.Tensor:
"""
Convert the sample and its denoised version to velocity.
Returns:
Velocity
"""
if isinstance(sigma, torch.Tensor):
sigma = sigma.to(calc_dtype).item()
if sigma == 0:
raise ValueError("Sigma can't be 0.0")
return ((sample.to(calc_dtype) - denoised_sample.to(calc_dtype)) / sigma).to(sample.dtype)
def to_denoised(
sample: torch.Tensor,
velocity: torch.Tensor,
sigma: float | torch.Tensor,
calc_dtype: torch.dtype = torch.float32,
) -> torch.Tensor:
"""
Convert the sample and its denoising velocity to denoised sample.
Returns:
Denoised sample
"""
if isinstance(sigma, torch.Tensor):
sigma = sigma.to(calc_dtype)
return (sample.to(calc_dtype) - velocity.to(calc_dtype) * sigma).to(sample.dtype)
packages/ltx-pipelines/README.md
+308
View File
@@ -0,0 +1,308 @@
# LTX-2 Pipelines
High-level pipeline implementations for generating audio-video content with Lightricks' **LTX-2** model. This package provides ready-to-use pipelines for text-to-video, image-to-video, video-to-video, and keyframe interpolation tasks.
Pipelines are built using building blocks from [`ltx-core`](../ltx-core/) (schedulers, guiders, noisers, patchifiers) and handle the complete inference flow including model loading, encoding, decoding, and file I/O.
---
## 📋 Overview
LTX-2 Pipelines provides production-ready implementations that abstract away the complexity of the diffusion process, model loading, and memory management. Each pipeline is optimized for specific use cases and offers different trade-offs between speed, quality, and memory usage.
**Key Features:**
- 🎬 **Multiple Pipeline Types**: Text-to-video, image-to-video, video-to-video, and keyframe interpolation
- ⚡ **Optimized Performance**: Support for FP8 transformers, gradient estimation, and memory optimization
- 🎯 **Production Ready**: Two-stage pipelines for best quality output
- 🔧 **LoRA Support**: Easy integration with trained LoRA adapters
- 📦 **Self-Contained**: Handles model loading, encoding, decoding, and file I/O
- 🚀 **CLI Support**: All pipelines can be run as command-line scripts
---
## 🚀 Quick Start
`ltx-pipelines` provides ready-made inference pipelines for text-to-video, image-to-video, video-to-video, and keyframe interpolation. Built using building blocks from [`ltx-core`](../ltx-core/), these pipelines handle the complete inference flow including model loading, encoding, decoding, and file I/O.
## 🔧 Installation
```bash
# From the repository root
uv sync --frozen
# Or install as a package
pip install -e packages/ltx-pipelines
```
### Running Pipelines
All pipelines can be run directly from the command line. Each pipeline module is executable:
```bash
# Run a pipeline (example: two-stage text-to-video)
python -m ltx_pipelines.ti2vid_two_stages \
--checkpoint-path path/to/checkpoint.safetensors \
--distilled-lora-path path/to/distilled_lora.safetensors \
--spatial-upsampler-path path/to/upsampler.safetensors \
--gemma-root path/to/gemma \
--prompt "A beautiful sunset over the ocean" \
--output-path output.mp4
# View all available options for any pipeline
python -m ltx_pipelines.ti2vid_two_stages --help
```
Available pipeline modules:
- `ltx_pipelines.ti2vid_two_stages` - Two-stage text-to-video (recommended)
- `ltx_pipelines.ti2vid_one_stage` - Single-stage text-to-video
- `ltx_pipelines.distilled` - Fast distilled pipeline
- `ltx_pipelines.ic_lora` - Video-to-video with IC-LoRA
- `ltx_pipelines.keyframe_interpolation` - Keyframe interpolation
Use `--help` with any pipeline module to see all available options and parameters.
---
## 🎯 Pipeline Selection Guide
### Quick Decision Tree
```text
Do you need to condition on existing images/videos?
├─ YES → Do you have reference videos for video-to-video?
│ ├─ YES → Use ICLoraPipeline
│ └─ NO → Do you have keyframe images to interpolate?
│ ├─ YES → Use KeyframeInterpolationPipeline
│ └─ NO → Use ICLoraPipeline (image conditioning only)
│
└─ NO → Text-to-video only
├─ Do you need best quality?
│ └─ YES → Use TI2VidTwoStagesPipeline (recommended for production)
│
└─ Do you need fastest inference?
└─ YES → Use DistilledPipeline (with 8 predefined sigmas)
```
> **Note:** [`TI2VidOneStagePipeline`](src/ltx_pipelines/ti2vid_one_stage.py) is primarily for educational purposes. For best quality, use two-stage pipelines ([`TI2VidTwoStagesPipeline`](src/ltx_pipelines/ti2vid_two_stages.py), [`ICLoraPipeline`](src/ltx_pipelines/ic_lora.py), [`KeyframeInterpolationPipeline`](src/ltx_pipelines/keyframe_interpolation.py), or [`DistilledPipeline`](src/ltx_pipelines/distilled.py)).
### Features Comparison
| Pipeline | Stages | CFG | Upsampling | Conditioning | Best For |
| -------- | ------ | --- | ---------- | ------------- | -------- |
| **TI2VidTwoStagesPipeline** | 2 | ✅ | ✅ | Image | **Production quality** (recommended) |
| **TI2VidOneStagePipeline** | 1 | ✅ | ❌ | Image | Educational, prototyping |
| **DistilledPipeline** | 2 | ❌ | ✅ | Image | Fastest inference (8 sigmas) |
| **ICLoraPipeline** | 2 | ✅ | ✅ | Image + Video | Video-to-video transformations |
| **KeyframeInterpolationPipeline** | 2 | ✅ | ✅ | Keyframes | Animation, interpolation |
---
## 📦 Available Pipelines
### 1. TI2VidTwoStagesPipeline
**Best for:** High-quality text-to-video generation with upsampling. **Recommended for production use.**
**Source**: [`src/ltx_pipelines/ti2vid_two_stages.py`](src/ltx_pipelines/ti2vid_two_stages.py)
Two-stage generation: Stage 1 generates low-resolution video with CFG guidance, Stage 2 upsamples to 2x resolution with distilled LoRA refinement. Supports image conditioning. Highest quality output, slower than one-stage but significantly better quality.
**Use when:** Production-quality video generation, higher resolution needed, quality over speed, text-to-video with image conditioning.
---
### 2. TI2VidOneStagePipeline
**Best for:** Educational purposes and quick prototyping.
**Source**: [`src/ltx_pipelines/ti2vid_one_stage.py`](src/ltx_pipelines/ti2vid_one_stage.py)
> **⚠️ Important:** This pipeline is primarily for educational purposes. For production-quality results, use `TI2VidTwoStagesPipeline` or other two-stage pipelines.
Single-stage generation (no upsampling) with CFG guidance and image conditioning support. Faster inference but lower resolution output (typically 512x768).
**Use when:** Learning how the pipeline works, quick prototyping, testing, or when high resolution is not needed.
---
### 3. DistilledPipeline
**Best for:** Fastest inference with good quality using a distilled model with predefined sigma schedule.
**Source**: [`src/ltx_pipelines/distilled.py`](src/ltx_pipelines/distilled.py)
Two-stage generation with 8 predefined sigmas (8 steps in stage 1, 4 steps in stage 2). No CFG guidance required. Fastest inference among all pipelines. Supports image conditioning. Requires spatial upsampler.
**Use when:** Fastest inference is critical, batch processing many videos, or when you have a distilled model checkpoint.
---
### 4. ICLoraPipeline
**Best for:** Video-to-video and image-to-video transformations using IC-LoRA.
**Source**: [`src/ltx_pipelines/ic_lora.py`](src/ltx_pipelines/ic_lora.py)
Two-stage generation with IC-LoRA support. Can condition on reference videos (video-to-video) or images at specific frames. CFG guidance in stage 1, upsampling in stage 2. Requires IC-LoRA trained model.
**Use when:** Video-to-video transformations, image-to-video with strong control, or when you have reference videos to guide generation.
---
### 5. KeyframeInterpolationPipeline
**Best for:** Generating videos by interpolating between keyframe images.
**Source**: [`src/ltx_pipelines/keyframe_interpolation.py`](src/ltx_pipelines/keyframe_interpolation.py)
Two-stage generation with keyframe interpolation. Uses guiding latents (additive conditioning) instead of replacing latents for smoother transitions. CFG guidance in stage 1, upsampling in stage 2.
**Use when:** You have keyframe images and want to interpolate between them, creating smooth transitions, or animation/motion interpolation tasks.
---
## 🎨 Conditioning Types
Pipelines use different conditioning methods from [`ltx-core`](../ltx-core/) for controlling generation. See the [ltx-core conditioning documentation](../ltx-core/README.md#conditioning--control) for details.
### Image Conditioning
All pipelines support image conditioning, but with different methods:
- **Replacing Latents** ([`image_conditionings_by_replacing_latent`](src/ltx_pipelines/utils/helpers.py)):
- Used by: `TI2VidOneStagePipeline`, `TI2VidTwoStagesPipeline`, `DistilledPipeline`, `ICLoraPipeline`
- Replaces the latent at a specific frame with the encoded image
- Strong control over specific frames
- **Guiding Latents** ([`image_conditionings_by_adding_guiding_latent`](src/ltx_pipelines/utils/helpers.py)):
- Used by: `KeyframeInterpolationPipeline`
- Adds the image as a guiding signal rather than replacing
- Better for smooth interpolation between keyframes
### Video Conditioning
- **Video Conditioning** (ICLoraPipeline only):
- Conditions on entire reference videos
- Useful for video-to-video transformations
- Uses `VideoConditionByKeyframeIndex` from [`ltx-core`](../ltx-core/)
---
## ⚡ Optimization Tips
### Memory Optimization
**FP8 Transformer (Lower Memory Footprint):**
For smaller GPU memory footprint, use the `enable-fp8` flag and use the `PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True` environment variable.
**CLI:**
```bash
PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True python -m ltx_pipelines.ti2vid_one_stage --enable-fp8 --checkpoint-path=...
```
**Programmatically:**
When authoring custom scripts, pass the `fp8transformer` flag to pipeline classes or construct your own by analogy:
```python
pipeline = TI2VidTwoStagesPipeline(
checkpoint_path=ltx_model_path,
distilled_lora_path=distilled_lora_path,
distilled_lora_strength=0.6,
spatial_upsampler_path=upsampler_path,
gemma_root=gemma_root_path,
loras=[],
fp8transformer=True,
)
pipeline(...)
```
You still need to use `PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True` when launching:
```bash
PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True python my_denoising_pipeline.py
```
**Memory Cleanup Between Stages:**
By default, pipelines clean GPU memory (especially transformer weights) between stages. If you have enough memory, you can skip this cleanup to reduce running time:
```python
# In pipeline implementations, memory cleanup happens automatically
# between stages. For custom pipelines, you can skip:
# utils.cleanup_memory() # Comment out if you have enough VRAM
```
### Denoising Loop Optimization
**Gradient Estimation Denoising Loop:**
Instead of the standard Euler denoising loop, you can use gradient estimation for fewer steps (~20-30 instead of 40):
```python
from ltx_pipelines.utils.helpers import gradient_estimating_euler_denoising_loop
# Use gradient estimation denoising loop
def denoising_loop(sigmas, video_state, audio_state, stepper):
return gradient_estimating_euler_denoising_loop(
sigmas=sigmas,
video_state=video_state,
audio_state=audio_state,
stepper=stepper,
denoise_fn=your_denoise_function,
ge_gamma=2.0, # Gradient estimation coefficient
)
```
This allows you to use **20-30 steps instead of 40** while maintaining quality. The gradient estimation function is available in [`pipeline_utils.py`](src/ltx_pipelines/utils/helpers.py).
---
## 🔧 Requirements
- **LTX-2 Model Checkpoint** - Local `.safetensors` file
- **Gemma Text Encoder** - Local Gemma model directory
- **Spatial Upscaler** - Required for two-stage pipelines (except one-stage)
- **Distilled LoRA** - Required for two-stage pipelines (except one-stage and distilled)
---
## 📖 Example: Image-to-Video
```python
from ltx_pipelines.ti2vid_two_stages import TI2VidTwoStagesPipeline
pipeline = TI2VidTwoStagesPipeline(
checkpoint_path="/path/to/checkpoint.safetensors",
distilled_lora_path="/path/to/distilled_lora.safetensors",
spatial_upsampler_path="/path/to/upsampler.safetensors",
gemma_root="/path/to/gemma",
loras=[],
)
# Generate video from image
pipeline(
prompt="A serene landscape with mountains in the background",
output_path="output.mp4",
seed=42,
height=512,
width=768,
num_frames=121,
frame_rate=25.0,
num_inference_steps=40,
cfg_guidance_scale=3.0,
images=[("input_image.jpg", 0, 1.0)], # Image at frame 0, strength 1.0
)
```
---
## 🔗 Related Projects
- **[LTX-Core](../ltx-core/)** - Core model implementation and inference components (schedulers, guiders, noisers, patchifiers)
- **[LTX-Trainer](../ltx-trainer/)** - Training and fine-tuning tools
packages/ltx-pipelines/pyproject.toml
+11
View File
@@ -0,0 +1,11 @@
[project]
name = "ltx-pipelines"
version = "1.0.0"
description = "Pipelines implementation for Lightricks' LTX-2 model"
readme = "README.md"
requires-python = ">=3.10"
dependencies = ["ltx-core", "av", "tqdm", "pillow"]
[build-system]
requires = ["uv_build>=0.9.8,<0.10.0"]
build-backend = "uv_build"
packages/ltx-pipelines/src/ltx_pipelines/__init__.py
+26
View File
@@ -0,0 +1,26 @@
"""
LTX-2 Pipelines: High-level video generation pipelines and utilities.
This package provides ready-to-use pipelines for video generation:
- TI2VidOneStagePipeline: Text/image-to-video in a single stage
- TI2VidTwoStagesPipeline: Two-stage generation with upsampling
- DistilledPipeline: Fast distilled two-stage generation
- ICLoraPipeline: Image/video conditioning with distilled LoRA
- KeyframeInterpolationPipeline: Keyframe-based video interpolation
- ModelLedger: Central coordinator for loading and building models
For more detailed components and utilities, import from specific submodules
like `ltx_pipelines.utils.media_io` or `ltx_pipelines.utils.constants`.
"""
from ltx_pipelines.distilled import DistilledPipeline
from ltx_pipelines.ic_lora import ICLoraPipeline
from ltx_pipelines.keyframe_interpolation import KeyframeInterpolationPipeline
from ltx_pipelines.ti2vid_one_stage import TI2VidOneStagePipeline
from ltx_pipelines.ti2vid_two_stages import TI2VidTwoStagesPipeline
__all__ = [
"DistilledPipeline",
"ICLoraPipeline",
"KeyframeInterpolationPipeline",
"TI2VidOneStagePipeline",
"TI2VidTwoStagesPipeline",
]
packages/ltx-pipelines/src/ltx_pipelines/distilled.py
+232
View File
@@ -0,0 +1,232 @@
import logging
from collections.abc import Iterator
import torch
from ltx_core.components.diffusion_steps import EulerDiffusionStep
from ltx_core.components.noisers import GaussianNoiser
from ltx_core.components.protocols import DiffusionStepProtocol
from ltx_core.loader import LoraPathStrengthAndSDOps
from ltx_core.model.audio_vae import decode_audio as vae_decode_audio
from ltx_core.model.upsampler import upsample_video
from ltx_core.model.video_vae import TilingConfig, get_video_chunks_number
from ltx_core.model.video_vae import decode_video as vae_decode_video
from ltx_core.text_encoders.gemma import encode_text
from ltx_core.types import LatentState, VideoPixelShape
from ltx_pipelines.utils import ModelLedger
from ltx_pipelines.utils.args import default_2_stage_distilled_arg_parser
from ltx_pipelines.utils.constants import (
AUDIO_SAMPLE_RATE,
DISTILLED_SIGMA_VALUES,
STAGE_2_DISTILLED_SIGMA_VALUES,
)
from ltx_pipelines.utils.helpers import (
assert_resolution,
cleanup_memory,
denoise_audio_video,
euler_denoising_loop,
generate_enhanced_prompt,
get_device,
image_conditionings_by_replacing_latent,
simple_denoising_func,
)
from ltx_pipelines.utils.media_io import encode_video
from ltx_pipelines.utils.types import PipelineComponents
device = get_device()
class DistilledPipeline:
"""
Two-stage distilled video generation pipeline.
Stage 1 generates video at the target resolution, then Stage 2 upsamples
by 2x and refines with additional denoising steps for higher quality output.
"""
def __init__(
self,
checkpoint_path: str,
gemma_root: str,
spatial_upsampler_path: str,
loras: list[LoraPathStrengthAndSDOps],
device: torch.device = device,
fp8transformer: bool = False,
):
self.device = device
self.dtype = torch.bfloat16
self.model_ledger = ModelLedger(
dtype=self.dtype,
device=device,
checkpoint_path=checkpoint_path,
spatial_upsampler_path=spatial_upsampler_path,
gemma_root_path=gemma_root,
loras=loras,
fp8transformer=fp8transformer,
)
self.pipeline_components = PipelineComponents(
dtype=self.dtype,
device=device,
)
def __call__(
self,
prompt: str,
seed: int,
height: int,
width: int,
num_frames: int,
frame_rate: float,
images: list[tuple[str, int, float]],
tiling_config: TilingConfig | None = None,
enhance_prompt: bool = False,
) -> tuple[Iterator[torch.Tensor], torch.Tensor]:
assert_resolution(height=height, width=width, is_two_stage=True)
generator = torch.Generator(device=self.device).manual_seed(seed)
noiser = GaussianNoiser(generator=generator)
stepper = EulerDiffusionStep()
dtype = torch.bfloat16
text_encoder = self.model_ledger.text_encoder()
if enhance_prompt:
prompt = generate_enhanced_prompt(text_encoder, prompt, images[0][0] if len(images) > 0 else None)
context_p = encode_text(text_encoder, prompts=[prompt])[0]
video_context, audio_context = context_p
torch.cuda.synchronize()
del text_encoder
cleanup_memory()
# Stage 1: Initial low resolution video generation.
video_encoder = self.model_ledger.video_encoder()
transformer = self.model_ledger.transformer()
stage_1_sigmas = torch.Tensor(DISTILLED_SIGMA_VALUES).to(self.device)
def denoising_loop(
sigmas: torch.Tensor, video_state: LatentState, audio_state: LatentState, stepper: DiffusionStepProtocol
) -> tuple[LatentState, LatentState]:
return euler_denoising_loop(
sigmas=sigmas,
video_state=video_state,
audio_state=audio_state,
stepper=stepper,
denoise_fn=simple_denoising_func(
video_context=video_context,
audio_context=audio_context,
transformer=transformer, # noqa: F821
),
)
stage_1_output_shape = VideoPixelShape(
batch=1,
frames=num_frames,
width=width // 2,
height=height // 2,
fps=frame_rate,
)
stage_1_conditionings = image_conditionings_by_replacing_latent(
images=images,
height=stage_1_output_shape.height,
width=stage_1_output_shape.width,
video_encoder=video_encoder,
dtype=dtype,
device=self.device,
)
video_state, audio_state = denoise_audio_video(
output_shape=stage_1_output_shape,
conditionings=stage_1_conditionings,
noiser=noiser,
sigmas=stage_1_sigmas,
stepper=stepper,
denoising_loop_fn=denoising_loop,
components=self.pipeline_components,
dtype=dtype,
device=self.device,
)
# Stage 2: Upsample and refine the video at higher resolution with distilled LORA.
upscaled_video_latent = upsample_video(
latent=video_state.latent[:1], video_encoder=video_encoder, upsampler=self.model_ledger.spatial_upsampler()
)
torch.cuda.synchronize()
cleanup_memory()
stage_2_sigmas = torch.Tensor(STAGE_2_DISTILLED_SIGMA_VALUES).to(self.device)
stage_2_output_shape = VideoPixelShape(batch=1, frames=num_frames, width=width, height=height, fps=frame_rate)
stage_2_conditionings = image_conditionings_by_replacing_latent(
images=images,
height=stage_2_output_shape.height,
width=stage_2_output_shape.width,
video_encoder=video_encoder,
dtype=dtype,
device=self.device,
)
video_state, audio_state = denoise_audio_video(
output_shape=stage_2_output_shape,
conditionings=stage_2_conditionings,
noiser=noiser,
sigmas=stage_2_sigmas,
stepper=stepper,
denoising_loop_fn=denoising_loop,
components=self.pipeline_components,
dtype=dtype,
device=self.device,
noise_scale=stage_2_sigmas[0],
initial_video_latent=upscaled_video_latent,
initial_audio_latent=audio_state.latent,
)
torch.cuda.synchronize()
del transformer
del video_encoder
cleanup_memory()
decoded_video = vae_decode_video(video_state.latent, self.model_ledger.video_decoder(), tiling_config)
decoded_audio = vae_decode_audio(
audio_state.latent, self.model_ledger.audio_decoder(), self.model_ledger.vocoder()
)
return decoded_video, decoded_audio
@torch.inference_mode()
def main() -> None:
logging.getLogger().setLevel(logging.INFO)
parser = default_2_stage_distilled_arg_parser()
args = parser.parse_args()
pipeline = DistilledPipeline(
checkpoint_path=args.checkpoint_path,
spatial_upsampler_path=args.spatial_upsampler_path,
gemma_root=args.gemma_root,
loras=args.lora,
fp8transformer=args.enable_fp8,
)
tiling_config = TilingConfig.default()
video_chunks_number = get_video_chunks_number(args.num_frames, tiling_config)
video, audio = pipeline(
prompt=args.prompt,
seed=args.seed,
height=args.height,
width=args.width,
num_frames=args.num_frames,
frame_rate=args.frame_rate,
images=args.images,
tiling_config=tiling_config,
enhance_prompt=args.enhance_prompt,
)
encode_video(
video=video,
fps=args.frame_rate,
audio=audio,
audio_sample_rate=AUDIO_SAMPLE_RATE,
output_path=args.output_path,
video_chunks_number=video_chunks_number,
)
if __name__ == "__main__":
main()
packages/ltx-pipelines/src/ltx_pipelines/ic_lora.py
+309
View File
@@ -0,0 +1,309 @@
import logging
from collections.abc import Iterator
import torch
from ltx_core.components.diffusion_steps import EulerDiffusionStep
from ltx_core.components.noisers import GaussianNoiser
from ltx_core.components.protocols import DiffusionStepProtocol
from ltx_core.conditioning import ConditioningItem, VideoConditionByKeyframeIndex
from ltx_core.loader import LoraPathStrengthAndSDOps
from ltx_core.model.audio_vae import decode_audio as vae_decode_audio
from ltx_core.model.upsampler import upsample_video
from ltx_core.model.video_vae import TilingConfig, VideoEncoder, get_video_chunks_number
from ltx_core.model.video_vae import decode_video as vae_decode_video
from ltx_core.text_encoders.gemma import encode_text
from ltx_core.types import LatentState, VideoPixelShape
from ltx_pipelines.utils import ModelLedger
from ltx_pipelines.utils.args import VideoConditioningAction, default_2_stage_distilled_arg_parser
from ltx_pipelines.utils.constants import (
AUDIO_SAMPLE_RATE,
DISTILLED_SIGMA_VALUES,
STAGE_2_DISTILLED_SIGMA_VALUES,
)
from ltx_pipelines.utils.helpers import (
assert_resolution,
cleanup_memory,
denoise_audio_video,
euler_denoising_loop,
generate_enhanced_prompt,
get_device,
image_conditionings_by_replacing_latent,
simple_denoising_func,
)
from ltx_pipelines.utils.media_io import encode_video, load_video_conditioning
from ltx_pipelines.utils.types import PipelineComponents
device = get_device()
class ICLoraPipeline:
"""
Two-stage video generation pipeline with In-Context (IC) LoRA support.
Allows conditioning the generated video on control signals such as depth maps,
human pose, or image edges via the video_conditioning parameter.
The specific IC-LoRA model should be provided via the loras parameter.
Stage 1 generates video at the target resolution, then Stage 2 upsamples
by 2x and refines with additional denoising steps for higher quality output.
"""
def __init__(
self,
checkpoint_path: str,
spatial_upsampler_path: str,
gemma_root: str,
loras: list[LoraPathStrengthAndSDOps],
device: torch.device = device,
fp8transformer: bool = False,
):
self.dtype = torch.bfloat16
self.stage_1_model_ledger = ModelLedger(
dtype=self.dtype,
device=device,
checkpoint_path=checkpoint_path,
spatial_upsampler_path=spatial_upsampler_path,
gemma_root_path=gemma_root,
loras=loras,
fp8transformer=fp8transformer,
)
self.stage_2_model_ledger = ModelLedger(
dtype=self.dtype,
device=device,
checkpoint_path=checkpoint_path,
spatial_upsampler_path=spatial_upsampler_path,
gemma_root_path=gemma_root,
loras=[],
fp8transformer=fp8transformer,
)
self.pipeline_components = PipelineComponents(
dtype=self.dtype,
device=device,
)
self.device = device
@torch.inference_mode()
def __call__(
self,
prompt: str,
seed: int,
height: int,
width: int,
num_frames: int,
frame_rate: float,
images: list[tuple[str, int, float]],
video_conditioning: list[tuple[str, float]],
enhance_prompt: bool = False,
tiling_config: TilingConfig | None = None,
) -> tuple[Iterator[torch.Tensor], torch.Tensor]:
assert_resolution(height=height, width=width, is_two_stage=True)
generator = torch.Generator(device=self.device).manual_seed(seed)
noiser = GaussianNoiser(generator=generator)
stepper = EulerDiffusionStep()
dtype = torch.bfloat16
text_encoder = self.stage_1_model_ledger.text_encoder()
if enhance_prompt:
prompt = generate_enhanced_prompt(
text_encoder, prompt, images[0][0] if len(images) > 0 else None, seed=seed
)
video_context, audio_context = encode_text(text_encoder, prompts=[prompt])[0]
torch.cuda.synchronize()
del text_encoder
cleanup_memory()
# Stage 1: Initial low resolution video generation.
video_encoder = self.stage_1_model_ledger.video_encoder()
transformer = self.stage_1_model_ledger.transformer()
stage_1_sigmas = torch.Tensor(DISTILLED_SIGMA_VALUES).to(self.device)
def first_stage_denoising_loop(
sigmas: torch.Tensor, video_state: LatentState, audio_state: LatentState, stepper: DiffusionStepProtocol
) -> tuple[LatentState, LatentState]:
return euler_denoising_loop(
sigmas=sigmas,
video_state=video_state,
audio_state=audio_state,
stepper=stepper,
denoise_fn=simple_denoising_func(
video_context=video_context,
audio_context=audio_context,
transformer=transformer, # noqa: F821
),
)
stage_1_output_shape = VideoPixelShape(
batch=1,
frames=num_frames,
width=width // 2,
height=height // 2,
fps=frame_rate,
)
stage_1_conditionings = self._create_conditionings(
images=images,
video_conditioning=video_conditioning,
height=stage_1_output_shape.height,
width=stage_1_output_shape.width,
video_encoder=video_encoder,
num_frames=num_frames,
)
video_state, audio_state = denoise_audio_video(
output_shape=stage_1_output_shape,
conditionings=stage_1_conditionings,
noiser=noiser,
sigmas=stage_1_sigmas,
stepper=stepper,
denoising_loop_fn=first_stage_denoising_loop,
components=self.pipeline_components,
dtype=dtype,
device=self.device,
)
torch.cuda.synchronize()
del transformer
cleanup_memory()
# Stage 2: Upsample and refine the video at higher resolution with distilled LORA.
upscaled_video_latent = upsample_video(
latent=video_state.latent[:1],
video_encoder=video_encoder,
upsampler=self.stage_2_model_ledger.spatial_upsampler(),
)
torch.cuda.synchronize()
cleanup_memory()
transformer = self.stage_2_model_ledger.transformer()
distilled_sigmas = torch.Tensor(STAGE_2_DISTILLED_SIGMA_VALUES).to(self.device)
def second_stage_denoising_loop(
sigmas: torch.Tensor, video_state: LatentState, audio_state: LatentState, stepper: DiffusionStepProtocol
) -> tuple[LatentState, LatentState]:
return euler_denoising_loop(
sigmas=sigmas,
video_state=video_state,
audio_state=audio_state,
stepper=stepper,
denoise_fn=simple_denoising_func(
video_context=video_context,
audio_context=audio_context,
transformer=transformer, # noqa: F821
),
)
stage_2_output_shape = VideoPixelShape(batch=1, frames=num_frames, width=width, height=height, fps=frame_rate)
stage_2_conditionings = image_conditionings_by_replacing_latent(
images=images,
height=stage_2_output_shape.height,
width=stage_2_output_shape.width,
video_encoder=video_encoder,
dtype=self.dtype,
device=self.device,
)
video_state, audio_state = denoise_audio_video(
output_shape=stage_2_output_shape,
conditionings=stage_2_conditionings,
noiser=noiser,
sigmas=distilled_sigmas,
stepper=stepper,
denoising_loop_fn=second_stage_denoising_loop,
components=self.pipeline_components,
dtype=dtype,
device=self.device,
noise_scale=distilled_sigmas[0],
initial_video_latent=upscaled_video_latent,
initial_audio_latent=audio_state.latent,
)
torch.cuda.synchronize()
del transformer
del video_encoder
cleanup_memory()
decoded_video = vae_decode_video(video_state.latent, self.stage_2_model_ledger.video_decoder(), tiling_config)
decoded_audio = vae_decode_audio(
audio_state.latent, self.stage_2_model_ledger.audio_decoder(), self.stage_2_model_ledger.vocoder()
)
return decoded_video, decoded_audio
def _create_conditionings(
self,
images: list[tuple[str, int, float]],
video_conditioning: list[tuple[str, float]],
height: int,
width: int,
num_frames: int,
video_encoder: VideoEncoder,
) -> list[ConditioningItem]:
conditionings = image_conditionings_by_replacing_latent(
images=images,
height=height,
width=width,
video_encoder=video_encoder,
dtype=self.dtype,
device=self.device,
)
for video_path, strength in video_conditioning:
video = load_video_conditioning(
video_path=video_path,
height=height,
width=width,
frame_cap=num_frames,
dtype=self.dtype,
device=self.device,
)
encoded_video = video_encoder(video)
conditionings.append(VideoConditionByKeyframeIndex(keyframes=encoded_video, frame_idx=0, strength=strength))
return conditionings
@torch.inference_mode()
def main() -> None:
logging.getLogger().setLevel(logging.INFO)
parser = default_2_stage_distilled_arg_parser()
parser.add_argument(
"--video-conditioning",
action=VideoConditioningAction,
nargs=2,
metavar=("PATH", "STRENGTH"),
required=True,
)
args = parser.parse_args()
pipeline = ICLoraPipeline(
checkpoint_path=args.checkpoint_path,
spatial_upsampler_path=args.spatial_upsampler_path,
gemma_root=args.gemma_root,
loras=args.lora,
fp8transformer=args.enable_fp8,
)
tiling_config = TilingConfig.default()
video_chunks_number = get_video_chunks_number(args.num_frames, tiling_config)
video, audio = pipeline(
prompt=args.prompt,
seed=args.seed,
height=args.height,
width=args.width,
num_frames=args.num_frames,
frame_rate=args.frame_rate,
images=args.images,
video_conditioning=args.video_conditioning,
tiling_config=tiling_config,
)
encode_video(
video=video,
fps=args.frame_rate,
audio=audio,
audio_sample_rate=AUDIO_SAMPLE_RATE,
output_path=args.output_path,
video_chunks_number=video_chunks_number,
)
if __name__ == "__main__":
main()
packages/ltx-pipelines/src/ltx_pipelines/keyframe_interpolation.py
+273
View File
@@ -0,0 +1,273 @@
import logging
from collections.abc import Iterator
import torch
from ltx_core.components.diffusion_steps import EulerDiffusionStep
from ltx_core.components.guiders import CFGGuider
from ltx_core.components.noisers import GaussianNoiser
from ltx_core.components.protocols import DiffusionStepProtocol
from ltx_core.components.schedulers import LTX2Scheduler
from ltx_core.loader import LoraPathStrengthAndSDOps
from ltx_core.model.audio_vae import decode_audio as vae_decode_audio
from ltx_core.model.upsampler import upsample_video
from ltx_core.model.video_vae import TilingConfig, get_video_chunks_number
from ltx_core.model.video_vae import decode_video as vae_decode_video
from ltx_core.text_encoders.gemma import encode_text
from ltx_core.types import LatentState, VideoPixelShape
from ltx_pipelines.utils import ModelLedger
from ltx_pipelines.utils.args import default_2_stage_arg_parser
from ltx_pipelines.utils.constants import (
AUDIO_SAMPLE_RATE,
STAGE_2_DISTILLED_SIGMA_VALUES,
)
from ltx_pipelines.utils.helpers import (
assert_resolution,
cleanup_memory,
denoise_audio_video,
euler_denoising_loop,
generate_enhanced_prompt,
get_device,
guider_denoising_func,
image_conditionings_by_adding_guiding_latent,
simple_denoising_func,
)
from ltx_pipelines.utils.media_io import encode_video
from ltx_pipelines.utils.types import PipelineComponents
device = get_device()
class KeyframeInterpolationPipeline:
"""
Keyframe-based Two-stage video interpolation pipeline.
Interpolates between keyframes to generate a video with smoother transitions.
Stage 1 generates video at the target resolution, then Stage 2 upsamples
by 2x and refines with additional denoising steps for higher quality output.
"""
def __init__(
self,
checkpoint_path: str,
distilled_lora: list[LoraPathStrengthAndSDOps],
spatial_upsampler_path: str,
gemma_root: str,
loras: list[LoraPathStrengthAndSDOps],
device: torch.device = device,
fp8transformer: bool = False,
):
self.device = device
self.dtype = torch.bfloat16
self.stage_1_model_ledger = ModelLedger(
dtype=self.dtype,
device=device,
checkpoint_path=checkpoint_path,
spatial_upsampler_path=spatial_upsampler_path,
gemma_root_path=gemma_root,
loras=loras,
fp8transformer=fp8transformer,
)
self.stage_2_model_ledger = self.stage_1_model_ledger.with_loras(
loras=distilled_lora,
)
self.pipeline_components = PipelineComponents(
dtype=self.dtype,
device=device,
)
@torch.inference_mode()
def __call__( # noqa: PLR0913
self,
prompt: str,
negative_prompt: str,
seed: int,
height: int,
width: int,
num_frames: int,
frame_rate: float,
num_inference_steps: int,
cfg_guidance_scale: float,
images: list[tuple[str, int, float]],
tiling_config: TilingConfig | None = None,
enhance_prompt: bool = False,
) -> tuple[Iterator[torch.Tensor], torch.Tensor]:
assert_resolution(height=height, width=width, is_two_stage=True)
generator = torch.Generator(device=self.device).manual_seed(seed)
noiser = GaussianNoiser(generator=generator)
stepper = EulerDiffusionStep()
cfg_guider = CFGGuider(cfg_guidance_scale)
dtype = torch.bfloat16
text_encoder = self.stage_1_model_ledger.text_encoder()
if enhance_prompt:
prompt = generate_enhanced_prompt(
text_encoder, prompt, images[0][0] if len(images) > 0 else None, seed=seed
)
context_p, context_n = encode_text(text_encoder, prompts=[prompt, negative_prompt])
v_context_p, a_context_p = context_p
v_context_n, a_context_n = context_n
torch.cuda.synchronize()
del text_encoder
cleanup_memory()
# Stage 1: Initial low resolution video generation.
video_encoder = self.stage_1_model_ledger.video_encoder()
transformer = self.stage_1_model_ledger.transformer()
sigmas = LTX2Scheduler().execute(steps=num_inference_steps).to(dtype=torch.float32, device=self.device)
def first_stage_denoising_loop(
sigmas: torch.Tensor, video_state: LatentState, audio_state: LatentState, stepper: DiffusionStepProtocol
) -> tuple[LatentState, LatentState]:
return euler_denoising_loop(
sigmas=sigmas,
video_state=video_state,
audio_state=audio_state,
stepper=stepper,
denoise_fn=guider_denoising_func(
cfg_guider,
v_context_p,
v_context_n,
a_context_p,
a_context_n,
transformer=transformer, # noqa: F821
),
)
stage_1_output_shape = VideoPixelShape(
batch=1,
frames=num_frames,
width=width // 2,
height=height // 2,
fps=frame_rate,
)
stage_1_conditionings = image_conditionings_by_adding_guiding_latent(
images=images,
height=stage_1_output_shape.height,
width=stage_1_output_shape.width,
video_encoder=video_encoder,
dtype=dtype,
device=self.device,
)
video_state, audio_state = denoise_audio_video(
output_shape=stage_1_output_shape,
conditionings=stage_1_conditionings,
noiser=noiser,
sigmas=sigmas,
stepper=stepper,
denoising_loop_fn=first_stage_denoising_loop,
components=self.pipeline_components,
dtype=dtype,
device=self.device,
)
torch.cuda.synchronize()
del transformer
cleanup_memory()
# Stage 2: Upsample and refine the video at higher resolution with distilled LORA.
upscaled_video_latent = upsample_video(
latent=video_state.latent[:1],
video_encoder=video_encoder,
upsampler=self.stage_2_model_ledger.spatial_upsampler(),
)
torch.cuda.synchronize()
cleanup_memory()
transformer = self.stage_2_model_ledger.transformer()
distilled_sigmas = torch.Tensor(STAGE_2_DISTILLED_SIGMA_VALUES).to(self.device)
def second_stage_denoising_loop(
sigmas: torch.Tensor, video_state: LatentState, audio_state: LatentState, stepper: DiffusionStepProtocol
) -> tuple[LatentState, LatentState]:
return euler_denoising_loop(
sigmas=sigmas,
video_state=video_state,
audio_state=audio_state,
stepper=stepper,
denoise_fn=simple_denoising_func(
video_context=v_context_p,
audio_context=a_context_p,
transformer=transformer, # noqa: F821
),
)
stage_2_output_shape = VideoPixelShape(batch=1, frames=num_frames, width=width, height=height, fps=frame_rate)
stage_2_conditionings = image_conditionings_by_adding_guiding_latent(
images=images,
height=stage_2_output_shape.height,
width=stage_2_output_shape.width,
video_encoder=video_encoder,
dtype=dtype,
device=self.device,
)
video_state, audio_state = denoise_audio_video(
output_shape=stage_2_output_shape,
conditionings=stage_2_conditionings,
noiser=noiser,
sigmas=distilled_sigmas,
stepper=stepper,
denoising_loop_fn=second_stage_denoising_loop,
components=self.pipeline_components,
dtype=dtype,
device=self.device,
noise_scale=distilled_sigmas[0],
initial_video_latent=upscaled_video_latent,
initial_audio_latent=audio_state.latent,
)
torch.cuda.synchronize()
del transformer
del video_encoder
cleanup_memory()
decoded_video = vae_decode_video(video_state.latent, self.stage_2_model_ledger.video_decoder(), tiling_config)
decoded_audio = vae_decode_audio(
audio_state.latent, self.stage_2_model_ledger.audio_decoder(), self.stage_2_model_ledger.vocoder()
)
return decoded_video, decoded_audio
@torch.inference_mode()
def main() -> None:
logging.getLogger().setLevel(logging.INFO)
parser = default_2_stage_arg_parser()
args = parser.parse_args()
pipeline = KeyframeInterpolationPipeline(
checkpoint_path=args.checkpoint_path,
distilled_lora=args.distilled_lora,
spatial_upsampler_path=args.spatial_upsampler_path,
gemma_root=args.gemma_root,
loras=args.lora,
fp8transformer=args.enable_fp8,
)
tiling_config = TilingConfig.default()
video_chunks_number = get_video_chunks_number(args.num_frames, tiling_config)
video, audio = pipeline(
prompt=args.prompt,
negative_prompt=args.negative_prompt,
seed=args.seed,
height=args.height,
width=args.width,
num_frames=args.num_frames,
frame_rate=args.frame_rate,
num_inference_steps=args.num_inference_steps,
cfg_guidance_scale=args.cfg_guidance_scale,
images=args.images,
tiling_config=tiling_config,
)
encode_video(
video=video,
fps=args.frame_rate,
audio=audio,
audio_sample_rate=AUDIO_SAMPLE_RATE,
output_path=args.output_path,
video_chunks_number=video_chunks_number,
)
if __name__ == "__main__":
main()
packages/ltx-pipelines/src/ltx_pipelines/ti2vid_one_stage.py
+193
View File
@@ -0,0 +1,193 @@
import logging
from collections.abc import Iterator
import torch
from ltx_core.components.diffusion_steps import EulerDiffusionStep
from ltx_core.components.guiders import CFGGuider
from ltx_core.components.noisers import GaussianNoiser
from ltx_core.components.protocols import DiffusionStepProtocol
from ltx_core.components.schedulers import LTX2Scheduler
from ltx_core.loader import LoraPathStrengthAndSDOps
from ltx_core.model.audio_vae import decode_audio as vae_decode_audio
from ltx_core.model.video_vae import decode_video as vae_decode_video
from ltx_core.text_encoders.gemma import encode_text
from ltx_core.types import LatentState, VideoPixelShape
from ltx_pipelines.utils import ModelLedger
from ltx_pipelines.utils.args import default_1_stage_arg_parser
from ltx_pipelines.utils.constants import AUDIO_SAMPLE_RATE
from ltx_pipelines.utils.helpers import (
assert_resolution,
cleanup_memory,
denoise_audio_video,
euler_denoising_loop,
generate_enhanced_prompt,
get_device,
guider_denoising_func,
image_conditionings_by_replacing_latent,
)
from ltx_pipelines.utils.media_io import encode_video
from ltx_pipelines.utils.types import PipelineComponents
device = get_device()
class TI2VidOneStagePipeline:
"""
Single-stage text/image-to-video generation pipeline.
Generates video at the target resolution in a single diffusion pass with
classifier-free guidance (CFG). Supports optional image conditioning via
the images parameter.
"""
def __init__(
self,
checkpoint_path: str,
gemma_root: str,
loras: list[LoraPathStrengthAndSDOps],
device: torch.device = device,
fp8transformer: bool = False,
):
self.dtype = torch.bfloat16
self.device = device
self.model_ledger = ModelLedger(
dtype=self.dtype,
device=device,
checkpoint_path=checkpoint_path,
gemma_root_path=gemma_root,
loras=loras,
fp8transformer=fp8transformer,
)
self.pipeline_components = PipelineComponents(
dtype=self.dtype,
device=device,
)
def __call__( # noqa: PLR0913
self,
prompt: str,
negative_prompt: str,
seed: int,
height: int,
width: int,
num_frames: int,
frame_rate: float,
num_inference_steps: int,
cfg_guidance_scale: float,
images: list[tuple[str, int, float]],
enhance_prompt: bool = False,
) -> tuple[Iterator[torch.Tensor], torch.Tensor]:
assert_resolution(height=height, width=width, is_two_stage=False)
generator = torch.Generator(device=self.device).manual_seed(seed)
noiser = GaussianNoiser(generator=generator)
stepper = EulerDiffusionStep()
cfg_guider = CFGGuider(cfg_guidance_scale)
dtype = torch.bfloat16
text_encoder = self.model_ledger.text_encoder()
if enhance_prompt:
prompt = generate_enhanced_prompt(
text_encoder, prompt, images[0][0] if len(images) > 0 else None, seed=seed
)
context_p, context_n = encode_text(text_encoder, prompts=[prompt, negative_prompt])
v_context_p, a_context_p = context_p
v_context_n, a_context_n = context_n
torch.cuda.synchronize()
del text_encoder
cleanup_memory()
# Stage 1: Initial low resolution video generation.
video_encoder = self.model_ledger.video_encoder()
transformer = self.model_ledger.transformer()
sigmas = LTX2Scheduler().execute(steps=num_inference_steps).to(dtype=torch.float32, device=self.device)
def first_stage_denoising_loop(
sigmas: torch.Tensor, video_state: LatentState, audio_state: LatentState, stepper: DiffusionStepProtocol
) -> tuple[LatentState, LatentState]:
return euler_denoising_loop(
sigmas=sigmas,
video_state=video_state,
audio_state=audio_state,
stepper=stepper,
denoise_fn=guider_denoising_func(
cfg_guider,
v_context_p,
v_context_n,
a_context_p,
a_context_n,
transformer=transformer, # noqa: F821
),
)
stage_1_output_shape = VideoPixelShape(batch=1, frames=num_frames, width=width, height=height, fps=frame_rate)
stage_1_conditionings = image_conditionings_by_replacing_latent(
images=images,
height=stage_1_output_shape.height,
width=stage_1_output_shape.width,
video_encoder=video_encoder,
dtype=dtype,
device=self.device,
)
video_state, audio_state = denoise_audio_video(
output_shape=stage_1_output_shape,
conditionings=stage_1_conditionings,
noiser=noiser,
sigmas=sigmas,
stepper=stepper,
denoising_loop_fn=first_stage_denoising_loop,
components=self.pipeline_components,
dtype=dtype,
device=self.device,
)
torch.cuda.synchronize()
del transformer
cleanup_memory()
decoded_video = vae_decode_video(video_state.latent, self.model_ledger.video_decoder())
decoded_audio = vae_decode_audio(
audio_state.latent, self.model_ledger.audio_decoder(), self.model_ledger.vocoder()
)
return decoded_video, decoded_audio
@torch.inference_mode()
def main() -> None:
logging.getLogger().setLevel(logging.INFO)
parser = default_1_stage_arg_parser()
args = parser.parse_args()
pipeline = TI2VidOneStagePipeline(
checkpoint_path=args.checkpoint_path,
gemma_root=args.gemma_root,
loras=args.lora,
fp8transformer=args.enable_fp8,
)
video, audio = pipeline(
prompt=args.prompt,
negative_prompt=args.negative_prompt,
seed=args.seed,
height=args.height,
width=args.width,
num_frames=args.num_frames,
frame_rate=args.frame_rate,
num_inference_steps=args.num_inference_steps,
cfg_guidance_scale=args.cfg_guidance_scale,
images=args.images,
)
encode_video(
video=video,
fps=args.frame_rate,
audio=audio,
audio_sample_rate=AUDIO_SAMPLE_RATE,
output_path=args.output_path,
video_chunks_number=1,
)
if __name__ == "__main__":
main()
packages/ltx-pipelines/src/ltx_pipelines/ti2vid_two_stages.py
+276
View File
@@ -0,0 +1,276 @@
import logging
from collections.abc import Iterator
import torch
from ltx_core.components.diffusion_steps import EulerDiffusionStep
from ltx_core.components.guiders import CFGGuider
from ltx_core.components.noisers import GaussianNoiser
from ltx_core.components.protocols import DiffusionStepProtocol
from ltx_core.components.schedulers import LTX2Scheduler
from ltx_core.loader import LoraPathStrengthAndSDOps
from ltx_core.model.audio_vae import decode_audio as vae_decode_audio
from ltx_core.model.upsampler import upsample_video
from ltx_core.model.video_vae import TilingConfig, get_video_chunks_number
from ltx_core.model.video_vae import decode_video as vae_decode_video
from ltx_core.text_encoders.gemma import encode_text
from ltx_core.types import LatentState, VideoPixelShape
from ltx_pipelines.utils import ModelLedger
from ltx_pipelines.utils.args import default_2_stage_arg_parser
from ltx_pipelines.utils.constants import (
AUDIO_SAMPLE_RATE,
STAGE_2_DISTILLED_SIGMA_VALUES,
)
from ltx_pipelines.utils.helpers import (
assert_resolution,
cleanup_memory,
denoise_audio_video,
euler_denoising_loop,
generate_enhanced_prompt,
get_device,
guider_denoising_func,
image_conditionings_by_replacing_latent,
simple_denoising_func,
)
from ltx_pipelines.utils.media_io import encode_video
from ltx_pipelines.utils.types import PipelineComponents
device = get_device()
class TI2VidTwoStagesPipeline:
"""
Two-stage text/image-to-video generation pipeline.
Stage 1 generates video at the target resolution with CFG guidance, then
Stage 2 upsamples by 2x and refines using a distilled LoRA for higher
quality output. Supports optional image conditioning via the images parameter.
"""
def __init__(
self,
checkpoint_path: str,
distilled_lora: list[LoraPathStrengthAndSDOps],
spatial_upsampler_path: str,
gemma_root: str,
loras: list[LoraPathStrengthAndSDOps],
device: str = device,
fp8transformer: bool = False,
):
self.device = device
self.dtype = torch.bfloat16
self.stage_1_model_ledger = ModelLedger(
dtype=self.dtype,
device=device,
checkpoint_path=checkpoint_path,
gemma_root_path=gemma_root,
spatial_upsampler_path=spatial_upsampler_path,
loras=loras,
fp8transformer=fp8transformer,
)
self.stage_2_model_ledger = self.stage_1_model_ledger.with_loras(
loras=distilled_lora,
)
self.pipeline_components = PipelineComponents(
dtype=self.dtype,
device=device,
)
@torch.inference_mode()
def __call__( # noqa: PLR0913
self,
prompt: str,
negative_prompt: str,
seed: int,
height: int,
width: int,
num_frames: int,
frame_rate: float,
num_inference_steps: int,
cfg_guidance_scale: float,
images: list[tuple[str, int, float]],
tiling_config: TilingConfig | None = None,
enhance_prompt: bool = False,
) -> tuple[Iterator[torch.Tensor], torch.Tensor]:
assert_resolution(height=height, width=width, is_two_stage=True)
generator = torch.Generator(device=self.device).manual_seed(seed)
noiser = GaussianNoiser(generator=generator)
stepper = EulerDiffusionStep()
cfg_guider = CFGGuider(cfg_guidance_scale)
dtype = torch.bfloat16
text_encoder = self.stage_1_model_ledger.text_encoder()
if enhance_prompt:
prompt = generate_enhanced_prompt(
text_encoder, prompt, images[0][0] if len(images) > 0 else None, seed=seed
)
context_p, context_n = encode_text(text_encoder, prompts=[prompt, negative_prompt])
v_context_p, a_context_p = context_p
v_context_n, a_context_n = context_n
torch.cuda.synchronize()
del text_encoder
cleanup_memory()
# Stage 1: Initial low resolution video generation.
video_encoder = self.stage_1_model_ledger.video_encoder()
transformer = self.stage_1_model_ledger.transformer()
sigmas = LTX2Scheduler().execute(steps=num_inference_steps).to(dtype=torch.float32, device=self.device)
def first_stage_denoising_loop(
sigmas: torch.Tensor, video_state: LatentState, audio_state: LatentState, stepper: DiffusionStepProtocol
) -> tuple[LatentState, LatentState]:
return euler_denoising_loop(
sigmas=sigmas,
video_state=video_state,
audio_state=audio_state,
stepper=stepper,
denoise_fn=guider_denoising_func(
cfg_guider,
v_context_p,
v_context_n,
a_context_p,
a_context_n,
transformer=transformer, # noqa: F821
),
)
stage_1_output_shape = VideoPixelShape(
batch=1,
frames=num_frames,
width=width // 2,
height=height // 2,
fps=frame_rate,
)
stage_1_conditionings = image_conditionings_by_replacing_latent(
images=images,
height=stage_1_output_shape.height,
width=stage_1_output_shape.width,
video_encoder=video_encoder,
dtype=dtype,
device=self.device,
)
video_state, audio_state = denoise_audio_video(
output_shape=stage_1_output_shape,
conditionings=stage_1_conditionings,
noiser=noiser,
sigmas=sigmas,
stepper=stepper,
denoising_loop_fn=first_stage_denoising_loop,
components=self.pipeline_components,
dtype=dtype,
device=self.device,
)
torch.cuda.synchronize()
del transformer
cleanup_memory()
# Stage 2: Upsample and refine the video at higher resolution with distilled LORA.
upscaled_video_latent = upsample_video(
latent=video_state.latent[:1],
video_encoder=video_encoder,
upsampler=self.stage_2_model_ledger.spatial_upsampler(),
)
torch.cuda.synchronize()
cleanup_memory()
transformer = self.stage_2_model_ledger.transformer()
distilled_sigmas = torch.Tensor(STAGE_2_DISTILLED_SIGMA_VALUES).to(self.device)
def second_stage_denoising_loop(
sigmas: torch.Tensor, video_state: LatentState, audio_state: LatentState, stepper: DiffusionStepProtocol
) -> tuple[LatentState, LatentState]:
return euler_denoising_loop(
sigmas=sigmas,
video_state=video_state,
audio_state=audio_state,
stepper=stepper,
denoise_fn=simple_denoising_func(
video_context=v_context_p,
audio_context=a_context_p,
transformer=transformer, # noqa: F821
),
)
stage_2_output_shape = VideoPixelShape(batch=1, frames=num_frames, width=width, height=height, fps=frame_rate)
stage_2_conditionings = image_conditionings_by_replacing_latent(
images=images,
height=stage_2_output_shape.height,
width=stage_2_output_shape.width,
video_encoder=video_encoder,
dtype=dtype,
device=self.device,
)
video_state, audio_state = denoise_audio_video(
output_shape=stage_2_output_shape,
conditionings=stage_2_conditionings,
noiser=noiser,
sigmas=distilled_sigmas,
stepper=stepper,
denoising_loop_fn=second_stage_denoising_loop,
components=self.pipeline_components,
dtype=dtype,
device=self.device,
noise_scale=distilled_sigmas[0],
initial_video_latent=upscaled_video_latent,
initial_audio_latent=audio_state.latent,
)
torch.cuda.synchronize()
del transformer
del video_encoder
cleanup_memory()
decoded_video = vae_decode_video(video_state.latent, self.stage_2_model_ledger.video_decoder(), tiling_config)
decoded_audio = vae_decode_audio(
audio_state.latent, self.stage_2_model_ledger.audio_decoder(), self.stage_2_model_ledger.vocoder()
)
return decoded_video, decoded_audio
@torch.inference_mode()
def main() -> None:
logging.getLogger().setLevel(logging.INFO)
parser = default_2_stage_arg_parser()
args = parser.parse_args()
pipeline = TI2VidTwoStagesPipeline(
checkpoint_path=args.checkpoint_path,
distilled_lora=args.distilled_lora,
spatial_upsampler_path=args.spatial_upsampler_path,
gemma_root=args.gemma_root,
loras=args.lora,
fp8transformer=args.enable_fp8,
)
tiling_config = TilingConfig.default()
video_chunks_number = get_video_chunks_number(args.num_frames, tiling_config)
video, audio = pipeline(
prompt=args.prompt,
negative_prompt=args.negative_prompt,
seed=args.seed,
height=args.height,
width=args.width,
num_frames=args.num_frames,
frame_rate=args.frame_rate,
num_inference_steps=args.num_inference_steps,
cfg_guidance_scale=args.cfg_guidance_scale,
images=args.images,
tiling_config=tiling_config,
)
encode_video(
video=video,
fps=args.frame_rate,
audio=audio,
audio_sample_rate=AUDIO_SAMPLE_RATE,
output_path=args.output_path,
video_chunks_number=video_chunks_number,
)
if __name__ == "__main__":
main()
packages/ltx-pipelines/src/ltx_pipelines/utils/__init__.py
+5
View File
@@ -0,0 +1,5 @@
from ltx_pipelines.utils.model_ledger import ModelLedger
__all__ = [
"ModelLedger",
]
packages/ltx-pipelines/src/ltx_pipelines/utils/args.py
+277
View File
@@ -0,0 +1,277 @@
import argparse
from pathlib import Path
from ltx_core.loader import LTXV_LORA_COMFY_RENAMING_MAP, LoraPathStrengthAndSDOps
from ltx_pipelines.utils.constants import (
DEFAULT_1_STAGE_HEIGHT,
DEFAULT_1_STAGE_WIDTH,
DEFAULT_2_STAGE_HEIGHT,
DEFAULT_2_STAGE_WIDTH,
DEFAULT_CFG_GUIDANCE_SCALE,
DEFAULT_FRAME_RATE,
DEFAULT_LORA_STRENGTH,
DEFAULT_NEGATIVE_PROMPT,
DEFAULT_NUM_FRAMES,
DEFAULT_NUM_INFERENCE_STEPS,
DEFAULT_SEED,
)
class VideoConditioningAction(argparse.Action):
def __call__(
self,
parser: argparse.ArgumentParser, # noqa: ARG002
namespace: argparse.Namespace,
values: list[str],
option_string: str | None = None, # noqa: ARG002
) -> None:
path, strength_str = values
resolved_path = resolve_path(path)
strength = float(strength_str)
current = getattr(namespace, self.dest) or []
current.append((resolved_path, strength))
setattr(namespace, self.dest, current)
class ImageAction(argparse.Action):
def __call__(
self,
parser: argparse.ArgumentParser, # noqa: ARG002
namespace: argparse.Namespace,
values: list[str],
option_string: str | None = None, # noqa: ARG002
) -> None:
path, frame_idx, strength_str = values
resolved_path = resolve_path(path)
frame_idx = int(frame_idx)
strength = float(strength_str)
current = getattr(namespace, self.dest) or []
current.append((resolved_path, frame_idx, strength))
setattr(namespace, self.dest, current)
class LoraAction(argparse.Action):
def __call__(
self,
parser: argparse.ArgumentParser, # noqa: ARG002
namespace: argparse.Namespace,
values: list[str],
option_string: str | None = None,
) -> None:
if len(values) > 2:
msg = f"{option_string} accepts at most 2 arguments (PATH and optional STRENGTH), got {len(values)} values"
raise argparse.ArgumentError(self, msg)
path = values[0]
strength_str = values[1] if len(values) > 1 else str(DEFAULT_LORA_STRENGTH)
resolved_path = resolve_path(path)
strength = float(strength_str)
current = getattr(namespace, self.dest) or []
current.append(LoraPathStrengthAndSDOps(resolved_path, strength, LTXV_LORA_COMFY_RENAMING_MAP))
setattr(namespace, self.dest, current)
def resolve_path(path: str) -> str:
return str(Path(path).expanduser().resolve().as_posix())
def basic_arg_parser() -> argparse.ArgumentParser:
parser = argparse.ArgumentParser()
parser.add_argument(
"--checkpoint-path",
type=resolve_path,
required=True,
help="Path to LTX-2 model checkpoint (.safetensors file).",
)
parser.add_argument(
"--gemma-root",
type=resolve_path,
required=True,
help="Path to the root directory containing the Gemma text encoder model files.",
)
parser.add_argument(
"--prompt",
type=str,
required=True,
help="Text prompt describing the desired video content to be generated by the model.",
)
parser.add_argument(
"--output-path",
type=resolve_path,
required=True,
help="Path to the output video file (MP4 format).",
)
parser.add_argument(
"--seed",
type=int,
default=DEFAULT_SEED,
help=(
f"Random seed value used to initialize the noise tensor for "
f"reproducible generation (default: {DEFAULT_SEED})."
),
)
parser.add_argument(
"--height",
type=int,
default=DEFAULT_1_STAGE_HEIGHT,
help=f"Height of the generated video in pixels, should be divisible by 32 (default: {DEFAULT_1_STAGE_HEIGHT}).",
)
parser.add_argument(
"--width",
type=int,
default=DEFAULT_1_STAGE_WIDTH,
help=f"Width of the generated video in pixels, should be divisible by 32 (default: {DEFAULT_1_STAGE_WIDTH}).",
)
parser.add_argument(
"--num-frames",
type=int,
default=DEFAULT_NUM_FRAMES,
help=f"Number of frames to generate in the output video sequence, num-frames = (8 x K) + 1, "
f"where k is a non-negative integer (default: {DEFAULT_NUM_FRAMES}).",
)
parser.add_argument(
"--frame-rate",
type=float,
default=DEFAULT_FRAME_RATE,
help=f"Frame rate of the generated video (fps) (default: {DEFAULT_FRAME_RATE}).",
)
parser.add_argument(
"--num-inference-steps",
type=int,
default=DEFAULT_NUM_INFERENCE_STEPS,
help=(
f"Number of denoising steps in the diffusion sampling process. "
f"Higher values improve quality but increase generation time (default: {DEFAULT_NUM_INFERENCE_STEPS})."
),
)
parser.add_argument(
"--image",
dest="images",
action=ImageAction,
nargs=3,
metavar=("PATH", "FRAME_IDX", "STRENGTH"),
default=[],
help=(
"Image conditioning input: path to image file, target frame index, "
"and conditioning strength (all three required). Default: empty list [] (no image conditioning). "
"Can be specified multiple times. Example: --image path/to/image1.jpg 0 0.8 "
"--image path/to/image2.jpg 160 0.9"
),
)
parser.add_argument(
"--lora",
dest="lora",
action=LoraAction,
nargs="+", # Accept 1-2 arguments per use (path and optional strength); validation is handled in LoraAction
metavar=("PATH", "STRENGTH"),
default=[],
help=(
"LoRA (Low-Rank Adaptation) model: path to model file and optional strength "
f"(default strength: {DEFAULT_LORA_STRENGTH}). Can be specified multiple times. "
"Example: --lora path/to/lora1.safetensors 0.8 --lora path/to/lora2.safetensors"
),
)
parser.add_argument(
"--enable-fp8",
action="store_true",
help="Enable FP8 mode to reduce memory footprint by keeping model in lower precision. "
"Note that calculations are still performed in bfloat16 precision.",
)
parser.add_argument("--enhance-prompt", action="store_true")
return parser
def default_1_stage_arg_parser() -> argparse.ArgumentParser:
parser = basic_arg_parser()
parser.add_argument(
"--cfg-guidance-scale",
type=float,
default=DEFAULT_CFG_GUIDANCE_SCALE,
help=(
f"Classifier-free guidance (CFG) scale controlling how strongly "
f"the model adheres to the prompt. Higher values increase prompt "
f"adherence but may reduce diversity (default: {DEFAULT_CFG_GUIDANCE_SCALE})."
),
)
parser.add_argument(
"--negative-prompt",
type=str,
default=DEFAULT_NEGATIVE_PROMPT,
help=(
"Negative prompt describing what should not appear in the generated video, "
"used to guide the diffusion process away from unwanted content. "
"Default: a comprehensive negative prompt covering common artifacts and quality issues."
),
)
return parser
def default_2_stage_arg_parser() -> argparse.ArgumentParser:
parser = default_1_stage_arg_parser()
parser.set_defaults(height=DEFAULT_2_STAGE_HEIGHT, width=DEFAULT_2_STAGE_WIDTH)
# Update help text to reflect 2-stage defaults
for action in parser._actions:
if "--height" in action.option_strings:
action.help = (
f"Height of the generated video in pixels, should be divisible by 64 "
f"(default: {DEFAULT_2_STAGE_HEIGHT})."
)
if "--width" in action.option_strings:
action.help = (
f"Width of the generated video in pixels, should be divisible by 64 (default: {DEFAULT_2_STAGE_WIDTH})."
)
parser.add_argument(
"--distilled-lora",
dest="distilled_lora",
action=LoraAction,
nargs="+", # Accept 1-2 arguments per use (path and optional strength); validation is handled in LoraAction
metavar=("PATH", "STRENGTH"),
required=True,
help=(
"Distilled LoRA (Low-Rank Adaptation) model used in the second stage (upscaling and refinement): "
f"path to model file and optional strength (default strength: {DEFAULT_LORA_STRENGTH}). "
"The second stage upsamples the video by 2x resolution and refines it using a distilled "
"denoising schedule (fewer steps, no CFG). The distilled LoRA is specifically trained "
"for this refinement process to improve quality at higher resolutions. "
"Example: --distilled-lora path/to/distilled_lora.safetensors 0.8"
),
)
parser.add_argument(
"--spatial-upsampler-path",
type=resolve_path,
required=True,
help=(
"Path to the spatial upsampler model used to increase the resolution "
"of the generated video in the latent space."
),
)
return parser
def default_2_stage_distilled_arg_parser() -> argparse.ArgumentParser:
parser = basic_arg_parser()
parser.set_defaults(height=DEFAULT_2_STAGE_HEIGHT, width=DEFAULT_2_STAGE_WIDTH)
# Update help text to reflect 2-stage defaults
for action in parser._actions:
if "--height" in action.option_strings:
action.help = (
f"Height of the generated video in pixels, should be divisible by 64 "
f"(default: {DEFAULT_2_STAGE_HEIGHT})."
)
if "--width" in action.option_strings:
action.help = (
f"Width of the generated video in pixels, should be divisible by 64 (default: {DEFAULT_2_STAGE_WIDTH})."
)
parser.add_argument(
"--spatial-upsampler-path",
type=resolve_path,
required=True,
help=(
"Path to the spatial upsampler model used to increase the resolution "
"of the generated video in the latent space."
),
)
return parser
packages/ltx-pipelines/src/ltx_pipelines/utils/constants.py
+77
View File
@@ -0,0 +1,77 @@
# =============================================================================
# Diffusion Schedule
# =============================================================================
# Noise schedule for the distilled pipeline. These sigma values control noise
# levels at each denoising step and were tuned to match the distillation process.
from ltx_core.types import SpatioTemporalScaleFactors
DISTILLED_SIGMA_VALUES = [1.0, 0.99375, 0.9875, 0.98125, 0.975, 0.909375, 0.725, 0.421875, 0.0]
# Reduced schedule for super-resolution stage 2 (subset of distilled values)
STAGE_2_DISTILLED_SIGMA_VALUES = [0.909375, 0.725, 0.421875, 0.0]
# =============================================================================
# Video Generation Defaults
# =============================================================================
DEFAULT_SEED = 10
DEFAULT_1_STAGE_HEIGHT = 512
DEFAULT_1_STAGE_WIDTH = 768
DEFAULT_2_STAGE_HEIGHT = DEFAULT_1_STAGE_HEIGHT * 2
DEFAULT_2_STAGE_WIDTH = DEFAULT_1_STAGE_WIDTH * 2
DEFAULT_NUM_FRAMES = 121
DEFAULT_FRAME_RATE = 24.0
DEFAULT_NUM_INFERENCE_STEPS = 40
DEFAULT_CFG_GUIDANCE_SCALE = 4.0
# =============================================================================
# Audio
# =============================================================================
AUDIO_SAMPLE_RATE = 24000
# =============================================================================
# LoRA
# =============================================================================
DEFAULT_LORA_STRENGTH = 1.0
# =============================================================================
# Video VAE Architecture
# =============================================================================
VIDEO_SCALE_FACTORS = SpatioTemporalScaleFactors.default()
VIDEO_LATENT_CHANNELS = 128
# =============================================================================
# Image Preprocessing
# =============================================================================
# CRF (Constant Rate Factor) for H.264 encoding used in image conditioning.
# Lower = higher quality, 0 = lossless. This mimics compression artifacts.
DEFAULT_IMAGE_CRF = 33
# =============================================================================
# Prompts
# =============================================================================
DEFAULT_NEGATIVE_PROMPT = (
"blurry, out of focus, overexposed, underexposed, low contrast, washed out colors, excessive noise, "
"grainy texture, poor lighting, flickering, motion blur, distorted proportions, unnatural skin tones, "
"deformed facial features, asymmetrical face, missing facial features, extra limbs, disfigured hands, "
"wrong hand count, artifacts around text, inconsistent perspective, camera shake, incorrect depth of "
"field, background too sharp, background clutter, distracting reflections, harsh shadows, inconsistent "
"lighting direction, color banding, cartoonish rendering, 3D CGI look, unrealistic materials, uncanny "
"valley effect, incorrect ethnicity, wrong gender, exaggerated expressions, wrong gaze direction, "
"mismatched lip sync, silent or muted audio, distorted voice, robotic voice, echo, background noise, "
"off-sync audio, incorrect dialogue, added dialogue, repetitive speech, jittery movement, awkward "
"pauses, incorrect timing, unnatural transitions, inconsistent framing, tilted camera, flat lighting, "
"inconsistent tone, cinematic oversaturation, stylized filters, or AI artifacts."
)

Some files were not shown because too many files have changed in this diff Show More