Instructions to use suryatmodulus/PrunaVAED with libraries, inference providers, notebooks, and local apps. Follow these links to get started.
- Libraries
- Diffusers
How to use suryatmodulus/PrunaVAED with Diffusers:
pip install -U diffusers transformers accelerate
import torch from diffusers import DiffusionPipeline # switch to "mps" for apple devices pipe = DiffusionPipeline.from_pretrained("suryatmodulus/PrunaVAED", torch_dtype=torch.bfloat16, device_map="cuda") prompt = "Astronaut in a jungle, cold color palette, muted colors, detailed, 8k" image = pipe(prompt).images[0] - LTX.io
How to use suryatmodulus/PrunaVAED with LTX.io:
# Install the LTX-2 pipelines git clone https://github.com/Lightricks/LTX-2.git cd LTX-2 uv sync --frozen
# Download the weights from this repo, plus the Gemma text encoder hf download suryatmodulus/PrunaVAED --local-dir models/PrunaVAED hf download google/gemma-3-12b-it-qat-q4_0-unquantized --local-dir models/gemma-3-12b
# Fast pipeline (distilled model, no distilled LoRA needed) uv run python -m ltx_pipelines.distilled \ --distilled-checkpoint-path models/PrunaVAED/<distilled-checkpoint>.safetensors \ --spatial-upsampler-path models/PrunaVAED/<spatial-upsampler>.safetensors \ --gemma-root models/gemma-3-12b \ --prompt "A beautiful sunset over the ocean" \ --output-path output.mp4 # For image-to-video, add: --image path/to/image.jpg 0 0.8# HQ pipeline (two-stage, higher quality) uv run python -m ltx_pipelines.ti2vid_two_stages_hq \ --checkpoint-path models/PrunaVAED/<checkpoint>.safetensors \ --distilled-lora models/PrunaVAED/<distilled-lora>.safetensors 0.8 \ --spatial-upsampler-path models/PrunaVAED/<spatial-upsampler>.safetensors \ --gemma-root models/gemma-3-12b \ --prompt "A beautiful sunset over the ocean" \ --output-path output.mp4 # For image-to-video, add: --image path/to/image.jpg 0 0.8 - Notebooks
- Google Colab
- Kaggle
Commit ·
93293db
0
Parent(s):
Duplicate from PrunaAI/PrunaVAED
Browse filesCo-authored-by: Mehdi SI-MOHAMMED <mehdiPrunaAI@users.noreply.huggingface.co>
- .gitattributes +75 -0
- README.md +289 -0
- demo/demo_distilled_decode.py +227 -0
- example/Ltx23VAED/04MhZs7YN08-Scene-0048.mp4 +3 -0
- example/Ltx23VAED/7y57zc-GvBA-Scene-0009.mp4 +3 -0
- example/Ltx23VAED/8YSha7iU2ZI-Scene-1025.mp4 +3 -0
- example/Ltx23VAED/97moBezrzLo-Scene-0076.mp4 +3 -0
- example/Ltx23VAED/9GDE5-rOOfk-Scene-0010.mp4 +3 -0
- example/Ltx23VAED/DZt-RtpcJZQ-Scene-0043.mp4 +3 -0
- example/Ltx23VAED/FAfMSWi0FSA-Scene-0200.mp4 +3 -0
- example/Ltx23VAED/a5M8hWgQbjU-Scene-0191.mp4 +3 -0
- example/Ltx23VAED/ijGuAu0lLNM-Scene-0008.mp4 +3 -0
- example/Ltx23VAED/zFehFw9SbDw-Scene-0034.mp4 +3 -0
- example/PrunaVAED/04MhZs7YN08-Scene-0048.mp4 +3 -0
- example/PrunaVAED/7y57zc-GvBA-Scene-0009.mp4 +3 -0
- example/PrunaVAED/8YSha7iU2ZI-Scene-1025.mp4 +3 -0
- example/PrunaVAED/97moBezrzLo-Scene-0076.mp4 +3 -0
- example/PrunaVAED/9GDE5-rOOfk-Scene-0010.mp4 +3 -0
- example/PrunaVAED/DZt-RtpcJZQ-Scene-0043.mp4 +3 -0
- example/PrunaVAED/FAfMSWi0FSA-Scene-0200.mp4 +3 -0
- example/PrunaVAED/a5M8hWgQbjU-Scene-0191.mp4 +3 -0
- example/PrunaVAED/ijGuAu0lLNM-Scene-0008.mp4 +3 -0
- example/PrunaVAED/zFehFw9SbDw-Scene-0034.mp4 +3 -0
- example/prompts.json +44 -0
- patch_diffusers.py +219 -0
- requirements-demo.txt +19 -0
- vae/config.json +93 -0
- vae/diffusion_pytorch_model.safetensors +3 -0
.gitattributes
ADDED
|
@@ -0,0 +1,75 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
*.7z filter=lfs diff=lfs merge=lfs -text
|
| 2 |
+
*.arrow filter=lfs diff=lfs merge=lfs -text
|
| 3 |
+
*.bin filter=lfs diff=lfs merge=lfs -text
|
| 4 |
+
*.bz2 filter=lfs diff=lfs merge=lfs -text
|
| 5 |
+
*.ckpt filter=lfs diff=lfs merge=lfs -text
|
| 6 |
+
*.ftz filter=lfs diff=lfs merge=lfs -text
|
| 7 |
+
*.gz filter=lfs diff=lfs merge=lfs -text
|
| 8 |
+
*.h5 filter=lfs diff=lfs merge=lfs -text
|
| 9 |
+
*.joblib filter=lfs diff=lfs merge=lfs -text
|
| 10 |
+
*.lfs.* filter=lfs diff=lfs merge=lfs -text
|
| 11 |
+
*.mlmodel filter=lfs diff=lfs merge=lfs -text
|
| 12 |
+
*.model filter=lfs diff=lfs merge=lfs -text
|
| 13 |
+
*.msgpack filter=lfs diff=lfs merge=lfs -text
|
| 14 |
+
*.npy filter=lfs diff=lfs merge=lfs -text
|
| 15 |
+
*.npz filter=lfs diff=lfs merge=lfs -text
|
| 16 |
+
*.onnx filter=lfs diff=lfs merge=lfs -text
|
| 17 |
+
*.ot filter=lfs diff=lfs merge=lfs -text
|
| 18 |
+
*.parquet filter=lfs diff=lfs merge=lfs -text
|
| 19 |
+
*.pb filter=lfs diff=lfs merge=lfs -text
|
| 20 |
+
*.pickle filter=lfs diff=lfs merge=lfs -text
|
| 21 |
+
*.pkl filter=lfs diff=lfs merge=lfs -text
|
| 22 |
+
*.pt filter=lfs diff=lfs merge=lfs -text
|
| 23 |
+
*.pth filter=lfs diff=lfs merge=lfs -text
|
| 24 |
+
*.rar filter=lfs diff=lfs merge=lfs -text
|
| 25 |
+
*.safetensors filter=lfs diff=lfs merge=lfs -text
|
| 26 |
+
saved_model/**/* filter=lfs diff=lfs merge=lfs -text
|
| 27 |
+
*.tar.* filter=lfs diff=lfs merge=lfs -text
|
| 28 |
+
*.tar filter=lfs diff=lfs merge=lfs -text
|
| 29 |
+
*.tflite filter=lfs diff=lfs merge=lfs -text
|
| 30 |
+
*.tgz filter=lfs diff=lfs merge=lfs -text
|
| 31 |
+
*.wasm filter=lfs diff=lfs merge=lfs -text
|
| 32 |
+
*.xz filter=lfs diff=lfs merge=lfs -text
|
| 33 |
+
*.zip filter=lfs diff=lfs merge=lfs -text
|
| 34 |
+
*.zst filter=lfs diff=lfs merge=lfs -text
|
| 35 |
+
*tfevents* filter=lfs diff=lfs merge=lfs -text
|
| 36 |
+
example/student/04MhZs7YN08-Scene-0048.mp4 filter=lfs diff=lfs merge=lfs -text
|
| 37 |
+
example/student/7y57zc-GvBA-Scene-0009.mp4 filter=lfs diff=lfs merge=lfs -text
|
| 38 |
+
example/student/8YSha7iU2ZI-Scene-1025.mp4 filter=lfs diff=lfs merge=lfs -text
|
| 39 |
+
example/student/97moBezrzLo-Scene-0076.mp4 filter=lfs diff=lfs merge=lfs -text
|
| 40 |
+
example/student/9GDE5-rOOfk-Scene-0010.mp4 filter=lfs diff=lfs merge=lfs -text
|
| 41 |
+
example/student/a5M8hWgQbjU-Scene-0191.mp4 filter=lfs diff=lfs merge=lfs -text
|
| 42 |
+
example/student/DZt-RtpcJZQ-Scene-0043_m2.mp4 filter=lfs diff=lfs merge=lfs -text
|
| 43 |
+
example/student/FAfMSWi0FSA-Scene-0200.mp4 filter=lfs diff=lfs merge=lfs -text
|
| 44 |
+
example/student/ijGuAu0lLNM-Scene-0008.mp4 filter=lfs diff=lfs merge=lfs -text
|
| 45 |
+
example/student/zFehFw9SbDw-Scene-0034.mp4 filter=lfs diff=lfs merge=lfs -text
|
| 46 |
+
example/teacher/04MhZs7YN08-Scene-0048.mp4 filter=lfs diff=lfs merge=lfs -text
|
| 47 |
+
example/teacher/7y57zc-GvBA-Scene-0009.mp4 filter=lfs diff=lfs merge=lfs -text
|
| 48 |
+
example/teacher/8YSha7iU2ZI-Scene-1025.mp4 filter=lfs diff=lfs merge=lfs -text
|
| 49 |
+
example/teacher/97moBezrzLo-Scene-0076.mp4 filter=lfs diff=lfs merge=lfs -text
|
| 50 |
+
example/teacher/9GDE5-rOOfk-Scene-0010.mp4 filter=lfs diff=lfs merge=lfs -text
|
| 51 |
+
example/teacher/a5M8hWgQbjU-Scene-0191.mp4 filter=lfs diff=lfs merge=lfs -text
|
| 52 |
+
example/teacher/DZt-RtpcJZQ-Scene-0043_teacher.mp4 filter=lfs diff=lfs merge=lfs -text
|
| 53 |
+
example/teacher/FAfMSWi0FSA-Scene-0200.mp4 filter=lfs diff=lfs merge=lfs -text
|
| 54 |
+
example/teacher/ijGuAu0lLNM-Scene-0008.mp4 filter=lfs diff=lfs merge=lfs -text
|
| 55 |
+
example/teacher/zFehFw9SbDw-Scene-0034.mp4 filter=lfs diff=lfs merge=lfs -text
|
| 56 |
+
example/Ltx23VAED/04MhZs7YN08-Scene-0048.mp4 filter=lfs diff=lfs merge=lfs -text
|
| 57 |
+
example/Ltx23VAED/7y57zc-GvBA-Scene-0009.mp4 filter=lfs diff=lfs merge=lfs -text
|
| 58 |
+
example/Ltx23VAED/8YSha7iU2ZI-Scene-1025.mp4 filter=lfs diff=lfs merge=lfs -text
|
| 59 |
+
example/Ltx23VAED/97moBezrzLo-Scene-0076.mp4 filter=lfs diff=lfs merge=lfs -text
|
| 60 |
+
example/Ltx23VAED/9GDE5-rOOfk-Scene-0010.mp4 filter=lfs diff=lfs merge=lfs -text
|
| 61 |
+
example/Ltx23VAED/a5M8hWgQbjU-Scene-0191.mp4 filter=lfs diff=lfs merge=lfs -text
|
| 62 |
+
example/Ltx23VAED/DZt-RtpcJZQ-Scene-0043.mp4 filter=lfs diff=lfs merge=lfs -text
|
| 63 |
+
example/Ltx23VAED/FAfMSWi0FSA-Scene-0200.mp4 filter=lfs diff=lfs merge=lfs -text
|
| 64 |
+
example/Ltx23VAED/ijGuAu0lLNM-Scene-0008.mp4 filter=lfs diff=lfs merge=lfs -text
|
| 65 |
+
example/Ltx23VAED/zFehFw9SbDw-Scene-0034.mp4 filter=lfs diff=lfs merge=lfs -text
|
| 66 |
+
example/PrunaVAED/04MhZs7YN08-Scene-0048.mp4 filter=lfs diff=lfs merge=lfs -text
|
| 67 |
+
example/PrunaVAED/7y57zc-GvBA-Scene-0009.mp4 filter=lfs diff=lfs merge=lfs -text
|
| 68 |
+
example/PrunaVAED/8YSha7iU2ZI-Scene-1025.mp4 filter=lfs diff=lfs merge=lfs -text
|
| 69 |
+
example/PrunaVAED/97moBezrzLo-Scene-0076.mp4 filter=lfs diff=lfs merge=lfs -text
|
| 70 |
+
example/PrunaVAED/9GDE5-rOOfk-Scene-0010.mp4 filter=lfs diff=lfs merge=lfs -text
|
| 71 |
+
example/PrunaVAED/a5M8hWgQbjU-Scene-0191.mp4 filter=lfs diff=lfs merge=lfs -text
|
| 72 |
+
example/PrunaVAED/DZt-RtpcJZQ-Scene-0043.mp4 filter=lfs diff=lfs merge=lfs -text
|
| 73 |
+
example/PrunaVAED/FAfMSWi0FSA-Scene-0200.mp4 filter=lfs diff=lfs merge=lfs -text
|
| 74 |
+
example/PrunaVAED/ijGuAu0lLNM-Scene-0008.mp4 filter=lfs diff=lfs merge=lfs -text
|
| 75 |
+
example/PrunaVAED/zFehFw9SbDw-Scene-0034.mp4 filter=lfs diff=lfs merge=lfs -text
|
README.md
ADDED
|
@@ -0,0 +1,289 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
---
|
| 2 |
+
license: other
|
| 3 |
+
license_name: ltx-2
|
| 4 |
+
license_link: https://github.com/Lightricks/LTX-2/blob/main/LICENSE
|
| 5 |
+
tags:
|
| 6 |
+
- video
|
| 7 |
+
- vae
|
| 8 |
+
- decoder
|
| 9 |
+
- ltx
|
| 10 |
+
- ltx-2.3
|
| 11 |
+
- pruned
|
| 12 |
+
- diffusers
|
| 13 |
+
- image-to-video
|
| 14 |
+
- text-to-video
|
| 15 |
+
library_name: diffusers
|
| 16 |
+
base_model: diffusers/LTX-2.3-Diffusers
|
| 17 |
+
pipeline_tag: text-to-video
|
| 18 |
+
---
|
| 19 |
+
<!-- header start -->
|
| 20 |
+
<!-- 200823 -->
|
| 21 |
+
<div style="width: auto; margin-left: auto; margin-right: auto">
|
| 22 |
+
<a href="https://docs.pruna.ai/en/latest/setup/pip.html" target="_blank" rel="noopener noreferrer">
|
| 23 |
+
<img src="https://imgur.com/rVAgqMY.png" alt="PrunaAI" style="width: 100%; min-width: 400px; display: block; margin: auto;">
|
| 24 |
+
</a>
|
| 25 |
+
</div>
|
| 26 |
+
<!-- header end -->
|
| 27 |
+
|
| 28 |
+
[](https://github.com/PrunaAI/pruna)
|
| 29 |
+
[](https://twitter.com/PrunaAI)
|
| 30 |
+
[](https://www.linkedin.com/company/pruna-ai)
|
| 31 |
+
[](https://discord.com/invite/JFQmtFKCjd)
|
| 32 |
+
[](https://dashboard.pruna.ai/login?utm_source=huggingface&utm_medium=org_card&utm_campaign=hf_traffic)
|
| 33 |
+
|
| 34 |
+
<div align="center">
|
| 35 |
+
|
| 36 |
+
<h1 style="color: #9334E9;">⚡ PrunaVAED</h1>
|
| 37 |
+
|
| 38 |
+
<h2>A drop-in replacement decoder for LTX-2.3</h2>
|
| 39 |
+
|
| 40 |
+
<h3>
|
| 41 |
+
<span style="color: #9334E9;">1.7× faster</span>
|
| 42 |
+
·
|
| 43 |
+
<span style="color: #9334E9;">~50% lower peak VRAM</span>
|
| 44 |
+
·
|
| 45 |
+
Near-original visual quality
|
| 46 |
+
</h3>
|
| 47 |
+
|
| 48 |
+
</div>
|
| 49 |
+
|
| 50 |
+
**PrunaVAED directly replaces the video VAE decoder in
|
| 51 |
+
[`diffusers/LTX-2.3-Diffusers`](https://huggingface.co/diffusers/LTX-2.3-Diffusers).
|
| 52 |
+
The encoder and latent format remain unchanged, making it a drop-in upgrade
|
| 53 |
+
for faster, more memory-efficient LTX-2.3 decoding.**
|
| 54 |
+
|
| 55 |
+
## Examples
|
| 56 |
+
|
| 57 |
+
Both columns decode the same LTX-2.3 latent from the distilled pipeline.
|
| 58 |
+
|
| 59 |
+
<table>
|
| 60 |
+
<tr>
|
| 61 |
+
<th align="center" width="50%">LTX-2.3 decoder</th>
|
| 62 |
+
<th align="center" width="50%">⚡ PrunaVAED</th>
|
| 63 |
+
</tr>
|
| 64 |
+
<tr>
|
| 65 |
+
<td width="50%">
|
| 66 |
+
<video controls autoplay muted loop playsinline
|
| 67 |
+
style="width:100%; display:block;"
|
| 68 |
+
src="https://huggingface.co/PrunaAI/PrunaVAED/resolve/main/example/Ltx23VAED/04MhZs7YN08-Scene-0048.mp4">
|
| 69 |
+
</video>
|
| 70 |
+
</td>
|
| 71 |
+
<td width="50%">
|
| 72 |
+
<video controls autoplay muted loop playsinline
|
| 73 |
+
style="width:100%; display:block;"
|
| 74 |
+
src="https://huggingface.co/PrunaAI/PrunaVAED/resolve/main/example/PrunaVAED/04MhZs7YN08-Scene-0048.mp4">
|
| 75 |
+
</video>
|
| 76 |
+
</td>
|
| 77 |
+
</tr>
|
| 78 |
+
<tr>
|
| 79 |
+
<td width="50%">
|
| 80 |
+
<video controls autoplay muted loop playsinline
|
| 81 |
+
style="width:100%; display:block;"
|
| 82 |
+
src="https://huggingface.co/PrunaAI/PrunaVAED/resolve/main/example/Ltx23VAED/FAfMSWi0FSA-Scene-0200.mp4">
|
| 83 |
+
</video>
|
| 84 |
+
</td>
|
| 85 |
+
<td width="50%">
|
| 86 |
+
<video controls autoplay muted loop playsinline
|
| 87 |
+
style="width:100%; display:block;"
|
| 88 |
+
src="https://huggingface.co/PrunaAI/PrunaVAED/resolve/main/example/PrunaVAED/FAfMSWi0FSA-Scene-0200.mp4">
|
| 89 |
+
</video>
|
| 90 |
+
</td>
|
| 91 |
+
</tr>
|
| 92 |
+
</table>
|
| 93 |
+
|
| 94 |
+
<details>
|
| 95 |
+
<summary><strong>More side-by-side comparisons</strong></summary>
|
| 96 |
+
|
| 97 |
+
<br>
|
| 98 |
+
|
| 99 |
+
<table>
|
| 100 |
+
<tr>
|
| 101 |
+
<th align="center" width="50%">LTX-2.3 decoder</th>
|
| 102 |
+
<th align="center" width="50%">⚡ PrunaVAED</th>
|
| 103 |
+
</tr>
|
| 104 |
+
<tr>
|
| 105 |
+
<td width="50%"><video controls autoplay muted loop playsinline style="width:100%; display:block;" src="https://huggingface.co/PrunaAI/PrunaVAED/resolve/main/example/Ltx23VAED/8YSha7iU2ZI-Scene-1025.mp4"></video></td>
|
| 106 |
+
<td width="50%"><video controls autoplay muted loop playsinline style="width:100%; display:block;" src="https://huggingface.co/PrunaAI/PrunaVAED/resolve/main/example/PrunaVAED/8YSha7iU2ZI-Scene-1025.mp4"></video></td>
|
| 107 |
+
</tr>
|
| 108 |
+
<tr>
|
| 109 |
+
<td width="50%"><video controls autoplay muted loop playsinline style="width:100%; display:block;" src="https://huggingface.co/PrunaAI/PrunaVAED/resolve/main/example/Ltx23VAED/DZt-RtpcJZQ-Scene-0043.mp4"></video></td>
|
| 110 |
+
<td width="50%"><video controls autoplay muted loop playsinline style="width:100%; display:block;" src="https://huggingface.co/PrunaAI/PrunaVAED/resolve/main/example/PrunaVAED/DZt-RtpcJZQ-Scene-0043.mp4"></video></td>
|
| 111 |
+
</tr>
|
| 112 |
+
<tr>
|
| 113 |
+
<td width="50%"><video controls autoplay muted loop playsinline style="width:100%; display:block;" src="https://huggingface.co/PrunaAI/PrunaVAED/resolve/main/example/Ltx23VAED/9GDE5-rOOfk-Scene-0010.mp4"></video></td>
|
| 114 |
+
<td width="50%"><video controls autoplay muted loop playsinline style="width:100%; display:block;" src="https://huggingface.co/PrunaAI/PrunaVAED/resolve/main/example/PrunaVAED/9GDE5-rOOfk-Scene-0010.mp4"></video></td>
|
| 115 |
+
</tr>
|
| 116 |
+
<tr>
|
| 117 |
+
<td width="50%"><video controls autoplay muted loop playsinline style="width:100%; display:block;" src="https://huggingface.co/PrunaAI/PrunaVAED/resolve/main/example/Ltx23VAED/7y57zc-GvBA-Scene-0009.mp4"></video></td>
|
| 118 |
+
<td width="50%"><video controls autoplay muted loop playsinline style="width:100%; display:block;" src="https://huggingface.co/PrunaAI/PrunaVAED/resolve/main/example/PrunaVAED/7y57zc-GvBA-Scene-0009.mp4"></video></td>
|
| 119 |
+
</tr>
|
| 120 |
+
</table>
|
| 121 |
+
</details>
|
| 122 |
+
|
| 123 |
+
|
| 124 |
+
## Benchmark
|
| 125 |
+
|
| 126 |
+
Metrics compare videos decoded from the **same latents** by the LTX-2.3 VAED and PrunaVAED. Two set of latents were
|
| 127 |
+
generated with the full
|
| 128 |
+
[`ti2vid_two_stages`](https://github.com/Lightricks/LTX-2/blob/main/packages/ltx-pipelines/src/ltx_pipelines/ti2vid_two_stages.py)
|
| 129 |
+
pipeline and the
|
| 130 |
+
[`distilled`](https://github.com/Lightricks/LTX-2/blob/main/packages/ltx-pipelines/src/ltx_pipelines/distilled.py)
|
| 131 |
+
pipeline (recommended).
|
| 132 |
+
|
| 133 |
+
<table class="w-full table-fixed">
|
| 134 |
+
<tr><th>Setting</th><th>Value</th></tr>
|
| 135 |
+
<tr><td>Model</td><td>PrunaVAED</td></tr>
|
| 136 |
+
<tr><td>Precision</td><td><code>bfloat16</code></td></tr>
|
| 137 |
+
<tr><td>Batch size</td><td>1</td></tr>
|
| 138 |
+
<tr><td>Device</td><td>NVIDIA H100 80GB</td></tr>
|
| 139 |
+
<tr><td>Decode</td><td>Full (no tiling / rolling / <code>torch.compile</code>)</td></tr>
|
| 140 |
+
</table>
|
| 141 |
+
|
| 142 |
+
<table class="w-full table-fixed">
|
| 143 |
+
<tr>
|
| 144 |
+
<th>Clip length</th><th>LTX-2.3 VAED</th><th>PrunaVAED</th><th>Speedup</th>
|
| 145 |
+
</tr>
|
| 146 |
+
<tr>
|
| 147 |
+
<td>5 s @ 720p</td><td align="right">841.3 ms</td><td align="right">500.2 ms</td><td align="right"><strong>1.68×</strong></td>
|
| 148 |
+
</tr>
|
| 149 |
+
<tr>
|
| 150 |
+
<td>10 s @ 720p</td><td align="right">1670.7 ms</td><td align="right">998.0 ms</td><td align="right"><strong>1.67×</strong></td>
|
| 151 |
+
</tr>
|
| 152 |
+
<tr>
|
| 153 |
+
<td>5 s @ 1080p</td><td align="right">1959.3 ms</td><td align="right">1152.1 ms</td><td align="right"><strong>1.7×</strong></td>
|
| 154 |
+
</tr>
|
| 155 |
+
<tr>
|
| 156 |
+
<td>10 s @ 1080p</td><td align="right">4903 ms*</td><td align="right">2356 ms</td><td align="right"><strong>2.08×</strong></td>
|
| 157 |
+
</tr>
|
| 158 |
+
</table>
|
| 159 |
+
|
| 160 |
+
\* Tiling required on H100; 14533ms otherwise
|
| 161 |
+
|
| 162 |
+
<table class="w-full table-fixed">
|
| 163 |
+
<tr>
|
| 164 |
+
<th>Latent set</th><th>Clips</th><th>PSNR ↑</th><th>LPIPS ↓</th><th>SSIM ↑</th><th>Δ-frame PSNR ↑</th>
|
| 165 |
+
</tr>
|
| 166 |
+
<tr>
|
| 167 |
+
<td>Distilled two-stage · 720p</td><td align="right">200</td><td align="right"><strong>39.23</strong></td><td align="right"><strong>0.0087</strong></td><td align="right"><strong>0.9811</strong></td><td align="right"><strong>38.32</strong></td>
|
| 168 |
+
</tr>
|
| 169 |
+
<tr>
|
| 170 |
+
<td>TI2Vid two-stage · 720p</td><td align="right">200</td><td align="right"><strong>40.34</strong></td><td align="right"><strong>0.0094</strong></td><td align="right"><strong>0.9823</strong></td><td align="right"><strong>40.06</strong></td>
|
| 171 |
+
</tr>
|
| 172 |
+
<tr>
|
| 173 |
+
<td>Distilled two-stage · 1080p</td><td align="right">200</td><td align="right"><strong>41.06</strong></td><td align="right"><strong>0.0052</strong></td><td align="right"><strong>0.9876</strong></td><td align="right"><strong>40.43</strong></td>
|
| 174 |
+
</tr>
|
| 175 |
+
</table>
|
| 176 |
+
|
| 177 |
+
**Δ-frame PSNR** is PSNR on consecutive-frame
|
| 178 |
+
differences (temporal consistency).
|
| 179 |
+
|
| 180 |
+
## Quickstart
|
| 181 |
+
|
| 182 |
+
End-to-end smoke demo: generate a short video with the diffusers LTX-2.3
|
| 183 |
+
**distilled** two-stage pipeline, then decode the same latent with the stock
|
| 184 |
+
LTX-2.3 VAE and with PrunaVAED. Writes two mp4s and prints decode time (ms).
|
| 185 |
+
Needs a **CUDA GPU**.
|
| 186 |
+
|
| 187 |
+
```bash
|
| 188 |
+
# 1. Install the Hugging Face CLI
|
| 189 |
+
pip install hf
|
| 190 |
+
|
| 191 |
+
# 2. Download this repo
|
| 192 |
+
hf download PrunaAI/PrunaVAED --local-dir PrunaVAED
|
| 193 |
+
cd PrunaVAED
|
| 194 |
+
|
| 195 |
+
# 3. Install dependencies
|
| 196 |
+
pip install -r requirements-demo.txt
|
| 197 |
+
|
| 198 |
+
# 4. Run the demo (~1080p, ~5 s @ 24 fps)
|
| 199 |
+
python demo/demo_distilled_decode.py
|
| 200 |
+
```
|
| 201 |
+
|
| 202 |
+
Outputs land in `outputs/demo_distilled/` (`ltx23.mp4`, `prunavaed.mp4`).
|
| 203 |
+
Edit `PROMPT` / resolution at the top of `demo/demo_distilled_decode.py` if needed.
|
| 204 |
+
|
| 205 |
+
|
| 206 |
+
## Architecture
|
| 207 |
+
|
| 208 |
+
PrunaVAED is bitwise identical to LTX 2.3 VAED from `conv_in` to `up_blocks.0`. Pruning starts
|
| 209 |
+
at `up_blocks.1`.
|
| 210 |
+
|
| 211 |
+
| Stage | Channel reduction |
|
| 212 |
+
|---|---:|
|
| 213 |
+
| `up_blocks.0` | unchanged |
|
| 214 |
+
| `up_blocks.1` | 25% |
|
| 215 |
+
| `up_blocks.2` | 50% |
|
| 216 |
+
| `up_blocks.3` | 50% |
|
| 217 |
+
|
| 218 |
+
| | Teacher (LTX-2.3) | PrunaVAED |
|
| 219 |
+
|---|---:|---:|
|
| 220 |
+
| Encoder params | 318.9 M | 318.9 M (unchanged) |
|
| 221 |
+
| Decoder params | 407.2 M | 345.0 M (−15%) |
|
| 222 |
+
| **Total VAE params** | **726.1 M** | **663.9 M** |
|
| 223 |
+
|
| 224 |
+
`prunavaed/patch_diffusers.py` provides the required compatibility shim for
|
| 225 |
+
the pinned diffusers version.
|
| 226 |
+
|
| 227 |
+
## Limitations
|
| 228 |
+
|
| 229 |
+
- Evaluated on 4-second clips at 24 fps.
|
| 230 |
+
- Benchmarked on one H100 80GB with bfloat16 and batch size 1.
|
| 231 |
+
- Speed and VRAM vary with hardware, resolution, batch size, and software.
|
| 232 |
+
- 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)
|
| 233 |
+
pipeline and the
|
| 234 |
+
[`distilled`](https://github.com/Lightricks/LTX-2/blob/main/packages/ltx-pipelines/src/ltx_pipelines/distilled.py))
|
| 235 |
+
- 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.
|
| 236 |
+
|
| 237 |
+
## What this is not
|
| 238 |
+
|
| 239 |
+
- Not a full T2V model.
|
| 240 |
+
- Not a replacement for the LTX denoiser/generator.
|
| 241 |
+
- Not compatible with arbitrary VAE latents.
|
| 242 |
+
- Not bit-exact with the original decoder.
|
| 243 |
+
|
| 244 |
+
## License
|
| 245 |
+
|
| 246 |
+
PrunaVAED is a derivative of LTX-2.3 and is distributed under the
|
| 247 |
+
**[LTX-2 Community License Agreement](LICENSE)**. Review its use restrictions
|
| 248 |
+
and commercial terms before using or redistributing the model.
|
| 249 |
+
|
| 250 |
+
Helper code adapted from Hugging Face diffusers retains its Apache-2.0
|
| 251 |
+
attribution; see [`NOTICE`](NOTICE).
|
| 252 |
+
|
| 253 |
+
## What's next?
|
| 254 |
+
|
| 255 |
+
- **Use PrunaVAED to speed up LTX-2.3 decoding.**
|
| 256 |
+
- Compress your own models with [Pruna](https://github.com/PrunaAI/pruna) and give us a ⭐️ for more efficiency!
|
| 257 |
+
- Want to use our optimized video model right away? Check
|
| 258 |
+
[P-Video](https://www.pruna.ai/p-video) and
|
| 259 |
+
[P-Video documentation](https://docs.pruna.ai/en/stable/docs_pruna_endpoints/performance_models/p-video.html).
|
| 260 |
+
|
| 261 |
+
<style>
|
| 262 |
+
.model-button {
|
| 263 |
+
display: inline-flex;
|
| 264 |
+
flex-direction: row;
|
| 265 |
+
justify-content: center;
|
| 266 |
+
align-items: center;
|
| 267 |
+
gap: 8px;
|
| 268 |
+
|
| 269 |
+
padding: 8px 20px;
|
| 270 |
+
border: none;
|
| 271 |
+
border-radius: 8px;
|
| 272 |
+
|
| 273 |
+
background: #9334e9;
|
| 274 |
+
color: #ffffff;
|
| 275 |
+
|
| 276 |
+
font-size: 14px;
|
| 277 |
+
font-weight: 400;
|
| 278 |
+
line-height: 1;
|
| 279 |
+
text-decoration: none;
|
| 280 |
+
white-space: nowrap;
|
| 281 |
+
cursor: pointer;
|
| 282 |
+
|
| 283 |
+
box-sizing: border-box;
|
| 284 |
+
overflow: visible;
|
| 285 |
+
opacity: 1;
|
| 286 |
+
}
|
| 287 |
+
</style>
|
| 288 |
+
|
| 289 |
+
<a href="https://dashboard.pruna.ai/login?utm_source=huggingface&utm_medium=model_card&utm_campaign=hf_traffic" class="model-button">Try our models</a>
|
demo/demo_distilled_decode.py
ADDED
|
@@ -0,0 +1,227 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
#!/usr/bin/env python3
|
| 2 |
+
"""Quick demo: generate a short LTX-2.3 video latent, then decode it twice.
|
| 3 |
+
|
| 4 |
+
Compares the stock LTX-2.3 VAE decoder with this repo's pruned decoder
|
| 5 |
+
(PrunaVAED) on the *same* latent. Prints decode time (ms) and peak VRAM (GiB),
|
| 6 |
+
and writes two mp4s.
|
| 7 |
+
|
| 8 |
+
Requires a CUDA GPU. From the repo root:
|
| 9 |
+
|
| 10 |
+
pip install -r requirements-test.txt
|
| 11 |
+
python demo/demo_distilled_decode.py
|
| 12 |
+
"""
|
| 13 |
+
|
| 14 |
+
from __future__ import annotations
|
| 15 |
+
|
| 16 |
+
import statistics
|
| 17 |
+
import sys
|
| 18 |
+
import time
|
| 19 |
+
from pathlib import Path
|
| 20 |
+
|
| 21 |
+
import imageio.v3 as iio
|
| 22 |
+
import torch
|
| 23 |
+
from diffusers import LTX2LatentUpsamplePipeline, LTX2Pipeline
|
| 24 |
+
from diffusers.models.autoencoders import AutoencoderKLLTX2Video
|
| 25 |
+
from diffusers.pipelines.ltx2.latent_upsampler import LTX2LatentUpsamplerModel
|
| 26 |
+
from diffusers.pipelines.ltx2.utils import (
|
| 27 |
+
DEFAULT_NEGATIVE_PROMPT,
|
| 28 |
+
DISTILLED_SIGMA_VALUES,
|
| 29 |
+
STAGE_2_DISTILLED_SIGMA_VALUES,
|
| 30 |
+
)
|
| 31 |
+
|
| 32 |
+
REPO_ROOT = Path(__file__).resolve().parents[1]
|
| 33 |
+
sys.path.insert(0, str(REPO_ROOT))
|
| 34 |
+
|
| 35 |
+
from patch_diffusers import patch_pruna_ltx2_decoder # noqa: E402
|
| 36 |
+
|
| 37 |
+
|
| 38 |
+
torch.backends.cuda.enable_cudnn_sdp(False)
|
| 39 |
+
|
| 40 |
+
# ---------------------------------------------------------------------------
|
| 41 |
+
# Settings (edit these if you want)
|
| 42 |
+
# ---------------------------------------------------------------------------
|
| 43 |
+
|
| 44 |
+
# Official Diffusers checkpoints for the 2-stage distilled recipe.
|
| 45 |
+
DISTILLED_MODEL = "diffusers/LTX-2.3-Distilled-Diffusers"
|
| 46 |
+
SPATIAL_UPSAMPLER = "dg845/LTX-2.3-Spatial-Upsampler-Diffusers"
|
| 47 |
+
LTX23_VAE = "diffusers/LTX-2.3-Diffusers" # stock decoder (baseline)
|
| 48 |
+
PRUNED_VAE = str(REPO_ROOT) # this repo's vae/ folder
|
| 49 |
+
|
| 50 |
+
# ~1080p for 5 s @ 24 fps. Height/width must be multiples of 64 (2-stage).
|
| 51 |
+
HEIGHT, WIDTH, NUM_FRAMES, FPS = 1088, 1920, 121, 24.0
|
| 52 |
+
SEED = 42
|
| 53 |
+
DECODE_WARMUP, DECODE_RUNS = 1, 3 # timing: 1 warm-up + median of 3 runs
|
| 54 |
+
|
| 55 |
+
PROMPT = (
|
| 56 |
+
"The video shows a hockey player in a green jersey and blue helmet skating on the ice with a hockey stick. The player is seen moving around the ice, passing the puck to another player who is also wearing a green jersey and blue helmet. The player in the green jersey is seen skating away from the camera, and then turning around to face the camera. The ice rink is surrounded by boards with advertisements, and there are other players in the background. The player in the green jersey is wearing black gloves and black skates. The player in the green jersey is also seen skating towards the camera and then away from the camera again."
|
| 57 |
+
)
|
| 58 |
+
|
| 59 |
+
DEVICE = "cuda"
|
| 60 |
+
DTYPE = torch.bfloat16
|
| 61 |
+
OUTPUT_DIR = REPO_ROOT / "outputs" / "demo_distilled"
|
| 62 |
+
|
| 63 |
+
|
| 64 |
+
# ---------------------------------------------------------------------------
|
| 65 |
+
# Helpers
|
| 66 |
+
# ---------------------------------------------------------------------------
|
| 67 |
+
|
| 68 |
+
def save_mp4(video: torch.Tensor, path: Path, fps: float) -> None:
|
| 69 |
+
"""Save a BCTHW tensor in [-1, 1] as an H.264 mp4."""
|
| 70 |
+
frames = video[0].permute(1, 2, 3, 0).clamp(-1, 1).float()
|
| 71 |
+
frames = ((frames + 1) / 2 * 255).round().byte().cpu().numpy()
|
| 72 |
+
path.parent.mkdir(parents=True, exist_ok=True)
|
| 73 |
+
iio.imwrite(path, frames, fps=fps, codec="libx264")
|
| 74 |
+
|
| 75 |
+
|
| 76 |
+
def load_vae(model_id: str) -> AutoencoderKLLTX2Video:
|
| 77 |
+
"""Load a video VAE decoder on GPU, tiling off so we time full-frame decodes."""
|
| 78 |
+
vae = AutoencoderKLLTX2Video.from_pretrained(
|
| 79 |
+
model_id, subfolder="vae", torch_dtype=DTYPE
|
| 80 |
+
)
|
| 81 |
+
vae.disable_tiling()
|
| 82 |
+
return vae.to(DEVICE).eval()
|
| 83 |
+
|
| 84 |
+
|
| 85 |
+
@torch.inference_mode()
|
| 86 |
+
def timed_decode(
|
| 87 |
+
vae: AutoencoderKLLTX2Video, latent: torch.Tensor
|
| 88 |
+
) -> tuple[torch.Tensor, float, float]:
|
| 89 |
+
"""Decode DECODE_RUNS times after warm-up.
|
| 90 |
+
|
| 91 |
+
Returns the video on CPU, the median latency in ms and the peak VRAM in GiB.
|
| 92 |
+
Deliberately no autocast: weights and latent are already bfloat16, and autocast
|
| 93 |
+
would run every PerChannelRMSNorm (``x**2``) in float32, costing time and VRAM.
|
| 94 |
+
"""
|
| 95 |
+
latent = latent.to(DEVICE, DTYPE)
|
| 96 |
+
|
| 97 |
+
for _ in range(DECODE_WARMUP):
|
| 98 |
+
vae.decode(latent, return_dict=False)
|
| 99 |
+
|
| 100 |
+
torch.cuda.synchronize()
|
| 101 |
+
torch.cuda.reset_peak_memory_stats()
|
| 102 |
+
times_ms = []
|
| 103 |
+
for _ in range(DECODE_RUNS):
|
| 104 |
+
video = None # keep the previous output out of the peak VRAM measurement
|
| 105 |
+
torch.cuda.synchronize()
|
| 106 |
+
t0 = time.perf_counter()
|
| 107 |
+
video = vae.decode(latent, return_dict=False)[0]
|
| 108 |
+
torch.cuda.synchronize()
|
| 109 |
+
times_ms.append((time.perf_counter() - t0) * 1000)
|
| 110 |
+
|
| 111 |
+
peak_gib = torch.cuda.max_memory_allocated() / 2**30
|
| 112 |
+
return video.cpu(), statistics.median(times_ms), peak_gib
|
| 113 |
+
|
| 114 |
+
|
| 115 |
+
@torch.inference_mode()
|
| 116 |
+
def generate_latent(prompt: str) -> torch.Tensor:
|
| 117 |
+
"""Run the official 2-stage distilled pipeline and return the final video latent.
|
| 118 |
+
|
| 119 |
+
Same layout as Lightricks DistilledPipeline:
|
| 120 |
+
stage 1 @ half-res → 2× latent upsample → stage 2 @ full-res.
|
| 121 |
+
"""
|
| 122 |
+
half_h, half_w = HEIGHT // 2, WIDTH // 2
|
| 123 |
+
generator = torch.Generator(DEVICE).manual_seed(SEED)
|
| 124 |
+
|
| 125 |
+
pipe = LTX2Pipeline.from_pretrained(DISTILLED_MODEL, torch_dtype=DTYPE)
|
| 126 |
+
pipe.enable_model_cpu_offload(device=DEVICE)
|
| 127 |
+
|
| 128 |
+
# Stage 1 — cheap draft at half resolution.
|
| 129 |
+
print(f"1/3 Stage 1 @ {half_w}×{half_h}")
|
| 130 |
+
video_latent, audio_latent = pipe(
|
| 131 |
+
prompt=prompt,
|
| 132 |
+
negative_prompt=DEFAULT_NEGATIVE_PROMPT,
|
| 133 |
+
height=half_h,
|
| 134 |
+
width=half_w,
|
| 135 |
+
num_frames=NUM_FRAMES,
|
| 136 |
+
frame_rate=FPS,
|
| 137 |
+
num_inference_steps=len(DISTILLED_SIGMA_VALUES),
|
| 138 |
+
sigmas=DISTILLED_SIGMA_VALUES,
|
| 139 |
+
guidance_scale=1.0,
|
| 140 |
+
generator=generator,
|
| 141 |
+
output_type="latent",
|
| 142 |
+
return_dict=False,
|
| 143 |
+
)
|
| 144 |
+
|
| 145 |
+
# Upsample — bring the latent to full resolution before refining.
|
| 146 |
+
print(f"2/3 Upsample → {WIDTH}×{HEIGHT}")
|
| 147 |
+
upsampler = LTX2LatentUpsamplerModel.from_pretrained(
|
| 148 |
+
SPATIAL_UPSAMPLER, subfolder="latent_upsampler", torch_dtype=DTYPE
|
| 149 |
+
)
|
| 150 |
+
upsample_pipe = LTX2LatentUpsamplePipeline(vae=pipe.vae, latent_upsampler=upsampler)
|
| 151 |
+
upsample_pipe.enable_model_cpu_offload(device=DEVICE)
|
| 152 |
+
video_latent = upsample_pipe(
|
| 153 |
+
latents=video_latent[:1],
|
| 154 |
+
height=half_h,
|
| 155 |
+
width=half_w,
|
| 156 |
+
num_frames=NUM_FRAMES,
|
| 157 |
+
output_type="latent",
|
| 158 |
+
return_dict=False,
|
| 159 |
+
)[0]
|
| 160 |
+
del upsample_pipe, upsampler
|
| 161 |
+
torch.cuda.empty_cache()
|
| 162 |
+
|
| 163 |
+
# Stage 2 — refine at full resolution (renoises from STAGE_2 schedule).
|
| 164 |
+
print(f"3/3 Stage 2 @ {WIDTH}×{HEIGHT}")
|
| 165 |
+
video_latent, _ = pipe(
|
| 166 |
+
latents=video_latent,
|
| 167 |
+
audio_latents=audio_latent,
|
| 168 |
+
prompt=prompt,
|
| 169 |
+
negative_prompt=DEFAULT_NEGATIVE_PROMPT,
|
| 170 |
+
height=HEIGHT,
|
| 171 |
+
width=WIDTH,
|
| 172 |
+
num_frames=NUM_FRAMES,
|
| 173 |
+
frame_rate=FPS,
|
| 174 |
+
num_inference_steps=len(STAGE_2_DISTILLED_SIGMA_VALUES),
|
| 175 |
+
noise_scale=STAGE_2_DISTILLED_SIGMA_VALUES[0],
|
| 176 |
+
sigmas=STAGE_2_DISTILLED_SIGMA_VALUES,
|
| 177 |
+
guidance_scale=1.0,
|
| 178 |
+
generator=generator,
|
| 179 |
+
output_type="latent",
|
| 180 |
+
return_dict=False,
|
| 181 |
+
)
|
| 182 |
+
|
| 183 |
+
latent = video_latent.detach().cpu()
|
| 184 |
+
del pipe
|
| 185 |
+
torch.cuda.empty_cache()
|
| 186 |
+
return latent
|
| 187 |
+
|
| 188 |
+
|
| 189 |
+
def compare_decoders(latent: torch.Tensor, out_dir: Path) -> None:
|
| 190 |
+
"""Decode the same latent with LTX-2.3 and PrunaVAED; print ms/VRAM and save mp4s."""
|
| 191 |
+
# Needed so Diffusers can build the pruned decoder graph correctly.
|
| 192 |
+
patch_pruna_ltx2_decoder()
|
| 193 |
+
|
| 194 |
+
results = {}
|
| 195 |
+
for name, model_id in (("ltx23", LTX23_VAE), ("prunavaed", PRUNED_VAE)):
|
| 196 |
+
print(f"Decoding with {name} …")
|
| 197 |
+
vae = load_vae(model_id)
|
| 198 |
+
video, ms, peak_gib = timed_decode(vae, latent)
|
| 199 |
+
results[name] = (ms, peak_gib)
|
| 200 |
+
path = out_dir / f"{name}.mp4"
|
| 201 |
+
save_mp4(video, path, FPS)
|
| 202 |
+
print(f" {name}: {ms:.1f} ms · {peak_gib:.2f} GiB peak → {path}")
|
| 203 |
+
del vae, video
|
| 204 |
+
torch.cuda.empty_cache()
|
| 205 |
+
|
| 206 |
+
(ltx_ms, ltx_gib), (pruna_ms, pruna_gib) = results["ltx23"], results["prunavaed"]
|
| 207 |
+
print(f"Speedup: {ltx_ms / pruna_ms:.2f}× Peak VRAM: {pruna_gib / ltx_gib:.0%} of LTX-2.3")
|
| 208 |
+
|
| 209 |
+
|
| 210 |
+
# ---------------------------------------------------------------------------
|
| 211 |
+
# Entry point
|
| 212 |
+
# ---------------------------------------------------------------------------
|
| 213 |
+
|
| 214 |
+
def main() -> None:
|
| 215 |
+
if not torch.cuda.is_available():
|
| 216 |
+
raise SystemExit("This demo needs a CUDA GPU.")
|
| 217 |
+
|
| 218 |
+
print(f"Prompt: {PROMPT[:80]}…")
|
| 219 |
+
print(f"Output: {OUTPUT_DIR}")
|
| 220 |
+
|
| 221 |
+
latent = generate_latent(PROMPT)
|
| 222 |
+
print(f"Latent shape: {tuple(latent.shape)}")
|
| 223 |
+
compare_decoders(latent, OUTPUT_DIR)
|
| 224 |
+
|
| 225 |
+
|
| 226 |
+
if __name__ == "__main__":
|
| 227 |
+
main()
|
example/Ltx23VAED/04MhZs7YN08-Scene-0048.mp4
ADDED
|
@@ -0,0 +1,3 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
version https://git-lfs.github.com/spec/v1
|
| 2 |
+
oid sha256:320db36560a3b9abac5c710bed668c0b25f8bc3b9cdd1935ba65ecec0ee2174f
|
| 3 |
+
size 1697633
|
example/Ltx23VAED/7y57zc-GvBA-Scene-0009.mp4
ADDED
|
@@ -0,0 +1,3 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
version https://git-lfs.github.com/spec/v1
|
| 2 |
+
oid sha256:004434c148e393a6be129a5682fd357b9340f035595d3f09348d764e487dd1e8
|
| 3 |
+
size 1697114
|
example/Ltx23VAED/8YSha7iU2ZI-Scene-1025.mp4
ADDED
|
@@ -0,0 +1,3 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
version https://git-lfs.github.com/spec/v1
|
| 2 |
+
oid sha256:ddabd909bc729ef15a643a0a63d4937d3846c5165a99d56b83e6e3f06b743179
|
| 3 |
+
size 828088
|
example/Ltx23VAED/97moBezrzLo-Scene-0076.mp4
ADDED
|
@@ -0,0 +1,3 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
version https://git-lfs.github.com/spec/v1
|
| 2 |
+
oid sha256:dba2ff092c0d4ed2fe369c9f3525d03f0e6b8d3d240f392a366e89e5b3149ad7
|
| 3 |
+
size 1848784
|
example/Ltx23VAED/9GDE5-rOOfk-Scene-0010.mp4
ADDED
|
@@ -0,0 +1,3 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
version https://git-lfs.github.com/spec/v1
|
| 2 |
+
oid sha256:a9cb3ab14edb2e0b772b5e472e1d3cb3ef802cba138095d7f52b119ea8a74c44
|
| 3 |
+
size 2239655
|
example/Ltx23VAED/DZt-RtpcJZQ-Scene-0043.mp4
ADDED
|
@@ -0,0 +1,3 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
version https://git-lfs.github.com/spec/v1
|
| 2 |
+
oid sha256:51facb9d24858bd48cdf7e565c4da9c863e875c9121e83b050233c31a59826b6
|
| 3 |
+
size 1174232
|
example/Ltx23VAED/FAfMSWi0FSA-Scene-0200.mp4
ADDED
|
@@ -0,0 +1,3 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
version https://git-lfs.github.com/spec/v1
|
| 2 |
+
oid sha256:505d374a8cc0caf0378903199684b2d0db7326ccac7c8f186d9b2eb4690f32ae
|
| 3 |
+
size 1864634
|
example/Ltx23VAED/a5M8hWgQbjU-Scene-0191.mp4
ADDED
|
@@ -0,0 +1,3 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
version https://git-lfs.github.com/spec/v1
|
| 2 |
+
oid sha256:b7629400d5d318681018cd22bb934d7ba728887c307e8be69853d0d1e01f1526
|
| 3 |
+
size 3360331
|
example/Ltx23VAED/ijGuAu0lLNM-Scene-0008.mp4
ADDED
|
@@ -0,0 +1,3 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
version https://git-lfs.github.com/spec/v1
|
| 2 |
+
oid sha256:b221c4b8606983128306ed1b6e7d049f300b5439b720a3f5c8bcbc9673f5b8b6
|
| 3 |
+
size 1711726
|
example/Ltx23VAED/zFehFw9SbDw-Scene-0034.mp4
ADDED
|
@@ -0,0 +1,3 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
version https://git-lfs.github.com/spec/v1
|
| 2 |
+
oid sha256:688f0d785cb9dafa858a35ac5192380df50483b81392dc1690ad32393ce536d2
|
| 3 |
+
size 2125596
|
example/PrunaVAED/04MhZs7YN08-Scene-0048.mp4
ADDED
|
@@ -0,0 +1,3 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
version https://git-lfs.github.com/spec/v1
|
| 2 |
+
oid sha256:43d29de4e52f013f14a82d1f4d41d78f7bb938184f30af48d2e53d09efffe558
|
| 3 |
+
size 1797361
|
example/PrunaVAED/7y57zc-GvBA-Scene-0009.mp4
ADDED
|
@@ -0,0 +1,3 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
version https://git-lfs.github.com/spec/v1
|
| 2 |
+
oid sha256:2de528d80672faf55079399ddf14533ef77b274099407d35fec5fec552b89ae3
|
| 3 |
+
size 1716725
|
example/PrunaVAED/8YSha7iU2ZI-Scene-1025.mp4
ADDED
|
@@ -0,0 +1,3 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
version https://git-lfs.github.com/spec/v1
|
| 2 |
+
oid sha256:c4c8ebfb36da16485337d9b2a48210efa707dedf8992c55be068629e1310c03d
|
| 3 |
+
size 898616
|
example/PrunaVAED/97moBezrzLo-Scene-0076.mp4
ADDED
|
@@ -0,0 +1,3 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
version https://git-lfs.github.com/spec/v1
|
| 2 |
+
oid sha256:c614003deb538d3b02d5e80010f0167b9c1a0fd55e62863479db9ab0e7134c89
|
| 3 |
+
size 1897236
|
example/PrunaVAED/9GDE5-rOOfk-Scene-0010.mp4
ADDED
|
@@ -0,0 +1,3 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
version https://git-lfs.github.com/spec/v1
|
| 2 |
+
oid sha256:845ad0a7d35633fe3ddf73da94534d468eb2ac36f6faa450db6c841439f785dd
|
| 3 |
+
size 2308198
|
example/PrunaVAED/DZt-RtpcJZQ-Scene-0043.mp4
ADDED
|
@@ -0,0 +1,3 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
version https://git-lfs.github.com/spec/v1
|
| 2 |
+
oid sha256:06f56792ff2d4793d12b4283c20788e4be6c22ea876729899b7e46614016c1f6
|
| 3 |
+
size 1214808
|
example/PrunaVAED/FAfMSWi0FSA-Scene-0200.mp4
ADDED
|
@@ -0,0 +1,3 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
version https://git-lfs.github.com/spec/v1
|
| 2 |
+
oid sha256:9442659a2cb5b6328bb8f0849b466a4e6b572243f1350efa93cc2cb70261ec10
|
| 3 |
+
size 1883423
|
example/PrunaVAED/a5M8hWgQbjU-Scene-0191.mp4
ADDED
|
@@ -0,0 +1,3 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
version https://git-lfs.github.com/spec/v1
|
| 2 |
+
oid sha256:af8032ecb4f9e3f5db93ff04293b296979cfba55351f73a1f309e1508d412c83
|
| 3 |
+
size 3051778
|
example/PrunaVAED/ijGuAu0lLNM-Scene-0008.mp4
ADDED
|
@@ -0,0 +1,3 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
version https://git-lfs.github.com/spec/v1
|
| 2 |
+
oid sha256:4eb2af16643f4402da2585878e41257e951141c491556e335b72abab3a9f26eb
|
| 3 |
+
size 1697225
|
example/PrunaVAED/zFehFw9SbDw-Scene-0034.mp4
ADDED
|
@@ -0,0 +1,3 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
version https://git-lfs.github.com/spec/v1
|
| 2 |
+
oid sha256:cf0ff506cb5ade7308c5ad0752bbe57e0b88b37462c5d4f20be10d6c3ff7961b
|
| 3 |
+
size 2129229
|
example/prompts.json
ADDED
|
@@ -0,0 +1,44 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
[
|
| 2 |
+
{
|
| 3 |
+
"vid": "04MhZs7YN08-Scene-0048",
|
| 4 |
+
"prompt": "In the video, a man is seen holding a small dog in his arms. The man is shirtless and wearing orange shorts. He is standing in a backyard with a pool and a swing set visible in the background. The dog appears to be a young puppy with a grey and white coat. The man is holding the dog close to his chest and appears to be speaking to the dog, although it is not clear what he is saying. The dog seems to be calm and comfortable in the man's arms. The man's tattoos are visible on his arms, and he has a muscular build. The video is taken during the daytime, and the lighting suggests it is either morning or late afternoon."
|
| 5 |
+
},
|
| 6 |
+
{
|
| 7 |
+
"vid": "7y57zc-GvBA-Scene-0009",
|
| 8 |
+
"prompt": "In the video, a person is seen petting two dogs. The first dog is brown and the second dog is black. The person is petting the brown dog first and then the black dog. The dogs seem to be enjoying the attention and are wagging their tails. The person is wearing a white shirt and is sitting on a white couch. The background is a white wall."
|
| 9 |
+
},
|
| 10 |
+
{
|
| 11 |
+
"vid": "8YSha7iU2ZI-Scene-1025",
|
| 12 |
+
"prompt": "The video shows a red cartoon character with a long tongue and two antennae on its head. The character is seen talking to a yellow cartoon character with a single antenna on its head. The red character is very angry and is yelling at the yellow character. The background is a plain wall."
|
| 13 |
+
},
|
| 14 |
+
{
|
| 15 |
+
"vid": "97moBezrzLo-Scene-0076",
|
| 16 |
+
"prompt": "The video shows a group of children wearing green jerseys and black pants. They are standing on a grassy field, and some of them are wearing pink hats. The children are huddled together, and one of them is holding a pink object. The children are smiling and seem to be enjoying themselves."
|
| 17 |
+
},
|
| 18 |
+
{
|
| 19 |
+
"vid": "9GDE5-rOOfk-Scene-0010",
|
| 20 |
+
"prompt": "The video shows a baseball game in progress. A player is seen running towards a base and sliding into it. The player is wearing a white uniform with a blue cap. The field is green with white lines marking the base paths. The player successfully reaches the base before being tagged out by the opposing team's player. The opposing player is wearing a black uniform with a white cap. The umpire is standing behind the player in white, wearing a black uniform. The scoreboard is visible in the background, showing the score and inning of the game."
|
| 21 |
+
},
|
| 22 |
+
{
|
| 23 |
+
"vid": "DZt-RtpcJZQ-Scene-0043",
|
| 24 |
+
"prompt": "The video shows a man wearing a blue t-shirt with the word \"drumeo\" written on it. He is holding drumsticks in his hands and is speaking to the camera. The background is a dark blue color, and the man is wearing glasses. He appears to be in a studio setting, and he is the only person in the frame. The man is gesturing with his hands as he speaks, and he seems to be explaining something about drumming."
|
| 25 |
+
},
|
| 26 |
+
{
|
| 27 |
+
"vid": "FAfMSWi0FSA-Scene-0200",
|
| 28 |
+
"prompt": "In the video, a person is seen wearing a full-body Iron Man costume. The costume is predominantly red with gold accents, and the person is wearing a helmet with a face shield. The person is standing in a room with white walls and a wooden floor. They are dancing and moving their arms and legs in a rhythmic manner. There is a yellow picture frame on the wall, and a chandelier hanging from the ceiling. The person appears to be enjoying themselves and is the only person in the room."
|
| 29 |
+
},
|
| 30 |
+
{
|
| 31 |
+
"vid": "a5M8hWgQbjU-Scene-0191",
|
| 32 |
+
"prompt": "The video shows a boat with a white canopy and red trim floating on a body of water. The boat is surrounded by calm water and the sun is setting in the background, casting a warm glow on the scene. The sky is clear with a few clouds, and the sun is partially obscured by the trees on the horizon. The boat appears to be stationary, and there are no visible people or movement on the boat. The water is reflecting the sunlight, creating a shimmering effect."
|
| 33 |
+
},
|
| 34 |
+
{
|
| 35 |
+
"vid": "ijGuAu0lLNM-Scene-0008",
|
| 36 |
+
"prompt": "In the video, a woman is driving a car with a baby in the back seat. The baby is wearing a white hoodie and is sitting in a car seat. The woman is talking to the baby and the baby is waving at the camera. The woman is also talking to the camera."
|
| 37 |
+
},
|
| 38 |
+
{
|
| 39 |
+
"vid": "zFehFw9SbDw-Scene-0034",
|
| 40 |
+
"prompt": "In the video, a person is seen washing a black dog in a sink. The person is using their hands to scrub the dog's fur, and the dog appears to be enjoying the bath. The sink is filled with water, and there are various objects visible in the background, including a yellow container and a bowl of food. The person's movements are gentle and caring, and the dog seems to be relaxed and comfortable. The video captures a heartwarming moment of a pet being bathed with love and care."
|
| 41 |
+
}
|
| 42 |
+
|
| 43 |
+
|
| 44 |
+
]
|
patch_diffusers.py
ADDED
|
@@ -0,0 +1,219 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
# Copyright 2026 Pruna AI
|
| 2 |
+
# Portions adapted from Hugging Face diffusers:
|
| 3 |
+
# Copyright 2026 The HuggingFace Team. Licensed under Apache-2.0.
|
| 4 |
+
|
| 5 |
+
"""Monkey-patch stock diffusers so pruned LTX-2.3 decoder weights load correctly.
|
| 6 |
+
|
| 7 |
+
Stock diffusers builds the decoder graph from ``decoder_block_out_channels`` using
|
| 8 |
+
the *nominal* width of each up block (divided by ``upsample_factor``). It assumes the
|
| 9 |
+
tensor leaving block N always matches that nominal width.
|
| 10 |
+
|
| 11 |
+
PrunaVAED keeps wider skip connections between blocks and only prunes inside each
|
| 12 |
+
block. Example on ``up_blocks.2``:
|
| 13 |
+
|
| 14 |
+
- tensor arriving from ``up_blocks.1``: **384** ch (after 512→384 projection)
|
| 15 |
+
- ResNets inside ``up_blocks.2``: **128** ch (50% pruned vs LTX-2.3 decoder 256)
|
| 16 |
+
- ``conv_in`` projection inside the block: **384 → 256** (feeds the upsampler)
|
| 17 |
+
|
| 18 |
+
Stock diffusers wires ``up_blocks.2`` as 192 → 128 and never creates the 384→256
|
| 19 |
+
``conv_in``, so ``from_pretrained`` fails with shape mismatches on PrunaVAED weights.
|
| 20 |
+
|
| 21 |
+
Fix (two ``__init__`` replacements; forward/decode unchanged)
|
| 22 |
+
------------------------------------------------------------
|
| 23 |
+
1. ``LTX2VideoDecoder3d``: track ``current_channels`` — the actual tensor width
|
| 24 |
+
between blocks — instead of recomputing from the previous block's nominal width.
|
| 25 |
+
2. ``LTX2VideoUpBlock3d``: insert ``conv_in`` when ``in_channels != out_channels *
|
| 26 |
+
upscale_factor`` (pre-upsampler width), not when ``in_channels != out_channels``.
|
| 27 |
+
|
| 28 |
+
The LTX-2.3 decoder is unaffected (no extra projections). Safe to call before any
|
| 29 |
+
``AutoencoderKLLTX2Video.from_pretrained``.
|
| 30 |
+
"""
|
| 31 |
+
|
| 32 |
+
from __future__ import annotations
|
| 33 |
+
|
| 34 |
+
import torch
|
| 35 |
+
import torch.nn as nn
|
| 36 |
+
from diffusers.models.autoencoders import autoencoder_kl_ltx2 as ltx2
|
| 37 |
+
|
| 38 |
+
|
| 39 |
+
def patch_pruna_ltx2_decoder() -> None:
|
| 40 |
+
if getattr(ltx2, "_PRUNA_LTX2_DECODER_PATCH", False):
|
| 41 |
+
return
|
| 42 |
+
|
| 43 |
+
Resnet = ltx2.LTX2VideoResnetBlock3d
|
| 44 |
+
Upsampler = ltx2.LTX2VideoUpsampler3d
|
| 45 |
+
CausalConv = ltx2.LTX2VideoCausalConv3d
|
| 46 |
+
MidBlock = ltx2.LTX2VideoMidBlock3d
|
| 47 |
+
UpBlock = ltx2.LTX2VideoUpBlock3d
|
| 48 |
+
TimeEmb = ltx2.PixArtAlphaCombinedTimestepSizeEmbeddings
|
| 49 |
+
RMSNorm = ltx2.PerChannelRMSNorm
|
| 50 |
+
|
| 51 |
+
def up_init(
|
| 52 |
+
self,
|
| 53 |
+
in_channels,
|
| 54 |
+
out_channels=None,
|
| 55 |
+
num_layers=1,
|
| 56 |
+
dropout=0.0,
|
| 57 |
+
resnet_eps=1e-6,
|
| 58 |
+
resnet_act_fn="swish",
|
| 59 |
+
spatio_temporal_scale=True,
|
| 60 |
+
upsample_type="spatiotemporal",
|
| 61 |
+
inject_noise=False,
|
| 62 |
+
timestep_conditioning=False,
|
| 63 |
+
upsample_residual=False,
|
| 64 |
+
upscale_factor=1,
|
| 65 |
+
spatial_padding_mode="zeros",
|
| 66 |
+
):
|
| 67 |
+
nn.Module.__init__(self)
|
| 68 |
+
out_channels = out_channels or in_channels
|
| 69 |
+
# Width right before the upsampler (ResNet width × upscale_factor).
|
| 70 |
+
# conv_in projects the incoming skip (e.g. 384) down to this width (e.g. 256).
|
| 71 |
+
pre = out_channels * upscale_factor
|
| 72 |
+
self.time_embedder = TimeEmb(in_channels * 4, 0) if timestep_conditioning else None
|
| 73 |
+
self.conv_in = None
|
| 74 |
+
if in_channels != pre: # stock diffusers compares in_channels != out_channels
|
| 75 |
+
self.conv_in = Resnet(
|
| 76 |
+
in_channels=in_channels,
|
| 77 |
+
out_channels=pre,
|
| 78 |
+
dropout=dropout,
|
| 79 |
+
eps=resnet_eps,
|
| 80 |
+
non_linearity=resnet_act_fn,
|
| 81 |
+
inject_noise=inject_noise,
|
| 82 |
+
timestep_conditioning=timestep_conditioning,
|
| 83 |
+
spatial_padding_mode=spatial_padding_mode,
|
| 84 |
+
)
|
| 85 |
+
self.upsamplers = None
|
| 86 |
+
if spatio_temporal_scale:
|
| 87 |
+
stride = {
|
| 88 |
+
"spatial": (1, 2, 2),
|
| 89 |
+
"temporal": (2, 1, 1),
|
| 90 |
+
"spatiotemporal": (2, 2, 2),
|
| 91 |
+
}[upsample_type]
|
| 92 |
+
self.upsamplers = nn.ModuleList(
|
| 93 |
+
[
|
| 94 |
+
Upsampler(
|
| 95 |
+
in_channels=pre,
|
| 96 |
+
stride=stride,
|
| 97 |
+
residual=upsample_residual,
|
| 98 |
+
upscale_factor=upscale_factor,
|
| 99 |
+
spatial_padding_mode=spatial_padding_mode,
|
| 100 |
+
),
|
| 101 |
+
]
|
| 102 |
+
)
|
| 103 |
+
self.resnets = nn.ModuleList(
|
| 104 |
+
[
|
| 105 |
+
Resnet(
|
| 106 |
+
in_channels=out_channels,
|
| 107 |
+
out_channels=out_channels,
|
| 108 |
+
dropout=dropout,
|
| 109 |
+
eps=resnet_eps,
|
| 110 |
+
non_linearity=resnet_act_fn,
|
| 111 |
+
inject_noise=inject_noise,
|
| 112 |
+
timestep_conditioning=timestep_conditioning,
|
| 113 |
+
spatial_padding_mode=spatial_padding_mode,
|
| 114 |
+
)
|
| 115 |
+
for _ in range(num_layers)
|
| 116 |
+
]
|
| 117 |
+
)
|
| 118 |
+
self.gradient_checkpointing = False
|
| 119 |
+
|
| 120 |
+
def decoder_init(
|
| 121 |
+
self,
|
| 122 |
+
in_channels=128,
|
| 123 |
+
out_channels=3,
|
| 124 |
+
block_out_channels=(256, 512, 1024),
|
| 125 |
+
spatio_temporal_scaling=(True, True, True),
|
| 126 |
+
layers_per_block=(5, 5, 5, 5),
|
| 127 |
+
upsample_type=("spatiotemporal", "spatiotemporal", "spatiotemporal"),
|
| 128 |
+
patch_size=4,
|
| 129 |
+
patch_size_t=1,
|
| 130 |
+
resnet_norm_eps=1e-6,
|
| 131 |
+
is_causal=False,
|
| 132 |
+
inject_noise=(False, False, False),
|
| 133 |
+
timestep_conditioning=False,
|
| 134 |
+
upsample_residual=(True, True, True),
|
| 135 |
+
upsample_factor=(2, 2, 2),
|
| 136 |
+
spatial_padding_mode="reflect",
|
| 137 |
+
):
|
| 138 |
+
nn.Module.__init__(self)
|
| 139 |
+
n = len(layers_per_block)
|
| 140 |
+
if isinstance(spatio_temporal_scaling, bool):
|
| 141 |
+
spatio_temporal_scaling = (spatio_temporal_scaling,) * (n - 1)
|
| 142 |
+
if isinstance(inject_noise, bool):
|
| 143 |
+
inject_noise = (inject_noise,) * n
|
| 144 |
+
if isinstance(upsample_residual, bool):
|
| 145 |
+
upsample_residual = (upsample_residual,) * (n - 1)
|
| 146 |
+
|
| 147 |
+
self.patch_size, self.patch_size_t = patch_size, patch_size_t
|
| 148 |
+
self.out_channels = out_channels * patch_size**2
|
| 149 |
+
self.is_causal = is_causal
|
| 150 |
+
|
| 151 |
+
ch = tuple(reversed(block_out_channels))
|
| 152 |
+
spatio_temporal_scaling = tuple(reversed(spatio_temporal_scaling))
|
| 153 |
+
layers_per_block = tuple(reversed(layers_per_block))
|
| 154 |
+
inject_noise = tuple(reversed(inject_noise))
|
| 155 |
+
upsample_residual = tuple(reversed(upsample_residual))
|
| 156 |
+
upsample_factor = tuple(reversed(upsample_factor))
|
| 157 |
+
upsample_type = tuple(upsample_type)
|
| 158 |
+
if len(upsample_type) == len(ch) - 1:
|
| 159 |
+
upsample_type = tuple(reversed(upsample_type))
|
| 160 |
+
|
| 161 |
+
width = ch[0]
|
| 162 |
+
self.conv_in = CausalConv(
|
| 163 |
+
in_channels=in_channels,
|
| 164 |
+
out_channels=width,
|
| 165 |
+
kernel_size=3,
|
| 166 |
+
stride=1,
|
| 167 |
+
spatial_padding_mode=spatial_padding_mode,
|
| 168 |
+
)
|
| 169 |
+
self.mid_block = MidBlock(
|
| 170 |
+
in_channels=width,
|
| 171 |
+
num_layers=layers_per_block[0],
|
| 172 |
+
resnet_eps=resnet_norm_eps,
|
| 173 |
+
inject_noise=inject_noise[0],
|
| 174 |
+
timestep_conditioning=timestep_conditioning,
|
| 175 |
+
spatial_padding_mode=spatial_padding_mode,
|
| 176 |
+
)
|
| 177 |
+
|
| 178 |
+
# Stock: input_channel = previous_nominal // factor (loses pruned skip widths).
|
| 179 |
+
# Patched: pass the real tensor width from the previous block as in_channels.
|
| 180 |
+
self.up_blocks = nn.ModuleList()
|
| 181 |
+
current = width
|
| 182 |
+
for i in range(len(ch)):
|
| 183 |
+
resnet_w = ch[i] // upsample_factor[i]
|
| 184 |
+
self.up_blocks.append(
|
| 185 |
+
UpBlock(
|
| 186 |
+
in_channels=current,
|
| 187 |
+
out_channels=resnet_w,
|
| 188 |
+
num_layers=layers_per_block[i + 1],
|
| 189 |
+
resnet_eps=resnet_norm_eps,
|
| 190 |
+
spatio_temporal_scale=spatio_temporal_scaling[i],
|
| 191 |
+
upsample_type=upsample_type[i],
|
| 192 |
+
inject_noise=inject_noise[i + 1],
|
| 193 |
+
timestep_conditioning=timestep_conditioning,
|
| 194 |
+
upsample_residual=upsample_residual[i],
|
| 195 |
+
upscale_factor=upsample_factor[i],
|
| 196 |
+
spatial_padding_mode=spatial_padding_mode,
|
| 197 |
+
)
|
| 198 |
+
)
|
| 199 |
+
current = resnet_w
|
| 200 |
+
|
| 201 |
+
self.norm_out = RMSNorm()
|
| 202 |
+
self.conv_act = nn.SiLU()
|
| 203 |
+
self.conv_out = CausalConv(
|
| 204 |
+
in_channels=current,
|
| 205 |
+
out_channels=self.out_channels,
|
| 206 |
+
kernel_size=3,
|
| 207 |
+
stride=1,
|
| 208 |
+
spatial_padding_mode=spatial_padding_mode,
|
| 209 |
+
)
|
| 210 |
+
self.time_embedder = self.scale_shift_table = self.timestep_scale_multiplier = None
|
| 211 |
+
if timestep_conditioning:
|
| 212 |
+
self.timestep_scale_multiplier = nn.Parameter(torch.tensor(1000.0))
|
| 213 |
+
self.time_embedder = TimeEmb(width * 2, 0)
|
| 214 |
+
self.scale_shift_table = nn.Parameter(torch.randn(2, width) / width**0.5)
|
| 215 |
+
self.gradient_checkpointing = False
|
| 216 |
+
|
| 217 |
+
UpBlock.__init__ = up_init
|
| 218 |
+
ltx2.LTX2VideoDecoder3d.__init__ = decoder_init
|
| 219 |
+
ltx2._PRUNA_LTX2_DECODER_PATCH = True
|
requirements-demo.txt
ADDED
|
@@ -0,0 +1,19 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
--extra-index-url https://download.pytorch.org/whl/cu130
|
| 2 |
+
torch==2.11.0+cu130
|
| 3 |
+
torchvision==0.26.0+cu130
|
| 4 |
+
torchaudio==2.11.0+cu130
|
| 5 |
+
diffusers==0.39.0
|
| 6 |
+
huggingface_hub==1.24.0
|
| 7 |
+
safetensors==0.8.0
|
| 8 |
+
accelerate==1.13.0
|
| 9 |
+
imageio==2.37.3
|
| 10 |
+
imageio-ffmpeg==0.6.0
|
| 11 |
+
|
| 12 |
+
transformers>=4.57.0
|
| 13 |
+
sentencepiece>=0.2.0
|
| 14 |
+
protobuf>=5.0.0
|
| 15 |
+
Pillow>=10.0.0
|
| 16 |
+
numpy>=1.24.0
|
| 17 |
+
tqdm>=4.66.0
|
| 18 |
+
av>=14.0.0
|
| 19 |
+
pytest>=8.0.0
|
vae/config.json
ADDED
|
@@ -0,0 +1,93 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
{
|
| 2 |
+
"_class_name": "AutoencoderKLLTX2Video",
|
| 3 |
+
"_diffusers_version": "0.38.0",
|
| 4 |
+
"block_out_channels": [
|
| 5 |
+
256,
|
| 6 |
+
512,
|
| 7 |
+
1024,
|
| 8 |
+
1024
|
| 9 |
+
],
|
| 10 |
+
"decoder_block_out_channels": [
|
| 11 |
+
128,
|
| 12 |
+
256,
|
| 13 |
+
384,
|
| 14 |
+
1024
|
| 15 |
+
],
|
| 16 |
+
"decoder_causal": false,
|
| 17 |
+
"decoder_inject_noise": [
|
| 18 |
+
false,
|
| 19 |
+
false,
|
| 20 |
+
false,
|
| 21 |
+
false,
|
| 22 |
+
false
|
| 23 |
+
],
|
| 24 |
+
"decoder_layers_per_block": [
|
| 25 |
+
4,
|
| 26 |
+
6,
|
| 27 |
+
4,
|
| 28 |
+
2,
|
| 29 |
+
2
|
| 30 |
+
],
|
| 31 |
+
"decoder_spatial_padding_mode": "zeros",
|
| 32 |
+
"decoder_spatio_temporal_scaling": [
|
| 33 |
+
true,
|
| 34 |
+
true,
|
| 35 |
+
true,
|
| 36 |
+
true
|
| 37 |
+
],
|
| 38 |
+
"down_block_types": [
|
| 39 |
+
"LTX2VideoDownBlock3D",
|
| 40 |
+
"LTX2VideoDownBlock3D",
|
| 41 |
+
"LTX2VideoDownBlock3D",
|
| 42 |
+
"LTX2VideoDownBlock3D"
|
| 43 |
+
],
|
| 44 |
+
"downsample_type": [
|
| 45 |
+
"spatial",
|
| 46 |
+
"temporal",
|
| 47 |
+
"spatiotemporal",
|
| 48 |
+
"spatiotemporal"
|
| 49 |
+
],
|
| 50 |
+
"encoder_causal": true,
|
| 51 |
+
"encoder_spatial_padding_mode": "zeros",
|
| 52 |
+
"in_channels": 3,
|
| 53 |
+
"latent_channels": 128,
|
| 54 |
+
"layers_per_block": [
|
| 55 |
+
4,
|
| 56 |
+
6,
|
| 57 |
+
4,
|
| 58 |
+
2,
|
| 59 |
+
2
|
| 60 |
+
],
|
| 61 |
+
"out_channels": 3,
|
| 62 |
+
"patch_size": 4,
|
| 63 |
+
"patch_size_t": 1,
|
| 64 |
+
"resnet_norm_eps": 1e-06,
|
| 65 |
+
"scaling_factor": 1.0,
|
| 66 |
+
"spatial_compression_ratio": 32,
|
| 67 |
+
"spatio_temporal_scaling": [
|
| 68 |
+
true,
|
| 69 |
+
true,
|
| 70 |
+
true,
|
| 71 |
+
true
|
| 72 |
+
],
|
| 73 |
+
"temporal_compression_ratio": 8,
|
| 74 |
+
"timestep_conditioning": false,
|
| 75 |
+
"upsample_factor": [
|
| 76 |
+
2,
|
| 77 |
+
2,
|
| 78 |
+
1,
|
| 79 |
+
2
|
| 80 |
+
],
|
| 81 |
+
"upsample_residual": [
|
| 82 |
+
false,
|
| 83 |
+
false,
|
| 84 |
+
false,
|
| 85 |
+
false
|
| 86 |
+
],
|
| 87 |
+
"upsample_type": [
|
| 88 |
+
"spatiotemporal",
|
| 89 |
+
"spatiotemporal",
|
| 90 |
+
"temporal",
|
| 91 |
+
"spatial"
|
| 92 |
+
]
|
| 93 |
+
}
|
vae/diffusion_pytorch_model.safetensors
ADDED
|
@@ -0,0 +1,3 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
version https://git-lfs.github.com/spec/v1
|
| 2 |
+
oid sha256:4cbf0cbe6c185514d62c6c58c35dc42d7ea15924f34391e08be12f44bfccdf1d
|
| 3 |
+
size 1327909418
|