Upload 14 files

.gitattributes

@@ -33,3 +33,5 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
 *.zip filter=lfs diff=lfs merge=lfs -text
 *.zst filter=lfs diff=lfs merge=lfs -text
 *tfevents* filter=lfs diff=lfs merge=lfs -text
+assets/mountain.mp4 filter=lfs diff=lfs merge=lfs -text
+assets/river.mp4 filter=lfs diff=lfs merge=lfs -text

LICENSE

@@ -0,0 +1,201 @@
+                                 Apache License
+                           Version 2.0, January 2004
+                        http://www.apache.org/licenses/
+
+   TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION
+
+   1. Definitions.
+
+      "License" shall mean the terms and conditions for use, reproduction,
+      and distribution as defined by Sections 1 through 9 of this document.
+
+      "Licensor" shall mean the copyright owner or entity authorized by
+      the copyright owner that is granting the License.
+
+      "Legal Entity" shall mean the union of the acting entity and all
+      other entities that control, are controlled by, or are under common
+      control with that entity. For the purposes of this definition,
+      "control" means (i) the power, direct or indirect, to cause the
+      direction or management of such entity, whether by contract or
+      otherwise, or (ii) ownership of fifty percent (50%) or more of the
+      outstanding shares, or (iii) beneficial ownership of such entity.
+
+      "You" (or "Your") shall mean an individual or Legal Entity
+      exercising permissions granted by this License.
+
+      "Source" form shall mean the preferred form for making modifications,
+      including but not limited to software source code, documentation
+      source, and configuration files.
+
+      "Object" form shall mean any form resulting from mechanical
+      transformation or translation of a Source form, including but
+      not limited to compiled object code, generated documentation,
+      and conversions to other media types.
+
+      "Work" shall mean the work of authorship, whether in Source or
+      Object form, made available under the License, as indicated by a
+      copyright notice that is included in or attached to the work
+      (an example is provided in the Appendix below).
+
+      "Derivative Works" shall mean any work, whether in Source or Object
+      form, that is based on (or derived from) the Work and for which the
+      editorial revisions, annotations, elaborations, or other modifications
+      represent, as a whole, an original work of authorship. For the purposes
+      of this License, Derivative Works shall not include works that remain
+      separable from, or merely link (or bind by name) to the interfaces of,
+      the Work and Derivative Works thereof.
+
+      "Contribution" shall mean any work of authorship, including
+      the original version of the Work and any modifications or additions
+      to that Work or Derivative Works thereof, that is intentionally
+      submitted to Licensor for inclusion in the Work by the copyright owner
+      or by an individual or Legal Entity authorized to submit on behalf of
+      the copyright owner. For the purposes of this definition, "submitted"
+      means any form of electronic, verbal, or written communication sent
+      to the Licensor or its representatives, including but not limited to
+      communication on electronic mailing lists, source code control systems,
+      and issue tracking systems that are managed by, or on behalf of, the
+      Licensor for the purpose of discussing and improving the Work, but
+      excluding communication that is conspicuously marked or otherwise
+      designated in writing by the copyright owner as "Not a Contribution."
+
+      "Contributor" shall mean Licensor and any individual or Legal Entity
+      on behalf of whom a Contribution has been received by Licensor and
+      subsequently incorporated within the Work.
+
+   2. Grant of Copyright License. Subject to the terms and conditions of
+      this License, each Contributor hereby grants to You a perpetual,
+      worldwide, non-exclusive, no-charge, royalty-free, irrevocable
+      copyright license to reproduce, prepare Derivative Works of,
+      publicly display, publicly perform, sublicense, and distribute the
+      Work and such Derivative Works in Source or Object form.
+
+   3. Grant of Patent License. Subject to the terms and conditions of
+      this License, each Contributor hereby grants to You a perpetual,
+      worldwide, non-exclusive, no-charge, royalty-free, irrevocable
+      (except as stated in this section) patent license to make, have made,
+      use, offer to sell, sell, import, and otherwise transfer the Work,
+      where such license applies only to those patent claims licensable
+      by such Contributor that are necessarily infringed by their
+      Contribution(s) alone or by combination of their Contribution(s)
+      with the Work to which such Contribution(s) was submitted. If You
+      institute patent litigation against any entity (including a
+      cross-claim or counterclaim in a lawsuit) alleging that the Work
+      or a Contribution incorporated within the Work constitutes direct
+      or contributory patent infringement, then any patent licenses
+      granted to You under this License for that Work shall terminate
+      as of the date such litigation is filed.
+
+   4. Redistribution. You may reproduce and distribute copies of the
+      Work or Derivative Works thereof in any medium, with or without
+      modifications, and in Source or Object form, provided that You
+      meet the following conditions:
+
+      (a) You must give any other recipients of the Work or
+          Derivative Works a copy of this License; and
+
+      (b) You must cause any modified files to carry prominent notices
+          stating that You changed the files; and
+
+      (c) You must retain, in the Source form of any Derivative Works
+          that You distribute, all copyright, patent, trademark, and
+          attribution notices from the Source form of the Work,
+          excluding those notices that do not pertain to any part of
+          the Derivative Works; and
+
+      (d) If the Work includes a "NOTICE" text file as part of its
+          distribution, then any Derivative Works that You distribute must
+          include a readable copy of the attribution notices contained
+          within such NOTICE file, excluding those notices that do not
+          pertain to any part of the Derivative Works, in at least one
+          of the following places: within a NOTICE text file distributed
+          as part of the Derivative Works; within the Source form or
+          documentation, if provided along with the Derivative Works; or,
+          within a display generated by the Derivative Works, if and
+          wherever such third-party notices normally appear. The contents
+          of the NOTICE file are for informational purposes only and
+          do not modify the License. You may add Your own attribution
+          notices within Derivative Works that You distribute, alongside
+          or as an addendum to the NOTICE text from the Work, provided
+          that such additional attribution notices cannot be construed
+          as modifying the License.
+
+      You may add Your own copyright statement to Your modifications and
+      may provide additional or different license terms and conditions
+      for use, reproduction, or distribution of Your modifications, or
+      for any such Derivative Works as a whole, provided Your use,
+      reproduction, and distribution of the Work otherwise complies with
+      the conditions stated in this License.
+
+   5. Submission of Contributions. Unless You explicitly state otherwise,
+      any Contribution intentionally submitted for inclusion in the Work
+      by You to the Licensor shall be under the terms and conditions of
+      this License, without any additional terms or conditions.
+      Notwithstanding the above, nothing herein shall supersede or modify
+      the terms of any separate license agreement you may have executed
+      with Licensor regarding such Contributions.
+
+   6. Trademarks. This License does not grant permission to use the trade
+      names, trademarks, service marks, or product names of the Licensor,
+      except as required for reasonable and customary use in describing the
+      origin of the Work and reproducing the content of the NOTICE file.
+
+   7. Disclaimer of Warranty. Unless required by applicable law or
+      agreed to in writing, Licensor provides the Work (and each
+      Contributor provides its Contributions) 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 the Work and assume any
+      risks associated with Your exercise of permissions under this License.
+
+   8. 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 any Contributor be
+      liable to You for damages, including any direct, indirect, special,
+      incidental, or consequential damages of any character arising as a
+      result of this License or out of the use or inability to use the
+      Work (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 such Contributor
+      has been advised of the possibility of such damages.
+
+   9. Accepting Warranty or Additional Liability. While redistributing
+      the Work or Derivative Works thereof, You may choose to offer,
+      and charge a fee for, acceptance of support, warranty, indemnity,
+      or other liability obligations and/or rights consistent with this
+      License. However, in accepting such obligations, You may act only
+      on Your own behalf and on Your sole responsibility, not on behalf
+      of any other Contributor, and only if You agree to indemnify,
+      defend, and hold each Contributor harmless for any liability
+      incurred by, or claims asserted against, such Contributor by reason
+      of your accepting any such warranty or additional liability.
+
+   END OF TERMS AND CONDITIONS
+
+   APPENDIX: How to apply the Apache License to your work.
+
+      To apply the Apache License to your work, attach the following
+      boilerplate notice, with the fields enclosed by brackets "[]"
+      replaced with your own identifying information. (Don't include
+      the brackets!)  The text should be enclosed in the appropriate
+      comment syntax for the file format. We also recommend that a
+      file or class name and description of purpose be included on the
+      same "printed page" as the copyright notice for easier
+      identification within third-party archives.
+
+   Copyright [yyyy] [name of copyright owner]
+
+   Licensed under the Apache License, Version 2.0 (the "License");
+   you may not use this file except in compliance with the License.
+   You may obtain a copy of the License at
+
+       http://www.apache.org/licenses/LICENSE-2.0
+
+   Unless required by applicable law or agreed to in writing, software
+   distributed under the License is distributed on an "AS IS" BASIS,
+   WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+   See the License for the specific language governing permissions and
+   limitations under the License.

README.md

@@ -1,13 +1,132 @@
----
-title: Q8 Ltx Video
-emoji: 👀
-colorFrom: red
-colorTo: indigo
-sdk: gradio
-sdk_version: 5.12.0
-app_file: app.py
-pinned: false
-short_description: Generate videos with LTX-Video fast!
----
-
-Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference
+# Q8 LTX-Video optimized for Ada
+
+This repository shows how to use the Q8 kernels from [`KONAKONA666/q8_kernels`](https://github.com/KONAKONA666/q8_kernels) with `diffusers` to optimize inference of [LTX-Video](https://huggingface.co/Lightricks/LTX-Video) on ADA GPUs. Go from 16.192 secs to 9.572 secs while reducing memory from 7GBs to 5GBs without quality loss 🤪 With `torch.compile()`, the time reduces further to 6.747 secs 🔥
+
+The Q8 transformer checkpoint is available here: [`sayakpaul/q8-ltx-video`](https://hf.co/sayakpaul/q8-ltx-video).
+
+## Getting started
+
+Install the dependencies:
+
+```bash
+pip install -U transformers accelerate
+git clone https://github.com/huggingface/diffusers && cd diffusers && pip install -e . && cd ..
+```
+
+Then install `q8_kernels`, following instructions from [here](https://github.com/KONAKONA666/q8_kernels/?tab=readme-ov-file#installation).
+
+To run inference with the Q8 kernels, we need some minor changes in `diffusers`. Apply [this patch](https://github.com/sayakpaul/q8-ltx-video/blob/368f549ca5136daf89049c9efe32748e73aca317/updates.patch) to take those into account:
+
+```bash
+git apply updates.patch
+```
+
+Now we can run inference:
+
+```bash
+python inference.py \
+    --prompt="A woman with long brown hair and light skin smiles at another woman with long blonde hair. The woman with brown hair wears a black jacket and has a small, barely noticeable mole on her right cheek. The camera angle is a close-up, focused on the woman with brown hair's face. The lighting is warm and natural, likely from the setting sun, casting a soft glow on the scene. The scene appears to be real-life footage"
+    --negative_prompt="worst quality, inconsistent motion, blurry, jittery, distorted" \
+    --q8_transformer_path="sayakpaul/q8-ltx-video"
+```
+
+## Why does the repo exist and some more details?
+
+There already exists [`KONAKONA666/LTX-Video`](https://github.com/KONAKONA666/LTX-Video). Then why this repo?
+
+That repo uses custom implementations of the LTX-Video pipeline components and can be hard to directly use in `diffusers`. This repo repurposes the kernels from the `q8_kernels` on the components directly from `diffusers`.
+
+<details>
+<summary>More details</summary>
+
+We do this by first converting the state dict of the original [LTX-Video transformer](https://huggingface.co/Lightricks/LTX-Video/tree/main/transformer). This includes FP8 quantization. This process also requires replacing:
+
+* linear layers of the model
+* RMSNorms of the model
+* GELUs of the model
+
+before the converted state dict is loaded into the model. Some layer params are kept in FP32 and some layers are not even quantized. Replacement utilities are in [`q8_ltx.py`](./q8_ltx.py).
+
+The model can then be serialized. The conversion and serialization are coded in [`conversion_utils.py`](./conversion_utils.py).
+
+During loading the model and using it for inference, we:
+
+* initialize the transformer model under a "meta" device
+* follow the same layer replacement scheme as detailed above
+* populate the converted state dict
+* replace the attention processors to use [the flash attention implementation](https://github.com/KONAKONA666/q8_kernels/blob/9cee3f3d4ca5ec8ab463179be32c8001e31f8f33/q8_kernels/functional/flash_attention.py) one from `q8_kernels`
+
+Refer [here](https://github.com/sayakpaul/q8-ltx-video/blob/368f549ca5136daf89049c9efe32748e73aca317/inference.py#L48) more details. Additionally, we leverage [flash-attention implementation](https://github.com/sayakpaul/q8-ltx-video/blob/368f549ca5136daf89049c9efe32748e73aca317/q8_attention_processors.py#L44) from `q8_kernels` which provides further speedup.
+
+</details>
+
+## Performance
+
+Below numbers were obtained for `max_sequence_length=512`, `num_inference_steps=50`, `num_frames=81`, `resolution=480x704`. Rest of the arguments were fixed at their default values as noticed in the [pipeline call signature of LTX-Video](https://github.com/huggingface/diffusers/blob/4b9f1c7d8c2e476eed38af3144b79105a5efcd93/src/diffusers/pipelines/ltx/pipeline_ltx.py#L496). The numbers also don't include the VAE decoding time to solely focus on the transformer.
+
+
+|  | **Time (Secs)** | **Memory (MB)** |
+|:-----------:|:-----------:|:-----------:|
+| Non Q8  | 16.192 | 7172.86  |
+| Non Q8 (+ compile)  | 16.205 | -  |
+| Q8  | 9.572  | 5413.51  |
+| Q8 (+ compile)  | 6.747  | -  |
+
+Benchmarking script is available in [`benchmark.py`](./benchmark.py). You would need to download the precomputed
+prompt embeddings from [here](https://huggingface.co/sayakpaul/q8-ltx-video/blob/main/prompt_embeds.pt) before running the benchmark.
+
+<details>
+<summary>Env</summary>
+
+```bash
++-----------------------------------------------------------------------------------------+
+| NVIDIA-SMI 560.35.05              Driver Version: 560.35.05      CUDA Version: 12.6     |
+|-----------------------------------------+------------------------+----------------------+
+| GPU  Name                 Persistence-M | Bus-Id          Disp.A | Volatile Uncorr. ECC |
+| Fan  Temp   Perf          Pwr:Usage/Cap |           Memory-Usage | GPU-Util  Compute M. |
+|                                         |                        |               MIG M. |
+|=========================================+========================+======================|
+|   0  NVIDIA GeForce RTX 4090        Off |   00000000:01:00.0 Off |                  Off |
+|  0%   46C    P8             18W /  450W |       2MiB /  24564MiB |      0%      Default |
+|                                         |                        |                  N/A |
++-----------------------------------------+------------------------+----------------------+
+```
+
+`diffusers-cli env`:
+
+```bash
+- 🤗 Diffusers version: 0.33.0.dev0
+- Platform: Linux-6.8.0-49-generic-x86_64-with-glibc2.39
+- Running on Google Colab?: No
+- Python version: 3.10.12
+- PyTorch version (GPU?): 2.5.1+cu124 (True)
+- Flax version (CPU?/GPU?/TPU?): not installed (NA)
+- Jax version: not installed
+- JaxLib version: not installed
+- Huggingface_hub version: 0.27.0
+- Transformers version: 4.47.1
+- Accelerate version: 1.2.1
+- PEFT version: 0.13.2
+- Bitsandbytes version: 0.44.1
+- Safetensors version: 0.4.4
+- xFormers version: 0.0.29.post1
+- Accelerator: NVIDIA GeForce RTX 4090, 24564 MiB
+NVIDIA GeForce RTX 4090, 24564 MiB
+- Using GPU in script?: <fill in>
+- Using distributed or parallel set-up in script?: <fill in>
+```
+
+</details>
+
+> [!NOTE]
+> The RoPE implementation from `q8_kernels` [isn't usable as of 1st Jan 2025](https://github.com/KONAKONA666/q8_kernels/blob/9cee3f3d4ca5ec8ab463179be32c8001e31f8f33/q8_kernels/functional/rope.py#L26). So, we resort to using [the one](https://github.com/huggingface/diffusers/blob/91008aabc4b8dbd96a356ab6f457f3bd84b10e8b/src/diffusers/models/transformers/transformer_ltx.py#L464) from `diffusers`.
+
+
+## Comparison
+
+Check out [this page](https://wandb.ai/sayakpaul/q8-ltx-video/runs/89h6ac5) on Weights and Biases that provides some comparative results. Generated videos are also available [here](./videos/).
+
+## Acknowledgement
+
+KONAKONA666's works on [`KONAKONA666/q8_kernels`](https://github.com/KONAKONA666/q8_kernels) and [KONAKONA666/LTX-Video](https://github.com/KONAKONA666/LTX-Video).
+

app.py

@@ -0,0 +1,156 @@
+import gradio as gr
+from app_utils import prepare_pipeline, compute_hash
+import os
+from diffusers.utils import export_to_video
+import tempfile
+import torch
+from inference import load_text_encoding_pipeline
+
+
+text_encoding_pipeline = load_text_encoding_pipeline()
+inference_pipeline = prepare_pipeline()
+
+
+def create_advanced_options():
+    with gr.Accordion("Advanced Options (Optional)", open=False):
+        seed = gr.Slider(label="Seed", minimum=0, maximum=1000000, step=1, value=646373)
+        inference_steps = gr.Slider(label="Inference Steps", minimum=1, maximum=50, step=1, value=30)
+        guidance_scale = gr.Slider(label="Guidance Scale", minimum=1.0, maximum=5.0, step=0.1, value=3.0)
+        max_sequence_length = gr.Slider(label="Maximum sequence length", minimum=128, maximum=512, step=1, value=128)
+        fps = gr.Slider(label="FPS", minimum=21, maximum=30, step=1, value=24)
+        return [
+            seed,
+            inference_steps,
+            guidance_scale,
+            max_sequence_length,
+            fps
+        ]
+
+
+@torch.no_grad()
+def generate_video_from_text(prompt, negative_prompt, seed, steps, guidance_scale, max_sequence_length, fps):
+    global text_encoding_pipeline
+    prompt_embeds, prompt_attention_mask, negative_prompt_embeds, negative_prompt_attention_mask = text_encoding_pipeline.encode_prompt(
+        prompt=prompt, negative_prompt=negative_prompt, max_sequence_length=max_sequence_length
+    )
+    global inference_pipeline
+    video = inference_pipeline(
+        prompt_embeds=prompt_embeds,
+        prompt_attention_mask=prompt_attention_mask,
+        negative_prompt_embeds=negative_prompt_embeds,
+        negative_prompt_attention_mask=negative_prompt_attention_mask,
+        guidance_scale=guidance_scale,
+        width=768,
+        height=512,
+        num_frames=121,
+        num_inference_steps=steps,
+        max_sequence_length=max_sequence_length,
+        generator=torch.manual_seed(seed),
+    ).frames[0]
+    out_path = tempfile.mkstemp(suffix=".mp4")
+    export_to_video(video, out_path, fps=fps)
+    return out_path
+
+
+with gr.Blocks(theme=gr.themes.Soft()) as iface:
+    with gr.Row(elem_id="title-row"):
+        gr.Markdown(
+            """
+        <div style="text-align: center; margin-bottom: 1em">
+            <h1 style="font-size: 2.5em; font-weight: 600; margin: 0.5em 0;">Fast Video Generation with <a href="https://github.com/sayakpaul/q8-ltx-video">Q8 LTX Video</a></h1>
+        </div>
+        """
+        )
+    with gr.Row(elem_id="title-row"):
+        gr.HTML(  # add technical report link
+            """
+        <div style="display:flex;column-gap:4px;">
+            <span>This space is modified from the original <a href="https://huggingface.co/spaces/Lightricks/LTX-Video-Playground">LTX-Video playground</a>. It uses optimized Q8 kernels along with torch.compile to allow for ultra-fast video
+            generation. As a result, it restricts generations to 121 frames with 512x768 resolution. For more details, refer to <a href="https://github.com/sayakpaul/q8-ltx-video">this link</a>.
+        </div>
+        """
+        )
+    with gr.Accordion(" 📖 Tips for Best Results", open=False, elem_id="instructions-accordion"):
+        gr.Markdown(
+            """
+        📝 Prompt Engineering
+        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
+        See examples for more inspiration.
+        🎮 Parameter Guide
+        - Resolution Preset: Higher resolutions for detailed scenes, lower for faster generation and simpler scenes
+        - Seed: Save seed values to recreate specific styles or compositions you like
+        - Guidance Scale: 3-3.5 are the recommended values
+        - Inference Steps: More steps (40+) for quality, fewer steps (20-30) for speed
+        - When using detailed prompts, use a higher `max_sequence_length` value.
+        """
+        )
+
+    with gr.Tabs():
+        # Text to Video Tab
+        with gr.TabItem("Text to Video"):
+            with gr.Row():
+                with gr.Column():
+                    txt2vid_prompt = gr.Textbox(
+                        label="Enter Your Prompt",
+                        placeholder="Describe the video you want to generate (minimum 50 characters)...",
+                        value="A woman with long brown hair and light skin smiles at another woman with long blonde hair. The woman with brown hair wears a black jacket and has a small, barely noticeable mole on her right cheek. The camera angle is a close-up, focused on the woman with brown hair's face. The lighting is warm and natural, likely from the setting sun, casting a soft glow on the scene. The scene appears to be real-life footage.",
+                        lines=5,
+                    )
+
+                    txt2vid_negative_prompt = gr.Textbox(
+                        label="Enter Negative Prompt",
+                        placeholder="Describe what you don't want in the video...",
+                        value="low quality, worst quality, deformed, distorted, disfigured, motion smear, motion artifacts, fused fingers, bad anatomy, weird hand, ugly",
+                        lines=2,
+                    )
+                    txt2vid_advanced = create_advanced_options()
+
+                    txt2vid_generate = gr.Button(
+                        "Generate Video",
+                        variant="primary",
+                        size="lg",
+                    )
+
+                with gr.Column():
+                    txt2vid_output = gr.Video(label="Generated Output")
+
+            with gr.Row():
+                gr.Examples(
+                    examples=[
+                        [
+                            "A clear, turquoise river flows through a rocky canyon, cascading over a small waterfall and forming a pool of water at the bottom.The river is the main focus of the scene, with its clear water reflecting the surrounding trees and rocks. The canyon walls are steep and rocky, with some vegetation growing on them. The trees are mostly pine trees, with their green needles contrasting with the brown and gray rocks. The overall tone of the scene is one of peace and tranquility.",
+                            "low quality, worst quality, deformed, distorted, disfigured, motion smear, motion artifacts, fused fingers, bad anatomy, weird hand, ugly",
+                            "assets/river.mp4",
+                        ],
+                        [
+                            "The camera pans over a snow-covered mountain range, revealing a vast expanse of snow-capped peaks and valleys.The mountains are covered in a thick layer of snow, with some areas appearing almost white while others have a slightly darker, almost grayish hue. The peaks are jagged and irregular, with some rising sharply into the sky while others are more rounded. The valleys are deep and narrow, with steep slopes that are also covered in snow. The trees in the foreground are mostly bare, with only a few leaves remaining on their branches. The sky is overcast, with thick clouds obscuring the sun. The overall impression is one of peace and tranquility, with the snow-covered mountains standing as a testament to the power and beauty of nature.",
+                            "low quality, worst quality, deformed, distorted, disfigured, motion smear, motion artifacts, fused fingers, bad anatomy, weird hand, ugly",
+                            "assets/mountain.mp4",
+                        ],
+                    ],
+                    inputs=[txt2vid_prompt, txt2vid_negative_prompt, txt2vid_output],
+                    label="Example Text-to-Video Generations",
+                )
+
+    txt2vid_generate.click(
+        fn=generate_video_from_text,
+        inputs=[
+            txt2vid_prompt,
+            txt2vid_negative_prompt,
+            *txt2vid_advanced,
+        ],
+        outputs=txt2vid_output,
+        concurrency_limit=1,
+        concurrency_id="generate_video_from_text",
+    )
+
+if __name__ == "__main__":
+    iface.queue(max_size=64, default_concurrency_limit=1, api_open=False).launch(share=True, show_api=False)

app_utils.py

@@ -0,0 +1,99 @@
+from inference import load_q8_transformer
+import hashlib
+from q8_kernels.graph.graph import make_dynamic_graphed_callable
+from argparse import Namespace
+from diffusers import LTXPipeline
+import types
+import torch
+
+# To account for the type-casting in `ff_output` of `LTXVideoTransformerBlock` 
+def patched_ltx_transformer_forward(
+    self,
+    hidden_states: torch.Tensor,
+    encoder_hidden_states: torch.Tensor,
+    temb: torch.Tensor,
+    image_rotary_emb = None,
+    encoder_attention_mask = None,
+) -> torch.Tensor:
+    batch_size = hidden_states.size(0)
+    norm_hidden_states = self.norm1(hidden_states)
+
+    num_ada_params = self.scale_shift_table.shape[0]
+    ada_values = self.scale_shift_table[None, None] + temb.reshape(batch_size, temb.size(1), num_ada_params, -1)
+    shift_msa, scale_msa, gate_msa, shift_mlp, scale_mlp, gate_mlp = ada_values.unbind(dim=2)
+    norm_hidden_states = norm_hidden_states * (1 + scale_msa) + shift_msa
+
+    attn_hidden_states = self.attn1(
+        hidden_states=norm_hidden_states,
+        encoder_hidden_states=None,
+        image_rotary_emb=image_rotary_emb,
+    )
+    hidden_states = hidden_states + attn_hidden_states * gate_msa
+
+    attn_hidden_states = self.attn2(
+        hidden_states,
+        encoder_hidden_states=encoder_hidden_states,
+        image_rotary_emb=None,
+        attention_mask=encoder_attention_mask,
+    )
+    hidden_states = hidden_states + attn_hidden_states
+    norm_hidden_states = self.norm2(hidden_states) * (1 + scale_mlp) + shift_mlp
+
+    ff_output = self.ff(norm_hidden_states).to(norm_hidden_states.dtype)
+    hidden_states = hidden_states + ff_output * gate_mlp
+
+    return hidden_states
+
+def load_transformer():
+    args = Namespace()
+    args.q8_transformer_path = "sayakpaul/q8-ltx-video"
+    transformer = load_q8_transformer(args)
+
+    transformer.to(torch.bfloat16)
+    for b in transformer.transformer_blocks:
+        b.to(dtype=torch.float)
+
+    for n, m in transformer.transformer_blocks.named_parameters():
+        if "scale_shift_table" in n:
+            m.data = m.data.to(torch.bfloat16)
+    
+    for b in transformer.transformer_blocks:
+        b.forward = types.MethodType(patched_ltx_transformer_forward, b)
+    
+    transformer.forward = make_dynamic_graphed_callable(transformer.forward)
+    return transformer 
+
+def warmup_transformer(pipe):
+    prompt_embeds = torch.load("prompt_embeds.pt", map_location="cuda", weights_only=True)
+    for _ in range(5):
+        _ = pipe(
+            **prompt_embeds,
+            output_type="latent",  
+            width=768,
+            height=512,
+            num_frames=121
+        )
+
+def prepare_pipeline():
+    pipe = LTXPipeline.from_pretrained("Lightricks/LTX-Video", text_encoder=None, torch_dtype=torch.bfloat16)
+    pipe.transformer = load_transformer()
+    pipe = pipe.to("cuda")
+    pipe.transformer.compile()
+    pipe.set_progress_bar_config(disable=True)
+
+    warmup_transformer(pipe)
+    return pipe
+
+
+def compute_hash(text: str) -> str:
+    # Encode the text to bytes
+    text_bytes = text.encode("utf-8")
+
+    # Create a SHA-256 hash object
+    hash_object = hashlib.sha256()
+
+    # Update the hash object with the text bytes
+    hash_object.update(text_bytes)
+
+    # Return the hexadecimal representation of the hash
+    return hash_object.hexdigest()

assets/mountain.mp4

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:540d193b2819996741ae23cf0bf106e1c3b226880b44ae1e5e54a8f38a181e9f
+size 1463129

assets/river.mp4

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:56f21a785ac95e20e0933e553a1f5f82946737f4077ff4d71bf11de5b0ba5162
+size 1788643

conversion_utils.py

@@ -0,0 +1,103 @@
+"""
+References:
+https://github.com/KONAKONA666/q8_kernels/blob/9cee3f3d4ca5ec8ab463179be32c8001e31f8f33/q8_kernels/utils/convert_weights.py
+"""
+
+import torch
+from q8_ltx import replace_gelu, replace_linear, replace_rms_norm, MODULES_TO_NOT_CONVERT
+import argparse
+from diffusers import LTXVideoTransformer3DModel
+from q8_kernels.functional.quantizer import quantize
+from q8_kernels.functional.fast_hadamard import hadamard_transform
+
+
+def convert_state_dict(orig_state_dict):
+    prefix = "transformer_blocks"
+    transformer_block_keys = []
+    non_transformer_block_keys = []
+    for k in orig_state_dict:
+        if prefix in k:
+            transformer_block_keys.append(k)
+        else:
+            non_transformer_block_keys.append(k)
+    attn_keys = []
+    ffn_keys = []
+    scale_shift_keys = []
+    for k in transformer_block_keys:
+        if "attn" in k:
+            attn_keys.append(k)
+    for k in transformer_block_keys:
+        if "ff" in k:
+            ffn_keys.append(k)
+    for k in transformer_block_keys:
+        if "scale_shift_table" in k:
+            scale_shift_keys.append(k)
+
+    assert len(attn_keys + ffn_keys + scale_shift_keys) == len(transformer_block_keys), "error"
+
+    new_state_dict = {}
+    for k in attn_keys:
+        new_key = k
+        if "norm" in k and "weight" in k:
+            new_state_dict[new_key] = orig_state_dict[k].float()
+        elif "bias" in k:
+            new_state_dict[new_key] = orig_state_dict[k].float()
+        elif "weight" in k:
+            w_quant, w_scales = quantize(hadamard_transform(orig_state_dict[k].cuda().to(torch.bfloat16)))
+            assert w_quant.dtype == torch.int8, k
+            new_state_dict[new_key] = w_quant
+            new_state_dict[new_key.replace("weight", "scales")] = w_scales
+
+    for k in ffn_keys:
+        new_key = k
+
+        if "bias" in k:
+            new_state_dict[new_key] = orig_state_dict[k].float()
+        elif "weight" in k:
+            w_quant, w_scales = quantize(hadamard_transform(orig_state_dict[k].cuda().to(torch.bfloat16)))
+            assert w_quant.dtype == torch.int8, k
+            new_state_dict[new_key] = w_quant
+            new_state_dict[new_key.replace("weight", "scales")] = w_scales
+
+    for k in scale_shift_keys:
+        new_state_dict[k] = orig_state_dict[k]
+
+    for k in non_transformer_block_keys:
+        new_state_dict[k] = orig_state_dict[k]
+
+    return new_state_dict
+
+
+@torch.no_grad()
+def main(args):
+    transformer = LTXVideoTransformer3DModel.from_pretrained(args.input_path, subfolder="transformer").to("cuda")
+    new_state_dict = convert_state_dict(transformer.state_dict())
+    transformer = replace_gelu(transformer)[0]
+    transformer = replace_linear(transformer)[0]
+    transformer = replace_rms_norm(transformer)[0]
+
+    m, u = transformer.load_state_dict(new_state_dict, strict=True)
+    for name, module in transformer.named_modules():
+        if any(n in name for n in MODULES_TO_NOT_CONVERT):
+            if hasattr(module, "weight"):
+                assert module.weight.dtype == torch.float32
+            elif hasattr(module, "linear"):
+                assert module.linear.weight.dtype == torch.float32
+        elif getattr(module, "weight", None) is not None:
+            print(f"Non FP32 {name=} {module.weight.dtype=}")
+            if "to_" in name:
+                assert module.weight.dtype != torch.float32, f"{name=}, {module.weight.dtype=}"
+
+    transformer.save_pretrained(args.output_path)
+    print(f"Model saved in {args.output_path}")
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser()
+
+    parser.add_argument("--input_path", type=str, required=True)
+    parser.add_argument("--output_path", type=str, required=True)
+
+    args = parser.parse_args()
+
+    main(args)

inference.py

@@ -0,0 +1,139 @@
+from diffusers import LTXPipeline, LTXVideoTransformer3DModel
+from huggingface_hub import hf_hub_download
+import argparse
+import os
+from q8_ltx import check_transformer_replaced_correctly, replace_gelu, replace_linear, replace_rms_norm
+import safetensors.torch
+from q8_kernels.graph.graph import make_dynamic_graphed_callable
+import torch
+import gc
+from diffusers.utils import export_to_video
+
+
+# Taken from
+# https://github.com/KONAKONA666/LTX-Video/blob/c8462ed2e359cda4dec7f49d98029994e850dc90/inference.py#L115C1-L138C28
+def convert_prompt_to_filename(text: str, max_len: int = 20) -> str:
+    # Remove non-letters and convert to lowercase
+    clean_text = "".join(char.lower() for char in text if char.isalpha() or char.isspace())
+    # Split into words
+    words = clean_text.split()
+
+    # Build result string keeping track of length
+    result = []
+    current_length = 0
+
+    for word in words:
+        # Add word length plus 1 for underscore (except for first word)
+        new_length = current_length + len(word)
+        if new_length <= max_len:
+            result.append(word)
+            current_length += len(word)
+        else:
+            break
+
+    return "-".join(result)
+
+
+def load_text_encoding_pipeline():
+    return LTXPipeline.from_pretrained(
+        "Lightricks/LTX-Video", transformer=None, vae=None, torch_dtype=torch.bfloat16
+    ).to("cuda")
+
+
+def encode_prompt(pipe, prompt, negative_prompt, max_sequence_length=128):
+    prompt_embeds, prompt_attention_mask, negative_prompt_embeds, negative_prompt_attention_mask = pipe.encode_prompt(
+        prompt=prompt, negative_prompt=negative_prompt, max_sequence_length=max_sequence_length
+    )
+    return prompt_embeds, prompt_attention_mask, negative_prompt_embeds, negative_prompt_attention_mask
+
+
+def load_q8_transformer(args):
+    with torch.device("meta"):
+        transformer_config = LTXVideoTransformer3DModel.load_config("Lightricks/LTX-Video", subfolder="transformer")
+        transformer = LTXVideoTransformer3DModel.from_config(transformer_config)
+
+    transformer = replace_gelu(transformer)[0]
+    transformer = replace_linear(transformer)[0]
+    transformer = replace_rms_norm(transformer)[0]
+
+    if os.path.isfile(f"{args.q8_transformer_path}/diffusion_pytorch_model.safetensors"):
+        state_dict = safetensors.torch.load_file(f"{args.q8_transformer_path}/diffusion_pytorch_model.safetensors")
+    else:
+        state_dict = safetensors.torch.load_file(
+            hf_hub_download(args.q8_transformer_path, "diffusion_pytorch_model.safetensors")
+        )
+    transformer.load_state_dict(state_dict, strict=True, assign=True)
+    check_transformer_replaced_correctly(transformer)
+    return transformer
+
+
+@torch.no_grad()
+def main(args):
+    text_encoding_pipeline = load_text_encoding_pipeline()
+    prompt_embeds, prompt_attention_mask, negative_prompt_embeds, negative_prompt_attention_mask = encode_prompt(
+        pipe=text_encoding_pipeline,
+        prompt=args.prompt,
+        negative_prompt=args.negative_prompt,
+        max_sequence_length=args.max_sequence_length,
+    )
+    del text_encoding_pipeline
+    torch.cuda.empty_cache()
+    torch.cuda.reset_peak_memory_stats()
+    gc.collect()
+
+    if args.q8_transformer_path:
+        transformer = load_q8_transformer(args)
+        pipe = LTXPipeline.from_pretrained("Lightricks/LTX-Video", transformer=None, text_encoder=None)
+        pipe.transformer = transformer
+
+        pipe.transformer = pipe.transformer.to(torch.bfloat16)
+        for b in pipe.transformer.transformer_blocks:
+            b.to(dtype=torch.float)
+
+        for n, m in pipe.transformer.transformer_blocks.named_parameters():
+            if "scale_shift_table" in n:
+                m.data = m.data.to(torch.bfloat16)
+
+        pipe.transformer.forward = make_dynamic_graphed_callable(pipe.transformer.forward)
+        pipe.vae = pipe.vae.to(torch.bfloat16)
+
+    else:
+        pipe = LTXPipeline.from_pretrained("Lightricks/LTX-Video", text_encoder=None, torch_dtype=torch.bfloat16)
+
+    pipe = pipe.to("cuda")
+
+    width, height = args.resolution.split("x")[::-1]
+    video = pipe(
+        prompt_embeds=prompt_embeds,
+        prompt_attention_mask=prompt_attention_mask,
+        negative_prompt_embeds=negative_prompt_embeds,
+        negative_prompt_attention_mask=negative_prompt_attention_mask,
+        width=int(width),
+        height=int(height),
+        num_frames=args.num_frames,
+        num_inference_steps=args.steps,
+        max_sequence_length=args.max_sequence_length,
+        generator=torch.manual_seed(2025),
+    ).frames[0]
+    print(f"Max memory: {torch.cuda.max_memory_allocated() / 1024 / 1024} MB.")
+
+    if args.out_path is None:
+        filename_from_prompt = convert_prompt_to_filename(args.prompt, max_len=30)
+        base_filename = f"{filename_from_prompt}_{args.num_frames}x{height}x{width}"
+        base_filename += "_q8" if args.q8_transformer_path is not None else ""
+        args.out_path = base_filename + ".mp4"
+    export_to_video(video, args.out_path, fps=24)
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--q8_transformer_path", type=str, default=None)
+    parser.add_argument("--prompt", type=str)
+    parser.add_argument("--negative_prompt", type=str, default=None)
+    parser.add_argument("--num_frames", type=int, default=81)
+    parser.add_argument("--resolution", type=str, default="480x704")
+    parser.add_argument("--steps", type=int, default=50)
+    parser.add_argument("--max_sequence_length", type=int, default=512)
+    parser.add_argument("--out_path", type=str, default=None)
+    args = parser.parse_args()
+    main(args)

pipeline_ltx.py

@@ -0,0 +1,811 @@
+# Copyright 2024 Lightricks and The HuggingFace Team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""
+To account for certain changes pertaining to Q8Linear.
+"""
+
+import inspect
+from typing import Any, Callable, Dict, List, Optional, Union
+
+import numpy as np
+import torch
+from transformers import T5EncoderModel, T5TokenizerFast
+
+from diffusers.callbacks import MultiPipelineCallbacks, PipelineCallback
+from diffusers.loaders import FromSingleFileMixin, LTXVideoLoraLoaderMixin
+from diffusers.models.autoencoders import AutoencoderKLLTXVideo
+from diffusers.models.transformers import LTXVideoTransformer3DModel
+from diffusers.schedulers import FlowMatchEulerDiscreteScheduler
+from diffusers.utils import is_torch_xla_available, logging, replace_example_docstring
+from diffusers.utils.torch_utils import randn_tensor
+from diffusers.video_processor import VideoProcessor
+from diffusers.pipelines.pipeline_utils import DiffusionPipeline
+from diffusers.pipelines.ltx.pipeline_output import LTXPipelineOutput
+
+try:
+    import q8_kernels # noqa
+    from q8_kernels.modules.linear import Q8Linear
+except:
+    Q8Linear = None
+
+
+if is_torch_xla_available():
+    import torch_xla.core.xla_model as xm
+
+    XLA_AVAILABLE = True
+else:
+    XLA_AVAILABLE = False
+
+logger = logging.get_logger(__name__)  # pylint: disable=invalid-name
+
+EXAMPLE_DOC_STRING = """
+    Examples:
+        ```py
+        >>> import torch
+        >>> from diffusers import LTXPipeline
+        >>> from diffusers.utils import export_to_video
+
+        >>> pipe = LTXPipeline.from_pretrained("Lightricks/LTX-Video", torch_dtype=torch.bfloat16)
+        >>> pipe.to("cuda")
+
+        >>> prompt = "A woman with long brown hair and light skin smiles at another woman with long blonde hair. The woman with brown hair wears a black jacket and has a small, barely noticeable mole on her right cheek. The camera angle is a close-up, focused on the woman with brown hair's face. The lighting is warm and natural, likely from the setting sun, casting a soft glow on the scene. The scene appears to be real-life footage"
+        >>> negative_prompt = "worst quality, inconsistent motion, blurry, jittery, distorted"
+
+        >>> video = pipe(
+        ...     prompt=prompt,
+        ...     negative_prompt=negative_prompt,
+        ...     width=704,
+        ...     height=480,
+        ...     num_frames=161,
+        ...     num_inference_steps=50,
+        ... ).frames[0]
+        >>> export_to_video(video, "output.mp4", fps=24)
+        ```
+"""
+
+
+# Copied from diffusers.pipelines.flux.pipeline_flux.calculate_shift
+def calculate_shift(
+    image_seq_len,
+    base_seq_len: int = 256,
+    max_seq_len: int = 4096,
+    base_shift: float = 0.5,
+    max_shift: float = 1.16,
+):
+    m = (max_shift - base_shift) / (max_seq_len - base_seq_len)
+    b = base_shift - m * base_seq_len
+    mu = image_seq_len * m + b
+    return mu
+
+
+# Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.retrieve_timesteps
+def retrieve_timesteps(
+    scheduler,
+    num_inference_steps: Optional[int] = None,
+    device: Optional[Union[str, torch.device]] = None,
+    timesteps: Optional[List[int]] = None,
+    sigmas: Optional[List[float]] = None,
+    **kwargs,
+):
+    r"""
+    Calls the scheduler's `set_timesteps` method and retrieves timesteps from the scheduler after the call. Handles
+    custom timesteps. Any kwargs will be supplied to `scheduler.set_timesteps`.
+
+    Args:
+        scheduler (`SchedulerMixin`):
+            The scheduler to get timesteps from.
+        num_inference_steps (`int`):
+            The number of diffusion steps used when generating samples with a pre-trained model. If used, `timesteps`
+            must be `None`.
+        device (`str` or `torch.device`, *optional*):
+            The device to which the timesteps should be moved to. If `None`, the timesteps are not moved.
+        timesteps (`List[int]`, *optional*):
+            Custom timesteps used to override the timestep spacing strategy of the scheduler. If `timesteps` is passed,
+            `num_inference_steps` and `sigmas` must be `None`.
+        sigmas (`List[float]`, *optional*):
+            Custom sigmas used to override the timestep spacing strategy of the scheduler. If `sigmas` is passed,
+            `num_inference_steps` and `timesteps` must be `None`.
+
+    Returns:
+        `Tuple[torch.Tensor, int]`: A tuple where the first element is the timestep schedule from the scheduler and the
+        second element is the number of inference steps.
+    """
+    if timesteps is not None and sigmas is not None:
+        raise ValueError("Only one of `timesteps` or `sigmas` can be passed. Please choose one to set custom values")
+    if timesteps is not None:
+        accepts_timesteps = "timesteps" in set(inspect.signature(scheduler.set_timesteps).parameters.keys())
+        if not accepts_timesteps:
+            raise ValueError(
+                f"The current scheduler class {scheduler.__class__}'s `set_timesteps` does not support custom"
+                f" timestep schedules. Please check whether you are using the correct scheduler."
+            )
+        scheduler.set_timesteps(timesteps=timesteps, device=device, **kwargs)
+        timesteps = scheduler.timesteps
+        num_inference_steps = len(timesteps)
+    elif sigmas is not None:
+        accept_sigmas = "sigmas" in set(inspect.signature(scheduler.set_timesteps).parameters.keys())
+        if not accept_sigmas:
+            raise ValueError(
+                f"The current scheduler class {scheduler.__class__}'s `set_timesteps` does not support custom"
+                f" sigmas schedules. Please check whether you are using the correct scheduler."
+            )
+        scheduler.set_timesteps(sigmas=sigmas, device=device, **kwargs)
+        timesteps = scheduler.timesteps
+        num_inference_steps = len(timesteps)
+    else:
+        scheduler.set_timesteps(num_inference_steps, device=device, **kwargs)
+        timesteps = scheduler.timesteps
+    return timesteps, num_inference_steps
+
+
+class LTXPipeline(DiffusionPipeline, FromSingleFileMixin, LTXVideoLoraLoaderMixin):
+    r"""
+    Pipeline for text-to-video generation.
+
+    Reference: https://github.com/Lightricks/LTX-Video
+
+    Args:
+        transformer ([`LTXVideoTransformer3DModel`]):
+            Conditional Transformer architecture to denoise the encoded video latents.
+        scheduler ([`FlowMatchEulerDiscreteScheduler`]):
+            A scheduler to be used in combination with `transformer` to denoise the encoded image latents.
+        vae ([`AutoencoderKLLTXVideo`]):
+            Variational Auto-Encoder (VAE) Model to encode and decode images to and from latent representations.
+        text_encoder ([`T5EncoderModel`]):
+            [T5](https://huggingface.co/docs/transformers/en/model_doc/t5#transformers.T5EncoderModel), specifically
+            the [google/t5-v1_1-xxl](https://huggingface.co/google/t5-v1_1-xxl) variant.
+        tokenizer (`CLIPTokenizer`):
+            Tokenizer of class
+            [CLIPTokenizer](https://huggingface.co/docs/transformers/en/model_doc/clip#transformers.CLIPTokenizer).
+        tokenizer (`T5TokenizerFast`):
+            Second Tokenizer of class
+            [T5TokenizerFast](https://huggingface.co/docs/transformers/en/model_doc/t5#transformers.T5TokenizerFast).
+    """
+
+    model_cpu_offload_seq = "text_encoder->transformer->vae"
+    _optional_components = []
+    _callback_tensor_inputs = ["latents", "prompt_embeds", "negative_prompt_embeds"]
+
+    def __init__(
+        self,
+        scheduler: FlowMatchEulerDiscreteScheduler,
+        vae: AutoencoderKLLTXVideo,
+        text_encoder: T5EncoderModel,
+        tokenizer: T5TokenizerFast,
+        transformer: LTXVideoTransformer3DModel,
+    ):
+        super().__init__()
+
+        self.register_modules(
+            vae=vae,
+            text_encoder=text_encoder,
+            tokenizer=tokenizer,
+            transformer=transformer,
+            scheduler=scheduler,
+        )
+
+        self.vae_spatial_compression_ratio = (
+            self.vae.spatial_compression_ratio if getattr(self, "vae", None) is not None else 32
+        )
+        self.vae_temporal_compression_ratio = (
+            self.vae.temporal_compression_ratio if getattr(self, "vae", None) is not None else 8
+        )
+        self.transformer_spatial_patch_size = (
+            self.transformer.config.patch_size if getattr(self, "transformer", None) is not None else 1
+        )
+        self.transformer_temporal_patch_size = (
+            self.transformer.config.patch_size_t if getattr(self, "transformer") is not None else 1
+        )
+
+        self.video_processor = VideoProcessor(vae_scale_factor=self.vae_spatial_compression_ratio)
+        self.tokenizer_max_length = (
+            self.tokenizer.model_max_length if getattr(self, "tokenizer", None) is not None else 128
+        )
+
+    def _get_t5_prompt_embeds(
+        self,
+        prompt: Union[str, List[str]] = None,
+        num_videos_per_prompt: int = 1,
+        max_sequence_length: int = 128,
+        device: Optional[torch.device] = None,
+        dtype: Optional[torch.dtype] = None,
+    ):
+        device = device or self._execution_device
+        dtype = dtype or self.text_encoder.dtype
+
+        prompt = [prompt] if isinstance(prompt, str) else prompt
+        batch_size = len(prompt)
+
+        text_inputs = self.tokenizer(
+            prompt,
+            padding="max_length",
+            max_length=max_sequence_length,
+            truncation=True,
+            add_special_tokens=True,
+            return_tensors="pt",
+        )
+        text_input_ids = text_inputs.input_ids
+        prompt_attention_mask = text_inputs.attention_mask
+        prompt_attention_mask = prompt_attention_mask.bool().to(device)
+
+        untruncated_ids = self.tokenizer(prompt, padding="longest", return_tensors="pt").input_ids
+
+        if untruncated_ids.shape[-1] >= text_input_ids.shape[-1] and not torch.equal(text_input_ids, untruncated_ids):
+            removed_text = self.tokenizer.batch_decode(untruncated_ids[:, max_sequence_length - 1 : -1])
+            logger.warning(
+                "The following part of your input was truncated because `max_sequence_length` is set to "
+                f" {max_sequence_length} tokens: {removed_text}"
+            )
+
+        prompt_embeds = self.text_encoder(text_input_ids.to(device))[0]
+        prompt_embeds = prompt_embeds.to(dtype=dtype, device=device)
+
+        # duplicate text embeddings for each generation per prompt, using mps friendly method
+        _, seq_len, _ = prompt_embeds.shape
+        prompt_embeds = prompt_embeds.repeat(1, num_videos_per_prompt, 1)
+        prompt_embeds = prompt_embeds.view(batch_size * num_videos_per_prompt, seq_len, -1)
+
+        prompt_attention_mask = prompt_attention_mask.view(batch_size, -1)
+        prompt_attention_mask = prompt_attention_mask.repeat(num_videos_per_prompt, 1)
+
+        return prompt_embeds, prompt_attention_mask
+
+    # Copied from diffusers.pipelines.mochi.pipeline_mochi.MochiPipeline.encode_prompt with 256->128
+    def encode_prompt(
+        self,
+        prompt: Union[str, List[str]],
+        negative_prompt: Optional[Union[str, List[str]]] = None,
+        do_classifier_free_guidance: bool = True,
+        num_videos_per_prompt: int = 1,
+        prompt_embeds: Optional[torch.Tensor] = None,
+        negative_prompt_embeds: Optional[torch.Tensor] = None,
+        prompt_attention_mask: Optional[torch.Tensor] = None,
+        negative_prompt_attention_mask: Optional[torch.Tensor] = None,
+        max_sequence_length: int = 128,
+        device: Optional[torch.device] = None,
+        dtype: Optional[torch.dtype] = None,
+    ):
+        r"""
+        Encodes the prompt into text encoder hidden states.
+
+        Args:
+            prompt (`str` or `List[str]`, *optional*):
+                prompt to be encoded
+            negative_prompt (`str` or `List[str]`, *optional*):
+                The prompt or prompts not to guide the image generation. If not defined, one has to pass
+                `negative_prompt_embeds` instead. Ignored when not using guidance (i.e., ignored if `guidance_scale` is
+                less than `1`).
+            do_classifier_free_guidance (`bool`, *optional*, defaults to `True`):
+                Whether to use classifier free guidance or not.
+            num_videos_per_prompt (`int`, *optional*, defaults to 1):
+                Number of videos that should be generated per prompt. torch device to place the resulting embeddings on
+            prompt_embeds (`torch.Tensor`, *optional*):
+                Pre-generated text embeddings. Can be used to easily tweak text inputs, *e.g.* prompt weighting. If not
+                provided, text embeddings will be generated from `prompt` input argument.
+            negative_prompt_embeds (`torch.Tensor`, *optional*):
+                Pre-generated negative text embeddings. Can be used to easily tweak text inputs, *e.g.* prompt
+                weighting. If not provided, negative_prompt_embeds will be generated from `negative_prompt` input
+                argument.
+            device: (`torch.device`, *optional*):
+                torch device
+            dtype: (`torch.dtype`, *optional*):
+                torch dtype
+        """
+        print(f"{max_sequence_length=}")
+        device = device or self._execution_device
+
+        prompt = [prompt] if isinstance(prompt, str) else prompt
+        if prompt is not None:
+            batch_size = len(prompt)
+        else:
+            batch_size = prompt_embeds.shape[0]
+
+        if prompt_embeds is None:
+            prompt_embeds, prompt_attention_mask = self._get_t5_prompt_embeds(
+                prompt=prompt,
+                num_videos_per_prompt=num_videos_per_prompt,
+                max_sequence_length=max_sequence_length,
+                device=device,
+                dtype=dtype,
+            )
+
+        if do_classifier_free_guidance and negative_prompt_embeds is None:
+            negative_prompt = negative_prompt or ""
+            negative_prompt = batch_size * [negative_prompt] if isinstance(negative_prompt, str) else negative_prompt
+
+            if prompt is not None and type(prompt) is not type(negative_prompt):
+                raise TypeError(
+                    f"`negative_prompt` should be the same type to `prompt`, but got {type(negative_prompt)} !="
+                    f" {type(prompt)}."
+                )
+            elif batch_size != len(negative_prompt):
+                raise ValueError(
+                    f"`negative_prompt`: {negative_prompt} has batch size {len(negative_prompt)}, but `prompt`:"
+                    f" {prompt} has batch size {batch_size}. Please make sure that passed `negative_prompt` matches"
+                    " the batch size of `prompt`."
+                )
+
+            negative_prompt_embeds, negative_prompt_attention_mask = self._get_t5_prompt_embeds(
+                prompt=negative_prompt,
+                num_videos_per_prompt=num_videos_per_prompt,
+                max_sequence_length=max_sequence_length,
+                device=device,
+                dtype=dtype,
+            )
+
+        return prompt_embeds, prompt_attention_mask, negative_prompt_embeds, negative_prompt_attention_mask
+
+    def check_inputs(
+        self,
+        prompt,
+        height,
+        width,
+        callback_on_step_end_tensor_inputs=None,
+        prompt_embeds=None,
+        negative_prompt_embeds=None,
+        prompt_attention_mask=None,
+        negative_prompt_attention_mask=None,
+    ):
+        if height % 32 != 0 or width % 32 != 0:
+            raise ValueError(f"`height` and `width` have to be divisible by 32 but are {height} and {width}.")
+
+        if callback_on_step_end_tensor_inputs is not None and not all(
+            k in self._callback_tensor_inputs for k in callback_on_step_end_tensor_inputs
+        ):
+            raise ValueError(
+                f"`callback_on_step_end_tensor_inputs` has to be in {self._callback_tensor_inputs}, but found {[k for k in callback_on_step_end_tensor_inputs if k not in self._callback_tensor_inputs]}"
+            )
+
+        if prompt is not None and prompt_embeds is not None:
+            raise ValueError(
+                f"Cannot forward both `prompt`: {prompt} and `prompt_embeds`: {prompt_embeds}. Please make sure to"
+                " only forward one of the two."
+            )
+        elif prompt is None and prompt_embeds is None:
+            raise ValueError(
+                "Provide either `prompt` or `prompt_embeds`. Cannot leave both `prompt` and `prompt_embeds` undefined."
+            )
+        elif prompt is not None and (not isinstance(prompt, str) and not isinstance(prompt, list)):
+            raise ValueError(f"`prompt` has to be of type `str` or `list` but is {type(prompt)}")
+
+        if prompt_embeds is not None and prompt_attention_mask is None:
+            raise ValueError("Must provide `prompt_attention_mask` when specifying `prompt_embeds`.")
+
+        if negative_prompt_embeds is not None and negative_prompt_attention_mask is None:
+            raise ValueError("Must provide `negative_prompt_attention_mask` when specifying `negative_prompt_embeds`.")
+
+        if prompt_embeds is not None and negative_prompt_embeds is not None:
+            if prompt_embeds.shape != negative_prompt_embeds.shape:
+                raise ValueError(
+                    "`prompt_embeds` and `negative_prompt_embeds` must have the same shape when passed directly, but"
+                    f" got: `prompt_embeds` {prompt_embeds.shape} != `negative_prompt_embeds`"
+                    f" {negative_prompt_embeds.shape}."
+                )
+            if prompt_attention_mask.shape != negative_prompt_attention_mask.shape:
+                raise ValueError(
+                    "`prompt_attention_mask` and `negative_prompt_attention_mask` must have the same shape when passed directly, but"
+                    f" got: `prompt_attention_mask` {prompt_attention_mask.shape} != `negative_prompt_attention_mask`"
+                    f" {negative_prompt_attention_mask.shape}."
+                )
+
+    @staticmethod
+    def _pack_latents(latents: torch.Tensor, patch_size: int = 1, patch_size_t: int = 1) -> torch.Tensor:
+        # Unpacked latents of shape are [B, C, F, H, W] are patched into tokens of shape [B, C, F // p_t, p_t, H // p, p, W // p, p].
+        # The patch dimensions are then permuted and collapsed into the channel dimension of shape:
+        # [B, F // p_t * H // p * W // p, C * p_t * p * p] (an ndim=3 tensor).
+        # dim=0 is the batch size, dim=1 is the effective video sequence length, dim=2 is the effective number of input features
+        batch_size, num_channels, num_frames, height, width = latents.shape
+        post_patch_num_frames = num_frames // patch_size_t
+        post_patch_height = height // patch_size
+        post_patch_width = width // patch_size
+        latents = latents.reshape(
+            batch_size,
+            -1,
+            post_patch_num_frames,
+            patch_size_t,
+            post_patch_height,
+            patch_size,
+            post_patch_width,
+            patch_size,
+        )
+        latents = latents.permute(0, 2, 4, 6, 1, 3, 5, 7).flatten(4, 7).flatten(1, 3)
+        return latents
+
+    @staticmethod
+    def _unpack_latents(
+        latents: torch.Tensor, num_frames: int, height: int, width: int, patch_size: int = 1, patch_size_t: int = 1
+    ) -> torch.Tensor:
+        # Packed latents of shape [B, S, D] (S is the effective video sequence length, D is the effective feature dimensions)
+        # are unpacked and reshaped into a video tensor of shape [B, C, F, H, W]. This is the inverse operation of
+        # what happens in the `_pack_latents` method.
+        batch_size = latents.size(0)
+        latents = latents.reshape(batch_size, num_frames, height, width, -1, patch_size_t, patch_size, patch_size)
+        latents = latents.permute(0, 4, 1, 5, 2, 6, 3, 7).flatten(6, 7).flatten(4, 5).flatten(2, 3)
+        return latents
+
+    @staticmethod
+    def _normalize_latents(
+        latents: torch.Tensor, latents_mean: torch.Tensor, latents_std: torch.Tensor, scaling_factor: float = 1.0
+    ) -> torch.Tensor:
+        # Normalize latents across the channel dimension [B, C, F, H, W]
+        latents_mean = latents_mean.view(1, -1, 1, 1, 1).to(latents.device, latents.dtype)
+        latents_std = latents_std.view(1, -1, 1, 1, 1).to(latents.device, latents.dtype)
+        latents = (latents - latents_mean) * scaling_factor / latents_std
+        return latents
+
+    @staticmethod
+    def _denormalize_latents(
+        latents: torch.Tensor, latents_mean: torch.Tensor, latents_std: torch.Tensor, scaling_factor: float = 1.0
+    ) -> torch.Tensor:
+        # Denormalize latents across the channel dimension [B, C, F, H, W]
+        latents_mean = latents_mean.view(1, -1, 1, 1, 1).to(latents.device, latents.dtype)
+        latents_std = latents_std.view(1, -1, 1, 1, 1).to(latents.device, latents.dtype)
+        latents = latents * latents_std / scaling_factor + latents_mean
+        return latents
+
+    def prepare_latents(
+        self,
+        batch_size: int = 1,
+        num_channels_latents: int = 128,
+        height: int = 512,
+        width: int = 704,
+        num_frames: int = 161,
+        dtype: Optional[torch.dtype] = None,
+        device: Optional[torch.device] = None,
+        generator: Optional[torch.Generator] = None,
+        latents: Optional[torch.Tensor] = None,
+    ) -> torch.Tensor:
+        if latents is not None:
+            return latents.to(device=device, dtype=dtype)
+
+        height = height // self.vae_spatial_compression_ratio
+        width = width // self.vae_spatial_compression_ratio
+        num_frames = (num_frames - 1) // self.vae_temporal_compression_ratio + 1
+
+        shape = (batch_size, num_channels_latents, num_frames, height, width)
+
+        if isinstance(generator, list) and len(generator) != batch_size:
+            raise ValueError(
+                f"You have passed a list of generators of length {len(generator)}, but requested an effective batch"
+                f" size of {batch_size}. Make sure the batch size matches the length of the generators."
+            )
+
+        latents = randn_tensor(shape, generator=generator, device=device, dtype=dtype)
+        latents = self._pack_latents(
+            latents, self.transformer_spatial_patch_size, self.transformer_temporal_patch_size
+        )
+        return latents
+
+    @property
+    def guidance_scale(self):
+        return self._guidance_scale
+
+    @property
+    def do_classifier_free_guidance(self):
+        return self._guidance_scale > 1.0
+
+    @property
+    def num_timesteps(self):
+        return self._num_timesteps
+
+    @property
+    def attention_kwargs(self):
+        return self._attention_kwargs
+
+    @property
+    def interrupt(self):
+        return self._interrupt
+
+    @torch.no_grad()
+    @replace_example_docstring(EXAMPLE_DOC_STRING)
+    def __call__(
+        self,
+        prompt: Union[str, List[str]] = None,
+        negative_prompt: Optional[Union[str, List[str]]] = None,
+        height: int = 512,
+        width: int = 704,
+        num_frames: int = 161,
+        frame_rate: int = 25,
+        num_inference_steps: int = 50,
+        timesteps: List[int] = None,
+        guidance_scale: float = 3,
+        num_videos_per_prompt: Optional[int] = 1,
+        generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None,
+        latents: Optional[torch.Tensor] = None,
+        prompt_embeds: Optional[torch.Tensor] = None,
+        prompt_attention_mask: Optional[torch.Tensor] = None,
+        negative_prompt_embeds: Optional[torch.Tensor] = None,
+        negative_prompt_attention_mask: Optional[torch.Tensor] = None,
+        decode_timestep: Union[float, List[float]] = 0.0,
+        decode_noise_scale: Optional[Union[float, List[float]]] = None,
+        output_type: Optional[str] = "pil",
+        return_dict: bool = True,
+        attention_kwargs: Optional[Dict[str, Any]] = None,
+        callback_on_step_end: Optional[Callable[[int, int, Dict], None]] = None,
+        callback_on_step_end_tensor_inputs: List[str] = ["latents"],
+        max_sequence_length: int = 128,
+    ):
+        r"""
+        Function invoked when calling the pipeline for generation.
+
+        Args:
+            prompt (`str` or `List[str]`, *optional*):
+                The prompt or prompts to guide the image generation. If not defined, one has to pass `prompt_embeds`.
+                instead.
+            height (`int`, defaults to `512`):
+                The height in pixels of the generated image. This is set to 480 by default for the best results.
+            width (`int`, defaults to `704`):
+                The width in pixels of the generated image. This is set to 848 by default for the best results.
+            num_frames (`int`, defaults to `161`):
+                The number of video frames to generate
+            num_inference_steps (`int`, *optional*, defaults to 50):
+                The number of denoising steps. More denoising steps usually lead to a higher quality image at the
+                expense of slower inference.
+            timesteps (`List[int]`, *optional*):
+                Custom timesteps to use for the denoising process with schedulers which support a `timesteps` argument
+                in their `set_timesteps` method. If not defined, the default behavior when `num_inference_steps` is
+                passed will be used. Must be in descending order.
+            guidance_scale (`float`, defaults to `3 `):
+                Guidance scale as defined in [Classifier-Free Diffusion Guidance](https://arxiv.org/abs/2207.12598).
+                `guidance_scale` is defined as `w` of equation 2. of [Imagen
+                Paper](https://arxiv.org/pdf/2205.11487.pdf). Guidance scale is enabled by setting `guidance_scale >
+                1`. Higher guidance scale encourages to generate images that are closely linked to the text `prompt`,
+                usually at the expense of lower image quality.
+            num_videos_per_prompt (`int`, *optional*, defaults to 1):
+                The number of videos to generate per prompt.
+            generator (`torch.Generator` or `List[torch.Generator]`, *optional*):
+                One or a list of [torch generator(s)](https://pytorch.org/docs/stable/generated/torch.Generator.html)
+                to make generation deterministic.
+            latents (`torch.Tensor`, *optional*):
+                Pre-generated noisy latents, sampled from a Gaussian distribution, to be used as inputs for image
+                generation. Can be used to tweak the same generation with different prompts. If not provided, a latents
+                tensor will ge generated by sampling using the supplied random `generator`.
+            prompt_embeds (`torch.Tensor`, *optional*):
+                Pre-generated text embeddings. Can be used to easily tweak text inputs, *e.g.* prompt weighting. If not
+                provided, text embeddings will be generated from `prompt` input argument.
+            prompt_attention_mask (`torch.Tensor`, *optional*):
+                Pre-generated attention mask for text embeddings.
+            negative_prompt_embeds (`torch.FloatTensor`, *optional*):
+                Pre-generated negative text embeddings. For PixArt-Sigma this negative prompt should be "". If not
+                provided, negative_prompt_embeds will be generated from `negative_prompt` input argument.
+            negative_prompt_attention_mask (`torch.FloatTensor`, *optional*):
+                Pre-generated attention mask for negative text embeddings.
+            decode_timestep (`float`, defaults to `0.0`):
+                The timestep at which generated video is decoded.
+            decode_noise_scale (`float`, defaults to `None`):
+                The interpolation factor between random noise and denoised latents at the decode timestep.
+            output_type (`str`, *optional*, defaults to `"pil"`):
+                The output format of the generate image. Choose between
+                [PIL](https://pillow.readthedocs.io/en/stable/): `PIL.Image.Image` or `np.array`.
+            return_dict (`bool`, *optional*, defaults to `True`):
+                Whether or not to return a [`~pipelines.ltx.LTXPipelineOutput`] instead of a plain tuple.
+            attention_kwargs (`dict`, *optional*):
+                A kwargs dictionary that if specified is passed along to the `AttentionProcessor` as defined under
+                `self.processor` in
+                [diffusers.models.attention_processor](https://github.com/huggingface/diffusers/blob/main/src/diffusers/models/attention_processor.py).
+            callback_on_step_end (`Callable`, *optional*):
+                A function that calls at the end of each denoising steps during the inference. The function is called
+                with the following arguments: `callback_on_step_end(self: DiffusionPipeline, step: int, timestep: int,
+                callback_kwargs: Dict)`. `callback_kwargs` will include a list of all tensors as specified by
+                `callback_on_step_end_tensor_inputs`.
+            callback_on_step_end_tensor_inputs (`List`, *optional*):
+                The list of tensor inputs for the `callback_on_step_end` function. The tensors specified in the list
+                will be passed as `callback_kwargs` argument. You will only be able to include variables listed in the
+                `._callback_tensor_inputs` attribute of your pipeline class.
+            max_sequence_length (`int` defaults to `128 `):
+                Maximum sequence length to use with the `prompt`.
+
+        Examples:
+
+        Returns:
+            [`~pipelines.ltx.LTXPipelineOutput`] or `tuple`:
+                If `return_dict` is `True`, [`~pipelines.ltx.LTXPipelineOutput`] is returned, otherwise a `tuple` is
+                returned where the first element is a list with the generated images.
+        """
+
+        if isinstance(callback_on_step_end, (PipelineCallback, MultiPipelineCallbacks)):
+            callback_on_step_end_tensor_inputs = callback_on_step_end.tensor_inputs
+
+        # 1. Check inputs. Raise error if not correct
+        self.check_inputs(
+            prompt=prompt,
+            height=height,
+            width=width,
+            callback_on_step_end_tensor_inputs=callback_on_step_end_tensor_inputs,
+            prompt_embeds=prompt_embeds,
+            negative_prompt_embeds=negative_prompt_embeds,
+            prompt_attention_mask=prompt_attention_mask,
+            negative_prompt_attention_mask=negative_prompt_attention_mask,
+        )
+
+        self._guidance_scale = guidance_scale
+        self._attention_kwargs = attention_kwargs
+        self._interrupt = False
+
+        # 2. Define call parameters
+        if prompt is not None and isinstance(prompt, str):
+            batch_size = 1
+        elif prompt is not None and isinstance(prompt, list):
+            batch_size = len(prompt)
+        else:
+            batch_size = prompt_embeds.shape[0]
+
+        device = self._execution_device
+
+        # 3. Prepare text embeddings
+        (
+            prompt_embeds,
+            prompt_attention_mask,
+            negative_prompt_embeds,
+            negative_prompt_attention_mask,
+        ) = self.encode_prompt(
+            prompt=prompt,
+            negative_prompt=negative_prompt,
+            do_classifier_free_guidance=self.do_classifier_free_guidance,
+            num_videos_per_prompt=num_videos_per_prompt,
+            prompt_embeds=prompt_embeds,
+            negative_prompt_embeds=negative_prompt_embeds,
+            prompt_attention_mask=prompt_attention_mask,
+            negative_prompt_attention_mask=negative_prompt_attention_mask,
+            max_sequence_length=max_sequence_length,
+            device=device,
+        )
+        if self.do_classifier_free_guidance:
+            prompt_embeds = torch.cat([negative_prompt_embeds, prompt_embeds], dim=0)
+            prompt_attention_mask = torch.cat([negative_prompt_attention_mask, prompt_attention_mask], dim=0)
+
+        if Q8Linear is not None and isinstance(self.transformer.transformer_blocks[0].attn1.to_q, Q8Linear):
+            prompt_attention_mask = prompt_attention_mask.to(torch.int64)
+            prompt_attention_mask = prompt_attention_mask.argmin(-1).int().squeeze()
+            prompt_attention_mask[prompt_attention_mask == 0] = max_sequence_length
+
+        # 4. Prepare latent variables
+        num_channels_latents = self.transformer.config.in_channels
+        latents = self.prepare_latents(
+            batch_size * num_videos_per_prompt,
+            num_channels_latents,
+            height,
+            width,
+            num_frames,
+            torch.float32,
+            device,
+            generator,
+            latents,
+        )
+
+        # 5. Prepare timesteps
+        latent_num_frames = (num_frames - 1) // self.vae_temporal_compression_ratio + 1
+        latent_height = height // self.vae_spatial_compression_ratio
+        latent_width = width // self.vae_spatial_compression_ratio
+        video_sequence_length = latent_num_frames * latent_height * latent_width
+        sigmas = np.linspace(1.0, 1 / num_inference_steps, num_inference_steps)
+        mu = calculate_shift(
+            video_sequence_length,
+            self.scheduler.config.base_image_seq_len,
+            self.scheduler.config.max_image_seq_len,
+            self.scheduler.config.base_shift,
+            self.scheduler.config.max_shift,
+        )
+        timesteps, num_inference_steps = retrieve_timesteps(
+            self.scheduler,
+            num_inference_steps,
+            device,
+            timesteps,
+            sigmas=sigmas,
+            mu=mu,
+        )
+        num_warmup_steps = max(len(timesteps) - num_inference_steps * self.scheduler.order, 0)
+        self._num_timesteps = len(timesteps)
+
+        # 6. Prepare micro-conditions
+        latent_frame_rate = frame_rate / self.vae_temporal_compression_ratio
+        rope_interpolation_scale = (
+            1 / latent_frame_rate,
+            self.vae_spatial_compression_ratio,
+            self.vae_spatial_compression_ratio,
+        )
+
+        # 7. Denoising loop
+        with self.progress_bar(total=num_inference_steps) as progress_bar:
+            for i, t in enumerate(timesteps):
+                if self.interrupt:
+                    continue
+
+                latent_model_input = torch.cat([latents] * 2) if self.do_classifier_free_guidance else latents
+                latent_model_input = latent_model_input.to(prompt_embeds.dtype)
+
+                # broadcast to batch dimension in a way that's compatible with ONNX/Core ML
+                timestep = t.expand(latent_model_input.shape[0])
+
+                noise_pred = self.transformer(
+                    hidden_states=latent_model_input,
+                    encoder_hidden_states=prompt_embeds,
+                    timestep=timestep,
+                    encoder_attention_mask=prompt_attention_mask,
+                    num_frames=latent_num_frames,
+                    height=latent_height,
+                    width=latent_width,
+                    rope_interpolation_scale=rope_interpolation_scale,
+                    attention_kwargs=attention_kwargs,
+                    return_dict=False,
+                )[0]
+                noise_pred = noise_pred.float()
+
+                if self.do_classifier_free_guidance:
+                    noise_pred_uncond, noise_pred_text = noise_pred.chunk(2)
+                    noise_pred = noise_pred_uncond + self.guidance_scale * (noise_pred_text - noise_pred_uncond)
+
+                # compute the previous noisy sample x_t -> x_t-1
+                latents = self.scheduler.step(noise_pred, t, latents, return_dict=False)[0]
+
+                if callback_on_step_end is not None:
+                    callback_kwargs = {}
+                    for k in callback_on_step_end_tensor_inputs:
+                        callback_kwargs[k] = locals()[k]
+                    callback_outputs = callback_on_step_end(self, i, t, callback_kwargs)
+
+                    latents = callback_outputs.pop("latents", latents)
+                    prompt_embeds = callback_outputs.pop("prompt_embeds", prompt_embeds)
+
+                # call the callback, if provided
+                if i == len(timesteps) - 1 or ((i + 1) > num_warmup_steps and (i + 1) % self.scheduler.order == 0):
+                    progress_bar.update()
+
+                if XLA_AVAILABLE:
+                    xm.mark_step()
+
+        if output_type == "latent":
+            video = latents
+        else:
+            latents = self._unpack_latents(
+                latents,
+                latent_num_frames,
+                latent_height,
+                latent_width,
+                self.transformer_spatial_patch_size,
+                self.transformer_temporal_patch_size,
+            )
+            latents = self._denormalize_latents(
+                latents, self.vae.latents_mean, self.vae.latents_std, self.vae.config.scaling_factor
+            )
+            latents = latents.to(prompt_embeds.dtype)
+
+            if not self.vae.config.timestep_conditioning:
+                timestep = None
+            else:
+                noise = torch.randn(latents.shape, generator=generator, device=device, dtype=latents.dtype)
+                if not isinstance(decode_timestep, list):
+                    decode_timestep = [decode_timestep] * batch_size
+                if decode_noise_scale is None:
+                    decode_noise_scale = decode_timestep
+                elif not isinstance(decode_noise_scale, list):
+                    decode_noise_scale = [decode_noise_scale] * batch_size
+
+                timestep = torch.tensor(decode_timestep, device=device, dtype=latents.dtype)
+                decode_noise_scale = torch.tensor(decode_noise_scale, device=device, dtype=latents.dtype)[
+                    :, None, None, None, None
+                ]
+                latents = (1 - decode_noise_scale) * latents + decode_noise_scale * noise
+
+            video = self.vae.decode(latents, timestep, return_dict=False)[0]
+            video = self.video_processor.postprocess_video(video, output_type=output_type)
+
+        # Offload all models
+        self.maybe_free_model_hooks()
+
+        if not return_dict:
+            return (video,)
+
+        return LTXPipelineOutput(frames=video)

prompt_embeds.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:2095fe12d573adf5adf6e2f0d163d494f792fe83deaeb6b20e67a28e6e44e913
+size 8391808

q8_attention_processors.py

@@ -0,0 +1,52 @@
+"""
+Reference:
+https://github.com/KONAKONA666/q8_kernels/blob/9cee3f3d4ca5ec8ab463179be32c8001e31f8f33/q8_kernels/modules/attention.py
+"""
+
+import torch
+import q8_kernels.functional as Q8F
+from diffusers.models.transformers.transformer_ltx import apply_rotary_emb
+from diffusers.models.attention import Attention
+
+NON_MM_PRECISION_TYPE = torch.bfloat16
+MM_PRECISION_TYPE = torch.bfloat16
+
+
+class LTXVideoQ8AttentionProcessor:
+    def __call__(
+        self,
+        attn: Attention,
+        hidden_states: torch.Tensor,
+        encoder_hidden_states=None,
+        attention_mask=None,
+        image_rotary_emb=None,
+    ) -> torch.Tensor:
+        if attention_mask is not None and attention_mask.ndim > 1:
+            attention_mask = attention_mask.argmin(-1).squeeze().int()
+
+        if encoder_hidden_states is None:
+            encoder_hidden_states = hidden_states
+
+        query = attn.to_q(hidden_states)
+        key = attn.to_k(encoder_hidden_states)
+        value = attn.to_v(encoder_hidden_states)
+
+        query = attn.norm_q(query, NON_MM_PRECISION_TYPE)
+        key = attn.norm_k(key, NON_MM_PRECISION_TYPE)
+
+        if image_rotary_emb is not None:
+            query = apply_rotary_emb(query, image_rotary_emb)
+            key = apply_rotary_emb(key, image_rotary_emb)
+
+        query = query.unflatten(2, (attn.heads, -1)).transpose(1, 2)
+        key = key.unflatten(2, (attn.heads, -1)).transpose(1, 2)
+        value = value.unflatten(2, (attn.heads, -1)).transpose(1, 2)
+
+        hidden_states = Q8F.flash_attention.flash_attn_func(
+            query, key, value, batch_mask=attention_mask, apply_qk_hadamard=True
+        )
+        hidden_states = hidden_states.transpose(1, 2).flatten(2, 3)
+
+        hidden_states = attn.to_out[0](hidden_states)
+        hidden_states = attn.to_out[1](hidden_states)
+        return hidden_states.to(NON_MM_PRECISION_TYPE)

q8_ltx.py

@@ -0,0 +1,141 @@
+"""
+References:
+https://github.com/KONAKONA666/q8_kernels/blob/9cee3f3d4ca5ec8ab463179be32c8001e31f8f33/q8_kernels/utils/convert_weights.py
+"""
+
+import torch
+import torch.nn as nn
+from q8_kernels.modules.rms_norm import RMSNorm as QRMSNorm
+from diffusers.models.normalization import RMSNorm
+from q8_kernels.modules.activations import GELU as QGELU
+from diffusers.models.activations import GELU
+from q8_kernels.modules.linear import Q8Linear
+from q8_attention_processors import LTXVideoQ8AttentionProcessor
+
+MODULES_TO_NOT_CONVERT = ["proj_in", "time_embed", "caption_projection", "proj_out"]
+
+
+def replace_linear(model, current_key_name=None, replaced=False):
+    for name, child in model.named_children():
+        if current_key_name is None:
+            current_key_name = []
+        current_key_name.append(name)
+
+        if isinstance(child, nn.Linear) and name not in MODULES_TO_NOT_CONVERT:
+            # Check if the current key is not in the `modules_to_not_convert`
+            current_key_name_str = ".".join(current_key_name)
+            if not any(
+                (key + "." in current_key_name_str) or (key == current_key_name_str) for key in MODULES_TO_NOT_CONVERT
+            ):
+                new_linear = Q8Linear(
+                    child.in_features, child.out_features, bias=child.bias is not None, device=child.weight.device
+                )
+                setattr(model, name, new_linear)
+                replaced = True
+        else:
+            replace_linear(model=child, current_key_name=current_key_name, replaced=replaced)
+
+        current_key_name.pop(-1)
+
+    return model, replaced
+
+
+def get_parent_module_and_attr(root, dotted_name: str):
+    """
+    Splits 'a.b.c' into:
+    - parent module = root.a.b
+    - attr_name = 'c'
+    """
+    parts = dotted_name.split(".")
+    *parent_parts, attr_name = parts
+    parent_module = root
+    for p in parent_parts:
+        parent_module = getattr(parent_module, p)
+    return parent_module, attr_name
+
+
+def replace_rms_norm(model):
+    modules_to_replace = []
+    for dotted_name, module in model.named_modules():
+        if isinstance(module, RMSNorm):
+            modules_to_replace.append((dotted_name, module))
+
+    replaced = False
+    for dotted_name, module in modules_to_replace:
+        parent, attr_name = get_parent_module_and_attr(model, dotted_name)
+        new_norm = QRMSNorm(
+            dim=module.dim,
+            elementwise_affine=module.elementwise_affine,
+        )
+        setattr(parent, attr_name, new_norm)
+        replaced = True
+
+    return model, replaced
+
+
+def replace_gelu(model, replaced=False):
+    for name, child in model.named_children():
+        if isinstance(child, GELU):
+            new_gelu = QGELU(
+                dim_in=child.proj.in_features,
+                dim_out=child.proj.out_features,
+                approximate=child.approximate,
+                bias=child.proj.bias is not None,
+            )
+            setattr(model, name, new_gelu)
+            replaced = True
+        else:
+            replace_gelu(model=child, replaced=replaced)
+
+    return model, replaced
+
+
+def set_attn_processors(model, processor):
+    def fn_recursive_attn_processor(name, module: torch.nn.Module, processor):
+        if hasattr(module, "set_processor"):
+            module.set_processor(processor)
+        for sub_name, child in module.named_children():
+            fn_recursive_attn_processor(f"{name}.{sub_name}", child, processor)
+
+    for name, module in model.named_children():
+        fn_recursive_attn_processor(name, module, processor)
+
+
+def attn_processors(model) -> dict:
+    processors = {}
+
+    def fn_recursive_add_processors(name: str, module: torch.nn.Module, processors: dict):
+        if hasattr(module, "get_processor"):
+            processors[f"{name}.processor"] = module.get_processor()
+
+        for sub_name, child in module.named_children():
+            fn_recursive_add_processors(f"{name}.{sub_name}", child, processors)
+
+        return processors
+
+    for name, module in model.named_children():
+        fn_recursive_add_processors(name, module, processors)
+
+    return processors
+
+
+def check_transformer_replaced_correctly(model):
+    for block in model.transformer_blocks:
+        assert isinstance(block.attn1.to_q, Q8Linear), f"{type(block.attn1.to_q)=} not linear."
+        assert isinstance(block.attn2.to_q, Q8Linear), f"{type(block.attn2.to_q)=} not linear."
+        assert block.attn1.to_q.weight.dtype == torch.int8, f"{block.attn1.to_q.weight.dtype=}."
+        assert block.attn2.to_q.weight.dtype == torch.int8, f"{name=} {block.attn2.to_q.weight.dtype=}."
+
+    for name, module in model.named_modules():
+        if "norm" in name and "norm_out" not in name:
+            assert isinstance(module, QRMSNorm), f"{name=}, {type(module)=}"
+
+    for block in model.transformer_blocks:
+        assert isinstance(block.ff.net[0], QGELU), f"{type(block.ff.net[0])=}"
+        if getattr(block.ff.net[0], "proj", None) is not None:
+            assert block.ff.net[0].proj.weight.dtype == torch.int8, f"{block.ff.net[0].proj.weight.dtype=}."
+
+    set_attn_processors(model, LTXVideoQ8AttentionProcessor())
+    all_attn_processors = attn_processors(model)
+    for k, v in all_attn_processors.items():
+        assert isinstance(v, LTXVideoQ8AttentionProcessor), f"{name} is not of type LTXVideoQ8AttentionProcessor."

requirements.txt

@@ -0,0 +1,7 @@
+torch
+diffusers
+transformers
+accelerate
+imageio
+Pillow
+-f git+git://github.com/KONAKONA666/q8_kernels.git