Upload folder using huggingface_hub

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

@@ -1,3 +1,123 @@
----
-license: apache-2.0
----
+# 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=data:image/png;base64,iVBORw0KGgoAAAANSUhEUgAAABQAAAAUCAYAAACNiR0NAAAAAXNSR0IArs4c6QAAAERlWElmTU0AKgAAAAgAAYdpAAQAAAABAAAAGgAAAAAAA6ABAAMAAAABAAEAAKACAAQAAAABAAAAFKADAAQAAAABAAAAFAAAAACy3fD9AAACmElEQVQ4Ea1VP2haYRA/fRo0mESRIIqb2IwxuNUl0CGFQBC6OAWcikMottCpqYtDQIgdQsBFhAjZqiQhbhmySJBOgmNU0EGCg9r61Bivd0ffoykE0iQH37/77n7f3e/uqQFIPB7P/N3d3QeDwfAFEedZ91ghnyH5JM1m87dWq6UavF6vdTKZfDcajW/p4rE49+wIFMj33Gq1vlNo+kxg758KpiETqP/29vaXweVyqaS0aBfPXEfGFwTjWCwM+KBQoWA4HAJx/KDNvxcmTTGbzYAH8SljOp2C2+2GjY0NqNfrcHFxAXNzc2LDfCuKIq78KBdFOwsgGzidTnA4HHBzcwO9Xg8sFgtsbm7C3t4eVCoVaDQa0O12YXl5GUwmk5z5cZ/PB6PRCNrttgADFQUXFhbw8PAQVVXF3d1dJAeMx+P0zn0Jh8OYz+eRADCRSGAqlcLxeIz7+/u4tLSEjKUDZrNZ8U4mk0jR4fr6Op6enoru+voa0+k0rq2tYTAYxE6ng9QiSLRgrVZDv9+PFLkA6kUhT+GEC8C8XF5ewtHRkejICShiaDabwPvj42NJm3k7ODiQdDl9Fr0ocqJpdXUVIpEIdz7Y7XZRr6ysQDQahXK5LORvbW1p5rC9vQ2UifAooBqHuVxO0vt72tnZwWq1qqtisRgWCgU5ZzIZPDk50fdUUEmZvxTmAgKBgAxunT/fJpRKJWmhUCgEVDi4uroSG46kWCzC4uKitNVgMICzszOhSgA5fiJZhp4Lbbh1KARpbF65D/lx3vMdP05Vlkf5zKIDyukFJi7N6AVwNAhVsdlsM+LsjaZ56sq8kyQUqs4P6rsAKV49B4x4Padf7Y9Kv9+fEmiBQH8S4Gsa5v8EHpL9VwL7xH8BvwEcd4ccVf02KQAAAABJRU5ErkJggg==)](https://console.ltx.video/playground)
+[![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-1.1.safetensors`](https://huggingface.co/Lightricks/LTX-2.3/blob/main/ltx-2.3-22b-distilled-1.1.safetensors) - [Download](https://huggingface.co/Lightricks/LTX-2.3/resolve/main/ltx-2.3-22b-distilled-1.1.safetensors)
+
+**Spatial Upscaler** - Required for current two-stage pipeline implementations in this repository
+  * [`ltx-2.3-spatial-upscaler-x2-1.1.safetensors`](https://huggingface.co/Lightricks/LTX-2.3/blob/main/ltx-2.3-spatial-upscaler-x2-1.1.safetensors) - [Download](https://huggingface.co/Lightricks/LTX-2.3/resolve/main/ltx-2.3-spatial-upscaler-x2-1.1.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, ICLoraPipeline, and LipDubPipeline)
+  * [`ltx-2.3-22b-distilled-lora-384-1.1.safetensors`](https://huggingface.co/Lightricks/LTX-2.3/blob/main/ltx-2.3-22b-distilled-lora-384-1.1.safetensors) - [Download](https://huggingface.co/Lightricks/LTX-2.3/resolve/main/ltx-2.3-22b-distilled-lora-384-1.1.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-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)
+  * [`LTX-2.3-22b-IC-LoRA-HDR`](https://huggingface.co/Lightricks/LTX-2.3-22b-IC-LoRA-HDR) - HDR IC-LoRA and pre-computed text embeddings for `HDRICLoraPipeline`
+  * [`LTX-2.3-22b-IC-LoRA-LipDub`](https://huggingface.co/Lightricks/LTX-2.3-22b-IC-LoRA-LipDub) - [Download](https://huggingface.co/Lightricks/LTX-2.3-22b-IC-LoRA-LipDub/resolve/main/ltx-2.3-22b-ic-lora-lipdub-0.9.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
+* **[HDRICLoraPipeline](packages/ltx-pipelines/src/ltx_pipelines/hdr_ic_lora.py)** - Video-to-video with HDR output (linear float frames via LogC3 inverse decode, suitable for EXR export and tonemapping)
+* **[LipDubPipeline](packages/ltx-pipelines/src/ltx_pipelines/lipdub.py)** - Lip dubbing, rephrasing, matching speaker identity (distilled model, single IC-LoRA, Two stages).
+
+### ⚡ 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). Fp8-cast should be used with bf16 checkpoints, it shall downcast them on the fly. For Hopper GPUs with TensorRT-LLM, use `--quantization fp8-scaled-mm` for FP8 scaled matrix multiplication. Fp8-scaled-mm should be used with fp8 checkpoints.
+* **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,413 @@
+# 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
+from ltx_core.loader import SDOps
+
+lora_sd_ops = SDOps(name="identity").with_matching()  # or a model-specific key-renaming SDOps
+
+builder = (
+    SingleGPUModelBuilder(
+        model_class_configurator=MyModelConfigurator,
+        model_path="/path/to/model.safetensors",
+    )
+    .lora("/path/to/lora_a.safetensors", 0.8, lora_sd_ops)
+    .lora("/path/to/lora_b.safetensors", 0.5, lora_sd_ops)
+)
+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", 1.0, lora_sd_ops)
+
+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/pyproject.toml

@@ -0,0 +1,55 @@
+[project]
+name = "ltx-core"
+version = "1.1.3"
+description = "Core implementation of Lightricks' LTX-2 model"
+readme = "README.md"
+requires-python = ">=3.10"
+dependencies = [
+    "torch~=2.7",
+    "torchaudio",
+    "einops",
+    "numpy",
+    "transformers>=4.52",
+    "safetensors",
+    "accelerate",
+    "scipy>=1.14",
+]
+
+[project.optional-dependencies]
+xformers = ["xformers"]
+fp8-trtllm = [
+    "tensorrt-llm==1.0.0",
+    "onnx>=1.16.0,<1.20.0",
+    "openmpi",
+]
+
+[tool.uv]
+conflicts = [
+  [
+    { extra = "xformers" },
+    { extra = "fp8-trtllm" },
+  ],
+]
+
+[tool.uv.sources]
+xformers = { index = "pytorch" }
+tensorrt-llm = { index = "nvidia" }
+
+[[tool.uv.index]]
+name = "pytorch"
+url = "https://download.pytorch.org/whl/cu129"
+explicit = true
+
+[[tool.uv.index]]
+name = "nvidia"
+url = "https://pypi.nvidia.com/"
+explicit = true
+
+[build-system]
+requires = ["uv_build>=0.9.8,<0.10.0"]
+build-backend = "uv_build"
+
+[dependency-groups]
+dev = [
+    "scikit-image>=0.25.2",
+]

packages/ltx-core/src/ltx_core/batch_split.py

@@ -0,0 +1,95 @@
+"""Batch-splitting adapter for the transformer.
+Wraps an ``X0Model`` (or ``BlockStreamingWrapper``) and splits batched inputs
+into smaller chunks before forwarding, then concatenates the results. This
+controls peak activation memory at the cost of more forward passes.
+The adapter is transparent — it has the same ``forward`` signature as
+``X0Model`` and proxies attribute access to the wrapped model.
+Example
+-------
+>>> from ltx_core.batch_split import BatchSplitAdapter
+>>> adapter = BatchSplitAdapter(model, max_batch_size=1)
+>>> # Receives B=4, runs 4xB=1 internally, returns B=4
+>>> denoised_video, denoised_audio = adapter(video=v_b4, audio=a_b4, perturbations=ptb)
+"""
+
+from __future__ import annotations
+
+from typing import Any
+
+import torch
+from torch import nn
+
+from ltx_core.guidance.perturbations import BatchedPerturbationConfig
+from ltx_core.model.transformer.modality import Modality
+
+
+def _split_perturbations(config: BatchedPerturbationConfig, sizes: list[int]) -> list[BatchedPerturbationConfig]:
+    """Split a ``BatchedPerturbationConfig`` along the batch dimension."""
+    it = iter(config.perturbations)
+    return [BatchedPerturbationConfig([next(it) for _ in range(s)]) for s in sizes]
+
+
+def _merge_tensors(tensors: list[torch.Tensor | None]) -> torch.Tensor | None:
+    """Concatenate tensors along batch dim, or return None if all are None."""
+    non_none = [t for t in tensors if t is not None]
+    if not non_none:
+        return None
+    return torch.cat(non_none, dim=0)
+
+
+class BatchSplitAdapter(nn.Module):
+    """Wraps a model and splits batched forward calls into smaller chunks.
+    Has the same ``forward`` signature as ``X0Model``:
+    ``(video, audio, perturbations) -> (denoised_video, denoised_audio)``.
+    Args:
+        model: The model to wrap (``X0Model``, ``BlockStreamingWrapper``, etc.).
+        max_batch_size: Maximum batch size per forward pass. Input batches
+            larger than this are split into sequential chunks.
+    """
+
+    def __init__(self, model: nn.Module, max_batch_size: int) -> None:
+        if max_batch_size < 1:
+            raise ValueError(f"max_batch_size must be >= 1, got {max_batch_size}")
+        super().__init__()
+        self._model = model
+        self._max_batch_size = max_batch_size
+
+    def _get_chunk_sizes(self, batch_size: int) -> list[int]:
+        full, remainder = divmod(batch_size, self._max_batch_size)
+        sizes = [self._max_batch_size] * full
+        if remainder:
+            sizes.append(remainder)
+        return sizes
+
+    def forward(
+        self,
+        video: Modality | None,
+        audio: Modality | None,
+        perturbations: BatchedPerturbationConfig,
+    ) -> tuple[torch.Tensor | None, torch.Tensor | None]:
+        batch_size = (video or audio).latent.shape[0]
+
+        if batch_size <= self._max_batch_size:
+            return self._model(video=video, audio=audio, perturbations=perturbations)
+
+        sizes = self._get_chunk_sizes(batch_size)
+        n = len(sizes)
+
+        v_chunks = video.split(sizes) if video is not None else [None] * n
+        a_chunks = audio.split(sizes) if audio is not None else [None] * n
+        p_chunks = _split_perturbations(perturbations, sizes)
+
+        chunk_results = [
+            self._model(video=vc, audio=ac, perturbations=pc)
+            for vc, ac, pc in zip(v_chunks, a_chunks, p_chunks, strict=True)
+        ]
+
+        results_v, results_a = zip(*chunk_results, strict=True)
+        return _merge_tensors(list(results_v)), _merge_tensors(list(results_a))
+
+    def __getattr__(self, name: str) -> Any:  # noqa: ANN401
+        """Proxy attribute access to the wrapped model."""
+        try:
+            return super().__getattr__(name)
+        except AttributeError:
+            return getattr(self._model, name)

packages/ltx-core/src/ltx_core/block_streaming/__init__.py

@@ -0,0 +1,19 @@
+"""Block streaming: memory-efficient sequential-block inference.
+Streams transformer blocks from safetensors to GPU one at a time.
+Block weights are provided by a :class:`WeightsProvider` which handles
+CPU-to-GPU copies, caching, and stream synchronization.  Two weight
+source strategies are available:
+- **RAM streaming** (default): all blocks pre-loaded into pinned CPU
+  buffers with LoRA fusion at build time.  Fast, higher CPU memory.
+- **Disk streaming** (``cpu_slots < num_blocks``): blocks read from
+  disk on demand with FIFO eviction.  Slower, lower CPU memory.
+"""
+
+from ltx_core.block_streaming.builder import DISK_CPU_SLOTS, StreamingModelBuilder
+from ltx_core.block_streaming.wrapper import BlockStreamingWrapper
+
+__all__ = [
+    "DISK_CPU_SLOTS",
+    "BlockStreamingWrapper",
+    "StreamingModelBuilder",
+]

packages/ltx-core/src/ltx_core/block_streaming/builder.py

@@ -0,0 +1,351 @@
+"""Builder that constructs a BlockStreamingWrapper from safetensors checkpoints."""
+
+from __future__ import annotations
+
+import logging
+from collections.abc import Callable
+from dataclasses import dataclass, field, replace
+from typing import Generic
+
+import safetensors
+import torch
+from torch import nn
+
+from ltx_core.block_streaming.disk import DiskBlockReader, DiskTensorReader, LoraSource
+from ltx_core.block_streaming.pool import WeightPool
+from ltx_core.block_streaming.provider import WeightsProvider
+from ltx_core.block_streaming.source import DiskWeightSource, PinnedWeightSource, WeightSource
+from ltx_core.block_streaming.utils import allocate_layout_views, derive_layout, make_block_key, resolve_attr
+from ltx_core.block_streaming.wrapper import BlockStreamingWrapper
+from ltx_core.loader.fuse_loras import aggregate_lora_products, fuse_lora_weights
+from ltx_core.loader.helpers import create_meta_model, load_state_dict, read_model_config
+from ltx_core.loader.module_ops import ModuleOps
+from ltx_core.loader.primitives import (
+    LoraPathStrengthAndSDOps,
+    LoraStateDictWithStrength,
+    ModelBuilderProtocol,
+    StateDictLoader,
+)
+from ltx_core.loader.registry import DummyRegistry, Registry
+from ltx_core.loader.sd_ops import SDOps
+from ltx_core.loader.sft_loader import SafetensorsModelStateDictLoader
+from ltx_core.model.model_protocol import ModelConfigurator, ModelType
+
+logger = logging.getLogger(__name__)
+
+DISK_CPU_SLOTS = 2
+_DEFAULT_GPU_SLOTS = 2
+
+
+@dataclass(frozen=True)
+class StreamingModelBuilder(Generic[ModelType], ModelBuilderProtocol[ModelType]):
+    """Immutable builder for :class:`BlockStreamingWrapper`.
+    Reads block weights from safetensors on demand.  ``cpu_slots`` and
+    ``gpu_slots`` control the memory/speed trade-off (see :meth:`build`).
+    Args:
+        model_class_configurator: Creates the model from a config dict.
+        model_path: One or more ``.safetensors`` checkpoint paths.
+        model_sd_ops: Key remapping applied to safetensors keys.
+        module_ops: Module-level mutations for the meta model.
+        loras: LoRA adapters fused into weights at load time.
+        model_loader: Strategy for reading checkpoint metadata.
+        registry: Shared cache for loaded state dicts.
+        blocks_attr: Dotted path to the ``nn.ModuleList`` (e.g.
+            ``"velocity_model.transformer_blocks"``).
+        blocks_prefix: State-dict key prefix for block weights
+            (e.g. ``"transformer_blocks"``).
+        state_dict_prefix: Wrapper offset prepended to keys when loading into
+            the meta model (e.g. ``"velocity_model."`` when wrapped by ``X0Model``).
+        model_wrapper: Optional callable wrapping the model
+            (e.g. ``X0Model``).
+    """
+
+    model_class_configurator: type[ModelConfigurator[ModelType]]
+    model_path: str | tuple[str, ...]
+    model_sd_ops: SDOps | None = None
+    module_ops: tuple[ModuleOps, ...] = field(default_factory=tuple)
+    loras: tuple[LoraPathStrengthAndSDOps, ...] = field(default_factory=tuple)
+    model_loader: StateDictLoader = field(default_factory=SafetensorsModelStateDictLoader)
+    registry: Registry = field(default_factory=DummyRegistry)
+
+    # Streaming-specific
+    blocks_attr: str = ""
+    blocks_prefix: str = ""
+    state_dict_prefix: str = ""
+    model_wrapper: Callable[[ModelType], nn.Module] | None = None
+
+    def with_sd_ops(self, sd_ops: SDOps | None) -> StreamingModelBuilder:
+        return replace(self, model_sd_ops=sd_ops)
+
+    def with_module_ops(self, module_ops: tuple[ModuleOps, ...]) -> StreamingModelBuilder:
+        return replace(self, module_ops=module_ops)
+
+    def with_loras(self, loras: tuple[LoraPathStrengthAndSDOps, ...]) -> StreamingModelBuilder:
+        return replace(self, loras=loras)
+
+    def model_config(self) -> dict:
+        """Read model configuration from the checkpoint metadata."""
+        return read_model_config(self.model_path, self.model_loader)
+
+    def meta_model(self, config: dict, module_ops: tuple[ModuleOps, ...]) -> ModelType:
+        """Create a model on the meta device and apply module operations."""
+        return create_meta_model(self.model_class_configurator, config, module_ops)
+
+    def build(
+        self,
+        target_device: torch.device,
+        dtype: torch.dtype,
+        cpu_slots_count: int | None = None,
+        gpu_slots_count: int | None = None,
+        **_kwargs: object,
+    ) -> BlockStreamingWrapper:
+        """Build and return a ready-to-use :class:`BlockStreamingWrapper`.
+        Args:
+            target_device: GPU device for compute.
+            dtype: Weight dtype (e.g. ``torch.bfloat16``).
+            cpu_slots_count: Number of pinned CPU buffer slots.
+                ``None`` = RAM streaming (all blocks pre-loaded with LoRA fusion).
+            gpu_slots_count: Number of GPU buffer slots.
+                ``None`` = ``_DEFAULT_GPU_SLOTS`` (2).
+        """
+        if not self.blocks_prefix:
+            raise ValueError("blocks_prefix must be non-empty for streaming")
+
+        config = read_model_config(self.model_path, self.model_loader)
+        meta_model: nn.Module = create_meta_model(self.model_class_configurator, config, self.module_ops)
+        if self.model_wrapper is not None:
+            meta_model = self.model_wrapper(meta_model)
+        meta_model.eval()
+
+        blocks = resolve_attr(meta_model, self.blocks_attr)
+
+        checkpoint_paths = list(self.model_path) if isinstance(self.model_path, tuple) else [self.model_path]
+        block_key_map, non_block_keys = _scan_checkpoint_keys(checkpoint_paths, self.model_sd_ops, self.blocks_prefix)
+
+        cpu_slots_count = cpu_slots_count if cpu_slots_count is not None else len(blocks)
+        gpu_slots_count = gpu_slots_count if gpu_slots_count is not None else _DEFAULT_GPU_SLOTS
+
+        if cpu_slots_count >= len(blocks):
+            source, lora_sources = self._build_pinned_source(
+                meta_model, target_device, dtype, cpu_slots_count, block_key_map, non_block_keys
+            )
+        else:
+            reader = DiskTensorReader(checkpoint_paths)
+            source, lora_sources = self._build_disk_source(
+                meta_model, target_device, dtype, cpu_slots_count, reader, block_key_map, non_block_keys
+            )
+
+        copy_stream = torch.cuda.Stream(device=target_device)
+        gpu_pool = WeightPool(
+            source.block_layout,
+            gpu_slots_count,
+            target_device,
+            reuse_barrier=lambda event: copy_stream.wait_event(event),
+        )
+        provider = WeightsProvider(gpu_pool, copy_stream, target_device, source, lora_sources, self.blocks_prefix)
+        return BlockStreamingWrapper(
+            model=meta_model,
+            blocks=blocks,
+            provider=provider,
+            target_device=target_device,
+        )
+
+    def _build_pinned_source(
+        self,
+        meta_model: nn.Module,
+        target_device: torch.device,
+        dtype: torch.dtype,
+        cpu_slots_count: int,
+        block_key_map: dict[int, list[tuple[str, str]]],
+        non_block_keys: list[tuple[str, str]],
+    ) -> tuple[WeightSource, list[LoraSource]]:
+        """Pre-load all blocks into pinned CPU buffers with LoRA fusion."""
+        model_sd = load_state_dict(
+            self.model_path, self.model_loader, self.registry, torch.device("cpu"), self.model_sd_ops
+        )
+
+        lora_sd_and_strengths = [
+            LoraStateDictWithStrength(
+                load_state_dict([lora.path], self.model_loader, self.registry, torch.device("cpu"), lora.sd_ops),
+                lora.strength,
+            )
+            for lora in self.loras
+        ]
+
+        for block_idx in block_key_map:
+            if block_idx >= cpu_slots_count:
+                raise ValueError(
+                    f"Pinned source requires one CPU slot per block; "
+                    f"got block index {block_idx} with only {cpu_slots_count} slots."
+                )
+
+        blocks = resolve_attr(meta_model, self.blocks_attr)
+        block_tensors: dict[str, torch.Tensor] = {}
+        for block_idx, entries in block_key_map.items():
+            block_params = dict(blocks[block_idx].named_parameters())
+            for _sft_key, param_name in entries:
+                key = make_block_key(self.blocks_prefix, block_idx, param_name)
+                block_tensors[key] = block_params[param_name]
+        blocks_layout = derive_layout(block_tensors, dtype)
+        pinned_blocks = allocate_layout_views(blocks_layout, pin_memory=True)
+
+        should_sync = False
+        for key, fused in fuse_lora_weights(model_sd, lora_sd_and_strengths, dtype=None, preserve_input_device=False):
+            if key in pinned_blocks:
+                pinned_blocks[key].copy_(fused, non_blocking=True)
+                model_sd.sd[key] = None
+                should_sync = True
+            else:
+                model_sd.sd[key] = fused
+        if should_sync:
+            torch.cuda.synchronize()
+
+        # Fill remaining pinned keys from the source state dict.
+        for key in blocks_layout:
+            if model_sd.sd[key] is None:
+                continue
+            pinned_blocks[key].copy_(model_sd.sd[key])
+            model_sd.sd[key] = None
+
+        pinned: dict[int, dict[str, torch.Tensor]] = {
+            block_idx: {
+                param_name: pinned_blocks[make_block_key(self.blocks_prefix, block_idx, param_name)]
+                for _sft_key, param_name in entries
+            }
+            for block_idx, entries in block_key_map.items()
+        }
+
+        non_block_sd: dict[str, torch.Tensor] = {
+            self.state_dict_prefix + model_key: model_sd.sd[model_key].to(device=target_device, dtype=dtype)
+            for _sft_key, model_key in non_block_keys
+        }
+
+        meta_model.load_state_dict(non_block_sd, strict=False, assign=True)
+
+        return PinnedWeightSource(pinned), []
+
+    def _build_disk_source(
+        self,
+        meta_model: nn.Module,
+        target_device: torch.device,
+        dtype: torch.dtype,
+        cpu_slots_count: int,
+        reader: DiskTensorReader,
+        block_key_map: dict[int, list[tuple[str, str]]],
+        non_block_keys: list[tuple[str, str]],
+    ) -> tuple[WeightSource, list[LoraSource]]:
+        """Create a DiskWeightSource backed by a DiskBlockReader for lazy loading.
+        Derives the shared pool layout from the meta model's block 0 — this
+        relies on module_ops (e.g. fp8_cast) leaving the meta param dtype in
+        sync with the post-sd_ops checkpoint dtype.
+        """
+        lora_sources = [LoraSource(lora.path, lora.sd_ops, lora.strength) for lora in self.loras]
+
+        self._load_non_block_weights(
+            reader,
+            non_block_keys,
+            meta_model,
+            target_device,
+            dtype,
+            sd_ops=self.model_sd_ops,
+            key_prefix=self.state_dict_prefix,
+            lora_sources=lora_sources,
+        )
+
+        blocks = resolve_attr(meta_model, self.blocks_attr)
+        layout = derive_layout(dict(blocks[0].named_parameters()), dtype)
+
+        cpu_pool = WeightPool(
+            layout,
+            cpu_slots_count,
+            torch.device("cpu"),
+            reuse_barrier=lambda event: event.synchronize(),
+            pin_memory=True,
+        )
+        block_reader = DiskBlockReader(
+            reader=reader,
+            block_key_map=block_key_map,
+            sd_ops=self.model_sd_ops,
+            blocks_prefix=self.blocks_prefix,
+        )
+        source = DiskWeightSource(cpu_pool, block_reader)
+        return source, lora_sources
+
+    # ------------------------------------------------------------------
+    # Helpers
+    # ------------------------------------------------------------------
+
+    @staticmethod
+    def _fuse_lora_delta(
+        model_key: str,
+        tensor: torch.Tensor,
+        lora_sources: list[LoraSource],
+    ) -> torch.Tensor:
+        """Add all matching LoRA deltas to *tensor* in-place via ``addmm_``."""
+        if not lora_sources or not model_key.endswith(".weight"):
+            return tensor
+        prefix = model_key[: -len(".weight")]
+        products = (
+            ab
+            for ab in (s.get_ab(prefix, device=tensor.device, dtype=tensor.dtype) for s in lora_sources)
+            if ab is not None
+        )
+        aggregate_lora_products(products, out=tensor)
+        return tensor
+
+    @staticmethod
+    @torch.inference_mode()
+    def _load_non_block_weights(
+        reader: DiskTensorReader,
+        non_block_keys: list[tuple[str, str]],
+        model: nn.Module,
+        device: torch.device,
+        dtype: torch.dtype,
+        sd_ops: SDOps | None = None,
+        key_prefix: str = "",
+        lora_sources: list[LoraSource] | None = None,
+    ) -> None:
+        """Load non-block weights into *model* on *device*."""
+        state_dict: dict[str, torch.Tensor] = {}
+        sources = lora_sources or []
+        for sft_key, model_key in non_block_keys:
+            tensor = reader.get_tensor(sft_key).to(device=device, dtype=dtype)
+            tensor = StreamingModelBuilder._fuse_lora_delta(model_key, tensor, sources)
+            if sd_ops is not None:
+                for kv in sd_ops.apply_to_key_value(model_key, tensor):
+                    state_dict[key_prefix + kv.new_key] = kv.new_value
+                continue
+            state_dict[key_prefix + model_key] = tensor
+        model.load_state_dict(state_dict, strict=False, assign=True)
+
+
+def _scan_checkpoint_keys(
+    checkpoint_paths: list[str],
+    sd_ops: SDOps | None,
+    blocks_prefix: str,
+) -> tuple[dict[int, list[tuple[str, str]]], list[tuple[str, str]]]:
+    """Partition checkpoint keys into per-block and non-block lists.
+    Opens the safetensors files for header-only key enumeration; no tensor data
+    is read.
+    """
+    block_key_map: dict[int, list[tuple[str, str]]] = {}
+    non_block_keys: list[tuple[str, str]] = []
+    prefix_dot = blocks_prefix + "."
+    for path in checkpoint_paths:
+        with safetensors.safe_open(path, framework="pt", device="cpu") as handle:
+            for sft_key in handle.keys():  # noqa: SIM118
+                model_key = sd_ops.apply_to_key(sft_key) if sd_ops else sft_key
+                if model_key is None:
+                    continue
+                if model_key.startswith(prefix_dot):
+                    rest = model_key[len(prefix_dot) :]
+                    idx_str, _, param_name = rest.partition(".")
+                    try:
+                        block_idx = int(idx_str)
+                    except ValueError:
+                        non_block_keys.append((sft_key, model_key))
+                        continue
+                    block_key_map.setdefault(block_idx, []).append((sft_key, param_name))
+                else:
+                    non_block_keys.append((sft_key, model_key))
+    return block_key_map, non_block_keys

packages/ltx-core/src/ltx_core/block_streaming/disk.py

@@ -0,0 +1,148 @@
+"""Safetensors I/O and LoRA fusion for block streaming."""
+
+from __future__ import annotations
+
+from collections.abc import Iterator
+
+import safetensors
+import torch
+
+from ltx_core.block_streaming.utils import allocate_layout_views, make_block_key
+from ltx_core.loader.fuse_loras import LoraProduct
+from ltx_core.loader.sd_ops import SDOps
+
+_SAFETENSORS_DTYPE_TO_TORCH: dict[str, torch.dtype] = {
+    "F64": torch.float64,
+    "F32": torch.float32,
+    "F16": torch.float16,
+    "BF16": torch.bfloat16,
+}
+
+
+class DiskTensorReader:
+    """Key-based tensor accessor over one or more safetensors files."""
+
+    def __init__(self, paths: list[str]) -> None:
+        self._handles: list[safetensors.safe_open] = []
+        self._key_to_handle_idx: dict[str, int] = {}
+        for path in paths:
+            handle = safetensors.safe_open(path, framework="pt", device="cpu")
+            handle_idx = len(self._handles)
+            self._handles.append(handle)
+            for sft_key in handle.keys():  # noqa: SIM118
+                self._key_to_handle_idx[sft_key] = handle_idx
+
+    def get_tensor(self, key: str) -> torch.Tensor:
+        return self._handles[self._key_to_handle_idx[key]].get_tensor(key)
+
+    def close(self) -> None:
+        self._handles.clear()
+        self._key_to_handle_idx.clear()
+
+    def __contains__(self, key: str) -> bool:
+        return key in self._key_to_handle_idx
+
+    def __iter__(self) -> Iterator[str]:
+        return iter(self._key_to_handle_idx)
+
+
+class DiskBlockReader:
+    """Reads one block at a time from safetensors into provided buffers."""
+
+    def __init__(
+        self,
+        reader: DiskTensorReader,
+        block_key_map: dict[int, list[tuple[str, str]]],
+        sd_ops: SDOps | None = None,
+        blocks_prefix: str = "",
+    ) -> None:
+        self._reader = reader
+        self._block_key_map = block_key_map
+        self._sd_ops = sd_ops
+        self._blocks_prefix = blocks_prefix
+
+    def read_into(self, target: dict[str, torch.Tensor], block_idx: int) -> None:
+        block_prefix = make_block_key(self._blocks_prefix, block_idx, "")
+        for sft_key, param_name in self._block_key_map[block_idx]:
+            tensor = self._reader.get_tensor(sft_key)
+            if self._sd_ops is None:
+                target[param_name].copy_(tensor)
+                continue
+            full_key = make_block_key(self._blocks_prefix, block_idx, param_name)
+            for result in self._sd_ops.apply_to_key_value(full_key, tensor):
+                if not result.new_key.startswith(block_prefix):
+                    raise ValueError(
+                        f"SDOps output key '{result.new_key}' is outside block {block_idx} "
+                        f"(expected prefix '{block_prefix}'); cannot route to a per-block buffer."
+                    )
+                target[result.new_key[len(block_prefix) :]].copy_(result.new_value)
+
+    def cleanup(self) -> None:
+        self._reader.close()
+
+
+class LoraSource:
+    """Pinned-memory cache of matched LoRA A/B factors backed by a single buffer."""
+
+    def __init__(self, path: str, sd_ops: SDOps | None, strength: float) -> None:
+        self.strength = strength
+        self._pinned_ab: dict[str, tuple[torch.Tensor, torch.Tensor]] = {}
+
+        a_keys: dict[str, str] = {}
+        b_keys: dict[str, str] = {}
+        with safetensors.safe_open(path, framework="pt", device="cpu") as handle:
+            for sft_key in handle.keys():  # noqa: SIM118
+                model_key = sd_ops.apply_to_key(sft_key) if sd_ops is not None else sft_key
+                if model_key is None:
+                    continue
+                if model_key.endswith(".lora_A.weight"):
+                    a_keys[model_key[: -len(".lora_A.weight")]] = sft_key
+                elif model_key.endswith(".lora_B.weight"):
+                    b_keys[model_key[: -len(".lora_B.weight")]] = sft_key
+
+            matched_prefixes = list(a_keys.keys() & b_keys.keys())
+
+            # Build the layout from safetensors header metadata only — no tensor data is read.
+            layout: dict[str, tuple[torch.Size, torch.dtype]] = {}
+            for prefix in matched_prefixes:
+                a_slice_view = handle.get_slice(a_keys[prefix])
+                b_slice_view = handle.get_slice(b_keys[prefix])
+                layout[f"{prefix}.A"] = (
+                    torch.Size(a_slice_view.get_shape()),
+                    _SAFETENSORS_DTYPE_TO_TORCH[a_slice_view.get_dtype()],
+                )
+                layout[f"{prefix}.B"] = (
+                    torch.Size(b_slice_view.get_shape()),
+                    _SAFETENSORS_DTYPE_TO_TORCH[b_slice_view.get_dtype()],
+                )
+
+            all_views = allocate_layout_views(layout, pin_memory=True)
+
+            for prefix in matched_prefixes:
+                a_view = all_views[f"{prefix}.A"]
+                b_view = all_views[f"{prefix}.B"]
+                a_view.copy_(handle.get_tensor(a_keys[prefix]))
+                b_view.copy_(handle.get_tensor(b_keys[prefix]))
+                self._pinned_ab[prefix] = (a_view, b_view)
+
+    def get_ab(
+        self,
+        param_prefix: str,
+        device: torch.device | None = None,
+        dtype: torch.dtype | None = None,
+    ) -> LoraProduct | None:
+        """Return the :class:`LoraProduct` for *param_prefix*, or ``None``."""
+        pair = self._pinned_ab.get(param_prefix)
+        if pair is None:
+            return None
+        a, b = pair
+        if device is not None and device.type == "cuda":
+            a = a.to(device=device, non_blocking=True)
+            b = b.to(device=device, non_blocking=True)
+        if dtype is not None:
+            a = a.to(dtype=dtype)
+            b = b.to(dtype=dtype)
+        return LoraProduct(a, b, self.strength)
+
+    def cleanup(self) -> None:
+        self._pinned_ab.clear()

packages/ltx-core/src/ltx_core/block_streaming/pool.py

@@ -0,0 +1,75 @@
+"""Weight buffer pool for block streaming."""
+
+from __future__ import annotations
+
+from collections import deque
+from typing import Callable
+
+import torch
+
+from ltx_core.block_streaming.utils import allocate_layout_views
+from ltx_core.loader.primitives import TensorLayout
+
+
+class WeightPool:
+    """Fixed pool of pre-allocated weight buffers with event-based reuse.
+    All slots share a single buffer (CPU or GPU); each slot is a
+    contiguous slice carved out of it via :func:`allocate_layout_views`.
+    Args:
+        buffer_layout: ``{name: (shape, dtype)}`` for each buffer.
+        capacity: Number of buffers to pre-allocate.
+        device: Device for allocation.
+        reuse_barrier: Called with the pending event before a buffer is reused.
+        pin_memory: Pin buffers (for async H2D copies from CPU).
+    """
+
+    def __init__(
+        self,
+        buffer_layout: TensorLayout,
+        capacity: int,
+        device: torch.device,
+        reuse_barrier: Callable[[torch.cuda.Event], None],
+        pin_memory: bool = False,
+    ) -> None:
+        self._buffer_layout = buffer_layout
+        self._capacity = capacity
+        self._free: deque[dict[str, torch.Tensor]] = deque()
+        self._events: dict[int, torch.cuda.Event] = {}
+        self._reuse_barrier = reuse_barrier
+        memory_layout = {
+            _make_key(slot, name): (shape, dtype)
+            for slot in range(capacity)
+            for name, (shape, dtype) in buffer_layout.items()
+        }
+        all_views = allocate_layout_views(memory_layout, device=device, pin_memory=pin_memory)
+        for slot in range(capacity):
+            self._free.append({name: all_views[_make_key(slot, name)] for name in buffer_layout})
+
+    @property
+    def capacity(self) -> int:
+        return self._capacity
+
+    @property
+    def buffer_layout(self) -> TensorLayout:
+        return self._buffer_layout
+
+    def acquire(self) -> dict[str, torch.Tensor]:
+        """Take a free buffer, waiting any pending event before returning."""
+        weights = self._free.popleft()
+        event = self._events.pop(id(weights), None)
+        if event is not None:
+            self._reuse_barrier(event)
+        return weights
+
+    def release(self, weights: dict[str, torch.Tensor], event: torch.cuda.Event | None = None) -> None:
+        """Return a buffer to the free list.
+        If *event* is given it is waited on the next :meth:`acquire`
+        of this buffer, ensuring the prior operation has completed.
+        """
+        if event is not None:
+            self._events[id(weights)] = event
+        self._free.append(weights)
+
+
+def _make_key(slot: int, name: str) -> str:
+    return f"{slot}/{name}"

packages/ltx-core/src/ltx_core/block_streaming/provider.py

@@ -0,0 +1,147 @@
+"""GPU weights provider for block streaming."""
+
+from __future__ import annotations
+
+from collections import OrderedDict
+
+import torch
+
+from ltx_core.block_streaming.disk import LoraSource
+from ltx_core.block_streaming.pool import WeightPool
+from ltx_core.block_streaming.source import WeightSource
+from ltx_core.block_streaming.utils import FP8_DTYPES
+from ltx_core.loader.fuse_loras import aggregate_lora_products, fuse_cast_fp8_weight
+
+
+def _contiguous_byte_view(weights: dict[str, torch.Tensor]) -> torch.Tensor | None:
+    """Return a ``uint8`` view spanning every tensor in *weights*, or ``None`` if
+    they don't share one contiguous storage region."""
+    tensors = list(weights.values())
+    if not tensors:
+        return None
+    storage = tensors[0].untyped_storage()
+    storage_ptr = storage.data_ptr()
+    start = end = tensors[0].storage_offset() * tensors[0].element_size()
+    for t in tensors:
+        if t.untyped_storage().data_ptr() != storage_ptr or not t.is_contiguous():
+            return None
+        offset = t.storage_offset() * t.element_size()
+        nbytes = t.numel() * t.element_size()
+        start = min(start, offset)
+        end = max(end, offset + nbytes)
+    view = torch.empty(0, dtype=torch.uint8, device=tensors[0].device)
+    view.set_(storage, start, (end - start,), (1,))
+    return view
+
+
+class WeightsProvider:
+    """Provides GPU-ready block weights via H2D copy from a pinned CPU weight source.
+    Args:
+        pool: Pre-allocated GPU weight buffer pool.
+        copy_stream: Dedicated CUDA stream for async H2D copies.
+        target_device: GPU device for compute.
+        source: Pinned CPU weight source.
+        lora_sources: LoRA adapters fused on H2D copy.
+        blocks_prefix: State-dict prefix for LoRA key matching.
+    """
+
+    def __init__(
+        self,
+        pool: WeightPool,
+        copy_stream: torch.cuda.Stream,
+        target_device: torch.device,
+        source: WeightSource,
+        lora_sources: list[LoraSource] | None = None,
+        blocks_prefix: str = "",
+    ) -> None:
+        self._copy_stream = copy_stream
+        self._pool = pool
+        self._cache: OrderedDict[int, dict[str, torch.Tensor]] = OrderedDict()
+        self._events: dict[int, torch.cuda.Event] = {}
+        self._target_device = target_device
+        self._source = source
+        self._lora_sources = lora_sources or []
+        self._blocks_prefix = blocks_prefix
+
+    def get(self, idx: int) -> dict[str, torch.Tensor]:
+        """Return GPU weights for block *idx*. Does H2D copy on miss."""
+        if idx in self._cache:
+            return self._cache[idx]
+
+        # Evict oldest GPU buffer if at capacity.
+        if len(self._cache) >= self._pool.capacity:
+            evicted_idx, evicted_weights = self._cache.popitem(last=False)
+            self._pool.release(evicted_weights, event=self._events.pop(evicted_idx, None))
+
+        gpu_weights = self._pool.acquire()
+        cpu_weights = self._source.get(idx)
+
+        h2d_event = self._copy_to_gpu(idx, gpu_weights, cpu_weights)
+        self._source.release(idx, event=h2d_event)
+
+        self._cache[idx] = gpu_weights
+        return gpu_weights
+
+    def _copy_to_gpu(
+        self,
+        idx: int,
+        gpu_weights: dict[str, torch.Tensor],
+        cpu_weights: dict[str, torch.Tensor],
+    ) -> torch.cuda.Event:
+        """Enqueue H2D copy + LoRA fusion on the copy stream and wait on compute.
+        The wait is intentionally inside this method so callers -- and
+        instrumentation regions wrapping it -- observe the full transfer time.
+        """
+        with torch.cuda.stream(self._copy_stream):
+            gpu_view = _contiguous_byte_view(gpu_weights)
+            cpu_view = _contiguous_byte_view(cpu_weights)
+            if gpu_view is not None and cpu_view is not None and gpu_view.numel() == cpu_view.numel():
+                gpu_view.copy_(cpu_view, non_blocking=True)
+            else:
+                for name, gpu_tensor in gpu_weights.items():
+                    gpu_tensor.copy_(cpu_weights[name], non_blocking=True)
+            if self._lora_sources:
+                self._fuse_block_loras(idx, gpu_weights)
+            h2d_event = torch.cuda.Event()
+            h2d_event.record(self._copy_stream)
+
+        torch.cuda.current_stream(self._target_device).wait_event(h2d_event)
+        return h2d_event
+
+    def release(self, idx: int, event: torch.cuda.Event) -> None:
+        """Attach a compute-done event -- waited before this buffer is recycled."""
+        self._events[idx] = event
+
+    def cleanup(self) -> None:
+        """Synchronize streams and release all resources."""
+        self._copy_stream.synchronize()
+        torch.cuda.current_stream(self._target_device).synchronize()
+        self._cache.clear()
+        self._events.clear()
+        self._source.cleanup()
+        for lora in self._lora_sources:
+            lora.cleanup()
+
+    def __len__(self) -> int:
+        return len(self._cache)
+
+    def _fuse_block_loras(self, idx: int, weights: dict[str, torch.Tensor]) -> None:
+        """Fuse LoRA deltas directly into GPU block weights."""
+        for name, tensor in weights.items():
+            if not name.endswith(".weight"):
+                continue
+            prefix = f"{self._blocks_prefix}.{idx}.{name}".removesuffix(".weight")
+            is_fp8 = tensor.dtype in FP8_DTYPES
+            agg_dtype = torch.bfloat16 if is_fp8 else tensor.dtype
+            products = (
+                ab
+                for ab in (s.get_ab(prefix, device=self._target_device, dtype=agg_dtype) for s in self._lora_sources)
+                if ab is not None
+            )
+            aggregated = aggregate_lora_products(products, agg_dtype)
+            if aggregated is None:
+                continue
+            if is_fp8:
+                tensor.copy_(fuse_cast_fp8_weight(aggregated, tensor, tensor.dtype))
+            else:
+                tensor.add_(aggregated)

packages/ltx-core/src/ltx_core/block_streaming/source.py

@@ -0,0 +1,102 @@
+"""Weight sources for block streaming: protocol and implementations."""
+
+from __future__ import annotations
+
+from collections import OrderedDict
+from typing import Protocol
+
+import torch
+
+from ltx_core.block_streaming.disk import DiskBlockReader
+from ltx_core.block_streaming.pool import WeightPool
+from ltx_core.loader.primitives import TensorLayout
+
+
+class WeightSource(Protocol):
+    """Provides pinned CPU weights for a given block index.
+    Assumes all buffers share an identical layout across all block indices.
+    """
+
+    @property
+    def block_layout(self) -> TensorLayout:
+        """Shared per-block buffer layout (shape + dtype for each param)."""
+        ...
+
+    def get(self, idx: int) -> dict[str, torch.Tensor]:
+        """Return CPU weights for block *idx*."""
+        ...
+
+    def release(self, idx: int, event: torch.cuda.Event) -> None:
+        """Signal that an async operation using these weights is guarded by *event*."""
+        ...
+
+    def cleanup(self) -> None:
+        """Release all resources (buffers, readers, events)."""
+        ...
+
+
+class DiskWeightSource(WeightSource):
+    """Reads block weights from disk into pinned CPU buffers on demand."""
+
+    def __init__(self, pool: WeightPool, reader: DiskBlockReader) -> None:
+        self._pool = pool
+        self._cache: OrderedDict[int, dict[str, torch.Tensor]] = OrderedDict()
+        self._events: dict[int, torch.cuda.Event] = {}
+        self._reader = reader
+
+    @property
+    def block_layout(self) -> TensorLayout:
+        return self._pool.buffer_layout
+
+    def get(self, idx: int) -> dict[str, torch.Tensor]:
+        """Return CPU weights for block *idx*. Reads from disk on miss."""
+        if idx in self._cache:
+            return self._cache[idx]
+
+        if len(self._cache) >= self._pool.capacity:
+            evicted_idx, evicted_weights = self._cache.popitem(last=False)
+            self._pool.release(evicted_weights, event=self._events.pop(evicted_idx, None))
+
+        weights = self._pool.acquire()
+        self._reader.read_into(weights, idx)
+        self._cache[idx] = weights
+        return weights
+
+    def release(self, idx: int, event: torch.cuda.Event) -> None:
+        """Attach an H2D event -- waited before this buffer is recycled."""
+        self._events[idx] = event
+
+    def cleanup(self) -> None:
+        """Clear cache and close the disk reader."""
+        self._cache.clear()
+        self._events.clear()
+        self._reader.cleanup()
+
+    def __len__(self) -> int:
+        return len(self._cache)
+
+
+class PinnedWeightSource(WeightSource):
+    """Pre-loaded pinned CPU weights."""
+
+    def __init__(self, weights: dict[int, dict[str, torch.Tensor]]) -> None:
+        if not weights:
+            raise ValueError("PinnedWeightSource requires at least one block")
+        self._weights = weights
+
+    @property
+    def block_layout(self) -> TensorLayout:
+        first_block = self._weights[min(self._weights)]
+        return {name: (t.shape, t.dtype) for name, t in first_block.items()}
+
+    def get(self, idx: int) -> dict[str, torch.Tensor]:
+        return self._weights[idx]
+
+    def release(self, idx: int, event: torch.cuda.Event) -> None:
+        pass
+
+    def cleanup(self) -> None:
+        self._weights.clear()
+
+    def __len__(self) -> int:
+        return len(self._weights)

packages/ltx-core/src/ltx_core/block_streaming/utils.py

@@ -0,0 +1,134 @@
+"""Shared utilities for the block_streaming package."""
+
+from __future__ import annotations
+
+import math
+import weakref
+from dataclasses import dataclass
+from typing import Any
+
+import torch
+from torch import nn
+
+from ltx_core.loader.primitives import TensorLayout
+
+FP8_DTYPES = frozenset({torch.float8_e4m3fn, torch.float8_e5m2})
+
+_BUFFER_ALIGN = 16
+
+
+def make_block_key(blocks_prefix: str, block_idx: int, param_name: str) -> str:
+    """Return the state-dict key for *param_name* under block *block_idx*."""
+    return f"{blocks_prefix}.{block_idx}.{param_name}"
+
+
+def resolve_attr(module: nn.Module, dotted_path: str) -> nn.ModuleList:
+    """Resolve a dotted attribute path like ``'model.language_model.layers'``."""
+    obj: Any = module
+    for part in dotted_path.split("."):
+        obj = getattr(obj, part)
+    if not isinstance(obj, nn.ModuleList):
+        raise TypeError(f"Expected nn.ModuleList at '{dotted_path}', got {type(obj).__name__}")
+    return obj
+
+
+def assign_tensor_to_module(root: nn.Module, dotted_name: str, tensor: torch.Tensor) -> None:
+    """Assign *tensor* to the parameter/buffer at *dotted_name* inside *root*.
+    Unlike ``param.data = tensor``, this works even when the existing parameter
+    lives on the ``meta`` device (which has an incompatible storage type).
+    """
+    parts = dotted_name.split(".")
+    parent = root
+    for part in parts[:-1]:
+        parent = getattr(parent, part)
+    leaf = parts[-1]
+    if leaf in parent._parameters:
+        parent._parameters[leaf] = nn.Parameter(tensor, requires_grad=False)
+    elif leaf in parent._buffers:
+        parent._buffers[leaf] = tensor
+    else:
+        raise AttributeError(f"{leaf} is not a parameter or buffer of {type(parent).__name__}")
+
+
+def derive_layout(tensors: dict[str, torch.Tensor], dtype: torch.dtype | None = None) -> TensorLayout:
+    """Derive a layout from a ``{name: tensor}`` dict.
+    If ``dtype`` is given, non-FP8 dtypes are coerced to it (FP8 preserved). If
+    ``None``, the source dtype is preserved as-is.
+    """
+    return {
+        name: (t.shape, t.dtype if dtype is None or t.dtype in FP8_DTYPES else dtype) for name, t in tensors.items()
+    }
+
+
+def _align_up(offset: int, alignment: int) -> int:
+    return (offset + alignment - 1) & ~(alignment - 1)
+
+
+def _alloc_pinned_exact(nbytes: int) -> torch.Tensor | None:
+    """Allocate exactly ``nbytes`` of pinned host memory via ``cudaHostRegister``.
+    Bypasses PyTorch's ``CachingHostAllocator``, which rounds every
+    ``pin_memory=True`` request up to ``PowerOf2Ceil(N)`` (see
+    ``aten/src/ATen/core/CachingHostAllocator.h``). Returns ``None`` if
+    registration fails. The unregister hook is bound to the storage (not the
+    tensor) so views of the buffer keep the registration alive until the
+    memory is actually freed. Caller is responsible for ensuring CUDA is
+    available.
+    """
+    cudart = torch.cuda.cudart()
+    buf = torch.empty(nbytes, dtype=torch.uint8)
+    ptr = buf.data_ptr()
+    err = int(cudart.cudaHostRegister(ptr, nbytes, 0))
+    if err != 0:
+        return None
+    weakref.finalize(buf.untyped_storage(), lambda p=ptr: cudart.cudaHostUnregister(p))
+    return buf
+
+
+def _alloc_buffer(nbytes: int, device: torch.device | None, pin_memory: bool) -> torch.Tensor:
+    """Allocate one ``uint8`` buffer for :func:`allocate_layout_views`.
+    For pinned host buffers, prefer ``cudaHostRegister`` to dodge the caching
+    allocator's power-of-2 rounding. Falls back to the caching allocator if
+    registration fails. Raises if pinning is requested without a CUDA runtime,
+    since pinning is fundamentally a CUDA driver operation.
+    """
+    if pin_memory and (device is None or torch.device(device).type == "cpu"):
+        if not torch.cuda.is_available():
+            raise RuntimeError("pin_memory=True requires CUDA, which is not available")
+        buf = _alloc_pinned_exact(nbytes)
+        if buf is not None:
+            return buf
+    return torch.empty(nbytes, dtype=torch.uint8, device=device, pin_memory=pin_memory)
+
+
+@dataclass(frozen=True)
+class _TensorSlice:
+    """Location of a single tensor view within the buffer."""
+
+    offset: int
+    shape: torch.Size
+    dtype: torch.dtype
+
+    def size(self) -> int:
+        return math.prod(self.shape) * self.dtype.itemsize
+
+
+def allocate_layout_views(
+    layout: TensorLayout,
+    device: torch.device | None = None,
+    pin_memory: bool = False,
+) -> dict[str, torch.Tensor]:
+    """Allocate a single ``uint8`` buffer and return per-key tensor views into it.
+    All keys in *layout* live in one contiguous allocation; each returned
+    tensor is a non-overlapping slice of that buffer reinterpreted at the
+    requested shape and dtype. The views keep the underlying storage alive
+    via PyTorch refcounting — drop them all to release the memory.
+    """
+    slices: dict[str, _TensorSlice] = {}
+    cursor = 0
+    for key, (shape, dtype) in layout.items():
+        cursor = _align_up(cursor, _BUFFER_ALIGN)
+        slices[key] = _TensorSlice(offset=cursor, shape=shape, dtype=dtype)
+        cursor += slices[key].size()
+    # Allocate at least one byte so empty layouts still produce a valid buffer.
+    buffer = _alloc_buffer(max(_align_up(cursor, _BUFFER_ALIGN), 1), device, pin_memory)
+    return {key: buffer[s.offset : s.offset + s.size()].view(s.dtype).view(s.shape) for key, s in slices.items()}

packages/ltx-core/src/ltx_core/block_streaming/wrapper.py

@@ -0,0 +1,96 @@
+"""Block streaming wrapper: streams transformer blocks through a WeightsProvider."""
+
+from __future__ import annotations
+
+import itertools
+from typing import Any
+
+import torch
+from torch import nn
+
+from ltx_core.block_streaming.provider import WeightsProvider
+from ltx_core.block_streaming.utils import assign_tensor_to_module
+
+
+class BlockStreamingWrapper(nn.Module):
+    """Streams sequential model blocks through GPU buffer caches.
+    The wrapper delegates all weight management to a :class:`WeightsProvider`
+    which handles CPU-to-GPU copies, caching, LoRA fusion, and stream
+    synchronization internally.
+    Use :class:`StreamingModelBuilder` to construct this wrapper -- it
+    handles checkpoint parsing, source selection, and provider creation.
+    Args:
+        model: The wrapped model (non-block params already on GPU).
+        blocks: Sequential blocks to stream (``nn.ModuleList``).
+        provider: Provides GPU-ready weights on demand.
+        target_device: GPU device for compute.
+    """
+
+    def __init__(
+        self,
+        model: nn.Module,
+        blocks: nn.ModuleList,
+        provider: WeightsProvider,
+        target_device: torch.device,
+    ) -> None:
+        super().__init__()
+        self._model = model
+        self._blocks = blocks
+        self._target_device = target_device
+        self._provider = provider
+
+        self._hooks: list[torch.utils.hooks.RemovableHandle] = []
+        self._register_hooks()
+
+    # ------------------------------------------------------------------
+    # Hook registration
+    # ------------------------------------------------------------------
+
+    def _pre_hook(self, block_idx: int) -> None:
+        """Load GPU weights for a block and inject them into its parameters."""
+        gpu_weights = self._provider.get(block_idx)
+
+        block = self._blocks[block_idx]
+        for name, _param in itertools.chain(block.named_parameters(), block.named_buffers()):
+            assign_tensor_to_module(block, name, gpu_weights[name])
+
+    def _post_hook(self, block_idx: int) -> None:
+        """Record a compute-done event and release the block weights."""
+        compute_done = torch.cuda.Event()
+        compute_done.record(torch.cuda.current_stream(self._target_device))
+        self._provider.release(block_idx, event=compute_done)
+
+    def _register_hooks(self) -> None:
+        for idx, block in enumerate(self._blocks):
+            pre = block.register_forward_pre_hook(
+                lambda _mod, _args, *, idx=idx: self._pre_hook(idx),
+            )
+            post = block.register_forward_hook(
+                lambda _mod, _args, _out, *, idx=idx: self._post_hook(idx),
+            )
+            self._hooks.extend([pre, post])
+
+    # ------------------------------------------------------------------
+    # Teardown
+    # ------------------------------------------------------------------
+
+    def teardown(self) -> None:
+        """Remove hooks and release all resources."""
+        for h in self._hooks:
+            h.remove()
+        self._hooks.clear()
+        self._provider.cleanup()
+
+    # ------------------------------------------------------------------
+    # Forward and attribute delegation
+    # ------------------------------------------------------------------
+
+    def forward(self, *args: Any, **kwargs: Any) -> Any:  # noqa: ANN401
+        return self._model(*args, **kwargs)
+
+    def __getattr__(self, name: str) -> Any:  # noqa: ANN401
+        """Proxy attribute access to the wrapped model."""
+        try:
+            return super().__getattr__(name)
+        except AttributeError:
+            return getattr(self._model, name)

packages/ltx-core/src/ltx_core/components/__init__.py

@@ -0,0 +1,10 @@
+"""
+Diffusion pipeline components.
+Submodules:
+    diffusion_steps - Diffusion stepping algorithms (EulerDiffusionStep)
+    guiders         - Guidance strategies (CFGGuider, STGGuider, APG variants)
+    noisers         - Noise samplers (GaussianNoiser)
+    patchifiers     - Latent patchification (VideoLatentPatchifier, AudioPatchifier)
+    protocols       - Protocol definitions (Patchifier, etc.)
+    schedulers      - Sigma schedulers (LTX2Scheduler, LinearQuadraticScheduler)
+"""

packages/ltx-core/src/ltx_core/components/diffusion_steps.py

@@ -0,0 +1,186 @@
+import torch
+
+from ltx_core.components.protocols import DiffusionStepProtocol
+from ltx_core.utils import to_velocity
+
+
+def _get_ancestral_step(
+    sigma_from: torch.Tensor,
+    sigma_to: torch.Tensor,
+    eta: float = 1.0,
+) -> tuple[torch.Tensor, torch.Tensor]:
+    """Compute ``(sigma_down, sigma_up)`` for one DDIM ancestral sampling step.
+    Both inputs are in the rescaled parameterization ``sigma / alpha``.
+    Returns ``sigma_down`` (deterministic component) and ``sigma_up``
+    (stochastic component) in the same rescaled space.
+    """
+    if not eta:
+        return sigma_to, torch.zeros_like(sigma_to)
+    variance = sigma_to**2 * (sigma_from**2 - sigma_to**2).clamp(min=0) / sigma_from**2
+    sigma_up = (eta * variance**0.5).clamp(max=sigma_to)
+    sigma_down = (sigma_to**2 - sigma_up**2).clamp(min=0) ** 0.5
+    return sigma_down, sigma_up
+
+
+class EulerDiffusionStep(DiffusionStepProtocol):
+    """
+    First-order Euler method for diffusion sampling.
+    Takes a single step from the current noise level (sigma) to the next by
+    computing velocity from the denoised prediction and applying: sample + velocity * dt.
+    """
+
+    def step(
+        self, sample: torch.Tensor, denoised_sample: torch.Tensor, sigmas: torch.Tensor, step_index: int, **_kwargs
+    ) -> torch.Tensor:
+        sigma = sigmas[step_index]
+        sigma_next = sigmas[step_index + 1]
+        dt = sigma_next - sigma
+        velocity = to_velocity(sample, sigma, denoised_sample)
+
+        return (sample.to(torch.float32) + velocity.to(torch.float32) * dt).to(sample.dtype)
+
+
+class Res2sDiffusionStep(DiffusionStepProtocol):
+    """
+    Second-order diffusion step for res_2s sampling with SDE noise injection.
+    Used by the res_2s denoising loop. Advances the sample from the current
+    sigma to the next by mixing a deterministic update (from the denoised
+    prediction) with injected noise via ``get_sde_coeff``, producing
+    variance-preserving transitions.
+    """
+
+    @staticmethod
+    def get_sde_coeff(
+        sigma_next: torch.Tensor,
+        sigma_up: torch.Tensor | None = None,
+        sigma_down: torch.Tensor | None = None,
+        sigma_max: torch.Tensor | None = None,
+    ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+        """
+        Compute SDE coefficients (alpha_ratio, sigma_down, sigma_up) for the step.
+        Given either ``sigma_down`` or ``sigma_up``, returns the mixing
+        coefficients used for variance-preserving noise injection. If
+        ``sigma_up`` is provided, ``sigma_down`` and ``alpha_ratio`` are
+        derived; if ``sigma_down`` is provided, ``sigma_up`` and
+        ``alpha_ratio`` are derived.
+        """
+        if sigma_down is not None:
+            alpha_ratio = (1 - sigma_next) / (1 - sigma_down)
+            sigma_up = (sigma_next**2 - sigma_down**2 * alpha_ratio**2).clamp(min=0) ** 0.5
+        elif sigma_up is not None:
+            # Fallback to avoid sqrt(neg_num)
+            sigma_up.clamp_(max=sigma_next * 0.9999)
+            sigmax = sigma_max if sigma_max is not None else torch.ones_like(sigma_next)
+            sigma_signal = sigmax - sigma_next
+            sigma_residual = (sigma_next**2 - sigma_up**2).clamp(min=0) ** 0.5
+            alpha_ratio = sigma_signal + sigma_residual
+            sigma_down = sigma_residual / alpha_ratio
+        else:
+            alpha_ratio = torch.ones_like(sigma_next)
+            sigma_down = sigma_next
+            sigma_up = torch.zeros_like(sigma_next)
+
+        sigma_up = torch.nan_to_num(sigma_up if sigma_up is not None else torch.zeros_like(sigma_next), 0.0)
+        # Replace NaNs in sigma_down with corresponding sigma_next elements (float32)
+        nan_mask = torch.isnan(sigma_down)
+        sigma_down[nan_mask] = sigma_next[nan_mask].to(sigma_down.dtype)
+        alpha_ratio = torch.nan_to_num(alpha_ratio, 1.0)
+
+        return alpha_ratio, sigma_down, sigma_up
+
+    def step(
+        self,
+        sample: torch.Tensor,
+        denoised_sample: torch.Tensor,
+        sigmas: torch.Tensor,
+        step_index: int,
+        noise: torch.Tensor,
+        eta: float = 0.5,
+    ) -> torch.Tensor:
+        """Advance one step with SDE noise injection via get_sde_coeff.
+        Args:
+            sample: Current noisy sample.
+            denoised_sample: Denoised prediction from the model.
+            sigmas: Noise schedule tensor.
+            step_index: Current step index in the schedule.
+            noise: Random noise tensor for stochastic injection.
+            eta: Controls stochastic noise injection strength (0=deterministic, 1=maximum). Default 0.5.
+        Returns:
+            Next sample with SDE noise injection applied.
+        """
+        sigma = sigmas[step_index]
+        sigma_next = sigmas[step_index + 1]
+        alpha_ratio, sigma_down, sigma_up = self.get_sde_coeff(sigma_next, sigma_up=sigma_next * eta)
+        output_dtype = denoised_sample.dtype
+        if torch.any(sigma_up == 0) or torch.any(sigma_next == 0):
+            return denoised_sample
+
+        # Extract epsilon prediction
+        eps_next = (sample - denoised_sample) / (sigma - sigma_next)
+        denoised_next = sample - sigma * eps_next
+
+        # Mix deterministic and stochastic components
+        x_noised = alpha_ratio * (denoised_next + sigma_down * eps_next) + sigma_up * noise
+        return x_noised.to(output_dtype)
+
+
+class EulerCfgPpDiffusionStep(DiffusionStepProtocol):
+    """Euler step using the CFG++ correction for the ODE derivative.
+    Instead of the standard velocity formula, the ODE derivative is computed
+    from the unconditioned prediction, keeping the conditioned prediction as
+    the target denoised state.  Ancestral (DDIM) noise injection is applied
+    in the rescaled sigma parameterization (sigma / alpha).
+    All diffusion quantities (alpha, ODE derivative, ancestral coefficients)
+    are computed internally from ``sigmas`` and ``uncond_denoised``.
+    Reference: CFG++ (https://arxiv.org/abs/2406.08070).
+    """
+
+    def __init__(self, eta: float = 1.0, s_noise: float = 1.0) -> None:
+        self.eta = eta
+        self.s_noise = s_noise
+
+    def step(
+        self,
+        sample: torch.Tensor,
+        denoised_sample: torch.Tensor,
+        sigmas: torch.Tensor,
+        step_index: int,
+        uncond_denoised: torch.Tensor,
+        noise: torch.Tensor | None = None,
+        **_kwargs,
+    ) -> torch.Tensor:
+        """Advance one CFG++ Euler step.
+        Args:
+            sample: Current noisy latent x_t.
+            denoised_sample: Conditioned denoised prediction x_0^cond.
+            sigmas: Full sigma schedule tensor.
+            step_index: Current step index.
+            uncond_denoised: Unconditioned denoised prediction x_0^uncond,
+                used to compute the ODE derivative direction.
+            noise: Noise tensor for stochastic injection; ignored when
+                ``eta=0`` or ``s_noise=0``.
+        Returns:
+            Updated latent x_{t-1}.
+        """
+        sigma_s = sigmas[step_index].to(torch.float32)
+        sigma_t = sigmas[step_index + 1].to(torch.float32)
+        _eps = torch.finfo(torch.float32).eps
+        # Clamp to avoid division by zero when sigma == 1.0 exactly.
+        alpha_s = (1.0 - sigma_s).clamp(min=_eps)
+        alpha_t = (1.0 - sigma_t).clamp(min=_eps)
+
+        x = sample.to(torch.float32)
+        denoised = denoised_sample.to(torch.float32)
+        uncond = uncond_denoised.to(torch.float32)
+
+        # ODE derivative: direction toward noise using uncond prediction (CFG++ correction)
+        d = (x - alpha_s * uncond) / sigma_s
+
+        # Ancestral step in rescaled sigma space (sigma / alpha)
+        sigma_down, sigma_up = _get_ancestral_step(sigma_s / alpha_s, sigma_t / alpha_t, eta=self.eta)
+        sigma_down = alpha_t * sigma_down
+
+        x_next = alpha_t * denoised + sigma_down * d
+        if noise is not None and self.eta > 0 and self.s_noise > 0:
+            x_next = x_next + alpha_t * noise.to(torch.float32) * self.s_noise * sigma_up
+        return x_next.to(sample.dtype)

packages/ltx-core/src/ltx_core/components/guiders.py

@@ -0,0 +1,364 @@
+import math
+from collections.abc import Mapping, Sequence
+from dataclasses import dataclass, field
+
+import torch
+
+from ltx_core.components.protocols import GuiderProtocol
+
+
+@dataclass(frozen=True)
+class CFGGuider(GuiderProtocol):
+    """
+    Classifier-free guidance (CFG) guider.
+    Computes the guidance delta as (scale - 1) * (cond - uncond), steering the
+    denoising process toward the conditioned prediction.
+    Attributes:
+        scale: Guidance strength. 1.0 means no guidance, higher values increase
+            adherence to the conditioning.
+    """
+
+    scale: float
+
+    def delta(self, cond: torch.Tensor, uncond: torch.Tensor) -> torch.Tensor:
+        return (self.scale - 1) * (cond - uncond)
+
+    def enabled(self) -> bool:
+        return self.scale != 1.0
+
+
+@dataclass(frozen=True)
+class CFGStarRescalingGuider(GuiderProtocol):
+    """
+    Calculates the CFG delta between conditioned and unconditioned samples.
+    To minimize offset in the denoising direction and move mostly along the
+    conditioning axis within the distribution, the unconditioned sample is
+    rescaled in accordance with the norm of the conditioned sample.
+    Attributes:
+        scale (float):
+            Global guidance strength. A value of 1.0 corresponds to no extra
+            guidance beyond the base model prediction. Values > 1.0 increase
+            the influence of the conditioned sample relative to the
+            unconditioned one.
+    """
+
+    scale: float
+
+    def delta(self, cond: torch.Tensor, uncond: torch.Tensor) -> torch.Tensor:
+        rescaled_neg = projection_coef(cond, uncond) * uncond
+        return (self.scale - 1) * (cond - rescaled_neg)
+
+    def enabled(self) -> bool:
+        return self.scale != 1.0
+
+
+@dataclass(frozen=True)
+class STGGuider(GuiderProtocol):
+    """
+    Calculates the STG delta between conditioned and perturbed denoised samples.
+    Perturbed samples are the result of the denoising process with perturbations,
+    e.g. attentions acting as passthrough for certain layers and modalities.
+    Attributes:
+        scale (float):
+            Global strength of the STG guidance. A value of 0.0 disables the
+            guidance. Larger values increase the correction applied in the
+            direction of (pos_denoised - perturbed_denoised).
+    """
+
+    scale: float
+
+    def delta(self, pos_denoised: torch.Tensor, perturbed_denoised: torch.Tensor) -> torch.Tensor:
+        return self.scale * (pos_denoised - perturbed_denoised)
+
+    def enabled(self) -> bool:
+        return self.scale != 0.0
+
+
+@dataclass(frozen=True)
+class LtxAPGGuider(GuiderProtocol):
+    """
+    Calculates the APG (adaptive projected guidance) delta between conditioned
+    and unconditioned samples.
+    To minimize offset in the denoising direction and move mostly along the
+    conditioning axis within the distribution, the (cond - uncond) delta is
+    decomposed into components parallel and orthogonal to the conditioned
+    sample. The `eta` parameter weights the parallel component, while `scale`
+    is applied to the orthogonal component. Optionally, a norm threshold can
+    be used to suppress guidance when the magnitude of the correction is small.
+    Attributes:
+        scale (float):
+            Strength applied to the component of the guidance that is orthogonal
+            to the conditioned sample. Controls how aggressively we move in
+            directions that change semantics but stay consistent with the
+            conditioning manifold.
+        eta (float):
+            Weight of the component of the guidance that is parallel to the
+            conditioned sample. A value of 1.0 keeps the full parallel
+            component; values in [0, 1] attenuate it, and values > 1.0 amplify
+            motion along the conditioning direction.
+        norm_threshold (float):
+            Minimum L2 norm of the guidance delta below which the guidance
+            can be reduced or ignored (depending on implementation).
+            This is useful for avoiding noisy or unstable updates when the
+            guidance signal is very small.
+    """
+
+    scale: float
+    eta: float = 1.0
+    norm_threshold: float = 0.0
+
+    def delta(self, cond: torch.Tensor, uncond: torch.Tensor) -> torch.Tensor:
+        guidance = cond - uncond
+        if self.norm_threshold > 0:
+            ones = torch.ones_like(guidance)
+            guidance_norm = guidance.norm(p=2, dim=[-1, -2, -3], keepdim=True)
+            scale_factor = torch.minimum(ones, self.norm_threshold / guidance_norm)
+            guidance = guidance * scale_factor
+        proj_coeff = projection_coef(guidance, cond)
+        g_parallel = proj_coeff * cond
+        g_orth = guidance - g_parallel
+        g_apg = g_parallel * self.eta + g_orth
+
+        return g_apg * (self.scale - 1)
+
+    def enabled(self) -> bool:
+        return self.scale != 1.0
+
+
+@dataclass(frozen=False)
+class LegacyStatefulAPGGuider(GuiderProtocol):
+    """
+    Calculates the APG (adaptive projected guidance) delta between conditioned
+    and unconditioned samples.
+    To minimize offset in the denoising direction and move mostly along the
+    conditioning axis within the distribution, the (cond - uncond) delta is
+    decomposed into components parallel and orthogonal to the conditioned
+    sample. The `eta` parameter weights the parallel component, while `scale`
+    is applied to the orthogonal component. Optionally, a norm threshold can
+    be used to suppress guidance when the magnitude of the correction is small.
+    Attributes:
+        scale (float):
+            Strength applied to the component of the guidance that is orthogonal
+            to the conditioned sample. Controls how aggressively we move in
+            directions that change semantics but stay consistent with the
+            conditioning manifold.
+        eta (float):
+            Weight of the component of the guidance that is parallel to the
+            conditioned sample. A value of 1.0 keeps the full parallel
+            component; values in [0, 1] attenuate it, and values > 1.0 amplify
+            motion along the conditioning direction.
+        norm_threshold (float):
+            Minimum L2 norm of the guidance delta below which the guidance
+            can be reduced or ignored (depending on implementation).
+            This is useful for avoiding noisy or unstable updates when the
+            guidance signal is very small.
+        momentum (float):
+            Exponential moving-average coefficient for accumulating guidance
+            over time. running_avg = momentum * running_avg + guidance
+    """
+
+    scale: float
+    eta: float
+    norm_threshold: float = 5.0
+    momentum: float = 0.0
+    # it is user's responsibility not to use same APGGuider for several denoisings or different modalities
+    # in order not to share accumulated average across different denoisings or modalities
+    running_avg: torch.Tensor | None = None
+
+    def delta(self, cond: torch.Tensor, uncond: torch.Tensor) -> torch.Tensor:
+        guidance = cond - uncond
+        if self.momentum != 0:
+            if self.running_avg is None:
+                self.running_avg = guidance.clone()
+            else:
+                self.running_avg = self.momentum * self.running_avg + guidance
+            guidance = self.running_avg
+
+        if self.norm_threshold > 0:
+            ones = torch.ones_like(guidance)
+            guidance_norm = guidance.norm(p=2, dim=[-1, -2, -3], keepdim=True)
+            scale_factor = torch.minimum(ones, self.norm_threshold / guidance_norm)
+            guidance = guidance * scale_factor
+
+        proj_coeff = projection_coef(guidance, cond)
+        g_parallel = proj_coeff * cond
+        g_orth = guidance - g_parallel
+        g_apg = g_parallel * self.eta + g_orth
+
+        return g_apg * self.scale
+
+    def enabled(self) -> bool:
+        return self.scale != 0.0
+
+
+@dataclass(frozen=True)
+class MultiModalGuiderParams:
+    """
+    Parameters for the multi-modal guider.
+    """
+
+    cfg_scale: float = 1.0
+    "CFG (Classifier-free guidance) scale controlling how strongly the model adheres to the prompt."
+    stg_scale: float = 0.0
+    "STG (Spatio-Temporal Guidance) scale controls how strongly the model reacts to the perturbation of the modality."
+    stg_blocks: list[int] | None = field(default_factory=list)
+    "Which transformer blocks to perturb for STG."
+    rescale_scale: float = 0.0
+    "Rescale scale controlling how strongly the model rescales the modality after applying other guidance."
+    modality_scale: float = 1.0
+    "Modality scale controlling how strongly the model reacts to the perturbation of the modality."
+    skip_step: int = 0
+    "Skip step controlling how often the model skips the step."
+
+
+def _params_for_sigma_from_sorted_dict(
+    sigma: float, params_by_sigma: Sequence[tuple[float, MultiModalGuiderParams]]
+) -> MultiModalGuiderParams:
+    """
+    Return params for the given sigma from a sorted (sigma_upper_bound -> params) structure.
+    Keys are sorted descending (bin upper bounds). Bin i is (key_{i+1}, key_i].
+    Get all keys >= sigma; use last in list (smallest such key = upper bound of bin containing sigma),
+    or last entry in the sequence if list is empty (sigma above max key).
+    """
+    if not params_by_sigma:
+        raise ValueError("params_by_sigma must be non-empty")
+    sigma = float(sigma)
+    keys_desc = [k for k, _ in params_by_sigma]
+    keys_ge_sigma = [k for k in keys_desc if k >= sigma]
+    # sigma above all keys: use first bin (max key)
+    key = keys_ge_sigma[-1] if keys_ge_sigma else keys_desc[0]
+    return next(p for k, p in params_by_sigma if k == key)
+
+
+@dataclass(frozen=True)
+class MultiModalGuider:
+    """
+    Multi-modal guider with constant params per instance.
+    For sigma-dependent params, use MultiModalGuiderFactory.build_from_sigma(sigma) to
+    obtain a guider for each step.
+    """
+
+    params: MultiModalGuiderParams
+    negative_context: torch.Tensor | None = None
+
+    def calculate(
+        self,
+        cond: torch.Tensor,
+        uncond_text: torch.Tensor | float,
+        uncond_perturbed: torch.Tensor | float,
+        uncond_modality: torch.Tensor | float,
+    ) -> torch.Tensor:
+        """
+        The guider calculates the guidance delta as (scale - 1) * (cond - uncond) for cfg and modality cfg,
+        and as scale * (cond - uncond) for stg, steering the denoising process away from the unconditioned
+        prediction.
+        """
+        pred = (
+            cond
+            + (self.params.cfg_scale - 1) * (cond - uncond_text)
+            + self.params.stg_scale * (cond - uncond_perturbed)
+            + (self.params.modality_scale - 1) * (cond - uncond_modality)
+        )
+
+        if self.params.rescale_scale != 0:
+            factor = cond.std() / pred.std()
+            factor = self.params.rescale_scale * factor + (1 - self.params.rescale_scale)
+            pred = pred * factor
+
+        return pred
+
+    def do_unconditional_generation(self) -> bool:
+        """Returns True if the guider is doing unconditional generation."""
+        return not math.isclose(self.params.cfg_scale, 1.0)
+
+    def do_perturbed_generation(self) -> bool:
+        """Returns True if the guider is doing perturbed generation."""
+        return not math.isclose(self.params.stg_scale, 0.0)
+
+    def do_isolated_modality_generation(self) -> bool:
+        """Returns True if the guider is doing isolated modality generation."""
+        return not math.isclose(self.params.modality_scale, 1.0)
+
+    def should_skip_step(self, step: int) -> bool:
+        """Returns True if the guider should skip the step."""
+        if self.params.skip_step == 0:
+            return False
+        return step % (self.params.skip_step + 1) != 0
+
+
+@dataclass(frozen=True)
+class MultiModalGuiderFactory:
+    """
+    Factory that creates a MultiModalGuider for a given sigma.
+    Single source of truth: _params_by_sigma (schedule). Use constant() for
+    one params for all sigma, from_dict() for sigma-binned params.
+    """
+
+    negative_context: torch.Tensor | None = None
+    _params_by_sigma: tuple[tuple[float, MultiModalGuiderParams], ...] = ()
+
+    @classmethod
+    def constant(
+        cls,
+        params: MultiModalGuiderParams,
+        negative_context: torch.Tensor | None = None,
+    ) -> "MultiModalGuiderFactory":
+        """Build a factory with constant params (same guider for all sigma)."""
+        return cls(
+            negative_context=negative_context,
+            _params_by_sigma=((float("inf"), params),),
+        )
+
+    @classmethod
+    def from_dict(
+        cls,
+        sigma_to_params: Mapping[float, MultiModalGuiderParams],
+        negative_context: torch.Tensor | None = None,
+    ) -> "MultiModalGuiderFactory":
+        """
+        Build a factory from a dict of sigma_value -> MultiModalGuiderParams.
+        Keys are sorted descending and used for bin lookup in params(sigma).
+        """
+        if not sigma_to_params:
+            raise ValueError("sigma_to_params must be non-empty")
+        sorted_items = tuple(sorted(sigma_to_params.items(), key=lambda x: x[0], reverse=True))
+        return cls(negative_context=negative_context, _params_by_sigma=sorted_items)
+
+    def params(self, sigma: float | torch.Tensor) -> MultiModalGuiderParams:
+        """Return params effective for the given sigma (getter; single source of truth)."""
+        sigma_val = float(sigma.item() if isinstance(sigma, torch.Tensor) else sigma)
+        return _params_for_sigma_from_sorted_dict(sigma_val, self._params_by_sigma)
+
+    def build_from_sigma(self, sigma: float | torch.Tensor) -> MultiModalGuider:
+        """Return a MultiModalGuider with params effective for the given sigma."""
+        return MultiModalGuider(
+            params=self.params(sigma),
+            negative_context=self.negative_context,
+        )
+
+
+def create_multimodal_guider_factory(
+    params: MultiModalGuiderParams | MultiModalGuiderFactory,
+    negative_context: torch.Tensor | None = None,
+) -> MultiModalGuiderFactory:
+    """
+    Create or return a MultiModalGuiderFactory. Pass constant params for a
+    single-params factory (uses MultiModalGuiderFactory.constant), or an existing
+    MultiModalGuiderFactory. When given a factory, returns it as-is unless
+    negative_context is provided. For sigma-dependent params use
+    MultiModalGuiderFactory.from_dict(...) and pass that as params.
+    """
+    if isinstance(params, MultiModalGuiderFactory):
+        if negative_context is not None and params.negative_context is not negative_context:
+            return MultiModalGuiderFactory.from_dict(dict(params._params_by_sigma), negative_context=negative_context)
+        return params
+    return MultiModalGuiderFactory.constant(params, negative_context=negative_context)
+
+
+def projection_coef(to_project: torch.Tensor, project_onto: torch.Tensor) -> torch.Tensor:
+    batch_size = to_project.shape[0]
+    positive_flat = to_project.reshape(batch_size, -1)
+    negative_flat = project_onto.reshape(batch_size, -1)
+    dot_product = torch.sum(positive_flat * negative_flat, dim=1, keepdim=True)
+    squared_norm = torch.sum(negative_flat**2, dim=1, keepdim=True) + 1e-8
+    return dot_product / squared_norm

packages/ltx-core/src/ltx_core/components/noisers.py

@@ -0,0 +1,35 @@
+from dataclasses import replace
+from typing import Protocol
+
+import torch
+
+from ltx_core.types import LatentState
+
+
+class Noiser(Protocol):
+    """Protocol for adding noise to a latent state during diffusion."""
+
+    def __call__(self, latent_state: LatentState, noise_scale: float) -> LatentState: ...
+
+
+class GaussianNoiser(Noiser):
+    """Adds Gaussian noise to a latent state, scaled by the denoise mask."""
+
+    def __init__(self, generator: torch.Generator):
+        super().__init__()
+
+        self.generator = generator
+
+    def __call__(self, latent_state: LatentState, noise_scale: float = 1.0) -> LatentState:
+        noise = torch.randn(
+            *latent_state.latent.shape,
+            device=latent_state.latent.device,
+            dtype=latent_state.latent.dtype,
+            generator=self.generator,
+        )
+        scaled_mask = latent_state.denoise_mask * noise_scale
+        latent = noise * scaled_mask + latent_state.latent * (1 - scaled_mask)
+        return replace(
+            latent_state,
+            latent=latent.to(latent_state.latent.dtype),
+        )

packages/ltx-core/src/ltx_core/components/patchifiers.py

@@ -0,0 +1,353 @@
+import math
+from typing import Optional, Tuple
+
+import einops
+import torch
+
+from ltx_core.components.protocols import Patchifier
+from ltx_core.types import AudioLatentShape, SpatioTemporalScaleFactors, VideoLatentShape
+
+
+class VideoLatentPatchifier(Patchifier):
+    def __init__(self, patch_size: int):
+        # Patch sizes for video latents.
+        self._patch_size = (
+            1,  # temporal dimension
+            patch_size,  # height dimension
+            patch_size,  # width dimension
+        )
+
+    @property
+    def patch_size(self) -> Tuple[int, int, int]:
+        return self._patch_size
+
+    def get_token_count(self, tgt_shape: VideoLatentShape) -> int:
+        return math.prod(tgt_shape.to_torch_shape()[2:]) // math.prod(self._patch_size)
+
+    def patchify(
+        self,
+        latents: torch.Tensor,
+    ) -> torch.Tensor:
+        latents = einops.rearrange(
+            latents,
+            "b c (f p1) (h p2) (w p3) -> b (f h w) (c p1 p2 p3)",
+            p1=self._patch_size[0],
+            p2=self._patch_size[1],
+            p3=self._patch_size[2],
+        )
+
+        return latents
+
+    def unpatchify(
+        self,
+        latents: torch.Tensor,
+        output_shape: VideoLatentShape,
+    ) -> torch.Tensor:
+        assert self._patch_size[0] == 1, "Temporal patch size must be 1 for symmetric patchifier"
+
+        patch_grid_frames = output_shape.frames // self._patch_size[0]
+        patch_grid_height = output_shape.height // self._patch_size[1]
+        patch_grid_width = output_shape.width // self._patch_size[2]
+
+        latents = einops.rearrange(
+            latents,
+            "b (f h w) (c p q) -> b c f (h p) (w q)",
+            f=patch_grid_frames,
+            h=patch_grid_height,
+            w=patch_grid_width,
+            p=self._patch_size[1],
+            q=self._patch_size[2],
+        )
+
+        return latents
+
+    def get_patch_grid_bounds(
+        self,
+        output_shape: AudioLatentShape | VideoLatentShape,
+        device: Optional[torch.device] = None,
+    ) -> torch.Tensor:
+        """
+        Return the per-dimension bounds [inclusive start, exclusive end) for every
+        patch produced by `patchify`. The bounds are expressed in the original
+        video grid coordinates: frame/time, height, and width.
+        The resulting tensor is shaped `[batch_size, 3, num_patches, 2]`, where:
+            - axis 1 (size 3) enumerates (frame/time, height, width) dimensions
+            - axis 3 (size 2) stores `[start, end)` indices within each dimension
+        Args:
+            output_shape: Video grid description containing frames, height, and width.
+            device: Device of the latent tensor.
+        """
+        if not isinstance(output_shape, VideoLatentShape):
+            raise ValueError("VideoLatentPatchifier expects VideoLatentShape when computing coordinates")
+
+        frames = output_shape.frames
+        height = output_shape.height
+        width = output_shape.width
+        batch_size = output_shape.batch
+
+        # Validate inputs to ensure positive dimensions
+        assert frames > 0, f"frames must be positive, got {frames}"
+        assert height > 0, f"height must be positive, got {height}"
+        assert width > 0, f"width must be positive, got {width}"
+        assert batch_size > 0, f"batch_size must be positive, got {batch_size}"
+
+        # Generate grid coordinates for each dimension (frame, height, width)
+        # We use torch.arange to create the starting coordinates for each patch.
+        # indexing='ij' ensures the dimensions are in the order (frame, height, width).
+        grid_coords = torch.meshgrid(
+            torch.arange(start=0, end=frames, step=self._patch_size[0], device=device),
+            torch.arange(start=0, end=height, step=self._patch_size[1], device=device),
+            torch.arange(start=0, end=width, step=self._patch_size[2], device=device),
+            indexing="ij",
+        )
+
+        # Stack the grid coordinates to create the start coordinates tensor.
+        # Shape becomes (3, grid_f, grid_h, grid_w)
+        patch_starts = torch.stack(grid_coords, dim=0)
+
+        # Create a tensor containing the size of a single patch:
+        # (frame_patch_size, height_patch_size, width_patch_size).
+        # Reshape to (3, 1, 1, 1) to enable broadcasting when adding to the start coordinates.
+        patch_size_delta = torch.tensor(
+            self._patch_size,
+            device=patch_starts.device,
+            dtype=patch_starts.dtype,
+        ).view(3, 1, 1, 1)
+
+        # Calculate end coordinates: start + patch_size
+        # Shape becomes (3, grid_f, grid_h, grid_w)
+        patch_ends = patch_starts + patch_size_delta
+
+        # Stack start and end coordinates together along the last dimension
+        # Shape becomes (3, grid_f, grid_h, grid_w, 2), where the last dimension is [start, end]
+        latent_coords = torch.stack((patch_starts, patch_ends), dim=-1)
+
+        # Broadcast to batch size and flatten all spatial/temporal dimensions into one sequence.
+        # Final Shape: (batch_size, 3, num_patches, 2)
+        latent_coords = einops.repeat(
+            latent_coords,
+            "c f h w bounds -> b c (f h w) bounds",
+            b=batch_size,
+            bounds=2,
+        )
+
+        return latent_coords
+
+
+def get_pixel_coords(
+    latent_coords: torch.Tensor,
+    scale_factors: SpatioTemporalScaleFactors,
+    causal_fix: bool = False,
+) -> torch.Tensor:
+    """
+    Map latent-space `[start, end)` coordinates to their pixel-space equivalents by scaling
+    each axis (frame/time, height, width) with the corresponding VAE downsampling factors.
+    Optionally compensate for causal encoding that keeps the first frame at unit temporal scale.
+    Args:
+        latent_coords: Tensor of latent bounds shaped `(batch, 3, num_patches, 2)`.
+        scale_factors: SpatioTemporalScaleFactors tuple `(temporal, height, width)` with integer scale factors applied
+        per axis.
+        causal_fix: When True, rewrites the temporal axis of the first frame so causal VAEs
+            that treat frame zero differently still yield non-negative timestamps.
+    """
+    # Broadcast the VAE scale factors so they align with the `(batch, axis, patch, bound)` layout.
+    # Axis 1 of `latent_coords` is ordered (frame/time, height, width) — match that explicitly by
+    # pulling fields from the NamedTuple rather than relying on tuple iteration order.
+    broadcast_shape = [1] * latent_coords.ndim
+    broadcast_shape[1] = -1  # axis dimension corresponds to (frame/time, height, width)
+    scale_tensor = torch.tensor(
+        [scale_factors.time, scale_factors.height, scale_factors.width],
+        device=latent_coords.device,
+    ).view(*broadcast_shape)
+
+    # Apply per-axis scaling to convert latent bounds into pixel-space coordinates.
+    pixel_coords = latent_coords * scale_tensor
+
+    if causal_fix:
+        # VAE temporal stride for the very first frame is 1 instead of `scale_factors.time`.
+        # Shift and clamp to keep the first-frame timestamps causal and non-negative.
+        pixel_coords[:, 0, ...] = (pixel_coords[:, 0, ...] + 1 - scale_factors.time).clamp(min=0)
+
+    return pixel_coords
+
+
+class AudioPatchifier(Patchifier):
+    def __init__(
+        self,
+        patch_size: int,
+        sample_rate: int = 16000,
+        hop_length: int = 160,
+        audio_latent_downsample_factor: int = 4,
+        is_causal: bool = True,
+        shift: int = 0,
+    ):
+        """
+        Patchifier tailored for spectrogram/audio latents.
+        Args:
+            patch_size: Number of mel bins combined into a single patch. This
+                controls the resolution along the frequency axis.
+            sample_rate: Original waveform sampling rate. Used to map latent
+                indices back to seconds so downstream consumers can align audio
+                and video cues.
+            hop_length: Window hop length used for the spectrogram. Determines
+                how many real-time samples separate two consecutive latent frames.
+            audio_latent_downsample_factor: Ratio between spectrogram frames and
+                latent frames; compensates for additional downsampling inside the
+                VAE encoder.
+            is_causal: When True, timing is shifted to account for causal
+                receptive fields so timestamps do not peek into the future.
+            shift: Integer offset applied to the latent indices. Enables
+                constructing overlapping windows from the same latent sequence.
+        """
+        self.hop_length = hop_length
+        self.sample_rate = sample_rate
+        self.audio_latent_downsample_factor = audio_latent_downsample_factor
+        self.is_causal = is_causal
+        self.shift = shift
+        self._patch_size = (1, patch_size, patch_size)
+
+    @property
+    def patch_size(self) -> Tuple[int, int, int]:
+        return self._patch_size
+
+    def get_token_count(self, tgt_shape: AudioLatentShape) -> int:
+        return tgt_shape.frames
+
+    def _get_audio_latent_time_in_sec(
+        self,
+        start_latent: int,
+        end_latent: int,
+        dtype: torch.dtype,
+        device: Optional[torch.device] = None,
+    ) -> torch.Tensor:
+        """
+        Converts latent indices into real-time seconds while honoring causal
+        offsets and the configured hop length.
+        Args:
+            start_latent: Inclusive start index inside the latent sequence. This
+                sets the first timestamp returned.
+            end_latent: Exclusive end index. Determines how many timestamps get
+                generated.
+            dtype: Floating-point dtype used for the returned tensor, allowing
+                callers to control precision.
+            device: Target device for the timestamp tensor. When omitted the
+                computation occurs on CPU to avoid surprising GPU allocations.
+        """
+        if device is None:
+            device = torch.device("cpu")
+
+        audio_latent_frame = torch.arange(start_latent, end_latent, dtype=dtype, device=device)
+
+        audio_mel_frame = audio_latent_frame * self.audio_latent_downsample_factor
+
+        if self.is_causal:
+            # Frame offset for causal alignment.
+            # The "+1" ensures the timestamp corresponds to the first sample that is fully available.
+            causal_offset = 1
+            audio_mel_frame = (audio_mel_frame + causal_offset - self.audio_latent_downsample_factor).clip(min=0)
+
+        return audio_mel_frame * self.hop_length / self.sample_rate
+
+    def _compute_audio_timings(
+        self,
+        batch_size: int,
+        num_steps: int,
+        device: Optional[torch.device] = None,
+    ) -> torch.Tensor:
+        """
+        Builds a `(B, 1, T, 2)` tensor containing timestamps for each latent frame.
+        This helper method underpins `get_patch_grid_bounds` for the audio patchifier.
+        Args:
+            batch_size: Number of sequences to broadcast the timings over.
+            num_steps: Number of latent frames (time steps) to convert into timestamps.
+            device: Device on which the resulting tensor should reside.
+        """
+        resolved_device = device
+        if resolved_device is None:
+            resolved_device = torch.device("cpu")
+
+        start_timings = self._get_audio_latent_time_in_sec(
+            self.shift,
+            num_steps + self.shift,
+            torch.float32,
+            resolved_device,
+        )
+        start_timings = start_timings.unsqueeze(0).expand(batch_size, -1).unsqueeze(1)
+
+        end_timings = self._get_audio_latent_time_in_sec(
+            self.shift + 1,
+            num_steps + self.shift + 1,
+            torch.float32,
+            resolved_device,
+        )
+        end_timings = end_timings.unsqueeze(0).expand(batch_size, -1).unsqueeze(1)
+
+        return torch.stack([start_timings, end_timings], dim=-1)
+
+    def patchify(
+        self,
+        audio_latents: torch.Tensor,
+    ) -> torch.Tensor:
+        """
+        Flattens the audio latent tensor along time. Use `get_patch_grid_bounds`
+        to derive timestamps for each latent frame based on the configured hop
+        length and downsampling.
+        Args:
+            audio_latents: Latent tensor to patchify.
+        Returns:
+            Flattened patch tokens tensor. Use `get_patch_grid_bounds` to compute the
+            corresponding timing metadata when needed.
+        """
+        audio_latents = einops.rearrange(
+            audio_latents,
+            "b c t f -> b t (c f)",
+        )
+
+        return audio_latents
+
+    def unpatchify(
+        self,
+        audio_latents: torch.Tensor,
+        output_shape: AudioLatentShape,
+    ) -> torch.Tensor:
+        """
+        Restores the `(B, C, T, F)` spectrogram tensor from flattened patches.
+        Use `get_patch_grid_bounds` to recompute the timestamps that describe each
+        frame's position in real time.
+        Args:
+            audio_latents: Latent tensor to unpatchify.
+            output_shape: Shape of the unpatched output tensor.
+        Returns:
+            Unpatched latent tensor. Use `get_patch_grid_bounds` to compute the timing
+            metadata associated with the restored latents.
+        """
+        # audio_latents shape: (batch, time, freq * channels)
+        audio_latents = einops.rearrange(
+            audio_latents,
+            "b t (c f) -> b c t f",
+            c=output_shape.channels,
+            f=output_shape.mel_bins,
+        )
+
+        return audio_latents
+
+    def get_patch_grid_bounds(
+        self,
+        output_shape: AudioLatentShape | VideoLatentShape,
+        device: Optional[torch.device] = None,
+    ) -> torch.Tensor:
+        """
+        Return the temporal bounds `[inclusive start, exclusive end)` for every
+        patch emitted by `patchify`. For audio this corresponds to timestamps in
+        seconds aligned with the original spectrogram grid.
+        The returned tensor has shape `[batch_size, 1, time_steps, 2]`, where:
+            - axis 1 (size 1) represents the temporal dimension
+            - axis 3 (size 2) stores the `[start, end)` timestamps per patch
+        Args:
+            output_shape: Audio grid specification describing the number of time steps.
+            device: Target device for the returned tensor.
+        """
+        if not isinstance(output_shape, AudioLatentShape):
+            raise ValueError("AudioPatchifier expects AudioLatentShape when computing coordinates")
+
+        return self._compute_audio_timings(output_shape.batch, output_shape.frames, device)

packages/ltx-core/src/ltx_core/components/protocols.py

@@ -0,0 +1,101 @@
+from typing import Protocol, Tuple
+
+import torch
+
+from ltx_core.types import AudioLatentShape, VideoLatentShape
+
+
+class Patchifier(Protocol):
+    """
+    Protocol for patchifiers that convert latent tensors into patches and assemble them back.
+    """
+
+    def patchify(
+        self,
+        latents: torch.Tensor,
+    ) -> torch.Tensor:
+        ...
+        """
+        Convert latent tensors into flattened patch tokens.
+        Args:
+            latents: Latent tensor to patchify.
+        Returns:
+            Flattened patch tokens tensor.
+        """
+
+    def unpatchify(
+        self,
+        latents: torch.Tensor,
+        output_shape: AudioLatentShape | VideoLatentShape,
+    ) -> torch.Tensor:
+        """
+        Converts latent tensors between spatio-temporal formats and flattened sequence representations.
+        Args:
+            latents: Patch tokens that must be rearranged back into the latent grid constructed by `patchify`.
+            output_shape: Shape of the output tensor. Note that output_shape is either AudioLatentShape or
+            VideoLatentShape.
+        Returns:
+            Dense latent tensor restored from the flattened representation.
+        """
+
+    @property
+    def patch_size(self) -> Tuple[int, int, int]:
+        ...
+        """
+        Returns the patch size as a tuple of (temporal, height, width) dimensions
+        """
+
+    def get_patch_grid_bounds(
+        self,
+        output_shape: AudioLatentShape | VideoLatentShape,
+        device: torch.device | None = None,
+    ) -> torch.Tensor:
+        ...
+        """
+        Compute metadata describing where each latent patch resides within the
+        grid specified by `output_shape`.
+        Args:
+            output_shape: Target grid layout for the patches.
+            device: Target device for the returned tensor.
+        Returns:
+            Tensor containing patch coordinate metadata such as spatial or temporal intervals.
+        """
+
+
+class SchedulerProtocol(Protocol):
+    """
+    Protocol for schedulers that provide a sigmas schedule tensor for a
+    given number of steps. Device is cpu.
+    """
+
+    def execute(self, steps: int, **kwargs) -> torch.FloatTensor: ...
+
+
+class GuiderProtocol(Protocol):
+    """
+    Protocol for guiders that compute a delta tensor given conditioning inputs.
+    The returned delta should be added to the conditional output (cond), enabling
+    multiple guiders to be chained together by accumulating their deltas.
+    """
+
+    scale: float
+
+    def delta(self, cond: torch.Tensor, uncond: torch.Tensor) -> torch.Tensor: ...
+
+    def enabled(self) -> bool:
+        """
+        Returns whether the corresponding perturbation is enabled. E.g. for CFG, this should return False if the scale
+        is 1.0.
+        """
+        ...
+
+
+class DiffusionStepProtocol(Protocol):
+    """
+    Protocol for diffusion steps that provide a next sample tensor for a given current sample tensor,
+    current denoised sample tensor, and sigmas tensor.
+    """
+
+    def step(
+        self, sample: torch.Tensor, denoised_sample: torch.Tensor, sigmas: torch.Tensor, step_index: int, **kwargs
+    ) -> torch.Tensor: ...

packages/ltx-core/src/ltx_core/components/schedulers.py

@@ -0,0 +1,130 @@
+import math
+from functools import lru_cache
+
+import numpy
+import scipy
+import torch
+
+from ltx_core.components.protocols import SchedulerProtocol
+
+BASE_SHIFT_ANCHOR = 1024
+MAX_SHIFT_ANCHOR = 4096
+
+
+class LTX2Scheduler(SchedulerProtocol):
+    """
+    Default scheduler for LTX-2 diffusion sampling.
+    Generates a sigma schedule with token-count-dependent shifting and optional
+    stretching to a terminal value.
+    """
+
+    def execute(
+        self,
+        steps: int,
+        latent: torch.Tensor | None = None,
+        max_shift: float = 2.05,
+        base_shift: float = 0.95,
+        stretch: bool = True,
+        terminal: float = 0.1,
+        default_number_of_tokens: int = MAX_SHIFT_ANCHOR,
+        **_kwargs,
+    ) -> torch.FloatTensor:
+        tokens = math.prod(latent.shape[2:]) if latent is not None else default_number_of_tokens
+        sigmas = torch.linspace(1.0, 0.0, steps + 1)
+
+        x1 = BASE_SHIFT_ANCHOR
+        x2 = MAX_SHIFT_ANCHOR
+        mm = (max_shift - base_shift) / (x2 - x1)
+        b = base_shift - mm * x1
+        sigma_shift = (tokens) * mm + b
+
+        power = 1
+        sigmas = torch.where(
+            sigmas != 0,
+            math.exp(sigma_shift) / (math.exp(sigma_shift) + (1 / sigmas - 1) ** power),
+            0,
+        )
+
+        # Stretch sigmas so that its final value matches the given terminal value.
+        if stretch:
+            non_zero_mask = sigmas != 0
+            non_zero_sigmas = sigmas[non_zero_mask]
+            one_minus_z = 1.0 - non_zero_sigmas
+            scale_factor = one_minus_z[-1] / (1.0 - terminal)
+            stretched = 1.0 - (one_minus_z / scale_factor)
+            sigmas[non_zero_mask] = stretched
+
+        return sigmas.to(torch.float32)
+
+
+class LinearQuadraticScheduler(SchedulerProtocol):
+    """
+    Scheduler with linear steps followed by quadratic steps.
+    Produces a sigma schedule that transitions linearly up to a threshold,
+    then follows a quadratic curve for the remaining steps.
+    """
+
+    def execute(
+        self, steps: int, threshold_noise: float = 0.025, linear_steps: int | None = None, **_kwargs
+    ) -> torch.FloatTensor:
+        if steps == 1:
+            return torch.FloatTensor([1.0, 0.0])
+
+        if linear_steps is None:
+            linear_steps = steps // 2
+        linear_sigma_schedule = [i * threshold_noise / linear_steps for i in range(linear_steps)]
+        threshold_noise_step_diff = linear_steps - threshold_noise * steps
+        quadratic_steps = steps - linear_steps
+        quadratic_sigma_schedule = []
+        if quadratic_steps > 0:
+            quadratic_coef = threshold_noise_step_diff / (linear_steps * quadratic_steps**2)
+            linear_coef = threshold_noise / linear_steps - 2 * threshold_noise_step_diff / (quadratic_steps**2)
+            const = quadratic_coef * (linear_steps**2)
+            quadratic_sigma_schedule = [
+                quadratic_coef * (i**2) + linear_coef * i + const for i in range(linear_steps, steps)
+            ]
+        sigma_schedule = linear_sigma_schedule + quadratic_sigma_schedule + [1.0]
+        sigma_schedule = [1.0 - x for x in sigma_schedule]
+        return torch.FloatTensor(sigma_schedule)
+
+
+class BetaScheduler(SchedulerProtocol):
+    """
+    Scheduler using a beta distribution to sample timesteps.
+    Based on: https://arxiv.org/abs/2407.12173
+    """
+
+    shift = 2.37
+    timesteps_length = 10000
+
+    def execute(self, steps: int, alpha: float = 0.6, beta: float = 0.6) -> torch.FloatTensor:
+        """
+        Execute the beta scheduler.
+        Args:
+            steps: The number of steps to execute the scheduler for.
+            alpha: The alpha parameter for the beta distribution.
+            beta: The beta parameter for the beta distribution.
+        Warnings:
+            The number of steps within `sigmas` theoretically might be less than `steps+1`,
+            because of the deduplication of the identical timesteps
+        Returns:
+            A tensor of sigmas.
+        """
+        model_sampling_sigmas = _precalculate_model_sampling_sigmas(self.shift, self.timesteps_length)
+        total_timesteps = len(model_sampling_sigmas) - 1
+        ts = 1 - numpy.linspace(0, 1, steps, endpoint=False)
+        ts = numpy.rint(scipy.stats.beta.ppf(ts, alpha, beta) * total_timesteps).tolist()
+        ts = list(dict.fromkeys(ts))
+
+        sigmas = [float(model_sampling_sigmas[int(t)]) for t in ts] + [0.0]
+        return torch.FloatTensor(sigmas)
+
+
+@lru_cache(maxsize=5)
+def _precalculate_model_sampling_sigmas(shift: float, timesteps_length: int) -> torch.Tensor:
+    timesteps = torch.arange(1, timesteps_length + 1, 1) / timesteps_length
+    return torch.Tensor([flux_time_shift(shift, 1.0, t) for t in timesteps])
+
+
+def flux_time_shift(mu: float, sigma: float, t: float) -> float:
+    return math.exp(mu) / (math.exp(mu) + (1 / t - 1) ** sigma)

packages/ltx-core/src/ltx_core/conditioning/__init__.py

@@ -0,0 +1,21 @@
+"""Conditioning utilities: latent state, tools, and conditioning types."""
+
+from ltx_core.conditioning.exceptions import ConditioningError
+from ltx_core.conditioning.item import ConditioningItem
+from ltx_core.conditioning.types import (
+    AudioConditionByReferenceLatent,
+    ConditioningItemAttentionStrengthWrapper,
+    VideoConditionByKeyframeIndex,
+    VideoConditionByLatentIndex,
+    VideoConditionByReferenceLatent,
+)
+
+__all__ = [
+    "AudioConditionByReferenceLatent",
+    "ConditioningError",
+    "ConditioningItem",
+    "ConditioningItemAttentionStrengthWrapper",
+    "VideoConditionByKeyframeIndex",
+    "VideoConditionByLatentIndex",
+    "VideoConditionByReferenceLatent",
+]

packages/ltx-core/src/ltx_core/conditioning/exceptions.py

@@ -0,0 +1,4 @@
+class ConditioningError(Exception):
+    """
+    Class for conditioning-related errors.
+    """

packages/ltx-core/src/ltx_core/conditioning/item.py

@@ -0,0 +1,20 @@
+from typing import Protocol
+
+from ltx_core.tools import LatentTools
+from ltx_core.types import LatentState
+
+
+class ConditioningItem(Protocol):
+    """Protocol for conditioning items that modify latent state during diffusion."""
+
+    def apply_to(self, latent_state: LatentState, latent_tools: LatentTools) -> LatentState:
+        """
+        Apply the conditioning to the latent state.
+        Args:
+            latent_state: The latent state to apply the conditioning to. This is state always patchified.
+        Returns:
+            The latent state after the conditioning has been applied.
+        IMPORTANT: If the conditioning needs to add extra tokens to the latent, it should add them to the end of the
+        latent.
+        """
+        ...

packages/ltx-core/src/ltx_core/conditioning/mask_utils.py

@@ -0,0 +1,210 @@
+"""Utilities for building 2D self-attention masks for conditioning items."""
+
+from __future__ import annotations
+
+from typing import TYPE_CHECKING
+
+import torch
+
+if TYPE_CHECKING:
+    from ltx_core.types import LatentState
+
+
+def resolve_cross_mask(
+    attention_mask: float | int | torch.Tensor,
+    num_new_tokens: int,
+    batch_size: int,
+    device: torch.device,
+    dtype: torch.dtype,
+) -> torch.Tensor:
+    """Convert an attention_mask (scalar or tensor) to a (B, M) cross_mask tensor.
+    Args:
+        attention_mask: Scalar value applied uniformly, 1D tensor of shape (M,)
+            broadcast across batch, or 2D tensor of shape (B, M).
+        num_new_tokens: Number of new conditioning tokens M.
+        batch_size: Batch size B.
+        device: Device for the output tensor.
+        dtype: Data type for the output tensor.
+    Returns:
+        Cross-mask tensor of shape (B, M).
+    """
+    if isinstance(attention_mask, (int, float)):
+        return torch.full(
+            (batch_size, num_new_tokens),
+            fill_value=float(attention_mask),
+            device=device,
+            dtype=dtype,
+        )
+    mask = attention_mask.to(device=device, dtype=dtype)
+
+    # Handle scalar (0-D) tensor like a Python scalar.
+    if mask.dim() == 0:
+        return torch.full(
+            (batch_size, num_new_tokens),
+            fill_value=float(mask.item()),
+            device=device,
+            dtype=dtype,
+        )
+
+    if mask.dim() == 1:
+        if mask.shape[0] != num_new_tokens:
+            raise ValueError(
+                f"1-D attention_mask length must equal num_new_tokens ({num_new_tokens}), got shape {tuple(mask.shape)}"
+            )
+        mask = mask.unsqueeze(0).expand(batch_size, -1)
+    elif mask.dim() == 2:
+        b, m = mask.shape
+        if m != num_new_tokens:
+            raise ValueError(
+                f"2-D attention_mask second dimension must equal num_new_tokens ({num_new_tokens}), "
+                f"got shape {tuple(mask.shape)}"
+            )
+        if b not in (batch_size, 1):
+            raise ValueError(
+                f"2-D attention_mask batch dimension must equal batch_size ({batch_size}) or 1, "
+                f"got shape {tuple(mask.shape)}"
+            )
+        if b == 1 and batch_size > 1:
+            mask = mask.expand(batch_size, -1)
+    else:
+        raise ValueError(
+            f"attention_mask tensor must be 0-D, 1-D, or 2-D, got {mask.dim()}-D with shape {tuple(mask.shape)}"
+        )
+    return mask
+
+
+def update_attention_mask(
+    latent_state: LatentState,
+    attention_mask: float | torch.Tensor | None,
+    num_noisy_tokens: int,
+    num_new_tokens: int,
+    batch_size: int,
+    device: torch.device,
+    dtype: torch.dtype,
+) -> torch.Tensor | None:
+    """Build or update the self-attention mask for newly appended conditioning tokens.
+    If *attention_mask* is ``None`` and no existing mask is present, returns
+    ``None``.  If *attention_mask* is ``None`` but an existing mask is present,
+    the mask is expanded with full attention (1s) for the new tokens so that
+    its dimensions stay consistent with the growing latent sequence.  Otherwise,
+    resolves *attention_mask* to a per-token cross-mask and expands the 2-D
+    attention mask via :func:`build_attention_mask`.
+    Args:
+        latent_state: Current latent state (provides the existing mask and total
+            existing-token count).
+        attention_mask: Per-token attention weight. Scalar, 1-D ``(M,)``, 2-D
+            ``(B, M)`` tensor, or ``None`` (no-op).
+        num_noisy_tokens: Number of original noisy tokens (from
+            ``latent_tools.target_shape.token_count()``).
+        num_new_tokens: Number of new conditioning tokens being appended.
+        batch_size: Batch size.
+        device: Device for the output tensor.
+        dtype: Data type for the output tensor.
+    Returns:
+        Updated attention mask of shape ``(B, N+M, N+M)``, or ``None`` if no
+        masking is needed.
+    """
+    if attention_mask is None:
+        if latent_state.attention_mask is None:
+            return None
+        # Existing mask present but no new mask requested: pad with 1s (full
+        # attention) so the mask dimensions stay consistent with the growing
+        # latent sequence.
+        cross_mask = torch.ones(batch_size, num_new_tokens, device=device, dtype=dtype)
+        return build_attention_mask(
+            existing_mask=latent_state.attention_mask,
+            num_noisy_tokens=num_noisy_tokens,
+            num_new_tokens=num_new_tokens,
+            num_existing_tokens=latent_state.latent.shape[1],
+            cross_mask=cross_mask,
+            device=device,
+            dtype=dtype,
+        )
+
+    cross_mask = resolve_cross_mask(attention_mask, num_new_tokens, batch_size, device, dtype)
+    return build_attention_mask(
+        existing_mask=latent_state.attention_mask,
+        num_noisy_tokens=num_noisy_tokens,
+        num_new_tokens=num_new_tokens,
+        num_existing_tokens=latent_state.latent.shape[1],
+        cross_mask=cross_mask,
+        device=device,
+        dtype=dtype,
+    )
+
+
+def build_attention_mask(
+    existing_mask: torch.Tensor | None,
+    num_noisy_tokens: int,
+    num_new_tokens: int,
+    num_existing_tokens: int,
+    cross_mask: torch.Tensor,
+    device: torch.device,
+    dtype: torch.dtype,
+) -> torch.Tensor:
+    """
+    Expand the attention mask to include newly appended conditioning tokens.
+    Each conditioning item appends M new reference tokens to the sequence. This function
+    builds a (B, N+M, N+M) attention mask with the following block structure:
+                     noisy      prev_ref    new_ref
+                   (N_noisy)   (N-N_noisy)    (M)
+                 ┌───────────┬───────────┬───────────┐
+        noisy    │           │           │           │
+       (N_noisy) │  existing │  existing │   cross   │
+                 │           │           │           │
+                 ├───────────┼───────────┼───────────┤
+       prev_ref  │           │           │           │
+      (N-N_noisy)│  existing │  existing │     0     │
+                 │           │           │           │
+                 ├───────────┼───────────┼───────────┤
+       new_ref   │           │           │           │
+         (M)     │   cross   │     0     │     1     │
+                 │           │           │           │
+                 └───────────┴───────────┴───────────┘
+    Where:
+      - **existing**: preserved from the previous mask (or 1.0 if first conditioning)
+      - **cross**: values from *cross_mask* (shape B, M), in [0, 1]
+      - **0**: no attention between different reference groups
+    Args:
+        existing_mask: Current attention mask of shape (B, N, N), or None if no mask exists yet.
+            When None, the top-left NxN block is filled with 1s (full attention between all
+            existing tokens including any prior reference tokens that had no mask).
+        num_noisy_tokens: Number of original noisy tokens (always at positions [0:num_noisy_tokens]).
+        num_new_tokens: Number of new conditioning tokens M being appended.
+        num_existing_tokens: Total number of current tokens N (noisy + any prior conditioning tokens).
+        cross_mask: Per-token attention weight of shape (B, M) controlling attention between
+            new reference tokens and noisy tokens. Values in [0, 1].
+        device: Device for the output tensor.
+        dtype: Data type for the output tensor.
+    Returns:
+        Attention mask of shape (B, N+M, N+M) with values in [0, 1].
+    """
+    batch_size = cross_mask.shape[0]
+    total = num_existing_tokens + num_new_tokens
+
+    # Start with zeros
+    mask = torch.zeros((batch_size, total, total), device=device, dtype=dtype)
+
+    # Top-left: preserve existing mask or fill with 1s for noisy tokens
+    if existing_mask is not None:
+        mask[:, :num_existing_tokens, :num_existing_tokens] = existing_mask
+    else:
+        mask[:, :num_existing_tokens, :num_existing_tokens] = 1.0
+
+    # Bottom-right: new reference tokens fully attend to themselves
+    mask[:, num_existing_tokens:, num_existing_tokens:] = 1.0
+
+    # Cross-attention between noisy tokens and new reference tokens
+    # cross_mask shape: (B, M) -> broadcast to (B, N_noisy, M) and (B, M, N_noisy)
+
+    # Noisy tokens attending to new reference tokens: [0:N_noisy, N:N+M]
+    # Each column j in this block gets cross_mask[:, j]
+    mask[:, :num_noisy_tokens, num_existing_tokens:] = cross_mask.unsqueeze(1)
+
+    # New reference tokens attending to noisy tokens: [N:N+M, 0:N_noisy]
+    # Each row i in this block gets cross_mask[:, i]
+    mask[:, num_existing_tokens:, :num_noisy_tokens] = cross_mask.unsqueeze(2)
+
+    # [N_noisy:N, N:N+M] and [N:N+M, N_noisy:N] remain 0 (no cross-ref attention)
+
+    return mask

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

@@ -0,0 +1,15 @@
+"""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_audio_cond import AudioConditionByReferenceLatent
+from ltx_core.conditioning.types.reference_video_cond import VideoConditionByReferenceLatent
+
+__all__ = [
+    "AudioConditionByReferenceLatent",
+    "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,83 @@
+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).
+        num_pixel_frames: Number of pixel frames the keyframe latent originally encodes.
+    """
+
+    def __init__(
+        self,
+        keyframes: torch.Tensor,
+        frame_idx: int,
+        strength: float,
+        num_pixel_frames: int = 1,
+    ):
+        self.keyframes = keyframes
+        self.frame_idx = frame_idx
+        self.strength = strength
+        self.num_pixel_frames = num_pixel_frames
+
+    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
+        # If the keyframe latent encodes a single pixel frame,
+        # narrow the temporal end to [start, start + 1) instead of the
+        # VAE-scaled range.
+        if self.num_pixel_frames == 1:
+            positions[:, 0, ..., 1:] = positions[:, 0, ..., :1] + 1
+        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/noise_mask_cond.py

@@ -0,0 +1,45 @@
+from dataclasses import dataclass
+
+from ltx_core.components.patchifiers import get_pixel_coords
+from ltx_core.conditioning.item import ConditioningItem
+from ltx_core.tools import LatentTools, SpatioTemporalScaleFactors
+from ltx_core.types import AudioLatentShape, LatentState, VideoLatentShape
+
+
+@dataclass(frozen=True)
+class TemporalRegionMask(ConditioningItem):
+    """Conditioning item that sets ``denoise_mask = 0`` outside a time range
+    and ``1`` inside, so only the specified temporal region is regenerated.
+    Uses ``start_time`` and ``end_time`` in seconds. Works in *patchified*
+    (token) space using the patchifier's ``get_patch_grid_bounds``: for video
+    coords are latent frame indices (converted from seconds via ``fps``), for
+    audio coords are already in seconds.
+    """
+
+    start_time: float  # seconds, inclusive
+    end_time: float  # seconds, exclusive
+    fps: float
+
+    def apply_to(self, latent_state: LatentState, latent_tools: LatentTools) -> LatentState:
+        coords = latent_tools.patchifier.get_patch_grid_bounds(
+            latent_tools.target_shape, device=latent_state.denoise_mask.device
+        )
+        if isinstance(latent_tools.target_shape, AudioLatentShape):
+            # Audio: patchifier get_patch_grid_bounds returns seconds
+            t_boundaries = coords[:, 0]
+        elif isinstance(latent_tools.target_shape, VideoLatentShape):
+            # Video: patchifier get_patch_grid_bounds returns latent bounds, converting to frame numbers & pixel bounds
+            scale_factors = getattr(latent_tools, "scale_factors", SpatioTemporalScaleFactors.default())
+            pixel_bounds = get_pixel_coords(coords, scale_factors, causal_fix=getattr(latent_tools, "causal_fix", True))
+            # converting frame numbers to seconds
+            t_boundaries = pixel_bounds[:, 0] / self.fps
+        else:
+            raise ValueError("Unsupported LatentShape type, expected AudioLatentShape or VideoLatentShape")
+        t_start, t_end = t_boundaries.unbind(dim=-1)  # [B, N]
+        in_region = (t_end > self.start_time) & (t_start < self.end_time)
+        state = latent_state.clone()
+        mask_val = in_region.to(state.denoise_mask.dtype)
+        if state.denoise_mask.dim() == 3:
+            mask_val = mask_val.unsqueeze(-1)
+        state.denoise_mask.copy_(mask_val)
+        return state

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

@@ -0,0 +1,59 @@
+"""Audio reference conditioning items."""
+
+from __future__ import annotations
+
+import torch
+
+from ltx_core.conditioning.mask_utils import update_attention_mask
+from ltx_core.tools import LatentTools
+from ltx_core.types import LatentState
+
+
+class AudioConditionByReferenceLatent:
+    """Append patchified reference audio tokens after the target audio sequence.
+    Mirrors :class:`ltx_core.conditioning.types.reference_video_cond.VideoConditionByReferenceLatent`
+    but for audio. The reference tokens are appended so the target audio tokens stay
+    in the first ``num_noisy_tokens`` positions and can be kept by
+    :meth:`ltx_core.tools.LatentTools.clear_conditioning`.
+    Args:
+        patchified: Patchified reference latent ``[B, T_ref, C]``.
+        positions: RoPE positions for reference tokens, ``[B, 1, T_ref, 2]``.
+        strength: 1.0 keeps reference clean; 0.0 would fully denoise it.
+    """
+
+    def __init__(
+        self,
+        patchified: torch.Tensor,
+        positions: torch.Tensor,
+        strength: float = 1.0,
+    ) -> None:
+        self.patchified = patchified
+        self.positions = positions.to(dtype=torch.float32)
+        self.strength = strength
+
+    def apply_to(self, latent_state: LatentState, latent_tools: LatentTools) -> LatentState:
+        tokens = self.patchified
+        denoise_mask = torch.full(
+            size=(*tokens.shape[:2], 1),
+            fill_value=1.0 - self.strength,
+            device=tokens.device,
+            dtype=tokens.dtype,
+        )
+
+        new_attention_mask = update_attention_mask(
+            latent_state=latent_state,
+            attention_mask=None,
+            num_noisy_tokens=latent_tools.patchifier.get_token_count(latent_tools.target_shape),
+            num_new_tokens=tokens.shape[1],
+            batch_size=tokens.shape[0],
+            device=tokens.device,
+            dtype=tokens.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, self.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/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/guidance/__init__.py

@@ -0,0 +1,15 @@
+"""Guidance and perturbation utilities for attention manipulation."""
+
+from ltx_core.guidance.perturbations import (
+    BatchedPerturbationConfig,
+    Perturbation,
+    PerturbationConfig,
+    PerturbationType,
+)
+
+__all__ = [
+    "BatchedPerturbationConfig",
+    "Perturbation",
+    "PerturbationConfig",
+    "PerturbationType",
+]

packages/ltx-core/src/ltx_core/guidance/perturbations.py

@@ -0,0 +1,79 @@
+from dataclasses import dataclass
+from enum import Enum
+
+import torch
+from torch._prims_common import DeviceLikeType
+
+
+class PerturbationType(Enum):
+    """Types of attention perturbations for STG (Spatio-Temporal Guidance)."""
+
+    SKIP_A2V_CROSS_ATTN = "skip_a2v_cross_attn"
+    SKIP_V2A_CROSS_ATTN = "skip_v2a_cross_attn"
+    SKIP_VIDEO_SELF_ATTN = "skip_video_self_attn"
+    SKIP_AUDIO_SELF_ATTN = "skip_audio_self_attn"
+
+
+@dataclass(frozen=True)
+class Perturbation:
+    """A single perturbation specifying which attention type to skip and in which blocks."""
+
+    type: PerturbationType
+    blocks: list[int] | None  # None means all blocks
+
+    def is_perturbed(self, perturbation_type: PerturbationType, block: int) -> bool:
+        if self.type != perturbation_type:
+            return False
+
+        if self.blocks is None:
+            return True
+
+        return block in self.blocks
+
+
+@dataclass(frozen=True)
+class PerturbationConfig:
+    """Configuration holding a list of perturbations for a single sample."""
+
+    perturbations: list[Perturbation] | None
+
+    def is_perturbed(self, perturbation_type: PerturbationType, block: int) -> bool:
+        if self.perturbations is None:
+            return False
+
+        return any(perturbation.is_perturbed(perturbation_type, block) for perturbation in self.perturbations)
+
+    @staticmethod
+    def empty() -> "PerturbationConfig":
+        return PerturbationConfig([])
+
+
+@dataclass(frozen=True)
+class BatchedPerturbationConfig:
+    """Perturbation configurations for a batch, with utilities for generating attention masks."""
+
+    perturbations: list[PerturbationConfig]
+
+    def mask(
+        self, perturbation_type: PerturbationType, block: int, device: DeviceLikeType, dtype: torch.dtype
+    ) -> torch.Tensor:
+        mask = torch.ones((len(self.perturbations),), device=device, dtype=dtype)
+        for batch_idx, perturbation in enumerate(self.perturbations):
+            if perturbation.is_perturbed(perturbation_type, block):
+                mask[batch_idx] = 0
+
+        return mask
+
+    def mask_like(self, perturbation_type: PerturbationType, block: int, values: torch.Tensor) -> torch.Tensor:
+        mask = self.mask(perturbation_type, block, values.device, values.dtype)
+        return mask.view(mask.numel(), *([1] * len(values.shape[1:])))
+
+    def any_in_batch(self, perturbation_type: PerturbationType, block: int) -> bool:
+        return any(perturbation.is_perturbed(perturbation_type, block) for perturbation in self.perturbations)
+
+    def all_in_batch(self, perturbation_type: PerturbationType, block: int) -> bool:
+        return all(perturbation.is_perturbed(perturbation_type, block) for perturbation in self.perturbations)
+
+    @staticmethod
+    def empty(batch_size: int) -> "BatchedPerturbationConfig":
+        return BatchedPerturbationConfig([PerturbationConfig.empty() for _ in range(batch_size)])

packages/ltx-core/src/ltx_core/hdr.py

@@ -0,0 +1,71 @@
+"""HDR utilities: LogC3 compression for HDR IC-LoRA training and inference.
+Provides compress/decompress and postprocess helpers for HDR video generation.
+Used by ltx-pipelines for HDR IC-LoRA and by ltx-trainer for HDR validation.
+"""
+
+from __future__ import annotations
+
+from typing import Literal
+
+import torch
+from torch import Tensor
+
+
+class LogC3:
+    """ARRI LogC3 (EI 800) HDR compression.
+    Maps linear [0, ∞) <-> LogC3 [0, 1] via the camera log curve. The log
+    curve allocates more precision to shadows/midtones and compresses
+    highlights smoothly. Callers are responsible for mapping the [0, 1]
+    output to the VAE's [-1, 1] input range.
+    """
+
+    name = "LogC3"
+    A = 5.555556
+    B = 0.052272
+    C = 0.247190
+    D = 0.385537
+    E = 5.367655
+    F = 0.092809
+    CUT = 0.010591
+
+    def compress(self, hdr: Tensor) -> Tensor:
+        """Compress linear HDR [0, ∞) → LogC3 [0, 1]."""
+        x = torch.clamp(hdr, min=0.0)
+        log_part = self.C * torch.log10(self.A * x + self.B) + self.D
+        lin_part = self.E * x + self.F
+        logc = torch.where(x >= self.CUT, log_part, lin_part)
+        return torch.clamp(logc, 0.0, 1.0)
+
+    def compress_ldr(self, ldr: Tensor) -> Tensor:
+        """Compress LDR [0, 1] → [0, 1] (no log curve, just clamp)."""
+        return torch.clamp(ldr, 0.0, 1.0)
+
+    def decompress(self, logc: Tensor) -> Tensor:
+        """Decompress LogC3 [0, 1] → linear HDR [0, ∞)."""
+        logc = torch.clamp(logc, 0.0, 1.0)
+        cut_log = self.E * self.CUT + self.F
+        lin_from_log = (torch.pow(10.0, (logc - self.D) / self.C) - self.B) / self.A
+        lin_from_lin = (logc - self.F) / self.E
+        return torch.where(logc >= cut_log, lin_from_log, lin_from_lin)
+
+    def decompress_ldr(self, logc: Tensor) -> Tensor:
+        """Decompress [0, 1] → LDR [0, 1] (identity clamp)."""
+        return torch.clamp(logc, 0.0, 1.0)
+
+
+def apply_hdr_decode_postprocess(
+    decoded_video: Tensor,
+    transform: Literal["logc3"] = "logc3",
+) -> Tensor:
+    """Apply HDR decompress to VAE decode output for HDR recovery.
+    Args:
+        decoded_video: Tensor from VAE decode in [0, 1], shape [B, C, F, H, W].
+            Must be float32 for sufficient color resolution.
+        transform: "logc3".
+    Returns:
+        HDR video tensor float32.
+    """
+    decoded_video = decoded_video.float()
+    if transform == "logc3":
+        return LogC3().decompress(decoded_video)
+    raise ValueError(f"Unsupported HDR transform: {transform}")

packages/ltx-core/src/ltx_core/loader/__init__.py

@@ -0,0 +1,56 @@
+"""Loader utilities for model weights, LoRAs, and safetensor operations."""
+
+from ltx_core.loader.fuse_loras import apply_loras
+from ltx_core.loader.helpers import (
+    create_meta_model,
+    load_state_dict,
+    read_model_config,
+)
+from ltx_core.loader.module_ops import ModuleOps
+from ltx_core.loader.primitives import (
+    LoRAAdaptableProtocol,
+    LoraPathStrengthAndSDOps,
+    LoraStateDictWithStrength,
+    ModelBuilderProtocol,
+    StateDict,
+    StateDictLoader,
+)
+from ltx_core.loader.registry import DummyRegistry, Registry, StateDictRegistry
+from ltx_core.loader.sd_ops import (
+    LTXV_LORA_COMFY_RENAMING_MAP,
+    ContentMatching,
+    ContentReplacement,
+    KeyValueOperation,
+    KeyValueOperationResult,
+    SDKeyValueOperation,
+    SDOps,
+)
+from ltx_core.loader.sft_loader import SafetensorsModelStateDictLoader, SafetensorsStateDictLoader
+from ltx_core.loader.single_gpu_model_builder import SingleGPUModelBuilder
+
+__all__ = [
+    "LTXV_LORA_COMFY_RENAMING_MAP",
+    "ContentMatching",
+    "ContentReplacement",
+    "DummyRegistry",
+    "KeyValueOperation",
+    "KeyValueOperationResult",
+    "LoRAAdaptableProtocol",
+    "LoraPathStrengthAndSDOps",
+    "LoraStateDictWithStrength",
+    "ModelBuilderProtocol",
+    "ModuleOps",
+    "Registry",
+    "SDKeyValueOperation",
+    "SDOps",
+    "SafetensorsModelStateDictLoader",
+    "SafetensorsStateDictLoader",
+    "SingleGPUModelBuilder",
+    "StateDict",
+    "StateDictLoader",
+    "StateDictRegistry",
+    "apply_loras",
+    "create_meta_model",
+    "load_state_dict",
+    "read_model_config",
+]

packages/ltx-core/src/ltx_core/loader/fuse_loras.py

@@ -0,0 +1,162 @@
+from collections.abc import Iterable, Iterator
+from typing import NamedTuple
+
+import torch
+
+from ltx_core.loader.kernels import TRITON_AVAILABLE
+from ltx_core.loader.primitives import LoraStateDictWithStrength, StateDict
+from ltx_core.quantization.fp8_cast import fused_add_round_launch
+from ltx_core.quantization.fp8_scaled_mm import quantize_weight_to_fp8_per_tensor
+
+
+class LoraProduct(NamedTuple):
+    """A LoRA's ``A``, ``B`` factors and its strength scalar."""
+
+    a: torch.Tensor
+    b: torch.Tensor
+    strength: float
+
+
+def _get_device() -> torch.device:
+    if torch.cuda.is_available():
+        return torch.device("cuda", torch.cuda.current_device())
+    return torch.device("cpu")
+
+
+def aggregate_lora_products(
+    products: Iterable[LoraProduct],
+    dtype: torch.dtype | None = None,
+    *,
+    out: torch.Tensor | None = None,
+) -> torch.Tensor | None:
+    """Accumulate ``sum((B * strength) @ A)`` across :class:`LoraProduct` items.
+    If ``out`` is provided, ``addmm_`` accumulates directly into it — caller
+    ensures A/B dtypes and devices match ``out``. Otherwise the first product
+    materializes the ``(out, in)``-shape aggregator at ``dtype``; subsequent
+    products use ``addmm_`` to avoid allocating the full intermediate delta.
+    Returns ``out`` (or the new aggregator), or ``None`` if ``products`` was empty
+    and ``out`` was not given.
+    """
+    aggregated = out
+    for product in products:
+        if aggregated is None:
+            aggregated = torch.matmul(product.b * product.strength, product.a).to(dtype=dtype)
+        else:
+            aggregated.addmm_(product.b, product.a, alpha=product.strength)
+    return aggregated
+
+
+def fuse_cast_fp8_weight(
+    delta_bf16: torch.Tensor,
+    weight_fp8: torch.Tensor,
+    target_dtype: torch.dtype,
+) -> torch.Tensor:
+    """Return ``(delta_bf16 + dequantize(weight_fp8)).to(target_dtype)``.
+    CUDA with Triton uses stochastic rounding; otherwise uses a deterministic bf16 add.
+    ``delta_bf16`` is the bf16 accumulator and is mutated in place.
+    """
+    if delta_bf16.dtype != torch.bfloat16:
+        raise ValueError(f"delta_bf16 must be bfloat16, got {delta_bf16.dtype}")
+    if str(weight_fp8.device).startswith("cuda") and TRITON_AVAILABLE:
+        fused_add_round_launch(delta_bf16, weight_fp8, seed=0)
+    else:
+        delta_bf16.add_(weight_fp8.to(dtype=torch.bfloat16))
+    return delta_bf16.to(dtype=target_dtype)
+
+
+def fuse_lora_weights(
+    model_sd: StateDict,
+    lora_sd_and_strengths: list[LoraStateDictWithStrength],
+    dtype: torch.dtype | None = None,
+    preserve_input_device: bool = True,
+) -> Iterator[tuple[str, torch.Tensor]]:
+    """Yield ``(key, fused_tensor)`` for each weight modified by at least one LoRA.
+    For scaled-FP8 weights, this includes both the updated ``.weight`` tensor
+    and its corresponding ``.weight_scale`` tensor.
+    When ``preserve_input_device`` is False, fused tensors are yielded on the device
+    used for fusion; caller is responsible for moving them to their final
+    destination.
+    """
+    for key, original_weight in model_sd.sd.items():
+        if original_weight is None or key.endswith(".weight_scale"):
+            continue
+        original_device = original_weight.device
+        weight = original_weight.to(device=_get_device())
+        target_dtype = dtype if dtype is not None else weight.dtype
+        deltas_dtype = target_dtype if target_dtype not in [torch.float8_e4m3fn, torch.float8_e5m2] else torch.bfloat16
+
+        deltas = _aggregate_deltas(lora_sd_and_strengths, key, deltas_dtype, weight.device)
+        if deltas is None:
+            continue
+
+        scale_key = key.replace(".weight", ".weight_scale") if key.endswith(".weight") else None
+        is_scaled_fp8 = scale_key is not None and scale_key in model_sd.sd
+
+        if weight.dtype == torch.float8_e4m3fn:
+            if is_scaled_fp8:
+                fused = _fuse_delta_with_scaled_fp8(deltas, weight, key, scale_key, model_sd)
+            else:
+                fused = {key: fuse_cast_fp8_weight(deltas, weight, target_dtype)}
+        elif weight.dtype == torch.bfloat16:
+            deltas.add_(weight)
+            fused = {key: deltas.to(dtype=target_dtype)}
+        else:
+            raise ValueError(f"Unsupported dtype: {weight.dtype}")
+
+        for k, v in fused.items():
+            yield k, v.to(device=original_device) if preserve_input_device else v
+
+
+def apply_loras(
+    model_sd: StateDict,
+    lora_sd_and_strengths: list[LoraStateDictWithStrength],
+    dtype: torch.dtype | None = None,
+    destination_sd: StateDict | None = None,
+) -> StateDict:
+    """Fuse LoRAs into ``model_sd`` and place the results in ``destination_sd``.
+    When ``destination_sd`` is provided, the fused tensors are placed directly into it.
+    """
+    if destination_sd is not None:
+        for key, fused in fuse_lora_weights(model_sd, lora_sd_and_strengths, dtype):
+            destination_sd.sd[key] = fused
+        return destination_sd
+
+    fused = dict(fuse_lora_weights(model_sd, lora_sd_and_strengths, dtype))
+    sd = {k: (fused[k] if k in fused else v.clone()) for k, v in model_sd.sd.items()}
+    return StateDict(sd, model_sd.device, model_sd.size, model_sd.dtype)
+
+
+def _aggregate_deltas(
+    lora_sd_and_strengths: list[LoraStateDictWithStrength], key: str, dtype: torch.dtype, device: torch.device
+) -> torch.Tensor | None:
+    prefix = key[: -len(".weight")]
+    key_a = f"{prefix}.lora_A.weight"
+    key_b = f"{prefix}.lora_B.weight"
+
+    def _ab_products() -> Iterator[LoraProduct]:
+        for lsd, coef in lora_sd_and_strengths:
+            if key_a not in lsd.sd or key_b not in lsd.sd:
+                continue
+            a = lsd.sd[key_a].to(device=device, dtype=dtype, non_blocking=True)
+            b = lsd.sd[key_b].to(device=device, dtype=dtype, non_blocking=True)
+            yield LoraProduct(a, b, coef)
+
+    return aggregate_lora_products(_ab_products(), dtype)
+
+
+def _fuse_delta_with_scaled_fp8(
+    deltas: torch.Tensor,
+    weight: torch.Tensor,
+    key: str,
+    scale_key: str,
+    model_sd: StateDict,
+) -> dict[str, torch.Tensor]:
+    """Dequantize scaled FP8 weight, add LoRA delta, and re-quantize."""
+    weight_scale = model_sd.sd[scale_key]
+
+    original_weight = weight.to(torch.float32) * weight_scale
+
+    new_weight = original_weight + deltas.to(torch.float32)
+
+    new_fp8_weight, new_weight_scale = quantize_weight_to_fp8_per_tensor(new_weight)
+    return {key: new_fp8_weight, scale_key: new_weight_scale}

packages/ltx-core/src/ltx_core/loader/helpers.py

@@ -0,0 +1,61 @@
+"""Shared model-construction helpers used by both SingleGPUModelBuilder and StreamingModelBuilder."""
+
+from __future__ import annotations
+
+from typing import TypeVar
+
+import torch
+from torch import nn
+
+from ltx_core.loader.module_ops import ModuleOps
+from ltx_core.loader.primitives import StateDict, StateDictLoader
+from ltx_core.loader.registry import Registry
+from ltx_core.loader.sd_ops import SDOps
+from ltx_core.model.model_protocol import ModelConfigurator
+
+_M = TypeVar("_M", bound=nn.Module)
+
+
+def load_state_dict(
+    paths: str | tuple[str, ...] | list[str],
+    loader: StateDictLoader,
+    registry: Registry,
+    device: torch.device | None,
+    sd_ops: SDOps | None = None,
+) -> StateDict:
+    """Load a state dict from disk, using registry caching."""
+    if isinstance(paths, str):
+        path_list = [paths]
+    elif isinstance(paths, tuple):
+        path_list = list(paths)
+    else:
+        path_list = paths
+    cached = registry.get(path_list, sd_ops)
+    if cached is not None:
+        return cached
+    result = loader.load(path_list, sd_ops=sd_ops, device=device)
+    registry.add(path_list, sd_ops=sd_ops, state_dict=result)
+    return result
+
+
+def read_model_config(
+    model_path: str | tuple[str, ...],
+    loader: StateDictLoader,
+) -> dict:
+    """Read metadata from the first shard of a checkpoint."""
+    first = model_path[0] if isinstance(model_path, tuple) else model_path
+    return loader.metadata(first)
+
+
+def create_meta_model(
+    configurator: type[ModelConfigurator[_M]],
+    config: dict,
+    module_ops: tuple[ModuleOps, ...] = (),
+) -> _M:
+    """Create a model on the meta device and apply module operations."""
+    with torch.device("meta"):
+        model = configurator.from_config(config)
+    for op in module_ops:
+        if op.matcher(model):
+            model = op.mutator(model)
+    return model

packages/ltx-core/src/ltx_core/loader/kernels.py

@@ -0,0 +1,79 @@
+# ruff: noqa: ANN001, ANN201, ERA001, N803, N806
+try:
+    import triton
+    import triton.language as tl
+
+    TRITON_AVAILABLE = True
+except (ImportError, OSError):
+    TRITON_AVAILABLE = False
+
+
+if TRITON_AVAILABLE:
+
+    @triton.jit
+    def fused_add_round_kernel(
+        x_ptr,
+        output_ptr,  # contents will be added to the output
+        seed,
+        n_elements,
+        EXPONENT_BIAS,
+        MANTISSA_BITS,
+        BLOCK_SIZE: tl.constexpr,
+    ):
+        """
+        A kernel to upcast 8bit quantized weights to bfloat16 with stochastic rounding
+        and add them to bfloat16 output weights. Might be used to upcast original model weights
+        and to further add them to precalculated deltas coming from LoRAs.
+        """
+        # Get program ID and compute offsets
+        pid = tl.program_id(axis=0)
+        block_start = pid * BLOCK_SIZE
+        offsets = block_start + tl.arange(0, BLOCK_SIZE)
+        mask = offsets < n_elements
+
+        # Load data
+        x = tl.load(x_ptr + offsets, mask=mask)
+        rand_vals = tl.rand(seed, offsets) - 0.5
+
+        x = tl.cast(x, tl.float16)
+        delta = tl.load(output_ptr + offsets, mask=mask)
+        delta = tl.cast(delta, tl.float16)
+        x = x + delta
+
+        x_bits = tl.cast(x, tl.int16, bitcast=True)
+
+        # Calculate the exponent. Unbiased fp16 exponent is ((x_bits & 0x7C00) >> 10) - 15 for
+        # normal numbers and -14 for subnormals.
+        fp16_exponent_bits = (x_bits & 0x7C00) >> 10
+        fp16_normals = fp16_exponent_bits > 0
+        fp16_exponent = tl.where(fp16_normals, fp16_exponent_bits - 15, -14)
+
+        # Add the target dtype's exponent bias and clamp to the target dtype's exponent range.
+        exponent = fp16_exponent + EXPONENT_BIAS
+        MAX_EXPONENT = 2 * EXPONENT_BIAS + 1
+        exponent = tl.where(exponent > MAX_EXPONENT, MAX_EXPONENT, exponent)
+        exponent = tl.where(exponent < 0, 0, exponent)
+
+        # Normal ULP exponent, expressed as an fp16 exponent field:
+        # (exponent - EXPONENT_BIAS - MANTISSA_BITS) + 15
+        # Simplifies to: fp16_exponent - MANTISSA_BITS + 15
+        # See https://en.wikipedia.org/wiki/Unit_in_the_last_place
+        eps_exp = tl.maximum(0, tl.minimum(31, exponent - EXPONENT_BIAS - MANTISSA_BITS + 15))
+
+        # Calculate epsilon in the target dtype
+        eps_normal = tl.cast(tl.cast(eps_exp << 10, tl.int16), tl.float16, bitcast=True)
+
+        # Subnormal ULP: 2^(1 - EXPONENT_BIAS - MANTISSA_BITS) ->
+        # fp16 exponent bits: (1 - EXPONENT_BIAS - MANTISSA_BITS) + 15 =
+        # 16 - EXPONENT_BIAS - MANTISSA_BITS
+        eps_subnormal = tl.cast(tl.cast((16 - EXPONENT_BIAS - MANTISSA_BITS) << 10, tl.int16), tl.float16, bitcast=True)
+        eps = tl.where(exponent > 0, eps_normal, eps_subnormal)
+
+        # Apply zero mask to epsilon
+        eps = tl.where(x == 0, 0.0, eps)
+
+        # Apply stochastic rounding
+        output = tl.cast(x + rand_vals * eps, tl.bfloat16)
+
+        # Store the result
+        tl.store(output_ptr + offsets, output, mask=mask)

packages/ltx-core/src/ltx_core/loader/module_ops.py

@@ -0,0 +1,14 @@
+from typing import Callable, NamedTuple
+
+import torch
+
+
+class ModuleOps(NamedTuple):
+    """
+    Defines a named operation for matching and mutating PyTorch modules.
+    Used to selectively transform modules in a model (e.g., replacing layers with quantized versions).
+    """
+
+    name: str
+    matcher: Callable[[torch.nn.Module], bool]
+    mutator: Callable[[torch.nn.Module], torch.nn.Module]

packages/ltx-core/src/ltx_core/loader/primitives.py

@@ -0,0 +1,158 @@
+from __future__ import annotations
+
+from dataclasses import dataclass
+from typing import TYPE_CHECKING, NamedTuple, Protocol
+
+import torch
+
+from ltx_core.loader.module_ops import ModuleOps
+from ltx_core.loader.sd_ops import SDOps
+from ltx_core.model.model_protocol import ModelType
+
+if TYPE_CHECKING:
+    from ltx_core.loader.registry import Registry
+
+
+# Per-key shape and dtype description for a flat collection of tensors.
+TensorLayout = dict[str, tuple[torch.Size, torch.dtype]]
+
+
+@dataclass(frozen=True)
+class StateDict:
+    """
+    Immutable container for a PyTorch state dictionary.
+    Contains:
+    - sd: Dictionary of tensors (weights, buffers, etc.)
+    - device: Device where tensors are stored
+    - size: Total memory footprint in bytes
+    - dtype: Set of tensor dtypes present
+    """
+
+    sd: dict
+    device: torch.device
+    size: int
+    dtype: set[torch.dtype]
+
+    def footprint(self) -> tuple[int, torch.device]:
+        return self.size, self.device
+
+
+class StateDictLoader(Protocol):
+    """
+    Protocol for loading state dictionaries from various sources.
+    Implementations must provide:
+    - metadata: Extract model metadata from a single path
+    - load: Load state dict from path(s) and apply SDOps transformations
+    """
+
+    def metadata(self, path: str) -> dict:
+        """
+        Load metadata from path
+        """
+
+    def load(self, path: str | list[str], sd_ops: SDOps | None = None, device: torch.device | None = None) -> StateDict:
+        """
+        Load state dict from path or paths (for sharded model storage) and apply sd_ops
+        """
+
+
+class BuilderProtocol(Protocol[ModelType]):
+    """Protocol for model builders that produce a model via ``build()``."""
+
+    def build(
+        self, device: torch.device | None = None, dtype: torch.dtype | None = None, **kwargs: object
+    ) -> ModelType: ...
+
+
+class ModelBuilderProtocol(BuilderProtocol[ModelType], Protocol[ModelType]):
+    """
+    Protocol for building PyTorch models from configuration dictionaries.
+    Implementations must provide:
+    - meta_model: Create a model from configuration dictionary and apply module operations
+    - build: Create and initialize a model from state dictionary and apply dtype transformations
+    """
+
+    model_sd_ops: SDOps | None
+    module_ops: tuple[ModuleOps, ...]
+    loras: tuple["LoraPathStrengthAndSDOps", ...]
+    registry: "Registry"
+
+    def meta_model(self, config: dict, module_ops: list[ModuleOps] | None = None) -> ModelType:
+        """
+        Create a model on the meta device from a configuration dictionary.
+        This decouples model creation from weight loading, allowing the model
+        architecture to be instantiated without allocating memory for parameters.
+        Args:
+            config: Model configuration dictionary.
+            module_ops: Optional list of module operations to apply (e.g., quantization).
+        Returns:
+            Model instance on meta device (no actual memory allocated for parameters).
+        """
+        ...
+
+    def with_sd_ops(self, sd_ops: SDOps | None) -> "ModelBuilderProtocol[ModelType]":
+        """Return a copy of this builder with the given state-dict key remapping ops."""
+        ...
+
+    def with_module_ops(self, module_ops: tuple[ModuleOps, ...]) -> "ModelBuilderProtocol[ModelType]":
+        """Return a copy of this builder with the given module operations (e.g. quantization)."""
+        ...
+
+    def with_loras(self, loras: tuple["LoraPathStrengthAndSDOps", ...]) -> "ModelBuilderProtocol[ModelType]":
+        """Return a copy of this builder with the given LoRAs to fuse at build time."""
+        ...
+
+    def with_registry(self, registry: "Registry") -> "ModelBuilderProtocol[ModelType]":
+        """Return a copy of this builder using the given weight registry for allocation."""
+        ...
+
+    def with_lora_load_device(self, device: torch.device) -> "ModelBuilderProtocol[ModelType]":
+        """Return a copy of this builder that loads LoRA weights onto the given device."""
+        ...
+
+    def build(
+        self, device: torch.device | None = None, dtype: torch.dtype | None = None, **kwargs: object
+    ) -> ModelType:
+        """
+        Build the model
+        Args:
+            device: Target device for the model
+            dtype: Target dtype for the model, if None, uses the dtype of the model_path model
+        Returns:
+            Model instance
+        """
+        ...
+
+    def model_config(self) -> dict:
+        """Return the model configuration dictionary extracted from the checkpoint metadata."""
+        ...
+
+
+class LoRAAdaptableProtocol(Protocol):
+    """
+    Protocol for models that can be adapted with LoRAs.
+    Implementations must provide:
+    - lora: Add a LoRA to the model
+    """
+
+    def lora(self, lora_path: str, strength: float) -> "LoRAAdaptableProtocol":
+        pass
+
+
+class LoraPathStrengthAndSDOps(NamedTuple):
+    """
+    Tuple containing a LoRA path, strength, and SDOps for applying to the LoRA state dict.
+    """
+
+    path: str
+    strength: float
+    sd_ops: SDOps
+
+
+class LoraStateDictWithStrength(NamedTuple):
+    """
+    Tuple containing a LoRA state dict and strength for applying to the model.
+    """
+
+    state_dict: StateDict
+    strength: float

packages/ltx-core/src/ltx_core/loader/registry.py

@@ -0,0 +1,84 @@
+import hashlib
+import threading
+from dataclasses import dataclass, field
+from pathlib import Path
+from typing import Protocol
+
+from ltx_core.loader.primitives import StateDict
+from ltx_core.loader.sd_ops import SDOps
+
+
+class Registry(Protocol):
+    """
+    Protocol for managing state dictionaries in a registry.
+    It is used to store state dictionaries and reuse them later without loading them again.
+    Implementations must provide:
+    - add: Add a state dictionary to the registry
+    - pop: Remove a state dictionary from the registry
+    - get: Retrieve a state dictionary from the registry
+    - clear: Clear all state dictionaries from the registry
+    """
+
+    def add(self, paths: list[str], sd_ops: SDOps | None, state_dict: StateDict) -> None: ...
+
+    def pop(self, paths: list[str], sd_ops: SDOps | None) -> StateDict | None: ...
+
+    def get(self, paths: list[str], sd_ops: SDOps | None) -> StateDict | None: ...
+
+    def clear(self) -> None: ...
+
+
+class DummyRegistry(Registry):
+    """
+    Dummy registry that does not store state dictionaries.
+    """
+
+    def add(self, paths: list[str], sd_ops: SDOps | None, state_dict: StateDict) -> None:
+        pass
+
+    def pop(self, paths: list[str], sd_ops: SDOps | None) -> StateDict | None:
+        pass
+
+    def get(self, paths: list[str], sd_ops: SDOps | None) -> StateDict | None:
+        pass
+
+    def clear(self) -> None:
+        pass
+
+
+@dataclass
+class StateDictRegistry(Registry):
+    """
+    Registry that stores state dictionaries in a dictionary.
+    """
+
+    _state_dicts: dict[str, StateDict] = field(default_factory=dict)
+    _lock: threading.Lock = field(default_factory=threading.Lock)
+
+    def _generate_id(self, paths: list[str], sd_ops: SDOps) -> str:
+        m = hashlib.sha256()
+        parts = [str(Path(p).resolve()) for p in paths]
+        if sd_ops is not None:
+            parts.append(sd_ops.name)
+        m.update("\0".join(parts).encode("utf-8"))
+        return m.hexdigest()
+
+    def add(self, paths: list[str], sd_ops: SDOps | None, state_dict: StateDict) -> str:
+        sd_id = self._generate_id(paths, sd_ops)
+        with self._lock:
+            if sd_id in self._state_dicts:
+                raise ValueError(f"State dict retrieved from {paths} with {sd_ops} already added, check with get first")
+            self._state_dicts[sd_id] = state_dict
+        return sd_id
+
+    def pop(self, paths: list[str], sd_ops: SDOps | None) -> StateDict | None:
+        with self._lock:
+            return self._state_dicts.pop(self._generate_id(paths, sd_ops), None)
+
+    def get(self, paths: list[str], sd_ops: SDOps | None) -> StateDict | None:
+        with self._lock:
+            return self._state_dicts.get(self._generate_id(paths, sd_ops), None)
+
+    def clear(self) -> None:
+        with self._lock:
+            self._state_dicts.clear()

packages/ltx-core/src/ltx_core/loader/sd_ops.py

@@ -0,0 +1,139 @@
+from dataclasses import dataclass, replace
+from typing import NamedTuple, Protocol
+
+import torch
+
+
+@dataclass(frozen=True, slots=True)
+class ContentReplacement:
+    """
+    Represents a content replacement operation.
+    Used to replace a specific content with a replacement in a state dict key.
+    """
+
+    content: str
+    replacement: str
+
+
+@dataclass(frozen=True, slots=True)
+class ContentMatching:
+    """
+    Represents a content matching operation.
+    Used to match a specific prefix and suffix in a state dict key.
+    """
+
+    prefix: str = ""
+    suffix: str = ""
+
+
+class KeyValueOperationResult(NamedTuple):
+    """
+    Represents the result of a key-value operation.
+    Contains the new key and value after the operation has been applied.
+    """
+
+    new_key: str
+    new_value: torch.Tensor
+
+
+class KeyValueOperation(Protocol):
+    """
+    Protocol for key-value operations.
+    Used to apply operations to a specific key and value in a state dict.
+    """
+
+    def __call__(self, tensor_key: str, tensor_value: torch.Tensor) -> list[KeyValueOperationResult]: ...
+
+
+@dataclass(frozen=True, slots=True)
+class SDKeyValueOperation:
+    """
+    Represents a key-value operation.
+    Used to apply operations to a specific key and value in a state dict.
+    """
+
+    key_matcher: ContentMatching
+    kv_operation: KeyValueOperation
+
+
+@dataclass(frozen=True, slots=True)
+class SDOps:
+    """Immutable class representing state dict key operations."""
+
+    name: str
+    mapping: tuple[
+        ContentReplacement | ContentMatching | SDKeyValueOperation, ...
+    ] = ()  # Immutable tuple of (key, value) pairs
+    allowed_keys: frozenset[str] | None = None
+
+    def with_replacement(self, content: str, replacement: str) -> "SDOps":
+        """Create a new SDOps instance with the specified replacement added to the mapping."""
+
+        new_mapping = (*self.mapping, ContentReplacement(content, replacement))
+        return replace(self, mapping=new_mapping)
+
+    def with_matching(self, prefix: str = "", suffix: str = "") -> "SDOps":
+        """Create a new SDOps instance with the specified prefix and suffix matching added to the mapping."""
+
+        new_mapping = (*self.mapping, ContentMatching(prefix, suffix))
+        return replace(self, mapping=new_mapping)
+
+    def with_additional_allowed_keys(self, keys: frozenset[str]) -> "SDOps":
+        """Create a new SDOps instance that only passes keys present in *keys* (post-replacement).
+        If allowed_keys already exists, the sets are merged via union.
+        """
+        merged = frozenset(keys) | self.allowed_keys if self.allowed_keys is not None else frozenset(keys)
+        return replace(self, allowed_keys=merged)
+
+    def with_kv_operation(
+        self,
+        operation: KeyValueOperation,
+        key_prefix: str = "",
+        key_suffix: str = "",
+    ) -> "SDOps":
+        """Create a new SDOps instance with the specified value operation added to the mapping."""
+        key_matcher = ContentMatching(key_prefix, key_suffix)
+        sd_kv_operation = SDKeyValueOperation(key_matcher, operation)
+        new_mapping = (*self.mapping, sd_kv_operation)
+        return replace(self, mapping=new_mapping)
+
+    def apply_to_key(self, key: str) -> str | None:
+        """Apply the mapping to the given name."""
+        matchers = [content for content in self.mapping if isinstance(content, ContentMatching)]
+        valid = any(key.startswith(f.prefix) and key.endswith(f.suffix) for f in matchers)
+        if not valid:
+            return None
+
+        for replacement in self.mapping:
+            if not isinstance(replacement, ContentReplacement):
+                continue
+            if replacement.content in key:
+                key = key.replace(replacement.content, replacement.replacement)
+
+        if self.allowed_keys is not None and key not in self.allowed_keys:
+            return None
+
+        return key
+
+    def apply_to_key_value(self, key: str, value: torch.Tensor) -> list[KeyValueOperationResult]:
+        """Apply the value operation to the given name and associated value."""
+        for operation in self.mapping:
+            if not isinstance(operation, SDKeyValueOperation):
+                continue
+            if key.startswith(operation.key_matcher.prefix) and key.endswith(operation.key_matcher.suffix):
+                return operation.kv_operation(key, value)
+        return [KeyValueOperationResult(key, value)]
+
+
+# Predefined SDOps instances
+LTXV_LORA_COMFY_RENAMING_MAP = (
+    SDOps("LTXV_LORA_COMFY_PREFIX_MAP").with_matching().with_replacement("diffusion_model.", "")
+)
+
+LTXV_LORA_COMFY_TARGET_MAP = (
+    SDOps("LTXV_LORA_COMFY_TARGET_MAP")
+    .with_matching()
+    .with_replacement("diffusion_model.", "")
+    .with_replacement(".lora_A.weight", ".weight")
+    .with_replacement(".lora_B.weight", ".weight")
+)

packages/ltx-core/src/ltx_core/loader/sft_loader.py

@@ -0,0 +1,66 @@
+import json
+
+import safetensors
+import torch
+
+from ltx_core.loader.primitives import StateDict, StateDictLoader
+from ltx_core.loader.sd_ops import SDOps
+
+
+class SafetensorsStateDictLoader(StateDictLoader):
+    """
+    Loads weights from safetensors files without metadata support.
+    Use this for loading raw weight files. For model files that include
+    configuration metadata, use SafetensorsModelStateDictLoader instead.
+    """
+
+    def metadata(self, path: str) -> dict:
+        raise NotImplementedError("Not implemented")
+
+    def load(self, path: str | list[str], sd_ops: SDOps, device: torch.device | None = None) -> StateDict:
+        """
+        Load state dict from path or paths (for sharded model storage) and apply sd_ops
+        """
+        sd = {}
+        size = 0
+        dtype = set()
+        device = device or torch.device("cpu")
+        model_paths = path if isinstance(path, list) else [path]
+        for shard_path in model_paths:
+            with safetensors.safe_open(shard_path, framework="pt", device=str(device)) as f:
+                safetensor_keys = f.keys()
+                for name in safetensor_keys:
+                    expected_name = name if sd_ops is None else sd_ops.apply_to_key(name)
+                    if expected_name is None:
+                        continue
+                    value = f.get_tensor(name).to(device=device, non_blocking=True, copy=False)
+                    key_value_pairs = ((expected_name, value),)
+                    if sd_ops is not None:
+                        key_value_pairs = sd_ops.apply_to_key_value(expected_name, value)
+                    for key, value in key_value_pairs:
+                        size += value.nbytes
+                        dtype.add(value.dtype)
+                        sd[key] = value
+
+        return StateDict(sd=sd, device=device, size=size, dtype=dtype)
+
+
+class SafetensorsModelStateDictLoader(StateDictLoader):
+    """
+    Loads weights and configuration metadata from safetensors model files.
+    Unlike SafetensorsStateDictLoader, this loader can read model configuration
+    from the safetensors file metadata via the metadata() method.
+    """
+
+    def __init__(self, weight_loader: SafetensorsStateDictLoader | None = None):
+        self.weight_loader = weight_loader if weight_loader is not None else SafetensorsStateDictLoader()
+
+    def metadata(self, path: str) -> dict:
+        with safetensors.safe_open(path, framework="pt") as f:
+            meta = f.metadata()
+            if meta is None or "config" not in meta:
+                return {}
+            return json.loads(meta["config"])
+
+    def load(self, path: str | list[str], sd_ops: SDOps | None = None, device: torch.device | None = None) -> StateDict:
+        return self.weight_loader.load(path, sd_ops, device)

packages/ltx-core/src/ltx_core/loader/single_gpu_model_builder.py

@@ -0,0 +1,162 @@
+import logging
+from dataclasses import dataclass, field, replace
+from typing import Generic
+
+import torch
+from torch import nn
+
+from ltx_core.loader.fuse_loras import apply_loras
+from ltx_core.loader.helpers import create_meta_model, load_state_dict, read_model_config
+from ltx_core.loader.module_ops import ModuleOps
+from ltx_core.loader.primitives import (
+    LoRAAdaptableProtocol,
+    LoraPathStrengthAndSDOps,
+    LoraStateDictWithStrength,
+    ModelBuilderProtocol,
+    StateDict,
+    StateDictLoader,
+)
+from ltx_core.loader.registry import DummyRegistry, Registry
+from ltx_core.loader.sd_ops import SDOps
+from ltx_core.loader.sft_loader import SafetensorsModelStateDictLoader
+from ltx_core.model.model_protocol import ModelConfigurator, ModelType
+
+logger: logging.Logger = logging.getLogger(__name__)
+
+
+def _check_uninitialized(model: nn.Module) -> list[str]:
+    """Return names of any parameters/buffers still on meta device."""
+    names = []
+    for name, param in model.named_parameters():
+        if str(param.device) == "meta":
+            names.append(name)
+    for name, buf in model.named_buffers():
+        if str(buf.device) == "meta":
+            names.append(name)
+    return names
+
+
+def _load_model_weights(
+    meta_model: nn.Module,
+    model_path: str | tuple[str, ...],
+    loras: tuple[LoraPathStrengthAndSDOps, ...],
+    loader: StateDictLoader,
+    registry: Registry,
+    device: torch.device,
+    dtype: torch.dtype | None,
+    model_sd_ops: SDOps | None = None,
+    lora_load_device: torch.device | None = None,
+) -> None:
+    """Load base weights and fuse LoRAs into *meta_model* in-place."""
+    if lora_load_device is None:
+        lora_load_device = device
+
+    model_sd = load_state_dict(model_path, loader, registry, device, model_sd_ops)
+
+    lora_strengths = [lora.strength for lora in loras]
+    if not lora_strengths or (min(lora_strengths) == 0 and max(lora_strengths) == 0):
+        sd = model_sd.sd
+        if dtype is not None:
+            sd = {key: value.to(dtype=dtype) for key, value in model_sd.sd.items()}
+        meta_model.load_state_dict(sd, strict=False, assign=True)
+        return
+
+    lora_state_dicts = [load_state_dict([lora.path], loader, registry, lora_load_device, lora.sd_ops) for lora in loras]
+    lora_sd_and_strengths = [
+        LoraStateDictWithStrength(sd, strength) for sd, strength in zip(lora_state_dicts, lora_strengths, strict=True)
+    ]
+    final_sd = apply_loras(
+        model_sd=model_sd,
+        lora_sd_and_strengths=lora_sd_and_strengths,
+        dtype=dtype,
+        destination_sd=model_sd if isinstance(registry, DummyRegistry) else None,
+    )
+    meta_model.load_state_dict(final_sd.sd, strict=False, assign=True)
+
+
+@dataclass(frozen=True)
+class SingleGPUModelBuilder(Generic[ModelType], ModelBuilderProtocol[ModelType], LoRAAdaptableProtocol):
+    """
+    Builder for PyTorch models residing on a single GPU.
+    Attributes:
+        model_class_configurator: Class responsible for constructing the model from a config dict.
+        model_path: Path (or tuple of shard paths) to the model's `.safetensors` checkpoint(s).
+        model_sd_ops: Optional state-dict operations applied when loading the model weights.
+        module_ops: Sequence of module-level mutations applied to the meta model before weight loading.
+        loras: Sequence of LoRA adapters (path, strength, optional sd_ops) to fuse into the model.
+        model_loader: Strategy for loading state dicts from disk. Defaults to
+            :class:`SafetensorsModelStateDictLoader`.
+        registry: Cache for already-loaded state dicts. Defaults to :class:`DummyRegistry` (no caching).
+        lora_load_device: Device used when loading LoRA weight tensors from disk. Defaults to
+            ``torch.device("cpu")``, which keeps LoRA weights in CPU memory and transfers them to
+            the target GPU sequentially during fusion, reducing peak GPU memory usage compared to
+            loading all LoRA weights directly onto the GPU at once.
+    """
+
+    model_class_configurator: type[ModelConfigurator[ModelType]]
+    model_path: str | tuple[str, ...]
+    model_sd_ops: SDOps | None = None
+    module_ops: tuple[ModuleOps, ...] = field(default_factory=tuple)
+    loras: tuple[LoraPathStrengthAndSDOps, ...] = field(default_factory=tuple)
+    model_loader: StateDictLoader = field(default_factory=SafetensorsModelStateDictLoader)
+    registry: Registry = field(default_factory=DummyRegistry)
+    lora_load_device: torch.device = field(default_factory=lambda: torch.device("cpu"))
+
+    def lora(self, lora_path: str, strength: float, sd_ops: SDOps) -> "SingleGPUModelBuilder":
+        return replace(self, loras=(*self.loras, LoraPathStrengthAndSDOps(lora_path, strength, sd_ops)))
+
+    def with_sd_ops(self, sd_ops: SDOps | None) -> "SingleGPUModelBuilder":
+        return replace(self, model_sd_ops=sd_ops)
+
+    def with_module_ops(self, module_ops: tuple[ModuleOps, ...]) -> "SingleGPUModelBuilder":
+        return replace(self, module_ops=module_ops)
+
+    def with_loras(self, loras: tuple[LoraPathStrengthAndSDOps, ...]) -> "SingleGPUModelBuilder":
+        return replace(self, loras=loras)
+
+    def with_registry(self, registry: Registry) -> "SingleGPUModelBuilder":
+        return replace(self, registry=registry)
+
+    def with_lora_load_device(self, device: torch.device) -> "SingleGPUModelBuilder":
+        return replace(self, lora_load_device=device)
+
+    def model_config(self) -> dict:
+        return read_model_config(self.model_path, self.model_loader)
+
+    def meta_model(self, config: dict, module_ops: tuple[ModuleOps, ...]) -> ModelType:
+        return create_meta_model(self.model_class_configurator, config, module_ops)
+
+    def load_sd(
+        self, paths: list[str], registry: Registry, device: torch.device | None, sd_ops: SDOps | None = None
+    ) -> StateDict:
+        return load_state_dict(paths, self.model_loader, registry, device, sd_ops)
+
+    def _return_model(self, meta_model: ModelType, device: torch.device) -> ModelType:
+        uninitialized = _check_uninitialized(meta_model)
+        if uninitialized:
+            logger.warning(f"Uninitialized parameters or buffers: {uninitialized}")
+            return meta_model
+        return meta_model.to(device)
+
+    def build(
+        self,
+        device: torch.device | None = None,
+        dtype: torch.dtype | None = None,
+        **kwargs: object,  # noqa: ARG002
+    ) -> ModelType:
+        device = torch.device("cuda") if device is None else device
+        config = self.model_config()
+        meta_model = self.meta_model(config, self.module_ops)
+
+        _load_model_weights(
+            meta_model=meta_model,
+            model_path=self.model_path,
+            loras=self.loras,
+            loader=self.model_loader,
+            registry=self.registry,
+            device=device,
+            dtype=dtype,
+            model_sd_ops=self.model_sd_ops,
+            lora_load_device=self.lora_load_device,
+        )
+        return self._return_model(meta_model, device)

packages/ltx-core/src/ltx_core/modality_tiling.py

@@ -0,0 +1,234 @@
+"""Video modality tiling helpers.
+Provides :class:`VideoModalityTilingHelper` — a stateless helper that
+tiles and blends video :class:`Modality` token sequences by
+spatial/temporal region.  Tile geometry is represented by the existing
+:class:`Tile` NamedTuple from :mod:`ltx_core.tiling`; no distributed
+primitives are required.
+"""
+
+from __future__ import annotations
+
+from dataclasses import dataclass, replace
+
+import torch
+
+from ltx_core.model.transformer.modality import Modality
+from ltx_core.tiling import Tile, TileCountConfig, create_tiles, identity_mapping_operation, split_by_count
+from ltx_core.tools import VideoLatentTools
+from ltx_core.types import VideoLatentShape
+
+
+@dataclass(frozen=True)
+class TilingContext:
+    """Opaque context produced by :meth:`VideoModalityTilingHelper.tile_modality`.
+    Carries the token-level keep mask and per-conditioning-token blend
+    weights needed by :meth:`~VideoModalityTilingHelper.blend`.
+    """
+
+    keep_mask: torch.Tensor
+    cond_blend_weights: torch.Tensor | None
+    """``(num_kept_cond,)`` — weight for each kept conditioning token,
+    equal to ``1 / num_tiles_that_keep_this_token``.  ``None`` when
+    there are no conditioning tokens."""
+
+
+class VideoModalityTilingHelper:
+    """Stateless helper that tiles and blends video :class:`Modality` sequences.
+    Constructed once with a :class:`TileCountConfig` and
+    :class:`VideoLatentTools`.  Tiles are computed at construction and
+    available via the :attr:`tiles` property.  Use :meth:`tile_modality`
+    and :meth:`blend` with any tile from that list.
+    Usage::
+        helper = VideoModalityTilingHelper(tiling, video_tools)
+        for tile in helper.tiles:
+            tiled_mod, ctx = helper.tile_modality(modality, tile)
+            result = run_model(tiled_mod)
+            helper.blend(result, tile, ctx, output=output)
+    """
+
+    def __init__(self, tiling: TileCountConfig, video_tools: VideoLatentTools) -> None:
+        self._patchifier = video_tools.patchifier
+        self._latent_shape = video_tools.target_shape
+        self._num_generated_tokens = self._patchifier.get_token_count(self._latent_shape)
+        self._tiles = create_tiles(
+            torch.Size([self._latent_shape.frames, self._latent_shape.height, self._latent_shape.width]),
+            splitters=[
+                split_by_count(tiling.frames.num_tiles, tiling.frames.overlap),
+                split_by_count(tiling.height.num_tiles, tiling.height.overlap),
+                split_by_count(tiling.width.num_tiles, tiling.width.overlap),
+            ],
+            mappers=[identity_mapping_operation] * 3,
+        )
+
+    @property
+    def tiles(self) -> list[Tile]:
+        """All tiles for the configured tiling layout."""
+        return self._tiles
+
+    # -- tile modality -----------------------------------------------------
+
+    def tile_modality(
+        self, modality: Modality, tile: Tile, *, normalize_positions: bool = True
+    ) -> tuple[Modality, TilingContext]:
+        """Slice *modality* to the tokens covered by *tile*.
+        Selects generated tokens belonging to the tile's spatial region
+        and conditioning tokens that overlap with the tile (or have
+        negative time coordinates).
+        Args:
+            normalize_positions: When True, shift all positions so the
+                tile's generated tokens start at zero in every dimension.
+        Returns:
+            A ``(tiled_modality, context)`` tuple.  Pass *context* to
+            :meth:`blend` together with the model output.
+        """
+        keep_mask = self._keep_mask(modality, tile)
+
+        tile_attention_mask = None
+        if modality.attention_mask is not None:
+            keep_indices = keep_mask.nonzero(as_tuple=False).squeeze(1)
+            tile_attention_mask = modality.attention_mask[:, keep_indices, :][:, :, keep_indices]
+
+        positions = modality.positions[:, :, keep_mask, :]
+        if normalize_positions:
+            num_tile_gen = self._tile_generated_token_count(tile)
+            gen_pos = positions[:, :, :num_tile_gen, :]  # (B, 3, num_tile_gen, 2)
+            offset = gen_pos[..., 0].amin(dim=2, keepdim=True).unsqueeze(-1)  # (B, 3, 1, 1)
+            positions = positions - offset
+
+        tiled = replace(
+            modality,
+            latent=modality.latent[:, keep_mask, :],
+            timesteps=modality.timesteps[:, keep_mask],
+            positions=positions,
+            attention_mask=tile_attention_mask,
+        )
+
+        cond_blend_weights = None
+        num_total = modality.latent.shape[1]
+        if num_total > self._num_generated_tokens:
+            cond_keep = keep_mask[self._num_generated_tokens :]
+            # Count how many tiles keep each conditioning token.
+            cond_counts = torch.zeros(cond_keep.sum(), dtype=torch.float32)
+            for t in self._tiles:
+                other_mask = self._keep_mask(modality, t)
+                other_cond = other_mask[self._num_generated_tokens :]
+                # Map other tile's kept cond tokens into this tile's kept subset.
+                cond_counts += other_cond[cond_keep].float()
+            cond_blend_weights = 1.0 / cond_counts
+
+        return tiled, TilingContext(keep_mask=keep_mask, cond_blend_weights=cond_blend_weights)
+
+    # -- blend -------------------------------------------------------------
+
+    def blend(
+        self,
+        tile_to_blend: torch.Tensor,
+        tile: Tile,
+        context: TilingContext,
+        output: torch.Tensor | None = None,
+    ) -> torch.Tensor:
+        """Blend-weight tile results and accumulate into the full token space.
+        Premultiplied (blend-weighted) data is **added** to *output*,
+        allowing multiple tiles to be accumulated into the same buffer.
+        Args:
+            tile_to_blend: Denoised tile tensor ``(B, num_tile_tokens, D)``,
+                where the first ``_tile_generated_token_count(tile)``
+                entries are generated tokens and the remainder are
+                conditioning tokens.
+            tile: The :class:`Tile` that was used in :meth:`tile_modality`.
+            context: The :class:`TilingContext` returned by :meth:`tile_modality`.
+            output: Optional pre-allocated output tensor.  When provided
+                its shape must be ``(B, num_total_tokens, D)`` and the
+                blended tile is **added** into it.  When ``None`` a new
+                zero-filled tensor is created.
+        Returns:
+            The output tensor with the blended tile added at the correct
+            positions.
+        """
+        batch, _, dim = tile_to_blend.shape
+        num_tile_gen = self._tile_generated_token_count(tile)
+        gen_indices = self._generated_token_indices(tile)
+
+        num_total_tokens = context.keep_mask.shape[0]
+        expected_shape = (batch, num_total_tokens, dim)
+
+        if output is not None:
+            if output.shape != expected_shape:
+                raise ValueError(f"Expected output shape {expected_shape}, got {output.shape}")
+            result = output
+        else:
+            result = torch.zeros(*expected_shape, device=tile_to_blend.device, dtype=tile_to_blend.dtype)
+
+        # Blend mask is (tile_F, tile_H, tile_W) — one weight per token in row-major order.
+        blend_weights = tile.blend_mask.reshape(-1).to(device=tile_to_blend.device, dtype=tile_to_blend.dtype)
+        tile_gen = tile_to_blend[:, :num_tile_gen, :] * blend_weights[None, :, None]
+
+        result[:, gen_indices, :] += tile_gen
+
+        # Scatter kept conditioning tokens, weighted by 1/N where N is
+        # the number of tiles that keep each token (so they sum to 1).
+        if num_total_tokens > self._num_generated_tokens and context.cond_blend_weights is not None:
+            cond_keep = context.keep_mask[self._num_generated_tokens :]
+            cond_indices = self._num_generated_tokens + cond_keep.nonzero(as_tuple=False).squeeze(1)
+            weights = context.cond_blend_weights.to(device=tile_to_blend.device, dtype=tile_to_blend.dtype)
+            result[:, cond_indices, :] += tile_to_blend[:, num_tile_gen:, :] * weights[None, :, None]
+
+        return result
+
+    # -- private -----------------------------------------------------------
+
+    def _tile_generated_token_count(self, tile: Tile) -> int:
+        """Number of generated tokens in *tile*."""
+        frame_slice, height_slice, width_slice = tile.in_coords
+        tile_shape = VideoLatentShape(
+            batch=self._latent_shape.batch,
+            channels=self._latent_shape.channels,
+            frames=frame_slice.stop - frame_slice.start,
+            height=height_slice.stop - height_slice.start,
+            width=width_slice.stop - width_slice.start,
+        )
+        return self._patchifier.get_token_count(tile_shape)
+
+    def _generated_token_indices(self, tile: Tile) -> torch.Tensor:
+        """Flat token indices of *tile*'s generated tokens in the full sequence."""
+        frame_slice, height_slice, width_slice = tile.in_coords
+        f = torch.arange(frame_slice.start, frame_slice.stop)
+        h = torch.arange(height_slice.start, height_slice.stop)
+        w = torch.arange(width_slice.start, width_slice.stop)
+        return (
+            f[:, None, None] * self._latent_shape.height * self._latent_shape.width
+            + h[None, :, None] * self._latent_shape.width
+            + w[None, None, :]
+        ).reshape(-1)
+
+    def _keep_mask(self, modality: Modality, tile: Tile) -> torch.Tensor:
+        """Boolean mask ``(num_total_tokens,)`` — True for tokens the tile processes.
+        Generated tokens are selected by grid position.  Conditioning
+        tokens are kept when their ``[start, end)`` intervals overlap
+        the tile in all three dimensions, or when they have a negative
+        time coordinate (reference tokens).
+        """
+        num_total = modality.latent.shape[1]
+        mask = torch.zeros(num_total, dtype=torch.bool)
+
+        gen_indices = self._generated_token_indices(tile)
+        mask[gen_indices] = True
+
+        if num_total > self._num_generated_tokens:
+            gen_positions = modality.positions[:, :, gen_indices, :]  # (B, 3, num_tile_gen, 2)
+            tile_start = gen_positions[..., 0].amin(dim=2)  # (B, 3)
+            tile_end = gen_positions[..., 1].amax(dim=2)  # (B, 3)
+
+            cond_positions = modality.positions[:, :, self._num_generated_tokens :, :]  # (B, 3, num_cond, 2)
+
+            overlaps = (cond_positions[..., 0] < tile_end.unsqueeze(2)) & (
+                cond_positions[..., 1] > tile_start.unsqueeze(2)
+            )  # (B, 3, num_cond)
+            overlaps_all_dims = overlaps.all(dim=1)  # (B, num_cond)
+
+            has_negative_time = cond_positions[:, 0, :, 0] < 0  # (B, num_cond)
+
+            keep_cond = (overlaps_all_dims | has_negative_time).any(dim=0)  # (num_cond,)
+            mask[self._num_generated_tokens :] = keep_cond
+
+        return mask

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

@@ -0,0 +1,8 @@
+"""Model definitions for LTX-2."""
+
+from ltx_core.model.model_protocol import ModelConfigurator, ModelType
+
+__all__ = [
+    "ModelConfigurator",
+    "ModelType",
+]

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

@@ -0,0 +1,29 @@
+"""Audio VAE model components."""
+
+from ltx_core.model.audio_vae.audio_vae import AudioDecoder, AudioEncoder, decode_audio, encode_audio
+from ltx_core.model.audio_vae.model_configurator import (
+    AUDIO_VAE_DECODER_COMFY_KEYS_FILTER,
+    AUDIO_VAE_ENCODER_COMFY_KEYS_FILTER,
+    VOCODER_COMFY_KEYS_FILTER,
+    AudioDecoderConfigurator,
+    AudioEncoderConfigurator,
+    VocoderConfigurator,
+)
+from ltx_core.model.audio_vae.ops import AudioProcessor
+from ltx_core.model.audio_vae.vocoder import Vocoder, VocoderWithBWE
+
+__all__ = [
+    "AUDIO_VAE_DECODER_COMFY_KEYS_FILTER",
+    "AUDIO_VAE_ENCODER_COMFY_KEYS_FILTER",
+    "VOCODER_COMFY_KEYS_FILTER",
+    "AudioDecoder",
+    "AudioDecoderConfigurator",
+    "AudioEncoder",
+    "AudioEncoderConfigurator",
+    "AudioProcessor",
+    "Vocoder",
+    "VocoderConfigurator",
+    "VocoderWithBWE",
+    "decode_audio",
+    "encode_audio",
+]

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

@@ -0,0 +1,71 @@
+from enum import Enum
+
+import torch
+
+from ltx_core.model.common.normalization import NormType, build_normalization_layer
+
+
+class AttentionType(Enum):
+    """Enum for specifying the attention mechanism type."""
+
+    VANILLA = "vanilla"
+    LINEAR = "linear"
+    NONE = "none"
+
+
+class AttnBlock(torch.nn.Module):
+    def __init__(
+        self,
+        in_channels: int,
+        norm_type: NormType = NormType.GROUP,
+    ) -> None:
+        super().__init__()
+        self.in_channels = in_channels
+
+        self.norm = build_normalization_layer(in_channels, normtype=norm_type)
+        self.q = torch.nn.Conv2d(in_channels, in_channels, kernel_size=1, stride=1, padding=0)
+        self.k = torch.nn.Conv2d(in_channels, in_channels, kernel_size=1, stride=1, padding=0)
+        self.v = torch.nn.Conv2d(in_channels, in_channels, kernel_size=1, stride=1, padding=0)
+        self.proj_out = torch.nn.Conv2d(in_channels, in_channels, kernel_size=1, stride=1, padding=0)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        h_ = x
+        h_ = self.norm(h_)
+        q = self.q(h_)
+        k = self.k(h_)
+        v = self.v(h_)
+
+        # compute attention
+        b, c, h, w = q.shape
+        q = q.reshape(b, c, h * w).contiguous()
+        q = q.permute(0, 2, 1).contiguous()  # b,hw,c
+        k = k.reshape(b, c, h * w).contiguous()  # b,c,hw
+        w_ = torch.bmm(q, k).contiguous()  # b,hw,hw    w[b,i,j]=sum_c q[b,i,c]k[b,c,j]
+        w_ = w_ * (int(c) ** (-0.5))
+        w_ = torch.nn.functional.softmax(w_, dim=2)
+
+        # attend to values
+        v = v.reshape(b, c, h * w).contiguous()
+        w_ = w_.permute(0, 2, 1).contiguous()  # b,hw,hw (first hw of k, second of q)
+        h_ = torch.bmm(v, w_).contiguous()  # b, c,hw (hw of q) h_[b,c,j] = sum_i v[b,c,i] w_[b,i,j]
+        h_ = h_.reshape(b, c, h, w).contiguous()
+
+        h_ = self.proj_out(h_)
+
+        return x + h_
+
+
+def make_attn(
+    in_channels: int,
+    attn_type: AttentionType = AttentionType.VANILLA,
+    norm_type: NormType = NormType.GROUP,
+) -> torch.nn.Module:
+    match attn_type:
+        case AttentionType.VANILLA:
+            return AttnBlock(in_channels, norm_type=norm_type)
+        case AttentionType.NONE:
+            return torch.nn.Identity()
+        case AttentionType.LINEAR:
+            raise NotImplementedError(f"Attention type {attn_type.value} is not supported yet.")
+        case _:
+            raise ValueError(f"Unknown attention type: {attn_type}")

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

@@ -0,0 +1,508 @@
+from typing import Set, Tuple
+
+import torch
+import torch.nn.functional as F
+
+from ltx_core.components.patchifiers import AudioPatchifier
+from ltx_core.model.audio_vae.attention import AttentionType, make_attn
+from ltx_core.model.audio_vae.causal_conv_2d import make_conv2d
+from ltx_core.model.audio_vae.causality_axis import CausalityAxis
+from ltx_core.model.audio_vae.downsample import build_downsampling_path
+from ltx_core.model.audio_vae.ops import AudioProcessor, PerChannelStatistics
+from ltx_core.model.audio_vae.resnet import ResnetBlock
+from ltx_core.model.audio_vae.upsample import build_upsampling_path
+from ltx_core.model.audio_vae.vocoder import Vocoder
+from ltx_core.model.common.normalization import NormType, build_normalization_layer
+from ltx_core.types import Audio, AudioLatentShape
+
+LATENT_DOWNSAMPLE_FACTOR = 4
+
+
+def build_mid_block(
+    channels: int,
+    temb_channels: int,
+    dropout: float,
+    norm_type: NormType,
+    causality_axis: CausalityAxis,
+    attn_type: AttentionType,
+    add_attention: bool,
+) -> torch.nn.Module:
+    """Build the middle block with two ResNet blocks and optional attention."""
+    mid = torch.nn.Module()
+    mid.block_1 = ResnetBlock(
+        in_channels=channels,
+        out_channels=channels,
+        temb_channels=temb_channels,
+        dropout=dropout,
+        norm_type=norm_type,
+        causality_axis=causality_axis,
+    )
+    mid.attn_1 = make_attn(channels, attn_type=attn_type, norm_type=norm_type) if add_attention else torch.nn.Identity()
+    mid.block_2 = ResnetBlock(
+        in_channels=channels,
+        out_channels=channels,
+        temb_channels=temb_channels,
+        dropout=dropout,
+        norm_type=norm_type,
+        causality_axis=causality_axis,
+    )
+    return mid
+
+
+def run_mid_block(mid: torch.nn.Module, features: torch.Tensor) -> torch.Tensor:
+    """Run features through the middle block."""
+    features = mid.block_1(features, temb=None)
+    features = mid.attn_1(features)
+    return mid.block_2(features, temb=None)
+
+
+class AudioEncoder(torch.nn.Module):
+    """
+    Encoder that compresses audio spectrograms into latent representations.
+    The encoder uses a series of downsampling blocks with residual connections,
+    attention mechanisms, and configurable causal convolutions.
+    """
+
+    def __init__(  # noqa: PLR0913
+        self,
+        *,
+        ch: int,
+        ch_mult: Tuple[int, ...] = (1, 2, 4, 8),
+        num_res_blocks: int,
+        attn_resolutions: Set[int],
+        dropout: float = 0.0,
+        resamp_with_conv: bool = True,
+        in_channels: int,
+        resolution: int,
+        z_channels: int,
+        double_z: bool = True,
+        attn_type: AttentionType = AttentionType.VANILLA,
+        mid_block_add_attention: bool = True,
+        norm_type: NormType = NormType.GROUP,
+        causality_axis: CausalityAxis = CausalityAxis.WIDTH,
+        sample_rate: int = 16000,
+        mel_hop_length: int = 160,
+        n_fft: int = 1024,
+        is_causal: bool = True,
+        mel_bins: int = 64,
+        **_ignore_kwargs,
+    ) -> None:
+        """
+        Initialize the Encoder.
+        Args:
+            Arguments are configuration parameters, loaded from the audio VAE checkpoint config
+            (audio_vae.model.params.ddconfig):
+            ch: Base number of feature channels used in the first convolution layer.
+            ch_mult: Multiplicative factors for the number of channels at each resolution level.
+            num_res_blocks: Number of residual blocks to use at each resolution level.
+            attn_resolutions: Spatial resolutions (e.g., in time/frequency) at which to apply attention.
+            resolution: Input spatial resolution of the spectrogram (height, width).
+            z_channels: Number of channels in the latent representation.
+            norm_type: Normalization layer type to use within the network (e.g., group, batch).
+            causality_axis: Axis along which convolutions should be causal (e.g., time axis).
+            sample_rate: Audio sample rate in Hz for the input signals.
+            mel_hop_length: Hop length used when computing the mel spectrogram.
+            n_fft: FFT size used to compute the spectrogram.
+            mel_bins: Number of mel-frequency bins in the input spectrogram.
+            in_channels: Number of channels in the input spectrogram tensor.
+            double_z: If True, predict both mean and log-variance (doubling latent channels).
+            is_causal: If True, use causal convolutions suitable for streaming setups.
+            dropout: Dropout probability used in residual and mid blocks.
+            attn_type: Type of attention mechanism to use in attention blocks.
+            resamp_with_conv: If True, perform resolution changes using strided convolutions.
+            mid_block_add_attention: If True, add an attention block in the mid-level of the encoder.
+        """
+        super().__init__()
+
+        self.per_channel_statistics = PerChannelStatistics(latent_channels=ch)
+        self.sample_rate = sample_rate
+        self.mel_hop_length = mel_hop_length
+        self.n_fft = n_fft
+        self.is_causal = is_causal
+        self.mel_bins = mel_bins
+
+        self.patchifier = AudioPatchifier(
+            patch_size=1,
+            audio_latent_downsample_factor=LATENT_DOWNSAMPLE_FACTOR,
+            sample_rate=sample_rate,
+            hop_length=mel_hop_length,
+            is_causal=is_causal,
+        )
+
+        self.ch = ch
+        self.temb_ch = 0
+        self.num_resolutions = len(ch_mult)
+        self.num_res_blocks = num_res_blocks
+        self.resolution = resolution
+        self.in_channels = in_channels
+        self.z_channels = z_channels
+        self.double_z = double_z
+        self.norm_type = norm_type
+        self.causality_axis = causality_axis
+        self.attn_type = attn_type
+
+        # downsampling
+        self.conv_in = make_conv2d(
+            in_channels,
+            self.ch,
+            kernel_size=3,
+            stride=1,
+            causality_axis=self.causality_axis,
+        )
+
+        self.non_linearity = torch.nn.SiLU()
+
+        self.down, block_in = build_downsampling_path(
+            ch=ch,
+            ch_mult=ch_mult,
+            num_resolutions=self.num_resolutions,
+            num_res_blocks=num_res_blocks,
+            resolution=resolution,
+            temb_channels=self.temb_ch,
+            dropout=dropout,
+            norm_type=self.norm_type,
+            causality_axis=self.causality_axis,
+            attn_type=self.attn_type,
+            attn_resolutions=attn_resolutions,
+            resamp_with_conv=resamp_with_conv,
+        )
+
+        self.mid = build_mid_block(
+            channels=block_in,
+            temb_channels=self.temb_ch,
+            dropout=dropout,
+            norm_type=self.norm_type,
+            causality_axis=self.causality_axis,
+            attn_type=self.attn_type,
+            add_attention=mid_block_add_attention,
+        )
+
+        self.norm_out = build_normalization_layer(block_in, normtype=self.norm_type)
+        self.conv_out = make_conv2d(
+            block_in,
+            2 * z_channels if double_z else z_channels,
+            kernel_size=3,
+            stride=1,
+            causality_axis=self.causality_axis,
+        )
+
+    def forward(self, spectrogram: torch.Tensor) -> torch.Tensor:
+        """
+        Encode audio spectrogram into latent representations.
+        Args:
+            spectrogram: Input spectrogram of shape (batch, channels, time, frequency)
+        Returns:
+            Encoded latent representation of shape (batch, channels, frames, mel_bins)
+        """
+        h = self.conv_in(spectrogram)
+        h = self._run_downsampling_path(h)
+        h = run_mid_block(self.mid, h)
+        h = self._finalize_output(h)
+
+        return self._normalize_latents(h)
+
+    def _run_downsampling_path(self, h: torch.Tensor) -> torch.Tensor:
+        for level in range(self.num_resolutions):
+            stage = self.down[level]
+            for block_idx in range(self.num_res_blocks):
+                h = stage.block[block_idx](h, temb=None)
+                if stage.attn:
+                    h = stage.attn[block_idx](h)
+
+            if level != self.num_resolutions - 1:
+                h = stage.downsample(h)
+
+        return h
+
+    def _finalize_output(self, h: torch.Tensor) -> torch.Tensor:
+        h = self.norm_out(h)
+        h = self.non_linearity(h)
+        return self.conv_out(h)
+
+    def _normalize_latents(self, latent_output: torch.Tensor) -> torch.Tensor:
+        """
+        Normalize encoder latents using per-channel statistics.
+        When the encoder is configured with ``double_z=True``, the final
+        convolution produces twice the number of latent channels, typically
+        interpreted as two concatenated tensors along the channel dimension
+        (e.g., mean and variance or other auxiliary parameters).
+        This method intentionally uses only the first half of the channels
+        (the "mean" component) as input to the patchifier and normalization
+        logic. The remaining channels are left unchanged by this method and
+        are expected to be consumed elsewhere in the VAE pipeline.
+        If ``double_z=False``, the encoder output already contains only the
+        mean latents and the chunking operation simply returns that tensor.
+        """
+        means = torch.chunk(latent_output, 2, dim=1)[0]
+        latent_shape = AudioLatentShape(
+            batch=means.shape[0],
+            channels=means.shape[1],
+            frames=means.shape[2],
+            mel_bins=means.shape[3],
+        )
+        latent_patched = self.patchifier.patchify(means)
+        latent_normalized = self.per_channel_statistics.normalize(latent_patched)
+        return self.patchifier.unpatchify(latent_normalized, latent_shape)
+
+
+def encode_audio(
+    audio: Audio,
+    audio_encoder: AudioEncoder,
+    audio_processor: AudioProcessor | None = None,
+) -> torch.Tensor:
+    """Encode audio waveform into latent representation.
+    Args:
+        audio: Audio container with waveform tensor of shape (batch, channels, samples) and sampling rate.
+        audio_encoder: Audio encoder model
+        audio_processor: Audio processor model (optional, if not provided, it will be created from the audio encoder)
+    """
+    dtype = next(audio_encoder.parameters()).dtype
+    device = next(audio_encoder.parameters()).device
+
+    if audio_processor is None:
+        audio_processor = AudioProcessor(
+            target_sample_rate=audio_encoder.sample_rate,
+            mel_bins=audio_encoder.mel_bins,
+            mel_hop_length=audio_encoder.mel_hop_length,
+            n_fft=audio_encoder.n_fft,
+        ).to(device=device)
+
+    mel_spectrogram = audio_processor.waveform_to_mel(audio.to(device=device))
+
+    latent = audio_encoder(mel_spectrogram.to(dtype=dtype))
+    return latent
+
+
+class AudioDecoder(torch.nn.Module):
+    """
+    Symmetric decoder that reconstructs audio spectrograms from latent features.
+    The decoder mirrors the encoder structure with configurable channel multipliers,
+    attention resolutions, and causal convolutions.
+    """
+
+    def __init__(  # noqa: PLR0913
+        self,
+        *,
+        ch: int,
+        out_ch: int,
+        ch_mult: Tuple[int, ...] = (1, 2, 4, 8),
+        num_res_blocks: int,
+        attn_resolutions: Set[int],
+        resolution: int,
+        z_channels: int,
+        norm_type: NormType = NormType.GROUP,
+        causality_axis: CausalityAxis = CausalityAxis.WIDTH,
+        dropout: float = 0.0,
+        mid_block_add_attention: bool = True,
+        sample_rate: int = 16000,
+        mel_hop_length: int = 160,
+        is_causal: bool = True,
+        mel_bins: int | None = None,
+    ) -> None:
+        """
+        Initialize the Decoder.
+        Args:
+            Arguments are configuration parameters, loaded from the audio VAE checkpoint config
+            (audio_vae.model.params.ddconfig):
+            - ch, out_ch, ch_mult, num_res_blocks, attn_resolutions
+            - resolution, z_channels
+            - norm_type, causality_axis
+        """
+        super().__init__()
+
+        # Internal behavioural defaults that are not driven by the checkpoint.
+        resamp_with_conv = True
+        attn_type = AttentionType.VANILLA
+
+        # Per-channel statistics for denormalizing latents
+        self.per_channel_statistics = PerChannelStatistics(latent_channels=ch)
+        self.sample_rate = sample_rate
+        self.mel_hop_length = mel_hop_length
+        self.is_causal = is_causal
+        self.mel_bins = mel_bins
+        self.patchifier = AudioPatchifier(
+            patch_size=1,
+            audio_latent_downsample_factor=LATENT_DOWNSAMPLE_FACTOR,
+            sample_rate=sample_rate,
+            hop_length=mel_hop_length,
+            is_causal=is_causal,
+        )
+
+        self.ch = ch
+        self.temb_ch = 0
+        self.num_resolutions = len(ch_mult)
+        self.num_res_blocks = num_res_blocks
+        self.resolution = resolution
+        self.out_ch = out_ch
+        self.give_pre_end = False
+        self.tanh_out = False
+        self.norm_type = norm_type
+        self.z_channels = z_channels
+        self.channel_multipliers = ch_mult
+        self.attn_resolutions = attn_resolutions
+        self.causality_axis = causality_axis
+        self.attn_type = attn_type
+
+        base_block_channels = ch * self.channel_multipliers[-1]
+        base_resolution = resolution // (2 ** (self.num_resolutions - 1))
+        self.z_shape = (1, z_channels, base_resolution, base_resolution)
+
+        self.conv_in = make_conv2d(
+            z_channels, base_block_channels, kernel_size=3, stride=1, causality_axis=self.causality_axis
+        )
+        self.non_linearity = torch.nn.SiLU()
+        self.mid = build_mid_block(
+            channels=base_block_channels,
+            temb_channels=self.temb_ch,
+            dropout=dropout,
+            norm_type=self.norm_type,
+            causality_axis=self.causality_axis,
+            attn_type=self.attn_type,
+            add_attention=mid_block_add_attention,
+        )
+        self.up, final_block_channels = build_upsampling_path(
+            ch=ch,
+            ch_mult=ch_mult,
+            num_resolutions=self.num_resolutions,
+            num_res_blocks=num_res_blocks,
+            resolution=resolution,
+            temb_channels=self.temb_ch,
+            dropout=dropout,
+            norm_type=self.norm_type,
+            causality_axis=self.causality_axis,
+            attn_type=self.attn_type,
+            attn_resolutions=attn_resolutions,
+            resamp_with_conv=resamp_with_conv,
+            initial_block_channels=base_block_channels,
+        )
+
+        self.norm_out = build_normalization_layer(final_block_channels, normtype=self.norm_type)
+        self.conv_out = make_conv2d(
+            final_block_channels, out_ch, kernel_size=3, stride=1, causality_axis=self.causality_axis
+        )
+
+    def forward(self, sample: torch.Tensor) -> torch.Tensor:
+        """
+        Decode latent features back to audio spectrograms.
+        Args:
+            sample: Encoded latent representation of shape (batch, channels, frames, mel_bins)
+        Returns:
+            Reconstructed audio spectrogram of shape (batch, channels, time, frequency)
+        """
+        sample, target_shape = self._denormalize_latents(sample)
+
+        h = self.conv_in(sample)
+        h = run_mid_block(self.mid, h)
+        h = self._run_upsampling_path(h)
+        h = self._finalize_output(h)
+
+        return self._adjust_output_shape(h, target_shape)
+
+    def _denormalize_latents(self, sample: torch.Tensor) -> tuple[torch.Tensor, AudioLatentShape]:
+        latent_shape = AudioLatentShape(
+            batch=sample.shape[0],
+            channels=sample.shape[1],
+            frames=sample.shape[2],
+            mel_bins=sample.shape[3],
+        )
+
+        sample_patched = self.patchifier.patchify(sample)
+        sample_denormalized = self.per_channel_statistics.un_normalize(sample_patched)
+        sample = self.patchifier.unpatchify(sample_denormalized, latent_shape)
+
+        target_frames = latent_shape.frames * LATENT_DOWNSAMPLE_FACTOR
+        if self.causality_axis != CausalityAxis.NONE:
+            target_frames = max(target_frames - (LATENT_DOWNSAMPLE_FACTOR - 1), 1)
+
+        target_shape = AudioLatentShape(
+            batch=latent_shape.batch,
+            channels=self.out_ch,
+            frames=target_frames,
+            mel_bins=self.mel_bins if self.mel_bins is not None else latent_shape.mel_bins,
+        )
+
+        return sample, target_shape
+
+    def _adjust_output_shape(
+        self,
+        decoded_output: torch.Tensor,
+        target_shape: AudioLatentShape,
+    ) -> torch.Tensor:
+        """
+        Adjust output shape to match target dimensions for variable-length audio.
+        This function handles the common case where decoded audio spectrograms need to be
+        resized to match a specific target shape.
+        Args:
+            decoded_output: Tensor of shape (batch, channels, time, frequency)
+            target_shape: AudioLatentShape describing (batch, channels, time, mel bins)
+        Returns:
+            Tensor adjusted to match target_shape exactly
+        """
+        # Current output shape: (batch, channels, time, frequency)
+        _, _, current_time, current_freq = decoded_output.shape
+        target_channels = target_shape.channels
+        target_time = target_shape.frames
+        target_freq = target_shape.mel_bins
+
+        # Step 1: Crop first to avoid exceeding target dimensions
+        decoded_output = decoded_output[
+            :, :target_channels, : min(current_time, target_time), : min(current_freq, target_freq)
+        ]
+
+        # Step 2: Calculate padding needed for time and frequency dimensions
+        time_padding_needed = target_time - decoded_output.shape[2]
+        freq_padding_needed = target_freq - decoded_output.shape[3]
+
+        # Step 3: Apply padding if needed
+        if time_padding_needed > 0 or freq_padding_needed > 0:
+            # PyTorch padding format: (pad_left, pad_right, pad_top, pad_bottom)
+            # For audio: pad_left/right = frequency, pad_top/bottom = time
+            padding = (
+                0,
+                max(freq_padding_needed, 0),  # frequency padding (left, right)
+                0,
+                max(time_padding_needed, 0),  # time padding (top, bottom)
+            )
+            decoded_output = F.pad(decoded_output, padding)
+
+        # Step 4: Final safety crop to ensure exact target shape
+        decoded_output = decoded_output[:, :target_channels, :target_time, :target_freq]
+
+        return decoded_output
+
+    def _run_upsampling_path(self, h: torch.Tensor) -> torch.Tensor:
+        for level in reversed(range(self.num_resolutions)):
+            stage = self.up[level]
+            for block_idx, block in enumerate(stage.block):
+                h = block(h, temb=None)
+                if stage.attn:
+                    h = stage.attn[block_idx](h)
+
+            if level != 0 and hasattr(stage, "upsample"):
+                h = stage.upsample(h)
+
+        return h
+
+    def _finalize_output(self, h: torch.Tensor) -> torch.Tensor:
+        if self.give_pre_end:
+            return h
+
+        h = self.norm_out(h)
+        h = self.non_linearity(h)
+        h = self.conv_out(h)
+        return torch.tanh(h) if self.tanh_out else h
+
+
+def decode_audio(latent: torch.Tensor, audio_decoder: "AudioDecoder", vocoder: "Vocoder") -> Audio:
+    """
+    Decode an audio latent representation using the provided audio decoder and vocoder.
+    Args:
+        latent: Input audio latent tensor.
+        audio_decoder: Model to decode the latent to waveform features.
+        vocoder: Model to convert decoded features to audio waveform.
+    Returns:
+        Decoded audio with waveform and sampling rate.
+    """
+    decoded_audio = audio_decoder(latent)
+    waveform = vocoder(decoded_audio).squeeze(0).float()
+    return Audio(waveform=waveform, sampling_rate=vocoder.output_sampling_rate)