Duplicate from mehdiPrunaAI/LTX2_3_PrunedDecoder

.gitattributes

@@ -0,0 +1,55 @@
+*.7z filter=lfs diff=lfs merge=lfs -text
+*.arrow filter=lfs diff=lfs merge=lfs -text
+*.bin filter=lfs diff=lfs merge=lfs -text
+*.bz2 filter=lfs diff=lfs merge=lfs -text
+*.ckpt filter=lfs diff=lfs merge=lfs -text
+*.ftz filter=lfs diff=lfs merge=lfs -text
+*.gz filter=lfs diff=lfs merge=lfs -text
+*.h5 filter=lfs diff=lfs merge=lfs -text
+*.joblib filter=lfs diff=lfs merge=lfs -text
+*.lfs.* filter=lfs diff=lfs merge=lfs -text
+*.mlmodel filter=lfs diff=lfs merge=lfs -text
+*.model filter=lfs diff=lfs merge=lfs -text
+*.msgpack filter=lfs diff=lfs merge=lfs -text
+*.npy filter=lfs diff=lfs merge=lfs -text
+*.npz filter=lfs diff=lfs merge=lfs -text
+*.onnx filter=lfs diff=lfs merge=lfs -text
+*.ot filter=lfs diff=lfs merge=lfs -text
+*.parquet filter=lfs diff=lfs merge=lfs -text
+*.pb filter=lfs diff=lfs merge=lfs -text
+*.pickle filter=lfs diff=lfs merge=lfs -text
+*.pkl filter=lfs diff=lfs merge=lfs -text
+*.pt filter=lfs diff=lfs merge=lfs -text
+*.pth filter=lfs diff=lfs merge=lfs -text
+*.rar filter=lfs diff=lfs merge=lfs -text
+*.safetensors filter=lfs diff=lfs merge=lfs -text
+saved_model/**/* filter=lfs diff=lfs merge=lfs -text
+*.tar.* filter=lfs diff=lfs merge=lfs -text
+*.tar filter=lfs diff=lfs merge=lfs -text
+*.tflite filter=lfs diff=lfs merge=lfs -text
+*.tgz filter=lfs diff=lfs merge=lfs -text
+*.wasm filter=lfs diff=lfs merge=lfs -text
+*.xz filter=lfs diff=lfs merge=lfs -text
+*.zip filter=lfs diff=lfs merge=lfs -text
+*.zst filter=lfs diff=lfs merge=lfs -text
+*tfevents* filter=lfs diff=lfs merge=lfs -text
+example/student/04MhZs7YN08-Scene-0048.mp4 filter=lfs diff=lfs merge=lfs -text
+example/student/7y57zc-GvBA-Scene-0009.mp4 filter=lfs diff=lfs merge=lfs -text
+example/student/8YSha7iU2ZI-Scene-1025.mp4 filter=lfs diff=lfs merge=lfs -text
+example/student/97moBezrzLo-Scene-0076.mp4 filter=lfs diff=lfs merge=lfs -text
+example/student/9GDE5-rOOfk-Scene-0010.mp4 filter=lfs diff=lfs merge=lfs -text
+example/student/a5M8hWgQbjU-Scene-0191.mp4 filter=lfs diff=lfs merge=lfs -text
+example/student/DZt-RtpcJZQ-Scene-0043_m2.mp4 filter=lfs diff=lfs merge=lfs -text
+example/student/FAfMSWi0FSA-Scene-0200.mp4 filter=lfs diff=lfs merge=lfs -text
+example/student/ijGuAu0lLNM-Scene-0008.mp4 filter=lfs diff=lfs merge=lfs -text
+example/student/zFehFw9SbDw-Scene-0034.mp4 filter=lfs diff=lfs merge=lfs -text
+example/teacher/04MhZs7YN08-Scene-0048.mp4 filter=lfs diff=lfs merge=lfs -text
+example/teacher/7y57zc-GvBA-Scene-0009.mp4 filter=lfs diff=lfs merge=lfs -text
+example/teacher/8YSha7iU2ZI-Scene-1025.mp4 filter=lfs diff=lfs merge=lfs -text
+example/teacher/97moBezrzLo-Scene-0076.mp4 filter=lfs diff=lfs merge=lfs -text
+example/teacher/9GDE5-rOOfk-Scene-0010.mp4 filter=lfs diff=lfs merge=lfs -text
+example/teacher/a5M8hWgQbjU-Scene-0191.mp4 filter=lfs diff=lfs merge=lfs -text
+example/teacher/DZt-RtpcJZQ-Scene-0043_teacher.mp4 filter=lfs diff=lfs merge=lfs -text
+example/teacher/FAfMSWi0FSA-Scene-0200.mp4 filter=lfs diff=lfs merge=lfs -text
+example/teacher/ijGuAu0lLNM-Scene-0008.mp4 filter=lfs diff=lfs merge=lfs -text
+example/teacher/zFehFw9SbDw-Scene-0034.mp4 filter=lfs diff=lfs merge=lfs -text

README.md

@@ -0,0 +1,212 @@
+---
+base_model:
+- Lightricks/LTX-2.3
+pipeline_tag: text-to-video
+---
+<!-- header start -->
+<!-- 200823 -->
+<div style="width: auto; margin-left: auto; margin-right: auto">
+    <a href="https://docs.pruna.ai/en/latest/setup/pip.html" target="_blank" rel="noopener noreferrer">
+        <img src="https://imgur.com/rVAgqMY.png" alt="PrunaAI" style="width: 100%; min-width: 400px; display: block; margin: auto;">
+    </a>
+</div>
+<!-- header end -->
+
+[![GitHub](https://img.shields.io/badge/GitHub-PrunaAI-9334E9?style=plastic&logo=github&logoColor=white)](https://github.com/PrunaAI/pruna) &nbsp;
+[![Twitter/X](https://img.shields.io/badge/Twitter%2FX-@PrunaAI-9334E9?style=plastic&logo=x&logoColor=white)](https://twitter.com/PrunaAI) &nbsp;
+[![LinkedIn](https://img.shields.io/badge/LinkedIn-PrunaAI-9334E9?style=plastic&logo=linkedin&logoColor=white)](https://www.linkedin.com/company/pruna-ai) &nbsp;
+[![Discord](https://img.shields.io/badge/Discord-Join%20us-9334E9?style=plastic&logo=discord&logoColor=white)](https://discord.com/invite/JFQmtFKCjd) &nbsp;
+[![P-Models](https://img.shields.io/badge/Performance%20Models-Try%20them%20now-9334E9?style=plastic&logo=data:image/svg%2bxml;base64,PHN2ZyB4bWxucz0iaHR0cDovL3d3dy53My5vcmcvMjAwMC9zdmciIHdpZHRoPSI0MTEiIGhlaWdodD0iNDc3IiBmaWxsPSJub25lIj48cGF0aCBmaWxsPSIjZmZmIiBkPSJNMjMuNzAzIDMwMC41MzJjLTEuODQyLS42NzQtNC4xNDItMS4wNTctNS43MDMtMi4wMjktMTIuNTU4LTcuODQ2LTguODM3LTI1LjE3MiA1LjI1NS0yOC4zNjFhLjY2LjY2IDAgMCAwIC4zNDgtLjIwMS42LjYgMCAwIDAgLjE1Ni0uMzZxNi40OS04MS45ODYgNjguNTc3LTEzNC42MjlhLjM0LjM0IDAgMCAwIC4wOTItLjEyNS4zMy4zMyAwIDAgMC0uMDI2LS4zLjMzLjMzIDAgMCAwLS4xMTItLjEwOGMtNC4xMTUtMi40OTctOC42NC00Ljc1LTEyLjM1My03LjQ0M2ExOTAgMTkwIDAgMCAxLTE4Ljc5NC0xNS41MjJjLTIuOTM2LTIuNzY4LTMuNzAzLTUuODkxLjU4LTcuNzdhLjUuNSAwIDAgMCAuMjgtLjYwOEM1MyA3NS45NSA0OC40MTcgNDguNTk5IDQ1LjUxOCAxOS42NXEtLjM5My0zLjg1Mi0uMjcxLTYuMzJhLjc1Ljc1IDAgMCAxIC4xMzQtLjQuNy43IDAgMCAxIC4zMjQtLjI1NXEyLjEwNC0uODQyIDUuODQ0LTEuNDMxYzU4LjExNC05LjEzNiAxMjcuMDA5IDguNDI1IDE0NC4zMDcgNzIuNjQ1cS4xMDMuMzY0LjE3OCAwIDUuMDk2LTI0LjA5NyA0LjYzOC00OC4xOTMtLjA0Ny0yLjc4NyAxLjg3OS00Ljc4OCAxMC40MDctMTAuOCAyNC43OTgtOC42ODdjNy40NzEgMS4wOTUgNy43NDIgNC42MzggNy4yMDkgMTEuMWEzMzQgMzM0IDAgMCAxLTExLjg4NCA2NC43MDYuMzg4LjM4OCAwIDAgMCAuMTgyLjQ0NC40LjQgMCAwIDAgLjE2NC4wNTFjOTIuNDIxIDguNDUzIDE1OC44NzUgNzkuNDg5IDE2NS4wODQgMTcxLjI3NGExIDEgMCAwIDAgLjEyOS40MjQuODMuODMgMCAwIDAgLjI5MS4yOTZjMy43NSAyLjE3IDYuMiAyLjIxNiA4Ljg1NSA2LjU5MiA1LjA2OCA4LjMzMiAxLjQ1IDIwLjQzMS04LjU5MyAyMi40N2ExLjcxIDEuNzEgMCAwIDAtMS4zNjUgMS40NjhxLTkuMTczIDc1LjM4My02NC43NjIgMTIzLjM3MWMtMjguNjIyIDI0LjcwNC02NS43NDQgMzkuNzQ5LTEwMy40NjQgNDIuMzMtMTAwLjM1OSA2Ljg2My0xODQuNDg2LTY1LjYyMi0xOTUuMDA1LTE2NS41N2EuOC44IDAgMCAwLS4xNTQtLjM5NS43NC43NCAwIDAgMC0uMzMzLS4yNU0xNTEuMzM4IDk3LjY0M2wxOS4wNjYgMTkuMzc1YS41LjUgMCAwIDAgLjE1LjEuNS41IDAgMCAwIC4xNzYuMDMxLjQzLjQzIDAgMCAwIC4zMTktLjE1cTEwLjg0Ny0xMi40NjUgNy40NDMtMjguMTU1Yy00LjUzNS0yMC44Ni0xOC41OTgtMzcuNjY0LTM3LjI3MS00OC4zODktMjEuMDY3LTEyLjY4OS00Ny44NjUtMTUuOTA1LTcyLjI4LTE0Ljg0cS0uNjE3LjAzLS4xNC40MjJsNjMuNDcyIDUzLjMyNmEuNC40IDAgMCAwIC4xNDQuMDcyLjM0LjM0IDAgMCAwIC4zMDEtLjA2OC4zLjMgMCAwIDAgLjA4OC0uMTI2IDQ3IDQ3IDAgMCAwIDMuMjgyLTEzLjA4MXEuMDIzLS4yMzMuMTI0LS40MzN0LjI2OC0uMzM0YzIuOTU1LTIuNDc4IDQuMTE1LjQ1OCA0LjY5NCAyLjg1MnEyLjMyIDkuNTM3LjE4NyAxOS4yNDRhMS40OCAxLjQ4IDAgMCAwIC40NCAxLjQwM3ptLTYuNTI2IDkuMjk1LTkuMDI0LTEwLjI3NmEuNy43IDAgMCAwLS4zNDMtLjIwNy44Ni44NiAwIDAgMC0uNDMzIDBxLTExLjY4IDMuMDIxLTIzLjExNC0uODIyYy01LjgxNi0zLjg1Mi0uNzU4LTUuNTE3IDIuNTktNS44MTZxNS45MTgtLjUzMyAxMS43MzUtMi41MDZhLjMuMyAwIDAgMCAuMTI3LS4wOC4zLjMgMCAwIDAgLjA3My0uMTMzLjM0LjM0IDAgMCAwIC4wMDEtLjE1NC4zNC4zNCAwIDAgMC0uMDctLjEzOGMtMTkuNDU5LTIxLjE1LTQwLjc2LTQwLjI0NS02Mi4yNzUtNTkuMjY0cS0uMzU2LS4zMTgtLjMxOC4xNWEzMDggMzA4IDAgMCAwIDguMzg3IDUxLjgwMmM0LjU5MiAxOC4zMDggMTMuNzU1IDM0LjE1NyAzMS44NDggNDEuODcycTI2LjI0NyAxMS4xODIgNTMuNzI5IDIuODMzYS41ODMuNTgzIDAgMCAwIC4yOTktLjg4OCAyMzUgMjM1IDAgMCAwLTEzLjIxMi0xNi4zNzNtNDUuMTI1IDQ5Ljc0NXEtMTIuNTk1IDEuMzE4LTIzLjQ0MS01LjU4M2MtMy4wNzctNC44OTkgMi4wNjYtOC41OTMgNS40NDItMTEuMTgzYS41ODQuNTg0IDAgMCAwIC4xODctLjcwMSA4LjQ2IDguNDYgMCAwIDAtMy40ODgtMy44NjIuODEuODEgMCAwIDAtLjgyMy4wMDljLTE2LjY5MSA5Ljk3Ny0zNC40ODUgMTAuNzgyLTUzLjE4NiA2Ljk2N3EtMS41Ny0uMzItMi44NTIuNjQ1LTYxLjQ1MSA0Ni4zNy02OS42NyAxMjMuOTEzLS42ODMgNi40MzMtMS44MDYgOC44MjdjLTIuOTQ1IDYuMzAyLTcuNzcgNy44NDUtMTQuMjIyIDkuNDcycS0uNDQ5LjExMyAwIC4yMzRjOS4yNzYgMi4zMSAxNC42OCA2LjE0MyAxNi4wNzQgMTYuNDAxIDguNDYyIDYyLjQ5IDQ4Ljk0IDExNy44NjQgMTEwLjM4MyAxMzguMzQxYTEwNCAxMDQgMCAwIDAgMTEuMDA2IDMuNTA3YzkyLjY0NSAyMy45MDkgMTgxLjIzMi0zMy42NjIgMjAyLjgwNC0xMjUuMDU0cTEuNjA5LTYuNzkgMy40NS0xOS4wODVjMS4xMDQtNy4zOTYgNS42OTUtMTIuMTE4IDEzLjE0Ny0xMy41NjguNDEyLS4wODQgMS45MDgtLjc1Ny4yODEtLjk5MS04LjA1MS0xLjE0MS0xMi4yNC02LjM4Ni0xMi45ODgtMTQuMTI4cS0xLjA3NS0xMS4yNS0yLjUxNi0xOS4wMzhjLTExLjc3Mi02My44OTItNTYuMzgzLTExNC41NzItMTE5LjI4NS0xMzEuNDZxLTE1LjEyLTMuODgtMzAuOTg3LTUuMjQ1YS4zODYuMzg2IDAgMCAwLS40MDIuMjUyIDE2OC42IDE2OC42IDAgMCAxLTExLjM1MiAyNC44NTRjLTQuMTMzIDcuNDA2LTguMTYzIDE1LjQ5NC0xNS43NTYgMTYuNDc2Ii8+PGNpcmNsZSBjeD0iMTQwLjA2NSIgY3k9IjI3My4wMyIgcj0iNTAuNjE4IiBzdHJva2U9IiNmZmYiIHN0cm9rZS13aWR0aD0iOS4yMDMiLz48cGF0aCBmaWxsPSIjZmZmIiBkPSJNMTQ2Ljk2NyAyNDIuMDkyYzE1LjY3MiAwIDI4LjM3NiAxMy45NjMgMjguMzc2IDMxLjE4OHMtMTIuNzA0IDMxLjE4OS0yOC4zNzYgMzEuMTg5LTI4LjM3Ny0xMy45NjQtMjguMzc3LTMxLjE4OWMwLTIuNjI1LjI5Ny01LjE3NC44NTItNy42MSAxLjYwNyAzLjcyNyA1LjMxMyA2LjMzOCA5LjYyOSA2LjMzOCA1Ljc4OSAwIDEwLjQ4MS00LjY5MyAxMC40ODItMTAuNDgycy00LjY5My0xMC40ODEtMTAuNDgyLTEwLjQ4MWMtLjc2MSAwLTEuNTA0LjA4My0yLjIxOS4yMzcgNS4xMzktNS42NzYgMTIuMjUzLTkuMTkgMjAuMTE1LTkuMTkiLz48Y2lyY2xlIGN4PSIyNjkuOTMyIiBjeT0iMjczLjAzIiByPSI1MC42MTgiIHN0cm9rZT0iI2ZmZiIgc3Ryb2tlLXdpZHRoPSI5LjIwMyIvPjxwYXRoIGZpbGw9IiNmZmYiIGQ9Ik0yNzYuODM0IDI0Mi4wOTJjMTUuNjcyIDAgMjguMzc2IDEzLjk2MyAyOC4zNzYgMzEuMTg4cy0xMi43MDQgMzEuMTg5LTI4LjM3NiAzMS4xODktMjguMzc3LTEzLjk2NC0yOC4zNzctMzEuMTg5YzAtMi42MjUuMjk3LTUuMTc0Ljg1My03LjYxIDEuNjA2IDMuNzI3IDUuMzEyIDYuMzM4IDkuNjI4IDYuMzM4IDUuNzg5IDAgMTAuNDgyLTQuNjkzIDEwLjQ4Mi0xMC40ODJzLTQuNjkzLTEwLjQ4MS0xMC40ODItMTAuNDgxYy0uNzYxIDAtMS41MDQuMDgzLTIuMjE5LjIzNyA1LjEzOS01LjY3NiAxMi4yNTQtOS4xOSAyMC4xMTUtOS4xOSIvPjxwYXRoIHN0cm9rZT0iI2ZmZiIgc3Ryb2tlLXdpZHRoPSI5LjIwMyIgZD0iTTE3Ny4xMzUgMzQ1LjMyOGMuODk5LS4yNCAyLjgwMS0uMTk3IDYuMzA0LjQ2NCA2LjMxMSAxLjE5IDE2LjU4NSA0LjIyMSAyNi4xMjIgNC4yMjEgOS40NTIgMCAxNy4xMTgtMS45MjYgMjEuMzgzLTIuNjAzLjc3LS4xMjMgMS4zODUtLjE5OSAxLjg3OC0uMjM3LS4zNCAxMy44NDMtMTIuODkzIDI3LjcxLTI3Ljg2OCAyNy43MS0xNS4xNzIgMC0yNy44NzYtMTMuODM1LTI3Ljg3Ni0yOC44MTEgMC0uMzU3LjAyOS0uNTk2LjA1Ny0uNzQ0WiIvPjwvc3ZnPg==)](https://dashboard.pruna.ai/login?utm_source=huggingface&utm_medium=org_card&utm_campaign=hf_traffic)
+
+<div align="center">
+
+<h1 style="color: #9334E9;">⚡ PrunaVAED</h1>
+
+<h2>A drop-in replacement decoder for LTX-2.3</h2>
+
+<h3>
+  <span style="color: #9334E9;">1.7× faster</span>
+  &nbsp;·&nbsp;
+  <span style="color: #9334E9;">~50% lower peak VRAM</span>
+  &nbsp;·&nbsp;
+  Near-original visual quality
+</h3>
+
+</div>
+
+**PrunaVAED directly replaces the video VAE decoder in
+[`diffusers/LTX-2.3-Diffusers`](https://huggingface.co/diffusers/LTX-2.3-Diffusers).
+The encoder and latent format remain unchanged, making it a drop-in upgrade
+for faster, more memory-efficient LTX-2.3 decoding.**
+
+## Examples
+
+Both columns decode the same LTX-2.3 latent from the distilled pipeline. 
+
+<table>
+  <tr>
+    <th align="center">LTX-2.3 decoder</th>
+    <th align="center">⚡ PrunaVAED</th>
+  </tr>
+  <tr>
+    <td>
+      <video controls autoplay muted loop playsinline
+             src="https://huggingface.co/mehdiPrunaAI/PrunaVAED/resolve/main/example/teacher/04MhZs7YN08-Scene-0048.mp4">
+      </video>
+    </td>
+    <td>
+      <video controls autoplay muted loop playsinline
+             src="https://huggingface.co/mehdiPrunaAI/PrunaVAED/resolve/main/example/student/04MhZs7YN08-Scene-0048.mp4">
+      </video>
+    </td>
+  </tr>
+  <tr>
+    <td>
+      <video controls autoplay muted loop playsinline
+             src="https://huggingface.co/mehdiPrunaAI/PrunaVAED/resolve/main/example/teacher/FAfMSWi0FSA-Scene-0200.mp4">
+      </video>
+    </td>
+    <td>
+      <video controls autoplay muted loop playsinline
+             src="https://huggingface.co/mehdiPrunaAI/PrunaVAED/resolve/main/example/student/FAfMSWi0FSA-Scene-0200.mp4">
+      </video>
+    </td>
+  </tr>
+</table>
+
+<details>
+<summary><strong>More side-by-side comparisons</strong></summary>
+
+<br>
+
+<table>
+  <tr>
+    <th align="center">LTX-2.3 decoder</th>
+    <th align="center">⚡ PrunaVAED</th>
+  </tr>
+  <tr>
+    <td><video controls autoplay muted loop playsinline src="https://huggingface.co/mehdiPrunaAI/PrunaVAED/resolve/main/example/teacher/8YSha7iU2ZI-Scene-1025.mp4"></video></td>
+    <td><video controls autoplay muted loop playsinline src="https://huggingface.co/mehdiPrunaAI/PrunaVAED/resolve/main/example/student/8YSha7iU2ZI-Scene-1025.mp4"></video></td>
+  </tr>
+  <tr>
+    <td><video controls autoplay muted loop playsinline src="https://huggingface.co/mehdiPrunaAI/PrunaVAED/resolve/main/example/teacher/DZt-RtpcJZQ-Scene-0043_teacher.mp4"></video></td>
+    <td><video controls autoplay muted loop playsinline src="https://huggingface.co/mehdiPrunaAI/PrunaVAED/resolve/main/example/student/DZt-RtpcJZQ-Scene-0043_m2.mp4"></video></td>
+  </tr>
+  <tr>
+    <td><video controls autoplay muted loop playsinline src="https://huggingface.co/mehdiPrunaAI/PrunaVAED/resolve/main/example/teacher/9GDE5-rOOfk-Scene-0010.mp4"></video></td>
+    <td><video controls autoplay muted loop playsinline src="https://huggingface.co/mehdiPrunaAI/PrunaVAED/resolve/main/example/student/9GDE5-rOOfk-Scene-0010.mp4"></video></td>
+  </tr>
+  <tr>
+    <td><video controls autoplay muted loop playsinline src="https://huggingface.co/mehdiPrunaAI/PrunaVAED/resolve/main/example/teacher/7y57zc-GvBA-Scene-0009.mp4"></video></td>
+    <td><video controls autoplay muted loop playsinline src="https://huggingface.co/mehdiPrunaAI/PrunaVAED/resolve/main/example/student/7y57zc-GvBA-Scene-0009.mp4"></video></td>
+  </tr>
+</table>
+
+</details>
+
+
+## Benchmark
+
+Metrics compare videos decoded from the **same latents** by the LTX-2.3 VAED and PrunaVAED. Two set of latents were
+generated with the full
+[`ti2vid_two_stages`](https://github.com/Lightricks/LTX-2/blob/main/packages/ltx-pipelines/src/ltx_pipelines/ti2vid_two_stages.py)
+pipeline and the
+[`distilled`](https://github.com/Lightricks/LTX-2/blob/main/packages/ltx-pipelines/src/ltx_pipelines/distilled.py)
+pipeline (recommended). (~4 s @ 24 fps)
+
+| Setting | Value |
+|---|---|
+| Model | PrunaVAED |
+| Precision | `bfloat16` |
+| Batch size | 1 |
+| Device | NVIDIA H100 80GB |
+| Decode | Full (no tiling / rolling / `torch.compile`) |
+
+| Latent set | Clips | PSNR ↑ | LPIPS ↓ | SSIM ↑ | Δ-frame PSNR ↑ |
+|---|---:|---:|---:|---:|---:|
+| Distilled two-stage  · 720p | 200 | **37.87** | **0.0108** | **0.9898** | **37.28** |
+| TI2Vid two-stage · 720p | 200 | **38.86** | **0.0123** | **0.9803** | **38.98** |
+| Distilled two-stage · 1080p | 200 | **39.59** | **0.0067** | **0.9927** | **39.37** |
+ 
+ **Δ-frame PSNR** is PSNR on consecutive-frame
+differences (temporal consistency).
+
+## Test
+
+End-to-end smoke test: generate a short video with the diffusers LTX-2.3
+**distilled** two-stage pipeline, then decode the same latent with the stock
+LTX-2.3 VAE and with PrunaVAED. Writes two mp4s and prints decode time (ms).
+Needs a **CUDA GPU**.
+
+```bash
+# 1. Install the Hugging Face CLI
+pip install hf
+
+# 2. Download this repo
+hf download PrunaAI/LTX23-PrunedDecoder --local-dir prunaVAED
+cd prunaVAED
+
+# 3. Install dependencies
+pip install -r requirements-test.txt
+
+# 4. Run the test (~720p, ~5 s @ 24 fps)
+python tests/test_distilled_decode.py
+```
+
+Outputs land in `outputs/test_distilled/` (`ltx23.mp4`, `prunavaed.mp4`).
+Edit `PROMPT` / resolution at the top of `tests/test_distilled_decode.py` if needed.
+
+
+## Architecture
+
+PrunaVAED is bitwise identical to LTX 2.3 VAED from `conv_in` to `up_blocks.0`. Pruning starts
+at `up_blocks.1`.
+
+| Stage | Channel reduction |
+|---|---:|
+| `up_blocks.0` | unchanged |
+| `up_blocks.1` | 25% |
+| `up_blocks.2` | 50% |
+| `up_blocks.3` | 50% |
+
+| | Teacher (LTX-2.3) | PrunaVAED |
+|---|---:|---:|
+| Encoder params | 318.9 M | 318.9 M (unchanged) |
+| Decoder params | 407.2 M | 345.0 M (−15%) |
+| **Total VAE params** | **726.1 M** | **663.9 M** |
+
+`prunavaed/patch_diffusers.py` provides the required compatibility shim for
+the pinned diffusers version.
+
+## Limitations
+
+- Evaluated on 4-second clips at 24 fps.
+- Benchmarked on one H100 80GB with bfloat16 and batch size 1.
+- Speed and VRAM vary with hardware, resolution, batch size, and software.
+- Evaluated only on two differents pipeline of LTX 2.3 ([`ti2vid_two_stages`](https://github.com/Lightricks/LTX-2/blob/main/packages/ltx-pipelines/src/ltx_pipelines/ti2vid_two_stages.py)
+pipeline and the
+[`distilled`](https://github.com/Lightricks/LTX-2/blob/main/packages/ltx-pipelines/src/ltx_pipelines/distilled.py))
+- The weights can be adapted to the custom [`LTX-2`](https://github.com/Lightricks/LTX-2) library without retraining, but no implementation is provided for this yet.
+
+## What this is not
+
+- Not a full T2V model.
+- Not a replacement for the LTX denoiser/generator.
+- Not compatible with arbitrary VAE latents.
+- Not bit-exact with the original decoder.
+
+## License
+
+PrunaVAED is a derivative of LTX-2.3 and is distributed under the
+**[LTX-2 Community License Agreement](LICENSE)**. Review its use restrictions
+and commercial terms before using or redistributing the model.
+
+Helper code adapted from Hugging Face diffusers retains its Apache-2.0
+attribution; see [`NOTICE`](NOTICE).
+
+## Try next
+
+- **Use PrunaVAED to speed up LTX-2.3 decoding.**
+- Try Pruna's end-to-end video models:
+  [P-Video](https://www.pruna.ai/p-video) and
+  [P-Video documentation](https://docs.pruna.ai/en/stable/docs_pruna_endpoints/performance_models/p-video.html).

example/student/04MhZs7YN08-Scene-0048.mp4

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:43d29de4e52f013f14a82d1f4d41d78f7bb938184f30af48d2e53d09efffe558
+size 1797361

example/student/7y57zc-GvBA-Scene-0009.mp4

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:2de528d80672faf55079399ddf14533ef77b274099407d35fec5fec552b89ae3
+size 1716725

example/student/8YSha7iU2ZI-Scene-1025.mp4

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

example/student/97moBezrzLo-Scene-0076.mp4

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

example/student/9GDE5-rOOfk-Scene-0010.mp4

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:845ad0a7d35633fe3ddf73da94534d468eb2ac36f6faa450db6c841439f785dd
+size 2308198

example/student/DZt-RtpcJZQ-Scene-0043_m2.mp4

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:06f56792ff2d4793d12b4283c20788e4be6c22ea876729899b7e46614016c1f6
+size 1214808

example/student/FAfMSWi0FSA-Scene-0200.mp4

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:9442659a2cb5b6328bb8f0849b466a4e6b572243f1350efa93cc2cb70261ec10
+size 1883423

example/student/a5M8hWgQbjU-Scene-0191.mp4

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

example/student/ijGuAu0lLNM-Scene-0008.mp4

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:4eb2af16643f4402da2585878e41257e951141c491556e335b72abab3a9f26eb
+size 1697225

example/student/zFehFw9SbDw-Scene-0034.mp4

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

example/teacher/04MhZs7YN08-Scene-0048.mp4

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:320db36560a3b9abac5c710bed668c0b25f8bc3b9cdd1935ba65ecec0ee2174f
+size 1697633

example/teacher/7y57zc-GvBA-Scene-0009.mp4

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:004434c148e393a6be129a5682fd357b9340f035595d3f09348d764e487dd1e8
+size 1697114

example/teacher/8YSha7iU2ZI-Scene-1025.mp4

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

example/teacher/97moBezrzLo-Scene-0076.mp4

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

example/teacher/9GDE5-rOOfk-Scene-0010.mp4

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

example/teacher/DZt-RtpcJZQ-Scene-0043_teacher.mp4

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:51facb9d24858bd48cdf7e565c4da9c863e875c9121e83b050233c31a59826b6
+size 1174232

example/teacher/FAfMSWi0FSA-Scene-0200.mp4

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:505d374a8cc0caf0378903199684b2d0db7326ccac7c8f186d9b2eb4690f32ae
+size 1864634

example/teacher/a5M8hWgQbjU-Scene-0191.mp4

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

example/teacher/ijGuAu0lLNM-Scene-0008.mp4

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

example/teacher/zFehFw9SbDw-Scene-0034.mp4

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:688f0d785cb9dafa858a35ac5192380df50483b81392dc1690ad32393ce536d2
+size 2125596

patch_diffusers.py

@@ -0,0 +1,219 @@
+# Copyright 2026 Pruna AI
+# Portions adapted from Hugging Face diffusers:
+# Copyright 2026 The HuggingFace Team. Licensed under Apache-2.0.
+
+"""Monkey-patch stock diffusers so pruned LTX-2.3 decoder weights load correctly.
+
+Stock diffusers builds the decoder graph from ``decoder_block_out_channels`` using
+the *nominal* width of each up block (divided by ``upsample_factor``). It assumes the
+tensor leaving block N always matches that nominal width.
+
+PrunaVAED keeps wider skip connections between blocks and only prunes inside each
+block. Example on ``up_blocks.2``:
+
+- tensor arriving from ``up_blocks.1``: **384** ch (after 512→384 projection)
+- ResNets inside ``up_blocks.2``: **128** ch (50% pruned vs LTX-2.3 decoder 256)
+- ``conv_in`` projection inside the block: **384 → 256** (feeds the upsampler)
+
+Stock diffusers wires ``up_blocks.2`` as 192 → 128 and never creates the 384→256
+``conv_in``, so ``from_pretrained`` fails with shape mismatches on PrunaVAED weights.
+
+Fix (two ``__init__`` replacements; forward/decode unchanged)
+------------------------------------------------------------
+1. ``LTX2VideoDecoder3d``: track ``current_channels`` — the actual tensor width
+   between blocks — instead of recomputing from the previous block's nominal width.
+2. ``LTX2VideoUpBlock3d``: insert ``conv_in`` when ``in_channels != out_channels *
+   upscale_factor`` (pre-upsampler width), not when ``in_channels != out_channels``.
+
+The LTX-2.3 decoder is unaffected (no extra projections). Safe to call before any
+``AutoencoderKLLTX2Video.from_pretrained``.
+"""
+
+from __future__ import annotations
+
+import torch
+import torch.nn as nn
+from diffusers.models.autoencoders import autoencoder_kl_ltx2 as ltx2
+
+
+def patch_pruna_ltx2_decoder() -> None:
+    if getattr(ltx2, "_PRUNA_LTX2_DECODER_PATCH", False):
+        return
+
+    Resnet = ltx2.LTX2VideoResnetBlock3d
+    Upsampler = ltx2.LTX2VideoUpsampler3d
+    CausalConv = ltx2.LTX2VideoCausalConv3d
+    MidBlock = ltx2.LTX2VideoMidBlock3d
+    UpBlock = ltx2.LTX2VideoUpBlock3d
+    TimeEmb = ltx2.PixArtAlphaCombinedTimestepSizeEmbeddings
+    RMSNorm = ltx2.PerChannelRMSNorm
+
+    def up_init(
+        self,
+        in_channels,
+        out_channels=None,
+        num_layers=1,
+        dropout=0.0,
+        resnet_eps=1e-6,
+        resnet_act_fn="swish",
+        spatio_temporal_scale=True,
+        upsample_type="spatiotemporal",
+        inject_noise=False,
+        timestep_conditioning=False,
+        upsample_residual=False,
+        upscale_factor=1,
+        spatial_padding_mode="zeros",
+    ):
+        nn.Module.__init__(self)
+        out_channels = out_channels or in_channels
+        # Width right before the upsampler (ResNet width × upscale_factor).
+        # conv_in projects the incoming skip (e.g. 384) down to this width (e.g. 256).
+        pre = out_channels * upscale_factor
+        self.time_embedder = TimeEmb(in_channels * 4, 0) if timestep_conditioning else None
+        self.conv_in = None
+        if in_channels != pre:  # stock diffusers compares in_channels != out_channels
+            self.conv_in = Resnet(
+                in_channels=in_channels,
+                out_channels=pre,
+                dropout=dropout,
+                eps=resnet_eps,
+                non_linearity=resnet_act_fn,
+                inject_noise=inject_noise,
+                timestep_conditioning=timestep_conditioning,
+                spatial_padding_mode=spatial_padding_mode,
+            )
+        self.upsamplers = None
+        if spatio_temporal_scale:
+            stride = {
+                "spatial": (1, 2, 2),
+                "temporal": (2, 1, 1),
+                "spatiotemporal": (2, 2, 2),
+            }[upsample_type]
+            self.upsamplers = nn.ModuleList(
+                [
+                    Upsampler(
+                        in_channels=pre,
+                        stride=stride,
+                        residual=upsample_residual,
+                        upscale_factor=upscale_factor,
+                        spatial_padding_mode=spatial_padding_mode,
+                    ),
+                ]
+            )
+        self.resnets = nn.ModuleList(
+            [
+                Resnet(
+                    in_channels=out_channels,
+                    out_channels=out_channels,
+                    dropout=dropout,
+                    eps=resnet_eps,
+                    non_linearity=resnet_act_fn,
+                    inject_noise=inject_noise,
+                    timestep_conditioning=timestep_conditioning,
+                    spatial_padding_mode=spatial_padding_mode,
+                )
+                for _ in range(num_layers)
+            ]
+        )
+        self.gradient_checkpointing = False
+
+    def decoder_init(
+        self,
+        in_channels=128,
+        out_channels=3,
+        block_out_channels=(256, 512, 1024),
+        spatio_temporal_scaling=(True, True, True),
+        layers_per_block=(5, 5, 5, 5),
+        upsample_type=("spatiotemporal", "spatiotemporal", "spatiotemporal"),
+        patch_size=4,
+        patch_size_t=1,
+        resnet_norm_eps=1e-6,
+        is_causal=False,
+        inject_noise=(False, False, False),
+        timestep_conditioning=False,
+        upsample_residual=(True, True, True),
+        upsample_factor=(2, 2, 2),
+        spatial_padding_mode="reflect",
+    ):
+        nn.Module.__init__(self)
+        n = len(layers_per_block)
+        if isinstance(spatio_temporal_scaling, bool):
+            spatio_temporal_scaling = (spatio_temporal_scaling,) * (n - 1)
+        if isinstance(inject_noise, bool):
+            inject_noise = (inject_noise,) * n
+        if isinstance(upsample_residual, bool):
+            upsample_residual = (upsample_residual,) * (n - 1)
+
+        self.patch_size, self.patch_size_t = patch_size, patch_size_t
+        self.out_channels = out_channels * patch_size**2
+        self.is_causal = is_causal
+
+        ch = tuple(reversed(block_out_channels))
+        spatio_temporal_scaling = tuple(reversed(spatio_temporal_scaling))
+        layers_per_block = tuple(reversed(layers_per_block))
+        inject_noise = tuple(reversed(inject_noise))
+        upsample_residual = tuple(reversed(upsample_residual))
+        upsample_factor = tuple(reversed(upsample_factor))
+        upsample_type = tuple(upsample_type)
+        if len(upsample_type) == len(ch) - 1:
+            upsample_type = tuple(reversed(upsample_type))
+
+        width = ch[0]
+        self.conv_in = CausalConv(
+            in_channels=in_channels,
+            out_channels=width,
+            kernel_size=3,
+            stride=1,
+            spatial_padding_mode=spatial_padding_mode,
+        )
+        self.mid_block = MidBlock(
+            in_channels=width,
+            num_layers=layers_per_block[0],
+            resnet_eps=resnet_norm_eps,
+            inject_noise=inject_noise[0],
+            timestep_conditioning=timestep_conditioning,
+            spatial_padding_mode=spatial_padding_mode,
+        )
+
+        # Stock: input_channel = previous_nominal // factor (loses pruned skip widths).
+        # Patched: pass the real tensor width from the previous block as in_channels.
+        self.up_blocks = nn.ModuleList()
+        current = width
+        for i in range(len(ch)):
+            resnet_w = ch[i] // upsample_factor[i]
+            self.up_blocks.append(
+                UpBlock(
+                    in_channels=current,
+                    out_channels=resnet_w,
+                    num_layers=layers_per_block[i + 1],
+                    resnet_eps=resnet_norm_eps,
+                    spatio_temporal_scale=spatio_temporal_scaling[i],
+                    upsample_type=upsample_type[i],
+                    inject_noise=inject_noise[i + 1],
+                    timestep_conditioning=timestep_conditioning,
+                    upsample_residual=upsample_residual[i],
+                    upscale_factor=upsample_factor[i],
+                    spatial_padding_mode=spatial_padding_mode,
+                )
+            )
+            current = resnet_w
+
+        self.norm_out = RMSNorm()
+        self.conv_act = nn.SiLU()
+        self.conv_out = CausalConv(
+            in_channels=current,
+            out_channels=self.out_channels,
+            kernel_size=3,
+            stride=1,
+            spatial_padding_mode=spatial_padding_mode,
+        )
+        self.time_embedder = self.scale_shift_table = self.timestep_scale_multiplier = None
+        if timestep_conditioning:
+            self.timestep_scale_multiplier = nn.Parameter(torch.tensor(1000.0))
+            self.time_embedder = TimeEmb(width * 2, 0)
+            self.scale_shift_table = nn.Parameter(torch.randn(2, width) / width**0.5)
+        self.gradient_checkpointing = False
+
+    UpBlock.__init__ = up_init
+    ltx2.LTX2VideoDecoder3d.__init__ = decoder_init
+    ltx2._PRUNA_LTX2_DECODER_PATCH = True

requirements-test.txt

@@ -0,0 +1,19 @@
+--extra-index-url https://download.pytorch.org/whl/cu130
+torch==2.11.0+cu130
+torchvision==0.26.0+cu130
+torchaudio==2.11.0+cu130
+diffusers==0.39.0
+huggingface_hub==1.24.0
+safetensors==0.8.0
+accelerate==1.13.0
+imageio==2.37.3
+imageio-ffmpeg==0.6.0
+
+transformers>=4.57.0
+sentencepiece>=0.2.0
+protobuf>=5.0.0
+Pillow>=10.0.0
+numpy>=1.24.0
+tqdm>=4.66.0
+av>=14.0.0
+pytest>=8.0.0

tests/test_distilled_decode.py

@@ -0,0 +1,223 @@
+#!/usr/bin/env python3
+"""Quick demo: generate a short LTX-2.3 video latent, then decode it twice.
+
+Compares the stock LTX-2.3 VAE decoder with this repo's pruned decoder
+(PrunaVAED) on the *same* latent. Prints decode time (ms) and writes two mp4s.
+
+Requires a CUDA GPU. From the repo root:
+
+    pip install -r requirements-test.txt
+    python tests/test_distilled_decode.py
+"""
+
+from __future__ import annotations
+
+import statistics
+import sys
+import time
+from pathlib import Path
+
+import imageio.v3 as iio
+import torch
+from diffusers import LTX2LatentUpsamplePipeline, LTX2Pipeline
+from diffusers.models.autoencoders import AutoencoderKLLTX2Video
+from diffusers.pipelines.ltx2.latent_upsampler import LTX2LatentUpsamplerModel
+from diffusers.pipelines.ltx2.utils import (
+    DEFAULT_NEGATIVE_PROMPT,
+    DISTILLED_SIGMA_VALUES,
+    STAGE_2_DISTILLED_SIGMA_VALUES,
+)
+
+torch.backends.cuda.enable_cudnn_sdp(False)
+
+# ---------------------------------------------------------------------------
+# Settings (edit these if you want)
+# ---------------------------------------------------------------------------
+
+REPO_ROOT = Path(__file__).resolve().parents[1]
+sys.path.insert(0, str(REPO_ROOT))
+
+from patch_diffusers import patch_pruna_ltx2_decoder  # noqa: E402
+
+# Official Diffusers checkpoints for the 2-stage distilled recipe.
+DISTILLED_MODEL = "diffusers/LTX-2.3-Distilled-Diffusers"
+SPATIAL_UPSAMPLER = "dg845/LTX-2.3-Spatial-Upsampler-Diffusers"
+LTX23_VAE = "diffusers/LTX-2.3-Diffusers"  # stock decoder (baseline)
+PRUNED_VAE = REPO_ROOT  # this repo's vae/ folder
+
+# ~720p for 5 s @ 24 fps. Height/width must be multiples of 64 (2-stage).
+HEIGHT, WIDTH, NUM_FRAMES, FPS = 704, 1280, 121, 24.0
+SEED = 42
+DECODE_WARMUP, DECODE_RUNS = 1, 3  # timing: 1 warm-up + median of 3 runs
+
+PROMPT = (
+    "In the video, a man is seen riding a skateboard in an indoor skate park. "
+    "He is wearing a black cap, a grey t-shirt, and black pants. The skate park "
+    "has a wooden floor and a ramp. The man performs a trick by jumping and "
+    "flipping the skateboard under his feet while in the air. He lands on the "
+    "skateboard and continues to ride it."
+)
+
+DEVICE = "cuda"
+DTYPE = torch.bfloat16
+OUTPUT_DIR = REPO_ROOT / "outputs" / "test_distilled"
+
+
+# ---------------------------------------------------------------------------
+# Helpers
+# ---------------------------------------------------------------------------
+
+def save_mp4(video: torch.Tensor, path: Path) -> None:
+    """Save a BCTHW tensor in [-1, 1] as an H.264 mp4."""
+    frames = video[0].permute(1, 2, 3, 0).clamp(-1, 1)
+    frames = ((frames + 1) / 2 * 255).byte().cpu().numpy()
+    path.parent.mkdir(parents=True, exist_ok=True)
+    iio.imwrite(path, frames, fps=FPS, codec="libx264")
+
+
+def load_vae(model_id: str) -> AutoencoderKLLTX2Video:
+    """Load a video VAE decoder and put it on GPU."""
+    vae = AutoencoderKLLTX2Video.from_pretrained(
+        model_id, subfolder="vae", torch_dtype=DTYPE
+    )
+    vae = vae.to(DEVICE).eval()
+    vae.use_tiling = False
+    if hasattr(vae, "disable_tiling"):
+        vae.disable_tiling()
+    return vae
+
+
+@torch.inference_mode()
+def timed_decode(vae: AutoencoderKLLTX2Video, latent: torch.Tensor) -> tuple[torch.Tensor, float]:
+    """Decode once after warm-up; return (video on CPU, median latency in ms)."""
+    latent = latent.to(DEVICE, DTYPE)
+
+    with torch.autocast(DEVICE, dtype=DTYPE):
+        for _ in range(DECODE_WARMUP):
+            vae.decode(latent, return_dict=False)
+
+    times_ms = []
+    video = None
+    for _ in range(DECODE_RUNS):
+        torch.cuda.synchronize()
+        t0 = time.perf_counter()
+        with torch.autocast(DEVICE, dtype=DTYPE):
+            video = vae.decode(latent, return_dict=False)[0]
+        torch.cuda.synchronize()
+        times_ms.append((time.perf_counter() - t0) * 1000)
+
+    return video.cpu(), statistics.median(times_ms)
+
+
+@torch.inference_mode()
+def generate_latent(prompt: str) -> torch.Tensor:
+    """Run the official 2-stage distilled pipeline and return the final video latent.
+
+    Same layout as Lightricks DistilledPipeline:
+      stage 1 @ half-res → 2× latent upsample → stage 2 @ full-res.
+    """
+    half_h, half_w = HEIGHT // 2, WIDTH // 2
+    generator = torch.Generator(DEVICE).manual_seed(SEED)
+
+    pipe = LTX2Pipeline.from_pretrained(DISTILLED_MODEL, torch_dtype=DTYPE)
+    pipe.enable_model_cpu_offload(device=DEVICE)
+
+    # Stage 1 — cheap draft at half resolution.
+    print(f"1/3  Stage 1 @ {half_w}×{half_h}")
+    video_latent, audio_latent = pipe(
+        prompt=prompt,
+        negative_prompt=DEFAULT_NEGATIVE_PROMPT,
+        height=half_h,
+        width=half_w,
+        num_frames=NUM_FRAMES,
+        frame_rate=FPS,
+        num_inference_steps=len(DISTILLED_SIGMA_VALUES),
+        sigmas=DISTILLED_SIGMA_VALUES,
+        guidance_scale=1.0,
+        generator=generator,
+        output_type="latent",
+        return_dict=False,
+    )
+
+    # Upsample — bring the latent to full resolution before refining.
+    print(f"2/3  Upsample → {WIDTH}×{HEIGHT}")
+    upsampler = LTX2LatentUpsamplerModel.from_pretrained(
+        SPATIAL_UPSAMPLER, subfolder="latent_upsampler", torch_dtype=DTYPE
+    )
+    upsample_pipe = LTX2LatentUpsamplePipeline(vae=pipe.vae, latent_upsampler=upsampler)
+    upsample_pipe.enable_model_cpu_offload(device=DEVICE)
+    video_latent = upsample_pipe(
+        latents=video_latent[:1],
+        height=half_h,
+        width=half_w,
+        num_frames=NUM_FRAMES,
+        output_type="latent",
+        return_dict=False,
+    )[0]
+    del upsample_pipe, upsampler
+    torch.cuda.empty_cache()
+
+    # Stage 2 — refine at full resolution (renoises from STAGE_2 schedule).
+    print(f"3/3  Stage 2 @ {WIDTH}×{HEIGHT}")
+    video_latent, _ = pipe(
+        latents=video_latent,
+        audio_latents=audio_latent,
+        prompt=prompt,
+        negative_prompt=DEFAULT_NEGATIVE_PROMPT,
+        height=HEIGHT,
+        width=WIDTH,
+        num_frames=NUM_FRAMES,
+        frame_rate=FPS,
+        num_inference_steps=len(STAGE_2_DISTILLED_SIGMA_VALUES),
+        noise_scale=STAGE_2_DISTILLED_SIGMA_VALUES[0],
+        sigmas=STAGE_2_DISTILLED_SIGMA_VALUES,
+        guidance_scale=1.0,
+        generator=generator,
+        output_type="latent",
+        return_dict=False,
+    )
+
+    latent = video_latent.detach().cpu()
+    del pipe
+    torch.cuda.empty_cache()
+    return latent
+
+
+def compare_decoders(latent: torch.Tensor, out_dir: Path) -> None:
+    """Decode the same latent with LTX-2.3 and PrunaVAED; print ms and save mp4s."""
+    # Needed so Diffusers can build the pruned decoder graph correctly.
+    patch_pruna_ltx2_decoder()
+
+    results = {}
+    for name, model_id in (("ltx23", LTX23_VAE), ("prunavaed", str(PRUNED_VAE))):
+        print(f"Decoding with {name} …")
+        vae = load_vae(model_id)
+        video, ms = timed_decode(vae, latent)
+        results[name] = ms
+        path = out_dir / f"{name}.mp4"
+        save_mp4(video, path)
+        print(f"  {name}: {ms:.1f} ms  →  {path}")
+        del vae, video
+        torch.cuda.empty_cache()
+
+    print(f"Speedup: {results['ltx23'] / results['prunavaed']:.2f}×")
+
+
+# ---------------------------------------------------------------------------
+# Entry point
+# ---------------------------------------------------------------------------
+
+def main() -> None:
+    if not torch.cuda.is_available():
+        raise SystemExit("This demo needs a CUDA GPU.")
+
+    print(f"Prompt: {PROMPT[:80]}…")
+    print(f"Output: {OUTPUT_DIR}")
+
+    latent = generate_latent(PROMPT)
+    print(f"Latent shape: {tuple(latent.shape)}")
+    compare_decoders(latent, OUTPUT_DIR)
+
+
+if __name__ == "__main__":
+    main()

vae/config.json

@@ -0,0 +1,96 @@
+{
+  "_class_name": "AutoencoderKLLTX2Video",
+  "_diffusers_version": "0.38.0",
+  "aligned_blks_indices": null,
+  "aligned_feature_projection_dim": null,
+  "aligned_feature_projection_mode": null,
+  "block_out_channels": [
+    256,
+    512,
+    1024,
+    1024
+  ],
+  "decoder_block_out_channels": [
+    128,
+    256,
+    384,
+    1024
+  ],
+  "decoder_causal": false,
+  "decoder_inject_noise": [
+    false,
+    false,
+    false,
+    false,
+    false
+  ],
+  "decoder_layers_per_block": [
+    4,
+    6,
+    4,
+    2,
+    2
+  ],
+  "decoder_spatial_padding_mode": "zeros",
+  "decoder_spatio_temporal_scaling": [
+    true,
+    true,
+    true,
+    true
+  ],
+  "down_block_types": [
+    "LTX2VideoDownBlock3D",
+    "LTX2VideoDownBlock3D",
+    "LTX2VideoDownBlock3D",
+    "LTX2VideoDownBlock3D"
+  ],
+  "downsample_type": [
+    "spatial",
+    "temporal",
+    "spatiotemporal",
+    "spatiotemporal"
+  ],
+  "encoder_causal": true,
+  "encoder_spatial_padding_mode": "zeros",
+  "in_channels": 3,
+  "latent_channels": 128,
+  "layers_per_block": [
+    4,
+    6,
+    4,
+    2,
+    2
+  ],
+  "out_channels": 3,
+  "patch_size": 4,
+  "patch_size_t": 1,
+  "resnet_norm_eps": 1e-06,
+  "scaling_factor": 1.0,
+  "spatial_compression_ratio": 32,
+  "spatio_temporal_scaling": [
+    true,
+    true,
+    true,
+    true
+  ],
+  "temporal_compression_ratio": 8,
+  "timestep_conditioning": false,
+  "upsample_factor": [
+    2,
+    2,
+    1,
+    2
+  ],
+  "upsample_residual": [
+    false,
+    false,
+    false,
+    false
+  ],
+  "upsample_type": [
+    "spatiotemporal",
+    "spatiotemporal",
+    "temporal",
+    "spatial"
+  ]
+}

vae/diffusion_pytorch_model.safetensors

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