Add files using upload-large-folder tool

.gitattributes

@@ -1,35 +1,9 @@
-*.7z filter=lfs diff=lfs merge=lfs -text
-*.arrow filter=lfs diff=lfs merge=lfs -text
-*.bin filter=lfs diff=lfs merge=lfs -text
-*.bz2 filter=lfs diff=lfs merge=lfs -text
-*.ckpt filter=lfs diff=lfs merge=lfs -text
-*.ftz filter=lfs diff=lfs merge=lfs -text
-*.gz filter=lfs diff=lfs merge=lfs -text
-*.h5 filter=lfs diff=lfs merge=lfs -text
-*.joblib filter=lfs diff=lfs merge=lfs -text
-*.lfs.* filter=lfs diff=lfs merge=lfs -text
-*.mlmodel filter=lfs diff=lfs merge=lfs -text
-*.model filter=lfs diff=lfs merge=lfs -text
-*.msgpack filter=lfs diff=lfs merge=lfs -text
-*.npy filter=lfs diff=lfs merge=lfs -text
-*.npz filter=lfs diff=lfs merge=lfs -text
-*.onnx filter=lfs diff=lfs merge=lfs -text
-*.ot filter=lfs diff=lfs merge=lfs -text
-*.parquet filter=lfs diff=lfs merge=lfs -text
-*.pb filter=lfs diff=lfs merge=lfs -text
-*.pickle filter=lfs diff=lfs merge=lfs -text
-*.pkl filter=lfs diff=lfs merge=lfs -text
-*.pt filter=lfs diff=lfs merge=lfs -text
-*.pth filter=lfs diff=lfs merge=lfs -text
-*.rar filter=lfs diff=lfs merge=lfs -text
+*.gif filter=lfs diff=lfs merge=lfs -text
 *.safetensors filter=lfs diff=lfs merge=lfs -text
-saved_model/**/* filter=lfs diff=lfs merge=lfs -text
-*.tar.* filter=lfs diff=lfs merge=lfs -text
-*.tar filter=lfs diff=lfs merge=lfs -text
-*.tflite filter=lfs diff=lfs merge=lfs -text
-*.tgz filter=lfs diff=lfs merge=lfs -text
-*.wasm filter=lfs diff=lfs merge=lfs -text
-*.xz filter=lfs diff=lfs merge=lfs -text
-*.zip filter=lfs diff=lfs merge=lfs -text
-*.zst filter=lfs diff=lfs merge=lfs -text
-*tfevents* filter=lfs diff=lfs merge=lfs -text
+*.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

@@ -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

@@ -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](https://ltx.io/model/licensing).
+      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

@@ -0,0 +1,120 @@
+# 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.3)
+[![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-PDF-EC1C24?logo=adobeacrobatreader&logoColor=white)](https://arxiv.org/abs/2601.03233)
+[![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.
+
+<div align="center">
+  <video src="https://github.com/user-attachments/assets/4414adc0-086c-43de-b367-9362eeb20228" width="70%" poster=""> </video>
+</div>
+
+## 🚀 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.3 HuggingFace repository](https://huggingface.co/Lightricks/LTX-2.3):
+
+**LTX-2.3 Model Checkpoint** (choose and download one of the following)
+  * [`ltx-2.3-22b-dev.safetensors`](https://huggingface.co/Lightricks/LTX-2.3/blob/main/ltx-2.3-22b-dev.safetensors) - [Download](https://huggingface.co/Lightricks/LTX-2.3/resolve/main/ltx-2.3-22b-dev.safetensors)
+  * [`ltx-2.3-22b-distilled.safetensors`](https://huggingface.co/Lightricks/LTX-2.3/blob/main/ltx-2.3-22b-distilled.safetensors) - [Download](https://huggingface.co/Lightricks/LTX-2.3/resolve/main/ltx-2.3-22b-distilled.safetensors)
+
+**Spatial Upscaler** - Required for current two-stage pipeline implementations in this repository
+  * [`ltx-2.3-spatial-upscaler-x2-1.0.safetensors`](https://huggingface.co/Lightricks/LTX-2.3/blob/main/ltx-2.3-spatial-upscaler-x2-1.0.safetensors) - [Download](https://huggingface.co/Lightricks/LTX-2.3/resolve/main/ltx-2.3-spatial-upscaler-x2-1.0.safetensors)
+  * [`ltx-2.3-spatial-upscaler-x1.5-1.0.safetensors`](https://huggingface.co/Lightricks/LTX-2.3/blob/main/ltx-2.3-spatial-upscaler-x1.5-1.0.safetensors) - [Download](https://huggingface.co/Lightricks/LTX-2.3/resolve/main/ltx-2.3-spatial-upscaler-x1.5-1.0.safetensors)
+
+**Temporal Upscaler** - Supported by the model and will be required for future pipeline implementations
+  * [`ltx-2.3-temporal-upscaler-x2-1.0.safetensors`](https://huggingface.co/Lightricks/LTX-2.3/blob/main/ltx-2.3-temporal-upscaler-x2-1.0.safetensors) - [Download](https://huggingface.co/Lightricks/LTX-2.3/resolve/main/ltx-2.3-temporal-upscaler-x2-1.0.safetensors)
+
+**Distilled LoRA** - Required for current two-stage pipeline implementations in this repository (except DistilledPipeline and ICLoraPipeline)
+  * [`ltx-2.3-22b-distilled-lora-384.safetensors`](https://huggingface.co/Lightricks/LTX-2.3/blob/main/ltx-2.3-22b-distilled-lora-384.safetensors) - [Download](https://huggingface.co/Lightricks/LTX-2.3/resolve/main/ltx-2.3-22b-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.3-22b-IC-LoRA-Union-Control`](https://huggingface.co/Lightricks/LTX-2.3-22b-IC-LoRA-Union-Control) - [Download](https://huggingface.co/Lightricks/LTX-2.3-22b-IC-LoRA-Union-Control/resolve/main/ltx-2.3-22b-ic-lora-union-control-ref0.5.safetensors)
+  * [`LTX-2.3-22b-IC-LoRA-Inpainting`](https://huggingface.co/Lightricks/LTX-2.3-22b-IC-LoRA-Inpainting) - [Download](https://huggingface.co/Lightricks/LTX-2.3-22b-IC-LoRA-Inpainting/resolve/main/ltx-2.3-22b-ic-lora-inpainting.safetensors)
+  * [`LTX-2.3-22b-IC-LoRA-Motion-Track-Control`](https://huggingface.co/Lightricks/LTX-2.3-22b-IC-LoRA-Motion-Track-Control) - [Download](https://huggingface.co/Lightricks/LTX-2.3-22b-IC-LoRA-Motion-Track-Control/resolve/main/ltx-2.3-22b-ic-lora-motion-track-control-ref0.5.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/image-to-video with 2x upsampling (recommended)
+* **[TI2VidTwoStagesHQPipeline](packages/ltx-pipelines/src/ltx_pipelines/ti2vid_two_stages_hq.py)** - Same two-stage flow as above but uses the res_2s second-order sampler (fewer steps, better quality)
+* **[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 (uses distilled model.)
+* **[KeyframeInterpolationPipeline](packages/ltx-pipelines/src/ltx_pipelines/keyframe_interpolation.py)** - Interpolate between keyframe images
+* **[A2VidPipelineTwoStage](packages/ltx-pipelines/src/ltx_pipelines/a2vid_two_stage.py)** - Audio-to-video generation conditioned on an input audio file
+* **[RetakePipeline](packages/ltx-pipelines/src/ltx_pipelines/retake.py)** - Regenerate a specific time region of an existing video
+
+### ⚡ Optimization Tips
+
+* **Use DistilledPipeline** - Fastest inference with only 8 predefined sigmas (8 steps stage 1, 4 steps stage 2)
+* **Enable FP8 quantization** - Enables lower memory footprint: `--quantization fp8-cast` (CLI) or `quantization=QuantizationPolicy.fp8_cast()` (Python). For Hopper GPUs with TensorRT-LLM, use `--quantization fp8-scaled-mm` for FP8 scaled matrix multiplication.
+* **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

@@ -0,0 +1,409 @@
+# 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
+- **`quantization/`**: FP8 quantization backends (FP8-TensorRT-LLM scaled MM, FP8 cast) for reduced memory footprint.
+
+## 🚀 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 asymmetric dual-stream LTX-2 transformer (14B-parameter video stream, 5B-parameter audio stream) with bidirectional 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 3-based multilingual encoder with multi-layer feature extraction and thinking tokens 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
+- **Quantization** ([`quantization/`](src/ltx_core/quantization/)): FP8 quantization backends for reduced memory footprint and faster inference
+
+### Loader
+
+The `loader/` module provides `SingleGPUModelBuilder`, a frozen dataclass that loads a PyTorch model from `.safetensors` checkpoints and optionally fuses one or more LoRA adapters.
+
+#### Basic usage
+
+```python
+from ltx_core.loader import SingleGPUModelBuilder
+
+builder = SingleGPUModelBuilder(
+    model_class_configurator=MyModelConfigurator,
+    model_path="/path/to/model.safetensors",
+)
+model = builder.build(device=torch.device("cuda"))
+```
+
+#### Loading LoRA adapters
+
+Use the `.lora()` method to attach one or more LoRA adapters before calling `.build()`:
+
+```python
+builder = (
+    SingleGPUModelBuilder(
+        model_class_configurator=MyModelConfigurator,
+        model_path="/path/to/model.safetensors",
+    )
+    .lora("/path/to/lora_a.safetensors", strength=0.8)
+    .lora("/path/to/lora_b.safetensors", strength=0.5)
+)
+model = builder.build(device=torch.device("cuda"))
+```
+
+#### Memory-efficient LoRA loading (`lora_load_device`)
+
+By default, LoRA weights are loaded onto the **CPU** (`lora_load_device=torch.device("cpu")`).  This means each LoRA adapter is kept in CPU memory and transferred to the GPU sequentially during weight fusion, which keeps peak GPU memory low even when fusing large adapters.
+
+If all adapters fit comfortably in GPU memory you can skip the CPU staging by setting `lora_load_device` to the target CUDA device:
+
+```python
+import torch
+from ltx_core.loader import SingleGPUModelBuilder
+
+# Load LoRA weights directly onto the GPU (faster, but uses more GPU memory)
+builder = SingleGPUModelBuilder(
+    model_class_configurator=MyModelConfigurator,
+    model_path="/path/to/model.safetensors",
+    lora_load_device=torch.device("cuda"),
+).lora("/path/to/lora.safetensors", strength=1.0)
+
+model = builder.build(device=torch.device("cuda"))
+```
+
+### Quantization
+
+The `quantization/` module provides FP8 quantization support for the LTX-2 transformer, significantly reducing memory usage while maintaining quality. Two backends are available:
+
+#### FP8 Scaled MM (TensorRT-LLM)
+
+Uses NVIDIA TensorRT-LLM's `cublas_scaled_mm` for efficient FP8 matrix multiplication. Weights are stored in FP8 format with per-tensor scaling, and inputs are quantized dynamically (or statically with calibration data).
+
+**Requirements**: `uv sync --frozen --extra fp8-trtllm`
+
+**Usage with QuantizationPolicy:**
+
+```python
+from ltx_core.quantization import QuantizationPolicy
+
+# Dynamic input quantization (no calibration needed)
+policy = QuantizationPolicy.fp8_scaled_mm()
+
+# Static input quantization with calibration file
+policy = QuantizationPolicy.fp8_scaled_mm(calibration_amax_path="/path/to/amax.json")
+```
+
+The policy provides `sd_ops` and `module_ops` that can be passed to the model builder:
+
+```python
+from ltx_core.loader import SingleGPUModelBuilder
+
+builder = SingleGPUModelBuilder(
+    model=model,
+    device=device,
+    sd_ops=policy.sd_ops,
+    module_ops=policy.module_ops,
+)
+builder.load(checkpoint_path)
+```
+
+**Calibration File Format** (for static input quantization):
+
+```json
+{
+  "amax_values": {
+    "transformer_blocks.0.attn.to_q.input_quantizer": 12.5,
+    "transformer_blocks.0.attn.to_k.input_quantizer": 8.3,
+    ...
+  }
+}
+```
+
+#### FP8 Cast
+
+A simpler approach that casts weights to FP8 for storage and upcasts during inference:
+
+```python
+policy = QuantizationPolicy.fp8_cast()
+```
+
+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 an **asymmetric dual-stream diffusion transformer** that jointly models the text-conditioned distribution of video and audio signals, capturing true joint dependencies (unlike sequential T2V→V2A pipelines).
+
+### Key Design Principles
+
+- **Decoupled Latent Representations**: Separate modality-specific VAEs enable 3D RoPE (video) vs 1D RoPE (audio), independent compression optimization, and native V2A/A2V editing workflows
+- **Asymmetric Dual-Stream**: 14B-parameter video stream (spatiotemporal dynamics) + 5B-parameter audio stream (1D temporal), sharing 48 transformer blocks but differing in width
+- **Bidirectional Cross-Modal Attention**: 1D temporal RoPE enables sub-frame alignment, mapping visual cues to auditory events (lip-sync, foley, environmental acoustics)
+- **Cross-Modality AdaLN**: Scaling/shift parameters conditioned on the other modality's hidden states for synchronization across differing diffusion timesteps/temporal resolutions
+
+```text
+┌─────────────────────────────────────────────────────────────┐
+│                    INPUT PREPARATION                        │
+│                                                             │
+│  Video Pixels → Video VAE Encoder → Video Latents           │
+│  Audio Waveform → Audio VAE Encoder → Audio Latents         │
+│  Text Prompt → Gemma 3 Encoder → Text Embeddings            │
+└─────────────────────────────────────────────────────────────┘
+                            ↓
+┌─────────────────────────────────────────────────────────────┐
+│     LTX-2 ASYMMETRIC DUAL-STREAM TRANSFORMER (48 Blocks)    │
+│                                                             │
+│  ┌──────────────────────┐      ┌──────────────────────┐     │
+│  │  Video Stream (14B)  │      │  Audio Stream (5B)   │     │
+│  │                      │      │                      │     │
+│  │  3D RoPE (x,y,t)     │      │  1D RoPE (temporal)  │     │
+│  │                      │      │                      │     │
+│  │  Self-Attn           │      │  Self-Attn           │     │
+│  │  Text Cross-Attn     │      │  Text Cross-Attn     │     │
+│  │                      │◄────►│                      │     │
+│  │  A↔V Cross-Attn      │      │  A↔V Cross-Attn      │     │
+│  │  (1D temporal RoPE)  │      │  (1D temporal RoPE)  │     │
+│  │  Cross-modality      │      │  Cross-modality      │     │
+│  │  AdaLN               │      │  AdaLN               │     │
+│  │  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 (24 kHz)        │
+└─────────────────────────────────────────────────────────────┘
+```
+
+---
+
+## The Transformer
+
+The core of LTX-2 is an **asymmetric dual-stream diffusion transformer** with 48 layers that processes both video and audio tokens simultaneously. The architecture allocates 14B parameters to the video stream and 5B parameters to the audio stream, reflecting the different information densities of the two modalities.
+
+### 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)
+
+Each dual-stream block performs four operations sequentially:
+
+1. **Self-Attention**: Within-modality attention for each stream
+2. **Text Cross-Attention**: Textual prompt conditioning for both streams
+3. **Audio-Visual Cross-Attention**: Bidirectional inter-modal exchange
+4. **Feed-Forward Network (FFN)**: Feature refinement
+
+```text
+┌─────────────────────────────────────────────────────────────┐
+│                    TRANSFORMER BLOCK                        │
+│                                                             │
+│  VIDEO (14B): Input → RMSNorm → AdaLN → Self-Attn →         │
+│              RMSNorm → Text Cross-Attn →                    │
+│              RMSNorm → AdaLN → A↔V Cross-Attn (1D RoPE) →   │
+│              RMSNorm → AdaLN → FFN → Output                 │
+│                                                             │
+│  AUDIO (5B):  Input → RMSNorm → AdaLN → Self-Attn →         │
+│              RMSNorm → Text Cross-Attn →                    │
+│              RMSNorm → AdaLN → A↔V Cross-Attn (1D RoPE) →   │
+│              RMSNorm → AdaLN → FFN → Output                 │
+│                                                             │
+│  RoPE: Video=3D (x,y,t), Audio=1D (t), Cross-Attn=1D (t)    │
+│  AdaLN: Timestep-conditioned, cross-modality for A↔V CA     │
+└─────────────────────────────────────────────────────────────┘
+```
+
+### Audio-Visual Cross-Attention Details
+
+Bidirectional cross-attention enables tight temporal alignment: video and audio streams exchange information bidirectionally using 1D temporal RoPE (synchronization only, no spatial alignment). AdaLN gates condition on each modality's timestep for cross-modal synchronization.
+
+### 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
+
+Compact neural audio representation optimized for diffusion-based training. Natively supports stereo: processes two-channel mel-spectrograms (16 kHz input) with channel concatenation before encoding.
+
+- **Encoder**: `[B, mel_bins, T]` → `[B, 8, T/4, 16]` latents (4× temporal downsampling, 8 channels, 16 mel bins in latent space, ~1/25s per token, 128-dim feature vector)
+- **Decoder**: `[B, 8, T, 16]` → `[B, mel_bins, T*4]` mel spectrogram
+- **Vocoder**: HiFi-GAN-based, modified for stereo synthesis and upsampling (16 kHz mel → 24 kHz waveform, doubled generator capacity for stereo)
+
+**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 3** (Gemma 3-12B) as the multilingual text encoder backbone, located in [`src/ltx_core/text_encoders/gemma/`](src/ltx_core/text_encoders/gemma/). Advanced text understanding is critical not only for global language support but for the phonetic and semantic accuracy of generated speech.
+
+### Text Encoder Architecture
+
+The text conditioning pipeline consists of three stages:
+
+1. **Gemma 3 Backbone**: Decoder-only LLM processes text tokens → embeddings across all layers `[B, T, D, L]`
+2. **Multi-Layer Feature Extractor**: Aggregates features from all decoder layers (not just final layer), applies mean-centered scaling, flattens to `[B, T, D×L]`, and projects via learnable matrix W (jointly optimized with LTX-2, LLM weights frozen)
+3. **Text Connector**: Bidirectional transformer blocks with learnable registers (replacing padded positions, also referred to as "thinking tokens" in the paper) for contextual mixing. Separate connectors for video and audio streams (`Embeddings1DConnector`)
+
+**Encoders**:
+
+- `AVGemmaTextEncoderModel`: Audio-video generation (two connectors → `AVGemmaEncoderOutput` with separate video/audio contexts)
+- `VideoGemmaTextEncoderModel`: Video-only generation (single connector → `VideoGemmaEncoderOutput`)
+
+### 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 and enables the model to synthesize speech that is synchronized with visual lip movement while being natural in cadence, accent, and emotional tone.
+
+**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/src/ltx_core/conditioning/types/__init__.py

@@ -0,0 +1,13 @@
+"""Conditioning type implementations."""
+
+from ltx_core.conditioning.types.attention_strength_wrapper import ConditioningItemAttentionStrengthWrapper
+from ltx_core.conditioning.types.keyframe_cond import VideoConditionByKeyframeIndex
+from ltx_core.conditioning.types.latent_cond import VideoConditionByLatentIndex
+from ltx_core.conditioning.types.reference_video_cond import VideoConditionByReferenceLatent
+
+__all__ = [
+    "ConditioningItemAttentionStrengthWrapper",
+    "VideoConditionByKeyframeIndex",
+    "VideoConditionByLatentIndex",
+    "VideoConditionByReferenceLatent",
+]

packages/ltx-core/src/ltx_core/conditioning/types/attention_strength_wrapper.py

@@ -0,0 +1,71 @@
+"""Wrapper conditioning item that adds attention masking to any inner conditioning."""
+
+from dataclasses import replace
+
+import torch
+
+from ltx_core.conditioning.item import ConditioningItem
+from ltx_core.conditioning.mask_utils import update_attention_mask
+from ltx_core.tools import LatentTools
+from ltx_core.types import LatentState
+
+
+class ConditioningItemAttentionStrengthWrapper(ConditioningItem):
+    """Wraps a conditioning item to add an attention mask for its tokens.
+    Separates the *attention-masking* concern from the underlying conditioning
+    logic (token layout, positional encoding, denoise strength).  The inner
+    conditioning item appends tokens to the latent sequence as usual, and this
+    wrapper then builds or updates the self-attention mask so that the newly
+    added tokens interact with the noisy tokens according to *attention_mask*.
+    Args:
+        conditioning: Any conditioning item that appends tokens to the latent.
+        attention_mask: Per-token attention weight controlling how strongly the
+            new conditioning tokens attend to/from noisy tokens.  Can be a
+            scalar (float) applied uniformly, or a tensor of shape ``(B, M)``
+            for spatial control, where ``M = F * H * W`` is the number of
+            patchified conditioning tokens.  Values in ``[0, 1]``.
+    Example::
+        cond = ConditioningItemAttentionStrengthWrapper(
+            VideoConditionByReferenceLatent(latent=ref, strength=1.0),
+            attention_mask=0.5,
+        )
+        state = cond.apply_to(latent_state, latent_tools)
+    """
+
+    def __init__(
+        self,
+        conditioning: ConditioningItem,
+        attention_mask: float | torch.Tensor,
+    ):
+        self.conditioning = conditioning
+        self.attention_mask = attention_mask
+
+    def apply_to(
+        self,
+        latent_state: LatentState,
+        latent_tools: LatentTools,
+    ) -> LatentState:
+        """Apply inner conditioning, then build the attention mask for its tokens."""
+        # Snapshot the original state for mask building
+        original_state = latent_state
+
+        # Inner conditioning appends tokens (positions, denoise mask, etc.)
+        new_state = self.conditioning.apply_to(latent_state, latent_tools)
+
+        num_new_tokens = new_state.latent.shape[1] - original_state.latent.shape[1]
+        if num_new_tokens == 0:
+            return new_state
+
+        # Build the attention mask using the *original* state as the reference
+        # so that the block structure is computed correctly.
+        new_attention_mask = update_attention_mask(
+            latent_state=original_state,
+            attention_mask=self.attention_mask,
+            num_noisy_tokens=latent_tools.target_shape.token_count(),
+            num_new_tokens=num_new_tokens,
+            batch_size=new_state.latent.shape[0],
+            device=new_state.latent.device,
+            dtype=new_state.latent.dtype,
+        )
+
+        return replace(new_state, attention_mask=new_attention_mask)

packages/ltx-core/src/ltx_core/conditioning/types/keyframe_cond.py

@@ -0,0 +1,70 @@
+import torch
+
+from ltx_core.components.patchifiers import get_pixel_coords
+from ltx_core.conditioning.item import ConditioningItem
+from ltx_core.conditioning.mask_utils import update_attention_mask
+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.
+    To add attention masking, wrap with :class:`ConditioningItemAttentionStrengthWrapper`.
+    Args:
+        keyframes: Keyframe latents [B, C, F, H, W].
+        frame_idx: Frame index offset for positional encoding.
+        strength: Conditioning strength (1.0 = clean, 0.0 = fully denoised).
+    """
+
+    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,
+        )
+
+        new_attention_mask = update_attention_mask(
+            latent_state=latent_state,
+            attention_mask=None,
+            num_noisy_tokens=latent_tools.target_shape.token_count(),
+            num_new_tokens=tokens.shape[1],
+            batch_size=tokens.shape[0],
+            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),
+            attention_mask=new_attention_mask,
+        )

packages/ltx-core/src/ltx_core/conditioning/types/latent_cond.py

@@ -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/conditioning/types/reference_video_cond.py

@@ -0,0 +1,91 @@
+"""Reference video conditioning for IC-LoRA inference."""
+
+import torch
+
+from ltx_core.components.patchifiers import get_pixel_coords
+from ltx_core.conditioning.item import ConditioningItem
+from ltx_core.conditioning.mask_utils import update_attention_mask
+from ltx_core.tools import VideoLatentTools
+from ltx_core.types import LatentState, VideoLatentShape
+
+
+class VideoConditionByReferenceLatent(ConditioningItem):
+    """
+    Conditions video generation on a reference video latent for IC-LoRA inference.
+    IC-LoRAs are trained by concatenating reference (control signal) and target tokens,
+    learning to attend across both. This class replicates that setup at inference by
+    appending reference tokens to the latent sequence.
+    IC-LoRAs can be trained with lower-resolution references than the target (e.g., 384px
+    reference for 768px output) for efficiency and better generalization. The
+    `downscale_factor` scales reference positions to match target coordinates, preserving
+    the learned positional relationships. This must match the factor used during training
+    (stored in LoRA metadata).
+    To add attention masking, wrap with :class:`ConditioningItemAttentionStrengthWrapper`.
+    Args:
+        latent: Reference video latents [B, C, F, H, W]
+        downscale_factor: Target/reference resolution ratio (e.g., 2 = half-resolution
+            reference). Spatial positions are scaled by this factor.
+        strength: Conditioning strength. 1.0 = full (reference kept clean),
+            0.0 = none (reference denoised). Default 1.0.
+    """
+
+    def __init__(
+        self,
+        latent: torch.Tensor,
+        downscale_factor: int = 1,
+        strength: float = 1.0,
+    ):
+        self.latent = latent
+        self.downscale_factor = downscale_factor
+        self.strength = strength
+
+    def apply_to(
+        self,
+        latent_state: LatentState,
+        latent_tools: VideoLatentTools,
+    ) -> LatentState:
+        """Append reference video tokens with scaled positions."""
+        tokens = latent_tools.patchifier.patchify(self.latent)
+
+        # Compute positions for the reference video's actual dimensions
+        latent_coords = latent_tools.patchifier.get_patch_grid_bounds(
+            output_shape=VideoLatentShape.from_torch_shape(self.latent.shape),
+            device=self.latent.device,
+        )
+        positions = get_pixel_coords(
+            latent_coords=latent_coords,
+            scale_factors=latent_tools.scale_factors,
+            causal_fix=latent_tools.causal_fix,
+        )
+        positions = positions.to(dtype=torch.float32)
+        positions[:, 0, ...] /= latent_tools.fps
+
+        # Scale spatial positions to match target coordinate space
+        if self.downscale_factor != 1:
+            positions[:, 1, ...] *= self.downscale_factor  # height axis
+            positions[:, 2, ...] *= self.downscale_factor  # width axis
+
+        denoise_mask = torch.full(
+            size=(*tokens.shape[:2], 1),
+            fill_value=1.0 - self.strength,
+            device=self.latent.device,
+            dtype=self.latent.dtype,
+        )
+
+        new_attention_mask = update_attention_mask(
+            latent_state=latent_state,
+            attention_mask=None,
+            num_noisy_tokens=latent_tools.target_shape.token_count(),
+            num_new_tokens=tokens.shape[1],
+            batch_size=tokens.shape[0],
+            device=self.latent.device,
+            dtype=self.latent.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),
+            attention_mask=new_attention_mask,
+        )

packages/ltx-core/src/ltx_core/model/audio_vae/causality_axis.py

@@ -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/common/__init__.py

@@ -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

@@ -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/transformer/__init__.py

@@ -0,0 +1,18 @@
+"""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,
+    LTXModelConfigurator,
+    LTXVideoOnlyModelConfigurator,
+)
+
+__all__ = [
+    "LTXV_MODEL_COMFY_RENAMING_MAP",
+    "LTXModel",
+    "LTXModelConfigurator",
+    "LTXVideoOnlyModelConfigurator",
+    "Modality",
+    "X0Model",
+]

packages/ltx-core/src/ltx_core/model/transformer/adaln.py

@@ -0,0 +1,45 @@
+from typing import Optional, Tuple
+
+import torch
+
+from ltx_core.model.transformer.timestep_embedding import PixArtAlphaCombinedTimestepSizeEmbeddings
+
+# Number of AdaLN modulation parameters per transformer block.
+# Base: 2 params (shift + scale) x 3 norms (self-attn, feed-forward, output).
+ADALN_NUM_BASE_PARAMS = 6
+# Cross-attention AdaLN adds 3 more (scale, shift, gate) for the CA norm.
+ADALN_NUM_CROSS_ATTN_PARAMS = 3
+
+
+def adaln_embedding_coefficient(cross_attention_adaln: bool) -> int:
+    """Total number of AdaLN parameters per block."""
+    return ADALN_NUM_BASE_PARAMS + (ADALN_NUM_CROSS_ATTN_PARAMS if cross_attention_adaln else 0)
+
+
+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

@@ -0,0 +1,249 @@
+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,
+        apply_gated_attention: bool = False,
+    ) -> 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)
+
+        # Optional per-head gating
+        if apply_gated_attention:
+            self.to_gate_logits = torch.nn.Linear(query_dim, heads, bias=True)
+        else:
+            self.to_gate_logits = None
+
+        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,
+        perturbation_mask: torch.Tensor | None = None,
+        all_perturbed: bool = False,
+    ) -> torch.Tensor:
+        """Multi-head attention with optional RoPE, perturbation masking, and per-head gating.
+        When ``perturbation_mask`` is all zeros, the expensive query/key path
+        (linear projections, RMSNorm, RoPE) is skipped entirely and only the
+        value projection is used as a pass-through.
+        Args:
+            x: Query input tensor of shape ``(B, T, query_dim)``.
+            context: Key/value context tensor of shape ``(B, S, context_dim)``.
+                Falls back to ``x`` (self-attention) when *None*.
+            mask: Optional attention mask. Interpretation depends on the attention
+                backend (additive bias for xformers/PyTorch SDPA).
+            pe: Rotary positional embeddings applied to both ``q`` and ``k``.
+            k_pe: Separate rotary positional embeddings for ``k`` only. When
+                *None*, ``pe`` is reused for keys.
+            perturbation_mask: Optional mask in ``[0, 1]`` that
+                blends the attention output with the raw value projection:
+                ``out = attn_out * mask + v * (1 - mask)``.
+                **1** keeps the full attention output, **0** bypasses attention
+                and passes the value projection through unchanged.
+                *None* or all-ones means standard attention; all-zeros skips
+                the query/key path entirely for efficiency.
+            all_perturbed: Whether all perturbations are active for this block.
+        Returns:
+            Output tensor of shape ``(B, T, query_dim)``.
+        """
+        context = x if context is None else context
+        use_attention = not all_perturbed
+
+        v = self.to_v(context)
+
+        if not use_attention:
+            out = v
+        else:
+            q = self.to_q(x)
+            k = self.to_k(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)
+
+            out = self.attention_function(q, k, v, self.heads, mask)  # (B, T, H*D)
+
+            if perturbation_mask is not None:
+                out = out * perturbation_mask + v * (1 - perturbation_mask)
+
+        # Apply per-head gating if enabled
+        if self.to_gate_logits is not None:
+            gate_logits = self.to_gate_logits(x)  # (B, T, H)
+            b, t, _ = out.shape
+            # Reshape to (B, T, H, D) for per-head gating
+            out = out.view(b, t, self.heads, self.dim_head)
+            # Apply gating: 2 * sigmoid(x) so that zero-init gives identity (2 * 0.5 = 1.0)
+            gates = 2.0 * torch.sigmoid(gate_logits)  # (B, T, H)
+            out = out * gates.unsqueeze(-1)  # (B, T, H, D) * (B, T, H, 1)
+            # Reshape back to (B, T, H*D)
+            out = out.view(b, t, self.heads * self.dim_head)
+
+        return self.to_out(out)

packages/ltx-core/src/ltx_core/model/transformer/gelu_approx.py

@@ -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

@@ -0,0 +1,40 @@
+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.
+    Attributes:
+        latent: Patchified latent tokens, shape ``(B, T, D)`` where *B* is
+            the batch size, *T* is the total number of tokens (noisy +
+            conditioning), and *D* is the input dimension.
+        timesteps: Per-token timestep embeddings, shape ``(B, T)``.
+        positions: Positional coordinates, shape ``(B, 3, T)`` for video
+            (time, height, width) or ``(B, 1, T)`` for audio.
+        context: Text conditioning embeddings from the prompt encoder.
+        enabled: Whether this modality is active in the current forward pass.
+        context_mask: Optional mask for the text context tokens.
+        attention_mask: Optional 2-D self-attention mask, shape ``(B, T, T)``.
+            Values in ``[0, 1]`` where ``1`` = full attention and ``0`` = no
+            attention. ``None`` means unrestricted (full) attention between
+            all tokens. Built incrementally by conditioning items; see
+            :class:`~ltx_core.conditioning.types.attention_strength_wrapper.ConditioningItemAttentionStrengthWrapper`.
+    """
+
+    latent: (
+        torch.Tensor
+    )  # Shape: (B, T, D) where B is the batch size, T is the number of tokens, and D is input dimension
+    sigma: torch.Tensor  # Shape: (B,). Current sigma value, used for cross-attention timestep calculation.
+    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
+    attention_mask: torch.Tensor | None = None

packages/ltx-core/src/ltx_core/model/transformer/model.py

@@ -0,0 +1,486 @@
+from enum import Enum
+
+import torch
+
+from ltx_core.guidance.perturbations import BatchedPerturbationConfig
+from ltx_core.model.transformer.adaln import AdaLayerNormSingle, adaln_embedding_coefficient
+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.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,
+        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,
+        apply_gated_attention: bool = False,
+        caption_projection: torch.nn.Module | None = None,
+        audio_caption_projection: torch.nn.Module | None = None,
+        cross_attention_adaln: bool = False,
+    ):
+        super().__init__()
+        self._enable_gradient_checkpointing = False
+        self.cross_attention_adaln = cross_attention_adaln
+        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,
+                norm_eps=norm_eps,
+                caption_projection=caption_projection,
+            )
+
+        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,
+                norm_eps=norm_eps,
+                caption_projection=audio_caption_projection,
+            )
+
+        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,
+            apply_gated_attention=apply_gated_attention,
+        )
+
+    @property
+    def _adaln_embedding_coefficient(self) -> int:
+        return adaln_embedding_coefficient(self.cross_attention_adaln)
+
+    def _init_video(
+        self,
+        in_channels: int,
+        out_channels: int,
+        norm_eps: float,
+        caption_projection: torch.nn.Module | None = None,
+    ) -> None:
+        """Initialize video-specific components."""
+        # Video input components
+        self.patchify_proj = torch.nn.Linear(in_channels, self.inner_dim, bias=True)
+        if caption_projection is not None:
+            self.caption_projection = caption_projection
+
+        self.adaln_single = AdaLayerNormSingle(self.inner_dim, embedding_coefficient=self._adaln_embedding_coefficient)
+
+        self.prompt_adaln_single = (
+            AdaLayerNormSingle(self.inner_dim, embedding_coefficient=2) if self.cross_attention_adaln else None
+        )
+
+        # 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,
+        norm_eps: float,
+        caption_projection: torch.nn.Module | None = None,
+    ) -> None:
+        """Initialize audio-specific components."""
+
+        # Audio input components
+        self.audio_patchify_proj = torch.nn.Linear(in_channels, self.audio_inner_dim, bias=True)
+        if caption_projection is not None:
+            self.audio_caption_projection = caption_projection
+
+        self.audio_adaln_single = AdaLayerNormSingle(
+            self.audio_inner_dim,
+            embedding_coefficient=self._adaln_embedding_coefficient,
+        )
+
+        self.audio_prompt_adaln_single = (
+            AdaLayerNormSingle(self.audio_inner_dim, embedding_coefficient=2) if self.cross_attention_adaln else None
+        )
+
+        # 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,
+                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,
+                caption_projection=getattr(self, "caption_projection", None),
+                prompt_adaln=getattr(self, "prompt_adaln_single", None),
+            )
+            self.audio_args_preprocessor = MultiModalTransformerArgsPreprocessor(
+                patchify_proj=self.audio_patchify_proj,
+                adaln=self.audio_adaln_single,
+                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,
+                caption_projection=getattr(self, "audio_caption_projection", None),
+                prompt_adaln=getattr(self, "audio_prompt_adaln_single", None),
+            )
+        elif self.model_type.is_video_enabled():
+            self.video_args_preprocessor = TransformerArgsPreprocessor(
+                patchify_proj=self.patchify_proj,
+                adaln=self.adaln_single,
+                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,
+                caption_projection=getattr(self, "caption_projection", None),
+                prompt_adaln=getattr(self, "prompt_adaln_single", None),
+            )
+        elif self.model_type.is_audio_enabled():
+            self.audio_args_preprocessor = TransformerArgsPreprocessor(
+                patchify_proj=self.audio_patchify_proj,
+                adaln=self.audio_adaln_single,
+                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,
+                caption_projection=getattr(self, "audio_caption_projection", None),
+                prompt_adaln=getattr(self, "audio_prompt_adaln_single", None),
+            )
+
+    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,
+        apply_gated_attention: bool,
+    ) -> 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,
+                apply_gated_attention=apply_gated_attention,
+                cross_attention_adaln=self.cross_attention_adaln,
+            )
+            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,
+                apply_gated_attention=apply_gated_attention,
+                cross_attention_adaln=self.cross_attention_adaln,
+            )
+            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, audio) if video is not None else None
+        audio_args = self.audio_args_preprocessor.prepare(audio, video) 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

@@ -0,0 +1,152 @@
+import torch
+
+from ltx_core.loader.sd_ops import 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.model.transformer.text_projection import create_caption_projection
+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:
+        # Build caption projections for 19B models (projection handled in transformer).
+        caption_projection, audio_caption_projection = _build_caption_projections(config, is_av=True)
+
+        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")),
+            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",
+            apply_gated_attention=config.get("apply_gated_attention", False),
+            caption_projection=caption_projection,
+            audio_caption_projection=audio_caption_projection,
+            cross_attention_adaln=config.get("cross_attention_adaln", False),
+        )
+
+
+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:
+        # Build caption projection for 19B model (projection handled in transformer).
+        caption_projection, _ = _build_caption_projections(config, is_av=False)
+
+        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")),
+            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",
+            apply_gated_attention=config.get("apply_gated_attention", False),
+            caption_projection=caption_projection,
+            cross_attention_adaln=config.get("cross_attention_adaln", False),
+        )
+
+
+def _build_caption_projections(
+    config: dict,
+    is_av: bool,
+) -> tuple[torch.nn.Module | None, torch.nn.Module | None]:
+    """Build caption projections for the transformer when projection is NOT in the text encoder.
+    19B models: projection is in the transformer (caption_proj_before_connector=False).
+    22B models: projection is in the text encoder, so no projections are created here.
+    Args:
+        config: Full model config dict (must contain "transformer" key).
+        is_av: Whether this is an audio-video model. When False, audio projection is skipped.
+    Returns:
+        Tuple of (video_caption_projection, audio_caption_projection), both None for 22B models.
+    """
+    transformer_config = config.get("transformer", {})
+    if transformer_config.get("caption_proj_before_connector", False):
+        return None, None
+
+    with torch.device("meta"):
+        caption_projection = create_caption_projection(transformer_config)
+        audio_caption_projection = create_caption_projection(transformer_config, audio=True) if is_av else None
+    return caption_projection, audio_caption_projection
+
+
+LTXV_MODEL_COMFY_RENAMING_MAP = (
+    SDOps("LTXV_MODEL_COMFY_PREFIX_MAP")
+    .with_matching(prefix="model.diffusion_model.")
+    .with_replacement("model.diffusion_model.", "")
+)

packages/ltx-core/src/ltx_core/model/transformer/rope.py

@@ -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

@@ -0,0 +1,38 @@
+import torch
+
+
+class PixArtAlphaTextProjection(torch.nn.Module):
+    """
+    Projects caption embeddings using dual linear layers.
+    Flow: linear_1 → activation → linear_2
+    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
+
+
+def create_caption_projection(transformer_config: dict, audio: bool = False) -> PixArtAlphaTextProjection:
+    """Create a caption projection for the transformer (V1/19B only)."""
+    caption_channels = transformer_config["caption_channels"]
+    if audio:
+        inner_dim = transformer_config["audio_num_attention_heads"] * transformer_config["audio_attention_head_dim"]
+    else:
+        inner_dim = transformer_config["num_attention_heads"] * transformer_config["attention_head_dim"]
+    return PixArtAlphaTextProjection(in_features=caption_channels, hidden_size=inner_dim)

packages/ltx-core/src/ltx_core/model/transformer/timestep_embedding.py

@@ -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

@@ -0,0 +1,398 @@
+from dataclasses import dataclass, replace
+
+import torch
+
+from ltx_core.guidance.perturbations import BatchedPerturbationConfig, PerturbationType
+from ltx_core.model.transformer.adaln import adaln_embedding_coefficient
+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
+    apply_gated_attention: bool = False
+    cross_attention_adaln: bool = False
+
+
+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,
+                apply_gated_attention=video.apply_gated_attention,
+            )
+            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,
+                apply_gated_attention=video.apply_gated_attention,
+            )
+            self.ff = FeedForward(video.dim, dim_out=video.dim)
+            video_sst_size = adaln_embedding_coefficient(video.cross_attention_adaln)
+            self.scale_shift_table = torch.nn.Parameter(torch.empty(video_sst_size, 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,
+                apply_gated_attention=audio.apply_gated_attention,
+            )
+            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,
+                apply_gated_attention=audio.apply_gated_attention,
+            )
+            self.audio_ff = FeedForward(audio.dim, dim_out=audio.dim)
+            audio_sst_size = adaln_embedding_coefficient(audio.cross_attention_adaln)
+            self.audio_scale_shift_table = torch.nn.Parameter(torch.empty(audio_sst_size, 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,
+                apply_gated_attention=video.apply_gated_attention,
+            )
+
+            # 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,
+                apply_gated_attention=audio.apply_gated_attention,
+            )
+
+            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.cross_attention_adaln = (video is not None and video.cross_attention_adaln) or (
+            audio is not None and audio.cross_attention_adaln
+        )
+
+        if self.cross_attention_adaln and video is not None:
+            self.prompt_scale_shift_table = torch.nn.Parameter(torch.empty(2, video.dim))
+        if self.cross_attention_adaln and audio is not None:
+            self.audio_prompt_scale_shift_table = torch.nn.Parameter(torch.empty(2, audio.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,
+        scale_shift_indices: slice,
+        num_scale_shift_values: int = 4,
+    ) -> tuple[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, scale_shift_indices
+        )
+        gate_ada_values = self.get_ada_values(
+            scale_shift_table[num_scale_shift_values:, :], batch_size, gate_timestep, slice(None, None)
+        )
+
+        scale, shift = (t.squeeze(2) for t in scale_shift_ada_values)
+        (gate,) = (t.squeeze(2) for t in gate_ada_values)
+
+        return scale, shift, gate
+
+    def _apply_text_cross_attention(
+        self,
+        x: torch.Tensor,
+        context: torch.Tensor,
+        attn: AttentionCallable,
+        scale_shift_table: torch.Tensor,
+        prompt_scale_shift_table: torch.Tensor | None,
+        timestep: torch.Tensor,
+        prompt_timestep: torch.Tensor | None,
+        context_mask: torch.Tensor | None,
+        cross_attention_adaln: bool = False,
+    ) -> torch.Tensor:
+        """Apply text cross-attention, with optional AdaLN modulation."""
+        if cross_attention_adaln:
+            shift_q, scale_q, gate = self.get_ada_values(scale_shift_table, x.shape[0], timestep, slice(6, 9))
+            return apply_cross_attention_adaln(
+                x,
+                context,
+                attn,
+                shift_q,
+                scale_q,
+                gate,
+                prompt_scale_shift_table,
+                prompt_timestep,
+                context_mask,
+                self.norm_eps,
+            )
+        return attn(rms_norm(x, eps=self.norm_eps), context=context, mask=context_mask)
+
+    def forward(  # noqa: PLR0915
+        self,
+        video: TransformerArgs | None,
+        audio: TransformerArgs | None,
+        perturbations: BatchedPerturbationConfig | None = None,
+    ) -> tuple[TransformerArgs | None, TransformerArgs | None]:
+        if video is None and audio is None:
+            raise ValueError("At least one of video or audio must be provided")
+
+        batch_size = (video or audio).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 ax.numel() > 0)
+        run_v2a = run_ax and (video is not None 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)
+            )
+            norm_vx = rms_norm(vx, eps=self.norm_eps) * (1 + vscale_msa) + vshift_msa
+            del vshift_msa, vscale_msa
+
+            all_perturbed = perturbations.all_in_batch(PerturbationType.SKIP_VIDEO_SELF_ATTN, self.idx)
+            none_perturbed = not perturbations.any_in_batch(PerturbationType.SKIP_VIDEO_SELF_ATTN, self.idx)
+            v_mask = (
+                perturbations.mask_like(PerturbationType.SKIP_VIDEO_SELF_ATTN, self.idx, vx)
+                if not all_perturbed and not none_perturbed
+                else None
+            )
+            vx = (
+                vx
+                + self.attn1(
+                    norm_vx,
+                    pe=video.positional_embeddings,
+                    mask=video.self_attention_mask,
+                    perturbation_mask=v_mask,
+                    all_perturbed=all_perturbed,
+                )
+                * vgate_msa
+            )
+            del vgate_msa, norm_vx, v_mask
+            vx = vx + self._apply_text_cross_attention(
+                vx,
+                video.context,
+                self.attn2,
+                self.scale_shift_table,
+                getattr(self, "prompt_scale_shift_table", None),
+                video.timesteps,
+                video.prompt_timestep,
+                video.context_mask,
+                cross_attention_adaln=self.cross_attention_adaln,
+            )
+
+        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)
+            )
+
+            norm_ax = rms_norm(ax, eps=self.norm_eps) * (1 + ascale_msa) + ashift_msa
+            del ashift_msa, ascale_msa
+            all_perturbed = perturbations.all_in_batch(PerturbationType.SKIP_AUDIO_SELF_ATTN, self.idx)
+            none_perturbed = not perturbations.any_in_batch(PerturbationType.SKIP_AUDIO_SELF_ATTN, self.idx)
+            a_mask = (
+                perturbations.mask_like(PerturbationType.SKIP_AUDIO_SELF_ATTN, self.idx, ax)
+                if not all_perturbed and not none_perturbed
+                else None
+            )
+            ax = (
+                ax
+                + self.audio_attn1(
+                    norm_ax,
+                    pe=audio.positional_embeddings,
+                    mask=audio.self_attention_mask,
+                    perturbation_mask=a_mask,
+                    all_perturbed=all_perturbed,
+                )
+                * agate_msa
+            )
+            del agate_msa, norm_ax, a_mask
+            ax = ax + self._apply_text_cross_attention(
+                ax,
+                audio.context,
+                self.audio_attn2,
+                self.audio_scale_shift_table,
+                getattr(self, "audio_prompt_scale_shift_table", None),
+                audio.timesteps,
+                audio.prompt_timestep,
+                audio.context_mask,
+                cross_attention_adaln=self.cross_attention_adaln,
+            )
+
+        # 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)
+
+            if run_a2v and not perturbations.all_in_batch(PerturbationType.SKIP_A2V_CROSS_ATTN, self.idx):
+                scale_ca_video_a2v, shift_ca_video_a2v, 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,
+                    slice(0, 2),
+                )
+                vx_scaled = vx_norm3 * (1 + scale_ca_video_a2v) + shift_ca_video_a2v
+                del scale_ca_video_a2v, shift_ca_video_a2v
+
+                scale_ca_audio_a2v, shift_ca_audio_a2v, _ = 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,
+                    slice(0, 2),
+                )
+                ax_scaled = ax_norm3 * (1 + scale_ca_audio_a2v) + shift_ca_audio_a2v
+                del scale_ca_audio_a2v, shift_ca_audio_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
+                )
+                del gate_out_a2v, a2v_mask, vx_scaled, ax_scaled
+
+            if run_v2a and not perturbations.all_in_batch(PerturbationType.SKIP_V2A_CROSS_ATTN, self.idx):
+                scale_ca_audio_v2a, shift_ca_audio_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,
+                    slice(2, 4),
+                )
+                ax_scaled = ax_norm3 * (1 + scale_ca_audio_v2a) + shift_ca_audio_v2a
+                del scale_ca_audio_v2a, shift_ca_audio_v2a
+                scale_ca_video_v2a, shift_ca_video_v2a, _ = 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,
+                    slice(2, 4),
+                )
+                vx_scaled = vx_norm3 * (1 + scale_ca_video_v2a) + shift_ca_video_v2a
+                del scale_ca_video_v2a, shift_ca_video_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_v2a, v2a_mask, ax_scaled, vx_scaled
+
+            del vx_norm3, ax_norm3
+
+        if run_vx:
+            vshift_mlp, vscale_mlp, vgate_mlp = self.get_ada_values(
+                self.scale_shift_table, vx.shape[0], video.timesteps, slice(3, 6)
+            )
+            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, vx_scaled
+
+        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, 6)
+            )
+            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, ax_scaled
+
+        return replace(video, x=vx) if video is not None else None, replace(audio, x=ax) if audio is not None else None
+
+
+def apply_cross_attention_adaln(
+    x: torch.Tensor,
+    context: torch.Tensor,
+    attn: AttentionCallable,
+    q_shift: torch.Tensor,
+    q_scale: torch.Tensor,
+    q_gate: torch.Tensor,
+    prompt_scale_shift_table: torch.Tensor,
+    prompt_timestep: torch.Tensor,
+    context_mask: torch.Tensor | None = None,
+    norm_eps: float = 1e-6,
+) -> torch.Tensor:
+    batch_size = x.shape[0]
+    shift_kv, scale_kv = (
+        prompt_scale_shift_table[None, None].to(device=x.device, dtype=x.dtype)
+        + prompt_timestep.reshape(batch_size, prompt_timestep.shape[1], 2, -1)
+    ).unbind(dim=2)
+    attn_input = rms_norm(x, eps=norm_eps) * (1 + q_scale) + q_shift
+    encoder_hidden_states = context * (1 + scale_kv) + shift_kv
+    return attn(attn_input, context=encoder_hidden_states, mask=context_mask) * q_gate

packages/ltx-core/src/ltx_core/model/transformer/transformer_args.py

@@ -0,0 +1,297 @@
+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,
+)
+
+
+@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
+    prompt_timestep: torch.Tensor | None = None
+    self_attention_mask: torch.Tensor | None = (
+        None  # Additive log-space self-attention bias (B, 1, T, T), None = full attention
+    )
+
+
+class TransformerArgsPreprocessor:
+    def __init__(  # noqa: PLR0913
+        self,
+        patchify_proj: torch.nn.Linear,
+        adaln: AdaLayerNormSingle,
+        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,
+        caption_projection: torch.nn.Module | None = None,
+        prompt_adaln: AdaLayerNormSingle | None = None,
+    ) -> None:
+        self.patchify_proj = patchify_proj
+        self.adaln = adaln
+        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
+        self.caption_projection = caption_projection
+        self.prompt_adaln = prompt_adaln
+
+    def _prepare_timestep(
+        self, timestep: torch.Tensor, adaln: AdaLayerNormSingle, batch_size: int, hidden_dtype: torch.dtype
+    ) -> tuple[torch.Tensor, torch.Tensor]:
+        """Prepare timestep embeddings."""
+        timestep_scaled = timestep * self.timestep_scale_multiplier
+        timestep, embedded_timestep = adaln(
+            timestep_scaled.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,
+    ) -> torch.Tensor:
+        """Prepare context for transformer blocks."""
+        if self.caption_projection is not None:
+            context = self.caption_projection(context)
+        batch_size = x.shape[0]
+        return context.view(batch_size, -1, x.shape[-1])
+
+    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_self_attention_mask(
+        self, attention_mask: torch.Tensor | None, x_dtype: torch.dtype
+    ) -> torch.Tensor | None:
+        """Prepare self-attention mask by converting [0,1] values to additive log-space bias.
+        Input shape: (B, T, T) with values in [0, 1].
+        Output shape: (B, 1, T, T) with 0.0 for full attention and a large negative value
+        for masked positions.
+        Positions with attention_mask <= 0 are fully masked (mapped to the dtype's minimum
+        representable value). Strictly positive entries are converted via log-space for
+        smooth attenuation, with small values clamped for numerical stability.
+        Returns None if input is None (no masking).
+        """
+        if attention_mask is None:
+            return None
+
+        # Convert [0, 1] attention mask to additive log-space bias:
+        #   1.0 -> log(1.0) = 0.0  (no bias, full attention)
+        #   0.0 -> finfo.min        (fully masked)
+        finfo = torch.finfo(x_dtype)
+        eps = finfo.tiny
+
+        bias = torch.full_like(attention_mask, finfo.min, dtype=x_dtype)
+        positive = attention_mask > 0
+        if positive.any():
+            bias[positive] = torch.log(attention_mask[positive].clamp(min=eps)).to(x_dtype)
+
+        return bias.unsqueeze(1)  # (B, 1, T, T) for head broadcast
+
+    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,
+        cross_modality: Modality | None = None,  # noqa: ARG002
+    ) -> TransformerArgs:
+        x = self.patchify_proj(modality.latent)
+        batch_size = x.shape[0]
+        timestep, embedded_timestep = self._prepare_timestep(
+            modality.timesteps, self.adaln, batch_size, modality.latent.dtype
+        )
+        prompt_timestep = None
+        if self.prompt_adaln is not None:
+            prompt_timestep, _ = self._prepare_timestep(
+                modality.sigma, self.prompt_adaln, batch_size, modality.latent.dtype
+            )
+        context = self._prepare_context(modality.context, x)
+        attention_mask = self._prepare_attention_mask(modality.context_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,
+        )
+        self_attention_mask = self._prepare_self_attention_mask(modality.attention_mask, 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,
+            prompt_timestep=prompt_timestep,
+            self_attention_mask=self_attention_mask,
+        )
+
+
+class MultiModalTransformerArgsPreprocessor:
+    def __init__(  # noqa: PLR0913
+        self,
+        patchify_proj: torch.nn.Linear,
+        adaln: AdaLayerNormSingle,
+        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,
+        caption_projection: torch.nn.Module | None = None,
+        prompt_adaln: AdaLayerNormSingle | None = None,
+    ) -> None:
+        self.simple_preprocessor = TransformerArgsPreprocessor(
+            patchify_proj=patchify_proj,
+            adaln=adaln,
+            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,
+            caption_projection=caption_projection,
+            prompt_adaln=prompt_adaln,
+        )
+        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,
+        cross_modality: Modality | None = None,
+    ) -> TransformerArgs:
+        transformer_args = self.simple_preprocessor.prepare(modality)
+        if cross_modality is None:
+            return transformer_args
+
+        if cross_modality.sigma.numel() > 1:
+            if cross_modality.sigma.shape[0] != modality.timesteps.shape[0]:
+                raise ValueError("Cross modality sigma must have the same batch size as the modality")
+            if cross_modality.sigma.ndim != 1:
+                raise ValueError("Cross modality sigma must be a 1D tensor")
+
+        cross_timestep = cross_modality.sigma.view(
+            modality.timesteps.shape[0], 1, *[1] * len(modality.timesteps.shape[2:])
+        )
+
+        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=cross_timestep,
+            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 | None,
+        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/video_vae/__init__.py

@@ -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

@@ -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/model_configurator.py

@@ -0,0 +1,79 @@
+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)
+        spatial_padding_mode = PaddingModeType(config.get("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=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)
+        spatial_padding_mode = PaddingModeType(config.get("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)
+        base_channels = config.get("decoder_base_channels", 128)
+
+        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=spatial_padding_mode,
+            base_channels=base_channels,
+        )
+
+
+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/resnet.py

@@ -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.zeros(4, in_channels))
+
+    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/tiling.py

@@ -0,0 +1,291 @@
+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)
+
+
+def compute_rectangular_mask_1d(
+    length: int,
+    left_ramp: int,
+    right_ramp: int,
+) -> torch.Tensor:
+    """
+    Generate a 1D rectangular (pulse) mask.
+    Args:
+        length: Output length of the mask.
+        left_ramp: Number of elements at the start of the mask to set to 0.
+        right_ramp: Number of elements at the end of the mask to set to 0.
+    Returns:
+        A 1D tensor of shape `(length,)` with values 0 or 1.
+    """
+    if length <= 0:
+        raise ValueError("Mask length must be positive.")
+
+    mask = torch.ones(length)
+    if left_ramp > 0:
+        mask[:left_ramp] = 0
+    if right_ramp > 0:
+        mask[-right_ramp:] = 0
+    return mask
+
+
+@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:
+    """Defines how a single dimension is split into overlapping intervals (tiles).
+    Each list has length N where N is the number of intervals. The i-th element
+    of each list describes the i-th interval.
+    Attributes:
+        starts: Start index of each interval (inclusive).
+        ends: End index of each interval (exclusive).
+        left_ramps: Length of the left blend ramp for each interval.
+            Used to create masks that fade in from 0 to 1.
+        right_ramps: Length of the right blend ramp for each interval.
+            Used to create masks that fade out from 1 to 0.
+    """
+
+    starts: List[int]
+    ends: List[int]
+    left_ramps: List[int]
+    right_ramps: List[int]
+
+
+@dataclass(frozen=True)
+class TensorTilingSpec:
+    """Specifies how a tensor of a given shape is split into intervals (tiles) along each dimension.
+    Attributes:
+        original_shape: Shape of the tensor being tiled.
+        dimension_intervals: Per-dimension intervals (starts, ends, ramps) for each axis.
+    """
+
+    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: TensorTilingSpec,
+    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(
+    tensor_shape: torch.Size,
+    splitters: List[SplitOperation],
+    mappers: List[MappingOperation],
+) -> List[Tile]:
+    if len(splitters) != len(tensor_shape):
+        raise ValueError(
+            f"Number of splitters must be equal to number of dimensions in tensor shape, "
+            f"got {len(splitters)} and {len(tensor_shape)}"
+        )
+    if len(mappers) != len(tensor_shape):
+        raise ValueError(
+            f"Number of mappers must be equal to number of dimensions in tensor shape, "
+            f"got {len(mappers)} and {len(tensor_shape)}"
+        )
+    intervals = [splitter(length) for splitter, length in zip(splitters, tensor_shape, strict=True)]
+    tiling_spec = TensorTilingSpec(original_shape=tensor_shape, dimension_intervals=tuple(intervals))
+    return create_tiles_from_intervals_and_mappers(tiling_spec, mappers)

packages/ltx-core/src/ltx_core/model/video_vae/video_vae.py

@@ -0,0 +1,1219 @@
+import logging
+from dataclasses import replace
+from typing import Any, Callable, Iterator, List, 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_rectangular_mask_1d,
+    compute_trapezoidal_mask_1d,
+    create_tiles,
+)
+from ltx_core.types import VIDEO_SCALE_FACTORS, SpatioTemporalScaleFactors, VideoLatentShape
+
+logger: logging.Logger = logging.getLogger(__name__)
+
+
+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 should be 1 + 8*k (e.g., 1, 9, 17, 25, 33...).
+                If not, the encoder crops the last frames to the nearest valid length.
+        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 (crop to nearest valid length if needed)
+        frames_count = sample.shape[2]
+        if ((frames_count - 1) % 8) != 0:
+            frames_to_crop = (frames_count - 1) % 8
+            logger.warning(
+                "Invalid number of frames %s for encode; cropping last %s frames to satisfy 1 + 8*k.",
+                frames_count,
+                frames_to_crop,
+            )
+            sample = sample[:, :, :-frames_to_crop, ...]
+
+        # 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 tiled_encode(
+        self,
+        video: torch.Tensor,
+        tiling_config: TilingConfig | None = None,
+    ) -> torch.Tensor:
+        """Encode video to latent using tiled processing of the given video tensor.
+        Device Handling:
+            - Input video can be on CPU or GPU
+            - Accumulation buffers are created on model's device
+            - Each tile is automatically moved to model's device before encoding
+            - Output latent is returned on model's device
+        Args:
+            video: Input video tensor (B, 3, F, H, W) in range [-1, 1]
+            tiling_config: Tiling configuration for the video tensor
+        Returns:
+            Latent tensor (B, 128, F', H', W') on model's device
+            where F' = 1 + (F-1)/8, H' = H/32, W' = W/32
+        """
+        # Detect model device and dtype
+        model_device = next(self.parameters()).device
+        model_dtype = next(self.parameters()).dtype
+
+        # Extract shape components
+        batch, _, frames, height, width = video.shape
+
+        # Check frame count and crop if needed
+        if (frames - 1) % VIDEO_SCALE_FACTORS.time != 0:
+            frames_to_crop = (frames - 1) % VIDEO_SCALE_FACTORS.time
+            logger.warning(
+                f"Number of frames {frames} of input video is not ({VIDEO_SCALE_FACTORS.time} * k + 1), "
+                f"last {frames_to_crop} frames will be cropped"
+            )
+            video = video[:, :, :-frames_to_crop, ...]
+            # Update frames after cropping
+            frames = video.shape[2]
+
+        # Calculate output latent shape (inverse of upscale)
+        latent_shape = VideoLatentShape(
+            batch=batch,
+            channels=self.latent_channels,  # 128 for standard VAE
+            frames=(frames - 1) // VIDEO_SCALE_FACTORS.time + 1,
+            height=height // VIDEO_SCALE_FACTORS.height,
+            width=width // VIDEO_SCALE_FACTORS.width,
+        )
+
+        # Prepare tiles (operates on VIDEO dimensions)
+        tiles = prepare_tiles_for_encoding(video, tiling_config)
+
+        # Initialize accumulation buffers on model device
+        latent_buffer = torch.zeros(
+            latent_shape.to_torch_shape(),
+            device=model_device,
+            dtype=model_dtype,
+        )
+        weights_buffer = torch.zeros_like(latent_buffer)
+
+        # Process each tile
+        for tile in tiles:
+            # Extract video tile from input (may be on CPU)
+            video_tile = video[tile.in_coords]
+
+            # Move tile to model device if needed
+            if video_tile.device != model_device or video_tile.dtype != model_dtype:
+                video_tile = video_tile.to(device=model_device, dtype=model_dtype)
+
+            # Encode tile to latent (output on model device)
+            latent_tile = self.forward(video_tile)
+
+            # Move blend mask to model device
+            mask = tile.blend_mask.to(
+                device=model_device,
+                dtype=model_dtype,
+            )
+
+            # Weighted accumulation in latent space
+            latent_buffer[tile.out_coords] += latent_tile * mask
+            weights_buffer[tile.out_coords] += mask
+
+            del latent_tile, mask, video_tile
+
+        # Normalize by accumulated weights
+        weights_buffer = weights_buffer.clamp(min=1e-8)
+        return latent_buffer / weights_buffer
+
+
+def prepare_tiles_for_encoding(
+    video: torch.Tensor,
+    tiling_config: TilingConfig | None = None,
+) -> List[Tile]:
+    """Prepare tiles for VAE encoding.
+    Args:
+        video: Input video tensor (B, 3, F, H, W) in range [-1, 1]
+        tiling_config: Tiling configuration for the video tensor
+    Returns:
+        List of tiles for the video tensor
+    """
+
+    splitters = [DEFAULT_SPLIT_OPERATION] * len(video.shape)
+    mappers = [DEFAULT_MAPPING_OPERATION] * len(video.shape)
+    minimum_spatial_overlap_px = 64
+    minimum_temporal_overlap_frames = 16
+
+    if tiling_config is not None and tiling_config.spatial_config is not None:
+        cfg = tiling_config.spatial_config
+
+        tile_size_px = cfg.tile_size_in_pixels
+        overlap_px = cfg.tile_overlap_in_pixels
+
+        # Set minimum spatial overlap to 64 pixels in order to allow cutting padding from
+        # the front and back of the tiles and concatenate tiles without artifacts.
+        # The encoder uses symmetric padding (pad=1) in H and W at each conv layer. At tile
+        # boundaries, convs see padding (zeros/reflect) instead of real neighbor pixels, causing
+        # incorrect context near edges.
+        # For each overlap we discard 1 latent per edge (32px at scale 32) and concatenate tiles at a
+        # shared region with the next tile.
+        if overlap_px < minimum_spatial_overlap_px:
+            logger.warning(
+                f"Overlap pixels {overlap_px} in spatial tiling is less than \
+            {minimum_spatial_overlap_px}, setting to minimum required {minimum_spatial_overlap_px}"
+            )
+            overlap_px = minimum_spatial_overlap_px
+
+        # Define split and map operations for the spatial dimensions
+
+        # Height axis (H)
+        splitters[3] = split_with_symmetric_overlaps(tile_size_px, overlap_px)
+        mappers[3] = make_mapping_operation(map_spatial_interval_to_latent, scale=VIDEO_SCALE_FACTORS.height)
+
+        # Width axis (W)
+        splitters[4] = split_with_symmetric_overlaps(tile_size_px, overlap_px)
+        mappers[4] = make_mapping_operation(map_spatial_interval_to_latent, scale=VIDEO_SCALE_FACTORS.width)
+
+    if tiling_config is not None and tiling_config.temporal_config is not None:
+        cfg = tiling_config.temporal_config
+        tile_size_frames = cfg.tile_size_in_frames
+        overlap_frames = cfg.tile_overlap_in_frames
+
+        if overlap_frames < minimum_temporal_overlap_frames:
+            logger.warning(f"Overlap frames {overlap_frames} is less than 16, setting to minimum required 16")
+            overlap_frames = minimum_temporal_overlap_frames
+
+        splitters[2] = split_temporal_frames(tile_size_frames, overlap_frames)
+        mappers[2] = make_mapping_operation(map_temporal_interval_to_latent, scale=VIDEO_SCALE_FACTORS.time)
+
+    return create_tiles(video.shape, splitters, mappers)
+
+
+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":
+        out_channels = in_channels // block_config.get("multiplier", 1)
+        block = DepthToSpaceUpsample(
+            dims=convolution_dimensions,
+            in_channels=in_channels,
+            stride=(2, 1, 1),
+            out_channels_reduction_factor=block_config.get("multiplier", 1),
+            spatial_padding_mode=spatial_padding_mode,
+        )
+    elif block_name == "compress_space":
+        out_channels = in_channels // block_config.get("multiplier", 1)
+        block = DepthToSpaceUpsample(
+            dims=convolution_dimensions,
+            in_channels=in_channels,
+            stride=(1, 2, 2),
+            out_channels_reduction_factor=block_config.get("multiplier", 1),
+            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,
+        base_channels: int = 128,
+    ):
+        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
+
+        # LTX VAE decoder architecture uses 3 upsampler blocks with multiplier equals to 2.
+        # Hence the total feature_channels is multiplied by 8 (2^3).
+        feature_channels = base_channels * 8
+
+        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,
+        timestep: torch.Tensor | None = None,
+        generator: torch.Generator | None = 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_with_symmetric_overlaps(tile_size, overlap)
+                mappers[axis_idx] = make_mapping_operation(map_spatial_interval_to_pixel, scale=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_temporal_latents(tile_size, overlap)
+            mappers[2] = make_mapping_operation(map_temporal_interval_to_frame, scale=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: torch.Tensor | None = None,
+        generator: torch.Generator | None = 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: torch.Tensor | None,
+        generator: torch.Generator | None,
+    ) -> 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,
+    generator: torch.Generator | 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.
+        generator: Optional random generator for deterministic decoding.
+    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, generator=generator):
+            yield convert_to_uint8(frames)
+    else:
+        decoded_video = video_decoder(latent, generator=generator)
+        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_with_symmetric_overlaps(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_temporal_latents(size: int, overlap: int) -> SplitOperation:
+    """Split a temporal axis into overlapping tiles with causal handling.
+    Example with size=24, overlap=8 (units are whatever axis you split):
+        Non-causal split would produce:
+            Tile 0: [0, 24), left_ramp=0,  right_ramp=8
+            Tile 1: [16, 40), left_ramp=8, right_ramp=8
+            Tile 2: [32, 56), left_ramp=8, right_ramp=0
+        Causal split produces:
+            Tile 0: [0, 24), left_ramp=0,  right_ramp=8  (unchanged - starts at anchor)
+            Tile 1: [15, 40), left_ramp=9, right_ramp=8  (shifted back 1, ramp +1)
+            Tile 2: [31, 56), left_ramp=9, right_ramp=0  (shifted back 1, ramp +1)
+    This ensures each tile can causally depend on frames from previous tiles while maintaining
+    proper temporal continuity through the blend ramps.
+    Args:
+        size: Tile size in *axis units* (latent steps for LTX time tiling)
+        overlap: Overlap between tiles in the same units
+    Returns:
+        Split operation that divides temporal dimension with causal handling
+    """
+    non_causal_split = split_with_symmetric_overlaps(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:]]
+
+        # Extend blend ramps by 1 for non-first tiles to blend over the extra frame
+        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 split_temporal_frames(tile_size_frames: int, overlap_frames: int) -> SplitOperation:
+    """Split a temporal axis in video frame space into overlapping tiles.
+    Args:
+        tile_size_frames: Tile length in frames.
+        overlap_frames: Overlap between consecutive tiles in frames.
+    Returns:
+        Split operation that takes frame count and returns DimensionIntervals in frame indices.
+    """
+    non_causal_split = split_with_symmetric_overlaps(tile_size_frames, overlap_frames)
+
+    def split(dimension_size: int) -> DimensionIntervals:
+        if dimension_size <= tile_size_frames:
+            return DEFAULT_SPLIT_OPERATION(dimension_size)
+        intervals = non_causal_split(dimension_size)
+        ends = intervals.ends
+        ends[:-1] = [e + 1 for e in ends[:-1]]
+        right_ramps = [0] * len(intervals.right_ramps)
+        return replace(intervals, ends=ends, right_ramps=right_ramps)
+
+    return split
+
+
+def make_mapping_operation(
+    map_func: Callable[[int, int, int, int, int], Tuple[slice, torch.Tensor | None]],
+    scale: int,
+) -> MappingOperation:
+    """Create a mapping operation over a set of tiling intervals.
+    The given mapping function is applied to each interval in the input dimension. The result function is used for
+    creating tiles in the output dimension.
+    Args:
+        map_func: Mapping function to create the mapping operation from
+        scale: Scale factor for the transformation, used as an argument for the mapping function
+    Returns:
+        Mapping operation that takes a set of tiling intervals and returns a set of slices and masks in the output
+        dimension.
+    """
+
+    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_interval_to_frame(
+    begin: int,
+    end: int,
+    left_ramp: int,
+    right_ramp: int,
+    scale: int,
+) -> Tuple[slice, torch.Tensor]:
+    """Map temporal interval in latent space to video frame space.
+    Args:
+        begin: Start position in latent space
+        end: End position in latent space
+        left_ramp: Left ramp size in latent space
+        right_ramp: Right ramp size in latent space
+        scale: Scale factor for transformation
+    Returns:
+        Tuple of (output_slice, blend_mask)
+    """
+    start = begin * scale
+    stop = 1 + (end - 1) * scale
+
+    left_ramp_frames = 0 if left_ramp == 0 else 1 + (left_ramp - 1) * scale
+    right_ramp_frames = right_ramp * scale
+
+    mask_1d = compute_trapezoidal_mask_1d(stop - start, left_ramp_frames, right_ramp_frames, True)
+    return slice(start, stop), mask_1d
+
+
+def map_temporal_interval_to_latent(
+    begin: int, end: int, left_ramp: int, right_ramp: int | None = None, scale: int = 1
+) -> Tuple[slice, torch.Tensor]:
+    """
+    Map temporal interval in video frame space to latent space.
+    Args:
+        begin: Start position in video frame space
+        end: End position in video frame space
+        left_ramp: Left ramp size in video frame space
+        right_ramp: Right ramp size in video frame space
+        scale: Scale factor for transformation
+    Returns:
+        Tuple of (output_slice, blend_mask)
+    """
+    start = begin // scale
+    stop = (end - 1) // scale + 1
+
+    left_ramp_latents = 0 if left_ramp == 0 else 1 + (left_ramp - 1) // scale
+    right_ramp_latents = right_ramp // scale
+
+    if right_ramp_latents != 0:
+        raise ValueError("For tiled encoding, temporal tiles are expected to have a right ramp equal to 0")
+
+    mask_1d = compute_rectangular_mask_1d(stop - start, left_ramp_latents, right_ramp_latents)
+
+    return slice(start, stop), mask_1d
+
+
+def map_spatial_interval_to_pixel(
+    begin: int,
+    end: int,
+    left_ramp: int,
+    right_ramp: int,
+    scale: int,
+) -> Tuple[slice, torch.Tensor]:
+    """Map spatial interval in latent space to pixel space.
+    Args:
+        begin: Start position in latent space
+        end: End position in latent space
+        left_ramp: Left ramp size in latent space
+        right_ramp: Right ramp size in latent space
+        scale: Scale factor for transformation
+    """
+    start = begin * scale
+    stop = end * scale
+    mask_1d = compute_trapezoidal_mask_1d(stop - start, left_ramp * scale, right_ramp * scale, False)
+    return slice(start, stop), mask_1d
+
+
+def map_spatial_interval_to_latent(
+    begin: int,
+    end: int,
+    left_ramp: int,
+    right_ramp: int,
+    scale: int,
+) -> Tuple[slice, torch.Tensor]:
+    """Map spatial interval in pixel space to latent space.
+       Args:
+        begin: Start position in pixel space
+        end: End position in pixel space
+        left_ramp: Left ramp size in pixel space
+        right_ramp: Right ramp size in pixel space
+        scale: Scale factor for transformation
+    Returns:
+        Tuple of (output_slice, blend_mask)
+    """
+    start = begin // scale
+    stop = end // scale
+    left_ramp = max(0, left_ramp // scale - 1)
+
+    right_ramp = 0 if right_ramp == 0 else 1
+
+    mask_1d = compute_rectangular_mask_1d(stop - start, left_ramp, right_ramp)
+    return slice(start, stop), mask_1d